import unittest
from fractions import Fraction as Q
import numpy as np
import sympy as sp
from example import source,MARKER
from lineage import *
from inference import *
from certkit import kl_interval,kl_lower,I,evaluate
from prospective_boxes import decide


class ScientificChecks(unittest.TestCase):
    def test_exact_full_source_inverse_and_resolvent(self):
        s=source('.85');E=MarkerChannel(MARKER)
        for rate in ['1/100','1','1000']:
            d=RandomDeadline(rate);r=d.recover(d.law(s,E)['law'],E)
            self.assertEqual(r['killed'],s.killed);self.assertEqual(r['exit_mass'],s.exits);self.assertEqual(r['preparation'],s.preparation)
        A=(sp.eye(2)-s.killed).inv();B=(2*sp.eye(2)-s.killed).inv()
        self.assertEqual(A-B,A*B)
        # Rectangular channel, singular H and nondividing states are allowed.
        s=BranchingSource(['1/2','1/2'],[[0,0],[0,0]],[0,0],[0,1],[[1,0,0,0],[0,0,0,1]])
        E=MarkerChannel([['1/2',0],[0,'1/2'],['1/2','1/2']]);d=RandomDeadline()
        r=d.recover(d.law(s,E)['law'],E)
        self.assertEqual(r['killed'],s.killed);self.assertEqual(r['sisters'],[None,None])

    def test_delayed_pair_recovery_retains_all_interruption_records(self):
        s=source('.85');E=MarkerChannel(MARKER);d=RandomDeadline('1',2)
        law=d.law(s,E);self.assertEqual(law['law'].shape,(2,12));self.assertAlmostEqual(law['law'].sum(),1)
        r=d.recover(law['law'],E)
        self.assertLess(np.max(np.abs(r['exit_mass']-np.array(s.exits,float))),1e-11)
        self.assertLess(np.max(np.abs(r['killed']-np.array(s.killed,float))),1e-12)
        with self.assertRaises(ValueError):d.recover(law['law'][:,:7],E)

    def test_calibration_mixtures_and_rank_failure(self):
        E=MarkerChannel(MARKER);P=matrix([['4/5','1/4'],['1/5','3/4']])
        self.assertEqual(known_mixture_calibration(E.emissions*P,P).emissions,E.emissions)
        with self.assertRaises(ValueError):MarkerChannel([['1/2','1/2'],['1/2','1/2']])
        with self.assertRaises(ValueError):BranchingSource([1,0],[[0,1],[1,0]],[1,1],[1,1],[[1,0,0,0],[0,0,0,1]])

    def test_kl_endpoints_use_correct_unscaled_threshold(self):
        for k,n in [(0,1000),(1000,1000),(4654,50000),(19497,50000),(1,20)]:
            L=Q(25,4);s=kl_interval(Q(k,n),n,L)
            for p in [s.a,s.b]:
                if 0<p<1:self.assertGreaterEqual(n*kl_lower(Q(k,n),p),L)
            self.assertLessEqual(s.a,Q(k,n));self.assertGreaterEqual(s.b,Q(k,n))

    def test_retained_counts_inference_and_unresolved_conflict(self):
        u=FounderCounts((19497,4654,4270,10288),(7826,3465));t=FounderCounts((13937,10895,3754,11178),(6926,3310));cal=CalibrationCounts((10211,85049),100000)
        r=InductionInference().contrast(u,t,cal)
        self.assertEqual(r['status'],'resolved');self.assertGreater(Q(r['lower']),Q(578,1000));self.assertLess(Q(r['upper']),Q(1092,1000))
        small=FounderCounts((12,1,1,3),(1,2));self.assertEqual(InductionInference().arm(small,cal)['status'],'unresolved')
        with self.assertRaises(ValueError):evaluate([[I('.3'),I(0)],[I(0),I('.3')]],I('.4'),I('.1'),I('.85'),I(1),'cancelled')
        with self.assertRaises(ValueError):FounderCounts((1,2,3,-1),(0,0))

    def test_concordance_same_means_different_extinction_and_one_daughter_alias(self):
        s=source('.18');more=s.concordance_shift(['.05','.05'])
        self.assertEqual(s.mean_generator(),more.mean_generator())
        for i in range(2):
            self.assertEqual(sum(s.sisters[i,2:]),sum(more.sisters[i,2:]))
        f,p,m=s.population(np.linspace(0,24,121));g,q,n=more.population(np.linspace(0,24,121))
        self.assertLess(np.max(abs(m-n)),1e-12);self.assertGreater(q[-1]-p[-1],.00054)
        self.assertGreaterEqual(np.min(g-f),-1e-12)
        x,y=sp.symbols('x y');delta=(more.sisters-s.sisters)*sp.Matrix([x*x,x*y,y*x,y*y])
        self.assertEqual(sp.expand(delta[0]-sp.Rational(1,20)*(x-y)**2),0)

    def test_continuous_prospective_boxes(self):
        self.assertFalse(decide(Q(1),140000,140000)['ok'])
        r=decide(Q(1),150000,150000)
        self.assertTrue(r['ok']);self.assertGreater(r['alt'].a,Q(2,5));self.assertLess(r['null'].b,Q(2,5))
        self.assertTrue(decide(Q(1),1000000,1000000,'hoeffding','matrix')['ok'])


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