src/fpgm/geometry/velocity.py

"""Differentiate lifted 3D tracks into linear velocity, with a robust object aggregate.

Deliberately out of scope: RANSAC-based inlier selection and Kabsch/Umeyama rigid-body
fitting for angular velocity. Both are deferred, not forgotten -- a component-wise
median with MAD trimming is a simpler, cheaper robust estimator that is adequate for
"is this object moving and how fast", and adding a 6-DoF rigid-body fit (which needs a
minimum point count and is far more sensitive to a bad correspondence) is a
meaningfully larger step that should be justified by an actual downstream need for
angular velocity, not bundled in here speculatively.
"""

from __future__ import annotations

import numpy as np
from scipy.signal import savgol_filter

from fpgm.config import VelocityConfig
from fpgm.types import Track3D, VelocityEstimate

_MAD_TO_STD = 1.4826  # scale factor making MAD a consistent estimator of std under normality


class VelocityEstimator:
    """Turns a :class:`~fpgm.types.Track3D` into a :class:`~fpgm.types.VelocityEstimate`."""

    def estimate(self, track3d: Track3D, cfg: VelocityConfig) -> VelocityEstimate:
        """Differentiate ``track3d`` per point, then robustly aggregate to one object velocity.

        Per point: low-confidence samples are dropped, small gaps are bridged by
        linearly interpolating *position* before differentiating, and gaps longer
        than ``cfg.max_gap_frames`` split the point's track into independent valid
        runs with no velocity fabricated across the gap -- exactly where the object
        may have changed direction (e.g. a grasp-and-lift), so interpolating through
        it would assert something we have no evidence for.

        Args:
            track3d: Lifted per-point 3D positions with per-sample confidence.
            cfg: Differentiation and aggregation tunables.

        Returns:
            A :class:`~fpgm.types.VelocityEstimate` with ``frame = "world"``.
        """
        n_frames, n_points = track3d.valid.shape
        times = np.asarray(track3d.timestamps, dtype=np.float64)

        linear_velocity = np.full((n_frames, n_points, 3), np.nan, dtype=np.float64)
        sample_confidence = np.zeros((n_frames, n_points), dtype=np.float32)

        base_valid = track3d.valid & (track3d.confidence >= cfg.min_confidence)

        # A Python loop over points is unavoidable here (not a vectorisation lapse):
        # each point's valid/gap structure is ragged -- different points have gaps
        # at different frames of different lengths -- so there is no shared axis to
        # batch the run-splitting logic over. Each run's actual differentiation is
        # vectorised across (frames, xyz).
        for q in range(n_points):
            valid_q = base_valid[:, q]
            if not np.any(valid_q):
                continue
            positions_q = track3d.xyz_world[:, q, :]
            confidence_q = track3d.confidence[:, q].astype(np.float64)
            for idx, pos, conf in _bridge_and_split_runs(
                times, valid_q, positions_q, confidence_q, cfg.max_gap_frames
            ):
                if idx.size < 2:
                    continue
                vel = _differentiate(pos, times[idx], cfg)
                linear_velocity[idx, q, :] = vel
                sample_confidence[idx, q] = conf.astype(np.float32)

        object_linear_velocity = np.full((n_frames, 3), np.nan, dtype=np.float64)
        object_speed = np.full(n_frames, np.nan, dtype=np.float64)
        object_velocity_confidence = np.zeros(n_frames, dtype=np.float32)
        inlier_mask = np.zeros((n_frames, n_points), dtype=bool)

        for t in range(n_frames):
            available = ~np.isnan(linear_velocity[t, :, 0])
            if not np.any(available):
                continue
            avail_idx = np.flatnonzero(available)
            vel_t = linear_velocity[t, avail_idx, :]
            inlier_local = _mad_trim(vel_t, cfg.mad_trim_factor)

            object_linear_velocity[t] = np.median(vel_t[inlier_local], axis=0)
            object_speed[t] = float(np.linalg.norm(object_linear_velocity[t]))
            inlier_mask[t, avail_idx[inlier_local]] = True
            mean_inlier_conf = float(np.mean(sample_confidence[t, avail_idx[inlier_local]]))
            object_velocity_confidence[t] = (inlier_local.sum() / n_points) * mean_inlier_conf

        return VelocityEstimate(
            point_id=track3d.point_id,
            timestamps=times,
            linear_velocity=linear_velocity,
            object_linear_velocity=object_linear_velocity,
            object_speed=object_speed,
            object_velocity_confidence=object_velocity_confidence,
            inlier_mask=inlier_mask,
            frame="world",
        )


def _mad_trim(vel: np.ndarray, mad_trim_factor: float) -> np.ndarray:
    """Return a boolean inlier mask over ``vel`` (n, 3) via MAD-trimmed distance from the median.

    Distance from the component-wise median is used (not per-component MAD) so a
    single point that is an outlier in any direction is caught, not just axis-aligned
    outliers.
    """
    n = vel.shape[0]
    if n <= 2:
        return np.ones(n, dtype=bool)
    median = np.median(vel, axis=0)
    dist = np.linalg.norm(vel - median, axis=1)
    mad = float(np.median(dist))
    if mad < 1e-12:
        return np.ones(n, dtype=bool)
    inlier = dist <= mad_trim_factor * _MAD_TO_STD * mad
    if not np.any(inlier):
        return np.ones(n, dtype=bool)  # degenerate: trimming would drop everything
    return inlier


def _bridge_and_split_runs(
    times: np.ndarray,
    valid: np.ndarray,
    positions: np.ndarray,
    confidence: np.ndarray,
    max_gap_frames: int,
) -> list[tuple[np.ndarray, np.ndarray, np.ndarray]]:
    """Split one point's timeline into runs, bridging short gaps by linear interpolation.

    Returns a list of ``(frame_indices, positions, confidence)`` where each run's
    ``frame_indices`` is contiguous (every frame between its first and last valid
    sample is present, with short-gap frames filled in). Gaps longer than
    ``max_gap_frames`` are never bridged -- they end one run and start the next.
    """
    valid_idx = np.flatnonzero(valid)
    if valid_idx.size == 0:
        return []

    boundaries = [int(valid_idx[0])]
    prev = valid_idx[0]
    for cur in valid_idx[1:]:
        if int(cur) - int(prev) - 1 > max_gap_frames:
            boundaries.append(int(prev))
            boundaries.append(int(cur))
        prev = cur
    boundaries.append(int(valid_idx[-1]))

    runs: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = []
    for start, end in zip(boundaries[0::2], boundaries[1::2]):
        idx = np.arange(start, end + 1)
        pos = positions[idx].copy()
        conf = confidence[idx].copy()
        local_valid = valid[idx]
        missing = ~local_valid
        if np.any(missing):
            run_times = times[idx]
            known_times = run_times[local_valid]
            for dim in range(3):
                pos[missing, dim] = np.interp(
                    run_times[missing], known_times, pos[local_valid, dim]
                )
            # Bridged-sample confidence is linearly interpolated between the two
            # anchoring real detections, consistent with the position bridging.
            conf[missing] = np.interp(run_times[missing], known_times, conf[local_valid])
        runs.append((idx, pos, conf))
    return runs


def _differentiate(positions: np.ndarray, times: np.ndarray, cfg: VelocityConfig) -> np.ndarray:
    """Differentiate one contiguous run's ``(n, 3)`` positions w.r.t. ``times`` (n,).

    Uses a Savitzky-Golay derivative when the run is long enough and uniformly
    sampled; otherwise falls back to central differences (``np.gradient``), which is
    also what correctly handles non-uniform timestamps -- assuming a uniform
    ``delta`` for a Savitzky-Golay filter on non-uniform samples would silently
    misweight the fit.
    """
    n = positions.shape[0]
    dt_diffs = np.diff(times)
    uniform = bool(np.allclose(dt_diffs, dt_diffs[0], rtol=1e-3, atol=1e-9))
    long_enough = n >= cfg.savgol_window and cfg.savgol_window > cfg.savgol_polyorder
    if cfg.method == "savgol" and uniform and long_enough:
        dt = float(np.mean(dt_diffs))
        return savgol_filter(
            positions,
            window_length=cfg.savgol_window,
            polyorder=cfg.savgol_polyorder,
            deriv=1,
            delta=dt,
            axis=0,
            mode="interp",
        )
    return np.gradient(positions, times, axis=0)