"""동적 객체 분리 — 박스 기하, 판정, 보간, 투영."""

import math

import pytest

np = pytest.importorskip("numpy")

from nusim.data.index import Box, CameraCalib, Frame, SceneIndex
from nusim.data.rig import box_corners, camera_pose, points_in_box
from nusim.recon.dynamic import (
    _hull, box_at, boxes_at, build_tracks, dynamic_instances, dynamic_mask, end_displacement, masked_instances,
    project_box, static_instances,
)

_YAW90 = (math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5))  # z축 +90°


def test_box_corners_axes_follow_devkit():
    # size=(w,l,h). 회전 없음이면 x 방향 반길이가 l/2, y 가 w/2.
    c = box_corners((0, 0, 0), (2.0, 4.0, 1.0), (1, 0, 0, 0))
    assert np.allclose(np.abs(c).max(0), [2.0, 1.0, 0.5])
    # +90° 요 → 길이 방향이 전역 y 로
    c = box_corners((10, 0, 0), (2.0, 4.0, 1.0), _YAW90)
    assert np.allclose(np.abs(c - [10, 0, 0]).max(0), [1.0, 2.0, 0.5])


def test_points_in_box_matches_corners():
    rng = np.random.default_rng(0)
    t, size, q = (5.0, -3.0, 1.0), (1.8, 4.5, 1.6), _YAW90
    # 꼭짓점을 중심 쪽으로 살짝 당기면 안, 밖으로 밀면 밖
    c = box_corners(t, size, q)
    assert points_in_box(t + (c - t) * 0.99, t, size, q).all()
    assert not points_in_box(t + (c - t) * 1.01, t, size, q).any()
    # pad 는 모든 방향으로 늘린다
    assert points_in_box(t + (c - t) * 1.01, t, size, q, pad=0.1).all()
    # 무작위 점: 로컬 좌표 판정과 일치
    p = rng.uniform(-6, 6, size=(2000, 3)) + t
    inside = points_in_box(p, t, size, q)
    local = p - t
    ref = (np.abs(local[:, 1]) <= 2.25) & (np.abs(local[:, 0]) <= 0.9) & (np.abs(local[:, 2]) <= 0.8)
    assert np.array_equal(inside, ref)


def _scene_with_tracks():
    s = SceneIndex(token="t", name="scene-test", location="singapore-onenorth", description="",
                   log_token="l", date_captured="", vehicle="", nbr_samples=10)
    s.instances = [("a", "vehicle.car"), ("b", "vehicle.car"), ("c", "human.pedestrian.adult"),
                   ("d", "vehicle.car")]
    rng = np.random.default_rng(1)
    q = (1.0, 0.0, 0.0, 0.0)
    for k in range(10):
        ts = 1_000_000 + 500_000 * k
        # a: 주차 차량 — 제자리지만 어노테이션이 ±0.6 m 흔들린다
        s.boxes.append(Box(ts, 0, tuple(np.array([10.0, 0, 0]) + rng.uniform(-0.6, 0.6, 3)),
                           (1.8, 4.5, 1.6), q, "vehicle.parked"))
        # b: 초당 4 m 로 달리는 차
        s.boxes.append(Box(ts, 1, (20.0 + 2.0 * k, 5.0, 0.0), (1.8, 4.5, 1.6), q, "vehicle.moving"))
        # c: 서 있는 보행자
        s.boxes.append(Box(ts, 2, (0.0, 8.0, 0.0), (0.6, 0.6, 1.7), q, "pedestrian.standing"))
    # d: 한 키프레임만 스쳐 간 차 (변위 0 이지만 moving)
    s.boxes.append(Box(1_000_000, 3, (30.0, -5.0, 0.0), (1.8, 4.5, 1.6), q, "vehicle.moving"))
    return s


def test_dynamic_rule():
    s = _scene_with_tracks()
    tr = build_tracks(s)
    assert end_displacement(tr[0]) < 1.0  # 흔들림은 상쇄된다
    assert end_displacement(tr[1]) > 10.0
    assert dynamic_instances(s) == {1, 2, 3}  # 주차 차량만 배경


def test_masked_instances_agents_removes_parked_but_keeps_map_objects():
    s = _scene_with_tracks()
    s.instances.append(("e", "movable_object.trafficcone"))
    s.boxes.append(Box(1_000_000, 4, (5.0, 5.0, 0.0), (0.4, 0.4, 0.8), (1.0, 0.0, 0.0, 0.0), ""))
    assert masked_instances(s, "dynamic") == {1, 2, 3}
    assert masked_instances(s, "agents") == {0, 1, 2, 3}  # 주차 차량까지, 콘은 맵의 일부
    with pytest.raises(ValueError):
        masked_instances(s, "everything")


def test_box_interpolation():
    tr = build_tracks(_scene_with_tracks())[1]
    # 두 키프레임 한가운데
    pos, q = box_at(tr, 1_000_000 + 250_000)
    assert np.allclose(pos, [21.0, 5.0, 0.0])
    assert np.isclose(np.linalg.norm(q), 1.0)
    # 끝 너머 EDGE_HOLD 안쪽은 끝값, 바깥은 None
    assert np.allclose(box_at(tr, 1_000_000 - 200_000)[0], [20.0, 5.0, 0.0])
    assert box_at(tr, 1_000_000 - 300_000) is None


def _front_pose():
    calib = CameraCalib(translation=(0.0, 0.0, 1.5), rotation=(0.5, -0.5, 0.5, -0.5),
                        intrinsic=((1000.0, 0, 800.0), (0, 1000.0, 450.0), (0, 0, 1)),
                        width=1600, height=900)
    fr = Frame(0, "x.jpg", (0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0), True)
    return camera_pose(fr, calib, "CAM_FRONT")  # ego +x 를 본다


def test_project_box_in_front():
    pose = _front_pose()
    poly = project_box(pose, box_corners((20.0, 0.0, 1.5), (2.0, 2.0, 2.0), (1, 0, 0, 0)))
    # 20 m 앞, 폭 2 m → 대략 ±1000*1/19 ~ ±1000*1/21 픽셀. 화면 중앙 둘레.
    assert poly is not None
    assert 700 < poly[:, 0].min() < 760 and 840 < poly[:, 0].max() < 900
    assert poly[:, 1].min() < 450 < poly[:, 1].max()


def test_project_box_behind_is_none_and_straddling_is_clipped():
    pose = _front_pose()
    assert project_box(pose, box_corners((-20.0, 0.0, 1.5), (2, 2, 2), (1, 0, 0, 0))) is None
    # 카메라를 감싸는 박스: 절반은 뒤에 있다. 자르지 않으면 부호가 뒤집혀
    # 화면 반대편에 찍힌다. 잘랐으면 화면 전체를 덮는 거대한 다각형이 나와야 한다.
    poly = project_box(pose, box_corners((0.0, 0.0, 1.5), (4, 4, 4), (1, 0, 0, 0)))
    assert poly[:, 0].min() < 0 and poly[:, 0].max() > 1600
    assert poly[:, 1].min() < 0 and poly[:, 1].max() > 900


def test_dynamic_mask_covers_box():
    pose = _front_pose()
    m = dynamic_mask(pose, [((20.0, 0.0, 1.5), (2.0, 2.0, 2.0), (1, 0, 0, 0))], 800, 450,
                     scale=2, pad=0.0)
    assert m[225, 400] and not m[10, 10]
    assert 0.001 < m.mean() < 0.05


def test_hull_square():
    pts = np.array([[0, 0], [1, 0], [1, 1], [0, 1], [0.5, 0.5], [0.2, 0.7]], dtype=float)
    h = _hull(pts)
    assert len(h) == 4
    assert {tuple(p) for p in h} == {(0, 0), (1, 0), (1, 1), (0, 1)}


def test_static_agents_masked_outside_annotation_window():
    """주차 차량은 어노테이션 구간(키프레임) 밖 시각에도 같은 자리에 가린다."""
    s = _scene_with_tracks()
    tr = build_tracks(s)
    late = 100_000_000  # 어노테이션이 끝나고 한참 뒤
    assert boxes_at(tr, {0}, late) == []
    assert static_instances(s, "dynamic") == set()
    st = static_instances(s, "agents")
    assert st == {0}  # 움직이는 차·보행자는 여전히 시간별 박스
    (pos, size, rot), = boxes_at(tr, {0}, late, static=st)
    assert np.linalg.norm(pos - [10.0, 0, 0]) < 0.7  # 흔들린 어노테이션의 중앙값