"""
test_person_tracker.py — Unit tests for the P0 person tracking pipeline.

Tests cover: IoU, Kalman filter, colour histogram, bearing geometry,
depth sampling, tracker state machine, and follow-target selection.
No camera, no detector, no ROS2 needed.
"""

import math

import numpy as np
import pytest

from saltybot_bringup._person_tracker import (
    BBox,
    CamParams,
    Detection,
    FollowTargetSelector,
    KalmanBoxFilter,
    PersonTrack,
    PersonTracker,
    TrackState,
    DEPTH_GOOD, DEPTH_MARGINAL, DEPTH_EXTRAPOLATED, DEPTH_INVALID,
    bearing_from_pixel,
    depth_at_bbox,
    extract_torso_hist,
    hist_similarity,
    iou,
)


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _det(x=10, y=20, w=60, h=150, conf=0.85, frame=None) -> Detection:
    return Detection(BBox(x, y, w, h), conf, frame)


def _solid_bgr(h=200, w=100, b=128, g=64, r=32) -> np.ndarray:
    """Uniform colour BGR image."""
    return np.full((h, w, 3), (b, g, r), dtype=np.uint8)


def _depth_image(h=480, w=640, val_mm=2000) -> np.ndarray:
    """Uniform uint16 depth image at val_mm mm."""
    return np.full((h, w), val_mm, dtype=np.uint16)


def _run_tracker_n(tracker, bbox, n=5, cam=None):
    """Feed the same detection to a tracker n times."""
    for _ in range(n):
        tracker.update([_det(*bbox)], cam=cam)


# ─────────────────────────────────────────────────────────────────────────────
# BBox
# ─────────────────────────────────────────────────────────────────────────────

def test_bbox_fields():
    b = BBox(10, 20, 60, 150)
    assert b.x == 10 and b.y == 20 and b.w == 60 and b.h == 150


def test_bbox_is_named_tuple():
    b = BBox(0, 0, 1, 1)
    assert isinstance(b, tuple)


# ─────────────────────────────────────────────────────────────────────────────
# IoU
# ─────────────────────────────────────────────────────────────────────────────

def test_iou_identical():
    b = BBox(0, 0, 100, 100)
    assert abs(iou(b, b) - 1.0) < 1e-6


def test_iou_no_overlap():
    a = BBox(0,   0, 50, 50)
    b = BBox(100, 0, 50, 50)
    assert iou(a, b) == pytest.approx(0.0, abs=1e-6)


def test_iou_partial_overlap():
    a = BBox(0, 0, 100, 100)
    b = BBox(50, 0, 100, 100)  # 50 % horizontal overlap
    result = iou(a, b)
    assert 0.0 < result < 1.0


def test_iou_exact_overlap_value():
    a = BBox(0, 0, 100, 100)     # area 10000
    b = BBox(0, 0, 50, 100)      # area 5000, inter=5000, union=10000
    assert abs(iou(a, b) - 0.5) < 1e-6


def test_iou_contained():
    a = BBox(0, 0, 100, 100)
    b = BBox(25, 25, 50, 50)     # b inside a
    result = iou(a, b)
    # inter = 2500, union = 10000+2500-2500 = 10000
    assert abs(result - 0.25) < 1e-6


def test_iou_symmetric():
    a = BBox(10, 10, 80, 80)
    b = BBox(50, 50, 80, 80)
    assert abs(iou(a, b) - iou(b, a)) < 1e-9


def test_iou_touching_edges():
    a = BBox(0, 0, 50, 50)
    b = BBox(50, 0, 50, 50)   # touch at x=50 → zero overlap
    assert iou(a, b) == pytest.approx(0.0, abs=1e-6)


# ─────────────────────────────────────────────────────────────────────────────
# KalmanBoxFilter
# ─────────────────────────────────────────────────────────────────────────────

def test_kalman_predict_returns_bbox():
    kf = KalmanBoxFilter(BBox(10, 20, 60, 150))
    pred = kf.predict()
    assert isinstance(pred, BBox)
    assert pred.w >= 1 and pred.h >= 1


def test_kalman_update_returns_bbox():
    kf = KalmanBoxFilter(BBox(10, 20, 60, 150))
    updated = kf.update(BBox(10, 20, 60, 150))
    assert isinstance(updated, BBox)


def test_kalman_converges_to_stationary():
    """After many identical measurements, Kalman should converge near them."""
    init  = BBox(100, 100, 80, 160)
    meas  = BBox(102, 98,  80, 160)
    kf    = KalmanBoxFilter(init)
    for _ in range(20):
        kf.predict()
        result = kf.update(meas)
    # Should be within ~10 px of measurement
    assert abs(result.x - meas.x) < 10
    assert abs(result.y - meas.y) < 10


def test_kalman_predict_advances_with_velocity():
    """Give the filter a few frames of rightward motion; prediction overshoots."""
    kf = KalmanBoxFilter(BBox(100, 100, 60, 150))
    for i in range(5):
        kf.predict()
        kf.update(BBox(100 + i * 10, 100, 60, 150))
    pred = kf.predict()
    # After motion right, predicted cx should be further right
    assert pred.x > 100 + 4 * 10 - 5  # at least near last position


def test_kalman_velocity_zero_initially():
    kf = KalmanBoxFilter(BBox(100, 100, 60, 150))
    assert kf.velocity_px == (0.0, 0.0)


def test_kalman_bbox_property():
    b = BBox(50, 50, 80, 120)
    kf = KalmanBoxFilter(b)
    out = kf.bbox
    assert isinstance(out, BBox)
    # Should be near initial
    assert abs(out.x - b.x) <= 2
    assert abs(out.y - b.y) <= 2


def test_kalman_width_height_stay_positive():
    kf = KalmanBoxFilter(BBox(10, 10, 5, 5))
    for _ in range(30):
        kf.predict()
    assert kf.bbox.w >= 1 and kf.bbox.h >= 1


# ─────────────────────────────────────────────────────────────────────────────
# extract_torso_hist
# ─────────────────────────────────────────────────────────────────────────────

def test_torso_hist_shape():
    bgr = _solid_bgr()
    h = extract_torso_hist(bgr, BBox(0, 0, 100, 200))
    assert h is not None
    assert h.shape == (128,)   # 16 * 8


def test_torso_hist_normalised():
    bgr = _solid_bgr()
    h = extract_torso_hist(bgr, BBox(0, 0, 100, 200))
    assert h is not None
    assert abs(h.sum() - 1.0) < 1e-5


def test_torso_hist_none_for_none_frame():
    assert extract_torso_hist(None, BBox(0, 0, 100, 200)) is None


def test_torso_hist_none_for_tiny_bbox():
    bgr = _solid_bgr()
    assert extract_torso_hist(bgr, BBox(0, 0, 2, 2)) is None


def test_torso_hist_different_colours():
    """Two differently coloured crops should produce different histograms."""
    bgr_red  = np.full((200, 100, 3), (0, 0, 255), dtype=np.uint8)
    bgr_blue = np.full((200, 100, 3), (255, 0, 0), dtype=np.uint8)
    h_red  = extract_torso_hist(bgr_red,  BBox(0, 0, 100, 200))
    h_blue = extract_torso_hist(bgr_blue, BBox(0, 0, 100, 200))
    assert h_red is not None and h_blue is not None
    assert not np.allclose(h_red, h_blue)


# ─────────────────────────────────────────────────────────────────────────────
# hist_similarity
# ─────────────────────────────────────────────────────────────────────────────

def test_hist_similarity_identical():
    h = np.ones(128, dtype=np.float32) / 128.0
    assert abs(hist_similarity(h, h) - 1.0) < 1e-6


def test_hist_similarity_orthogonal():
    """Non-overlapping histograms → 0 similarity."""
    h1 = np.zeros(128, dtype=np.float32)
    h2 = np.zeros(128, dtype=np.float32)
    h1[:64]  = 1.0 / 64
    h2[64:]  = 1.0 / 64
    assert abs(hist_similarity(h1, h2)) < 1e-6


def test_hist_similarity_range():
    rng = np.random.RandomState(0)
    for _ in range(20):
        h1 = rng.rand(128).astype(np.float32)
        h2 = rng.rand(128).astype(np.float32)
        h1 /= h1.sum(); h2 /= h2.sum()
        s = hist_similarity(h1, h2)
        assert 0.0 <= s <= 1.0 + 1e-6


def test_hist_similarity_symmetric():
    h1 = np.random.RandomState(1).rand(128).astype(np.float32)
    h2 = np.random.RandomState(2).rand(128).astype(np.float32)
    h1 /= h1.sum(); h2 /= h2.sum()
    assert abs(hist_similarity(h1, h2) - hist_similarity(h2, h1)) < 1e-6


# ─────────────────────────────────────────────────────────────────────────────
# bearing_from_pixel
# ─────────────────────────────────────────────────────────────────────────────

def test_bearing_centre_is_zero():
    """Pixel at principal point → bearing = 0°."""
    assert abs(bearing_from_pixel(320.0, 320.0, 615.0)) < 1e-9


def test_bearing_right_is_positive():
    b = bearing_from_pixel(400.0, 320.0, 615.0)
    assert b > 0.0


def test_bearing_left_is_negative():
    b = bearing_from_pixel(200.0, 320.0, 615.0)
    assert b < 0.0


def test_bearing_symmetric():
    b_right = bearing_from_pixel(400.0, 320.0, 615.0)
    b_left  = bearing_from_pixel(240.0, 320.0, 615.0)
    assert abs(b_right + b_left) < 1e-9


def test_bearing_approx_at_45_deg():
    """u - cx = fx → atan(1) = 45°."""
    bearing = bearing_from_pixel(935.0, 320.0, 615.0)
    assert abs(bearing - 45.0) < 0.1


def test_bearing_degrees_not_radians():
    """Result must be in degrees (much larger than a radian value)."""
    b = bearing_from_pixel(400.0, 320.0, 615.0)
    assert abs(b) > 0.5   # atan2(80/615) ≈ 7.4°


# ─────────────────────────────────────────────────────────────────────────────
# depth_at_bbox
# ─────────────────────────────────────────────────────────────────────────────

def test_depth_uniform_field():
    d = _depth_image(val_mm=2000)
    depth_m, quality = depth_at_bbox(d, BBox(200, 150, 80, 200))
    assert abs(depth_m - 2.0) < 0.01
    assert quality == DEPTH_GOOD


def test_depth_zero_image_invalid():
    d = np.zeros((480, 640), dtype=np.uint16)
    depth_m, quality = depth_at_bbox(d, BBox(200, 150, 80, 200))
    assert depth_m == 0.0
    assert quality == DEPTH_INVALID


def test_depth_partial_fill_marginal():
    d = np.zeros((480, 640), dtype=np.uint16)
    # Fill only 40 % of the central window with valid readings
    d[200:260, 280:360] = 1500
    _, quality = depth_at_bbox(d, BBox(200, 150, 160, 200), window_frac=1.0)
    assert quality in (DEPTH_MARGINAL, DEPTH_EXTRAPOLATED, DEPTH_GOOD)


def test_depth_scale_applied():
    d = _depth_image(val_mm=3000)
    depth_m, _ = depth_at_bbox(d, BBox(100, 100, 80, 150), depth_scale=0.001)
    assert abs(depth_m - 3.0) < 0.01


def test_depth_out_of_bounds_bbox():
    """Bbox outside image should return DEPTH_INVALID."""
    d = _depth_image()
    _, quality = depth_at_bbox(d, BBox(700, 500, 80, 100))  # off-screen
    assert quality == DEPTH_INVALID


def test_depth_at_5m():
    d = _depth_image(val_mm=5000)
    depth_m, _ = depth_at_bbox(d, BBox(200, 100, 80, 200))
    assert abs(depth_m - 5.0) < 0.05


# ─────────────────────────────────────────────────────────────────────────────
# PersonTracker — state machine
# ─────────────────────────────────────────────────────────────────────────────

def test_tracker_new_track_tentative():
    t = PersonTracker(min_hits=3)
    active = t.update([_det()])
    # min_hits=3: first frame → TENTATIVE, not in active output
    assert len(active) == 0
    assert len(t.tracks) == 1
    assert t.tracks[0].state == TrackState.TENTATIVE


def test_tracker_track_becomes_active():
    t = PersonTracker(min_hits=3)
    for _ in range(3):
        active = t.update([_det()])
    assert len(active) == 1
    assert active[0].state == TrackState.ACTIVE


def test_tracker_persistent_id():
    t = PersonTracker(min_hits=2)
    _run_tracker_n(t, (10, 20, 60, 150), n=4)
    assert len(t.tracks) == 1
    assert t.tracks[0].track_id == 0


def test_tracker_two_detections_two_tracks():
    t = PersonTracker(min_hits=2)
    # Two widely-separated boxes → two tracks
    dets = [
        Detection(BBox(10,  20, 60, 150), 0.9),
        Detection(BBox(400, 20, 60, 150), 0.9),
    ]
    for _ in range(3):
        t.update(dets)
    assert len(t.tracks) == 2


def test_tracker_ids_increment():
    t = PersonTracker(min_hits=1)
    t.update([_det(x=10)])
    t.update([_det(x=10), _det(x=300)])
    ids = {trk.track_id for trk in t.tracks}
    assert len(ids) == 2


def test_tracker_track_goes_lost():
    t = PersonTracker(min_hits=2, max_lost_frames=5)
    _run_tracker_n(t, (10, 20, 60, 150), n=3)
    # Now send no detections
    t.update([])
    lost = [tr for tr in t.tracks if tr.state == TrackState.LOST]
    assert len(lost) == 1


def test_tracker_track_removed_after_max_lost():
    t = PersonTracker(min_hits=2, max_lost_frames=3)
    _run_tracker_n(t, (10, 20, 60, 150), n=3)
    for _ in range(5):
        t.update([])
    assert len(t.tracks) == 0


def test_tracker_iou_matching_same_track():
    """Slightly moved detection should match the existing track (not create new)."""
    t = PersonTracker(min_hits=2, iou_threshold=0.3)
    _run_tracker_n(t, (10, 20, 60, 150), n=3)
    n_before = len(t.tracks)
    t.update([_det(x=15, y=20, w=60, h=150)])   # small shift, high IoU
    assert len(t.tracks) == n_before


def test_tracker_no_overlap_creates_new_track():
    t = PersonTracker(min_hits=2, iou_threshold=0.3)
    _run_tracker_n(t, (10, 20, 60, 150), n=3)
    t.update([_det(x=500, y=20, w=60, h=150)])   # far away → new track
    assert len(t.tracks) == 2   # old (lost) + new (tentative)


def test_tracker_reset():
    t = PersonTracker(min_hits=2)
    _run_tracker_n(t, (10, 20, 60, 150), n=3)
    t.reset()
    assert len(t.tracks) == 0
    assert t._next_id == 0


def test_tracker_bearing_set_with_cam():
    cam = CamParams(fx=615.0, cx=320.0)
    t = PersonTracker(min_hits=2)
    _run_tracker_n(t, (10, 20, 60, 150), n=3, cam=cam)
    active = [tr for tr in t.tracks if tr.state == TrackState.ACTIVE]
    assert len(active) > 0
    # bearing of a box at x=10..70 (cx ≈ 40, image cx=320) → negative (left)
    assert active[0].bearing < 0.0


def test_tracker_distance_set_with_depth():
    cam = CamParams(fx=615.0, cx=320.0)
    t = PersonTracker(min_hits=2)
    d = _depth_image(val_mm=2000)
    for _ in range(3):
        t.update([_det(x=200, y=100, w=80, h=200)], cam=cam, depth_u16=d)
    active = [tr for tr in t.tracks if tr.state == TrackState.ACTIVE]
    assert len(active) > 0
    assert abs(active[0].distance - 2.0) < 0.1


def test_tracker_no_depth_distance_zero():
    cam = CamParams()
    t = PersonTracker(min_hits=2)
    for _ in range(3):
        t.update([_det()], cam=cam, depth_u16=None)
    active = [tr for tr in t.tracks if tr.state == TrackState.ACTIVE]
    assert len(active) > 0
    assert active[0].distance == 0.0


# ─────────────────────────────────────────────────────────────────────────────
# PersonTracker — re-identification
# ─────────────────────────────────────────────────────────────────────────────

def test_reid_restores_track_id():
    """Person disappears for a few frames then reappears; same track_id."""
    # Blue person in centre
    bgr_blue = np.full((480, 640, 3), (200, 50, 0), dtype=np.uint8)
    bbox = (250, 50, 80, 200)
    t = PersonTracker(
        min_hits=2, max_lost_frames=10,
        reid_threshold=0.4, reid_max_dist=200.0,
    )

    # Establish track
    for _ in range(3):
        t.update([Detection(BBox(*bbox), 0.9, bgr_blue)])

    track_id_before = t.tracks[0].track_id

    # Disappear for 3 frames
    for _ in range(3):
        t.update([])

    # Re-appear at same position with same colour
    t.update([Detection(BBox(*bbox), 0.9, bgr_blue)])

    # Track should be re-identified with the same ID
    assert any(tr.track_id == track_id_before for tr in t.tracks)


def test_reid_different_colour_creates_new_track():
    """Person disappears; different colour appears same place → new track ID."""
    bgr_blue = np.full((480, 640, 3), (200, 50, 0), dtype=np.uint8)
    bgr_red  = np.full((480, 640, 3), (0, 50, 200), dtype=np.uint8)
    bbox = (250, 50, 80, 200)
    t = PersonTracker(
        min_hits=2, max_lost_frames=5,
        reid_threshold=0.85,   # strict threshold → red won't match blue
        reid_max_dist=200.0,
    )

    for _ in range(3):
        t.update([Detection(BBox(*bbox), 0.9, bgr_blue)])

    track_id_before = t.tracks[0].track_id

    for _ in range(2):
        t.update([])

    # Red person (different colour)
    for _ in range(3):
        t.update([Detection(BBox(*bbox), 0.9, bgr_red)])

    new_ids = {tr.track_id for tr in t.tracks}
    # Should contain a new ID (not only the original)
    assert track_id_before not in new_ids or len(new_ids) > 1


# ─────────────────────────────────────────────────────────────────────────────
# FollowTargetSelector
# ─────────────────────────────────────────────────────────────────────────────

def _make_track(track_id, x, y, w, h, dist=0.0) -> PersonTrack:
    trk = PersonTrack(
        track_id=track_id, state=TrackState.ACTIVE,
        bbox=BBox(x, y, w, h), distance=dist,
    )
    return trk


def test_selector_inactive_returns_none():
    sel = FollowTargetSelector()
    sel.stop()
    result = sel.update([_make_track(0, 100, 50, 80, 200, dist=3.0)])
    assert result is None


def test_selector_active_selects_nearest_by_distance():
    sel = FollowTargetSelector()
    sel.start()
    tracks = [
        _make_track(0, 100, 50, 80, 200, dist=4.0),
        _make_track(1, 300, 50, 80, 200, dist=1.5),
    ]
    result = sel.update(tracks, img_cx=320.0)
    assert result.track_id == 1   # closer


def test_selector_locks_on_same_track():
    sel = FollowTargetSelector()
    sel.start()
    tracks = [
        _make_track(0, 100, 50, 80, 200, dist=2.0),
        _make_track(1, 300, 50, 80, 200, dist=5.0),
    ]
    first = sel.update(tracks, img_cx=320.0)
    # Second call — switch distance but should keep locked ID
    tracks[0] = _make_track(0, 100, 50, 80, 200, dist=2.0)
    second = sel.update(tracks, img_cx=320.0)
    assert second.track_id == first.track_id


def test_selector_holds_last_when_tracks_empty():
    sel = FollowTargetSelector(hold_frames=5)
    sel.start()
    sel.update([_make_track(0, 100, 50, 80, 200, dist=2.0)])
    # Now empty — should hold for up to 5 frames
    result = sel.update([], img_cx=320.0)
    assert result is not None
    assert result.track_id == 0


def test_selector_stops_holding_after_hold_frames():
    sel = FollowTargetSelector(hold_frames=2)
    sel.start()
    sel.update([_make_track(0, 100, 50, 80, 200)])
    for _ in range(5):
        result = sel.update([])
    assert result is None


def test_selector_selects_by_centre_when_no_depth():
    sel = FollowTargetSelector()
    sel.start()
    # Track 0: x=0..80 → centre x=40 (far left of image)
    # Track 1: x=280..360 → centre x=320 (near image centre)
    tracks = [
        _make_track(0, 0,   50, 80, 200, dist=0.0),
        _make_track(1, 280, 50, 80, 200, dist=0.0),
    ]
    result = sel.update(tracks, img_cx=320.0)
    assert result.track_id == 1   # nearer to image centre


def test_selector_restart_reselects():
    sel = FollowTargetSelector()
    sel.start()
    t0 = _make_track(0, 100, 50, 80, 200, dist=2.0)
    sel.update([t0])
    sel.stop()
    sel.start()
    t1 = _make_track(1, 300, 50, 80, 200, dist=1.0)
    result = sel.update([t0, t1])
    assert result.track_id == 1   # re-selected nearest
