"""
_reid_gallery.py — Appearance gallery for person re-identification (no ROS2 deps).

Matches an incoming embedding against stored identities using cosine similarity.
New identities are created when the best match falls below the threshold.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import List, Tuple

import numpy as np


@dataclass
class Identity:
    identity_id: int
    embedding: np.ndarray          # shape (D,) L2-normalised
    last_seen: float = field(default_factory=time.monotonic)
    hit_count: int = 1

    def update(self, new_embedding: np.ndarray, alpha: float = 0.1) -> None:
        """EMA update of the stored embedding, re-normalised after blending."""
        merged = (1.0 - alpha) * self.embedding + alpha * new_embedding
        n = float(np.linalg.norm(merged))
        self.embedding = merged / n if n > 1e-6 else merged
        self.last_seen = time.monotonic()
        self.hit_count += 1


class ReidGallery:
    """
    Lightweight cosine-similarity re-ID gallery.

    Parameters
    ----------
    match_threshold : float
        Cosine similarity (dot product of unit vectors) required to accept a
        match.  Range [0, 1]; 0 = always new identity, 1 = perfect match only.
    max_age_s : float
        Identities not seen for this many seconds are pruned.
    """

    def __init__(
        self,
        match_threshold: float = 0.75,
        max_age_s: float = 300.0,
    ):
        self._threshold = match_threshold
        self._max_age_s = max_age_s
        self._identities: List[Identity] = []
        self._next_id = 1

    def match(self, embedding: np.ndarray) -> Tuple[int, float, bool]:
        """
        Match embedding against the gallery.

        Returns
        -------
        (identity_id, match_score, is_new)
            identity_id : assigned ID (new or existing)
            match_score : cosine similarity to best match (0.0 if new)
            is_new      : True if a new identity was created
        """
        self._prune()

        if not self._identities:
            return self._add_identity(embedding)

        scores = np.array(
            [float(np.dot(embedding, ident.embedding)) for ident in self._identities]
        )
        best_idx = int(np.argmax(scores))
        best_score = float(scores[best_idx])

        if best_score >= self._threshold:
            ident = self._identities[best_idx]
            ident.update(embedding)
            return ident.identity_id, best_score, False

        return self._add_identity(embedding)

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _add_identity(self, embedding: np.ndarray) -> Tuple[int, float, bool]:
        new_id = self._next_id
        self._next_id += 1
        self._identities.append(
            Identity(identity_id=new_id, embedding=embedding.copy())
        )
        return new_id, 0.0, True

    def _prune(self) -> None:
        now = time.monotonic()
        self._identities = [
            ident
            for ident in self._identities
            if now - ident.last_seen < self._max_age_s
        ]

    @property
    def size(self) -> int:
        return len(self._identities)
