calibration/imu_lidar/contracts.py

"""Shared contracts for the LiDAR–IMU calibration pipeline."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

import numpy as np


class TransformConvention(str, Enum):
    """The only transform convention used by this project."""

    T_A_B = "T_A_B maps points from frame B into frame A"


class CalibrationMode(str, Enum):
    ROTATION_ONLY = "rotation_only"
    FULL_SE3 = "full_se3"


class CalibrationStatus(str, Enum):
    NOT_RUN = "not_run"
    BLOCKED = "blocked"
    ROTATION_ONLY_ACCEPTED = "rotation_only_accepted"
    ROTATION_ONLY_PRIOR_CONSTRAINED = "rotation_only_prior_constrained"
    FULL_SE3_ACCEPTED = "full_se3_accepted"
    FULL_SE3_REJECTED = "full_se3_rejected_due_to_observability"


@dataclass(frozen=True)
class SessionInput:
    """Input paths for one independently recorded session."""

    session_id: str
    imu_source: Path
    lidar_source: Path
    board_configuration_id: str | None = None
    # Optional session-local override. The request-level value remains a
    # backward-compatible fallback for batches whose timelines are all aligned.
    fixed_time_offset_s: float | None = None


@dataclass(frozen=True)
class CalibrationRequest:
    """Top-level calibration request."""

    vehicle_config: Path | None
    sessions: tuple[SessionInput, ...] = ()
    requested_mode: CalibrationMode = CalibrationMode.ROTATION_ONLY
    output_directory: Path | None = None
    max_iterations: int = 2
    min_pair_rotation_deg: float = 3.0
    min_pair_translation_m: float = 0.3
    min_registration_fitness: float = 0.5
    max_imu_gap_s: float = 0.05
    max_lidar_gap_s: float = 1.0
    time_offset_search_s: float = 1.0
    # If set, skip |ω| search and use this constant (host-UTC-bridged sessions: 0).
    fixed_time_offset_s: float | None = None
    # Signed 3-axis refine after hand-eye; disable for already-bridged timelines.
    enable_signed_time_refine: bool = True
    # Reject signed refine steps that walk farther than this from the coarse δt.
    max_signed_refine_shift_s: float = 0.05


@dataclass
class CalibrationResult:
    """Result envelope written by finalize after pipeline gates."""

    status: CalibrationStatus = CalibrationStatus.NOT_RUN
    message: str = "Calibration has not been executed."
    details: dict[str, Any] = field(default_factory=dict)
    T_IMU_lidar: np.ndarray | None = None
    time_offset_s: float | None = None


@dataclass(frozen=True)
class ImuSeries:
    """Normalized IMU samples.

    ``t_s`` is the native IMU clock in seconds (need not match LiDAR epoch).
    Gyro must be rad/s; accelerometer must be m/s^2.
    """

    t_s: np.ndarray
    gyro_rad_s: np.ndarray
    acc_m_s2: np.ndarray

    def __post_init__(self) -> None:
        object.__setattr__(self, "t_s", np.asarray(self.t_s, dtype=float).reshape(-1))
        object.__setattr__(self, "gyro_rad_s", np.asarray(self.gyro_rad_s, dtype=float).reshape(-1, 3))
        object.__setattr__(self, "acc_m_s2", np.asarray(self.acc_m_s2, dtype=float).reshape(-1, 3))
        n = self.t_s.size
        if self.gyro_rad_s.shape != (n, 3) or self.acc_m_s2.shape != (n, 3):
            raise ValueError("IMU arrays must share the same length and have shape (N, 3)")


@dataclass(frozen=True)
class LidarFrame:
    """One LiDAR sweep in Cartesian sensor coordinates."""

    frame_id: str
    t_start_s: float
    t_end_s: float
    points_xyz: np.ndarray
    path: Path | None = None

    @property
    def t_mid_s(self) -> float:
        return 0.5 * (self.t_start_s + self.t_end_s)


@dataclass(frozen=True)
class MotionPair:
    """One relative-motion observation between keyframes i and j."""

    session_id: str
    i: int
    j: int
    t_i_s: float
    t_j_s: float
    R_A: np.ndarray
    R_B: np.ndarray
    t_A_m: np.ndarray | None = None
    t_B_m: np.ndarray | None = None
    fitness: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)