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

Covers:
  - _l2_norm helper
  - EmbeddingModel (histogram fallback — no model file needed)
  - ReidGallery cosine-similarity matching
"""

import sys
import os
import time

import numpy as np
import pytest

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

from saltybot_person_reid._embedding_model import EmbeddingModel, _l2_norm
from saltybot_person_reid._reid_gallery import ReidGallery


# ── _l2_norm ──────────────────────────────────────────────────────────────────

class TestL2Norm:

    def test_unit_vector_unchanged(self):
        v = np.array([1.0, 0.0, 0.0], dtype=np.float32)
        assert np.allclose(_l2_norm(v), v)

    def test_normalised_to_unit_norm(self):
        v = np.array([3.0, 4.0], dtype=np.float32)
        assert abs(np.linalg.norm(_l2_norm(v)) - 1.0) < 1e-6

    def test_zero_vector_does_not_crash(self):
        v = np.zeros(4, dtype=np.float32)
        result = _l2_norm(v)
        assert result.shape == (4,)


# ── EmbeddingModel ────────────────────────────────────────────────────────────

class TestEmbeddingModel:

    def test_histogram_fallback_shape(self):
        m = EmbeddingModel(model_path=None)
        bgr = np.random.randint(0, 255, (100, 50, 3), dtype=np.uint8)
        emb = m.embed(bgr)
        assert emb.shape == (128,)

    def test_embedding_is_unit_norm(self):
        m = EmbeddingModel(model_path=None)
        bgr = np.random.randint(0, 255, (80, 40, 3), dtype=np.uint8)
        emb = m.embed(bgr)
        assert abs(np.linalg.norm(emb) - 1.0) < 1e-5

    def test_empty_crop_returns_zero_vector(self):
        m = EmbeddingModel(model_path=None)
        emb = m.embed(np.zeros((0, 0, 3), dtype=np.uint8))
        assert emb.shape == (128,)
        assert np.all(emb == 0.0)

    def test_red_and_blue_crops_differ(self):
        m = EmbeddingModel(model_path=None)
        red  = np.full((80, 40, 3), (0,   0, 200), dtype=np.uint8)
        blue = np.full((80, 40, 3), (200, 0,   0), dtype=np.uint8)
        sim = float(np.dot(m.embed(red), m.embed(blue)))
        assert sim < 0.99

    def test_same_crop_deterministic(self):
        m = EmbeddingModel(model_path=None)
        bgr = np.random.randint(0, 255, (80, 40, 3), dtype=np.uint8)
        assert np.allclose(m.embed(bgr), m.embed(bgr))

    def test_embedding_float32(self):
        m = EmbeddingModel(model_path=None)
        bgr = np.random.randint(0, 255, (60, 30, 3), dtype=np.uint8)
        emb = m.embed(bgr)
        assert emb.dtype == np.float32


# ── ReidGallery ───────────────────────────────────────────────────────────────

def _unit(dim: int = 128, seed: int | None = None) -> np.ndarray:
    rng = np.random.default_rng(seed)
    v = rng.standard_normal(dim).astype(np.float32)
    return v / np.linalg.norm(v)


class TestReidGallery:

    def test_first_match_creates_identity(self):
        g = ReidGallery(match_threshold=0.75)
        uid, score, is_new = g.match(_unit(seed=0))
        assert uid == 1
        assert is_new is True
        assert score == pytest.approx(0.0)

    def test_identical_embedding_matches(self):
        g = ReidGallery(match_threshold=0.75)
        emb = _unit(seed=1)
        g.match(emb)
        uid2, score2, is_new2 = g.match(emb)
        assert uid2 == 1
        assert is_new2 is False
        assert score2 > 0.99

    def test_orthogonal_embeddings_create_new_id(self):
        g = ReidGallery(match_threshold=0.75)
        e1 = np.zeros(128, dtype=np.float32); e1[0]  = 1.0
        e2 = np.zeros(128, dtype=np.float32); e2[64] = 1.0
        uid1, _, new1 = g.match(e1)
        uid2, _, new2 = g.match(e2)
        assert uid1 != uid2
        assert new2 is True

    def test_ids_are_monotonically_increasing(self):
        # threshold > 1.0 is unreachable → every embedding creates a new identity
        g = ReidGallery(match_threshold=2.0)
        ids = [g.match(_unit(seed=i))[0] for i in range(5)]
        assert ids == list(range(1, 6))

    def test_gallery_size_increments_for_new_ids(self):
        g = ReidGallery(match_threshold=2.0)
        for i in range(4):
            g.match(_unit(seed=i))
        assert g.size == 4

    def test_prune_removes_stale_identities(self):
        g = ReidGallery(match_threshold=0.75, max_age_s=0.01)
        g.match(_unit(seed=0))
        time.sleep(0.05)
        g._prune()
        assert g.size == 0

    def test_empty_gallery_prune_is_safe(self):
        g = ReidGallery()
        g._prune()
        assert g.size == 0

    def test_match_below_threshold_increments_id(self):
        g = ReidGallery(match_threshold=0.99)
        # Two random unit vectors are almost certainly < 0.99 similar
        e1, e2 = _unit(seed=10), _unit(seed=20)
        uid1, _, _ = g.match(e1)
        uid2, _, _ = g.match(e2)
        # uid2 may or may not equal uid1 depending on random similarity,
        # but both must be valid positive integers
        assert uid1 >= 1
        assert uid2 >= 1

    def test_identity_update_does_not_change_id(self):
        g = ReidGallery(match_threshold=0.5)
        emb = _unit(seed=5)
        uid_first, _, _ = g.match(emb)
        for _ in range(10):
            g.match(emb)
        uid_last, _, _ = g.match(emb)
        assert uid_last == uid_first


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