src/fpgm/geometry/validation.py

"""Validate lifted velocity against scene_flows ground truth.

The trick that makes this a fair test rather than tautological: the ground-truth
velocity is produced by feeding scene_flows positions through the exact same
:class:`~fpgm.geometry.velocity.VelocityEstimator` used for our lifted track (same
gap-bridging, same differentiation, same MAD-trimmed aggregation). Any discrepancy is
therefore attributable to the lifting stage (camera convention, depth interpolation),
not to a difference in the velocity algorithm itself.

Magnitude error and direction (cosine) error are reported *separately* because they
point at different bugs: a magnitude-only error implies a depth-scale problem, while a
direction error implies a convention/axis-sign problem -- collapsing them into one
number would hide which stage to go fix.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from fpgm.config import VelocityConfig
from fpgm.geometry.camera import Camera
from fpgm.geometry.velocity import VelocityEstimator
from fpgm.types import ClipTiming, Track2D, Track3D, VelocityEstimate

_EPS = 1e-9


@dataclass
class ValidationReport:
    """Per-frame and aggregate agreement between our lifted velocity and ground truth."""

    n_frames: int
    n_gt_points: int
    n_compared_frames: int
    magnitude_rel_error: np.ndarray  # (T,) float, NaN where not comparable
    direction_cosine_error: np.ndarray  # (T,) float in [0, 2], NaN where not comparable
    per_axis_bias: np.ndarray  # (3,) mean signed (ours - gt) velocity error, m/s
    per_axis_noise_std: np.ndarray  # (3,) std of the signed error, m/s
    median_magnitude_rel_error: float
    median_direction_cosine_error: float
    frac_frames_under_threshold: float
    rel_error_threshold: float
    reprojection_error_px: np.ndarray  # (T,) float, NaN where not comparable
    median_reprojection_error_px: float

    def summary(self) -> str:
        """Render a readable table of the validation results."""
        lines = [
            "Scene-flow velocity validation",
            "===============================",
            f"frames total / compared        : {self.n_frames} / {self.n_compared_frames}",
            f"ground-truth points             : {self.n_gt_points}",
            "",
            "metric                          value",
            "-------------------------------  --------",
            f"{'median magnitude rel. error':<32}  {self.median_magnitude_rel_error:.4f}",
            f"{'median direction cosine error':<32}  {self.median_direction_cosine_error:.4f}",
            f"{'frac frames < ' + f'{self.rel_error_threshold:g}':<32}  "
            f"{self.frac_frames_under_threshold:.4f}",
            f"{'median reprojection error (px)':<32}  {self.median_reprojection_error_px:.4f}",
            "",
            "per-axis bias (ours - gt, m/s)  x         y         z",
            "-------------------------------  --------  --------  --------",
            (
                f"{'systematic (mean)':<32}  "
                f"{self.per_axis_bias[0]:8.4f}  {self.per_axis_bias[1]:8.4f}  "
                f"{self.per_axis_bias[2]:8.4f}"
            ),
            (
                f"{'noise (std)':<32}  "
                f"{self.per_axis_noise_std[0]:8.4f}  {self.per_axis_noise_std[1]:8.4f}  "
                f"{self.per_axis_noise_std[2]:8.4f}"
            ),
        ]
        return "\n".join(lines)


def _cosine_error(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """1 - cosine-similarity between rows of ``a`` and ``b`` (both (T, 3)); NaN for zero vectors."""
    na = np.linalg.norm(a, axis=-1)
    nb = np.linalg.norm(b, axis=-1)
    denom = na * nb
    cos_sim = np.full(a.shape[0], np.nan, dtype=np.float64)
    ok = denom > _EPS
    cos_sim[ok] = np.sum(a[ok] * b[ok], axis=-1) / denom[ok]
    return 1.0 - np.clip(cos_sim, -1.0, 1.0)


def validate_against_scene_flow(
    scene_flows: np.ndarray,
    scene_visibility: np.ndarray,
    scene_depth_valid: np.ndarray,
    object_point_indices: np.ndarray,
    camera: Camera,
    frame_is_world: bool,
    timing: ClipTiming,
    track2d: Track2D,
    lifted_velocity: VelocityEstimate,
    velocity_cfg: VelocityConfig,
    rel_error_threshold: float = 0.2,
) -> ValidationReport:
    """Compare our lifted-track velocity against a ground-truth velocity from scene_flows.

    Args:
        scene_flows: ``(T, N, 3)`` raw tracked 3D points from the h5.
        scene_visibility: ``(T, N)`` bool.
        scene_depth_valid: ``(T, N)`` bool.
        object_point_indices: ``(K,)`` indices into ``N`` selecting the object's own
            annotated scene-flow points (e.g. those inside the object mask).
        camera: Camera at the scene-flow annotation resolution.
        frame_is_world: Resolved scene-flow frame convention (see
            :mod:`fpgm.geometry.convention`).
        timing: Clip timing, used to build ground-truth timestamps.
        track2d: Our tracked 2D query points (clip-frame indexed, e.g. via
            :func:`fpgm.geometry.temporal.resample_track_to_clip_frames`), used only
            for the reprojection sanity check.
        lifted_velocity: The velocity estimate produced by our lifting pipeline
            (:func:`fpgm.geometry.lifting.lift_track_to_3d` followed by
            :class:`~fpgm.geometry.velocity.VelocityEstimator`).
        velocity_cfg: Passed to the *same* :class:`VelocityEstimator` used to build
            ``lifted_velocity``, so the ground-truth pass is an apples-to-apples
            comparison.
        rel_error_threshold: Relative magnitude-error threshold for the "fraction of
            good frames" summary statistic.

    Returns:
        A :class:`ValidationReport`.
    """
    object_point_indices = np.asarray(object_point_indices, dtype=np.int64)
    n_frames = scene_flows.shape[0]
    gt_points = scene_flows[:, object_point_indices, :].astype(np.float64)

    if frame_is_world:
        gt_world_xyz = gt_points
    else:
        gt_world_xyz = camera.cam_to_world(gt_points)

    gt_valid = scene_visibility[:, object_point_indices] & scene_depth_valid[:, object_point_indices]
    gt_confidence = gt_valid.astype(np.float32)

    frame_idx = np.arange(n_frames, dtype=np.int64)
    gt_timestamps = np.asarray(timing.clip_frame_to_seconds(frame_idx), dtype=np.float64)

    gt_track3d = Track3D(
        point_id=object_point_indices.astype(np.int32),
        timestamps=gt_timestamps,
        xyz_world=gt_world_xyz,
        valid=gt_valid,
        confidence=gt_confidence,
    )
    v_gt = VelocityEstimator().estimate(gt_track3d, velocity_cfg)

    magnitude_rel_error = np.full(n_frames, np.nan, dtype=np.float64)
    direction_cosine_error = np.full(n_frames, np.nan, dtype=np.float64)
    signed_error = np.full((n_frames, 3), np.nan, dtype=np.float64)

    t_common = min(n_frames, lifted_velocity.object_speed.shape[0])
    both_valid = np.zeros(n_frames, dtype=bool)
    both_valid[:t_common] = ~np.isnan(v_gt.object_speed[:t_common]) & ~np.isnan(
        lifted_velocity.object_speed[:t_common]
    )

    ours_speed = np.full(n_frames, np.nan)
    ours_speed[:t_common] = lifted_velocity.object_speed[:t_common]
    ours_vel = np.full((n_frames, 3), np.nan)
    ours_vel[:t_common] = lifted_velocity.object_linear_velocity[:t_common]

    gt_speed_safe = np.where(v_gt.object_speed > _EPS, v_gt.object_speed, _EPS)
    magnitude_rel_error[both_valid] = (
        np.abs(ours_speed[both_valid] - v_gt.object_speed[both_valid]) / gt_speed_safe[both_valid]
    )
    direction_cosine_error[both_valid] = _cosine_error(
        ours_vel[both_valid], v_gt.object_linear_velocity[both_valid]
    )
    signed_error[both_valid] = ours_vel[both_valid] - v_gt.object_linear_velocity[both_valid]

    n_compared = int(np.count_nonzero(both_valid))
    if n_compared > 0:
        per_axis_bias = np.nanmean(signed_error, axis=0)
        per_axis_noise_std = np.nanstd(signed_error, axis=0)
        median_mag = float(np.nanmedian(magnitude_rel_error))
        median_dir = float(np.nanmedian(direction_cosine_error))
        frac_under = float(np.mean(magnitude_rel_error[both_valid] < rel_error_threshold))
    else:
        per_axis_bias = np.full(3, np.nan)
        per_axis_noise_std = np.full(3, np.nan)
        median_mag = float("nan")
        median_dir = float("nan")
        frac_under = float("nan")

    reprojection_error_px = _reprojection_check(
        gt_world_xyz, gt_valid, camera, track2d, n_frames
    )
    finite_reproj = reprojection_error_px[np.isfinite(reprojection_error_px)]
    median_reproj = float(np.median(finite_reproj)) if finite_reproj.size else float("nan")

    return ValidationReport(
        n_frames=n_frames,
        n_gt_points=int(object_point_indices.size),
        n_compared_frames=n_compared,
        magnitude_rel_error=magnitude_rel_error,
        direction_cosine_error=direction_cosine_error,
        per_axis_bias=per_axis_bias,
        per_axis_noise_std=per_axis_noise_std,
        median_magnitude_rel_error=median_mag,
        median_direction_cosine_error=median_dir,
        frac_frames_under_threshold=frac_under,
        rel_error_threshold=rel_error_threshold,
        reprojection_error_px=reprojection_error_px,
        median_reprojection_error_px=median_reproj,
    )


def _reprojection_check(
    gt_world_xyz: np.ndarray,
    gt_valid: np.ndarray,
    camera: Camera,
    track2d: Track2D,
    n_frames: int,
) -> np.ndarray:
    """Differentiation-free check: does reprojected GT land near our tracked query points?

    Compares, per frame, the centroid of the reprojected ground-truth object points
    against the centroid of our own visible tracked query points. A large centroid
    offset means the camera/convention is wrong at the projection level, independent
    of anything to do with differentiation or depth interpolation -- so this check
    should be examined *before* trusting any velocity numbers at all.
    """
    cam = camera.rescaled(*track2d.resolution)
    frame_to_row = {int(f): i for i, f in enumerate(track2d.frames)}
    out = np.full(n_frames, np.nan, dtype=np.float64)
    for t in range(n_frames):
        valid_t = gt_valid[t]
        if not np.any(valid_t):
            continue
        row = frame_to_row.get(t)
        if row is None:
            continue
        query_visible = track2d.visible[row]
        if not np.any(query_visible):
            continue
        uv_gt, depth_gt = cam.project(gt_world_xyz[t, valid_t])
        in_front = depth_gt > 0
        if not np.any(in_front):
            continue
        gt_centroid = np.mean(uv_gt[in_front], axis=0)
        query_centroid = np.mean(track2d.uv[row, query_visible], axis=0)
        out[t] = float(np.linalg.norm(gt_centroid - query_centroid))
    return out


@dataclass
class GripperValidationReport:
    """Comparison of estimated object velocity against robot proprioception.

    While the object is grasped, it moves rigidly with the end-effector, so the
    gripper's own velocity is a strong reference for the object's. Crucially it is
    *independent* of everything the pipeline does: it comes from joint encoders,
    not from the camera, the masks, the tracks or the scene-flow annotations.
    """

    n_frames: int
    n_compared: int
    grasped_only: bool
    median_abs_error_mps: float
    median_rel_error: float
    frac_within_50pct: float
    ours_peak_mps: float
    gripper_peak_mps: float
    ours_mean_mps: float
    gripper_mean_mps: float

    def summary(self) -> str:
        return (
            "Gripper-proprioception validation\n"
            "=================================\n"
            f"frames compared            : {self.n_compared}/{self.n_frames}"
            f"{' (grasped only)' if self.grasped_only else ''}\n"
            f"peak speed  ours / gripper : {self.ours_peak_mps:.3f} / "
            f"{self.gripper_peak_mps:.3f} m/s\n"
            f"mean speed  ours / gripper : {self.ours_mean_mps:.3f} / "
            f"{self.gripper_mean_mps:.3f} m/s\n"
            f"median absolute error      : {self.median_abs_error_mps:.3f} m/s\n"
            f"median relative error      : {self.median_rel_error:.2f}\n"
            f"frames within 50% of truth : {self.frac_within_50pct:.0%}"
        )


def gripper_speed(gripper_pose: np.ndarray, fps: float) -> np.ndarray:
    """End-effector speed per frame (m/s) from ``gripper_pose`` positions.

    Returns an array of length ``T`` (the first sample repeats the second) so it
    aligns with the clip's frame axis.
    """
    pos = np.asarray(gripper_pose, dtype=np.float64)[:, :3]
    step = np.linalg.norm(np.diff(pos, axis=0), axis=1) * float(fps)
    if step.size == 0:
        return np.zeros(len(pos))
    return np.concatenate([step[:1], step])


def validate_against_gripper(
    velocity: VelocityEstimate,
    gripper_pose: np.ndarray,
    fps: float,
    gripper_open: np.ndarray | None = None,
    grasped_only: bool = True,
    object_xyz_world: np.ndarray | None = None,
    max_grasp_distance_m: float = 0.12,
) -> GripperValidationReport:
    """Score an estimated object velocity against the end-effector's own velocity.

    This is deliberately offered alongside the scene-flow check rather than instead
    of it, because the two fail differently: scene-flow validation is unavailable
    exactly where annotations are missing (which, on this dataset, is often the
    manipulated object itself), while proprioception is always present but is only
    a valid reference while the object is actually held.

    Args:
        velocity: pipeline output for the clip.
        gripper_pose: ``(T, 7)`` end-effector pose; only the position is used.
        fps: trajectory (control) rate, used to differentiate the pose.
        gripper_open: ``(T, 1)`` or ``(T,)`` boolean; frames where the gripper is
            open are excluded when ``grasped_only`` is set, since a released object
            no longer follows the end-effector.
        grasped_only: restrict the comparison to grasped frames.
    """
    ours = np.asarray(velocity.object_speed, dtype=np.float64)
    truth = gripper_speed(gripper_pose, fps)

    n = min(len(ours), len(truth))
    ours, truth = ours[:n], truth[:n]

    usable = np.isfinite(ours) & np.isfinite(truth)
    if grasped_only:
        # A closed gripper does NOT imply it is holding *this* object -- on real
        # episodes the hand closes and travels while the object still sits on the
        # table, and comparing those frames manufactures enormous fake errors.
        # Require the gripper to be closed AND co-located with the object.
        if gripper_open is not None:
            closed = ~np.asarray(gripper_open).reshape(-1)[:n].astype(bool)
            usable &= closed
        if object_xyz_world is not None:
            obj = np.asarray(object_xyz_world, dtype=np.float64)[:n]
            hand = np.asarray(gripper_pose, dtype=np.float64)[:n, :3]
            with np.errstate(invalid="ignore"):
                near = np.linalg.norm(obj - hand, axis=1) <= max_grasp_distance_m
            usable &= np.nan_to_num(near, nan=False).astype(bool)

    if not usable.any():
        return GripperValidationReport(
            n_frames=n,
            n_compared=0,
            grasped_only=grasped_only,
            median_abs_error_mps=float("nan"),
            median_rel_error=float("nan"),
            frac_within_50pct=float("nan"),
            ours_peak_mps=float("nan"),
            gripper_peak_mps=float(truth.max()) if len(truth) else float("nan"),
            ours_mean_mps=float("nan"),
            gripper_mean_mps=float(truth.mean()) if len(truth) else float("nan"),
        )

    a, b = ours[usable], truth[usable]
    abs_err = np.abs(a - b)
    rel_err = abs_err / np.maximum(b, 1e-3)
    return GripperValidationReport(
        n_frames=n,
        n_compared=int(usable.sum()),
        grasped_only=grasped_only,
        median_abs_error_mps=float(np.median(abs_err)),
        median_rel_error=float(np.median(rel_err)),
        frac_within_50pct=float(np.mean(rel_err < 0.5)),
        ours_peak_mps=float(a.max()),
        gripper_peak_mps=float(b.max()),
        ours_mean_mps=float(a.mean()),
        gripper_mean_mps=float(b.mean()),
    )