import itertools
import unittest
import mpmath as mp
import numpy as np
from example import (PolymerCatalogue,CatalyticEnvironment,CappedZipfSource,
                     KineticParameters,ReactionEvent,FedReactor,TheoremBounds)


class ScientificChecks(unittest.TestCase):
    def test_catalogue_closure_and_literal_raf_oracle(self):
        for n in range(2,7):
            c = PolymerCatalogue(n)
            self.assertEqual(len(c.words),2**(n+1)-2)
            self.assertEqual(len(c.channels),(n-2)*2**(n+1)+4)
        c = PolymerCatalogue(4)
        source = CatalyticEnvironment.specified(c)
        r = c.channel_index['00','11']
        self.assertEqual(source.max_raf({r}),frozenset({r}))
        self.assertTrue(source.witness_event())
        self.assertEqual(len(source.irreducible_raf()),1)
        self.assertFalse(CatalyticEnvironment(c,tuple(frozenset() for _ in c.words)).max_raf())
        # Exhaust every subset of the five catalysed channels; maximal RAF is
        # the union of all directly tested RAFs, not just a pruning assertion.
        candidates = sorted(set().union(*source.rows))
        union = set()
        for bits in itertools.product((False,True),repeat=len(candidates)):
            selected = {r for r,take in zip(candidates,bits) if take}
            closure = c.closure(selected)
            valid = selected and all(all(x in closure for x in
                (c.channels[r].left,c.channels[r].right,c.channels[r].product))
                and any(r in source.rows[x] for x in closure) for r in selected)
            if valid:
                union.update(selected)
        self.assertEqual(source.max_raf(),union)

    def test_capped_law_against_direct_finite_sum_and_conditioning(self):
        c = PolymerCatalogue(4)
        source = CappedZipfSource(c)
        R = len(c.channels)
        with mp.workdps(70):
            a = mp.mpf('1.5')
            tail = mp.zeta(a,R)/mp.zeta(a)
            mean = sum((k-1)*mp.power(k,-a)/mp.zeta(a) for k in range(1,R))+(R-1)*tail
            self.assertLess(abs(mp.mpf(CappedZipfSource.statistics(4)['mean_degree'])-mean),mp.mpf('1e-32'))
            self.assertAlmostEqual(source.weights[-1],float(tail),14)
        rng = np.random.default_rng(77)
        for _ in range(100):
            e = source.sample(rng,True)
            self.assertTrue(e.witness_event())
            self.assertEqual(len(e.max_raf({c.channel_index['00','11']})),1)
            self.assertTrue(all(len(row) <= R-1 for row in e.rows))
        # Conditional inclusion tilts the degree by d/R, not by adding a bit.
        biased = np.arange(R)*source.weights
        biased /= biased.sum()
        self.assertEqual(biased[0],0)
        self.assertGreater(np.dot(np.arange(R),biased),np.dot(np.arange(R),source.weights))

    def test_coincident_catalyst_falling_factorials_and_mass(self):
        c = PolymerCatalogue(2)
        e = CatalyticEnvironment.specified(c,(('0','0','0'),('00','0','0')))
        reactor = FedReactor(e,KineticParameters.sample(e))
        counts = np.arange(1,len(c.words)+1,dtype=np.int64)
        vector = reactor.rates(counts,7)
        for i,event in enumerate(reactor.events):
            self.assertAlmostEqual(vector[i],event.propensity(counts,7),12)
        triple = ReactionEvent((0,0,0),(1,0),2.,'catalytic')
        self.assertEqual(triple.propensity([2,0],5),0)
        self.assertAlmostEqual(triple.propensity([4,0],5),2*4*3*2/25)
        self.assertTrue(np.all(c.lengths@reactor.stoich == 0))
        for values in (np.ones(len(c.words)),np.arange(len(c.words))/7):
            derivative = reactor.rhs(0,np.r_[values,0,0,0,0])
            self.assertAlmostEqual(c.lengths@derivative[:len(values)],10-c.lengths@values,10)

    def test_deterministic_reference_and_two_independent_balances(self):
        c = PolymerCatalogue(4)
        e = CatalyticEnvironment.specified(c)
        p = KineticParameters.sample(e,reproduce_host_marks=True)
        times = np.linspace(0,100,201)
        for enabled in (True,False):
            model = FedReactor(e,p,enabled)
            solution = model.deterministic(times)
            N = len(c.words)
            self.assertLess(np.max(np.abs(solution[:,:N]@c.lengths-10)),1e-10)
            self.assertLess(np.max(np.abs(solution[:,:N]@c.nonfood+solution[:,N]-solution[:,N+1]-solution[:,N+2])),1e-9)
            output = solution[-1,N]-solution[2,N]
            if enabled:
                self.assertAlmostEqual(output,162.69,delta=.02)
            else:
                self.assertAlmostEqual(output,4.15e-5,delta=1e-7)
                self.assertEqual(float(np.max(np.abs(solution[:,N+1]))),0)

    def test_jump_rewards_and_budget_semantics(self):
        c = PolymerCatalogue(4)
        e = CatalyticEnvironment.specified(c)
        # Larger basal coefficient is explicitly a test fixture, not paper data.
        p = KineticParameters.sample(e,epsilon=.01)
        model = FedReactor(e,p)
        run = model.stochastic(10,99,end=2,start=.2,max_events=10000)
        self.assertTrue(run['completed'])
        self.assertTrue(run['integer_balances_exact'])
        self.assertAlmostEqual(run['final_nonfood_mass']+run['total_export'],
                               run['signed_catalytic_input']+run['signed_basal_input'],12)
        self.assertGreater(run['positive_basal_input'],0)
        self.assertLessEqual(run['window_export'],run['total_export'])
        limited = model.stochastic(10,99,end=2,start=.2,max_events=1)
        self.assertFalse(limited['completed'])
        self.assertIsNone(limited['productive'])
        with self.assertRaises(ValueError):
            model.stochastic(10**60,99)

    def test_theorem_hypotheses_and_log_domain_bounds(self):
        self.assertFalse(TheoremBounds.evaluate(4,20)['applicable'])
        bounds = TheoremBounds.evaluate(4,10**60*25)
        self.assertTrue(bounds['applicable'])
        self.assertEqual(float(bounds['log10_upper_bound_capped_at_one']),0)
        self.assertLess(float(bounds['log10_disabled_upper_bound_capped_at_one']),-1e10)
        stats = CappedZipfSource.statistics(1024)
        self.assertAlmostEqual(float(stats['X_times_p']),9/np.pi**2,delta=.03)
        budget = TheoremBounds.structural_budget(4)
        self.assertEqual((budget['h'],budget['L']),(2,16))
        self.assertAlmostEqual(float(budget['log2_budget']),2*np.log2(5*2**17+4**4*2**2),12)
        self.assertLess(float(TheoremBounds.structural_budget(1000000)['log2_budget_over_n']),.07)


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