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

Host doubles only: no kernels exist yet, so every positive-step run with a
moving body must be rejected before any device write.
"""
from types import SimpleNamespace

import pytest

from CelerisLab.body.motion import (MAX_STEP, R_MIN, MovingCircle, Translation,
                                    enumerate_donor_coverage)

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():
    c = circle()
    assert c.record() == ((50., 40.), (50., 40.), (0., 0.))  # hold: no translation velocity
    c.set_target((0., .25), (0., .05))
    assert c.record() == ((50., 40.), (50., 40.25), (0., .05))
    uploads = []
    c.advance(uploads.append, lambda: None, lambda: None)
    assert uploads == [((50., 40.), (50., 40.25), (0., .05))]
    assert c.applied == (0., .25) and c.pending is None  # offset-only commit
    c.advance(uploads.append, lambda: None, lambda: None)  # no new target: hold step
    assert uploads[-1] == ((50., 40.25), (50., 40.25), (0., 0.)) and c.applied == (0., .25)
    c.set_target((0., 0.), (0., -.05))  # return to zero offset is a real target
    assert c.pending == ((0., 0.), (0., -.05))
    c.advance(uploads.append, lambda: None, lambda: None)
    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(stage):
    c = circle()
    c.set_target((0., .25), (0., .05))

    def call(name):
        def run(*_):
            if name == stage:
                raise RuntimeError(f"{name} failed")
        return run
    with pytest.raises(RuntimeError, match=f"{stage} failed"):
        c.advance(call("upload"), call("step"), call("synchronize"))
    assert c.failed and c.applied == (0., 0.)
    for later in (c.record, lambda: c.set_target((0., .1), (0., 0.)),
                  lambda: c.advance(*(lambda *_: None,) * 3)):
        with pytest.raises(RuntimeError, match="unknown"):
            later()


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)


@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 make_simulation  # noqa: E402


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=[],
                                 has_pending_edit=lambda: False,
                                 _upload_action_async=lambda s: writes.append("action"),
                                 zero_obs_async=lambda s: writes.append("obs"))
    sim.stepper = SimpleNamespace(step=lambda *a, **k: writes.append("step"))
    return sim, writes


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))
    for call in (lambda: sim.run(1), lambda: sim.run(5), lambda: sim.step(),
                 lambda: sim.run(3, 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)
    assert writes == [] and sim.bodies.objects[0].motion.applied == (0., 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.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.set_body(bid, offset=(0., .3), vy=.05)
    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)


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():
    still = [body(100., 40., 5.)]
    mgr = ObjectManager(nx=120, ny=80, nz=1, nq=9, cfg=Cfg())
    reference = mgr.build_compact_lists_for(still)
    with_moving = mgr.build_compact_lists_for(
        still + [body(50., 40., 10., Translation((0., 0., -1., 1.), .1))])
    for a, b in zip(reference[:8], with_moving[:8]):
        assert np.asarray(a).tobytes() == np.asarray(b).tobytes()
    assert len(reference[0]) > 0


def test_failed_moving_step_marks_simulation_failed(monkeypatch):
    sim = routed_simulation(monkeypatch)
    bid = sim.add_body("circle", (50., 40.), 10., translation=Translation((0., 0., -1., 1.), .1))
    motion = sim.bodies.get(bid).motion
    motion.set_target((0., .2), (0., .05))

    def fail(*_):
        raise RuntimeError("sync failed")
    with pytest.raises(RuntimeError, match="sync failed"):
        motion.advance(lambda *_: None, lambda: None, 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):
        with pytest.raises(RuntimeError, match="Simulation failed"):
            call()
    sim.close()  # close still works; no re-entry afterwards