"""Root-channel host contract and actual-header CPU parity (no CUDA evidence)."""
from dataclasses import replace
import numpy as np
import pytest

from CelerisLab.config import LBMConfig
from CelerisLab.lbm.blocks import BlockLayout2D
from CelerisLab.lbm.channel import channel_masks, validate_channel_config, validate_channel_transfers
from CelerisLab.lbm.refinement import FixedPatch2D
from CelerisLab.lbm.refinement_tables import blocked_base_tables
from .test_numerical_operators import compile_run


@pytest.mark.parametrize("shift", [0, 1])
@pytest.mark.parametrize("profile", [0, 1])
@pytest.mark.parametrize("clamp,damping", [(1, ".5f"), (0, ".8f")])
def test_actual_kernel_uniform_parity(tmp_path, shift, profile, clamp, damping):
    """Nonuniform ghosts, corners, top wall, global profile and outlet halo donor."""
    defines = dict(USE_DDF_SHIFTING=shift, INLET_PROFILE=profile,
                   OUTLET_BACKFLOW_CLAMP=clamp, OUTLET_SRT_NEQ_DAMP=damping)
    uniform = np.array(compile_run(tmp_path, "channel_root.cpp", dict(
        defines, UNIFORM=1, NX=24, NY=20)), np.float32)
    assert np.isfinite(uniform).all()
    for by, bx in ((20, 24), (5, 6), (4, 1)):
        blocked = np.array(compile_run(tmp_path, "channel_root.cpp", dict(
            defines, UNIFORM=0, NX=bx+2, NY=by+2)), np.float32)
        np.testing.assert_array_equal(blocked, uniform)


def transfer(bounds, block=(4, 6)):
    cfg = LBMConfig(nx=24, ny=20, viscosity=.1)
    pair = FixedPatch2D.from_config(cfg, bounds)
    layout = BlockLayout2D.build(20, 24, *block)
    targets = blocked_base_tables(pair, layout, np.argsort(pair.coalesce_coarse))[1]
    return pair, layout, targets


def test_masks_and_minimal_footprint():
    fluid, bulk = channel_masks((20, 24))
    assert fluid.sum() == 18*22 and bulk.sum() == 16*22
    assert fluid[1, 1] and not bulk[1, 1]
    assert bulk[2, 1] and not fluid[0, 0]
    # No tile alignment or artificial large padding: closest legal x donor is x=1.
    validate_channel_transfers(*transfer((3, 4, 19, 16)))
    for bounds in ((2, 4, 19, 16), (3, 3, 19, 16), (3, 4, 22, 16), (3, 4, 19, 17)):
        with pytest.raises(ValueError, match="bulk"):
            validate_channel_transfers(*transfer(bounds))


def test_actual_closure_read_conflicts():
    pair, layout, targets = transfer((6, 6, 18, 14))
    validate_channel_transfers(pair, layout, targets)
    for target in (layout.slot(int(layout.owner[5, -1]) - 1, 2),  # integer under NumPy 1 and 2
                   layout.slot(layout.owner[1, 4], 4),
                   layout.slot(layout.owner[-2, 4], 3)):
        with pytest.raises(ValueError, match="immutable"):
            validate_channel_transfers(pair, layout, np.append(targets, target))


def test_channel_scope():
    cfg = LBMConfig(nx=24, ny=20, viscosity=.1)
    validate_channel_config(cfg, [(4, 6), (4, 4)])
    with pytest.raises(ValueError, match="all-blocked"):
        validate_channel_config(cfg, [None, (4, 4)])
    for change in (dict(inlet_scheme="equilibrium"), dict(outlet_mode="zero_gradient"),
                   dict(y_wall_bc="free_slip"), dict(wall_slot_width=2)):
        with pytest.raises(ValueError, match="channel requires"):
            validate_channel_config(replace(cfg, **change), [(4, 6), (4, 4)])


def test_evolved_ghost_snapshot_and_physical_download():
    from contextlib import nullcontext
    from unittest.mock import Mock
    from .test_refinement_three_level_multipatch import host_module
    host = host_module()
    storage = {}
    host.compiler.build_module.return_value.get_global.side_effect = lambda name: (name, 12 if name == "d_block_channel" else 20)
    host.cuda.mem_alloc.side_effect = lambda size: Mock()
    host.cuda.memcpy_htod.side_effect = lambda ptr, a: storage.__setitem__(id(ptr), np.array(a, copy=True))
    host.cuda.memcpy_dtoh.side_effect = lambda a, ptr: np.copyto(a, storage[id(ptr)])
    context = Mock(activate=nullcontext, sm_arch="sm_70")
    cfg = LBMConfig(nx=24, ny=20, viscosity=.1, velocity=.025)
    backend = host.FixedPatchStepper(cfg, context=context, coarse_boundary="channel",
                                    refinements=[(6, 6, 18, 14)], level_blocks=[(4, 6), (4, 4)])
    try:
        weights = np.array([4/9]+[1/9]*4+[1/36]*4)[:, None, None]
        arrays = [np.broadcast_to(weights, level.shape).copy() for level in backend._levels]
        arrays[0][:, 0, 0] *= 1.001
        backend.initialize(*arrays)
        snapshot = backend._download_raw()
        assert np.isfinite(snapshot.values[0][:, 0, :]).all()
        assert np.isnan(backend.download()[0][:, 0, :]).all()
        backend._initialize_raw(snapshot)
        for a, b in zip(snapshot.values, backend._download_raw().values):
            np.testing.assert_array_equal(a, b)
        arrays[0][0, 0, 0] = np.nan
        with pytest.raises(ValueError, match="finite"):
            backend.initialize(*arrays)
        assert backend._levels[0].kernel == host.compiler.build_module.return_value.get_function("BlockChannelStep")
        assert [call.kwargs["block_root_bc"] for call in host.compiler.build_module.call_args_list] == [2, 0]
    finally:
        backend.close()