"""
pid_controller.py — PID controller, gain scheduler, and instability detector
for the SaltyBot adaptive balance controller (Issue #136).

Pure module — no ROS2 dependency; fully unit-testable.

Classes
───────
  GainSet             — immutable (kp, ki, kd) triple with safety-clamp helper
  PIDController       — discrete PID with anti-windup and output clamping
  GainScheduler       — weight-class → gain table, exponential blending,
                        safety bounds, manual override, default revert
  InstabilityDetector — monitors tilt for threshold violations and oscillations
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass


# ── GainSet ──────────────────────────────────────────────────────────────────

@dataclass
class GainSet:
    kp: float
    ki: float
    kd: float

    def clamp(
        self,
        kp_min: float, kp_max: float,
        ki_min: float, ki_max: float,
        kd_min: float, kd_max: float,
    ) -> "GainSet":
        """Return a new GainSet with all gains clamped to the given bounds."""
        return GainSet(
            kp=max(kp_min, min(kp_max, self.kp)),
            ki=max(ki_min, min(ki_max, self.ki)),
            kd=max(kd_min, min(kd_max, self.kd)),
        )


# ── PIDController ─────────────────────────────────────────────────────────────

class PIDController:
    """
    Discrete PID controller with:
      - Anti-windup: integral clamped to ±i_clamp
      - Output clamping: result clamped to [output_min, output_max]
      - Derivative on error (D-on-error, not D-on-measurement)
      - Zero derivative on first step (no derivative kick)
    """

    def __init__(
        self,
        kp: float,
        ki: float,
        kd: float,
        output_min: float = -1e9,
        output_max: float = 1e9,
        i_clamp: float = 1e9,
    ):
        self._kp = kp
        self._ki = ki
        self._kd = kd
        self._output_min = output_min
        self._output_max = output_max
        self._i_clamp = abs(i_clamp)

        self._integral: float = 0.0
        self._prev_error: float | None = None

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

    @property
    def kp(self) -> float:
        return self._kp

    @property
    def ki(self) -> float:
        return self._ki

    @property
    def kd(self) -> float:
        return self._kd

    @property
    def integral(self) -> float:
        return self._integral

    # ── Mutators ──────────────────────────────────────────────────────────────

    def update_gains(self, kp: float, ki: float, kd: float) -> None:
        """Update gains without resetting integrator or derivative state."""
        self._kp = kp
        self._ki = ki
        self._kd = kd

    def reset(self) -> None:
        """Clear integrator and derivative memory (e.g., after mode change)."""
        self._integral = 0.0
        self._prev_error = None

    # ── Step ──────────────────────────────────────────────────────────────────

    def step(self, error: float, dt: float) -> float:
        """
        Compute one PID step.

        Parameters
        ----------
        error : setpoint − measurement
        dt    : elapsed time since last step (s); must be > 0

        Returns
        -------
        Control output clamped to [output_min, output_max].
        Returns 0.0 if dt ≤ 0.
        """
        if dt <= 0.0:
            return 0.0

        # Proportional
        p = self._kp * error

        # Integral with anti-windup
        self._integral += error * dt
        self._integral = max(-self._i_clamp, min(self._i_clamp, self._integral))
        i = self._ki * self._integral

        # Derivative (zero on first call to avoid derivative kick)
        if self._prev_error is None:
            d = 0.0
        else:
            d = self._kd * (error - self._prev_error) / dt
        self._prev_error = error

        output = p + i + d
        return max(self._output_min, min(self._output_max, output))


# ── GainScheduler ────────────────────────────────────────────────────────────

class GainScheduler:
    """
    Maps estimated robot weight to PID gains via a 3-class lookup table.
    Transitions between classes are blended exponentially so the active gains
    track the target smoothly.

    Weight classes (ascending):
        "empty" — up to weight_thresholds["empty_max"]  kg
        "light" — up to weight_thresholds["light_max"]  kg
        "heavy" — above light_max

    Safety clamping is applied to every blended gain before returning.
    """

    CLASSES = ("empty", "light", "heavy")

    def __init__(
        self,
        gain_table: dict[str, GainSet],
        weight_thresholds: dict[str, float],
        blend_alpha: float = 0.05,
        kp_min: float = 0.0,  kp_max: float = 100.0,
        ki_min: float = 0.0,  ki_max: float = 20.0,
        kd_min: float = 0.0,  kd_max: float = 20.0,
    ):
        """
        Parameters
        ----------
        gain_table         : {"empty": GainSet, "light": GainSet, "heavy": GainSet}
        weight_thresholds  : {"empty_max": float, "light_max": float}
        blend_alpha        : convergence rate ∈ (0, 1] per control tick
        """
        self._table = gain_table
        self._thresholds = weight_thresholds
        self._alpha = max(1e-6, min(1.0, blend_alpha))
        self._bounds = {
            "kp_min": kp_min, "kp_max": kp_max,
            "ki_min": ki_min, "ki_max": ki_max,
            "kd_min": kd_min, "kd_max": kd_max,
        }

        # Active (blended) gains start at "empty" table values
        self._active_class: str = "empty"
        self._override_active: bool = False
        raw = self._table["empty"]
        self._active = self._apply_bounds(GainSet(raw.kp, raw.ki, raw.kd))

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

    @property
    def active_class(self) -> str:
        return self._active_class

    @property
    def active_gains(self) -> GainSet:
        return self._active

    @property
    def override_active(self) -> bool:
        return self._override_active

    # ── Classification ────────────────────────────────────────────────────────

    def classify(self, weight_kg: float) -> str:
        """Map weight (kg) to class name."""
        if weight_kg <= self._thresholds.get("empty_max", 0.5):
            return "empty"
        if weight_kg <= self._thresholds.get("light_max", 2.0):
            return "light"
        return "heavy"

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

    def update(self, weight_kg: float) -> tuple[GainSet, str, bool]:
        """
        Blend active gains toward the target for the given weight.
        Manual override bypasses scheduling.

        Returns
        -------
        (active_gains, active_class, class_changed)
        """
        if self._override_active:
            return self._active, self._active_class, False

        target_class = self.classify(weight_kg)
        class_changed = target_class != self._active_class
        self._active_class = target_class

        target = self._table[target_class]
        a = self._alpha
        blended = GainSet(
            kp=(1.0 - a) * self._active.kp + a * target.kp,
            ki=(1.0 - a) * self._active.ki + a * target.ki,
            kd=(1.0 - a) * self._active.kd + a * target.kd,
        )
        self._active = self._apply_bounds(blended)
        return self._active, self._active_class, class_changed

    # ── Override / revert ────────────────────────────────────────────────────

    def override_gains(self, kp: float, ki: float, kd: float) -> None:
        """
        Apply manual gain override (dynamic reconfigure / operator tuning).
        Gains are clamped to safety bounds.
        """
        self._override_active = True
        self._active = self._apply_bounds(GainSet(kp, ki, kd))

    def clear_override(self) -> None:
        """Release manual override; scheduling resumes on next update()."""
        self._override_active = False

    def revert_to_defaults(self, weight_kg: float) -> None:
        """
        Immediately snap active gains to table values for the current weight
        class (no blending). Used after instability is detected.
        """
        self._override_active = False
        cls = self.classify(weight_kg)
        self._active_class = cls
        raw = self._table[cls]
        self._active = self._apply_bounds(GainSet(raw.kp, raw.ki, raw.kd))

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

    def _apply_bounds(self, gains: GainSet) -> GainSet:
        b = self._bounds
        return gains.clamp(
            b["kp_min"], b["kp_max"],
            b["ki_min"], b["ki_max"],
            b["kd_min"], b["kd_max"],
        )


# ── InstabilityDetector ───────────────────────────────────────────────────────

class InstabilityDetector:
    """
    Detects controller instability via two independent criteria:

      1. Tilt threshold  — more than 50 % of recent samples exceed
                           tilt_threshold_rad  →  unstable
      2. Oscillations    — sign-reversal count in window exceeds
                           oscillation_count  →  unstable (growing oscillations)
    """

    def __init__(
        self,
        tilt_threshold_rad: float = 0.4,
        oscillation_count: int = 6,
        window_size: int = 100,
    ):
        self._tilt_threshold = tilt_threshold_rad
        self._osc_count = oscillation_count
        self._tilts: deque = deque(maxlen=window_size)

    def update(self, tilt_rad: float) -> bool:
        """
        Record tilt sample and return True if instability is detected.
        False is always returned until the window has at least 4 samples.
        """
        self._tilts.append(tilt_rad)

        if len(self._tilts) < 4:
            return False

        tilts = list(self._tilts)

        # Criterion 1: tilt amplitude
        exceed = sum(1 for t in tilts if abs(t) > self._tilt_threshold)
        if exceed > len(tilts) * 0.5:
            return True

        # Criterion 2: sign reversals (oscillation detection)
        reversals = 0
        prev_sign: int | None = None
        for t in tilts:
            if abs(t) < 0.01:
                continue
            sign = 1 if t > 0 else -1
            if prev_sign is not None and sign != prev_sign:
                reversals += 1
            prev_sign = sign

        return reversals >= self._osc_count

    def reset(self) -> None:
        self._tilts.clear()
