"""
weight_estimator.py — Motor-current weight estimator and calibration routine
for the SaltyBot adaptive balance controller (Issue #136).

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

Classes
───────
  WeightEstimator    — rolling-window current → weight, change detection
  CalibrationRoutine — state-machine: IDLE → ROCKING → SETTLING → DONE/FAILED
"""

from __future__ import annotations

import math
from collections import deque
from enum import Enum


# ── WeightEstimator ───────────────────────────────────────────────────────────

class WeightEstimator:
    """
    Estimates payload weight from motor current at steady state.

    Model
    ─────
        weight_kg = max(0, (I_avg − I_idle) / k_per_kg)

    where I_avg is averaged over a rolling window of steady-state samples.
    A sample is accepted only when |tilt| ≤ steady_tilt_threshold_rad.

    Change detection
    ────────────────
    Returns `changed=True` when the new estimate differs from the previous
    by more than change_threshold_kg.  The caller uses this to trigger gain
    rescheduling.
    """

    def __init__(
        self,
        idle_current_a: float = 0.5,
        current_per_kg_a: float = 0.3,
        window_s: float = 3.0,
        control_rate_hz: float = 100.0,
        steady_tilt_threshold_rad: float = 0.05,
        change_threshold_kg: float = 0.3,
    ):
        self._idle = idle_current_a
        self._per_kg = max(1e-6, current_per_kg_a)   # avoid div-by-zero
        self._steady_threshold = steady_tilt_threshold_rad
        self._change_threshold = change_threshold_kg

        window_size = max(10, int(window_s * control_rate_hz))
        self._samples: deque = deque(maxlen=window_size)
        self._last_estimate: float | None = None

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

    @property
    def last_estimate(self) -> float | None:
        return self._last_estimate

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

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

    def update(
        self,
        current_a: float,
        tilt_rad: float,
    ) -> tuple[float | None, bool]:
        """
        Submit one sample (only accepted if in steady state).

        Returns
        -------
        (weight_kg, changed)
        weight_kg : None if fewer than min_samples collected yet
        changed   : True if estimate moved by ≥ change_threshold_kg
        """
        if abs(tilt_rad) <= self._steady_threshold:
            self._samples.append(current_a)

        min_samples = max(5, self._samples.maxlen // 4)
        if len(self._samples) < min_samples:
            return None, False

        avg_i = sum(self._samples) / len(self._samples)
        weight = max(0.0, (avg_i - self._idle) / self._per_kg)

        changed = (
            self._last_estimate is None
            or abs(weight - self._last_estimate) >= self._change_threshold
        )
        self._last_estimate = weight
        return weight, changed

    def reset(self) -> None:
        """Clear sample buffer and last estimate (e.g., after calibration)."""
        self._samples.clear()
        self._last_estimate = None


# ── CalibrationRoutine ────────────────────────────────────────────────────────

class CalibrationState(Enum):
    IDLE     = "idle"
    ROCKING  = "rocking"
    SETTLING = "settling"
    DONE     = "done"
    FAILED   = "failed"


class CalibrationRoutine:
    """
    On-demand calibration: rocks the robot back and forth to disturb the
    balance, lets it settle, then measures steady-state current to produce
    a weight estimate.

    State machine
    ─────────────
        IDLE ──(start)──► ROCKING ──(cycles done)──► SETTLING ──(samples ok)──► DONE
                                                               └──(timeout)────► FAILED

        abort() returns to IDLE from any active state.

    Usage
    ─────
        Each control tick, call step(current_a, tilt_rad, dt).
        The return value is a tilt offset (rad) to add to the balance
        set-point.  Returns 0.0 when not active.
        When state == DONE, read weight_estimate.
    """

    def __init__(
        self,
        amplitude_rad: float = 0.05,
        frequency_hz: float = 0.5,
        rock_cycles: int = 3,
        settle_s: float = 2.0,
        settle_tilt_threshold_rad: float = 0.03,
        idle_current_a: float = 0.5,
        current_per_kg_a: float = 0.3,
    ):
        self._amplitude = amplitude_rad
        self._frequency = max(1e-3, frequency_hz)
        self._rock_cycles = max(1, rock_cycles)
        self._settle_s = settle_s
        self._settle_threshold = settle_tilt_threshold_rad
        self._idle = idle_current_a
        self._per_kg = max(1e-6, current_per_kg_a)

        self._state = CalibrationState.IDLE
        self._elapsed: float = 0.0
        self._settle_elapsed: float = 0.0
        self._settle_samples: list[float] = []
        self._weight_estimate: float | None = None

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

    @property
    def state(self) -> CalibrationState:
        return self._state

    @property
    def weight_estimate(self) -> float | None:
        return self._weight_estimate

    @property
    def is_active(self) -> bool:
        return self._state in (CalibrationState.ROCKING, CalibrationState.SETTLING)

    # ── Control ───────────────────────────────────────────────────────────────

    def start(self) -> bool:
        """
        Start calibration routine.
        Returns False if already running (ROCKING or SETTLING).
        """
        if self.is_active:
            return False
        self._state = CalibrationState.ROCKING
        self._elapsed = 0.0
        self._settle_elapsed = 0.0
        self._settle_samples = []
        self._weight_estimate = None
        return True

    def abort(self) -> None:
        """Immediately abort and return to IDLE."""
        self._state = CalibrationState.IDLE
        self._elapsed = 0.0
        self._settle_elapsed = 0.0

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

    def step(self, current_a: float, tilt_rad: float, dt: float) -> float:
        """
        Advance one control tick.

        Parameters
        ----------
        current_a : motor current reading (A)
        tilt_rad  : current tilt angle (rad)
        dt        : elapsed time since last step (s)

        Returns
        -------
        Tilt offset (rad) to superimpose on the balance set-point.
        0.0 when IDLE / DONE / FAILED.
        """
        if self._state == CalibrationState.ROCKING:
            self._elapsed += dt
            rock_duration = self._rock_cycles / self._frequency

            if self._elapsed >= rock_duration:
                self._state = CalibrationState.SETTLING
                self._settle_elapsed = 0.0
                return 0.0

            return self._amplitude * math.sin(
                2.0 * math.pi * self._frequency * self._elapsed
            )

        if self._state == CalibrationState.SETTLING:
            self._settle_elapsed += dt

            if abs(tilt_rad) <= self._settle_threshold:
                self._settle_samples.append(current_a)

            if self._settle_elapsed >= self._settle_s:
                if len(self._settle_samples) >= 5:
                    avg_i = sum(self._settle_samples) / len(self._settle_samples)
                    self._weight_estimate = max(
                        0.0, (avg_i - self._idle) / self._per_kg
                    )
                    self._state = CalibrationState.DONE
                else:
                    self._state = CalibrationState.FAILED

            return 0.0

        return 0.0
