import unittest
from fractions import Fraction as Q
import numpy as np
from mpmath import mp
from example import (Parameters,DrivenReactor,ObservationPolicy,OperatingMonitor,
    ProbabilityBudget,ProductionLedger,algebraic_certificate,A_WEIGHTS,B_WEIGHTS,Y_WEIGHTS,M_WEIGHTS,S_WEIGHTS,dot)


class ScientificChecks(unittest.TestCase):
    def test_literal_channels_marks_and_falling_factorial(self):
        on=DrivenReactor();off=DrivenReactor(enabled=False)
        self.assertEqual([c.label for c in on.channels],list(range(20)))
        self.assertEqual(set(on.by_label)-set(off.by_label),{6,7})
        for label in off.by_label:self.assertEqual(on.by_label[label],off.by_label[label])
        counts=[100,101,2,3,4,5];V=100
        self.assertEqual(on.by_label[9].count_rate(counts,V),Q(2,5))
        for x in (0,1):
            n=counts.copy();n[2]=x;self.assertEqual(on.by_label[9].count_rate(n,V),0)
        self.assertEqual(on.by_label[18].jump,on.by_label[1].jump)
        self.assertNotEqual(on.by_label[18].marks,on.by_label[1].marks)
        self.assertEqual(on.by_label[19].jump,on.by_label[0].jump)
        for c in on.channels:
            self.assertEqual(dot(M_WEIGHTS,c.jump)+c.marks[7],2*(c.marks[1]+c.marks[2]))
        cert=algebraic_certificate();self.assertTrue(all(cert['source_identities'].values()))
        self.assertEqual(cert['entry_exponent'],'8047/312500000000')
        self.assertEqual(cert['window_exponent'],'13/1080000')
        self.assertGreater(Q(cert['disabled_chernoff_slack']),Q(5,1000000))

    def test_local_growth_noise_box_and_low_counts(self):
        rng=np.random.default_rng(2501);V=100000000
        for K,r,d in ((8,18,'1/20'),(12,22,'1/100'),(10,20,'1/20')):
            reactor=DrivenReactor(Parameters(K,r,d))
            for _ in range(15):
                catalytic=[int(x) for x in rng.integers(0,10000,4)]
                counts=[0,0,*catalytic]
                counts[0]=int(Q(9,10)*V)-dot(A_WEIGHTS,counts)
                counts[1]=int(Q(9,10)*V)-dot(B_WEIGHTS,counts)
                Y=dot(Y_WEIGHTS,counts)
                self.assertLessEqual(Y,Q(V,1000))
                drift,noise=reactor.generator(counts,V,Y_WEIGHTS)
                self.assertGreaterEqual(drift,Q(39,100)*Y+Q(16,25)*reactor.p.epsilon*V)
                self.assertLessEqual(noise,5*Y+Q(25,16)*reactor.p.epsilon*V)
                self.assertLess(sum(c.count_rate(counts,V) for c in reactor.channels),3000*V)

    def test_deadline_endpoint_return_and_continuing_windows(self):
        reactor=DrivenReactor();policy=ObservationPolicy(5000,'3/2');monitor=OperatingMonitor(reactor,policy)
        monitor.event(499.8,0);monitor.event(500,0)
        self.assertEqual(monitor.entry_time,500);self.assertEqual(monitor.phase,'entered')
        monitor.event(500.1,0);monitor.event(501,14)
        self.assertEqual(monitor.completed,1);self.assertEqual(monitor.window_exports,[1])
        self.assertEqual(monitor.status,'active') # Reaching export cap did not stop chemistry.
        monitor.event(501.1,18) # X falls to exactly the residence floor in the final fraction.
        self.assertEqual(monitor.status,'return_failure');self.assertFalse(monitor.report()['event_verdict'])
        monitor.audit_ledger()
        missed=OperatingMonitor(reactor,ObservationPolicy(5000,1))
        missed.event(499.8,0);missed.event(500.1,0)
        self.assertEqual(missed.status,'missed_deadline');self.assertIsNone(missed.entry_time)
        history=OperatingMonitor(reactor,ObservationPolicy(5000,2))
        for i in range(5):history.event(i+1,0)
        history.advance(501) # First window fails, but the continuing chemical state is retained.
        self.assertEqual(history.status,'active');self.assertEqual(history.counts[2],5)
        history.event(502,14)
        self.assertEqual(history.window_exports,[0,1]);self.assertFalse(history.report()['event_verdict'])

    def test_supply_saturation_and_disabled_credit(self):
        policy=ObservationPolicy(5000,'3/2');cap=policy.supply_cap;threshold=policy.V*policy.end/8
        for a,b in ((0,0),(cap+1,0),(cap-1,cap+7),(1,int(threshold)),(int(threshold),0)):
            self.assertEqual(a+b>threshold,min(a,cap)+min(b,cap)>threshold)
        off=OperatingMonitor(DrivenReactor(enabled=False),ObservationPolicy(5000,1))
        for i in range(501):off.event((i+1)/1000,10)
        self.assertEqual(off.status,'resource_exit');self.assertTrue(off.report()['event_verdict'])
        self.assertEqual(off.counts[0],5501) # First exiting state is retained, never clipped.
        off.audit_ledger()
        path=DrivenReactor().stochastic(ObservationPolicy(100000000,1),limit=5,seed=1)
        self.assertEqual(path['status'],'unfinished');self.assertIsNone(path['event_verdict'])

    def test_interval_budget_sizing_and_fractional_horizon_scope(self):
        b=ProbabilityBudget();result=b.evaluate()
        self.assertLess(abs(mp.mpf(result['enabled_success_lower_ordinary'])-mp.mpf('0.9238492232933793636')),mp.mpf('1e-18'))
        self.assertTrue(b.interval_pass('0.0762'));self.assertFalse(b.interval_pass('0.01'))
        self.assertTrue(result['disabled_below_10_power_minus_21699'])
        sizes=b.sufficient_scale(100,'0.01')
        self.assertTrue(ProbabilityBudget(sizes['individually_interval_checked_V'],100).interval_pass('0.01'))
        self.assertEqual(sizes['simple_sufficient_V'],264915869)
        self.assertFalse(sizes['minimum_claimed'])
        fractional=ProbabilityBudget(100000000,'3/2').evaluate()
        self.assertEqual(fractional['complete_windows'],1);self.assertIsNone(fractional['disabled_log10_upper'])
        outside=ProbabilityBudget(100000000,'1e50').evaluate()
        self.assertFalse(outside['simplified_horizon_verified']);self.assertIsNone(outside['simplified_success_lower'])
        with self.assertRaises(ValueError):ProbabilityBudget(99999999,100)
        with self.assertRaises(ValueError):ProbabilityBudget(parameters=Parameters(delta='1/10'))

    def test_deterministic_ledger_and_independent_solver(self):
        reactor=DrivenReactor();a=reactor.deterministic(1);b=reactor.deterministic(1,method='BDF')
        times=np.linspace(0,501,101);x=a.sol(times).T;y=b.sol(times).T
        np.testing.assert_allclose(x,y,rtol=1e-6,atol=2e-8)
        np.testing.assert_allclose(x[:,:6]@A_WEIGHTS+x[:,11],1+x[:,7],atol=2e-9)
        np.testing.assert_allclose(x[:,:6]@B_WEIGHTS+x[:,12],1+x[:,8],atol=2e-9)
        np.testing.assert_allclose(x[:,:6]@M_WEIGHTS+x[:,13],4+2*(x[:,7]+x[:,8]),atol=5e-9)
        self.assertGreater(a.sol(501)[6]-a.sol(500)[6],1/5000)
        off=DrivenReactor(enabled=False).deterministic(1)
        self.assertLess(off.sol(501)[6]-off.sol(500)[6],1/5000)

    def test_output_geometry_units_and_gross_service(self):
        ledger=ProductionLedger(ObservationPolicy())
        report=ledger.guarantees();physical=ledger.physical()
        self.assertEqual(report['per_window_export_mass'],'20000')
        self.assertEqual(report['covalent_equivalents_per_window'],'5000')
        self.assertEqual(report['each_food_budget'],'120000000000')
        self.assertEqual(report['gross_drive_budget'],'7500000000')
        self.assertAlmostEqual(float(physical['volume_pL']),166.0539067,places=5)
        self.assertEqual(mp.mpf(physical['covalent_equivalent_rate_pM_per_hour']),50)
        self.assertEqual(ledger.aligned_interval('1.99'),0)
        self.assertEqual(ledger.aligned_interval(2),20000)
        self.assertEqual(ledger.aligned_interval('5.5'),80000)
        # Opposite driven jumps restore counts but still consume two gross service events.
        monitor=OperatingMonitor(DrivenReactor(),ObservationPolicy(5000,1))
        monitor.event(.1,0);before=monitor.counts.copy()
        monitor.event(.2,18);monitor.event(.3,19)
        self.assertEqual(monitor.counts,before)
        self.assertEqual(monitor.ledger[3:5],[1,1]);monitor.audit_ledger()


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