"""Translating-circle host contract (docs/moving_bodies.md 4a, 5, 7).

Real Simulation entrypoints with device doubles; no CUDA context or GPU work.
"""
import math
from types import SimpleNamespace

import pytest

from CelerisLab.body.motion import MAX_STEP, R_MIN, MovingCircle, Translation

BAD = (float("nan"), float("inf"))


def circle(box=(-2., 2., -3., 3.), speed=.1):
    return MovingCircle((50., 40.), 10., Translation(box, speed))


def test_translation_and_radius_validation():
    for box in ((0., 1., -1., 1.)[:3], (.5, 1., -1., 1.), (-1., 1., BAD[0], 1.)):
        with pytest.raises(ValueError):
            Translation(box, .1)
    for speed in (0., MAX_STEP + 1e-9) + BAD:
        with pytest.raises(ValueError):
            Translation((0., 0., -1., 1.), speed)
    with pytest.raises(ValueError):
        MovingCircle((0., 0.), R_MIN - .1, Translation((0., 0., -1., 1.), .1))


def test_targets_validate_against_applied_and_last_wins():
    c = circle()
    c.set_target((0., .3), (0., .05))
    c.set_target((0., .5), (0., .05))  # second setter replaces, checked vs applied (0, 0)
    assert c.pending == ((0., .5), (0., .05))
    for offset, velocity in (((0., .51), (0., 0.)),           # Euclidean bound
                             ((.36, .36), (0., 0.)),          # |dc|_2 = .509
                             ((0., 3.1), (0., 0.)),           # outside box
                             ((0., .1), (0., .11)),           # speed bound
                             ((0., 0.), (0., .01)),           # hold with sliding velocity
                             ((BAD[0], 0.), (0., 0.)), ((0., 0.), (0., BAD[1]))):
        with pytest.raises(ValueError):
            c.set_target(offset, velocity)
        assert c.pending == ((0., .5), (0., .05))  # invalid replacement keeps the old target
    c.set_target((.35, .35), (0., 0.))  # |dc|_2 = .495 inside the bound
    assert c.applied == (0., 0.)


def test_record_commit_and_hold(monkeypatch):
    sim, writes, c = routed_advancing_simulation(monkeypatch)
    assert c.record() == ((50., 40.), (50., 40.), (0., 0.))
    sim.set_body(0, offset=(0., .25), vy=.05)
    assert c.record() == ((50., 40.), (50., 40.25), (0., .05))
    seen_at_sync = []
    sim.run(1, stream=SimpleNamespace(synchronize=lambda: seen_at_sync.append(c.applied)))
    assert seen_at_sync == [(0., 0.), (0., .25)]  # commit between step sync and hold sync
    uploads = [w[1][0] for w in writes if isinstance(w, tuple) and w[0] == "upload"]
    assert uploads == [((50., 40.), (50., 40.25), (0., .05)),
                       ((50., 40.25), (50., 40.25), (0., 0.))]
    assert c.applied == (0., .25) and c.pending is None  # offset-only commit
    sim.run(1, stream=STREAM)  # no new target: hold step
    uploads = [w[1][0] for w in writes if isinstance(w, tuple) and w[0] == "upload"]
    assert uploads[-1] == ((50., 40.25), (50., 40.25), (0., 0.)) and c.applied == (0., .25)
    sim.set_body(0, offset=(0., 0.), vy=-.05)  # zero offset is a real target
    assert c.pending == ((0., 0.), (0., -.05))
    sim.run(1, stream=STREAM)
    assert c.applied == (0., 0.) and c.pending is None


@pytest.mark.parametrize("stage", ("upload", "step", "synchronize"))
def test_failure_after_upload_attempt_is_fail_closed(monkeypatch, stage):
    sim, writes, c = routed_advancing_simulation(monkeypatch, fail="upload" if stage == "upload" else None)
    sim.set_body(0, offset=(0., .25), vy=.05)

    def fail(*a, **k):
        raise RuntimeError(f"{stage} failed")
    if stage == "step":
        sim.stepper._step = fail
    stream = SimpleNamespace(synchronize=fail) if stage == "synchronize" else STREAM
    with pytest.raises(RuntimeError, match=f"{stage} failed"):
        sim.run(1, stream=stream)
    assert c.failed and c.applied == (0., 0.) and c.pending == ((0., .25), (0., .05))
    assert sim._failed is not None and sim.stepper.blocked == sim._failed
    for later in (c.record, lambda: sim.set_body(0, offset=(0., .1), vy=0.)):
        with pytest.raises(RuntimeError, match="unknown"):
            later()
    with pytest.raises(RuntimeError, match="Simulation failed"):
        sim.run(1, stream=STREAM)
    with pytest.raises(RuntimeError, match="moving circle step failed"):
        sim.stepper.step(1, action_gpu=None, obs_gpu=None)


def test_band_covers_offset_box_with_margin():
    assert circle().band() == (50 - 2 - 10 - 2, 50 + 2 + 10 + 2, 40 - 3 - 10 - 2, 40 + 3 + 10 + 2)


# D2Q9 non-rest directions, paired order of lbm/descriptors.
_CX = (1, -1, 0, 0, 1, -1, 1, -1)
_CY = (0, 0, 1, -1, 1, -1, -1, 1)


def _inside(x: int, y: int, c, r: float) -> bool:
    return (x - c[0]) ** 2 + (y - c[1]) ** 2 < r * r


def empty_donor_cells(c_old, c_new, radius: float):
    """Fresh cells (inside c_old, outside c_new) whose donor set V(x) is empty."""
    r = radius
    xs = range(math.floor(min(c_old[0], c_new[0]) - r) - 1, math.ceil(max(c_old[0], c_new[0]) + r) + 2)
    ys = range(math.floor(min(c_old[1], c_new[1]) - r) - 1, math.ceil(max(c_old[1], c_new[1]) + r) + 2)
    empty = []
    for x in xs:
        for y in ys:
            if not _inside(x, y, c_old, r) or _inside(x, y, c_new, r):
                continue
            if not any(not _inside(x + cx, y + cy, c_old, r) and not _inside(x + cx, y + cy, c_new, r)
                       for cx, cy in zip(_CX, _CY)):
                empty.append((x, y))
    return empty


def enumerate_donor_coverage(radius: float, positions: int = 8, directions: int = 16) -> int:
    """Count empty-donor fresh cells over sub-cell centres and moves of MAX_STEP.

    Evidence for the declared bounds (radius >= R_MIN, |dc|_2 <= MAX_STEP),
    not a proof; the device keeps the ``no_donor`` counter.
    """
    count = 0
    for i in range(positions):
        for j in range(positions):
            c_old = (i / positions, j / positions)
            for k in range(directions):
                a = 2.0 * math.pi * k / directions
                c_new = (c_old[0] + MAX_STEP * math.cos(a), c_old[1] + MAX_STEP * math.sin(a))
                count += len(empty_donor_cells(c_old, c_new, radius))
    return count


@pytest.mark.parametrize("radius", (R_MIN, 4.3, 7.5, 12.0))
def test_refill_donor_sets_nonempty_within_bounds(radius):
    # Evidence for R >= R_MIN and |dc|_2 <= MAX_STEP on a fixed sample, not a proof.
    assert enumerate_donor_coverage(radius) == 0


pytest.importorskip("pycuda.driver")
from CelerisLab import simulation as module  # noqa: E402
from tests.unit.test_simulation_lifecycle import Owner, make_simulation  # noqa: E402


class ClosableOwner(Owner):
    """Owner double that refuses activation after close, like CudaContext.activate."""
    def activate(self):
        if self.is_closed:
            raise RuntimeError("CUDA context is closed")
        return super().activate()


def moving_simulation(monkeypatch, *, moving=True):
    sim = make_simulation(monkeypatch)
    obj = SimpleNamespace(motion=circle() if moving else None, obj_id=0,
                          center=(50., 40.), radius=10.)
    writes = []
    sim._failed, sim._initialized = None, True
    sim.bodies = SimpleNamespace(objects=[obj], count=1, jets=[], action_gpu=None, obs_gpu=None,
                                 has_pending_edit=lambda: False,
                                 _upload_action_async=lambda s: writes.append("action"),
                                 zero_obs_async=lambda s: writes.append("obs"),
                                 increment_obs_steps=lambda n: writes.append("count"),
                                 download_obs_full_async=lambda s: writes.append("download"))
    sim.stepper = SimpleNamespace(step=lambda n, **k: writes.append(("step", n)), blocked=None,
                                  moving=None)
    sim._assert_runtime_contracts = lambda: None
    return sim, writes


STREAM = SimpleNamespace(synchronize=lambda: None)


def test_positive_steps_rejected_before_device_writes(monkeypatch):
    sim, writes = moving_simulation(monkeypatch)
    sim.bodies.objects[0].motion.set_target((0., .1), (0., .01))
    sim.bodies.objects[0].motion.applied = (0., 0.)
    for call in (lambda: sim.run(1, stream=STREAM), lambda: sim.run(5, stream=STREAM),
                 lambda: sim.step(), lambda: sim.run(3, stream=STREAM, checkpoint_interval=1)):
        with pytest.raises(RuntimeError, match="moving kernels"):
            call()
    with pytest.raises(RuntimeError, match="moving kernels"):  # nominal at rest as well
        sim.bodies.objects[0].motion.pending = None
        sim.run(1)
    sim.bodies.objects[0].state = SimpleNamespace(omega=.02)  # spin does not lift the guard
    with pytest.raises(RuntimeError, match="moving kernels"):
        sim.run(2, stream=STREAM)
    assert writes == [] and sim.bodies.objects[0].motion.applied == (0., 0.)


def test_run_zero_keeps_pending_and_static_runs_proceed(monkeypatch):
    sim, writes = moving_simulation(monkeypatch)
    motion = sim.bodies.objects[0].motion
    motion.set_target((0., .1), (0., .01))
    sim.run(0, stream=STREAM)  # existing action/obs work only; no motion upload or commit
    assert motion.pending == ((0., .1), (0., .01)) and motion.applied == (0., 0.)
    assert ("step", 0) in writes
    static, writes = moving_simulation(monkeypatch, moving=False)
    static.save_checkpoint = lambda *a: writes.append("checkpoint")
    static.run(3, stream=STREAM)
    static.run(4, stream=STREAM, checkpoint_interval=2)
    assert writes.count(("step", 3)) == 1 and writes.count(("step", 2)) == 2


def test_direct_stepper_entry_is_guarded():
    from CelerisLab.lbm.stepper import LBMStepper
    launches = []
    module = SimpleNamespace(get_function=lambda name: lambda *a, **k: launches.append(name))
    cfg = SimpleNamespace(streaming="double_buffer", threads_per_block=128, is_d2q9=True,
                          nx=64, ny=32, nz=1)
    field = SimpleNamespace(ddf_gpu="ddf", temp_gpu="temp", n_curved=0)
    stepper = LBMStepper(field, module, cfg)
    assert stepper.blocked is None
    stepper.blocked = "translating circles need the moving kernels"
    for n in (1, 3):
        with pytest.raises(RuntimeError, match="moving kernels"):
            stepper.step(n, action_gpu=None, obs_gpu=None)
    assert launches == [] and (field.ddf_gpu, field.temp_gpu) == ("ddf", "temp")
    assert stepper.step_count == 0
    stepper.step(0, action_gpu=None, obs_gpu=None)  # zero steps stay a no-op
    assert launches == [] and stepper.step_count == 0


def test_moving_body_rejects_topology_and_checkpoint(monkeypatch):
    sim, _ = moving_simulation(monkeypatch)
    for call in (sim.sync_bodies, sim.save_checkpoint, lambda: sim.load_checkpoint("x.h5")):
        with pytest.raises(ValueError, match="translating circles"):
            call()


# -- Real-object routing, initialize rejections, packing, lifecycle -------------
import numpy as np  # noqa: E402

from CelerisLab.body.geometry.circle import CircleGeometry  # noqa: E402
from CelerisLab.body.manager import ObjectManager  # noqa: E402
from CelerisLab.body.objects import SimObject  # noqa: E402
from CelerisLab.lbm.descriptors import FLUID, SOLID  # noqa: E402


class Cfg:
    dim, is_d3q19, is_d2q9 = 2, False, True
    streaming, ddf_shifting, store_precision, y_wall_bc = "double_buffer", False, "FP32", "bounce_back"


def body(x, y, r, translation=None, **kw):
    obj = SimObject(obj_id=-1, geometry=CircleGeometry(x, y, r), center=(x, y), radius=r, **kw)
    if translation is not None:
        obj.motion = MovingCircle((x, y), r, translation)
    return obj


def routed_simulation(monkeypatch):
    sim = make_simulation(monkeypatch)
    sim._failed, sim._initialized, sim.lbm_cfg = None, False, Cfg()
    sim.stepper = SimpleNamespace(blocked=None)
    sim.bodies = ObjectManager(nx=120, ny=80, nz=1, nq=9, cfg=Cfg())
    return sim


def test_add_body_and_set_body_route_targets(monkeypatch):
    sim = routed_simulation(monkeypatch)
    bid = sim.add_body("circle", (50., 40.), 10., translation=Translation((-2., 2., -3., 3.), .1))
    still = sim.add_body("circle", (100., 40.), 5.)
    motion = sim.bodies.get(bid).motion
    sim._moving_tables = "tables"
    sim._attach_moving()
    assert sim.stepper.blocked is None
    sim.set_body(bid, offset=(0., .3), vy=.05)
    assert sim.stepper.blocked == module.PENDING_BLOCKED  # direct steps never consume a target
    sim.set_body(bid, offset=(0., .4), vy=.05)  # last wins
    for kwargs in (dict(offset=(0., .6)), dict(vy=.05), dict(vx=.05),
                   dict(offset=(0., .1), vy=.2), dict(offset=(0., 0.), vy=.01)):
        with pytest.raises(ValueError):
            sim.set_body(bid, **kwargs)
        assert motion.pending == ((0., .4), (0., .05))  # invalid input keeps the target
    sim.set_body(bid, omega=.01)  # omega-only keeps the pending target
    obj = sim.bodies.get(bid)
    assert obj.state.omega == .01 and (obj.state.vx, obj.state.vy) == (0., 0.)
    assert motion.pending == ((0., .4), (0., .05)) and motion.applied == (0., 0.)
    with pytest.raises(ValueError, match="not declared"):
        sim.set_body(still, offset=(0., .1))
    sim.set_body(still, vx=.01)  # stationary API unchanged
    assert sim.bodies.get(still).state.vx == .01 and sim.bodies.get(still).motion is None
    for kwargs in (dict(type="sensor"), dict(surface_actuation=dict(angle=0., half_width=.2,
                                                                    normal_velocity=.01))):
        kind = kwargs.pop("type", "circle")
        with pytest.raises(ValueError, match="translation"):
            sim.add_body(kind, (20., 20.), 5., translation=Translation((0., 0., 0., 0.), .1),
                         **kwargs)


class Tables:
    """MovingTables double: records uploads, steps and diagnostics reads."""
    DIAG = ("no_donor",)

    def __init__(self, log, no_donor=0., fail=None):
        self.log, self.no_donor, self.fail = log, no_donor, fail

    def upload_records(self, circles, stream):
        records = [m.record() for m, _ in circles]
        self.log.append(("upload", records))
        uploads = sum(1 for w in self.log if isinstance(w, tuple) and w[0] == "upload")
        if self.fail == "upload" or (self.fail == "hold_upload" and uploads == 2):
            raise RuntimeError("upload failed")
        return any(r[0] != r[1] for r in records)

    def zero_diagnostics(self, stream):
        self.log.append("zero_diag")

    def download_diagnostics(self):
        return np.array([[self.no_donor]], np.float32)


def advancing_simulation(monkeypatch, **tables):
    sim, writes = moving_simulation(monkeypatch)
    sim.bodies.objects[0].motion = circle()
    sim.stepper = SimpleNamespace(
        moving=Tables(writes, **tables), blocked=None,
        _step=lambda n, advance=False, **k: writes.append(("step", n, advance)))
    from CelerisLab.lbm.stepper import LBMStepper
    sim.stepper.step = LBMStepper.step.__get__(sim.stepper)  # real public entry
    sim._moving_tables = sim.stepper.moving
    return sim, writes, sim.bodies.objects[0].motion


def routed_advancing_simulation(monkeypatch, **tables):
    """Real body registration, target routing, run and owner hook; double device I/O."""
    sim = routed_simulation(monkeypatch)
    bid = sim.add_body("circle", (50., 40.), 10., translation=Translation((-2., 2., -3., 3.), .1))
    sim._initialized = True
    sim._assert_runtime_contracts = lambda: None
    writes = []
    monkeypatch.setattr(sim.bodies, "_upload_action_async", lambda s: writes.append("action"))
    monkeypatch.setattr(sim.bodies, "zero_obs_async", lambda s: writes.append("obs"))
    monkeypatch.setattr(sim.bodies, "download_obs_full_async", lambda s: writes.append("download"))
    sim.stepper = SimpleNamespace(
        moving=Tables(writes, **tables), blocked=None,
        _step=lambda n, advance=False, **k: writes.append(("step", n, advance)))
    from CelerisLab.lbm.stepper import LBMStepper
    sim.stepper.step = LBMStepper.step.__get__(sim.stepper)
    sim._moving_tables = sim.stepper.moving
    sim._attach_moving()
    return sim, writes, sim.bodies.get(bid).motion


def test_run_one_advances_and_commits_after_sync(monkeypatch):
    sim, writes, motion = advancing_simulation(monkeypatch)
    motion.set_target((0., .25), (0., .05))
    with pytest.raises(ValueError, match="run\\(1\\)"):  # multi-step with a target
        sim.run(2, stream=STREAM)
    assert writes == [] and motion.pending is not None
    sim.run(0, stream=STREAM)  # zero steps never consume the target
    assert motion.pending is not None
    sim.run(1, stream=STREAM)
    uploads = [w for w in writes if isinstance(w, tuple) and w[0] == "upload"]
    assert uploads[0][1] == [((50., 40.), (50., 40.25), (0., .05))]
    assert uploads[1][1] == [((50., 40.25), (50., 40.25), (0., 0.))]  # first hold after move
    assert ("step", 1, True) in writes
    assert motion.applied == (0., .25) and motion.pending is None
    writes.clear()
    sim.run(3, stream=STREAM)  # hold steps: no advance
    assert ("step", 3, False) in writes and motion.applied == (0., .25)


def test_zero_motion_target_is_consumed_by_run_one(monkeypatch):
    sim, writes, motion = advancing_simulation(monkeypatch)
    motion.set_target((0., 0.), (0., 0.))          # equal to applied: a hold target
    sim._attach_moving()
    assert sim.stepper.blocked == module.PENDING_BLOCKED
    sim.run(1, stream=STREAM)
    assert motion.pending is None and motion.applied == (0., 0.)
    assert sim.stepper.blocked is None and ("step", 1, False) in writes
    writes.clear()
    sim.run(3, stream=STREAM)                      # no stale latch
    assert ("step", 3, False) in writes


@pytest.mark.parametrize("streaming,d2q9,bound", [("double_buffer", True, True),
                                                  ("double_buffer", False, False),
                                                  ("esopull", True, False)])
def test_moving_kernels_bound_only_for_d2q9_double_buffer(monkeypatch, streaming, d2q9, bound):
    from CelerisLab.lbm import stepper as stepper_module
    monkeypatch.setattr(stepper_module.cuda, "mem_alloc", lambda n: object())
    moving = {"MovingCircleLinks", "MovingCircleCells"}

    def get_function(name):
        if name in moving and not (d2q9 and streaming == "double_buffer"):
            raise KeyError(name)                 # absent from D3Q19 / EsoPull modules
        return lambda *a, **k: None
    cfg = SimpleNamespace(streaming=streaming, threads_per_block=128, is_d2q9=d2q9,
                          nx=16, ny=8, nz=4, nq=9 if d2q9 else 19)
    stepper = stepper_module.LBMStepper(SimpleNamespace(), SimpleNamespace(get_function=get_function), cfg)
    assert (stepper.moving_fn is not None) == bound and (stepper.moving_cells_fn is not None) == bound


def test_moving_diagnostics_empty_shape_matches_layout(monkeypatch):
    from CelerisLab.body.motion import MovingTables
    sim = make_simulation(monkeypatch)
    sim._failed, sim.stepper = None, SimpleNamespace(moving=None)
    assert sim.moving_diagnostics().shape == (0, MovingTables.DIAG_FLOATS)


@pytest.mark.parametrize("failure", ("upload", "first_sync", "no_donor", "hold_upload",
                                     "hold_sync"))
def test_failed_advance_fails_simulation(monkeypatch, failure):
    kwargs = dict(no_donor=1.) if failure == "no_donor" else dict(fail=failure)
    sim, writes, motion = advancing_simulation(monkeypatch, **kwargs)
    motion._on_failure = sim._fail_moving
    motion.set_target((0., .25), (0., .05))
    syncs = []

    def synchronize():  # stands in for an asynchronous failure surfacing at sync
        syncs.append(1)
        if (failure, len(syncs)) in (("first_sync", 1), ("hold_sync", 2)):
            raise RuntimeError("sync failed")
    with pytest.raises(RuntimeError):
        sim.run(1, stream=SimpleNamespace(synchronize=synchronize))
    # Before the step-group sync and no_donor check nothing is committed; the hold
    # record is staged after the commit (applied already matches the moved device).
    committed = failure in ("hold_upload", "hold_sync")
    assert motion.applied == ((0., .25) if committed else (0., 0.))
    assert (motion.pending is None) == committed
    assert motion.failed and sim._failed is not None
    assert sim.stepper.blocked == sim._failed  # direct stepper reuse refused too
    with pytest.raises(RuntimeError, match="Simulation failed"):
        sim.run(1, stream=STREAM)


def test_failed_hold_step_blocks_the_direct_stepper(monkeypatch):
    # No target is pending, so only the failure latch can refuse direct steps.
    from CelerisLab.lbm.stepper import LBMStepper
    sim, writes, motion = advancing_simulation(monkeypatch, fail="upload")
    motion._on_failure = sim._fail_moving
    launches = []
    kernels = SimpleNamespace(get_function=lambda name: lambda *a, **k: launches.append(name))
    cfg = SimpleNamespace(streaming="double_buffer", threads_per_block=128, is_d2q9=True,
                          nx=64, ny=32, nz=1)
    field = SimpleNamespace(ddf_gpu="ddf", temp_gpu="temp", flag_gpu="flag", n_curved=0)
    stepper = LBMStepper(field, kernels, cfg)
    stepper.moving = sim.stepper.moving
    sim.stepper = stepper
    assert motion.pending is None and stepper.blocked is None
    with pytest.raises(RuntimeError, match="upload failed"):
        sim.run(2, stream=STREAM)  # hold step
    assert sim._failed is not None and stepper.blocked == sim._failed
    with pytest.raises(RuntimeError, match="moving circle step failed"):
        stepper.step(1, action_gpu=None, obs_gpu=None)
    assert launches == [] and stepper.step_count == 0
    assert (field.ddf_gpu, field.temp_gpu) == ("ddf", "temp")


def test_moving_tables_layout_and_attachment(monkeypatch):
    from CelerisLab.body.motion import MovingTables
    a = MovingCircle((50., 40.), 10., Translation((-1., 1., -3., 3.), .1))
    b = MovingCircle((90.25, 30.5), 5., Translation((0., 0., 0., 0.), .1))
    band, geom = MovingTables.layout([(a, 2), (b, 5)])
    wa, ha = 63 - 37 + 1, 55 - 25 + 1
    assert band == [37, 25, wa, ha, 0, 2, 83, 23, 16, 16, wa * ha * 8, 5]
    assert geom == [50., 40., 50., 40., 10., 0., 0., 90.25, 30.5, 90.25, 30.5, 5., 0., 0.]  # hold
    sim = routed_simulation(monkeypatch)
    sim.add_body("circle", (50., 40.), 10., translation=Translation((0., 0., 0., 0.), .1))
    sim._moving_tables = "tables"
    sim._attach_moving()
    assert sim.stepper.moving == "tables" and sim.stepper.blocked is None
    sim._moving_tables = None
    sim._attach_moving()
    assert sim.stepper.moving is None and sim.stepper.blocked is not None


class CountingOwner(Owner):
    """Owner double that counts nested activations like CudaContext."""
    def __init__(self):
        super().__init__()
        self.active = 0

    def activate(self):
        from contextlib import contextmanager

        @contextmanager
        def scope():
            self.active += 1
            try:
                yield self
            finally:
                self.active -= 1
        return scope()


def test_public_diagnostics_keep_owner_active_through_sync_and_download(monkeypatch):
    sim = make_simulation(monkeypatch)
    sim.ctx, sim._failed = CountingOwner(), None
    seen = []
    sim._stream = SimpleNamespace(synchronize=lambda: seen.append(("sync", sim.ctx.active)))
    expected = np.arange(6, dtype=np.float32).reshape(1, 6)

    def download():
        seen.append(("download", sim.ctx.active))
        return expected
    sim.stepper = SimpleNamespace(moving=SimpleNamespace(download_diagnostics=download))
    assert sim.moving_diagnostics() is expected
    assert [name for name, _ in seen] == ["sync", "download"]
    assert all(active > 0 for _, active in seen) and sim.ctx.active == 0


def test_stepper_launches_moving_links_before_static_curved():
    from CelerisLab.lbm.stepper import LBMStepper
    launches = []
    module = SimpleNamespace(get_function=lambda name: lambda *a, **k: launches.append((name, a, k)))
    cfg = SimpleNamespace(streaming="double_buffer", threads_per_block=128, is_d2q9=True,
                          nx=64, ny=32, nz=1)
    curved = SimpleNamespace(**{f"{name}_gpu": name for name in (
        "fluid_idx", "dir", "q", "rx", "ry", "rz", "fallback_class", "body_id")},
        surface_velocity=None, surface_velocity_gpu=None, link_exchange_gpu=None)
    field = SimpleNamespace(ddf_gpu="ddf", temp_gpu="temp", flag_gpu="flag", n_curved=5,
                            curved=curved, invalidate_host_ddf_cache=lambda: None)
    stepper = LBMStepper(field, module, cfg)
    stepper._launch_force_region = stepper._launch_sensor = lambda *a, **k: None
    stepper.moving = SimpleNamespace(band_gpu="band", geom_gpu="geom", diag_gpu="diag",
                                     slot_debug_gpu=None, n_bodies=1, n_slots=300)
    stepper.step(2, action_gpu="act", obs_gpu="obs")
    names = [name for name, _, _ in launches]
    assert names == ["MovingCircleLinks", "CurvedBoundaryKernel", "OneStep"] * 2
    assert stepper.step_count == 2
    assert launches[1][1][0] == "ddf" and launches[3][1][0] == "temp"  # swap between steps
    args, kw = launches[0][1], launches[0][2]
    assert args[:4] == ("ddf", "flag", "band", "geom") and int(args[5]) == 300
    assert kw["grid"] == (3, 1, 1) and kw["block"] == (128, 1, 1)


def test_recompile_installs_a_guarded_stepper(monkeypatch):
    from CelerisLab.lbm.stepper import LBMStepper
    sim = routed_simulation(monkeypatch)
    sim.add_body("circle", (50., 40.), 10., translation=Translation((0., 0., -1., 1.), .1))
    launches = []
    kernel_module = SimpleNamespace(get_function=lambda name: lambda *a, **k: launches.append(name))
    monkeypatch.setattr(module.compiler, "build_module", lambda *a, **k: kernel_module)
    sim.lbm_cfg = SimpleNamespace(validate_runtime=lambda: None, viscosity=.01, omega=1.,
                                  streaming="double_buffer", threads_per_block=128,
                                  is_d2q9=True, nx=120, ny=80, nz=1)
    sim._compiled_viscosity = .01
    sim._assert_allocation_contract = lambda: None
    sim._resolve_compile_arch = lambda: "sm_70"
    sim._assert_object_count_contract = lambda **k: None
    sim.field = SimpleNamespace(module=None, invalidate_module_constant_cache=lambda: None,
                                _upload_params=lambda: None, _runtime_params={},
                                ddf_gpu="ddf", temp_gpu="temp", n_curved=0)
    sim.stepper = SimpleNamespace(_step_count=7, blocked=None)
    sim._recompile()
    assert isinstance(sim.stepper, LBMStepper) and sim.stepper.step_count == 7
    with pytest.raises(RuntimeError, match="moving kernels"):
        sim.stepper.step(1, action_gpu=None, obs_gpu=None)
    assert launches == [] and (sim.field.ddf_gpu, sim.field.temp_gpu) == ("ddf", "temp")
    assert sim.stepper.step_count == 7


def validate(monkeypatch, objects, flags=None, **cfg):
    sim = routed_simulation(monkeypatch)
    for obj in objects:
        sim.bodies.add(obj)
    for name, value in cfg.items():
        setattr(sim.lbm_cfg, name, value)
    base = np.full(120 * 80, FLUID, np.uint16) if flags is None else flags
    sim.field = SimpleNamespace(nx=120, ny=80, build_channel_flags=lambda: base.copy())
    sim._validate_moving_bodies()


def test_initialize_rejects_band_and_modes(monkeypatch):
    t = Translation((-1., 1., -3., 3.), .1)
    validate(monkeypatch, [body(50., 40., 10., t), body(100., 40., 5.)])  # accepted
    walls = np.full(120 * 80, FLUID, np.uint16)
    walls[60 + 30 * 120] = SOLID
    cases = [([body(10., 40., 10., t)], None, {}, "domain edge"),
             ([body(50., 40., 10., t)], walls, {}, "non-fluid"),
             ([body(50., 40., 10., t), body(68., 40., 4.)], None, {}, "overlaps"),
             ([body(50., 40., 10., t), body(50., 60., 3., is_sensor=True)], None, {}, "overlaps"),
             ([body(50., 40., 10., t), body(64., 40., 3., is_force_region=True)], None, {}, "overlaps"),
             ([body(50., 40., 10., t)], None, {"streaming": "esopull"}, "double_buffer"),
             ([body(50., 40., 10., t)], None, {"ddf_shifting": True}, "unshifted"),
             ([body(50., 40., 10., t)], None, {"is_d2q9": False}, "D2Q9")]
    for objects, flags, cfg, match in cases:
        with pytest.raises(ValueError, match=match):
            validate(monkeypatch, objects, flags, **cfg)


def test_initialize_validates_before_mutation(monkeypatch):
    sim = routed_simulation(monkeypatch)
    sim.bodies.add(body(10., 40., 10., Translation((0., 0., 0., 0.), .1)))
    uploads = []
    sim.field = SimpleNamespace(nx=120, ny=80, build_channel_flags=lambda: np.full(9600, FLUID, np.uint16),
                                upload_flags=lambda: uploads.append("flags"))
    sim._assert_allocation_contract = lambda: None
    with pytest.raises(ValueError, match="domain edge"):
        sim.initialize()
    assert uploads == [] and not sim._initialized


def test_moving_circle_excluded_from_static_links():
    def manager(*objects):
        mgr = ObjectManager(nx=120, ny=80, nz=1, nq=9, cfg=Cfg())
        ids = [mgr.add(obj) for obj in objects]
        return mgr, ids
    reference_mgr, _ = manager(body(100., 40., 5.), body(20., 40., 6.))
    mixed_mgr, ids = manager(body(100., 40., 5.), body(20., 40., 6.),
                             body(55., 40., 10., Translation((0., 0., -1., 1.), .1)))
    reference = reference_mgr.build_compact_lists_for(reference_mgr.objects)
    mixed = mixed_mgr.build_compact_lists_for(mixed_mgr.objects)
    for a, b in zip(reference[:8], mixed[:8]):  # static link columns byte-identical
        assert np.asarray(a).tobytes() == np.asarray(b).tobytes()
    assert len(reference[0]) > 0 and ids[2] not in set(np.asarray(mixed[3]).tolist())


def test_failed_moving_step_marks_simulation_failed(monkeypatch):
    sim, writes, motion = routed_advancing_simulation(monkeypatch)
    sim.ctx = ClosableOwner()
    sim.set_body(0, offset=(0., .2), vy=.05)

    def fail(*_):
        raise RuntimeError("sync failed")
    with pytest.raises(RuntimeError, match="sync failed"):
        sim.run(1, stream=SimpleNamespace(synchronize=fail))
    assert sim._failed is not None and motion.applied == (0., 0.)
    for call in (lambda: sim.run(0), lambda: sim.run(1), sim.step, sim.sync_bodies,
                 sim.save_checkpoint, lambda: sim.load_checkpoint("x.h5")):
        with pytest.raises(RuntimeError, match="Simulation failed"):
            call()
    with pytest.raises(RuntimeError, match="unknown"):  # no retry on the failed circle
        sim.set_body(0, offset=(0., .1), vy=.05)
    with sim.ctx.activate():  # a refused close keeps the failed state and resources
        with pytest.raises(RuntimeError, match="active CUDA"):
            sim.close()
    assert sim._failed is not None and sim.bodies is not None and not sim._closed
    sim.close()  # retry outside the activation succeeds
    assert sim._closed and sim._failed is not None
    with pytest.raises(RuntimeError, match="CUDA context is closed"):  # no re-entry after close
        sim.run(1)
    with pytest.raises(RuntimeError, match="unknown"):  # the failed circle stays refused
        motion.record()