"""MRT reuse GPU gates. Caller owns GPU admission (CUDA_VISIBLE_DEVICES).

Run: python -m unittest discover -s tests/validation/multigrid -p test_refinement_mrt_gpu.py
Parity compares the blocked MRT channel root with uniform MRT OneStep: the existing
pressure-NEQ outlet branch differs by collision model, so SRT is never the reference.
AMR checks are finite-evolution/split-invariance smoke tests, not benchmark validation.
"""
from dataclasses import replace
import unittest
import numpy as np

from CelerisLab.config import LBMConfig
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.population_storage import _encode_physical
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from test_channel_root_gpu import run_root, state
from test_refinement_blocked_base import raw, exact


class RefinementMRTGPU(unittest.TestCase):
    def test_uniform_matched_mrt_root(self):
        with CudaContext(0) as context:
            for shifted in (False, True):
                for profile in ("uniform", "parabolic"):
                    cfg = LBMConfig(nx=24, ny=20, viscosity=.1, velocity=.025, collision="MRT",
                                    ddf_shifting=shifted, inlet_profile=profile)
                    initial = _encode_physical(state(cfg.ny, cfg.nx), shifted)
                    expected = run_root(cfg, context, initial)
                    self.assertTrue(np.isfinite(expected).all())
                    for block in ((20, 24), (5, 6), (4, 1)):
                        np.testing.assert_array_equal(run_root(cfg, context, initial, block), expected)

    def test_amr_mrt_smoke(self):
        from CelerisLab.body.geometry.circle import CircleGeometry
        from CelerisLab.body.objects import SimObject
        center = (24., 16.)
        body = SimObject(0, geometry=CircleGeometry(*center, 2.15), center=center, radius=2.15)
        with CudaContext(0) as context:
            for boundary in ("periodic", "channel"):
                for shifted in (False, True):
                    cfg = LBMConfig(nx=48, ny=32, viscosity=.1, velocity=.015 if boundary == "channel" else 0.,
                                    collision="MRT", ddf_shifting=shifted)
                    reference = None
                    for block in ((32, 48), (8, 8), (4, 1)):
                        with FixedPatchStepper(cfg, (8, 6, 40, 26), context=context, coarse_boundary=boundary,
                                               base_block=block, fine_block=(8, 8), fine_objects=(body,)) as backend:
                            backend.initialize(*[state(*level.shape[1:]) for level in backend._levels])
                            backend.step(8)
                            raw = backend._download_raw()
                            for a, mask in zip(raw.values, backend._state):
                                self.assertTrue(np.isfinite(a[:, mask]).all())
                            if reference is None: reference = raw.values
                            else:
                                for a, b in zip(reference, raw.values): np.testing.assert_array_equal(a, b)

    def test_dense_versus_blocked_mrt_two_level(self):
        """Unshifted MRT: all-dense PatchCoarseStep/PatchFineStep vs mixed/all-blocked storage."""
        from test_refinement_blocked_base import BASE_BLOCKS, FINE_BLOCK
        from test_refinement_blocked_fine import NX, NY, BOUNDS
        import test_refinement_blocked_fine as fixture
        cfg = LBMConfig(nx=NX, ny=NY, viscosity=.1, collision="MRT")
        with CudaContext(0) as context:
            results = {}
            layouts = ((None, None), *((base, FINE_BLOCK) for base in (None, *BASE_BLOCKS)))
            for base_block, fine_block in layouts:
                with FixedPatchStepper(cfg, BOUNDS, context=context, fine_block=fine_block,
                                       base_block=base_block) as backend:
                    backend.initialize(*[fixture.state(*level.shape[1:], seed) for seed, level in enumerate(backend._levels)])
                    backend.step(4)
                    results[base_block, fine_block] = raw(backend)
            for layout in layouts[1:]:
                for level, (a, b) in enumerate(zip(results[layout], results[None, None])):
                    exact(self, a, b, f"layout={layout} level={level}")

    def test_shifted_all_blocked_three_level_mrt_smoke(self):
        """Per-level shear rate omega_l = 1/(3*2**l*nu + 1/2); finite ACTIVE leaves."""
        cfg = LBMConfig(nx=16, ny=16, viscosity=.1, collision="MRT", ddf_shifting=True)
        kw = dict(fine_block=(4, 4), base_block=(4, 4), child=(8, 8, 16, 16), child_block=(4, 4))
        with CudaContext(0) as context:
            with FixedPatchStepper(cfg, (4, 4, 12, 12), context=context, **kw) as backend:
                self.assertEqual(len(backend._level_cfgs), 3)
                for level, level_cfg in enumerate(backend._level_cfgs):
                    checked = replace(level_cfg)
                    checked.validate()  # omega is derived only by validate()
                    self.assertEqual(checked.collision, "MRT")
                    self.assertTrue(checked.ddf_shifting)
                    self.assertTrue(np.isclose(checked.omega, 1 / (3 * 2**level * .1 + .5), rtol=1e-14))
                backend.initialize(*[state(*level.shape[1:]) for level in backend._levels])
                backend.step(4)
                values = backend._download_raw().values
                for a, mask in zip(values, backend._state):
                    self.assertTrue(np.isfinite(a[:, mask]).all())


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