"""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.
"""
from contextlib import nullcontext
from dataclasses import replace
import unittest
from unittest.mock import Mock, 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))


def child_body(spec=None):
    from CelerisLab.body.geometry.circle import CircleGeometry
    from CelerisLab.body.objects import SimObject
    return SimObject(99, geometry=CircleGeometry(8., 8., .8), center=(8., 8.), radius=.8,
                     surface_actuation=spec)


def child_jet(speed=.002):
    from CelerisLab.body.jets import Jet, Tube, Parallel
    return Jet(Tube((6.5, 8.), 0., 3., .3), Parallel(0.), amplitude=speed)


def mocked_child():
    from CelerisLab.cuda import compiler_v2 as compiler
    context = Mock(activate=nullcontext, sm_arch="sm_70")
    module = Mock(get_function=Mock(side_effect=lambda name: Mock(name=name)))
    module.get_global.return_value = (Mock(), 20)
    with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
            patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \
            patch.object(compiler, "load_module", return_value=module), \
            patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \
            patch.object(cuda, "memcpy_htod"):
        return FixedPatchStepper(CFG, L1, context=context, child_objects=[child_body()],
                                 child_jets=[child_jet()], **KW)


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)
        for kw in (dict(child_objects=[child_body()]), dict(child_jets=[child_jet()]),
                   dict(KW, child_jets=[child_jet()])):
            with 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)


    def test_four_halo_bfl_steps_and_fresh_middle_feeds(self):
        run = mocked_child()
        run._status = "ready"
        f, g = run._levels[1:]
        b = g.boundary
        events, feeds = [], []
        advance, gather = run._advance_child, run._gather_halo
        def record_feed(buffer, **kw):
            self.assertIs(buffer, (f.current, f.spare)[len(feeds)])
            feeds.append(buffer)
            advance(buffer, **kw)
        def record_halo(level, buffer, **kw):
            if level is g:
                events.append(("halo", buffer))
            gather(level, buffer, **kw)
        b.kernel.side_effect = lambda buffer, *a, **k: events.append(("bfl", buffer))
        def record_step(src, dst, halo, flags, kind, **kw):
            self.assertIs(flags, b.flags)
            self.assertIs(kind, g.kind)
            events.append(("step", src))
        g.kernel.side_effect = record_step
        with patch.object(run, "_advance_child", side_effect=record_feed), \
                patch.object(run, "_gather_halo", side_effect=record_halo):
            run.step()
        self.assertEqual(feeds, [f.current, f.spare])
        expected = [(kind, buffer) for buffer in (g.current, g.spare, g.current, g.spare)
                    for kind in ("halo", "bfl", "step")]
        self.assertEqual(events, expected)
        self.assertGreater(len(np.unique(b.tables["block"])), 1)
        # Negative control: reusing L1's old source at the midpoint must fail.
        feeds.clear()
        with patch.object(run, "_advance_child", side_effect=lambda buffer, **kw:
                          record_feed(f.current, **kw)):
            with self.assertRaises(AssertionError):
                run.step()
        run.close()

    def test_child_update_isolated_and_failure_atomic(self):
        run = mocked_child()
        b = run._levels[2].boundary
        before = b.tables["surface_velocity"].copy()
        with self.assertRaises(IndexError):
            run.set_jet_amplitude(0, -.002)  # default stays L1
        for level in (0, 3, True):
            with self.assertRaises(ValueError):
                run.set_jet_amplitude(0, -.002, level=level)
        with patch.object(cuda, "memcpy_htod") as upload:
            run.set_jet_amplitude(0, -.002, level=2)
            self.assertEqual(upload.call_count, 1)
            self.assertIs(upload.call_args.args[0], b.device["surface_velocity"])
        np.testing.assert_array_equal(b.tables["surface_velocity"], -before)
        self.assertEqual(run._jets, ())
        with patch.object(cuda, "memcpy_htod", side_effect=RuntimeError("copy failed")):
            with self.assertRaisesRegex(RuntimeError, "copy failed"):
                run.set_jet_amplitude(0, .002, level=2)
        self.assertEqual(run._status, "failed")
        self.assertEqual(b.jets[0].amplitude, -.002)
        np.testing.assert_array_equal(b.tables["surface_velocity"], -before)
        run.close()

    def test_child_public_masks_include_solids(self):
        from CelerisLab.lbm.descriptors import SOLID
        run = mocked_child()
        solid = run.child_boundary.flags.reshape(run.child_patch.fine_shape) & SOLID != 0
        self.assertTrue(solid.any())
        self.assertFalse(run._child_leaf[solid].any())
        arrays = [np.ones(level.shape, np.float32) for level in run._levels]
        run._mask_nonleaves(*arrays)
        self.assertTrue(np.isnan(arrays[2][:, solid]).all())
        with patch.object(cuda, "memcpy_htod") as upload:
            run.initialize(*arrays)
        sent = upload.call_args_list[-1].args[1]
        dense = run.child_layout.to_dense(sent)
        self.assertTrue(np.isnan(dense[:, ~run._child_leaf]).all())
        run.close()


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())


class ThreeLevelBodyGPU(unittest.TestCase):
    """Frozen gates; CPU mocks do not establish these numerical results."""
    def run_case(self, mode, block, *, jets=(), spec=None, switch=False):
        from CelerisLab.cuda.context import CudaContext
        from CelerisLab.lbm.descriptors import SOLID
        cfg = CFG if mode == "periodic" else replace(
            CFG, velocity=0., wall_normal_velocity=.001, wall_slot_width=4,
            y_wall_bc="on_site_velocity", top_wall_bc="reservoir")
        with CudaContext(0) as ctx, FixedPatchStepper(
                cfg, L1, context=ctx, coarse_boundary=mode, child_objects=[child_body(spec)],
                child_jets=jets, **dict(KW, child_block=block)) as run:
            run.initialize(*(state(*level.shape[1:], i + 1) for i, level in enumerate(run._levels)))
            snapshots = []
            for stop in (0, 1, 2, 4, 8):
                run.step(stop - run.step_count)
                if stop == 4 and switch:
                    run.set_jet_amplitude(0, -.002, level=2)
                public, stored = run.download(), raw(run)
                leaves = [~run.patch.covered, run._fine_leaf, run._child_leaf]
                if mode != "periodic":
                    leaves[0][[0, -1]] = False
                for g, mask in zip(public, leaves):
                    self.assertTrue(np.isfinite(g[:, mask]).all())
                    self.assertTrue(np.isnan(g[:, ~mask]).all())
                invalid = [run.patch.covered.copy(), run._middle_kind == OUTER_GHOST,
                           (run.child_patch.fine_kind == OUTER_GHOST) |
                           ((run.child_boundary.flags.reshape(run.child_patch.fine_shape) & SOLID) != 0)]
                if mode != "periodic":
                    invalid[0][[0, -1]] = True
                for g, mask in zip(stored, invalid):
                    self.assertTrue(np.isnan(g[:, mask]).all())
                snapshots.append(public + tuple(stored))
            return snapshots

    def test_layouts_static_driven_switched_and_zero(self):
        for mode in ("periodic", "slot_reservoir"):
            cases = [self.run_case(mode, (4, 4), jets=jets, switch=switch)
                     for jets, switch in (((), False), ([child_jet()], False),
                                          ([child_jet()], True), ([child_jet(0.)], False))]
            for a, b in zip(cases[0], cases[3]):
                for x, y in zip(a, b):
                    np.testing.assert_array_equal(x, y)
            self.assertFalse(np.array_equal(cases[1][-1][2], cases[2][-1][2], equal_nan=True))
            for block in ((8, 4), (4, 8)):
                for expected, (jets, switch) in zip(cases, (((), False), ([child_jet()], False),
                                                          ([child_jet()], True), ([child_jet(0.)], False))):
                    actual = self.run_case(mode, block, jets=jets, switch=switch)
                    for a, b in zip(actual, expected):
                        for x, y in zip(a, b):
                            np.testing.assert_array_equal(x, y)

    def test_legacy_surface_equivalence(self):
        from CelerisLab.body.objects import SurfaceActuation
        spec = SurfaceActuation(.7, .85, .002)
        for mode in ("periodic", "slot_reservoir"):
            legacy = self.run_case(mode, (4, 4), spec=spec)[-1][2]
            driven = self.run_case(mode, (4, 4), jets=[spec.as_jet((8., 8.), .8)])[-1][2]
            mask = np.isfinite(legacy)
            np.testing.assert_array_equal(mask, np.isfinite(driven))
            self.assertLessEqual(np.max(np.abs(legacy[mask] - driven[mask])), 1e-6)


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