import unittest
from fractions import Fraction as F
from dataclasses import replace
import numpy as np
from branches import nominal
from dynamics import OperatingModel,MaintainedSupply,LinearDonor,SaturatingDonor
from certificates import ModalCertificate,TubeCertificate,fixed_band_stock,R,S
import tube_kernel as k


class ScientificTests(unittest.TestCase):
    def test_source_field_and_storage_identity(self):
        model=OperatingModel();u=k.P0;f=k.field(u,F(3,25))
        np.testing.assert_allclose(model.kinetics.rhs(0,np.array(u,float),.12),np.array(f,float),atol=2e-12)
        HG=F(21,100)*(50-u[2]-u[3]);HT=F(21,10)*(F(505,1000)-u[4])*u[7]
        self.assertEqual(f[1]+f[2]+f[3],HG-F(32,10)*u[0]*(u[1]-F(178,100)))
        self.assertEqual(sum(a*b for a,b in zip(k.ELL,f)),F(3,25)*k.R1(u[0])-HG-HT)
        num=model.kinetics.jacobian(np.array(u,float),.12)
        np.testing.assert_allclose(num,np.array(k.jacobian(u,F(3,25)),float),atol=1e-10)

    def test_stationarity_cannot_determine_repair_clock(self):
        base=nominal(True)
        for sigma in [F(1,100),F(1),F(100)]:self.assertEqual(replace(base.trx,b=base.trx.b*sigma,c=base.trx.c*sigma).response,base.trx.response)
        a=OperatingModel(repair_scale=1).necessary_repair_delay(.95*19.096);b=OperatingModel(repair_scale=.1).necessary_repair_delay(.95*19.096)
        self.assertAlmostEqual(b,10*a);self.assertGreater(a,70);self.assertLess(a,72)
        damaged=np.array(k.P0,float);damaged[5]=0;damaged[6]=.95*19.096;damaged[7]=0
        fast=OperatingModel().kinetics.rhs(0,damaged,.12);slow=OperatingModel(repair_scale=.1).kinetics.rhs(0,damaged,.12)
        self.assertAlmostEqual(slow[6],fast[6]/10);self.assertAlmostEqual(slow[0],fast[0])

    def test_nominal_faces_preparation_and_accounts(self):
        c=ModalCertificate();r=c.check();self.assertEqual(r['deadline'],F(1,250));self.assertEqual(r['shortfall_bound'],F(9,25000))
        self.assertLess(r['storage_range_S'],F(1,10));self.assertTrue(r['narrow_preparation']['contained']);self.assertFalse(r['independent_point_one_percent_box']['contained'])
        self.assertTrue(c.membership(k.C0,radius=S)['contained']);self.assertTrue(c.membership(k.P0)['contained'])
        self.assertLess(r['largest_uniform_halfwidth_around_p'],F(64,10**8))

    def test_exact_tube_replay_and_invalid_witness(self):
        t=TubeCertificate('tube_Q120');result=t.verify();self.assertTrue(result['mission_success']);self.assertGreater(F(result['exact']['min_LT_after_tau']),F(4034,1000))
        self.assertGreater(F(result['exact']['drawdown'][0]),F(7,10));self.assertLess(F(result['exact']['drawdown'][1]),F(751,1000))
        t.data['radii'][0][0]='-1'
        with self.assertRaises(AssertionError):t.validate_structure()

    def test_supply_laws_and_fixed_band(self):
        self.assertEqual(fixed_band_stock(F(251,250))['linear_sufficient_initial_stock'],1927680)
        for Q,K in [(200,F(1,100)),(1500,F(1))]:self.assertTrue(fixed_band_stock(F(251,250),Q,K)['saturating_passes'])
        self.assertFalse(fixed_band_stock(F(251,250),120,1)['saturating_passes'])
        for supply in [LinearDonor(),SaturatingDonor()]:self.assertAlmostEqual(supply.scale(0,supply.initial_stock),.12);self.assertEqual(supply.scale(0,0),0)

    def test_depleted_numerical_ledgers_and_different_endpoints(self):
        model=OperatingModel();p=np.array(k.P0,float);times=np.linspace(0,1.004,101)
        for stock,V,success in [(120,1000,True),(60,500,False)]:
            r=model.solve(p,LinearDonor(stock,V),times);self.assertEqual(r['services'][-1,1]>4,success);self.assertLess(r['storage_identity_residual'],1e-8);self.assertLess(r['donor_identity_residual'],1e-8)
            self.assertLess(r['scales'][-1],.11371266);self.assertTrue(np.all(np.diff(r['donor'])<=1e-10))
            self.assertAlmostEqual(float(r['integrated_services'][-1].sum()-r['regeneration'][-1]),float(r['storage'][0]-r['storage'][-1]),places=8)

    def test_invalid_models_are_rejected(self):
        with self.assertRaises(ValueError):OperatingModel(repair_scale=0)
        with self.assertRaises(ValueError):LinearDonor(0,1000)
        with self.assertRaises(ValueError):SaturatingDonor(120,0)
        with self.assertRaises(ValueError):MaintainedSupply(.12,.2)
        with self.assertRaises(ValueError):fixed_band_stock(1,10,1)
        with self.assertRaises(ValueError):TubeCertificate('unknown')
        p=np.array(k.P0,float);p[0]=31
        with self.assertRaises(ValueError):OperatingModel().solve(p,MaintainedSupply(),[0,1])


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