"""Periodic body-free block backend: CPU contract doubles plus one short GPU case.

Frozen GPU gates (32x24, 8 steps, positive low-Mach non-equilibrium state):
finite cores; bit-identical 1x1, 2x2 and 4x2 splits; max |G - FP64 oracle|
<= 3e-6. Static circle (centre (15.5, 11.5), R=4.3, on the 2x2 corner seam), same
case: halo writes and q<.5 halo donors asserted present; finite fluid / NaN
solid cores; bit-identical splits; max |G - FP64 BFL+pull/SRT| <= 3e-6 on
fluid cells, from the same FP32-packed link data and initial populations.
Mass/momentum drift is reported, not gated.
"""
from contextlib import nullcontext
from pathlib import Path
import resource
from threading import RLock
import unittest
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pycuda.driver as cuda

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.lbm import block_stepper as implementation
from CelerisLab.lbm.blocks import NO_DONOR, BlockLayout2D
from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY

ROOT = Path(__file__).resolve().parents[3]
NX, NY, NU, STEPS, ORACLE_ATOL = 32, 24, .1, 8, 3e-6
SPLITS = ((24, 32), (12, 16), (12, 8))  # 1x1, 2x2, 4x2 blocks (by, bx)
W = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4)
CX, CY = np.array(D2Q9_EX, float), np.array(D2Q9_EY, float)


def config(**kwargs):
    return LBMConfig(**dict(nx=NX, ny=NY, viscosity=NU, **kwargs))


def feq(rho, ux, uy):
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    return W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))


def initial_state():
    """Positive low-Mach equilibrium plus a seeded non-equilibrium part."""
    y, x = np.mgrid[0:NY, 0:NX]
    rho = 1 + 1e-3 * np.cos(2 * np.pi * x / NX)
    ux = .02 * np.sin(2 * np.pi * y / NY)
    uy = .01 * np.cos(2 * np.pi * x / NX)
    noise = np.random.default_rng(2026).uniform(-1, 1, (9, NY, NX)) * 1e-3 * W[:, None, None]
    return (feq(rho, ux, uy) + noise).astype(np.float32)


def oracle(g, steps):
    """FP64 periodic pull + SRT, from the same FP32-rounded G_0."""
    g = g.astype(np.float64)
    omega = 1 / (3 * NU + .5)
    for _ in range(steps):
        f = np.stack([np.roll(g[q], (D2Q9_EY[q], D2Q9_EX[q]), axis=(0, 1)) for q in range(9)])
        rho = f.sum(axis=0)
        ux, uy = (CX[:, None, None] * f).sum(axis=0) / rho, (CY[:, None, None] * f).sum(axis=0) / rho
        g = f + omega * (feq(rho, ux, uy) - f)
    return g


CENTER, RADIUS = (15.5, 11.5), 4.3


def circle(center=CENTER, radius=RADIUS):
    return SimObject(0, geometry=CircleGeometry(*center, radius), center=center, radius=radius)


def dense_flags():
    """FLUID everywhere, circle flags on its solids (as the stepper builds them)."""
    mask = circle().get_flag_mask(NX, NY)
    flags = np.full(NX * NY, 1, np.uint16)
    flags[(mask & 2) != 0] = mask[(mask & 2) != 0]
    return flags


def circle_oracle(g, steps, solid, links):
    """FP64 periodic static linear BFL -> pull -> SRT on fluid; solids NaN."""
    g = g.astype(np.float64)
    g[:, solid] = np.nan
    omega = 1 / (3 * NU + .5)
    rows = [(*divmod(int(link.fluid_idx), NX), int(link.dir), float(np.float32(link.q)),
             int(link.fallback)) for link in links]
    for _ in range(steps):
        staged = g.copy()
        for y, x, d, q, fallback in rows:
            ey, ex, o = D2Q9_EY[d], D2Q9_EX[d], d + 1 if d % 2 else d - 1
            toward = g[d, y, x]
            if fallback:
                value = toward
            elif q < .5:
                value = 2 * q * toward + (1 - 2 * q) * g[d, (y - ey) % NY, (x - ex) % NX]
            else:
                value = toward / (2 * q) + (1 - 1 / (2 * q)) * g[o, y, x]
            staged[o, (y + ey) % NY, (x + ex) % NX] = value
        f = np.stack([np.roll(staged[q], (D2Q9_EY[q], D2Q9_EX[q]), axis=(0, 1)) for q in range(9)])
        rho = f.sum(axis=0)
        ux, uy = (CX[:, None, None] * f).sum(axis=0) / rho, (CY[:, None, None] * f).sum(axis=0) / rho
        g = f + omega * (feq(rho, ux, uy) - f)
        g[:, solid] = np.nan
    return g


def totals(g):
    g = g.astype(np.float64)
    return g.sum(), (CX[:, None, None] * g).sum(), (CY[:, None, None] * g).sum()


class CpuContracts(unittest.TestCase):
    """No device: driver doubles prove rejection before allocation and lifecycle."""

    def forbidden_context(self):
        return Mock(activate=Mock(side_effect=AssertionError("CUDA reached")), sm_arch="sm_70")

    def test_unsupported_inputs_rejected_before_cuda(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        cases = ((config(), BlockLayout2D.build(NY, NX, 12, 16, periodic=(True, False))),
                 (config(), BlockLayout2D.build(NY, NX * 2, 12, 16)),
                 (config(collision="MRT"), layout),
                 (config(), "layout"))
        for cfg, lay in cases:
            with self.subTest(cfg=cfg.collision, layout=getattr(lay, "periodic", lay)):
                with self.assertRaises(ValueError):
                    PeriodicBlockStepper(cfg, lay, context=self.forbidden_context())
        moving = circle()
        moving.state.vx = .01
        build_links = CircleGeometry.build_cut_links
        for objects in ([moving], [circle((1.5, 11.5), 2.0)], [circle(), circle()]):
            with self.assertRaises(ValueError):
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context(), objects=objects)
        for center, radius in (((16.0, 12.0), -.1), ((16.0, 12.0), -.5), ((15.5, 11.5), np.nan),
                               ((np.inf, 11.5), 4.3), ((15.5, np.nan), 4.3)):
            with self.subTest(center=center, radius=radius), self.assertRaises(ValueError):
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context(),
                                     objects=[circle(center, radius)])
        with patch.object(implementation, "_static_circles") as prepare:
            with self.assertRaisesRegex(AssertionError, "CUDA reached"):  # body-free path
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context(), objects=iter(()))
            prepare.assert_not_called()
        with patch.object(CircleGeometry, "build_cut_links",
                          lambda geometry, *a, **k: build_links(geometry, *a, **k)[1:]):
            with self.assertRaisesRegex(ValueError, "cover every fluid pull"):
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context(),
                                     objects=[circle()])
        with patch.object(implementation, "MAX_GRID_DIM", 3):
            with self.assertRaisesRegex(ValueError, "launch limits"):
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context())

    def test_memory_and_param_layout_rejected_before_allocation(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        context = Mock(activate=nullcontext, sm_arch="sm_70")
        with patch.object(cuda, "mem_get_info", return_value=(1024, 1 << 30)), \
                patch.object(cuda, "Stream") as stream, patch.object(cuda, "mem_alloc") as alloc:
            with self.assertRaises(MemoryError):
                PeriodicBlockStepper(config(), layout, context=context)
            stream.assert_not_called()
            alloc.assert_not_called()
        module = Mock()
        module.get_global.return_value = (Mock(), 16)
        with patch.object(cuda, "mem_get_info", return_value=(1 << 30, 1 << 30)), \
                patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
                patch.object(compiler, "compile_kernel") as compile_kernel, \
                patch.object(compiler, "cleanup_build_dir"), \
                patch.object(compiler, "load_module", return_value=module), \
                patch.object(cuda, "mem_alloc") as alloc, patch.object(cuda, "memcpy_htod") as upload:
            with self.assertRaisesRegex(RuntimeError, "layout mismatch"):
                PeriodicBlockStepper(config(), layout, context=context)
            self.assertEqual(compile_kernel.call_args.kwargs["source"], "blocks.cu")
            alloc.assert_not_called()
            upload.assert_not_called()

    def test_static_circle_tables(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        tables = implementation._static_circles(layout, (circle(),))
        links = CircleGeometry(*CENTER, RADIUS).build_cut_links(NX, NY, domain_flags=dense_flags())
        self.assertEqual(len(tables["cell"]), len(links))
        np.testing.assert_array_equal(tables["q"], np.array([l.q for l in links], np.float32))
        solid_cells = layout.owner[tables["solid"]].astype(np.intp)
        self.assertTrue((tables["flags"][solid_cells] & 2).all())
        self.assertEqual(int(np.count_nonzero(tables["flags"] & 2)), int(tables["solid"].sum()))

    def fake_backend(self):
        backend = PeriodicBlockStepper.__new__(PeriodicBlockStepper)
        backend.layout = BlockLayout2D.build(NY, NX, 12, 16)
        backend.ctx = Mock(activate=nullcontext)
        backend._lock = RLock()
        backend._status = "ready"
        backend._step_count = 0
        backend._ddf_shifting = False  # unshifted storage, as for the default config
        backend._stream = Mock()
        backend._module = Mock()
        backend._copy = backend._step = Mock()
        backend.boundary, backend._curved, backend._tables, backend._links = None, None, {}, 0
        backend._dst, backend._src = Mock(), Mock()
        backend._block, backend._grid, backend._copy_grid = (32, 1, 1), (1, 1, 1), (1, 1)
        backend._allocations = [Mock(), Mock()]
        backend._current, backend._spare = backend._allocations
        return backend

    def test_initialize_rejects_bad_state_before_upload(self):
        backend = self.fake_backend()
        good = initial_state()
        bad = (good[:, :-1], np.where(np.arange(NX) == 3, np.nan, good), -good)
        with patch.object(cuda, "memcpy_htod") as upload:
            for g in bad:
                with self.assertRaises(ValueError):
                    backend.initialize(g)
            upload.assert_not_called()
            backend.initialize(good)
            self.assertEqual(upload.call_count, 2)
            self.assertTrue(np.isnan(upload.call_args_list[1].args[1]).all())
        backend.close()

    def test_failed_launch_and_download_make_backend_unusable(self):
        for failure in ("step", "download"):
            backend = self.fake_backend()
            if failure == "step":
                backend._step = Mock(side_effect=RuntimeError("launch failed"))
                with self.assertRaisesRegex(RuntimeError, "launch failed"):
                    backend.step()
            else:
                with patch.object(cuda, "memcpy_dtoh", side_effect=RuntimeError("read failed")):
                    with self.assertRaisesRegex(RuntimeError, "read failed"):
                        backend.download()
            for operation in (backend.download, backend.step):
                with self.assertRaisesRegex(RuntimeError, "failed"):
                    operation()
            backend.close()
            for allocation in backend_allocations(backend):
                allocation.free.assert_called_once()

    def test_activation_failure_retains_resources_for_retry(self):
        backend = self.fake_backend()
        allocations, stream, module = backend._allocations[:], backend._stream, backend._module
        unavailable = MagicMock()
        unavailable.__enter__.side_effect = RuntimeError("activation failed")
        backend.ctx.activate = Mock(side_effect=[unavailable, nullcontext()])
        with self.assertRaisesRegex(RuntimeError, "activation failed"):
            backend.close()
        self.assertEqual(backend._status, "failed")
        self.assertEqual(backend._allocations, allocations)
        self.assertIs(backend._stream, stream)
        self.assertIs(backend._module, module)
        backend.close()
        self.assertEqual(backend._status, "closed")
        self.assertIsNone(backend._module)
        self.assertIsNone(backend._copy)
        self.assertIsNone(backend._step)
        for allocation in allocations:
            allocation.free.assert_called_once()


def backend_allocations(backend):
    return (backend._current, backend._spare)


class PeriodicBlockGPU(unittest.TestCase):
    def setUp(self):
        self.assertEqual(Path(implementation.__file__).resolve(),
                         ROOT / "src/CelerisLab/lbm/block_stepper.py")

    def tearDown(self):
        rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
        print(f"host_peak_rss={rss:.3f} MiB", flush=True)

    def test_splits_bit_identical_and_fp64_oracle(self):
        from CelerisLab.cuda.context import CudaContext
        g0 = initial_state()
        results = []
        with CudaContext(0) as context:
            for by, bx in SPLITS:
                layout = BlockLayout2D.build(NY, NX, by, bx)
                with PeriodicBlockStepper(config(), layout, context=context) as backend:
                    backend.initialize(g0)
                    backend.step(STEPS)
                    self.assertEqual(backend.step_count, STEPS)
                    results.append(backend.download())
        for (by, bx), g in zip(SPLITS, results):
            self.assertTrue(np.isfinite(g).all(), (by, bx))
            np.testing.assert_array_equal(g, results[0], err_msg=f"split {(by, bx)}")
        error = np.abs(results[0].astype(np.float64) - oracle(g0, STEPS)).max()
        before, after = totals(g0), totals(results[0])
        print(f"max_abs_oracle_error={error:.3e} gate={ORACLE_ATOL:.0e}", flush=True)
        print("drift mass={:.3e} jx={:.3e} jy={:.3e}".format(
            *(a - b for a, b in zip(after, before))), flush=True)
        self.assertLessEqual(error, ORACLE_ATOL)

    def test_static_circle_across_block_seams(self):
        from CelerisLab.cuda.context import CudaContext
        flags = dense_flags()
        solid = ((flags & 2) != 0).reshape(NY, NX)
        links = CircleGeometry(*CENTER, RADIUS).build_cut_links(NX, NY, domain_flags=flags)
        for by, bx in SPLITS[1:]:  # fixture must exercise receiver-local halo staging
            layout = BlockLayout2D.build(NY, NX, by, bx)
            mapped = layout.map_cut_links(links)
            cells = int(np.prod(layout.local_shape))
            core = set(layout.owner.ravel().tolist())
            writes = {int(w) // (9 * cells) * cells + int(w) % cells for w in mapped.write}
            self.assertTrue(writes - core, (by, bx))
            self.assertTrue({int(d) for d in mapped.donor if d != NO_DONOR} - core, (by, bx))
        g0 = initial_state()
        results = []
        with CudaContext(0) as context:
            for by, bx in SPLITS:
                layout = BlockLayout2D.build(NY, NX, by, bx)
                with PeriodicBlockStepper(config(), layout, context=context, objects=[circle()]) as backend:
                    for key, values, dtype in (("q", [l.q for l in links], np.float32),
                                               ("dir", [l.dir for l in links], np.uint8),
                                               ("fallback", [l.fallback for l in links], np.uint8)):
                        np.testing.assert_array_equal(backend.boundary[key], np.array(values, dtype))
                    backend.initialize(g0)
                    backend.step(STEPS)
                    results.append(backend.download())
        for (by, bx), g in zip(SPLITS, results):
            self.assertTrue(np.isfinite(g[:, ~solid]).all(), (by, bx))
            self.assertTrue(np.isnan(g[:, solid]).all(), (by, bx))
            np.testing.assert_array_equal(g, results[0], err_msg=f"split {(by, bx)}")
        expected = circle_oracle(g0, STEPS, solid, links)
        error = np.abs(results[0][:, ~solid].astype(np.float64) - expected[:, ~solid]).max()
        start = g0.copy()
        start[:, solid] = 0
        before, after = totals(start), totals(np.nan_to_num(results[0]))
        print(f"circle max_abs_oracle_error={error:.3e} gate={ORACLE_ATOL:.0e} links={len(links)}", flush=True)
        print("circle drift mass={:.3e} jx={:.3e} jy={:.3e}".format(
            *(a - b for a, b in zip(after, before))), flush=True)
        self.assertLessEqual(error, ORACLE_ATOL)


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