"""
_embedding_model.py — Appearance embedding extractor (no ROS2 deps).

Primary:  MobileNetV2 ONNX torso crop → 128-dim L2-normalised embedding.
Fallback: 128-bin HSV histogram (H:16 × S:8) when no model file is available.
"""

from __future__ import annotations

import numpy as np
import cv2

# Top fraction of the bounding box height used as torso crop
_INPUT_SIZE = (128, 256)   # (W, H) fed to MobileNetV2


class EmbeddingModel:
    """
    Extract a 128-dim L2-normalised appearance embedding from a BGR crop.

    Parameters
    ----------
    model_path : str or None
        Path to a MobileNetV2+projection ONNX file.  When None (or file
        not found), falls back to a 128-bin HSV colour histogram.
    """

    def __init__(self, model_path: str | None = None):
        self._net = None
        if model_path:
            try:
                self._net = cv2.dnn.readNetFromONNX(model_path)
            except Exception:
                pass  # histogram fallback

    def embed(self, bgr_crop: np.ndarray) -> np.ndarray:
        """
        Parameters
        ----------
        bgr_crop : np.ndarray  shape (H, W, 3) uint8

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

        if self._net is not None:
            return self._mobilenet_embed(bgr_crop)
        return self._histogram_embed(bgr_crop)

    # ── MobileNetV2 path ──────────────────────────────────────────────────────

    def _mobilenet_embed(self, bgr: np.ndarray) -> np.ndarray:
        resized = cv2.resize(bgr, _INPUT_SIZE)
        blob = cv2.dnn.blobFromImage(
            resized,
            scalefactor=1.0 / 255.0,
            size=_INPUT_SIZE,
            mean=(0.485 * 255, 0.456 * 255, 0.406 * 255),
            swapRB=True,
            crop=False,
        )
        # Std normalisation: divide channel-wise
        blob[:, 0] /= 0.229
        blob[:, 1] /= 0.224
        blob[:, 2] /= 0.225

        self._net.setInput(blob)
        feat = self._net.forward().flatten().astype(np.float32)

        # Ensure 128-dim — average-pool if model output differs
        if feat.shape[0] != 128:
            n = feat.shape[0]
            block = max(1, n // 128)
            feat = feat[: block * 128].reshape(128, block).mean(axis=1)

        return _l2_norm(feat)

    # ── HSV histogram fallback ────────────────────────────────────────────────

    def _histogram_embed(self, bgr: np.ndarray) -> np.ndarray:
        """128-bin HSV histogram: 16 H-bins × 8 S-bins, concatenated."""
        hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
        hist = cv2.calcHist(
            [hsv], [0, 1], None,
            [16, 8], [0, 180, 0, 256],
        ).flatten().astype(np.float32)
        return _l2_norm(hist)


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