orisys-sdk/src/orisys/signal_processing.py

"""Signal helpers for force filtering and force calibration.

Batch ``lowpass_filter`` is the zero-phase Butterworth used by example charts.
``LowpassFilter`` is a causal, stateful filter for per-frame ``Sensor`` outputs.

Force calibration is owned by ``orisys.calibration``. These helpers remain as
convenience wrappers:

  1. Bundle JSON via ``calibration.load(sensor_type=...)``
  2. Hardcoded linear defaults for normal force when no file is present
  3. Shear helpers return raw when no curve is available

Prefer ``Sensor.calibration`` / ``read_info(FNORMAL, …)`` in new code.
"""

from __future__ import annotations

import collections
import time
from dataclasses import dataclass

import numpy as np
from scipy.signal import butter, filtfilt, sosfilt, sosfilt_zi

# Linear calibration from raw normal-force units to Newtons (ORY-BOX / ORY-MINI).
# Prefer sensors/<name>/force_calibration.json when available.
NORMAL_FORCE_CAL_SLOPE = 3.92685554668795e-05
NORMAL_FORCE_CAL_OFFSET = 0.0217596967278535

DEFAULT_SAMPLE_RATE_HZ = 30.0
DEFAULT_LOWPASS_ORDER = 2


def _parse_calibration_point(point) -> tuple[float, float]:
    """Legacy helper; prefer ``orisys.calibration``."""
    from .calibration import _parse_point

    return _parse_point(point)


@dataclass(frozen=True)
class ForceChannelCalibration:
    """Compatibility wrapper around ``calibration.ChannelCurve`` convert math."""

    method: str = "linear"
    slope: float | None = None
    offset: float | None = None
    points: tuple[tuple[float, float], ...] = ()
    extrapolate: bool = True
    clamp_output_min: float | None = 0.0
    input_mode: str = "magnitude"
    signed: bool = False
    unit: str = "N"

    @classmethod
    def linear(cls, slope: float, offset: float, *, unit: str = "N") -> ForceChannelCalibration:
        return cls(method="linear", slope=float(slope), offset=float(offset), unit=unit)

    @classmethod
    def from_channel_dict(
        cls,
        block: dict,
        *,
        unit: str = "N",
        default_method: str = "piecewise_linear",
    ) -> ForceChannelCalibration | None:
        from .calibration import _channel_from_block

        curve = _channel_from_block(
            "fn", block, unit=unit, default_method=default_method
        )
        if curve is None:
            return None
        return cls(
            method=curve.method,
            slope=curve.slope,
            offset=curve.offset,
            points=curve.points,
            extrapolate=curve.extrapolate,
            clamp_output_min=curve.clamp_output_min,
            input_mode=curve.input_mode,
            signed=curve.signed,
            unit=curve.unit,
        )

    def convert(self, raw: float) -> float:
        from .calibration import ChannelCurve

        return ChannelCurve(
            id="fn",
            method=self.method,
            points=self.points,
            slope=self.slope if self.slope is not None else NORMAL_FORCE_CAL_SLOPE,
            offset=self.offset if self.offset is not None else NORMAL_FORCE_CAL_OFFSET,
            signed=self.signed,
            extrapolate=self.extrapolate,
            clamp_output_min=self.clamp_output_min,
            input_mode=self.input_mode,
            unit=self.unit,
        ).convert(raw)

    def convert_component(self, raw: float) -> float:
        if not self.signed:
            return self.convert(raw)
        raw = float(raw)
        sign = -1.0 if raw < 0.0 else 1.0
        return sign * self.convert(abs(raw))


_DEFAULT_NORMAL = ForceChannelCalibration.linear(
    NORMAL_FORCE_CAL_SLOPE,
    NORMAL_FORCE_CAL_OFFSET,
)


def _load_calibration(sensor_name: str | None):
    if not sensor_name:
        return None
    try:
        from .calibration import load

        return load(sensor_type=sensor_name)
    except Exception:
        return None


def _normal_calibration(sensor_name: str | None = None) -> ForceChannelCalibration:
    cal = _load_calibration(sensor_name)
    if cal is not None:
        curve = cal.get("fn")
        if curve is not None:
            return ForceChannelCalibration(
                method=curve.method,
                slope=curve.slope,
                offset=curve.offset,
                points=curve.points,
                extrapolate=curve.extrapolate,
                clamp_output_min=curve.clamp_output_min,
                input_mode=curve.input_mode,
                signed=curve.signed,
                unit=curve.unit,
            )
    return _DEFAULT_NORMAL


def _shear_curve(sensor_name: str | None, channel: str):
    cal = _load_calibration(sensor_name)
    if cal is None:
        return None
    return cal.get(channel)


def _load_normal_force_coeffs(sensor_name: str | None = None) -> tuple[float, float]:
    """Legacy helper: linear slope/offset only (piecewise returns hardcoded defaults)."""
    cal = _normal_calibration(sensor_name)
    if cal.method == "linear" and cal.slope is not None and cal.offset is not None:
        return float(cal.slope), float(cal.offset)
    return NORMAL_FORCE_CAL_SLOPE, NORMAL_FORCE_CAL_OFFSET


def calibrate_normal_force(raw_value: float, *, sensor_name: str | None = None) -> float:
    """Convert a raw normal-force reading to Newtons.

    Uses ``sensors/<name>/force_calibration.json`` when present (piecewise or linear).
    Falls back to hardcoded linear defaults when missing.
    """
    return _normal_calibration(sensor_name).convert(raw_value)


def calibrate_shear_force(raw_value: float, *, sensor_name: str | None = None) -> float:
    """Convert a raw shear magnitude to Newtons using the shear channel curve.

    Prefers the average of FX/FY magnitude maps when present.
    Returns ``float(raw_value)`` when unavailable.
    """
    cal_x = _shear_curve(sensor_name, "fx")
    cal_y = _shear_curve(sensor_name, "fy")
    if cal_x is not None and cal_y is not None:
        return 0.5 * (
            cal_x.convert(abs(float(raw_value))) + cal_y.convert(abs(float(raw_value)))
        )
    if cal_x is not None:
        return cal_x.convert(abs(float(raw_value)))
    if cal_y is not None:
        return cal_y.convert(abs(float(raw_value)))
    return float(raw_value)


def calibrate_shear_components(
    shear_x: float,
    shear_y: float,
    *,
    sensor_name: str | None = None,
) -> tuple[float, float]:
    """Convert raw shear components to Newtons via ``calibration.load`` curves."""
    cal = _load_calibration(sensor_name)
    if cal is None:
        return float(shear_x), float(shear_y)
    return cal.convert(0.0, float(shear_x), float(shear_y))[1:]


def _nyquist_hz(sample_rate_hz: float) -> float:
    return 0.5 * float(sample_rate_hz)


def _validate_cutoff(cutoff_hz: float, sample_rate_hz: float) -> None:
    nyquist = _nyquist_hz(sample_rate_hz)
    if cutoff_hz < 0 or cutoff_hz >= nyquist:
        raise ValueError(
            f"cutoff_hz must be in (0, {nyquist}) for sample_rate_hz={sample_rate_hz}"
        )


def lowpass_filter(
    data,
    cutoff_hz: float = 5.0,
    sample_rate_hz: float = DEFAULT_SAMPLE_RATE_HZ,
    order: int = DEFAULT_LOWPASS_ORDER,
) -> np.ndarray:
    """Zero-phase Butterworth low-pass for uniformly sampled 1-D data.

    Matches the example/viewer chart filter: ``cutoff_hz == 0`` is a no-op.
    """
    y = np.asarray(data, dtype=np.float64)
    if cutoff_hz == 0:
        return y.copy()
    if y.size == 0:
        return y

    _validate_cutoff(cutoff_hz, sample_rate_hz)
    b, a = butter(int(order), cutoff_hz / _nyquist_hz(sample_rate_hz), btype="low")
    padlen = 3 * (max(len(a), len(b)) - 1)
    if y.size <= padlen:
        return y.copy()
    try:
        return filtfilt(b, a, y)
    except ValueError:
        return y.copy()


class LowpassFilter:
    """Causal Butterworth low-pass with persistent SOS state (one sample at a time)."""

    def __init__(
        self,
        cutoff_hz: float = 0.0,
        sample_rate_hz: float = DEFAULT_SAMPLE_RATE_HZ,
        order: int = DEFAULT_LOWPASS_ORDER,
        n_channels: int = 3,
    ):
        self.n_channels = int(n_channels)
        self._sos = None
        self._zi = None
        self.configure(
            cutoff_hz=cutoff_hz,
            sample_rate_hz=sample_rate_hz,
            order=order,
        )

    @property
    def enabled(self) -> bool:
        return self._sos is not None

    def configure(
        self,
        *,
        cutoff_hz: float | None = None,
        sample_rate_hz: float | None = None,
        order: int | None = None,
    ) -> None:
        if cutoff_hz is not None:
            self.cutoff_hz = float(cutoff_hz)
        if sample_rate_hz is not None:
            self.sample_rate_hz = float(sample_rate_hz)
        if order is not None:
            self.order = int(order)

        if self.cutoff_hz <= 0:
            self._sos = None
            self._zi = None
            return

        _validate_cutoff(self.cutoff_hz, self.sample_rate_hz)
        self._sos = butter(
            self.order,
            self.cutoff_hz / _nyquist_hz(self.sample_rate_hz),
            btype="low",
            output="sos",
        )
        self._zi = None

    def reset(self, x0=None) -> None:
        """Clear filter memory. Optional ``x0`` initializes to a steady-state DC."""
        if self._sos is None:
            self._zi = None
            return
        if x0 is None:
            self._zi = None
            return
        x = np.asarray(x0, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"x0 must have {self.n_channels} channels, got {x.size}")
        zi0 = sosfilt_zi(self._sos)
        self._zi = zi0[:, :, None] * x[None, None, :]

    def process(self, values) -> np.ndarray:
        """Filter one multi-channel sample. Returns float64 array of shape (n_channels,)."""
        x = np.asarray(values, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"expected {self.n_channels} channels, got {x.size}")
        if self._sos is None:
            return x.copy()
        if self._zi is None:
            self.reset(x)
        y, self._zi = sosfilt(self._sos, x[None, :], zi=self._zi, axis=0)
        return np.asarray(y[0], dtype=np.float64)


class CausalMedianFilter:
    """Causal sliding-median over the last ``window`` samples, one sample at a time.

    The optical-flow force chain produces heavy-tailed jitter (rare multi-sigma
    spikes). A linear low-pass spreads such a spike over several frames instead
    of deleting it, whereas a median removes it outright - measured on idle
    data, a 5-sample median cut the jitter 4x better than a moving average of
    the same width.

    ``window <= 1`` disables the filter (pass-through). Until ``window`` samples
    have arrived the median of the available samples is returned, so the output
    is well defined from the first frame.
    """

    def __init__(self, window: int = 1, n_channels: int = 3):
        self.n_channels = int(n_channels)
        self.window = 1
        self._buf = None
        self.configure(window)

    @property
    def enabled(self) -> bool:
        return self.window > 1

    def configure(self, window: int | None = None) -> None:
        if window is not None:
            try:
                self.window = max(1, int(window))
            except (TypeError, ValueError):
                self.window = 1
        self._buf = None

    def reset(self, x0=None) -> None:
        """Drop the window history (e.g. after a tracker reset)."""
        self._buf = None

    def process(self, values) -> np.ndarray:
        """Filter one multi-channel sample. Returns float64 array of shape (n_channels,)."""
        x = np.asarray(values, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"expected {self.n_channels} channels, got {x.size}")
        if not self.enabled:
            return x.copy()
        if self._buf is None or self._buf.shape[1] != x.size:
            self._buf = x[None, :].copy()
            return x.copy()
        self._buf = np.vstack((self._buf, x[None, :]))[-self.window:]
        return np.median(self._buf, axis=0)


class MultiAxisCausalFilter:
    """Causal median → low-pass over a vector, one independent state per axis.

    The force output needs the same chain twice (raw features and newtons), and
    an unavailable axis (NaN, or an axis a model never calibrated) must neither
    enter the filter history nor pull a valid axis towards zero. One
    :class:`CausalMedianFilter` and one :class:`LowpassFilter` per axis give that
    without a second implementation of either operator: callers keep one instance
    per processed triplet and feed a per-axis validity mask.

    ``median_window <= 1`` and ``lowpass_cutoff_hz <= 0`` each switch their own
    stage off; with both off the input is returned unchanged, NaN included.
    """

    def __init__(
        self,
        *,
        median_window: int = 1,
        lowpass_cutoff_hz: float = 0.0,
        sample_rate_hz: float = DEFAULT_SAMPLE_RATE_HZ,
        lowpass_order: int = DEFAULT_LOWPASS_ORDER,
        n_channels: int = 3,
    ):
        self.n_channels = max(1, int(n_channels))
        try:
            self._median_window = max(1, int(median_window))
        except (TypeError, ValueError):
            self._median_window = 1
        self._lowpass_cutoff_hz = float(lowpass_cutoff_hz or 0.0)
        self._sample_rate_hz = float(sample_rate_hz or DEFAULT_SAMPLE_RATE_HZ)
        try:
            self._lowpass_order = max(1, int(lowpass_order))
        except (TypeError, ValueError):
            self._lowpass_order = DEFAULT_LOWPASS_ORDER
        self._medians: list[CausalMedianFilter] = []
        self._lowpasses: list[LowpassFilter] = []
        self._build()

    @property
    def enabled(self) -> bool:
        return bool(self._medians and self._medians[0].enabled) or bool(
            self._lowpasses and self._lowpasses[0].enabled
        )

    @property
    def median_window(self) -> int:
        return int(self._medians[0].window) if self._medians else 1

    @property
    def lowpass_cutoff_hz(self) -> float:
        return float(self._lowpasses[0].cutoff_hz) if self._lowpasses else 0.0

    def configure(
        self,
        *,
        median_window: int | None = None,
        lowpass_cutoff_hz: float | None = None,
        sample_rate_hz: float | None = None,
        lowpass_order: int | None = None,
    ) -> None:
        """Re-tune the chain; every axis state is dropped (callers re-warm it)."""
        if median_window is not None:
            try:
                self._median_window = max(1, int(median_window))
            except (TypeError, ValueError):
                self._median_window = 1
        if lowpass_cutoff_hz is not None:
            self._lowpass_cutoff_hz = float(lowpass_cutoff_hz)
        if sample_rate_hz is not None:
            self._sample_rate_hz = float(sample_rate_hz)
        if lowpass_order is not None:
            try:
                self._lowpass_order = max(1, int(lowpass_order))
            except (TypeError, ValueError):
                self._lowpass_order = DEFAULT_LOWPASS_ORDER
        self._build()

    def _build(self) -> None:
        self._medians = [
            CausalMedianFilter(window=self._median_window, n_channels=1)
            for _ in range(self.n_channels)
        ]
        self._lowpasses = [
            LowpassFilter(
                cutoff_hz=self._lowpass_cutoff_hz,
                sample_rate_hz=self._sample_rate_hz,
                order=self._lowpass_order,
                n_channels=1,
            )
            for _ in range(self.n_channels)
        ]

    def reset(self) -> None:
        """Drop every axis' filter memory (new zero, model change, discontinuity)."""
        for median in self._medians:
            median.reset()
        for lowpass in self._lowpasses:
            lowpass.reset()

    def process(self, values, valid=None) -> np.ndarray:
        """Filter one vector; invalid axes stay NaN and do not advance their state."""
        x = np.asarray(values, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"expected {self.n_channels} channels, got {x.size}")
        if valid is None:
            usable = np.isfinite(x)
        else:
            usable = np.asarray(valid, dtype=bool).reshape(-1) & np.isfinite(x)
            if usable.size != self.n_channels:
                raise ValueError(
                    f"expected {self.n_channels} validity flags, got {usable.size}"
                )
        out = np.full(self.n_channels, np.nan, dtype=np.float64)
        for index in range(self.n_channels):
            if not usable[index]:
                continue
            value = self._medians[index].process(x[index : index + 1])
            out[index] = self._lowpasses[index].process(value)[0]
        return out

    def state(self) -> dict:
        return {
            "median_window": self.median_window,
            "median_applied": bool(self._medians and self._medians[0].enabled),
            "lowpass_cutoff_hz": self.lowpass_cutoff_hz,
            "lowpass_applied": bool(self._lowpasses and self._lowpasses[0].enabled),
            "lowpass_order": (
                int(self._lowpasses[0].order) if self._lowpasses else DEFAULT_LOWPASS_ORDER
            ),
            "lowpass_sample_rate_hz": (
                float(self._lowpasses[0].sample_rate_hz)
                if self._lowpasses
                else DEFAULT_SAMPLE_RATE_HZ
            ),
        }


class SampleRateEstimator:
    """Robust online estimate of the rate at which :meth:`tick` is called.

    Uses the *median* of the most recent inter-call intervals, so one stall or
    scheduling hiccup does not poison the estimate. :attr:`rate_hz` returns the
    seed value until ``min_samples`` intervals have been observed.

    This exists because a fixed floor filter designed for 30 Hz behaves
    differently at 60 Hz: the cutoff stays put in Hz but the per-frame
    smoothing halves. Measuring the achieved rate keeps the filter identical in
    the time domain whatever the deployment runs at.
    """

    def __init__(
        self,
        initial_hz: float = 30.0,
        window: int = 45,
        min_samples: int = 8,
        max_interval_s: float = 1.0,
    ):
        self.initial_hz = float(initial_hz) if initial_hz else 30.0
        self.window = max(3, int(window))
        self.min_samples = max(2, int(min_samples))
        self.max_interval_s = float(max_interval_s)
        self.reset()

    def reset(self, initial_hz: float | None = None) -> None:
        if initial_hz is not None:
            self.initial_hz = float(initial_hz) if initial_hz else 30.0
        self._dts: collections.deque = collections.deque(maxlen=self.window)
        self._last = None
        self._rate = self.initial_hz
        self._measured = False

    @property
    def rate_hz(self) -> float:
        return self._rate

    @property
    def measured(self) -> bool:
        """True once the estimate comes from real intervals."""
        return self._measured

    def tick(self, now: float | None = None) -> float:
        """Register one sample and return the current rate estimate."""
        now = time.perf_counter() if now is None else float(now)
        last, self._last = self._last, now
        if last is None:
            return self._rate
        dt = now - last
        if not (1e-5 < dt <= self.max_interval_s):
            # Stalls, pauses and the first frame carry no rate information.
            return self._rate
        self._dts.append(dt)
        if len(self._dts) >= self.min_samples:
            dt_med = float(np.median(self._dts))
            if dt_med > 0.0:
                self._rate = 1.0 / dt_med
                self._measured = True
        return self._rate