"""Scientific identities, exact pulse law, unconditioned histories and operating arithmetic."""
import unittest
from fractions import Fraction as Q
import math
import numpy as np
from mpmath import mp
from example import (Reactor,Intervention,MolecularPulse,PulseOutcome,CycleLedger,StochasticCycle,
    HistoryMission,ReturnedStockController,DensityExperiment,MissionCertificate,PhaseTransport,
    A,B,I,Y,V0,INITIAL_DENSITY,dot,restart,stopping_boundary_example)

class ScienceTests(unittest.TestCase):
    def test_source_generator_and_phase(self):
        for r,d in [(19,'1/50'),(21,'1/25')]:self.assertTrue(Reactor(r,d).symbolic_identities()['identities_exact'])
        model=Reactor();self.assertEqual(model.channels[9].propensity((0,0,1,0,0,0),100),0)
        self.assertEqual(model.channels[9].propensity((0,0,2,0,0,0),100),Q(2,5))
        self.assertTrue(PhaseTransport().certificates()['symbolic_recurrence'])
        for V in (1,100,10**6):
            for n in (89*V,90*V,91*V):
                weights=PhaseTransport().weights(V,n)
                self.assertLessEqual(sum(weights),1)
                for w,target in zip(weights,Y[2:]):self.assertGreaterEqual(w,mp.mpf(target.numerator)/target.denominator/35)

    def test_exact_molecular_pulse_law(self):
        N=(1,1,2,1,0,0);V=3;p=Intervention('1/4');law=list(MolecularPulse().enumerate(N,V,p))
        self.assertEqual(sum(mass for _,mass in law),1)
        self.assertEqual(sum(mass*dot(Y,out.retained) for out,mass in law),sum(Q(n)*p.probabilities(i)[0]*Y[i] for i,n in enumerate(N)))
        extinct=sum(mass for out,mass in law if dot(I,out.retained)==0)
        expected=math.prod((1-p.probabilities(i)[0])**n for i,n in enumerate(N) if I[i])
        self.assertEqual(extinct,expected);self.assertGreater(extinct,0)
        for out,mass in law:
            self.assertEqual(tuple(a+b+c for a,b,c in zip(out.retained,out.withdrawn,out.lost)),N)
            self.assertEqual(out.start[2:],out.retained[2:])
        self.assertEqual(p.doses(1000),(745,755))
        with self.assertRaises(ValueError):Intervention('1/5')
        with self.assertRaises(ValueError):Intervention(survival=('97/100',)*6)

    def test_inventory_marks_and_collection_boundary(self):
        original=(4,4,3,0,0,1);pulse=PulseOutcome((3,3,1,0,0,1),(1,0,1,0,0,0),(0,1,1,0,0,0),(2,2));ledger=CycleLedger(original,pulse,100);model=Reactor()
        for t,label in [(.1,8),(.2,18),(.3,19),(.4,2),(.5,4),(.6,6),(.7,8),(2.99,14),(3.,14),(3.1,10),(3.2,11)]:ledger.event(t,model.channels[label])
        ledger.audit();v=ledger.values
        self.assertEqual(v[:5],[1,1,3,3,2]);self.assertEqual(v[5],2)
        self.assertEqual(v[6],1);self.assertEqual(ledger.audit(),1)
        incomplete=ledger.record(False,3.2,11);self.assertIsNone(incomplete['success'])

    def test_history_keeps_failures_actual_returns_and_budget(self):
        boundary=stopping_boundary_example()
        self.assertEqual(boundary['physical_closed'],'1/2');self.assertEqual(boundary['stopped_closed'],'1')
        self.assertEqual(boundary['physical_strict'],boundary['stopped_strict']);self.assertEqual(boundary['physical_strict'],'1/4')
        class ScriptedCycle:
            def __init__(self):self.starts=[]
            def run(self,N,V,p,rng,event_budget):
                self.starts.append(N);endpoint=list(N);endpoint[2]+=1
                return dict(endpoint=endpoint,complete=True,events=1,success=False,counters={'collected_X':0})
        cycle=ScriptedCycle();seen=[]
        def controller(history):
            seen.append(len(history))
            if history:history[0]['endpoint'][2]=999999  # Cannot mutate stored history.
            return Intervention()
        result=HistoryMission(cycle).run((0,0,100,0,0,0),100,3,controller,17,10)
        self.assertEqual(seen,[0,1,2]);self.assertFalse(result['all_success']);self.assertEqual([n[2] for n in cycle.starts],[100,101,102]);self.assertEqual(result['history'][0]['endpoint'][2],101)
        actual=HistoryMission(StochasticCycle(Reactor())).run((0,0,100,0,0,0),100,2,ReturnedStockController(),17,5)
        self.assertFalse(actual['complete']);self.assertIsNone(actual['all_success']);self.assertIsNone(actual['history'][0]['success']);self.assertEqual(actual['total_events'],5)

    def test_full_budget_against_nested_formula_and_repetition(self):
        for V in (10**6,10**9,V0):
            v=mp.mpf(V);E=lambda denominator:mp.exp(-v/denominator)
            material=4*E(2000)+48000*v*mp.exp(-v/2000+mp.mpf(1)/50)
            recovery=mp.exp(-3*v/5000000)+E(2)+material+E(10000)+12000*v*mp.exp(-v/10000+mp.mpf(9)/500)
            free=E(40000000)+100*E(10**10)+E(1)+E(200)
            joint=free+E(100000)+2*E(300)+E(2000)
            terminal=material+recovery+4*(E(320000)+E(1))
            expected=mp.exp(-1177*v/10**9)+4*E(100000)+joint+terminal
            cert=MissionCertificate(V,Reactor());self.assertLess(abs(mp.exp(cert.log_error())/expected-1),mp.mpf('1e-60'))
        cert=MissionCertificate(V0,Reactor())
        self.assertGreaterEqual(Q(cert.evaluate(48000)['joint_product_lower_certified']),Q(99,100))
        self.assertGreaterEqual(Q(cert.evaluate(100)['joint_product_lower_certified']),Q(999979,10**6))
        self.assertEqual(cert.evaluate(0)['joint_product_lower_certified'],'1.000000000')
        self.assertEqual(MissionCertificate(10**6,Reactor()).evaluate(2)['joint_product_lower_certified'],'0.000000000')
        with self.assertRaises(ValueError):MissionCertificate(V0,Reactor(release=22))

    def test_sizing_rounding_and_net_synthesis(self):
        cert=MissionCertificate(V0,Reactor());rich=tuple(int(Q(c)*V0) for c in INITIAL_DENSITY)
        self.assertTrue(restart(rich,V0));self.assertTrue(restart((0,0,V0,0,0,0),V0))
        totals=cert.inventory(100);self.assertEqual(totals['collected_I'],357142857200);self.assertEqual(totals['collected_X'],18518518600)
        self.assertEqual(totals['net_synthesis'],163035714343);self.assertEqual(cert.inventory(100,rich)['net_synthesis'],351485714343)
        self.assertGreater(cert.inventory(55)['net_synthesis'],0);self.assertLessEqual(cert.inventory(54)['net_synthesis'],0);self.assertGreater(cert.inventory(2,rich)['net_synthesis'],0)
        for m in (100,10000,1000000):
            size=cert.sizing(m,'1/100');self.assertGreaterEqual(Q(size['taylor_sum']),Q(size['taylor_target']))
            self.assertGreaterEqual(Q(MissionCertificate(size['sufficient_V'],Reactor()).evaluate(m)['joint_product_lower_certified']),Q(99,100))
        demanded=cert.sizing(100,'1/100',free_demand=10**13)
        self.assertGreaterEqual(MissionCertificate(demanded['sufficient_V'],Reactor()).inventory(100)['collected_X'],10**13)

    def test_density_and_completed_small_trace(self):
        reactor=Reactor();a=DensityExperiment(reactor).run(INITIAL_DENSITY,2);b=DensityExperiment(reactor).run(INITIAL_DENSITY,2,'BDF')
        self.assertLess(np.max(abs(np.array(a['trajectory'])-np.array(b['trajectory']))),2e-7)
        self.assertLess(max(abs(r['inventory_residual']) for r in a['history']),1e-8)
        for record in a['history']:
            self.assertGreaterEqual(record['counters']['collected_I'],record['counters']['collected_X'])
        result=HistoryMission(StochasticCycle(reactor)).run((0,0,20,0,0,0),20,2,ReturnedStockController(),43,30000)
        self.assertTrue(result['complete']);self.assertEqual(result['history'][1]['start'],result['history'][0]['endpoint'])
        for r in result['history']:
            c=r['counters'];self.assertEqual(c['net_synthesis'],dot(I,r['endpoint'])-dot(I,r['start'])+c['all_wash_I']+r['withdrawn_I']+r['lost_I'])

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