import unittest
from fractions import Fraction as F
import numpy as np
from branching import *
from certificates import TwoSiteCertificate,SmoothSlopeCertificate
import deadline_validated as deadline


class ScientificTests(unittest.TestCase):
    def test_allocation_conservation_and_common_daughter_law(self):
        src=MolecularSource(4);i,j,k,p=src.pairs
        st=np.array(src.states)
        np.testing.assert_array_equal(st[j]+st[k],st[i])
        for state in range(len(st)):
            mask=i==state
            first=np.bincount(j[mask],weights=p[mask],minlength=len(st))
            second=np.bincount(k[mask],weights=p[mask],minlength=len(st))
            np.testing.assert_allclose(first,second,atol=1e-15)
            np.testing.assert_allclose(first,src.D[state],atol=1e-15)
        np.testing.assert_allclose(src.Q.sum(1),0,atol=1e-14)
        self.assertTrue(np.all((src.Q-np.diag(src.Q.diagonal()))>=0))

    def test_nonlinear_covariance_and_nonmonotone_obstruction(self):
        src=MolecularSource(2);i=src.index[2,0]
        f=np.array([1. if a in (0,2) else 0. for a,r in src.states])
        self.assertAlmostEqual((Complementary().product(src,f,f)-(src.D@f)**2)[i],.25)
        increasing=np.array([a/2 for a,r in src.states])
        self.assertLess((Complementary().product(src,increasing,increasing)-(src.D@increasing)**2)[i],0)
        np.testing.assert_allclose(Independent().product(src,f,f),(src.D@f)**2)

    def test_equal_means_ordered_pgf_variance_and_capture(self):
        src=MolecularSource(2);times=np.array([0.,5.,20.]);computed=[]
        for law in [Complementary(),Independent()]:
            model=BranchingPopulation(src,StepHazard(),law)
            mean,var=model.moments(times);u=model.pgf(times)
            cap=IndependentCapture(.3);z=cap.nondetection(model,times)
            self.assertTrue(np.all(z>=u-1e-11))
            cm,cv=cap.moments(mean,var);np.testing.assert_allclose(cm,.3*mean)
            np.testing.assert_allclose(cv,.09*var+.21*mean)
            computed.append((mean,var,u,z))
        np.testing.assert_allclose(computed[0][0],computed[1][0],rtol=1e-10)
        for i in [1,2,3]:self.assertTrue(np.all(computed[0][i]<=computed[1][i]+1e-8))

    def test_erlang_clock_and_pure_birth_limit(self):
        src=MolecularSource(2)
        class NoDeath:
            def values(self,states,sites):return np.zeros(len(states))
        one=BranchingPopulation(src,NoDeath(),Complementary())
        mean,var=one.moments([2.]);np.testing.assert_allclose(mean,np.exp(.2),rtol=1e-9)
        np.testing.assert_allclose(var,np.exp(.2)*(np.exp(.2)-1),rtol=1e-8)
        for phases in [2,4]:
            models=[BranchingPopulation(src,StepHazard(),law,phases=phases) for law in [Complementary(),Independent()]]
            a,b=[m.moments([5.])[0] for m in models];np.testing.assert_allclose(a,b,rtol=1e-9)
            a,b=[m.extinction()[0] for m in models];self.assertTrue(np.all(a<=b+1e-10))
        with self.assertRaises(RuntimeError):one.extinction(max_iterations=0)

    def test_exact_boxes_response_and_regrowth(self):
        cert=TwoSiteCertificate();r=cert.calculate()
        self.assertTrue(r['integrator_coefficients_verified'])
        self.assertGreater(r['response_gap_lower'][5],F(1144144,10**7))
        self.assertLess(r['response_gap_upper'][5],F(1144331,10**7))
        hit=cert.regrowth();self.assertGreater(hit['gap_lower'],F(98,1000))
        self.assertLess(cert.regrowth(2)['independent_eventual_upper'],F(1)+F(1,10**6))
        with self.assertRaises(ValueError):cert.regrowth(1)

    def test_fresh_slope_envelope_and_observation_limits(self):
        r=SmoothSlopeCertificate().calculate(F(7999,1000),F(8001,1000),1)
        self.assertGreater(r['gap_lower'],F(1427,100000))
        self.assertLess(r['all_active_gap_upper'],F(122,1000000))
        row=r['rows'][0]
        self.assertTrue(all(l<=u for l,u in zip(row['death_lower'],row['death_upper'])))
        # The endpoint hazard vectors are not ordered: statewise envelopes are necessary.
        src=MolecularSource(16);a=SmoothHazard(7.999).values(src.states,16);b=SmoothHazard(8.001).values(src.states,16)
        self.assertTrue(np.any(a>b) and np.any(a<b))
        self.assertEqual(PairedCloneObservation(.5,.01,.02).interval(.1),(.1-.25/4-.1,.2))
        self.assertGreater(balanced_survival_bound(np.log(128),np.log(128)/2.12),1)
        self.assertLess(balanced_survival_bound(1000,1000/2.12),1)

    def test_invalid_inputs_and_no_silent_probability_clipping(self):
        for kwargs in [dict(sites=0),dict(sites=2.5),dict(erasure=-.1)]:
            with self.assertRaises(ValueError):MolecularSource(**kwargs)
        src=MolecularSource()
        for kwargs in [dict(division_rate=0),dict(phases=0)]:
            with self.assertRaises(ValueError):BranchingPopulation(src,StepHazard(),Complementary(),**kwargs)
        with self.assertRaises(ValueError):IndependentCapture(1.01)
        with self.assertRaises(ValueError):BranchingPopulation(src,StepHazard(),Complementary()).pgf([1.],1.01)
        with self.assertRaises(ValueError):SmoothSlopeCertificate().calculate(7,10)
        with self.assertRaises(ValueError):PairedCloneObservation().interval(0,pairs=0)


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