"""Tiny CPU algebra/event tests; no full trajectory or GPU execution."""
from dataclasses import replace
import unittest
import numpy as np
import spacetime_reference as st


class SpaceTime(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.plan = st.Plan(12, (3, 3, 9, 8))
        cls.spec = st.Spec(n_coarse=12, patch_lo=1, halo=2,
                           tau_coarse=.5576, tau_fine=.6152, collision="srt")

    def state(self, tagged=False):
        arrays = []
        for n, mask in ((12,self.plan.maps.coarse_active),(24,self.plan.maps.fine_active)):
            q, y, x = np.indices((9,n,n))
            a = np.broadcast_to(st.cpu.W[:,None,None], (9,n,n)).copy()
            if tagged:
                a *= 1+1e-3*np.random.default_rng(n).normal(size=a.shape)
            a[:,~mask] = np.nan
            arrays.append(a)
        return st.cpu.State(*arrays)

    def test_events_and_corner_leftovers(self):
        groups = [g for g in self.plan.groups if None in g.events]
        self.assertEqual(len(groups), 8)
        self.assertTrue(all(g.events.count(None)==3 for g in groups))
        for m in (0,1):
            events = list(map(tuple,self.plan.outgoing[m]))
            self.assertEqual(len(events),len(set(events)))
        st.Plan(16,(4,5,11,12))

    def test_interpolation(self):
        for g in self.plan.groups:
            samples, a, e, coeff = st.interpolation(g.kind,g.points)
            np.testing.assert_allclose(g.weights@np.ones(len(samples)),1,rtol=0,atol=2e-14)
            np.testing.assert_allclose(g.weights.mean(axis=0), np.eye(len(samples))[1 if g.kind=='face' else 0],rtol=0,atol=2e-14)
            if g.kind=='face':
                beta=np.array([1.,.03,.02,0.])
            else:
                cx,cy=st.cpu.CX[g.q],st.cpu.CY[g.q]
                beta=np.array([1.,.03,.02,-cx*.03-cy*.02])
                steady=a@np.array([1.,.03,.02,0.])
                self.assertLess(abs((coeff@steady)[3]),2e-14)
            np.testing.assert_allclose(g.weights@(a@beta),e@beta,rtol=0,atol=2e-14)

    def test_rest(self):
        initial=self.state()
        final,d=st.advance(self.spec,self.plan,initial)
        for a,b in zip((initial.coarse_f,initial.fine_f),(final.coarse_f,final.fine_f)):
            np.testing.assert_allclose(a,b,rtol=0,atol=2e-14,equal_nan=True)

    def test_frozen_recovers_ivp(self):
        for collision in ('identity','srt'):
            spec=replace(self.spec,collision=collision)
            a=b=self.state(True)
            for _ in range(2):
                a,_=st.advance(spec,self.plan,a,'frozen')
                b,_=st.cpu.advance_coarse_once(spec,self.plan.maps,b)
                for aa,bb in zip((a.coarse_f,a.fine_f),(b.coarse_f,b.fine_f)):
                    np.testing.assert_allclose(aa,bb,rtol=0,atol=2e-14,equal_nan=True)

    def test_second_substep_uses_characteristic_origin(self):
        # Review counterexample in (x,y): ghost (7,10), m=1 traces back
        # to (6,9), parent (3,4), not its current parent (3,5).
        plan=st.Plan(16,(4,4,12,12))
        owners=[g for g in plan.groups if g.q==5 and (1,10*32+7) in g.events]
        self.assertEqual(len(owners),1)
        self.assertEqual(owners[0].source,4*16+3)
        spec=st.Spec(n_coarse=16,patch_lo=2,halo=4,collision='identity')
        arrays=[]
        for n,mask in ((16,plan.maps.coarse_active),(32,plan.maps.fine_active)):
            a=np.broadcast_to(st.cpu.W[:,None,None],(9,n,n)).copy()
            a[:,~mask]=np.nan
            arrays.append(a)
        arrays[0][5,4,3]=1.; arrays[0][5,5,3]=2.
        final,_=st.advance(spec,plan,st.cpu.State(*arrays),'frozen')
        self.assertEqual(final.fine_f[5,11,8],1.)
        self.assertNotEqual(final.fine_f[5,11,8],arrays[0][5,5,3])

    def test_corner_affine_tradeoff_is_explicit(self):
        group=next(g for g in self.plan.groups if g.kind=='corner')
        _,a,e,coeff=st.interpolation(group.kind,group.points)
        linear_x=np.array([0.,1.,0.,0.])
        # The two paper constraints are consistent. Arbitrary affine exactness
        # is the incompatible additional demand.
        self.assertEqual(abs(float((e@linear_x).mean())),.125)
        fitted=coeff@(a@linear_x)
        self.assertLess(abs(fitted[3]),2e-14)
        self.assertLess(abs(float((e@fitted).mean())),2e-14)
        self.assertGreater(np.max(np.abs(e@fitted-e@linear_x)),.01)

    def test_directional_budgets(self):
        for collision in ('identity','srt'):
            spec=replace(self.spec,collision=collision)
            a=self.state(True)
            for _ in range(2):
                a,d=st.advance(spec,self.plan,a)
                np.testing.assert_allclose(d['directional_residual'],0,rtol=0,atol=2e-12)
                self.assertLess(d['group_error'],2e-14)

    def test_missing_packets_fail(self):
        for mutation in ('omit_leftover','omit_outgoing_0','omit_outgoing_1'):
            _,d=st.advance(self.spec,self.plan,self.state(True),mutation=mutation)
            self.assertGreater(np.max(np.abs(d['directional_residual'])),1e-5)

    def test_negative_interpolation_is_reported_not_clipped(self):
        initial=self.state()
        group=self.plan.groups[0]
        initial.coarse_f.reshape(9,-1)[group.q,group.source]=-.01
        spec=replace(self.spec,collision='identity')
        _,d=st.advance(spec,self.plan,initial,'frozen')
        self.assertGreaterEqual(d['negative_boundary_count'],4)
        self.assertEqual(d['min_boundary'],-.01)
        self.assertEqual(d['nonfinite_boundary_count'],0)


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