"""
_person_reid.py — Appearance-based person re-identification (no ROS2 deps).

Algorithm
---------
For each person detection (bounding-box crop from a camera image):

1. Extract a 2-D HS colour histogram:
     H axis: 16 bins over [0, 180]   (OpenCV hue)
     S axis:  8 bins over [0, 256]   (saturation)
   The Value channel is omitted for illumination robustness.
   Result: 128-dimensional float32 vector, L2-normalised.

2. Match against the gallery of active TrackEntry objects using cosine
   similarity (= dot product of L2-normalised vectors).  The best match
   whose similarity exceeds `threshold` is selected; otherwise a new track
   is created.

3. On a successful match the gallery entry is updated with an exponential
   moving average (α = 0.3) of the new feature to adapt gradually to
   appearance changes (lighting, pose) while staying robust to outliers.

4. Tracks not updated for `timeout_s` wall-clock seconds are flagged
   `is_stale`; stale tracks are pruned from the gallery on the next cycle.

Public API
----------
  H_BINS, S_BINS, FEAT_DIM             histogram shape constants
  TrackEntry                           mutable dataclass for one active track
  extract_hsv_histogram(bgr_crop)   -> np.ndarray (FEAT_DIM,) L2-normalised
  cosine_similarity(a, b)           -> float in [-1, 1]
  match_track(feature, tracks, threshold)
                                    -> (track_id, similarity) | (None, 0.0)
  TrackGallery                         mutable store of active tracks
    .add(camera_id, feature, bbox, now)   -> int  (new track_id)
    .update(track_id, camera_id, feature, bbox, confidence, now)
    .match(feature, threshold)        -> (track_id, similarity) | (None, 0.0)
    .mark_stale(now, timeout_s)
    .prune_stale()
    .tracks                           -> Dict[int, TrackEntry]
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Dict, Optional, Tuple

import cv2
import numpy as np


# ── Histogram constants ───────────────────────────────────────────────────────

H_BINS:   int = 16
S_BINS:   int = 8
FEAT_DIM: int = H_BINS * S_BINS   # 128

_EMA_ALPHA: float = 0.3   # weight of newest observation in EMA update


# ── Track entry ───────────────────────────────────────────────────────────────

@dataclass
class TrackEntry:
    """Mutable state for one active person track."""

    track_id:   int
    camera_id:  str
    feature:    np.ndarray              # (FEAT_DIM,) L2-normalised appearance
    confidence: float                   # last match similarity (1.0 for new)
    first_seen: float                   # time.time() wall-clock seconds
    last_seen:  float                   # time.time() wall-clock seconds
    bbox:       Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
    # (centre_x, centre_y, size_x, size_y) in image pixels
    is_stale:   bool = False


# ── Feature extraction ────────────────────────────────────────────────────────

def extract_hsv_histogram(bgr_crop: np.ndarray) -> np.ndarray:
    """
    Extract a 2-D HS colour histogram from a BGR image crop.

    The Value channel is excluded for illumination robustness.
    The result is L2-normalised; a zero-area or empty crop returns a zero vector.

    Parameters
    ----------
    bgr_crop : uint8 BGR image, any size ≥ 1×1.

    Returns
    -------
    np.ndarray, shape (FEAT_DIM,) = (128,), dtype float32, L2-normalised.
    """
    if bgr_crop is None or bgr_crop.size == 0:
        return np.zeros(FEAT_DIM, dtype=np.float32)

    hsv  = cv2.cvtColor(bgr_crop, cv2.COLOR_BGR2HSV)
    hist = cv2.calcHist(
        [hsv], [0, 1], None,
        [H_BINS, S_BINS],
        [0, 180, 0, 256],
    )
    flat = hist.flatten().astype(np.float32)
    norm = np.linalg.norm(flat)
    if norm > 0.0:
        flat /= norm
    return flat


# ── Similarity ────────────────────────────────────────────────────────────────

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """
    Cosine similarity between two vectors.

    Works correctly for both unit-normalised and unnormalised inputs.
    Returns 0.0 when either vector has (near-)zero norm.

    Returns
    -------
    float in [-1, 1].
    """
    na = float(np.linalg.norm(a))
    nb = float(np.linalg.norm(b))
    if na < 1e-12 or nb < 1e-12:
        return 0.0
    return float(np.dot(a, b) / (na * nb))


# ── Gallery matching ──────────────────────────────────────────────────────────

def match_track(
    feature:   np.ndarray,
    tracks:    Dict[int, TrackEntry],
    threshold: float,
) -> Tuple[Optional[int], float]:
    """
    Find the gallery entry with the highest cosine similarity to `feature`.

    Stale tracks are excluded from matching.

    Parameters
    ----------
    feature   : query appearance vector (FEAT_DIM,)
    tracks    : gallery dict keyed by track_id
    threshold : minimum cosine similarity required to accept a match

    Returns
    -------
    (track_id, similarity) if a match is found, else (None, 0.0).
    """
    best_id:  Optional[int] = None
    best_sim: float = threshold  # strict: new track if sim == threshold exactly

    for tid, entry in tracks.items():
        if entry.is_stale:
            continue
        sim = cosine_similarity(feature, entry.feature)
        if sim > best_sim:
            best_sim = sim
            best_id  = tid

    if best_id is not None:
        return best_id, best_sim
    return None, 0.0


# ── Gallery ───────────────────────────────────────────────────────────────────

class TrackGallery:
    """
    Mutable store of active person tracks.

    Thread-safety: not thread-safe; call from a single ROS2 callback thread.
    """

    def __init__(self) -> None:
        self._tracks:  Dict[int, TrackEntry] = {}
        self._next_id: int = 1

    @property
    def tracks(self) -> Dict[int, TrackEntry]:
        return self._tracks

    def add(
        self,
        camera_id: str,
        feature:   np.ndarray,
        bbox:      Tuple[float, float, float, float],
        now:       float,
    ) -> int:
        """Create a new track and return its assigned track_id."""
        tid = self._next_id
        self._next_id += 1
        self._tracks[tid] = TrackEntry(
            track_id=tid,
            camera_id=camera_id,
            feature=feature.copy(),
            confidence=1.0,
            first_seen=now,
            last_seen=now,
            bbox=bbox,
        )
        return tid

    def update(
        self,
        track_id:   int,
        camera_id:  str,
        feature:    np.ndarray,
        bbox:       Tuple[float, float, float, float],
        confidence: float,
        now:        float,
    ) -> None:
        """
        Update an existing track with a new observation.

        The appearance feature is blended via EMA (α=0.3) for robustness.
        """
        entry = self._tracks[track_id]

        # EMA blend: α * new + (1-α) * old, then re-normalise
        updated = _EMA_ALPHA * feature + (1.0 - _EMA_ALPHA) * entry.feature
        norm = np.linalg.norm(updated)
        if norm > 0.0:
            updated /= norm

        self._tracks[track_id] = TrackEntry(
            track_id=track_id,
            camera_id=camera_id,
            feature=updated,
            confidence=confidence,
            first_seen=entry.first_seen,
            last_seen=now,
            bbox=bbox,
            is_stale=False,
        )

    def match(
        self,
        feature:   np.ndarray,
        threshold: float,
    ) -> Tuple[Optional[int], float]:
        """Delegate to module-level match_track."""
        return match_track(feature, self._tracks, threshold)

    def mark_stale(self, now: float, timeout_s: float) -> None:
        """Flag tracks whose last_seen age exceeds timeout_s."""
        for entry in self._tracks.values():
            if (now - entry.last_seen) > timeout_s:
                entry.is_stale = True

    def prune_stale(self) -> None:
        """Remove all stale tracks from the gallery."""
        self._tracks = {
            tid: e for tid, e in self._tracks.items() if not e.is_stale
        }
