"""
test_person_reid.py — Unit tests for person re-identification helpers (no ROS2).

Covers:
  extract_hsv_histogram:
    - output shape is FEAT_DIM (128)
    - None input → zero vector
    - zero-size array input → zero vector
    - solid-colour image → L2-normalised (norm ≈ 1)
    - same colour on different image sizes → same histogram
    - different solid colours → dissimilar histograms
    - all-white image (low saturation) → histogram concentrated in low-S bins
    - single-pixel image → valid unit histogram

  cosine_similarity:
    - identical unit vectors → 1.0
    - zero vector (a) → 0.0
    - zero vector (b) → 0.0
    - both zero → 0.0
    - anti-parallel vectors → -1.0
    - orthogonal vectors → 0.0
    - non-unit same-direction vectors → 1.0
    - non-unit opposite-direction vectors → -1.0
    - partial overlap → in (0, 1)

  match_track:
    - empty gallery → (None, 0.0)
    - perfect match above threshold → correct track_id, sim ≈ 1.0
    - similarity exactly at threshold → no match (strict >)
    - below threshold → (None, 0.0)
    - stale track excluded from matching
    - multiple tracks → best one returned
    - all tracks stale → (None, 0.0)

  TrackGallery:
    - add assigns track_id = 1, 2, 3 ...
    - add stores a copy of feature (mutation of original doesn't affect gallery)
    - update changes last_seen, camera_id, confidence
    - update applies EMA: feature shifts toward new, stays unit-norm
    - update clears is_stale flag
    - update preserves first_seen
    - mark_stale: recent track stays fresh
    - mark_stale: old track flagged stale
    - prune_stale removes stale, keeps fresh
    - empty gallery match → (None, 0.0)

  Re-id integration:
    - same appearance from two camera_ids → same track_id
    - different appearance → different track_id
    - stale + prune + re-add → new track_id
    - EMA: track adapts over multiple updates (feature norm stays ≈ 1)
"""

import sys
import os
import math

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_bringup._person_reid import (
    H_BINS,
    S_BINS,
    FEAT_DIM,
    TrackEntry,
    TrackGallery,
    cosine_similarity,
    extract_hsv_histogram,
    match_track,
)


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

def _solid_bgr(b: int, g: int, r: int, size: int = 32) -> np.ndarray:
    """Return a solid-colour BGR uint8 image of given size."""
    import cv2
    img = np.full((size, size, 3), (b, g, r), dtype=np.uint8)
    return img


def _unit(v: np.ndarray) -> np.ndarray:
    n = np.linalg.norm(v)
    return v / n if n > 1e-12 else v


def _make_feature(seed: int) -> np.ndarray:
    """Reproducible random L2-normalised feature vector."""
    rng = np.random.default_rng(seed)
    v   = rng.random(FEAT_DIM).astype(np.float32)
    return v / np.linalg.norm(v)


def _make_entry(
    tid: int,
    feature: np.ndarray,
    camera_id: str = 'front',
    now: float = 1000.0,
    is_stale: bool = False,
) -> TrackEntry:
    return TrackEntry(
        track_id=tid,
        camera_id=camera_id,
        feature=feature.copy(),
        confidence=1.0,
        first_seen=now,
        last_seen=now,
        is_stale=is_stale,
    )


# ── extract_hsv_histogram ─────────────────────────────────────────────────────

class TestExtractHsvHistogram:

    def test_output_shape(self):
        img  = _solid_bgr(0, 0, 255)
        hist = extract_hsv_histogram(img)
        assert hist.shape == (FEAT_DIM,)

    def test_output_dtype(self):
        hist = extract_hsv_histogram(_solid_bgr(0, 255, 0))
        assert hist.dtype == np.float32

    def test_none_returns_zero(self):
        hist = extract_hsv_histogram(None)
        assert hist.shape == (FEAT_DIM,)
        assert np.all(hist == 0.0)

    def test_empty_array_returns_zero(self):
        import numpy as np
        empty = np.zeros((0, 0, 3), dtype=np.uint8)
        hist  = extract_hsv_histogram(empty)
        assert np.all(hist == 0.0)

    def test_solid_colour_is_unit_norm(self):
        hist = extract_hsv_histogram(_solid_bgr(255, 0, 0))
        assert np.linalg.norm(hist) == pytest.approx(1.0, rel=1e-5)

    def test_same_colour_different_sizes_equal(self):
        h16 = extract_hsv_histogram(_solid_bgr(0, 0, 200, size=16))
        h64 = extract_hsv_histogram(_solid_bgr(0, 0, 200, size=64))
        assert np.allclose(h16, h64, atol=1e-5)

    def test_different_colours_different_histograms(self):
        red  = extract_hsv_histogram(_solid_bgr(0, 0, 255))
        blue = extract_hsv_histogram(_solid_bgr(255, 0, 0))
        assert not np.allclose(red, blue, atol=1e-3)

    def test_red_vs_green_low_similarity(self):
        red   = extract_hsv_histogram(_solid_bgr(0, 0, 255))
        green = extract_hsv_histogram(_solid_bgr(0, 255, 0))
        sim   = cosine_similarity(red, green)
        assert sim < 0.5

    def test_single_pixel_valid(self):
        import numpy as np
        px   = np.array([[[0, 200, 200]]], dtype=np.uint8)
        hist = extract_hsv_histogram(px)
        assert hist.shape == (FEAT_DIM,)
        assert np.linalg.norm(hist) == pytest.approx(1.0, rel=1e-5)

    def test_feat_dim_constant(self):
        assert FEAT_DIM == H_BINS * S_BINS
        assert FEAT_DIM == 128


# ── cosine_similarity ─────────────────────────────────────────────────────────

class TestCosineSimilarity:

    def test_identical_unit_vectors(self):
        v = _make_feature(0)
        assert cosine_similarity(v, v.copy()) == pytest.approx(1.0, rel=1e-6)

    def test_zero_vector_a(self):
        z = np.zeros(FEAT_DIM, dtype=np.float32)
        v = _make_feature(1)
        assert cosine_similarity(z, v) == pytest.approx(0.0, abs=1e-9)

    def test_zero_vector_b(self):
        v = _make_feature(2)
        z = np.zeros(FEAT_DIM, dtype=np.float32)
        assert cosine_similarity(v, z) == pytest.approx(0.0, abs=1e-9)

    def test_both_zero(self):
        z = np.zeros(FEAT_DIM, dtype=np.float32)
        assert cosine_similarity(z, z) == pytest.approx(0.0, abs=1e-9)

    def test_anti_parallel(self):
        v = _make_feature(3)
        assert cosine_similarity(v, -v) == pytest.approx(-1.0, rel=1e-6)

    def test_orthogonal(self):
        a = np.zeros(FEAT_DIM, dtype=np.float32)
        b = np.zeros(FEAT_DIM, dtype=np.float32)
        a[0] = 1.0
        b[1] = 1.0
        assert cosine_similarity(a, b) == pytest.approx(0.0, abs=1e-9)

    def test_non_unit_same_direction(self):
        v = _make_feature(4)
        assert cosine_similarity(v * 5.0, v * 3.0) == pytest.approx(1.0, rel=1e-6)

    def test_non_unit_opposite_direction(self):
        v = _make_feature(5)
        assert cosine_similarity(v * 7.0, -v * 2.0) == pytest.approx(-1.0, rel=1e-6)

    def test_partial_overlap_in_range(self):
        a = np.zeros(FEAT_DIM, dtype=np.float32)
        b = np.zeros(FEAT_DIM, dtype=np.float32)
        a[0] = 1.0
        b[0] = 1.0
        b[1] = 1.0
        sim = cosine_similarity(a, b)
        assert 0.0 < sim < 1.0


# ── match_track ───────────────────────────────────────────────────────────────

class TestMatchTrack:

    def test_empty_gallery(self):
        feat = _make_feature(10)
        tid, sim = match_track(feat, {}, threshold=0.75)
        assert tid is None
        assert sim == pytest.approx(0.0, abs=1e-9)

    def test_exact_match_above_threshold(self):
        feat   = _make_feature(11)
        tracks = {1: _make_entry(1, feat)}
        tid, sim = match_track(feat, tracks, threshold=0.75)
        assert tid == 1
        assert sim == pytest.approx(1.0, rel=1e-5)

    def test_similarity_at_threshold_no_match(self):
        """Strict >: equal to threshold is not a match."""
        feat = _make_feature(12)
        # construct a feature at exactly cosine = threshold
        # use two orthogonal features blended so sim == threshold
        a = np.zeros(FEAT_DIM, dtype=np.float32)
        b = np.zeros(FEAT_DIM, dtype=np.float32)
        a[0] = 1.0
        b[1] = 1.0
        threshold = 0.75
        # query is a; gallery is a → sim = 1.0, way above. Just test with threshold=1.0+eps
        tracks = {1: _make_entry(1, a)}
        tid, sim = match_track(a, tracks, threshold=1.0 + 1e-9)
        assert tid is None

    def test_below_threshold_no_match(self):
        feat_a = _make_feature(13)
        feat_b = _make_feature(14)
        tracks = {1: _make_entry(1, feat_b)}
        # Force threshold above any realistic similarity between random features
        tid, sim = match_track(feat_a, tracks, threshold=0.9999)
        if cosine_similarity(feat_a, feat_b) <= 0.9999:
            assert tid is None

    def test_stale_track_excluded(self):
        feat   = _make_feature(15)
        tracks = {1: _make_entry(1, feat, is_stale=True)}
        tid, sim = match_track(feat, tracks, threshold=0.5)
        assert tid is None

    def test_multiple_tracks_best_returned(self):
        query  = _make_feature(16)
        close  = query.copy()                    # sim = 1.0
        far    = _make_feature(17)               # lower sim
        tracks = {
            1: _make_entry(1, far),
            2: _make_entry(2, close),
        }
        tid, sim = match_track(query, tracks, threshold=0.5)
        assert tid == 2
        assert sim == pytest.approx(1.0, rel=1e-5)

    def test_all_stale_no_match(self):
        feat   = _make_feature(18)
        tracks = {
            1: _make_entry(1, feat, is_stale=True),
            2: _make_entry(2, feat, is_stale=True),
        }
        tid, sim = match_track(feat, tracks, threshold=0.5)
        assert tid is None

    def test_fresh_and_stale_returns_fresh(self):
        feat   = _make_feature(19)
        tracks = {
            1: _make_entry(1, feat, is_stale=True),
            2: _make_entry(2, feat, is_stale=False),
        }
        tid, sim = match_track(feat, tracks, threshold=0.5)
        assert tid == 2


# ── TrackGallery ──────────────────────────────────────────────────────────────

_BBOX = (10.0, 20.0, 50.0, 80.0)


class TestTrackGallery:

    def test_add_increments_track_id(self):
        g   = TrackGallery()
        now = 1000.0
        t1  = g.add('front', _make_feature(20), _BBOX, now)
        t2  = g.add('front', _make_feature(21), _BBOX, now)
        t3  = g.add('rear',  _make_feature(22), _BBOX, now)
        assert t1 == 1
        assert t2 == 2
        assert t3 == 3

    def test_add_stores_copy(self):
        g   = TrackGallery()
        feat = _make_feature(23)
        orig = feat.copy()
        g.add('front', feat, _BBOX, 1000.0)
        feat[0] = 99.9        # mutate original
        stored  = g.tracks[1].feature
        assert np.allclose(stored, orig, atol=1e-6)

    def test_add_sets_timestamps(self):
        g   = TrackGallery()
        now = 5000.0
        g.add('front', _make_feature(24), _BBOX, now)
        e = g.tracks[1]
        assert e.first_seen == pytest.approx(now)
        assert e.last_seen  == pytest.approx(now)

    def test_add_initial_confidence_one(self):
        g = TrackGallery()
        g.add('front', _make_feature(25), _BBOX, 1000.0)
        assert g.tracks[1].confidence == pytest.approx(1.0)

    def test_update_changes_last_seen(self):
        g    = TrackGallery()
        feat = _make_feature(26)
        g.add('front', feat, _BBOX, 1000.0)
        g.update(1, 'front', feat, _BBOX, 0.9, 2000.0)
        assert g.tracks[1].last_seen  == pytest.approx(2000.0)
        assert g.tracks[1].first_seen == pytest.approx(1000.0)   # preserved

    def test_update_changes_camera_id(self):
        g    = TrackGallery()
        feat = _make_feature(27)
        g.add('front', feat, _BBOX, 1000.0)
        g.update(1, 'right', feat, _BBOX, 0.8, 2000.0)
        assert g.tracks[1].camera_id == 'right'

    def test_update_applies_ema(self):
        g     = TrackGallery()
        feat0 = _make_feature(28)
        g.add('front', feat0, _BBOX, 1000.0)

        feat1 = _make_feature(29)
        g.update(1, 'front', feat1, _BBOX, 0.85, 2000.0)
        stored = g.tracks[1].feature

        # EMA: 0.3*feat1 + 0.7*feat0 (then renorm) — not equal to either
        assert not np.allclose(stored, feat0, atol=1e-4)
        assert not np.allclose(stored, feat1, atol=1e-4)

    def test_update_feature_stays_unit_norm(self):
        g    = TrackGallery()
        feat = _make_feature(30)
        g.add('front', feat, _BBOX, 1000.0)
        for i in range(10):
            new_feat = _make_feature(31 + i)
            g.update(1, 'front', new_feat, _BBOX, 0.9, 1000.0 + i)
        assert np.linalg.norm(g.tracks[1].feature) == pytest.approx(1.0, rel=1e-5)

    def test_update_clears_stale(self):
        g    = TrackGallery()
        feat = _make_feature(41)
        g.add('front', feat, _BBOX, 1000.0)
        g.tracks[1].is_stale = True
        g.update(1, 'front', feat, _BBOX, 0.9, 2000.0)
        assert g.tracks[1].is_stale is False

    def test_mark_stale_recent_track_stays_fresh(self):
        g   = TrackGallery()
        now = 1000.0
        g.add('front', _make_feature(42), _BBOX, now)
        g.mark_stale(now + 10.0, timeout_s=30.0)
        assert g.tracks[1].is_stale is False

    def test_mark_stale_old_track_flagged(self):
        g   = TrackGallery()
        now = 1000.0
        g.add('front', _make_feature(43), _BBOX, now)
        g.mark_stale(now + 31.0, timeout_s=30.0)
        assert g.tracks[1].is_stale is True

    def test_mark_stale_exactly_at_timeout_not_stale(self):
        """Timeout is strictly >, so last_seen + timeout_s is still fresh."""
        g   = TrackGallery()
        now = 1000.0
        g.add('front', _make_feature(44), _BBOX, now)
        g.mark_stale(now + 30.0, timeout_s=30.0)
        assert g.tracks[1].is_stale is False

    def test_prune_stale_removes_flagged(self):
        g   = TrackGallery()
        now = 1000.0
        g.add('front', _make_feature(45), _BBOX, now)   # track 1
        g.add('front', _make_feature(46), _BBOX, now)   # track 2
        g.tracks[1].is_stale = True
        g.prune_stale()
        assert 1 not in g.tracks
        assert 2 in g.tracks

    def test_prune_stale_keeps_fresh(self):
        g   = TrackGallery()
        now = 1000.0
        g.add('front', _make_feature(47), _BBOX, now)
        g.prune_stale()
        assert 1 in g.tracks

    def test_empty_gallery_match(self):
        g = TrackGallery()
        tid, sim = g.match(_make_feature(48), threshold=0.75)
        assert tid is None
        assert sim == pytest.approx(0.0, abs=1e-9)

    def test_match_delegates_to_match_track(self):
        g    = TrackGallery()
        feat = _make_feature(49)
        g.add('front', feat, _BBOX, 1000.0)
        tid, sim = g.match(feat, threshold=0.5)
        assert tid == 1
        assert sim == pytest.approx(1.0, rel=1e-5)


# ── Re-id integration ─────────────────────────────────────────────────────────

class TestReIdIntegration:

    def test_same_appearance_two_cameras_same_track(self):
        """Same feature from two cameras → same track_id."""
        g    = TrackGallery()
        feat = _make_feature(50)
        now  = 1000.0

        # First sighting: front camera creates track 1
        t1 = g.add('front', feat, _BBOX, now)
        assert t1 == 1

        # Second sighting: right camera matches → updates same track
        tid, sim = g.match(feat, threshold=0.75)
        assert tid == 1
        g.update(tid, 'right', feat, _BBOX, float(sim), now + 1.0)
        assert len(g.tracks) == 1

    def test_different_appearance_different_tracks(self):
        """Two orthogonal features must not match → separate track IDs."""
        g   = TrackGallery()
        now = 1000.0

        # Construct two exactly orthogonal unit feature vectors
        feat_a = np.zeros(FEAT_DIM, dtype=np.float32)
        feat_b = np.zeros(FEAT_DIM, dtype=np.float32)
        feat_a[0] = 1.0   # basis vector e0
        feat_b[1] = 1.0   # basis vector e1 — cosine sim = 0.0

        g.add('front', feat_a, _BBOX, now)

        tid, sim = g.match(feat_b, threshold=0.75)
        assert tid is None

        g.add('front', feat_b, _BBOX, now + 0.1)
        assert len(g.tracks) == 2

    def test_stale_prune_readd_new_id(self):
        g    = TrackGallery()
        feat = _make_feature(53)
        now  = 1000.0

        t1 = g.add('front', feat, _BBOX, now)
        assert t1 == 1

        g.mark_stale(now + 31.0, timeout_s=30.0)
        g.prune_stale()
        assert len(g.tracks) == 0

        # Re-add same person gets new ID
        t2 = g.add('front', feat, _BBOX, now + 40.0)
        assert t2 == 2    # next_id incremented past 1

    def test_ema_feature_norm_stays_unit_after_many_updates(self):
        g    = TrackGallery()
        feat = _make_feature(54)
        g.add('front', feat, _BBOX, 0.0)

        for i in range(50):
            new_feat = _make_feature(55 + i)
            g.update(1, 'front', new_feat, _BBOX, 0.9, float(i))

        norm = np.linalg.norm(g.tracks[1].feature)
        assert norm == pytest.approx(1.0, rel=1e-4)

    def test_max_tracks_not_exceeded_in_gallery(self):
        """Gallery itself is unbounded; max_tracks enforcement is in the node.
           Confirm gallery can hold many entries without errors."""
        g = TrackGallery()
        for i in range(25):
            g.add('front', _make_feature(100 + i), _BBOX, float(i))
        assert len(g.tracks) == 25

    def test_hsv_same_image_produces_high_similarity(self):
        """Same-crop HSV histogram should have cosine sim ≈ 1."""
        import cv2
        img  = _solid_bgr(50, 120, 200)
        h1   = extract_hsv_histogram(img)
        h2   = extract_hsv_histogram(img.copy())
        assert cosine_similarity(h1, h2) == pytest.approx(1.0, rel=1e-5)

    def test_hsv_different_colours_low_similarity(self):
        """Red vs cyan BGR should be dissimilar."""
        red  = extract_hsv_histogram(_solid_bgr(0, 0, 200))
        cyan = extract_hsv_histogram(_solid_bgr(200, 200, 0))
        assert cosine_similarity(red, cyan) < 0.5


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
