import json,unittest
from pathlib import Path
from fractions import Fraction as F
from dataclasses import replace
import numpy as np
from population import *
from certificates import *
from example import phases_for,replay_sites


class ScientificChecks(unittest.TestCase):
    def test_conservative_source_and_correlated_inheritance(self):
        s=InheritedPopulation();Q=s.generator()
        self.assertEqual(s.states,((0,0),(0,1),(0,2),(1,0),(1,1),(2,0)))
        self.assertTrue(all(sum(row)==0 for row in Q));self.assertEqual(s.phi([F(1)]*6),[0]*6)
        for i,row in enumerate(s.pairs):
            self.assertEqual(sum(p for j,k,p in row),1)
            for j,k,p in row:self.assertEqual(tuple(a+b for a,b in zip(s.states[j],s.states[k])),s.states[i])
        x=list(map(F,['.1','.2','.3','.4','.5','.6']));pair=s.daughter(x,x)[5]
        marginal=sum(p*x[j] for j,k,p in s.pairs[5]);self.assertNotEqual(pair,marginal**2)
        # D Phi(1) is A, not -A (the prose sign in section 2.3 is a typo).
        epsilon=F(1,1000000);A=s.mean_matrix()
        for j in range(6):
            plus=[F(1)]*6;minus=plus.copy();plus[j]+=epsilon;minus[j]-=epsilon
            derivative=[(a-b)/(2*epsilon) for a,b in zip(s.phi(plus),s.phi(minus))]
            self.assertEqual(derivative,[A[i][j] for i in range(6)])

    def test_exact_barriers_and_rejection_on_changed_source(self):
        s=InheritedPopulation();q=tuple(map(F,['.94','.982','.987','.27','.735','.235']))
        c=ExposureCertificate(q,F(67,73));self.assertEqual(min(c.verify(s)['control_slack']),0)
        identity=population_barrier_generator_identity(s,q,[1,2,0,3,1,2],Action(F(29,100)))
        self.assertEqual(identity['direct'],identity['factored'])
        with self.assertRaises(ValueError):c.verify(InheritedPopulation(MolecularRates(protected_death=F(1))))
        with self.assertRaises(ValueError):AmplitudeCertificate(q,F(29,100)).verify(s)
        self.assertAlmostEqual(c.necessary_exposure(1,.01,5),4.725704981,places=5)
        self.assertGreater(c.survival_floor([0,0,0,0,0,100],3,s),c.survival_floor([0,0,0,0,0,1],3,s))

    def test_memory_size_certificates_and_repaired_endpoint_weights(self):
        data=json.loads(Path('site_certificates.json').read_text());r=replay_sites(data)
        self.assertEqual(r['amplitude']['N6_vmax29/100']['fully_marked_floor'],F(45962721,10**8))
        for row in r['mean_sign_endpoints'].values():
            self.assertGreater(row['lower_collatz'][0],0);self.assertLess(row['upper_collatz'][1],0)
        self.assertLess(r['spectral_at_point']['4']['bounds'][1],0)
        self.assertGreater(r['spectral_at_point']['6']['bounds'][0],0)

    def test_chronological_schedule_and_endpoint(self):
        s=InheritedPopulation();ph=phases_for(F(58,5),F(29,100),100,'early')
        self.assertEqual(sum(p.duration*p.action.eraser for p in ph),F(58,5))
        direct=s.numerical_flow(s.numerical_flow(np.zeros(6),ph[1]),ph[0])
        np.testing.assert_allclose(s.schedule(ph),direct,atol=0)
        self.assertGreater(s.schedule(ph,np.full(6,.1))[5],s.schedule(ph)[5])
        with self.assertRaises(ValueError):phases_for(30,F(29,100),100,'early')
        with self.assertRaises(ValueError):population_survival([.5],[-1])

    def test_fresh_short_interval_integrator_and_analytic_counterexample(self):
        from validated_pgf import certify
        s=InheritedPopulation();r=certify([(F(3,10),F(1))]);u=s.schedule([Phase(1,Action(F(29,100)))])
        for (a,b),x in zip(r['extinction_bounds'],u):
            # Floating solve is only cross-checking a narrow interval, not its proof.
            self.assertLess(abs(float((F(a)+F(b))/2)-x),1e-10)
        lo=s.taylor_coefficients(Action());hi=s.taylor_coefficients(Action(F(29,100)))
        self.assertEqual(hi[4][2]-lo[4][2],-F(49619,600000000))
        self.assertTrue(all(hi[k][2]==lo[k][2] for k in range(4)))

    def test_dimension_free_second_actuator_and_duration(self):
        for N in [1,2,6,8]:
            s=InheritedPopulation(MolecularRates(sites=N));w=WeightedDrift((F(1),)*s.size,F(1,5),Action(F(0),F(29,100)))
            self.assertEqual(w.verify(s),[0]*s.size)
        self.assertEqual(finite_time_floor(1,0,.3),1)
        self.assertEqual(finite_time_floor(0,10,.3),0)
        self.assertEqual(concentration_lower_bound(29,100,.29,1)['status'],'no_finite_AUC_can_meet_necessary_exposure')

    def test_feedback_pathwise_budget_and_resource_cap(self):
        s=InheritedPopulation();z=[0]*6;z[-1]=1
        sim=FeedbackSimulator(s,.29,.01,17);result=sim.run(z,10,lambda *args:.29)
        self.assertLessEqual(result['spent'],.01)
        result=FeedbackSimulator(s,.29,1).run(z,10,lambda *args:.29,max_events=0)
        self.assertEqual(result['status'],'unresolved_resource_cap')
        with self.assertRaises(ValueError):FeedbackSimulator(s,.29,1).run(z,10,lambda *args:.3)


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