"""Nested three-level FixedPatchStepper (child=..., single chain): CPU guards and GPU gates.

Fixture = reviewed CPU probe: L0 16x16 periodic, L1 (4,4,12,12) -> 24x24 window,
L2 (8,8,16,16) in L1 window cells -> 24x24 window; 4x4 blocks on every level.
Run CPU and GPU classes in separate processes. GPU gates (frozen before any run):
- uniform moving equilibrium stays on all three levels' ACTIVE leaves, |dG| <= 1e-6;
- seeded non-equilibrium: every level's ACTIVE leaves finite, raw owner cores NaN on
  OUTER/covered (INNER partly invalid by contract, not asserted);
- each L1 hop feeds L2 from the L1 buffer it is about to step (current, then the new
  G written by the first hop), never from L0: 2 child advances per coarse step.
Numerical comparison against the CPU probe (G vs C(f)) is done by the acceptance run.
"""
import unittest
from unittest.mock import patch

import numpy as np
import pycuda.driver as cuda

from CelerisLab.config import LBMConfig
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.refinement import ACTIVE, OUTER_GHOST
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from test_refinement_blocked_fine import state, circle
from test_refinement_blocked_base import raw

CFG = LBMConfig(nx=16, ny=16, viscosity=.1)
L1, L2 = (4, 4, 12, 12), (8, 8, 16, 16)
KW = dict(fine_block=(4, 4), base_block=(4, 4), child=L2, child_block=(4, 4))


class ThreeLevelCPU(unittest.TestCase):
    def test_rejections_before_cuda(self):
        context = object()  # never activated: every case must fail first
        bad = (dict(KW, child_block=None), dict(KW, fine_block=None),
               dict(KW, fine_objects=[circle()]), dict(KW, child=(4, 8, 16, 16)),
               dict(KW, child_block=(5, 4)))
        for kw in bad:
            with self.subTest(kw=kw), self.assertRaises(ValueError):
                FixedPatchStepper(CFG, L1, context=context, **kw)
        with self.assertRaises(ValueError):  # several L1 patches + child: not yet
            FixedPatchStepper(LBMConfig(nx=32, ny=16, viscosity=.1), ((4, 4, 12, 12), (20, 4, 28, 12)),
                              context=context, **KW)


class ThreeLevelGPU(unittest.TestCase):
    def test_equilibrium_is_kept_on_every_level(self):
        from CelerisLab.cuda.context import CudaContext
        feq = np.asarray(compute_feq_d2q9(1.0, .03, -.02), np.float32)[:, None, None]
        with CudaContext(0) as ctx, FixedPatchStepper(CFG, L1, context=ctx, **KW) as run:
            shapes = ((16, 16), run.patch.fine_shape, run.child_patch.fine_shape)
            run.initialize(*(np.broadcast_to(feq, (9, *s)).copy() for s in shapes))
            run.step(8)
            self.assertEqual(run.step_count, 8)
            for level, g in enumerate(run.download()):
                leaf = np.isfinite(g[0])
                self.assertTrue(leaf.any(), f"level {level}")
                np.testing.assert_allclose(g[:, leaf], np.broadcast_to(feq, g.shape)[:, leaf],
                                           atol=1e-6, rtol=0, err_msg=f"level {level}")

    def test_nonequilibrium_masks_and_fresh_middle_feeds(self):
        from CelerisLab.cuda.context import CudaContext
        with CudaContext(0) as ctx, FixedPatchStepper(CFG, L1, context=ctx, **KW) as run:
            p1, p2 = run.patch, run.child_patch
            run.initialize(state(16, 16, 1), state(*p1.fine_shape, 2), state(*p2.fine_shape, 3))
            f = run._levels[1]
            feeds, original = [], run._advance_child
            def record(buffer, **kw):
                feeds.append(int(buffer))
                original(buffer, **kw)
            with patch.object(run, "_advance_child", side_effect=record):
                cur, spare = int(f.current), int(f.spare)
                run.step(1)
            self.assertEqual(feeds, [cur, spare])  # current G, then the new G of hop 1
            run.step(7)
            middle = run._middle_kind
            for g, leaf in zip(run.download(), (~p1.covered, middle == ACTIVE, p2.fine_kind == ACTIVE)):
                self.assertTrue(np.isfinite(g[:, leaf]).all())
                self.assertTrue(np.isnan(g[:, ~leaf]).all())
            rc, rf, rg = raw(run)
            self.assertTrue(np.isnan(rc[:, p1.covered]).all())
            self.assertTrue(np.isfinite(rf[:, middle == ACTIVE]).all())
            self.assertTrue(np.isnan(rf[:, middle == OUTER_GHOST]).all())  # incl. child-covered
            self.assertTrue(np.isnan(rf[:, p2.covered]).all())
            self.assertTrue(np.isfinite(rg[:, p2.fine_kind == ACTIVE]).all())
            self.assertTrue(np.isnan(rg[:, p2.fine_kind == OUTER_GHOST]).all())


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