"""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.
"""
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


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)


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