"""
_floor_classifier.py — Floor surface type classifier (no ROS2 deps).

Classifies floor patches from D435i RGB frames into one of six categories:
  carpet  · tile  · wood  · concrete  · grass  · gravel

Algorithm
---------
1. Crop the bottom `roi_frac` of the image (the visible floor region).
2. Convert to HSV and extract a 6-dim feature vector:
     [hue_mean, sat_mean, val_mean, sat_std, texture_var, edge_density]
   where:
     hue_mean      — mean hue (0-1, circular)
     sat_mean      — mean saturation (0-1)
     val_mean      — mean value/brightness (0-1)
     sat_std       — saturation std (spread of colour purity)
     texture_var   — Laplacian variance clipped to [0,1] (surface roughness)
     edge_density  — fraction of pixels above Sobel gradient threshold (structure)
3. Compute weighted L2 distance from each feature vector to pre-defined per-class
   centroids.  Return the nearest class plus a softmax-based confidence score.

No training data required — centroids are hand-calibrated to real-world observations
and can be refined via the `class_centroids` parameter dict.

Public API
----------
  extract_features(bgr, roi_frac=0.4) → np.ndarray (6,)
  classify_floor_patch(bgr, ...) → ClassifyResult
  LabelSmoother — majority-vote temporal smoother

ClassifyResult
--------------
  label      : str     — floor surface class
  confidence : float   — 0–1 (1 = no competing class)
  features   : ndarray — (6,) raw features (for debugging)
  distances  : dict    — {label: L2 distance} to all class centroids
"""

from __future__ import annotations

import math
from collections import deque, Counter
from typing import Dict, List, NamedTuple, Optional

import numpy as np


# ── Default class centroids (6 features each) ─────────────────────────────────
#
# Feature order:  [hue_mean, sat_mean, val_mean, sat_std, texture_var, edge_density]
# Tuned for typical indoor/outdoor floor surfaces under D435i colour stream.
#
_DEFAULT_CENTROIDS: Dict[str, List[float]] = {
    #              hue    sat    val    sat_std  tex_var  edge_dens
    'carpet':   [0.05,  0.30,  0.45,  0.08,   0.03,   0.08],
    'tile':     [0.08,  0.08,  0.65,  0.04,   0.06,   0.35],
    'wood':     [0.08,  0.45,  0.50,  0.09,   0.05,   0.20],
    'concrete': [0.06,  0.05,  0.55,  0.02,   0.02,   0.08],
    'grass':    [0.33,  0.60,  0.45,  0.12,   0.07,   0.28],
    'gravel':   [0.07,  0.18,  0.42,  0.08,   0.09,   0.42],
}

# Per-feature weights: amplify dimensions whose natural range is smaller so that
# all features contribute comparably to the L2 distance.
_FEATURE_WEIGHTS = np.array([3.0, 2.0, 1.0, 5.0, 8.0, 2.0], dtype=np.float64)

# Laplacian normalisation reference: variance this large → texture_var = 1.0
_LAP_REF = 500.0


class ClassifyResult(NamedTuple):
    label:      str
    confidence: float          # 0–1
    features:   np.ndarray     # (6,) float64
    distances:  Dict[str, float]


def extract_features(bgr: np.ndarray, roi_frac: float = 0.40) -> np.ndarray:
    """
    Extract a 6-dim feature vector from the floor ROI of a BGR image.

    Parameters
    ----------
    bgr      : (H, W, 3) uint8 BGR image
    roi_frac : fraction of the image height to use as floor ROI (bottom)

    Returns
    -------
    (6,) float64 array: [hue_mean, sat_mean, val_mean, sat_std, texture_var, edge_density]
    """
    import cv2

    h = bgr.shape[0]
    roi_start = max(0, int(h * (1.0 - roi_frac)))
    roi = bgr[roi_start:, :, :]

    # ── HSV colour features ───────────────────────────────────────────────────
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV).astype(np.float64)
    hue = hsv[:, :, 0] / 179.0   # cv2 hue: 0-179 → normalise to 0-1
    sat = hsv[:, :, 1] / 255.0
    val = hsv[:, :, 2] / 255.0

    hue_mean = _circular_mean(hue)
    sat_mean = float(sat.mean())
    val_mean = float(val.mean())
    sat_std  = float(sat.std())

    # ── Texture: Laplacian variance ───────────────────────────────────────────
    grey    = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    lap     = cv2.Laplacian(grey, cv2.CV_64F)
    lap_var = float(lap.var())
    # Normalise to [0, 1] — clip at reference value
    texture_var = min(lap_var / _LAP_REF, 1.0)

    # ── Edges: fraction of pixels with strong Sobel gradient ─────────────────
    sx      = cv2.Sobel(grey, cv2.CV_64F, 1, 0, ksize=3)
    sy      = cv2.Sobel(grey, cv2.CV_64F, 0, 1, ksize=3)
    mag     = np.hypot(sx, sy)
    edge_density = float((mag > 30.0).mean())

    return np.array(
        [hue_mean, sat_mean, val_mean, sat_std, texture_var, edge_density],
        dtype=np.float64,
    )


def classify_floor_patch(
    bgr:             np.ndarray,
    roi_frac:        float = 0.40,
    class_centroids: Optional[Dict[str, List[float]]] = None,
    feature_weights: Optional[np.ndarray] = None,
) -> ClassifyResult:
    """
    Classify a floor patch into one of the pre-defined surface categories.

    Parameters
    ----------
    bgr             : (H, W, 3) uint8 BGR image
    roi_frac        : floor ROI fraction (bottom of image)
    class_centroids : override default centroid dict
    feature_weights : (6,) weights applied before L2 distance

    Returns
    -------
    ClassifyResult(label, confidence, features, distances)
    """
    centroids = class_centroids if class_centroids is not None else _DEFAULT_CENTROIDS
    weights   = feature_weights if feature_weights is not None else _FEATURE_WEIGHTS

    feats = extract_features(bgr, roi_frac=roi_frac)
    w_feats = feats * weights

    distances: Dict[str, float] = {}
    for label, centroid in centroids.items():
        w_centroid = np.asarray(centroid, dtype=np.float64) * weights
        distances[label] = float(np.linalg.norm(w_feats - w_centroid))

    best_label = min(distances, key=lambda k: distances[k])

    # Confidence: softmax over negative distances
    d_vals = np.array(list(distances.values()))
    softmax = np.exp(-d_vals) / np.exp(-d_vals).sum()
    best_idx = list(distances.keys()).index(best_label)
    confidence = float(softmax[best_idx])

    return ClassifyResult(
        label=best_label,
        confidence=confidence,
        features=feats,
        distances=distances,
    )


# ── Temporal smoother ─────────────────────────────────────────────────────────

class LabelSmoother:
    """
    Majority-vote smoother over the last N classification results.

    Usage:
        smoother = LabelSmoother(window=5)
        label = smoother.update('carpet')   # → smoothed label
    """

    def __init__(self, window: int = 5) -> None:
        self._window = window
        self._history: deque = deque(maxlen=window)

    def update(self, label: str) -> str:
        self._history.append(label)
        return Counter(self._history).most_common(1)[0][0]

    def clear(self) -> None:
        self._history.clear()

    @property
    def ready(self) -> bool:
        """True once the window is full."""
        return len(self._history) == self._window


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

def _circular_mean(hue_01: np.ndarray) -> float:
    """Circular mean of hue values in [0, 1]."""
    angles = hue_01 * 2.0 * math.pi
    sin_mean = float(np.sin(angles).mean())
    cos_mean = float(np.cos(angles).mean())
    angle = math.atan2(sin_mean, cos_mean)
    return (angle / (2.0 * math.pi)) % 1.0
