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

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

KERNELS = Path(__file__).resolve().parents[2] / "src/CelerisLab/lbm/kernels"


@pytest.mark.parametrize("collision", ["SRT", "MRT"])
def test_patch_accepts_implemented_models(collision):
    FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, collision=collision), (8, 6, 16, 14))


@pytest.mark.parametrize("options", [dict(collision="TRT"), dict(collision="MRT", les_enabled=True)])
def test_patch_rejects_unimplemented(options):
    with pytest.raises(ValueError):
        FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, **options), (8, 6, 16, 14))


def test_srt_statements_and_guards_preserved():
    """MRT is an additive compile-time branch; SRT statements and other guards remain."""
    blocks = (KERNELS / "blocks.cu").read_text()
    channel = (KERNELS / "step/block_channel.cu").read_text()
    refinement = (KERNELS / "refinement.cu").read_text()
    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);"):
        assert text in blocks
    assert "collide_srt(f, feq, forcing, omega);" in channel
    assert "collide_srt_no_force(f, feq, d_params.omega);" in refinement
    for source in (blocks, refinement):
        assert "(COLLISION_MODEL != 0 && COLLISION_MODEL != 2)" in source
        assert "USE_LES" in source and "STORE_PRECISION != 0" in source
    assert "USE_DDF_SHIFTING ||" in refinement  # dense/mixed shifted stays rejected
    for source in (blocks, channel, refinement):
        assert source.count("collide_mrt_no_force(") == 1  # one MRT-only collision site each
    for source in (blocks, refinement):
        assert '#if COLLISION_MODEL == 2\n#include "operators/collision_mrt.cuh"\n#endif' in source