"""
terrain_analyzer.py — IMU vibration analysis for surface classification (Issue #142).

Algorithm
─────────
  For each control tick, one IMU linear-acceleration triple (ax, ay, az) is
  submitted to a rolling window.  The RMS vibration is computed as:

      vibration_rms = sqrt( Var(ax) + Var(ay) + Var(az) )   [m/s²]

  Variance naturally removes the DC gravity component, so the signal
  captures only dynamic surface vibrations.

  Normalised roughness:
      roughness = min(1.0, vibration_rms / max_variance_ms2)

  Surface classification thresholds (roughness boundaries):
      smooth  [0.00, 0.12)   grip = 0.90
      tile    [0.12, 0.28)   grip = 0.85
      carpet  [0.28, 0.48)   grip = 0.75
      grass   [0.48, 0.68)   grip = 0.55
      gravel  [0.68, 1.00]   grip = 0.45

  Hysteresis prevents rapid class switching: the candidate class must persist
  for `hysteresis_count` consecutive samples before becoming current.

Pure module — no ROS2 dependency.
"""

from __future__ import annotations

import math
from collections import deque
from dataclasses import dataclass
from typing import Optional


# ── Surface classification table ─────────────────────────────────────────────
# (roughness_upper_bound, surface_type, grip_coefficient)
# Ordered from smoothest to roughest; last entry has no upper bound.
_SURFACE_TABLE: list[tuple[float, str, float]] = [
    (0.12, "smooth",  0.90),
    (0.28, "tile",    0.85),
    (0.48, "carpet",  0.75),
    (0.68, "grass",   0.55),
    (1.01, "gravel",  0.45),
]

SURFACE_TYPES: tuple[str, ...] = tuple(e[1] for e in _SURFACE_TABLE)


# ── SurfaceAnalysis ───────────────────────────────────────────────────────────

@dataclass
class SurfaceAnalysis:
    roughness:    float   # 0.0 (smooth) – 1.0 (rough), normalised
    surface_type: str     # "smooth" | "tile" | "carpet" | "grass" | "gravel" | "unknown"
    grip:         float   # 0.0 (no grip) – 1.0 (full grip)


# ── TerrainAnalyzer ───────────────────────────────────────────────────────────

class TerrainAnalyzer:
    """
    Rolling-window vibration analyser with hysteresis-protected classification.

    Parameters
    ----------
    max_variance_ms2    : normalisation ceiling (m/s²); maps to roughness=1.0
    window_s            : rolling window duration (s)
    control_rate_hz     : expected update rate (Hz) — sets window size
    hysteresis_count    : consecutive samples in new class before switching
    """

    def __init__(
        self,
        max_variance_ms2: float = 3.0,
        window_s: float = 0.5,
        control_rate_hz: float = 50.0,
        hysteresis_count: int = 5,
    ):
        self._max_var = max(1e-3, max_variance_ms2)
        window_size = max(5, int(window_s * control_rate_hz))
        self._samples: deque = deque(maxlen=window_size)
        self._hysteresis_count = hysteresis_count

        self._current_type: str = "unknown"
        self._current_grip: float = 0.6
        self._candidate_type: str = "unknown"
        self._candidate_count: int = 0

    # ── Properties ────────────────────────────────────────────────────────────

    @property
    def current_surface_type(self) -> str:
        return self._current_type

    @property
    def current_grip(self) -> float:
        return self._current_grip

    @property
    def sample_count(self) -> int:
        return len(self._samples)

    # ── Update ────────────────────────────────────────────────────────────────

    def update(
        self,
        ax: float,
        ay: float,
        az: float,
    ) -> Optional[SurfaceAnalysis]:
        """
        Submit one IMU acceleration sample (m/s²) and return analysis.

        Returns
        -------
        SurfaceAnalysis when ≥ 5 samples have been collected; None otherwise.
        """
        self._samples.append((ax, ay, az))

        if len(self._samples) < 5:
            return None

        roughness = self._compute_roughness()
        candidate_type, candidate_grip = _classify(roughness)

        # Hysteresis: require sustained presence in new class
        if candidate_type == self._candidate_type:
            self._candidate_count += 1
        else:
            self._candidate_type = candidate_type
            self._candidate_count = 1

        if self._candidate_count >= self._hysteresis_count:
            self._current_type = candidate_type
            self._current_grip = candidate_grip

        return SurfaceAnalysis(
            roughness=roughness,
            surface_type=self._current_type,
            grip=self._current_grip,
        )

    def reset(self) -> None:
        self._samples.clear()
        self._current_type = "unknown"
        self._current_grip = 0.6
        self._candidate_type = "unknown"
        self._candidate_count = 0

    # ── Internal ─────────────────────────────────────────────────────────────

    def _compute_roughness(self) -> float:
        n = len(self._samples)
        if n < 2:
            return 0.0

        ax_vals = [s[0] for s in self._samples]
        ay_vals = [s[1] for s in self._samples]
        az_vals = [s[2] for s in self._samples]

        rms = math.sqrt(_variance(ax_vals) + _variance(ay_vals) + _variance(az_vals))
        return min(1.0, rms / self._max_var)


# ── Module-level helpers ──────────────────────────────────────────────────────

def _variance(vals: list[float]) -> float:
    n = len(vals)
    if n < 2:
        return 0.0
    mean = sum(vals) / n
    return sum((v - mean) ** 2 for v in vals) / n


def _classify(roughness: float) -> tuple[str, float]:
    """Map normalised roughness to (surface_type, grip)."""
    for upper, stype, grip in _SURFACE_TABLE:
        if roughness < upper:
            return stype, grip
    return "gravel", 0.45


def grip_for_surface(surface_type: str) -> float:
    """Return grip coefficient for a given surface type string."""
    for _, stype, grip in _SURFACE_TABLE:
        if stype == surface_type:
            return grip
    return 0.6   # unknown
