"""emotion_classifier.py — 7-class facial expression classifier (Issue #161).

Pure Python, no ROS2 / TensorRT / OpenCV dependencies.
Wraps a TensorRT FP16 emotion CNN and provides:
  - On-device TRT inference on 48×48 grayscale face crops
  - Heuristic fallback from 5-point SCRFD facial landmarks
  - Per-person EMA temporal smoothing for stable outputs
  - Per-person opt-out registry

Emotion classes (index order matches CNN output layer)
------------------------------------------------------
  0 = happy
  1 = sad
  2 = angry
  3 = surprised
  4 = fearful
  5 = disgusted
  6 = neutral

Coordinate convention
---------------------
  Face crop: BGR uint8 ndarray, any size (resized to INPUT_SIZE internally).
  Landmarks (lm10): 10 floats from FaceDetection.landmarks —
    [left_eye_x, left_eye_y, right_eye_x, right_eye_y,
     nose_x, nose_y, left_mouth_x, left_mouth_y,
     right_mouth_x, right_mouth_y]
  All coordinates are normalised image-space (0.0–1.0).
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

# ── Constants ─────────────────────────────────────────────────────────────────

EMOTIONS: List[str] = [
    "happy", "sad", "angry", "surprised", "fearful", "disgusted", "neutral"
]
N_CLASSES: int = len(EMOTIONS)

# Landmark slot indices within the 10-value array
_LEX, _LEY = 0, 1   # left eye
_REX, _REY = 2, 3   # right eye
_NX,  _NY  = 4, 5   # nose
_LMX, _LMY = 6, 7   # left mouth corner
_RMX, _RMY = 8, 9   # right mouth corner

# CNN input size
INPUT_SIZE: int = 48


# ── Result type ───────────────────────────────────────────────────────────────


@dataclass
class ClassifiedEmotion:
    emotion: str            # dominant emotion label
    confidence: float       # smoothed softmax probability (0.0–1.0)
    scores: List[float]     # per-class scores, len == N_CLASSES
    source: str = "cnn_trt" # "cnn_trt" | "landmark_heuristic" | "opt_out"


# ── Math helpers ──────────────────────────────────────────────────────────────


def _softmax(logits: List[float]) -> List[float]:
    """Numerically stable softmax."""
    m = max(logits)
    exps = [math.exp(x - m) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]


def _argmax(values: List[float]) -> int:
    best = 0
    for i in range(1, len(values)):
        if values[i] > values[best]:
            best = i
    return best


def _uniform_scores() -> List[float]:
    return [1.0 / N_CLASSES] * N_CLASSES


# ── Per-person temporal smoother + opt-out registry ───────────────────────────


class PersonEmotionTracker:
    """Exponential moving-average smoother + opt-out registry per person.

    Args:
        alpha:   EMA weight for the newest observation (0.0 = frozen, 1.0 = no smooth).
        max_age: Frames without update before the EMA is expired and reset.
    """

    def __init__(self, alpha: float = 0.3, max_age: int = 16) -> None:
        self._alpha = alpha
        self._max_age = max_age
        self._ema: Dict[int, List[float]] = {}      # person_id → EMA scores
        self._age: Dict[int, int] = {}              # person_id → frames since last update
        self._opt_out: set = set()                  # person_ids that opted out

    # ── Opt-out management ────────────────────────────────────────────────────

    def set_opt_out(self, person_id: int, value: bool) -> None:
        if value:
            self._opt_out.add(person_id)
            self._ema.pop(person_id, None)
            self._age.pop(person_id, None)
        else:
            self._opt_out.discard(person_id)

    def is_opt_out(self, person_id: int) -> bool:
        return person_id in self._opt_out

    # ── Smoothing ─────────────────────────────────────────────────────────────

    def smooth(self, person_id: int, scores: List[float]) -> List[float]:
        """Apply EMA to raw scores.  Returns smoothed scores (length N_CLASSES)."""
        if person_id < 0:
            return scores[:]

        # Reset stale entries
        if person_id in self._age and self._age[person_id] > self._max_age:
            self.reset(person_id)

        if person_id not in self._ema:
            self._ema[person_id] = scores[:]
            self._age[person_id] = 0
            return scores[:]

        ema = self._ema[person_id]
        a = self._alpha
        smoothed = [a * s + (1.0 - a) * e for s, e in zip(scores, ema)]
        self._ema[person_id] = smoothed
        self._age[person_id] = 0
        return smoothed

    def tick(self) -> None:
        """Advance age counter for all tracked persons (call once per frame)."""
        for pid in list(self._age.keys()):
            self._age[pid] += 1

    def reset(self, person_id: int) -> None:
        self._ema.pop(person_id, None)
        self._age.pop(person_id, None)

    def reset_all(self) -> None:
        self._ema.clear()
        self._age.clear()

    @property
    def tracked_ids(self) -> List[int]:
        return list(self._ema.keys())


# ── TensorRT engine loader ────────────────────────────────────────────────────


class _TrtEngine:
    """Thin wrapper around a TensorRT FP16 emotion CNN engine.

    Expected engine:
      - Input binding:  name "input"  shape (1, 1, 48, 48)  float32
      - Output binding: name "output" shape (1, 7)           float32 (softmax logits)

    The engine is built offline from a MobileNetV2-based 48×48 grayscale CNN
    (FER+ or AffectNet trained) via:
        trtexec --onnx=emotion_cnn.onnx --fp16 --saveEngine=emotion_fp16.trt
    """

    def __init__(self) -> None:
        self._context = None
        self._h_input = None
        self._h_output = None
        self._d_input = None
        self._d_output = None
        self._stream = None
        self._bindings: List = []

    def load(self, engine_path: str) -> bool:
        """Load TRT engine from file.  Returns True on success."""
        try:
            import tensorrt as trt              # type: ignore
            import pycuda.autoinit             # type: ignore  # noqa: F401
            import pycuda.driver as cuda       # type: ignore
            import numpy as np

            TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
            with open(engine_path, "rb") as f:
                engine_data = f.read()

            runtime = trt.Runtime(TRT_LOGGER)
            engine = runtime.deserialize_cuda_engine(engine_data)
            self._context = engine.create_execution_context()

            # Allocate host + device buffers
            self._h_input = cuda.pagelocked_empty((1, 1, INPUT_SIZE, INPUT_SIZE),
                                                  dtype=np.float32)
            self._h_output = cuda.pagelocked_empty((1, N_CLASSES), dtype=np.float32)
            self._d_input = cuda.mem_alloc(self._h_input.nbytes)
            self._d_output = cuda.mem_alloc(self._h_output.nbytes)
            self._stream = cuda.Stream()
            self._bindings = [int(self._d_input), int(self._d_output)]
            return True

        except Exception:
            self._context = None
            return False

    @property
    def ready(self) -> bool:
        return self._context is not None

    def infer(self, face_bgr) -> List[float]:
        """Run TRT inference on a BGR face crop.  Returns 7 softmax scores."""
        import pycuda.driver as cuda   # type: ignore
        import numpy as np

        img = _preprocess(face_bgr)    # (1, 1, 48, 48) float32
        np.copyto(self._h_input, img)
        cuda.memcpy_htod_async(self._d_input, self._h_input, self._stream)
        self._context.execute_async_v2(self._bindings, self._stream.handle, None)
        cuda.memcpy_dtoh_async(self._h_output, self._d_output, self._stream)
        self._stream.synchronize()
        logits = list(self._h_output[0])
        return _softmax(logits)


# ── Image pre-processing helper (importable without cv2 in test mode) ─────────


def _preprocess(face_bgr) -> "np.ndarray":  # type: ignore
    """Resize BGR crop to 48×48 grayscale, normalise to [-1, 1].

    Returns ndarray shape (1, 1, 48, 48) float32.
    """
    import numpy as np  # type: ignore
    import cv2          # type: ignore

    gray = cv2.cvtColor(face_bgr, cv2.COLOR_BGR2GRAY)
    resized = cv2.resize(gray, (INPUT_SIZE, INPUT_SIZE),
                         interpolation=cv2.INTER_LINEAR)
    norm = resized.astype(np.float32) / 127.5 - 1.0
    return norm.reshape(1, 1, INPUT_SIZE, INPUT_SIZE)


# ── Landmark heuristic ────────────────────────────────────────────────────────


def classify_from_landmarks(lm10: List[float]) -> ClassifiedEmotion:
    """Estimate emotion from 5-point SCRFD landmarks (10 floats).

    Uses geometric ratios between eyes, nose, and mouth corners.
    Accuracy is limited — treats this as a soft prior, not a definitive label.

    Returns ClassifiedEmotion with source="landmark_heuristic".
    """
    if len(lm10) < 10:
        scores = _uniform_scores()
        scores[6] = 0.5           # bias neutral
        scores = _renorm(scores)
        return ClassifiedEmotion("neutral", scores[6], scores, "landmark_heuristic")

    eye_y = (lm10[_LEY] + lm10[_REY]) / 2.0
    nose_y = lm10[_NY]
    mouth_y = (lm10[_LMY] + lm10[_RMY]) / 2.0
    eye_span = max(abs(lm10[_REX] - lm10[_LEX]), 1e-4)
    mouth_width = abs(lm10[_RMX] - lm10[_LMX])
    mouth_asymm = abs(lm10[_LMY] - lm10[_RMY])

    face_h = max(mouth_y - eye_y, 1e-4)

    # Ratio of mouth span to interocular distance
    width_ratio = mouth_width / eye_span          # >1.0 = wide open / happy smile
    # How far mouth is below nose, relative to face height
    mouth_below_nose = (mouth_y - nose_y) / face_h  # ~0.3–0.6 typical
    # Relative asymmetry of mouth corners
    asym_ratio = mouth_asymm / face_h             # >0.05 = notable asymmetry

    # Build soft scores
    scores = _uniform_scores()  # start uniform

    if width_ratio > 0.85 and mouth_below_nose > 0.35:
        # Wide mouth, normal vertical position → happy
        scores[0] = 0.55 + 0.25 * min(1.0, (width_ratio - 0.85) / 0.5)
    elif mouth_below_nose < 0.20 and width_ratio < 0.7:
        # Tight, compressed mouth high up → surprised OR angry
        scores[3] = 0.35   # surprised
        scores[2] = 0.30   # angry
    elif asym_ratio > 0.06:
        # Asymmetric mouth → disgust or sadness
        scores[5] = 0.30   # disgusted
        scores[1] = 0.25   # sad
    elif width_ratio < 0.65 and mouth_below_nose < 0.30:
        # Tight and compressed → sad/angry
        scores[1] = 0.35   # sad
        scores[2] = 0.25   # angry
    else:
        # Default to neutral
        scores[6] = 0.45

    scores = _renorm(scores)
    top_idx = _argmax(scores)
    return ClassifiedEmotion(
        emotion=EMOTIONS[top_idx],
        confidence=round(scores[top_idx], 3),
        scores=[round(s, 4) for s in scores],
        source="landmark_heuristic",
    )


def _renorm(scores: List[float]) -> List[float]:
    """Re-normalise scores so they sum to 1.0."""
    total = sum(scores)
    if total <= 0:
        return _uniform_scores()
    return [s / total for s in scores]


# ── Public classifier ─────────────────────────────────────────────────────────


class EmotionClassifier:
    """Facade combining TRT inference + landmark fallback + per-person smoothing.

    Usage
    -----
    >>> clf = EmotionClassifier(engine_path="/models/emotion_fp16.trt")
    >>> clf.load()
    True
    >>> result = clf.classify_crop(face_bgr, person_id=42, tracker=tracker)
    >>> result.emotion, result.confidence
    ('happy', 0.87)
    """

    def __init__(
        self,
        engine_path: str = "",
        alpha: float = 0.3,
    ) -> None:
        self._engine_path = engine_path
        self._alpha = alpha
        self._engine = _TrtEngine()

    def load(self) -> bool:
        """Load TRT engine.  Returns True if engine is ready."""
        if not self._engine_path:
            return False
        return self._engine.load(self._engine_path)

    @property
    def ready(self) -> bool:
        return self._engine.ready

    def classify_crop(
        self,
        face_bgr,
        person_id: int = -1,
        tracker: Optional[PersonEmotionTracker] = None,
    ) -> ClassifiedEmotion:
        """Classify a BGR face crop.

        Args:
            face_bgr:  BGR ndarray from cv_bridge / direct crop.
            person_id: For temporal smoothing.  -1 = no smoothing.
            tracker:   Optional PersonEmotionTracker for EMA smoothing.

        Returns:
            ClassifiedEmotion with smoothed scores.
        """
        if self._engine.ready:
            raw = self._engine.infer(face_bgr)
            source = "cnn_trt"
        else:
            # Uniform fallback — no inference without engine
            raw = _uniform_scores()
            raw[6] = 0.40          # mild neutral bias
            raw = _renorm(raw)
            source = "landmark_heuristic"

        smoothed = raw
        if tracker is not None and person_id >= 0:
            smoothed = tracker.smooth(person_id, raw)

        top_idx = _argmax(smoothed)
        return ClassifiedEmotion(
            emotion=EMOTIONS[top_idx],
            confidence=round(smoothed[top_idx], 3),
            scores=[round(s, 4) for s in smoothed],
            source=source,
        )

    def classify_from_landmarks(
        self,
        lm10: List[float],
        person_id: int = -1,
        tracker: Optional[PersonEmotionTracker] = None,
    ) -> ClassifiedEmotion:
        """Classify using landmark geometry only (no crop required)."""
        result = classify_from_landmarks(lm10)
        if tracker is not None and person_id >= 0:
            smoothed = tracker.smooth(person_id, result.scores)
            top_idx = _argmax(smoothed)
            result = ClassifiedEmotion(
                emotion=EMOTIONS[top_idx],
                confidence=round(smoothed[top_idx], 3),
                scores=[round(s, 4) for s in smoothed],
                source=result.source,
            )
        return result
