"""Two coarse/four fine steps: isolated fine circle on blocked storage, no flux-accuracy claim.

Run only through the bounded GPU runner after a fresh resource check. The fine
step hook reads the actual output buffer; it never changes the stepping order.

Retired with the dense refinement route (accepted losses, see the D2 commit):
the dense-kernel per-hop flux ledger (fine/coarse transfer + boundary budget),
the CurvedLinkSoA reflection oracle and the CELERIS_REFINEMENT_AUDIT collision
counts. The blocked test keeps the four-hop comparison with the uniform solver,
the raw-exchange mass sign and passive == zero actuation; raw exchange covers
the boundary part only, not full population flux closure.
"""
import importlib.util
from pathlib import Path
from unittest.mock import Mock, patch

import numpy as np
import pycuda.driver as cuda
import pytest

import CelerisLab
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm import refinement_stepper as implementation
from CelerisLab.lbm.descriptors import FLUID, SOLID


ROOT = Path(__file__).resolve().parents[3]
CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1])
CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1])
W = np.array([4/9]+[1/9]*4+[1/36]*4)
MOMENTS = np.stack((np.ones(9), CX, CY))
BOUNDS = (4, 4, 20, 20)


def read(pointer, shape, dtype=np.float32):
    values = np.empty(shape, dtype)
    cuda.memcpy_dtoh(values, pointer)
    return values


def rest(shape):
    return np.broadcast_to(W[:, None, None], (9, *shape)).astype(np.float32).copy()


def surface(speed):
    return None if speed is None else dict(angle=.7, half_width=.85, normal_velocity=speed)


def object_spec(speed):
    return SimObject(0, CircleGeometry(12.25, 12.25, 2.15), center=(12.25, 12.25),
                     radius=2.15, surface_actuation=surface(speed))


def uniform_reference(speed, roi):
    # Reuse the existing explicit config helper without importing unrelated tests.
    spec = importlib.util.spec_from_file_location("fine_surface_cases", ROOT/"tests/integration/_cases.py")
    cases = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(cases)
    with cases.make_sim(nx=48, ny=48, viscosity=.1, velocity=0, rho=1) as sim:
        sim.add_body("circle", center=(24, 24), radius=4.3, surface_actuation=surface(speed))
        sim.initialize()
        # Uniform 48^2 has the same dx=dt=.5 representation. Its distant channel
        # boundaries cannot influence the circle ROI within these four hops.
        traces = []
        for _ in range(4):
            sim.run(1)
            traces.append(sim.get_ddf().reshape(9, 48, 48)[:, 4:44, 4:44][:, roi].copy())
        return traces


def run_blocked_patch(speed):
    """All-blocked run (whole-level tiles) of the same case; fine ROI after each fine hop."""
    cfg = LBMConfig(nx=24, ny=24, viscosity=.05, velocity=0, rho=1)
    source = object_spec(speed)
    with CudaContext(0) as context:
        with implementation.FixedPatchStepper(cfg, context=context, refinements=[BOUNDS],
                                              level_blocks=[(24, 24), (40, 40)],
                                              level_objects=[(), (source,)],
                                              observe_raw_exchange=True) as backend:
            c, f = backend._levels
            flags = backend.boundaries[1].flags.reshape(40, 40)
            active = np.zeros((40, 40), bool)
            active[4:-4, 4:-4] = True
            active &= ((flags & FLUID) != 0) & ((flags & SOLID) == 0)
            yy, xx = np.mgrid[:40, :40]
            roi = active & ((xx-20)**2+(yy-20)**2 <= 9**2)
            backend.initialize(rest((24, 24)), rest((40, 40)))
            traces, original = [], f.kernel

            def fine(*args, **kwargs):
                original(*args, **kwargs)
                backend._stream.synchronize()
                advanced = f.layout.to_dense(read(args[1], f.layout.blocked_shape))
                assert np.isfinite(advanced[:, active]).all()
                assert np.isnan(advanced[:, (flags & SOLID) != 0]).all()
                traces.append(advanced[:, roi].copy())

            with patch.object(f, "kernel", fine):
                backend.step(2)
            assert len(traces) == 4 and backend.step_count == 2
            mass = backend.read_raw_exchange(0, level=1)["total"]["mass_into_fluid"]
            return traces, roi, backend.download(), float(mass)


@pytest.mark.parametrize("speed", [0., .002, -.002])
def test_fine_surface_blocked_matches_uniform_four_hops(speed):
    """Blocked storage: fine ROI after each of four hops against the uniform 48^2 solver
    (same tolerances as the dense test); blowing/suction drives mass with its sign;
    zero actuation equals the passive circle bit for bit."""
    actual, roi, final, mass = run_blocked_patch(speed)
    reference = uniform_reference(speed, roi)
    for got, wanted in zip(actual, reference):
        np.testing.assert_allclose(got, wanted, atol=3e-6, rtol=0)
        np.testing.assert_allclose(MOMENTS@got.astype(float), MOMENTS@wanted.astype(float),
                                   atol=2e-5, rtol=0)
    if speed:
        assert speed*mass > 0
    else:
        passive, passive_roi, passive_final, _ = run_blocked_patch(None)
        np.testing.assert_array_equal(passive_roi, roi)
        for got, wanted in zip(actual+list(final), passive+list(passive_final)):
            np.testing.assert_array_equal(got, wanted)
    print(f"blocked speed={speed} four hops vs uniform PASS mass_into_fluid={mass:+.9g}", flush=True)


def test_cross_interface_donor_rejects_before_cuda():
    cfg = LBMConfig(nx=24, ny=24, viscosity=.05)
    source = SimObject(0, CircleGeometry(5.05, 12.25, .7),
                       center=(5.05, 12.25), radius=.7, surface_actuation=surface(.002))
    context = Mock()
    context.activate.side_effect = AssertionError("must reject before CUDA activation")
    with patch.object(cuda, "mem_alloc") as allocate:
        with pytest.raises(ValueError, match="donor crosses"):
            implementation.FixedPatchStepper(cfg, context=context, refinements=[BOUNDS],
                                             level_blocks=[(24, 24), (40, 40)],
                                             level_objects=[(), (source,)])
        context.activate.assert_not_called()
        allocate.assert_not_called()