"""Public practice checks. Uses independent table receipts and finite differences.
Passing these exposed tests is not independent implementation review.
Run: python3 reward_checks.py reward_reference.py (or reward_exercise.py).
"""
import copy
import importlib.util
import math
from pathlib import Path
import sys
import unittest

path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name('reward_exercise.py')
spec = importlib.util.spec_from_file_location('learner', path)
learner = importlib.util.module_from_spec(spec)
spec.loader.exec_module(learner)
sys.argv = [sys.argv[0]]
EDGES = [[0,1,3,0],[1,2,4,0],[2,2,5,1],[0,4,6,3],[1,5,7,3],[],[3,7,6,6],[4,8,7,6],[]]

def receipt(state, action):
    nxt = EDGES[state][action]
    return dict(state=state,action=action,next=nxt,reward=10 if nxt==8 else -10 if nxt==5 else -1,terminal=nxt in (5,8))

def weights(shift):
    return {k: [math.sin(i*.7+shift)*.2 for i in range(n)] for k,n in [('w1',72),('b1',8),('w2',32),('b2',4)]}

def prediction(n, state, action):
    return n['b2'][action]+sum(n['w2'][action*8+j]*math.tanh(n['w1'][j*9+state]+n['b1'][j]) for j in range(8))

def objective(n, target, batch):
    total = 0.
    for t in batch:
        goal = t['reward']
        if not t['terminal']:
            goal += .9*max(prediction(target,t['next'],a) for a in range(4))
        total += (prediction(n,t['state'],t['action'])-goal)**2/2
    return total/len(batch)

class PublicChecks(unittest.TestCase):
    def test_every_transition_and_terminal(self):
        for s in [0,1,2,3,4,6,7]:
            for a in range(4):
                self.assertEqual(learner.transition(s,a),receipt(s,a))
        for s,a in [(5,0),(8,1),(-1,0),(0,4),(1.5,1),(True,0)]:
            with self.assertRaises(ValueError): learner.transition(s,a)

    def test_stopped_targets_terminal_and_timeout(self):
        self.assertAlmostEqual(learner.stopped_target(receipt(4,1),[100,200,300,400]),-10)
        self.assertAlmostEqual(learner.stopped_target(receipt(7,1),[100,200,300,400]),10)
        # This boundary bump remains nonterminal even if a collection timer expires.
        self.assertAlmostEqual(learner.stopped_target(receipt(0,0),[2,3,1,0]),-1+.9*3)

    def test_all_parameter_derivatives_and_immutable_inputs(self):
        online,target=weights(.2),weights(.8)
        batch=[receipt(0,1),receipt(4,1),receipt(7,1),receipt(0,0),receipt(0,1)]
        before=copy.deepcopy((online,target,batch))
        g,loss=learner.gradient(online,target,batch)
        self.assertAlmostEqual(loss,objective(online,target,batch),places=11)
        h=1e-5
        for key,values in online.items():
            self.assertEqual(len(g[key]),len(values))
            for index in range(len(values)):
                plus,minus=copy.deepcopy(online),copy.deepcopy(online)
                plus[key][index]+=h;minus[key][index]-=h
                numerical=(objective(plus,target,batch)-objective(minus,target,batch))/(2*h)
                self.assertTrue(math.isfinite(g[key][index]))
                self.assertAlmostEqual(g[key][index],numerical,delta=2e-8)
        self.assertEqual((online,target,batch),before)

    def test_duplicate_receipts_average_instead_of_sum(self):
        n,t=weights(.1),weights(.6)
        one=learner.gradient(n,t,[receipt(3,2)])
        repeated=learner.gradient(n,t,[receipt(3,2)]*8)
        self.assertAlmostEqual(one[1],repeated[1],places=12)
        for key in n:
            for a,b in zip(one[0][key],repeated[0][key]): self.assertAlmostEqual(a,b,places=12)

    def test_gradient_rejects_forged_receipts_and_bounds(self):
        n,t=weights(.1),weights(.6)
        bad=receipt(0,0);bad['next']=8
        for batch in [[],[receipt(0,0)]*9,[bad]]:
            with self.assertRaises(ValueError): learner.gradient(n,t,batch)

if __name__=='__main__': unittest.main()