"""Body-independent Jet: true-hit spatial selection, fields, profiles, overlap, amplitude."""
import math
from types import SimpleNamespace

import numpy as np
import pytest

from CelerisLab.body.jets import (Jet, Parallel, Radial, Sector, Tube,
                                  edge_to_node_frame, fine_hits_to_edge_frame,
                                  jet_link_shapes)
from CelerisLab.body.objects import SurfaceActuation
from CelerisLab.lbm.curved_links import CurvedLinkSoA, surface_link_velocity
from .test_state_regressions import device, make_field, manager
from .test_surface_actuation import circle, projected

CENTER, RADIUS = (6., 5.), 2.999


def scene(spec=None):
    """Single circle links plus their true FP64 hits, row-aligned with the SoA."""
    obj = circle(spec)
    cl, cfg, flags = projected(obj)
    links = obj.geometry.build_cut_links(cfg.nx, cfg.ny, domain_flags=flags)
    hits = np.array([(link.hit_x, link.hit_y) for link in links])
    assert len(hits) == len(cl.dir)
    return obj, cl, cfg, flags, hits


def arc(hits, angle, half_width):
    d = hits - CENTER
    delta = np.angle(np.exp(1j*(np.arctan2(d[:, 1], d[:, 0]) - angle)))
    return np.abs(delta) < half_width


def test_concentric_sector_gives_inward_radial_and_tube_parallel_on_same_links():
    *_, hits = scene()
    half = math.radians(10)
    expected = arc(hits, 0., half)
    assert expected.sum() >= 2
    sector = Jet(Sector(CENTER, 0., half, RADIUS + .1), Radial(CENTER), "uniform", 1e-3)
    shape, jet_id = jet_link_shapes([sector], hits)
    assert np.array_equal(jet_id >= 0, expected)
    d = hits[expected] - CENTER
    np.testing.assert_allclose(shape[expected], -d/np.hypot(*d.T)[:, None], atol=1e-15)
    # Tube along +x through the same arc: lateral |y-5| < R sin(10 deg).
    tube = Jet(Tube((CENTER[0], CENTER[1]), 0., RADIUS + .1, RADIUS*math.sin(half)),
               Parallel(math.pi), "uniform", 1e-3)
    shape_t, id_t = jet_link_shapes([tube], hits)
    assert np.array_equal(id_t >= 0, expected)
    np.testing.assert_allclose(shape_t[expected], [[-1., 0.]]*expected.sum(), atol=1e-15)


def test_off_centre_radial_points_to_field_origin_not_body_centre():
    *_, hits = scene()
    origin = (12., 9.)
    jet = Jet(Sector(CENTER, .1, math.pi, 10.), Radial(origin, inward=False))
    shape, jet_id = jet_link_shapes([jet], hits)
    assert np.all(jet_id == 0)
    d = hits - origin
    np.testing.assert_allclose(shape, d/np.hypot(*d.T)[:, None], atol=1e-15)
    assert not np.allclose(shape, (hits - CENTER)/RADIUS, atol=1e-3)


@pytest.mark.parametrize("profile,expected", [("uniform", 1.), ("parabolic", .75), ("quartic", .5625)])
def test_profiles_on_lateral_coordinate(profile, expected):
    hits = np.array([[1., .5]])
    jet = Jet(Tube((0., 0.), 0., 2., 1.), Parallel(0.), profile)
    shape, _ = jet_link_shapes([jet], hits)
    np.testing.assert_allclose(shape[0], [expected, 0.], rtol=1e-15)


def test_selection_uses_true_hit_not_clipped_q():
    # Same link as test_aperture_uses_true_hit_not_fluid_node_or_clamped_q:
    # receiver (9,5), direction (-1,+1); true t < .01 but packed q == .01.
    t = (6-math.sqrt(36-8*(9-2.999**2)))/4
    angle = math.atan2(t, 3-t)
    _, cl, _, _, hits = scene()
    row = np.flatnonzero((cl.fluid_idx == 5*12 + 9) & (cl.dir == 8))
    assert len(row) == 1 and np.isclose(cl.q[row[0]], .01)
    clipped = np.array([9., 5.]) + .01*np.array([-1., 1.])
    jet = Jet(Sector(CENTER, angle, 1e-4, RADIUS + .1), Radial(CENTER))
    _, jet_id = jet_link_shapes([jet], hits)
    assert jet_id[row[0]] == 0
    with pytest.raises(ValueError, match="selects no"):
        jet_link_shapes([jet], clipped[None, :])


def test_fine_half_cell_frame():
    edge = fine_hits_to_edge_frame([[0., 0.], [3., 1.5]], (4., 6.), .5)
    np.testing.assert_array_equal(edge, [[4.25, 6.25], [5.75, 7.]])
    np.testing.assert_array_equal(edge_to_node_frame(edge), [[3.75, 5.75], [5.25, 6.5]])


def test_same_physical_hit_in_simulation_and_patch_frames():
    # One physical circle: Simulation node-frame centre C; patch edge-frame
    # centre C + 1/2; fine-local centre (C + 1/2 - origin)/dx - 1/2 as in
    # prepare_fine_boundary. Fine hits mapped edge -> node must lie on the
    # Simulation circle, and one node-frame Jet selects the same arc on both.
    from CelerisLab.body.geometry.circle import CircleGeometry
    C, R, origin, dx = np.array([10.3, 9.6]), 2.15, np.array([4., 3.]), .5
    local = (C + .5 - origin)/dx - .5
    fine = CircleGeometry(*local, R/dx).build_cut_links(40, 40)
    coarse = CircleGeometry(*C, R).build_cut_links(24, 20)
    fine_hits = edge_to_node_frame(fine_hits_to_edge_frame(
        [(l.hit_x, l.hit_y) for l in fine], origin, dx))
    coarse_hits = np.array([(l.hit_x, l.hit_y) for l in coarse])
    np.testing.assert_allclose(np.hypot(*(fine_hits - C).T), R, rtol=0, atol=1e-12)
    np.testing.assert_allclose(np.hypot(*(coarse_hits - C).T), R, rtol=0, atol=1e-12)
    jet = Jet(Sector(tuple(C), .4, .5, R + .01, R - .01), Radial(tuple(C), inward=False))
    for hits in (fine_hits, coarse_hits):
        shape, jet_id = jet_link_shapes([jet], hits)
        d = hits - C
        angle = np.arctan2(d[:, 1], d[:, 0])
        assert np.array_equal(jet_id == 0, np.abs(angle - .4) < .5)
        np.testing.assert_allclose(shape[jet_id == 0], d[jet_id == 0]/R, atol=1e-12)
    # Relabelling instead of converting moves every hit off the circle.
    wrong = fine_hits_to_edge_frame([(l.hit_x, l.hit_y) for l in fine], origin, dx)
    assert np.abs(np.hypot(*(wrong - C).T) - R).max() > .3


def test_links_are_distinct_by_receiver_and_direction_and_overlap_rejected():
    # Two physical links (different receiver/direction) sharing one hit point
    # are two rows, both driven; two jets on one link are rejected.
    hits = np.array([[1., 1.], [1., 1.]])
    a = Jet(Tube((0., 1.), 0., 2., .5), Parallel(0.))
    shape, jet_id = jet_link_shapes([a], hits)
    assert jet_id.tolist() == [0, 0]
    b = Jet(Sector((1., 0.), math.pi/2, .1, 2.), Parallel(0.))
    with pytest.raises(ValueError, match="overlaps"):
        jet_link_shapes([a, b], hits)
    obj, cl, cfg, flags, real = scene(SurfaceActuation(0., .3, .002))
    with pytest.raises(ValueError, match="SurfaceActuation"):
        surface_link_velocity([obj], cfg, flags, cl.fluid_idx, cl.dir, cl.body_id, cl.rx, cl.ry,
                              jets=[Jet(Sector(CENTER, 0., .1, 4.), Radial(CENTER))], hits=real)
    with pytest.raises(ValueError, match="row-aligned"):
        surface_link_velocity([], cfg, flags, cl.fluid_idx, cl.dir, cl.body_id, cl.rx, cl.ry,
                              jets=[a], hits=real[:-1])


def test_amplitude_update_reuses_cached_shapes():
    obj, cl, cfg, flags, hits = scene()
    jet = Jet(Sector(CENTER, 0., .5, 4.), Radial(CENTER), "parabolic", 1e-3)
    cl.assign_surface([obj], cfg, flags, jets=[jet], hits=hits)
    shape, ids = cl.jet_shape, cl.jet_id
    before = cl.surface_velocity.copy()
    cl.set_jet_amplitude(0, -2e-3)
    assert cl.jet_shape is shape and cl.jet_id is ids
    np.testing.assert_allclose(cl.surface_velocity, -2*before, rtol=1e-6, atol=0)
    assert np.array_equal(cl.aperture, ids >= 0)
    with pytest.raises(ValueError):
        cl.set_jet_amplitude(0, 1.5)
    assert cl.jets[0].amplitude == -2e-3 and jet.amplitude == 1e-3  # registered Jet untouched
    np.testing.assert_allclose(cl.surface_velocity, -2*before, rtol=1e-6, atol=0)


def test_jet_definition_is_immutable_and_amplitude_validated():
    jet = Jet(Sector(CENTER, 0., .5, 4.), Radial(CENTER), "parabolic", 1e-3)
    for name, value in (("field", Parallel(0.)), ("region", Tube(CENTER, 0., 1., 1.)),
                        ("profile", "uniform"), ("amplitude", float("nan"))):
        with pytest.raises(AttributeError):
            setattr(jet, name, value)
    for bad in (float("nan"), float("inf"), 1.0):
        with pytest.raises(ValueError):
            jet.with_amplitude(bad)
    assert jet.with_amplitude(-1e-3).field is jet.field and jet.amplitude == 1e-3
    with pytest.raises(ValueError, match="finite"):
        jet_link_shapes([jet], [[np.nan, 5.]])


@pytest.mark.parametrize("spec", [SurfaceActuation(0., .6, .002), SurfaceActuation(2., 1.1, -.003)])
def test_legacy_adapter_measured_equivalence(spec):
    obj, legacy, cfg, flags, hits = scene(spec)
    cl = CurvedLinkSoA()
    cl.assign_host(legacy.fluid_idx, legacy.dir, legacy.q, legacy.body_id,
                   legacy.rx, legacy.ry, legacy.rz, legacy.fallback_class)
    cl.assign_surface([circle()], cfg, flags, jets=[spec.as_jet(CENTER, RADIUS)], hits=hits)
    assert np.array_equal(cl.aperture, legacy.aperture)
    diff = np.abs(cl.surface_velocity - legacy.surface_velocity).max()
    print(f"legacy adapter max |dU| = {diff:.3e} (exact={diff == 0})")
    assert diff <= 4*np.finfo(np.float32).eps*abs(spec.normal_velocity)


def jet_case(device, jets):
    field = make_field(device)
    bodies = manager(field, [circle(radius=2.)])
    bodies.jets.extend(jets)
    bodies.sync_to_gpu(field)
    return field, bodies


def test_public_surface_exchange_for_jet_preserves_aperture_at_zero_and_reverse(device):
    from CelerisLab.simulation import Simulation
    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
    field, bodies = jet_case(device, [jet])
    sim = SimpleNamespace(bodies=bodies, field=field, _initialized=True,
                          stream=SimpleNamespace(synchronize=device.synchronize))
    sim._require_usable = Simulation._require_usable.__get__(sim)
    cl = field.curved
    assert bodies.get(0).surface_actuation is None
    aperture, velocity = cl.aperture, cl.surface_velocity.copy()
    assert aperture.any() and (~aperture).any()
    # Synthetic per-link exchanges test readback/splitting, not kernel physics.
    values = np.arange(cl.count * 3, dtype=np.float32).reshape(-1, 3)
    device.htod(cl.link_exchange_gpu, values)
    for amplitude in (2e-3, 0., -1e-3):
        Simulation.set_jet_amplitude(sim, 0, amplitude)
        result = Simulation.read_surface_exchange(sim, 0)
        assert cl.aperture is aperture
        np.testing.assert_array_equal(aperture, cl.jet_id >= 0)
        np.testing.assert_allclose(cl.surface_velocity, velocity * (amplitude / 2e-3))
        for name, mask in (("wall", ~aperture), ("aperture", aperture)):
            np.testing.assert_array_equal(result[name], values[mask].sum(axis=0, dtype=np.float64))


def test_public_surface_exchange_keeps_legacy_and_passive_guards(device):
    from CelerisLab.simulation import Simulation
    field = make_field(device)
    bodies = manager(field, [circle(SurfaceActuation(0., .6, .002), radius=2.)])
    bodies.sync_to_gpu(field)
    sim = SimpleNamespace(bodies=bodies, field=field,
                          stream=SimpleNamespace(synchronize=device.synchronize))
    sim._require_usable = Simulation._require_usable.__get__(sim)
    for value in Simulation.read_surface_exchange(sim, 0).values():
        np.testing.assert_array_equal(value, 0.)
    field.curved.free_gpu_columns()
    with pytest.raises(RuntimeError, match="diagnostic is allocated"):
        Simulation.read_surface_exchange(sim, 0)

    # An exchange buffer elsewhere in the scene must not admit a passive body.
    for jets in ([], [Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER))]):
        field, bodies = jet_case(device, jets)
        passive = bodies.add(circle(center=(2., 3.), radius=.6))
        bodies.sync_to_gpu(field)
        sim.bodies, sim.field = bodies, field
        with pytest.raises(ValueError, match="Body has no surface actuation"):
            Simulation.read_surface_exchange(sim, passive)
        if not jets:
            with pytest.raises(ValueError, match="Body has no surface actuation"):
                Simulation.read_surface_exchange(sim, 0)


def test_manager_passes_true_hits_and_reuploads_only_velocity(device):
    # CPU memory double only: checks host/device bookkeeping, not GPU numerics.
    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
    field, bodies = jet_case(device, [jet])
    cl = field.curved
    assert cl.jet_id is not None and (cl.jet_id == 0).sum() >= 2
    velocity_gpu = cl.surface_velocity_gpu
    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
                                  cl.surface_velocity)
    geometry = {name: bytes(getattr(cl, name + "_gpu").data)
                for name in ("fluid_idx", "dir", "q", "body_id", "rx", "ry")}
    writes, before = device.writes, cl.surface_velocity.copy()
    cl.set_jet_amplitude(0, -1e-3, SimpleNamespace(synchronize=lambda: None))
    assert device.writes == writes + 1 and cl.surface_velocity_gpu is velocity_gpu
    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
                                  cl.surface_velocity)
    assert all(bytes(getattr(cl, name + "_gpu").data) == data for name, data in geometry.items())
    np.testing.assert_allclose(cl.surface_velocity, -.5*before, rtol=1e-6, atol=0)
    np.testing.assert_array_equal(cl.surface_velocity[cl.jet_id < 0], 0.)


def test_manager_rejects_overlapping_jets_and_moving_bodies(device):
    a = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER))
    b = Jet(Tube(CENTER, 0., 3., .5), Parallel(0.))
    with pytest.raises(ValueError, match="overlaps"):
        jet_case(device, [a, b])
    field = make_field(device)
    moving = circle(radius=2.)
    moving.state.omega = 1e-3
    bodies = manager(field, [moving])
    bodies.jets.append(a)
    with pytest.raises(ValueError, match="stationary"):
        bodies.sync_to_gpu(field)


def test_origins_are_frozen_float_pairs():
    point = [6., 5.]
    array = np.array([6., 5.])
    jets = [Jet(Sector(point, 0., .5, 4.), Radial(array)), Jet(Tube(array, 0., 1., 1.), Parallel(0.))]
    point[0], array[0] = 100., 100.
    assert jets[0].region.origin == (6., 5.) and jets[0].field.origin == (6., 5.)
    assert jets[1].region.origin == (6., 5.)
    assert all(type(v) is float for v in jets[0].region.origin)
    for bad in ((1., 2., 3.), (1., float("nan"))):
        with pytest.raises(ValueError):
            Radial(bad)


def test_set_body_rejects_motion_of_jet_driven_body(device):
    from CelerisLab.simulation import Simulation
    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
    field, bodies = jet_case(device, [jet])
    sim = SimpleNamespace(bodies=bodies, field=field, lbm_cfg=field.cfg, _initialized=True)
    with pytest.raises(ValueError, match="stationary"):
        Simulation.set_body(sim, 0, omega=.01)
    with pytest.raises(ValueError, match="stationary"):
        Simulation.set_body(sim, 0, vx=.01)
    action = bodies.action.copy()
    with pytest.raises(ValueError, match="stationary"):
        bodies.set_body_state(0, omega=.01)  # direct manager entry
    assert bodies.get(0).state.omega == 0 and np.array_equal(bodies.action, action)
    Simulation.set_body(sim, 0, omega=0.)  # stationary updates stay allowed
    bodies.set_body_state(0, omega=0.)


def test_action_upload_rejects_motion_of_jet_driven_body(device):
    # Direct edits of the public action array are checked at the upload gate.
    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
    field, bodies = jet_case(device, [jet])
    stream = SimpleNamespace(synchronize=lambda: None)
    slot = 3*field.cfg.dim
    for index in (0, 1, slot - 1):  # vx, vy, omega
        bodies.action[:] = 0
        bodies.action[index] = .01
        writes = device.writes
        with pytest.raises(ValueError, match="stationary"):
            bodies._upload_action_async(stream)
        assert device.writes == writes  # nothing reached the device
    bodies.action[:] = 0
    bodies._upload_action_async(stream)  # zero motion stays accepted


def test_fine_boundary_keeps_edge_frame_hits_and_checks_jets():
    from CelerisLab.body.objects import SimObject
    from CelerisLab.body.geometry.circle import CircleGeometry
    from CelerisLab.config import LBMConfig
    from CelerisLab.lbm import refinement
    cfg = LBMConfig(nx=24, ny=24, viscosity=.05)
    cfg.validate()
    patch = refinement.FixedPatch2D.from_config(cfg, (4, 4, 20, 20))
    C, R = (12.25, 12.25), 2.15  # coarse cell-edge frame, as fine_objects
    body = SimObject(0, geometry=CircleGeometry(*C, R), center=C, radius=R)
    plan = refinement.prepare_fine_boundary(cfg, patch, [body])
    assert plan.curved_hits.shape == (len(plan.curved_host[0]), 2)
    assert not plan.curved_hits.flags.writeable
    np.testing.assert_allclose(np.hypot(*(plan.curved_hits - C).T), R, rtol=0, atol=1e-12)
    jet = Jet(Sector(C, 0., .5, R + .01), Radial(C, inward=False))
    refinement.prepare_fine_boundary(cfg, patch, [body], jets=[jet])
    away = Jet(Sector((2., 2.), 0., .1, .5), Parallel(0.))
    with pytest.raises(ValueError, match="selects no"):
        refinement.prepare_fine_boundary(cfg, patch, [body], jets=[away])
    with pytest.raises(ValueError, match="overlaps"):
        refinement.prepare_fine_boundary(cfg, patch, [body], jets=[jet, jet])
    with pytest.raises(ValueError, match="need fine boundary objects"):
        refinement.prepare_fine_boundary(cfg, patch, [], jets=[jet])
    # One fine jet over two physical bodies is rejected; one per body is accepted.
    A, B = (8.75, 12.25), (15.75, 12.25)
    pair = [SimObject(i, geometry=CircleGeometry(*c, 1.5), center=c, radius=1.5)
            for i, c in enumerate((A, B))]
    span = Jet(Tube((4.5, 12.25), 0., 15., .5), Parallel(0.))
    with pytest.raises(ValueError, match="one jet may drive only one body"):
        refinement.prepare_fine_boundary(cfg, patch, pair, jets=[span])
    refinement.prepare_fine_boundary(cfg, patch, pair, jets=[
        Jet(Sector(A, 0., .5, 2.), Radial(A, inward=False)),
        Jet(Sector(B, math.pi, .5, 2.), Radial(B, inward=False))])
    # Jet vs legacy SurfaceActuation overlap on the fine level (root's case).
    legacy = SimObject(0, geometry=CircleGeometry(*C, R), center=C, radius=R,
                       surface_actuation=SurfaceActuation(0., .5, .002))
    with pytest.raises(ValueError, match="SurfaceActuation"):
        refinement.prepare_fine_boundary(cfg, patch, [legacy], jets=[jet])
    disjoint = Jet(Sector(C, math.pi, .5, R + .01), Radial(C, inward=False))
    refinement.prepare_fine_boundary(cfg, patch, [legacy], jets=[disjoint])


def test_action_upload_keeps_cfg_none_fallback(device):
    # Default ObjectManager(cfg=None) keeps the 2D slot fallback, empty or not.
    from CelerisLab.body.manager import ObjectManager
    stream = SimpleNamespace(synchronize=lambda: None)
    bodies = ObjectManager(12, 10, 1, 9)
    bodies._upload_action_async(stream)
    bodies.add(circle(radius=2.))
    bodies._upload_action_async(stream)


def two_body_case(device, jets):
    field = make_field(device)
    bodies = manager(field, [circle(center=(3., 5.), radius=1.6),
                             circle(center=(8.5, 5.), radius=1.6)])
    bodies.jets.extend(jets)
    bodies.sync_to_gpu(field)
    return field, bodies


@pytest.mark.parametrize("amplitude", [2e-3, 0.0])  # geometry checked even at zero amplitude
def test_one_jet_spanning_two_bodies_rejected(device, amplitude):
    both = Jet(Tube((0., 5.), 0., 12., .5), Parallel(0.), amplitude=amplitude)
    with pytest.raises(ValueError, match="one jet may drive only one body"):
        two_body_case(device, [both])


def test_separate_jets_on_separate_bodies_accepted(device):
    left = Jet(Sector((3., 5.), 0., .6, 2.), Radial((3., 5.), inward=False), amplitude=1e-3)
    right = Jet(Sector((8.5, 5.), math.pi, .6, 2.), Radial((8.5, 5.), inward=False), amplitude=-1e-3)
    field, _ = two_body_case(device, [left, right])
    cl = field.curved
    for j in (0, 1):
        assert len(np.unique(cl.body_id[cl.jet_id == j])) == 1
    assert cl.body_id[cl.jet_id == 0][0] != cl.body_id[cl.jet_id == 1][0]


def test_one_body_with_noncontiguous_rows_accepted_and_mixed_rejected():
    # Synthetic rows: ownership is set membership over body_id, not row adjacency.
    from CelerisLab.body.geometry.circle import CircleGeometry
    from CelerisLab.body.objects import SimObject
    from CelerisLab.config import LBMConfig
    from CelerisLab.lbm.descriptors import FLUID
    cfg = LBMConfig(nx=12, ny=10, viscosity=.1)
    cfg.validate()
    objs = [SimObject(i, geometry=CircleGeometry(3. + 5*i, 5., 1.), center=(3. + 5*i, 5.), radius=1.)
            for i in (0, 1)]
    flags = np.full(cfg.nx*cfg.ny, FLUID, np.uint16)
    fluid_idx = np.array([10, 20, 30], np.uint32)
    dirs = np.array([1, 1, 1], np.uint8)
    zeros = np.zeros(3, np.float32)
    hits = np.array([[0., 0.], [5., 5.], [0., 1.]])
    jet = Jet(Tube((0., -.5), math.pi/2, 2., .5), Parallel(0.))  # selects rows 0 and 2
    for body_id, ok in ((np.array([0, 1, 0], np.int32), True), (np.array([0, 1, 1], np.int32), False)):
        call = lambda: surface_link_velocity(objs, cfg, flags, fluid_idx, dirs, body_id,
                                             zeros, zeros, jets=[jet], hits=hits)
        if ok:
            _, aperture = call()
            assert aperture.tolist() == [True, False, True]
        else:
            with pytest.raises(ValueError, match="one jet may drive only one body"):
                call()