"""Host guards and source contract for MRT reuse in refined/blocked stepping (no CUDA)."""
from pathlib import Path
import unittest

from CelerisLab.config import LBMConfig
from CelerisLab.lbm.refinement import FixedPatch2D

KERNELS = Path(__file__).resolve().parents[2] / "src/CelerisLab/lbm/kernels"
BOUNDS = (8, 6, 16, 14)


class RefinementMRTHost(unittest.TestCase):
    def test_patch_accepts_implemented_models(self):
        for collision in ("SRT", "MRT"):
            with self.subTest(collision=collision):
                FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, collision=collision), BOUNDS)

    def test_patch_rejects_unimplemented(self):
        for options in (dict(collision="TRT"), dict(collision="MRT", les_enabled=True)):
            with self.subTest(options=options), self.assertRaises(ValueError):
                FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, **options), BOUNDS)

    def test_srt_statements_and_guards_preserved(self):
        """MRT is an additive compile-time branch; SRT statements and other guards remain."""
        blocks = (KERNELS / "blocked.cu").read_text() + (KERNELS / "step/block_step.cu").read_text()
        channel = (KERNELS / "step/block_channel.cu").read_text()
        block_collision = (KERNELS / "operators/block_collision.cuh").read_text()
        self.assertIn('#include "operators/block_collision.cuh"', blocks)
        self.assertEqual(blocks.count("block_collide(f);"), 1)
        for text in ("compute_feq_from_momentum(drho, jx, jy, ux, uy, feq);",
                     "compute_feq(rho, ux, uy, feq);",
                     "collide_srt_no_force(f, feq, d_params.omega);"):
            self.assertIn(text, block_collision)
        self.assertIn("collide_srt(f, feq, forcing, omega);", channel)
        for source in (blocks,):
            self.assertIn("(COLLISION_MODEL != 0 && COLLISION_MODEL != 2)", source)
            self.assertIn("USE_LES", source)
            self.assertIn("STORE_PRECISION != 0", source)
            self.assertIn('#if COLLISION_MODEL == 2\n#include "operators/collision_mrt.cuh"\n#endif', source)
        # One shared unforced MRT body; each blocked/patch path calls it exactly once.
        mrt = (KERNELS / "operators/collision_mrt.cuh").read_text()
        self.assertEqual(mrt.count("void collide_mrt_unforced("), 1)
        for source in (blocks + "\n" + block_collision, channel):
            self.assertEqual(source.count("collide_mrt_unforced(f);"), 1)
            self.assertNotIn("collide_mrt_no_force(", source)
            self.assertNotIn("compute_rho_u(f, rho, ux, uy);\n            collide_mrt", source)
            self.assertNotRegex(source, r"compute_rho_u\(f, rho, ux, uy\);\s*collide_mrt")


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