"""Single blocked level, FixedPatchStepper(refinements=[]): CPU contracts plus short GPU cases.

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
import unittest
from unittest.mock import 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
from CelerisLab.lbm import refinement_stepper as implementation
from CelerisLab.lbm.blocks import NO_DONOR, BlockLayout2D, static_circle_tables
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
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()


def single(cfg=None, *, context, block=(12, 16), objects=(), **kw):
    """One periodic blocked level (the former PeriodicBlockStepper capability)."""
    return FixedPatchStepper(cfg or config(), context=context, refinements=[], level_blocks=[block],
                             level_objects=[tuple(objects)], **kw)


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

    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):
        for cfg, kw in ((None, dict(block=None)),
                        (None, dict(coarse_boundary="channel")),
                        (None, dict(coarse_boundary="slot_reservoir")),
                        (config(collision="MRT"), {}),
                        (config(les_enabled=True), {}),
                        (None, dict(explosion="projected_linear")),
                        (None, dict(observe_raw_exchange=True)),
                        (None, dict(observe_raw_exchange=True, observe_window_exchange=True)),
                        (None, dict(level_jets=[(object(),)]))):
            with self.subTest(kw=kw), self.assertRaises(ValueError):
                single(cfg, context=self.forbidden_context(), **kw)
        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):
                single(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):
                single(context=self.forbidden_context(), objects=[circle(center, radius)])
        with patch.object(implementation, "static_circle_tables") as prepare:
            with self.assertRaisesRegex(AssertionError, "CUDA reached"):  # body-free path
                single(context=self.forbidden_context(), objects=iter(()))
            prepare.assert_not_called()
        with self.assertRaisesRegex(AssertionError, "CUDA reached"):
            single(context=self.forbidden_context(), objects=[circle()])
        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"):
                single(context=self.forbidden_context(), objects=[circle()])
        with patch.object(implementation, "MAX_GRID_DIM", 3):
            with self.assertRaisesRegex(ValueError, "launch limits"):
                single(context=self.forbidden_context())

    def test_memory_and_param_layout_rejected_before_allocation(self):
        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.assertRaisesRegex(MemoryError, "exceed free device memory"):
                single(context=context, objects=[circle()])
            stream.assert_not_called()
            alloc.assert_not_called()
        module = Mock()
        module.get_global.return_value = (Mock(), 16)
        with patch.object(cuda, "Stream"), patch.object(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), 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"):
                single(context=context)
            self.assertEqual(compile_kernel.call_args.kwargs["source"], "blocked.cu")
            alloc.assert_not_called()
            upload.assert_not_called()

    def test_static_circle_tables(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        tables = static_circle_tables(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()))


class SingleLevelGPU(unittest.TestCase):
    def setUp(self):
        self.assertEqual(Path(implementation.__file__).resolve(),
                         ROOT / "src/CelerisLab/lbm/refinement_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 split in SPLITS:
                with single(context=context, block=split) as backend:
                    backend.initialize(g0)
                    backend.step(STEPS)
                    self.assertEqual(backend.step_count, STEPS)
                    (g,) = backend.download()
                    results.append(g)
        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 split in SPLITS:
                with single(context=context, block=split, objects=[circle()]) as backend:
                    tables = backend._boundaries[0].tables
                    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(tables[key], np.array(values, dtype))
                    backend.initialize(g0)
                    backend.step(STEPS)
                    (g,) = backend.download()
                    results.append(g)
        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()