"""Grok review cases for claude/jet (test-only).

The oracle and wrap tests are Grok's inline patch (Slack ts 1792052057.423679).
The rotated-tube and square-wall cases follow Grok's stated checks; their text
was clipped in transit, so they are re-written here, not copied verbatim.
"""
import math
import numpy as np
from CelerisLab.body.jets import Jet, Parallel, Radial, Sector, Tube, jet_link_shapes
from .test_surface_actuation import circle


def big_scene():
    from CelerisLab.config import LBMConfig
    from CelerisLab.body.manager import ObjectManager
    from CelerisLab.lbm.descriptors import FLUID
    obj = circle(center=(20., 15.), radius=6.5)
    cfg = LBMConfig(nx=40, ny=30, viscosity=.1); cfg.validate()
    flags = ObjectManager.build_flags_for([obj], np.full(40*30, FLUID, np.uint16), nx=40, ny=30, nz=1)
    links = obj.geometry.build_cut_links(40, 30, domain_flags=flags)
    hits = np.array([(l.hit_x, l.hit_y) for l in links])
    fl, dr = obj.get_curved_list(40, 30, cfg.nq, domain_flags=flags)[:2]
    return hits, np.asarray(fl), np.asarray(dr)


def test_frozen_5_link_sector_oracle_with_edge_margin():
    hits, fl, dr = big_scene()
    half = math.radians(5)
    sel, s = Sector((20., 15.), 0., half, 7.5, 5.5).select(hits)
    got = sorted((int(f % 40), int(f // 40), int(d)) for f, d in zip(fl[sel], dr[sel]))
    assert got == sorted([(27, 15, 2), (27, 14, 8), (27, 16, 6), (27, 15, 6), (27, 15, 8)])
    d = hits - (20., 15.)
    dth = np.abs(np.arctan2(d[:, 1], d[:, 0]))
    near = np.hypot(*d.T) < 7.5
    assert np.all(np.abs(dth[near] - half) >= 1e-3)   # no hit within 1e-3 rad of the edge


def test_sector_axis_pi_wraps():
    hits = np.array([[-1., 1e-3], [-1., -1e-3], [1., 0.]])
    sel, _ = Sector((0., 0.), math.pi, .01, 2.).select(hits)
    assert sel.tolist() == [True, True, False]
    sel, _ = Sector((0., 0.), -math.pi, .01, 2.).select(hits)
    assert sel.tolist() == [True, True, False]


def test_rotated_tube_30deg_selection_and_parallel_axis():
    a = math.radians(30)
    u, n = np.array([math.cos(a), math.sin(a)]), np.array([-math.sin(a), math.cos(a)])
    hits = np.array([2*u, 2*u + .999*n, 2*u - .999*n, 2*u + 1.001*n, -.1*u, 4.1*u])
    jet = Jet(Tube((0., 0.), a, 4., 1.), Parallel(a))
    shape, jet_id = jet_link_shapes([jet], hits)
    assert (jet_id == 0).tolist() == [True, True, True, False, False, False]
    np.testing.assert_allclose(shape[:3], np.tile(u, (3, 1)), rtol=0, atol=1e-15)
    np.testing.assert_array_equal(shape[3:], 0.)


def test_exact_edge_rule_axis_aligned():
    # Exactly representable samples: lateral |s| = 1 is excluded; Tube along
    # 0 and length and Sector r_min/r_max are included.
    tube = Tube((0., 0.), 0., 4., 1.)
    sel, s = tube.select(np.array([[2., 1.], [2., -1.], [2., .5], [0., 0.], [4., 0.]]))
    assert sel.tolist() == [False, False, True, True, True]
    assert s.tolist() == [1., -1., .5, 0., 0.]
    sel, _ = Sector((0., 0.), 0., .5, 2., 1.).select(np.array([[1., 0.], [2., 0.], [.5, 0.], [2.5, 0.]]))
    assert sel.tolist() == [True, True, False, False]


def test_synthetic_square_wall_off_centre_radial():
    # Synthetic hit sampling on a flat face x=10 (no square geometry, cut links or
    # kernel involved): selection is purely spatial and the field points away
    # from (5, 0), not along the face normal (1, 0).
    y = np.linspace(-2., 2., 9)
    hits = np.stack((np.full_like(y, 10.), y), axis=1)
    jet = Jet(Tube((10., -1.), math.pi/2, 2., .5), Radial((5., 0.), inward=False), amplitude=1e-3)
    shape, jet_id = jet_link_shapes([jet], hits)
    assert np.array_equal(jet_id == 0, np.abs(y) <= 1)
    d = hits[jet_id == 0] - (5., 0.)
    np.testing.assert_allclose(shape[jet_id == 0], d/np.hypot(*d.T)[:, None], rtol=0, atol=1e-15)
    assert np.abs(shape[jet_id == 0, 1]).max() > .1  # not the wall normal