import hashlib
import itertools
import unittest
from pathlib import Path
from fractions import Fraction as Q
import numpy as np
from example import (ComplementaryAllocation,CountBox,ConstructedSource,OffspringKernel,
                     RedundancyBounds,FourSpeciesModule,CoupledModules,ReturnEllipsoid,
                     SemenovModel,semenov_scalar_budget,REGION_SHA256)


class ScientificChecks(unittest.TestCase):
    def test_complementary_allocation_exhaustively(self):
        for n in range(9):
            assignments=[sum(bits) for bits in itertools.product((0,1),repeat=n)]
            for lo in range(5):
                for hi in range(lo,7):
                    exact=Q(sum(lo <= d <= hi and lo <= n-d <= hi for d in assignments),2**n)
                    self.assertEqual(ComplementaryAllocation.coordinate_return(n,lo,hi),exact)
        joint=ComplementaryAllocation.coordinate_return(4,1,3)
        self.assertEqual(joint,Q(7,8))
        self.assertNotEqual(joint,Q(7,8)**2)
        self.assertEqual(ComplementaryAllocation.joint_return((4,4),CountBox((1,1),(3,3))),Q(7,8)**2)

    def test_constructed_literal_rates_and_protocol_monitors(self):
        model=ConstructedSource()
        state=np.array((53,3,530,53,53))
        rates=model.rates(state)
        self.assertEqual(len(rates),29)
        self.assertAlmostEqual(rates[4],15*53*52/53)
        self.assertAlmostEqual(rates[5],15*53*3/53)
        self.assertEqual(rates[-1],1)
        self.assertEqual(model.kinds[11],'reverse_growth')
        self.assertEqual(model.jumps[10].tolist(),[-1,0,0,0,1])
        self.assertEqual(model.jumps[12].tolist(),[-1,1,0,0,0])
        self.assertEqual(model.run((0,0),limit=1)['status'],'unfinished')
        self.assertIsNone(model.run((0,0),limit=1)['both_return'])
        self.assertEqual(model.run((0,0),limit=1,quota=1)['status'],'event_quota')
        self.assertEqual(model.run((0,0),deadline=1e-20)['status'],'deadline')
        low,high=model.boxes
        self.assertTrue(low.contains((10,1)) and high.contains((400,40)))
        self.assertEqual(Q(10,11),Q(400,440))
        # Decoder-correct does not imply admission.
        self.assertFalse(high.contains((300,10)))
        self.assertGreater(300,212)

    def test_scalar_certificate_assembly(self):
        constructed=ConstructedSource.paper_budget()
        self.assertEqual(Q(constructed['premonitor']),Q('0.991028'))
        self.assertEqual(Q(constructed['assembled_lower_bound']),Q('0.9901518'))
        self.assertEqual(constructed['capacity_per_pool_molecules'],100100212000)
        for label in (0,1):
            budget=semenov_scalar_budget(label)
            self.assertLess(Q(budget['joint_failure_upper_bound']),Q(2,250))
        self.assertGreater(Q(9901,10000)**10,Q('.905'))
        self.assertGreater(Q(124,125)**10,Q('.922'))

    def test_state_dependent_joint_kernel_iteration(self):
        kernel=OffspringKernel(('a','b'),{'a':{('a','a'):Q(9,10),('b','b'):Q(9,100)},
                                        'b':{('a','b'):Q(1,2),('b','a'):Q(12,25)}})
        self.assertEqual(kernel.uniform_one_cycle,Q(49,50))
        for G in range(5):
            for family in (False,True):
                bound=kernel.uniform_one_cycle**((2**G-1) if family else G)
                self.assertTrue(all(p >= bound for p in kernel.success(G,family).values()))
        self.assertNotEqual(kernel.success(3),kernel.success(3,True))
        with self.assertRaises(ValueError):OffspringKernel(('a',),{'a':{('a','a'):Q(2)}})

    def test_four_species_source_and_physical_scaling(self):
        module=FourSpeciesModule()
        for lo,hi in (('0.99579401232','0.99579401233'),('2.97636724376','2.97636724377')):
            center,f0=module.stationary_reduction(Q(lo))
            _,f1=module.stationary_reduction(Q(hi))
            self.assertLess(f0*f1,0)
            self.assertLess(np.max(np.abs(module.fluid(np.array(center,dtype=float)))),1e-9)
        u=np.array((13.,20.,1.,9.));N=100;gamma=1e-11
        assembly=CoupledModules(1,[[0]],gamma)
        rates,jumps=assembly.rates_and_jumps((N*u).astype(int).reshape(1,4),N)
        drift=rates@jumps/N
        bias=np.array((2e-5*u[0],-1e-5*u[0],4*u[2],-2*u[2]))/N
        expected=module.fluid(u)+bias-np.array((0,0,gamma*u[2],0))
        np.testing.assert_allclose(drift[:4],expected,atol=1e-12)
        ring=CoupledModules(3,[[0,5e-12,5e-12],[5e-12,0,5e-12],[5e-12,5e-12,0]])
        self.assertTrue(ring.interaction_within_paper_allowance())
        dense=CoupledModules(4,np.ones((4,4))*5e-12-np.eye(4)*5e-12)
        self.assertFalse(dense.interaction_within_paper_allowance())

    def test_redundancy_horizons_and_scale(self):
        bounds=RedundancyBounds()
        first=bounds.required(1)
        self.assertTrue(3.1e22 < first['sufficient_N'] < 3.2e22)
        self.assertEqual(first['sufficient_N'],bounds.required(1,family=True)['sufficient_N'])
        self.assertGreater(bounds.required(10,family=True)['sufficient_N'],bounds.required(10)['sufficient_N'])
        self.assertLess(float(bounds.required(40,family=True)['necessary_real_lower_bound']),1)
        # Necessary real lower bound can be below 1; never relabel it as a
        # physically calibrated or sufficient molecule count.
        with self.assertRaises(ValueError):RedundancyBounds(gamma=Q(1,100))

    def test_exact_regions_and_nominal_recovery(self):
        path=Path(__file__).with_name('terminal_regions.json')
        self.assertEqual(hashlib.sha256(path.read_bytes()).hexdigest(),REGION_SHA256)
        model=SemenovModel()
        for label,region in enumerate(ReturnEllipsoid.load()):
            counts=[int(region.omega*z) for z in region.center]
            self.assertLess(region.exact_count_energy(counts),1)
            _,trajectory=model.recovery(region)
            self.assertLess(region.energy(trajectory[-1]),1e-4)
            self.assertEqual(sum(trajectory[-1,1:4]) > .005,bool(label))
            rates,jumps=model.jump_rates(counts,region.omega)
            self.assertEqual(len(rates),27)
            np.testing.assert_allclose(rates@jumps/region.omega,model.rhs(0,np.array(counts,dtype=float)/region.omega),atol=1e-16)
        low,high=ReturnEllipsoid.load()
        self.assertAlmostEqual(low.cube_tolerance()*1e9,16.366,delta=.002)
        self.assertAlmostEqual(high.cube_tolerance()*1e9,1.788,delta=.002)


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