"""Causal, dependency-free smoothing filters for head-pose angles.

All filters share one interface so they can be swapped from configuration::

    f = create_angle_filter("median_ema", alpha=0.4, window=5)
    smoothed = f.update(raw_angle_deg, elapsed_time_sec)

``update`` must be called once per valid sample in time order.  ``reset`` forgets
all state (call it when a track is lost).  Every filter is causal, so each adds
some delay; see the per-class notes.  Delay and noise should be judged together
(see ``src.yaw_noise_report``).

Kinds
-----
``none``             pass-through.
``ema``              exponential moving average, per sample (the original behaviour).
``median_ema``       sliding median over ``window`` samples, then EMA.  Removes
                     isolated spikes; step delay is about (window - 1) / 2 samples.
``hampel_ema``       Hampel identifier: a sample further than ``hampel_threshold``
                     robust sigmas from the window median is replaced by that median,
                     everything else passes unchanged, then EMA.  Keeps small real
                     changes sharper than a plain median.
``one_euro``         One Euro Filter (time based).  Beware: with very noisy input the
                     speed term inflates the cutoff, so keep ``beta`` near 0 or
                     put a median in front (``median_one_euro``).
``median_one_euro``  sliding median, then One Euro Filter.

Median-type windows hold ``window`` samples by default, so their duration depends
on the pose update rate (frame rate / ``face.process_every_n_frames``).  With
``window_sec > 0`` the window instead holds every sample from the last
``window_sec`` seconds (trimmed to an odd count), which keeps the filter strength
independent of the update rate.  A median removes a spike run of at most
(n - 1) / 2 samples out of n.
"""

from __future__ import annotations

import math
import statistics
from collections import deque
from typing import Protocol

FILTER_KINDS: tuple[str, ...] = ("none", "ema", "median_ema", "hampel_ema", "one_euro", "median_one_euro")

# Floor for the robust sigma used by the Hampel identifier [deg]; avoids flagging
# every sample as an outlier when the window is nearly constant.
HAMPEL_MIN_SCALE_DEG = 1.0
_MAD_TO_SIGMA = 1.4826


class _SampleWindow:
    """Recent samples bounded by count, or by age when ``duration_sec`` > 0."""

    MAX_SAMPLES_IN_DURATION_MODE = 256

    def __init__(self, size: int, duration_sec: float = 0.0) -> None:
        if size < 1:
            raise ValueError("window must be >= 1")
        if duration_sec < 0.0:
            raise ValueError("window_sec must be >= 0")
        self._size = size
        self._duration = duration_sec
        self._items: deque[tuple[float, float]] = deque()

    def add(self, t: float, value: float) -> list[float]:
        """Append a sample and return the values currently inside the window."""
        self._items.append((t, value))
        if self._duration > 0.0:
            while len(self._items) > 1 and t - self._items[0][0] > self._duration:
                self._items.popleft()
            while len(self._items) > self.MAX_SAMPLES_IN_DURATION_MODE:
                self._items.popleft()
            values = [item[1] for item in self._items]
            # An odd count makes the median an actual sample (and symmetric).
            return values[1:] if len(values) % 2 == 0 else values
        while len(self._items) > self._size:
            self._items.popleft()
        return [item[1] for item in self._items]

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


class AngleFilter(Protocol):
    def update(self, value: float, t: float) -> float: ...

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


class PassThroughFilter:
    def update(self, value: float, t: float) -> float:
        return value

    def reset(self) -> None:
        return None


class EmaFilter:
    """``y = alpha * x + (1 - alpha) * y_prev``; the first sample passes through."""

    def __init__(self, alpha: float) -> None:
        if not 0.0 < alpha <= 1.0:
            raise ValueError("alpha must be in (0, 1]")
        self.alpha = alpha
        self._state: float | None = None

    def update(self, value: float, t: float) -> float:
        self._state = value if self._state is None else self.alpha * value + (1.0 - self.alpha) * self._state
        return self._state

    def reset(self) -> None:
        self._state = None


class MedianEmaFilter:
    def __init__(self, window: int, alpha: float, window_sec: float = 0.0) -> None:
        self._window = _SampleWindow(window, window_sec)
        self._ema = EmaFilter(alpha)

    def update(self, value: float, t: float) -> float:
        return self._ema.update(statistics.median(self._window.add(t, value)), t)

    def reset(self) -> None:
        self._window.clear()
        self._ema.reset()


class HampelEmaFilter:
    def __init__(self, window: int, threshold: float, alpha: float, window_sec: float = 0.0) -> None:
        if window < 3:
            raise ValueError("hampel window must be >= 3")
        if threshold <= 0.0:
            raise ValueError("hampel threshold must be positive")
        self.threshold = threshold
        self._window = _SampleWindow(window, window_sec)
        self._ema = EmaFilter(alpha)

    def update(self, value: float, t: float) -> float:
        values = self._window.add(t, value)
        cleaned = value
        if len(values) >= 3:
            center = statistics.median(values)
            mad = statistics.median(abs(item - center) for item in values)
            scale = max(_MAD_TO_SIGMA * mad, HAMPEL_MIN_SCALE_DEG)
            if abs(value - center) > self.threshold * scale:
                cleaned = center
        return self._ema.update(cleaned, t)

    def reset(self) -> None:
        self._window.clear()
        self._ema.reset()


class OneEuroFilter:
    """One Euro Filter (Casiez et al., 2012) driven by real timestamps [s]."""

    def __init__(self, min_cutoff: float, beta: float, d_cutoff: float = 1.0) -> None:
        if min_cutoff <= 0.0 or d_cutoff <= 0.0 or beta < 0.0:
            raise ValueError("min_cutoff and d_cutoff must be > 0 and beta >= 0")
        self.min_cutoff = min_cutoff
        self.beta = beta
        self.d_cutoff = d_cutoff
        self._x: float | None = None
        self._dx = 0.0
        self._t = 0.0

    @staticmethod
    def _alpha(dt: float, cutoff: float) -> float:
        tau = 1.0 / (2.0 * math.pi * cutoff)
        return 1.0 / (1.0 + tau / dt)

    def update(self, value: float, t: float) -> float:
        if self._x is None:
            self._x, self._t, self._dx = value, t, 0.0
            return value
        dt = max(t - self._t, 1e-6)
        self._t = t
        raw_dx = (value - self._x) / dt
        a_d = self._alpha(dt, self.d_cutoff)
        self._dx = a_d * raw_dx + (1.0 - a_d) * self._dx
        cutoff = self.min_cutoff + self.beta * abs(self._dx)
        a = self._alpha(dt, cutoff)
        self._x = a * value + (1.0 - a) * self._x
        return self._x

    def reset(self) -> None:
        self._x = None
        self._dx = 0.0


class MedianOneEuroFilter:
    def __init__(
        self, window: int, min_cutoff: float, beta: float, d_cutoff: float = 1.0, window_sec: float = 0.0,
    ) -> None:
        self._window = _SampleWindow(window, window_sec)
        self._euro = OneEuroFilter(min_cutoff, beta, d_cutoff)

    def update(self, value: float, t: float) -> float:
        return self._euro.update(statistics.median(self._window.add(t, value)), t)

    def reset(self) -> None:
        self._window.clear()
        self._euro.reset()


class GapResettingFilter:
    """Reset the wrapped filter when samples are more than ``max_gap_sec`` apart.

    Prevents stale state from blending into the first samples after the face was
    lost.  ``max_gap_sec <= 0`` disables the behaviour.
    """

    def __init__(self, inner: AngleFilter, max_gap_sec: float) -> None:
        self._inner = inner
        self._max_gap = max_gap_sec
        self._last_t: float | None = None

    def update(self, value: float, t: float) -> float:
        if self._max_gap > 0.0 and self._last_t is not None and t - self._last_t > self._max_gap:
            self._inner.reset()
        self._last_t = t
        return self._inner.update(value, t)

    def reset(self) -> None:
        self._inner.reset()
        self._last_t = None


def create_angle_filter(
    kind: str,
    *,
    alpha: float = 0.4,
    window: int = 5,
    hampel_threshold: float = 3.0,
    one_euro_min_cutoff: float = 1.0,
    one_euro_beta: float = 0.0,
    one_euro_d_cutoff: float = 1.0,
    reset_gap_sec: float = 0.0,
    window_sec: float = 0.0,
) -> AngleFilter:
    """Build a filter by name; raises ``ValueError`` for unknown kinds."""
    inner: AngleFilter
    if kind == "none":
        inner = PassThroughFilter()
    elif kind == "ema":
        inner = EmaFilter(alpha)
    elif kind == "median_ema":
        inner = MedianEmaFilter(window, alpha, window_sec)
    elif kind == "hampel_ema":
        inner = HampelEmaFilter(window, hampel_threshold, alpha, window_sec)
    elif kind == "one_euro":
        inner = OneEuroFilter(one_euro_min_cutoff, one_euro_beta, one_euro_d_cutoff)
    elif kind == "median_one_euro":
        inner = MedianOneEuroFilter(window, one_euro_min_cutoff, one_euro_beta, one_euro_d_cutoff, window_sec)
    else:
        raise ValueError(f"unknown smoothing filter {kind!r}; expected one of {FILTER_KINDS}")
    return GapResettingFilter(inner, reset_gap_sec) if reset_gap_sec > 0.0 else inner