"""Public transfer checks. These are exposed practice, not a hidden exam or authorship check.
Usage: python3 planner_checks.py planner_exercise.py
The selected local Python file executes with normal Python privileges.
"""
import copy
import importlib.util
import itertools
import math
import random
import sys
import unittest

spec = importlib.util.spec_from_file_location('learner', sys.argv.pop(1) if len(sys.argv)>1 else 'planner_exercise.py')
learner = importlib.util.module_from_spec(spec)
spec.loader.exec_module(learner)


def path_oracle(graph, start, goal):
    """Independent exhaustive simple-path oracle for small nonnegative-cost graphs."""
    costs = []
    def visit(node, used, cost):
        if node == goal:
            costs.append(cost)
            return
        for dest, edge in graph[node]:
            if dest not in used:
                visit(dest, used | {dest}, cost + edge)
    visit(start, {start}, 0)
    return min(costs) if costs else None


class TransferChecks(unittest.TestCase):
    def test_reopening(self):
        graph = {'s':[('a',3),('b',1)], 'a':[('g',2)], 'b':[('a',1),('g',8)], 'g':[]}
        result = learner.astar(graph, {'s':0,'a':0,'b':3,'g':0}, 's','g')
        self.assertEqual(result['route'], ['s','b','a','g'])
        self.assertEqual(result['cost'], 4)
        self.assertEqual(result['reopens'], 1)

    def test_distinct_cutoff_and_impossibility(self):
        graph = {'s':[('a',1)], 'a':[('g',1)], 'g':[]}
        self.assertEqual(learner.astar(graph, dict.fromkeys(graph,0),'s','g',1)['status'],'budget')
        self.assertEqual(learner.astar(graph, dict.fromkeys(graph,0),'s','g',2)['status'],'found')
        graph['a']=[]
        self.assertEqual(learner.astar(graph, dict.fromkeys(graph,0),'s','g',2)['status'],'exhausted')

    def test_generated_graphs_against_enumeration(self):
        rng=random.Random(73)
        for _ in range(25):
            names=['s','a','b','c','g']
            graph={a:[(b,rng.randrange(0,7)) for b in names if a!=b and rng.random()<.35] for a in names}
            saved=copy.deepcopy(graph)
            result=learner.astar(graph,dict.fromkeys(graph,0),'s','g')
            self.assertEqual(result['cost'],path_oracle(graph,'s','g'))
            self.assertEqual(graph,saved)
            if result['status']=='found':
                self.assertEqual(result['route'][0],'s'); self.assertEqual(result['route'][-1],'g')
                cost=0
                for a,b in zip(result['route'],result['route'][1:]):
                    cost+=next(w for dest,w in graph[a] if dest==b)
                self.assertEqual(cost,result['cost'])

    def test_filter_against_latent_path_enumeration(self):
        prior=[.6,.4]; matrix=[[.75,.25],[.25,.75]]; readings=[[.8,.2],[1,1],[.3,.7]]
        original=copy.deepcopy((prior,matrix,readings))
        run=learner.filter_belief(prior,matrix,readings)
        for count,step in enumerate(run['steps'],1):
            masses=[0.,0.]
            for path in itertools.product(range(2),repeat=count+1):
                weight=prior[path[0]]
                for t in range(count):
                    weight*=matrix[path[t]][path[t+1]]*readings[t][path[t+1]]
                masses[path[-1]]+=weight
            for j in range(2):
                self.assertAlmostEqual(step['posterior'][j],masses[j]/sum(masses),places=12)
        self.assertAlmostEqual(run['logEvidence'],math.log(sum(masses)),places=12)
        self.assertEqual((prior,matrix,readings),original)

    def test_missing_and_impossible(self):
        missing=learner.filter_belief([1,0],[[.8,.2],[.2,.8]],[[1,1]])
        self.assertEqual(missing['steps'][0]['posterior'],[.8,.2])
        result=learner.filter_belief([1,0],[[1,0],[0,1]],[[0,1],[1,1]])
        self.assertEqual(result['status'],'impossible'); self.assertIsNone(result['logEvidence'])
        self.assertEqual(len(result['steps']),1); self.assertIsNone(result['steps'][0]['posterior'])

    def test_invalid_contracts(self):
        with self.assertRaises(ValueError):
            learner.astar({'s':[('g',-1)],'g':[]},{'s':0,'g':0},'s','g')
        with self.assertRaises(ValueError):
            learner.filter_belief([.7,.7],[[1,0],[0,1]],[])
        with self.assertRaises(ValueError):
            learner.filter_belief([1,0],[[1,0],[0,1]],[[1,1]]*17)


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