Blame for runtime/backends/actioncodec2/preprocessing.py

"""Versioned preprocessing of continuous segments, after discontinuity QC."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.ndimage import median_filter
from scipy.signal import medfilt, savgol_filter
from scipy.spatial.transform import Rotation

from .native_backend import native_kernel
from .semantic import discretize_gripper_actions


def step_norms(values: np.ndarray) -> np.ndarray:
    """L2 norm of per-step deltas, ``[T-1]`` for ``[T, D]`` values."""
    kernel = native_kernel("step_norms")
    if kernel is not None:
        values = np.asarray(values)
        result = kernel(values, values.dtype == np.float32)
        return result.astype(np.float32) if values.dtype == np.float32 else result
    return np.linalg.norm(np.diff(values, axis=0), axis=1)


def _legacy_boundary_trim_indices(
    values: np.ndarray,
    *,
    tau: float,
    min_idle_steps: int,
) -> tuple[int, int]:
    kernel = native_kernel("boundary_trim")
    if kernel is not None:
        return tuple(
            kernel(values, tau, min_idle_steps, 1, True, np.asarray(values).dtype == np.float32)
        )
    if len(values) < 2:
        return 0, len(values)
    idle = step_norms(values) <= tau
    moving = ~idle
    if not moving.any():
        return len(values), len(values)
    head = int(np.argmax(moving))
    trailing = int(np.argmax(moving[::-1]))
    tail = len(idle) - trailing

    start = head + 1 if head >= min_idle_steps else 0
    stop = tail + 1 if trailing >= min_idle_steps else len(values)
    if stop <= start:
        return start, start
    return start, stop


def boundary_trim_indices(
    values: np.ndarray,
    *,
    tau: float,
    min_idle_steps: int,
    context_frames: int = 1,
) -> tuple[int, int]:
    """Propose a boundary slice retaining the first moving edge and context."""
    if context_frames < 1:
        raise ValueError("context_frames must be positive")
    kernel = native_kernel("boundary_trim")
    if kernel is not None:
        return tuple(
            kernel(
                values,
                tau,
                min_idle_steps,
                context_frames,
                False,
                np.asarray(values).dtype == np.float32,
            )
        )
    if len(values) < 2:
        return 0, len(values)
    moving = step_norms(values) > tau
    if not moving.any():
        return 0, len(values)
    head = int(np.argmax(moving))
    trailing = int(np.argmax(moving[::-1]))
    start = max(0, head - context_frames + 1) if head >= min_idle_steps else 0
    stop = len(values) - trailing if trailing >= min_idle_steps else len(values)
    return start, max(stop, start + 1)


def trim_deviation(
    values: np.ndarray,
    start: int,
    stop: int,
    *,
    profile: str,
    gripper_column: int | None = None,
) -> dict[str, object]:
    """Maximum physical error from constant-filling removed boundary rows."""
    if not 0 <= start < stop <= len(values):
        raise ValueError("trim slice must retain at least one row")
    kernel = native_kernel("trim_deviation")
    if kernel is not None:
        return kernel(
            values,
            start,
            stop,
            profile,
            -1 if gripper_column is None else gripper_column % values.shape[1],
        )
    width = values.shape[1]
    motion = [i for i in range(width) if gripper_column is None or i != gripper_column % width]
    result: dict[str, object] = {
        "motion_max_error": [0.0] * len(motion),
        "translation_max_error_m": 0.0,
        "rotation_max_error_rad": 0.0,
        "gripper_changed": False,
    }
    for removed, anchor in ((values[:start], values[start]), (values[stop:], values[stop - 1])):
        if not len(removed):
            continue
        delta = np.abs(removed[:, motion].astype(np.float64) - anchor[motion])
        result["motion_max_error"] = np.maximum(
            result["motion_max_error"], delta.max(axis=0)
        ).tolist()
        if profile == "eef":
            result["translation_max_error_m"] = max(
                result["translation_max_error_m"], float(np.linalg.norm(delta[:, :3], axis=1).max())
            )
            rotations = Rotation.from_rotvec(removed[:, motion[3:6]])
            relative = Rotation.from_rotvec(anchor[motion[3:6]]).inv() * rotations
            result["rotation_max_error_rad"] = max(
                result["rotation_max_error_rad"], float(relative.magnitude().max())
            )
        if gripper_column is not None:
            result["gripper_changed"] |= bool(
                np.any(removed[:, gripper_column] != anchor[gripper_column])
            )
    return result


def savgol_smooth(
    values: np.ndarray,
    *,
    window: int,
    polyorder: int = 3,
) -> np.ndarray:
    """Zero-phase Savitzky-Golay smoothing along time, per dimension."""
    length = len(values)
    if length < 3 or window < 3:
        return np.asarray(values, dtype=np.float32)
    if window >= length:
        window = length if length % 2 == 1 else length - 1
    if polyorder >= window:
        polyorder = window - 2 if window >= 3 else 1
    kernel = native_kernel("savgol_smooth")
    if kernel is not None:
        return kernel(values, window, polyorder)
    smoothed = savgol_filter(np.asarray(values, dtype=np.float64), window, polyorder, axis=0)
    return smoothed.astype(np.float32)


def ease_overspeed(
    values: np.ndarray,
    *,
    limits: np.ndarray,
    horizon: int | None = None,
    max_passes: int | None = None,
) -> np.ndarray:
    """Slew-rate-limit per-step deltas, repaying the deficit downstream."""
    del horizon, max_passes
    out = np.array(values, dtype=np.float64, copy=True)
    if len(out) < 2:
        return out.astype(np.float32)
    limit = np.broadcast_to(np.asarray(limits, dtype=np.float64), (out.shape[1],))
    kernel = native_kernel("ease_overspeed")
    if kernel is not None:
        return kernel(out, limit)
    steps = np.diff(out, axis=0)
    violating_dims = np.flatnonzero(np.any(np.abs(steps) > limit, axis=0))
    for j in violating_dims:
        debt = 0.0
        for k in range(len(steps)):
            candidate = steps[k, j] + debt
            if abs(candidate) <= limit[j]:
                steps[k, j] = candidate
                debt = 0.0
            else:
                sign = 1.0 if candidate > 0 else -1.0
                steps[k, j] = sign * limit[j]
                debt = candidate - sign * limit[j]
    out[1:] = out[0] + np.cumsum(steps, axis=0)
    return out.astype(np.float32)


def _median(values: np.ndarray, window: int, *, nearest: bool) -> np.ndarray:
    kernel = native_kernel("median_filter")
    if kernel is not None:
        return kernel(values, window, nearest).astype(values.dtype)
    return (
        median_filter(values, size=window, mode="nearest")
        if nearest
        else medfilt(values, kernel_size=window)
    )


def _motion_change(values, reference, columns):
    kernel = native_kernel("max_change")
    if kernel is not None:
        return kernel(
            values, reference, columns, values.dtype == np.float32 and reference.dtype == np.float32
        )
    return (
        np.abs(values[:, columns] - reference[:, columns]).max(axis=0).tolist() if columns else []
    )


def _overspeed_steps(values, limits):
    kernel = native_kernel("overspeed_steps")
    if kernel is not None:
        return kernel(values, limits, values.dtype == np.float32)
    return int(np.count_nonzero(np.any(np.abs(np.diff(values, axis=0)) > limits, axis=1)))


@dataclass(frozen=True)
class PreprocessingSpec:
    """Parameters for the shared preprocessing hook."""

    tau: float = 0.002
    min_idle_steps: int = 15
    savgol_window: int = 0
    savgol_polyorder: int = 3
    ease_horizon: int = 0
    smooth_columns: tuple[int, ...] | None = None
    gripper_mode: str = "binary"
    gripper_threshold: float = 0.8
    gripper_median_window: int = 0
    version: int = 2
    trim_mode: str = "shadow"
    profile: str = "joint"
    context_frames: int = 1
    motion_limits: tuple[float, ...] | None = None
    translation_limit: float | None = None
    rotation_limit: float | None = None

    def __post_init__(self) -> None:
        if self.version not in (1, 2):
            raise ValueError("unsupported preprocessing version")
        if self.trim_mode not in ("off", "shadow", "bounded"):
            raise ValueError("trim_mode must be off, shadow or bounded")
        if self.profile not in ("joint", "eef") or self.context_frames < 1:
            raise ValueError("invalid profile or context_frames")
        if self.savgol_window != 0 and (self.savgol_window < 3 or self.savgol_window % 2 == 0):
            raise ValueError("savgol_window must be zero or odd and at least three")
        if self.savgol_polyorder < 1 or (
            self.savgol_window and self.savgol_polyorder >= self.savgol_window
        ):
            raise ValueError("invalid Savitzky-Golay polynomial order")
        if self.ease_horizon < 0 or (
            self.gripper_median_window >= 3 and self.gripper_median_window % 2 == 0
        ):
            raise ValueError("invalid ease or gripper median window")
        if not np.isfinite(self.tau) or self.tau < 0 or self.min_idle_steps < 1:
            raise ValueError("invalid idle candidate parameters")
        budgets = [*(self.motion_limits or ()), self.translation_limit, self.rotation_limit]
        if any(v is not None and (not np.isfinite(v) or v <= 0) for v in budgets):
            raise ValueError("physical trim budgets must be finite and positive")
        if self.version == 2 and self.trim_mode == "bounded":
            if self.profile == "joint" and not self.motion_limits:
                raise ValueError("bounded joint trimming requires motion_limits")
            if self.profile == "eef" and (
                self.translation_limit is None or self.rotation_limit is None
            ):
                raise ValueError("bounded EEF trimming requires translation and rotation budgets")
        if self.gripper_mode not in ("binary", "eleven", "twenty"):
            raise ValueError("gripper_mode must be 'binary', 'eleven' or 'twenty'")
        if not 0.0 < self.gripper_threshold < 1.0:
            raise ValueError("gripper_threshold must lie strictly inside (0, 1)")

    def apply(
        self,
        values: np.ndarray,
        *,
        ease_limits: np.ndarray | None = None,
        smooth_columns: tuple[int, ...] | None = None,
        gripper_column: int | None = None,
        stats: dict[str, object] | None = None,
        active_dims: tuple[int, ...] | None = None,
        gripper_is_discrete: bool = False,
    ) -> np.ndarray:
        """Transform one finite continuous segment; record alignment and loss."""
        values = np.asarray(values)
        if values.ndim != 2 or not values.shape[1] or not np.isfinite(values).all():
            raise ValueError("preprocessing requires finite [T,D] values")
        log = {} if stats is None else stats
        if not len(values):
            log.update(start=0, stop=0, ease_steps=0)
            return values.copy()
        if self.version == 1:
            return self._legacy_apply_recorded(
                values, ease_limits, smooth_columns, gripper_column, log
            )
        start, stop = self._trim_slice(values, gripper_column, active_dims, log)
        original = values[start:stop]
        out = np.array(original, dtype=np.result_type(original.dtype, np.float32), copy=True)
        columns = self.smooth_columns if smooth_columns is None else smooth_columns
        columns = tuple(range(out.shape[1])) if columns is None else tuple(columns)
        if self.profile == "eef":
            columns = tuple(c for c in columns if c in (0, 1, 2))
        columns = tuple(
            c
            for c in columns
            if gripper_column is None or c % out.shape[1] != gripper_column % out.shape[1]
        )
        if self.savgol_window >= 3 and columns:
            out[:, columns] = savgol_smooth(
                out[:, columns], window=self.savgol_window, polyorder=self.savgol_polyorder
            )
            out[: self.context_frames, columns] = original[: self.context_frames, columns]
        log["motion_max_change"] = _motion_change(out, original, columns)
        log["ease_steps"] = 0
        if self.ease_horizon >= 1 and ease_limits is not None:
            if self.profile == "eef":
                raise ValueError("additive easing cannot modify EEF rotations")
            ease_limits = np.broadcast_to(ease_limits, (out.shape[1],)).copy()
            if gripper_column is not None:
                ease_limits[gripper_column] = np.inf
            log["ease_steps"] = _overspeed_steps(out, ease_limits)
            eased = ease_overspeed(out, limits=ease_limits)
            log["ease_endpoint_reverted"] = not bool(
                np.allclose(eased[-1], out[-1], atol=1e-6, rtol=0)
                and np.allclose(
                    eased[: self.context_frames], out[: self.context_frames], atol=1e-6, rtol=0
                )
            )
            if not log["ease_endpoint_reverted"]:
                out = eased
        if gripper_column is not None:
            column = gripper_column % out.shape[1]
            raw = original[:, column : column + 1]
            reference = discretize_gripper_actions(
                raw, mode=self.gripper_mode, threshold=self.gripper_threshold
            )
            if (
                self.gripper_mode in ("eleven", "twenty")
                and not gripper_is_discrete
                and self.gripper_median_window >= 3
            ):
                out[:, column] = _median(raw[:, 0], self.gripper_median_window, nearest=True)
            out[:, column] = discretize_gripper_actions(
                out[:, column : column + 1],
                mode=self.gripper_mode,
                threshold=self.gripper_threshold,
            )[:, 0]
            log["gripper_changed_frames"] = int(np.count_nonzero(out[:, column] != reference[:, 0]))
        log["motion_max_change"] = _motion_change(out, original, columns)
        return out

    def _legacy_apply_recorded(self, values, ease_limits, smooth_columns, gripper_column, log):
        start, stop = _legacy_boundary_trim_indices(
            values, tau=self.tau, min_idle_steps=self.min_idle_steps
        )
        log.update(start=start, stop=stop, version=1)
        out = self._apply_v1(
            values,
            ease_limits=ease_limits,
            smooth_columns=smooth_columns,
            gripper_column=gripper_column,
            stats=log,
        )
        if gripper_column is not None and len(out):
            reference = discretize_gripper_actions(
                values[start:stop, gripper_column : gripper_column + 1]
                if gripper_column >= 0
                else values[start:stop, -1:],
                mode=self.gripper_mode,
                threshold=self.gripper_threshold,
            )
            log["gripper_changed_frames"] = int(
                np.count_nonzero(out[:, gripper_column] != reference[:, 0])
            )
        return out

    def _trim_slice(self, values, gripper_column, active_dims, log):
        start, stop = boundary_trim_indices(
            values,
            tau=self.tau,
            min_idle_steps=self.min_idle_steps,
            context_frames=self.context_frames,
        )
        metrics = trim_deviation(
            values, start, stop, profile=self.profile, gripper_column=gripper_column
        )
        allowed = not metrics["gripper_changed"]
        if self.profile == "joint" and self.motion_limits is not None:
            dims = tuple(range(values.shape[1])) if active_dims is None else active_dims
            local = [
                i
                for i in range(values.shape[1])
                if gripper_column is None or i != gripper_column % values.shape[1]
            ]
            limits = np.asarray(self.motion_limits)[[dims[i] for i in local]]
            allowed &= bool(np.all(np.asarray(metrics["motion_max_error"]) <= limits))
        elif (
            self.profile == "eef"
            and self.translation_limit is not None
            and self.rotation_limit is not None
        ):
            allowed &= metrics["translation_max_error_m"] <= self.translation_limit
            allowed &= metrics["rotation_max_error_rad"] <= self.rotation_limit
        else:
            allowed = False
        log.update(
            version=2,
            candidate_start=start,
            candidate_stop=stop,
            candidate_deviation=metrics,
            trim_approved=bool(allowed),
            trim_mode=self.trim_mode,
        )
        if self.trim_mode != "bounded" or not allowed:
            start, stop = 0, len(values)
        log.update(start=start, stop=stop)
        return start, stop

    def _apply_v1(
        self,
        values: np.ndarray,
        *,
        ease_limits: np.ndarray | None = None,
        smooth_columns: tuple[int, ...] | None = None,
        gripper_column: int | None = None,
        stats: dict[str, int] | None = None,
    ) -> np.ndarray:
        smooth_columns = smooth_columns if smooth_columns is not None else self.smooth_columns
        out = np.asarray(values)
        start, stop = _legacy_boundary_trim_indices(
            out, tau=self.tau, min_idle_steps=self.min_idle_steps
        )
        if stop <= start:
            if stats is not None:
                stats["ease_steps"] = 0
            return out[start:start]
        out = out[start:stop]
        if self.savgol_window >= 3:
            if smooth_columns is None:
                out = savgol_smooth(out, window=self.savgol_window, polyorder=self.savgol_polyorder)
            else:
                out = np.array(out, copy=True)
                columns = tuple(smooth_columns)
                out[:, columns] = savgol_smooth(
                    out[:, columns], window=self.savgol_window, polyorder=self.savgol_polyorder
                )
        if stats is not None:
            stats["ease_steps"] = 0
        if self.ease_horizon >= 1 and ease_limits is not None:
            limit = np.broadcast_to(np.asarray(ease_limits, dtype=np.float64), (out.shape[1],))
            if stats is not None:
                stats["ease_steps"] = _overspeed_steps(out, limit)
            out = ease_overspeed(out, limits=ease_limits, horizon=self.ease_horizon)
        if gripper_column is not None:
            out = np.array(out, copy=True)
            column = gripper_column % out.shape[1]
            if (
                self.gripper_mode in ("eleven", "twenty")
                and self.gripper_median_window >= 3
                and len(out) >= self.gripper_median_window
            ):
                out[:, column] = _median(out[:, column], self.gripper_median_window, nearest=False)
            out[:, column] = discretize_gripper_actions(
                out[:, column : column + 1],
                mode=self.gripper_mode,
                threshold=self.gripper_threshold,
            )[:, 0]
        return out

    def to_dict(self) -> dict[str, object]:
        """Serialize for ``preprocessing.json`` alongside a saved tokenizer."""
        payload = {
            "tau": self.tau,
            "min_idle_steps": self.min_idle_steps,
            "savgol_window": self.savgol_window,
            "savgol_polyorder": self.savgol_polyorder,
            "ease_horizon": self.ease_horizon,
            "smooth_columns": list(self.smooth_columns)
            if self.smooth_columns is not None
            else None,
            "gripper_mode": self.gripper_mode,
            "gripper_threshold": self.gripper_threshold,
            "gripper_median_window": self.gripper_median_window,
        }
        if self.version == 2:
            payload.update(
                version=2,
                trim_mode=self.trim_mode,
                profile=self.profile,
                context_frames=self.context_frames,
                motion_limits=list(self.motion_limits) if self.motion_limits is not None else None,
                translation_limit=self.translation_limit,
                rotation_limit=self.rotation_limit,
            )
        return payload

    @classmethod
    def from_dict(cls, payload: dict[str, object]) -> PreprocessingSpec:
        """Inverse of :meth:`to_dict`; raises on unknown keys."""
        fields = dict(payload)
        smooth = fields.pop("smooth_columns", None)
        version = fields.pop("version", 1)
        limits = fields.pop("motion_limits", None)
        spec = cls(
            tau=float(fields.pop("tau")),
            min_idle_steps=int(fields.pop("min_idle_steps")),
            savgol_window=int(fields.pop("savgol_window")),
            savgol_polyorder=int(fields.pop("savgol_polyorder")),
            ease_horizon=int(fields.pop("ease_horizon")),
            smooth_columns=tuple(int(c) for c in smooth) if smooth is not None else None,
            gripper_mode=str(fields.pop("gripper_mode", "binary")),
            gripper_threshold=float(fields.pop("gripper_threshold", 0.8)),
            gripper_median_window=int(
                fields.pop("gripper_median_window", 5 if version == 1 else 0)
            ),
            version=version,
            trim_mode=fields.pop("trim_mode", "shadow"),
            profile=fields.pop("profile", "joint"),
            context_frames=int(fields.pop("context_frames", 1)),
            motion_limits=None if limits is None else tuple(float(v) for v in limits),
            translation_limit=fields.pop("translation_limit", None),
            rotation_limit=fields.pop("rotation_limit", None),
        )
        if fields:
            raise ValueError(f"unknown preprocessing keys: {sorted(fields)}")
        return spec