calibration/imu_lidar/pipeline.py

"""Executable LiDAR–IMU calibration pipeline (V1)."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np

from .contracts import (
    CalibrationMode,
    CalibrationRequest,
    CalibrationResult,
    CalibrationStatus,
    SessionInput,
)
from .finalize import finalize_result
from .imu_audit import audit_imu
from .imu_io import load_imu_samples
from .joint_optimizer import solve_joint_extrinsic
from .keyframes import build_keyframes
from .lidar_deskew import deskew_lidar_frames
from .lidar_io import load_lidar_frames
from .motion_pairs import build_motion_pairs
from .rotation_handeye import solve_rotation_handeye
from .time_offset import TimeOffsetResult, estimate_time_offset, refine_time_offset_signed
from .timestamp_audit import audit_timestamps
from .vehicle_config import load_vehicle_config


def _merge_time_offset(previous: TimeOffsetResult, refined: TimeOffsetResult) -> TimeOffsetResult:
    return TimeOffsetResult(
        delta_t_s=refined.delta_t_s,
        correlation_peak=refined.correlation_peak,
        search_s=previous.search_s,
        notes=tuple(list(previous.notes) + list(refined.notes)),
        ok=True,
    )


@dataclass(frozen=True)
class PipelineStage:
    name: str
    responsibility: str


STAGES = (
    PipelineStage("vehicle_config", "加载并校验当前车辆安装配置"),
    PipelineStage("timestamp_audit", "审查 IMU 与 LiDAR 时间域"),
    PipelineStage("imu_audit", "审查单位、轴向启发与静止零偏"),
    PipelineStage("time_offset", "粗估 δt，并用 R 做有符号三轴精修"),
    PipelineStage("lidar_motion", "关键帧、可选去畸变与 LiDAR 相对运动"),
    PipelineStage("motion_pairs", "IMU 预积分与雷达配准，构造相对运动对"),
    PipelineStage("rotation_handeye", "加权求解旋转外参"),
    PipelineStage("joint_optimizer", "联合精修；完整模式下可估计平移"),
    PipelineStage("finalize", "写出结果与质量报告"),
)


def describe_pipeline(_: CalibrationRequest) -> tuple[PipelineStage, ...]:
    """Return the planned stages."""

    return STAGES


def _build_pairs_and_handeye(
    *,
    session_id: str,
    working_frames,
    imu,
    delta_t_s: float,
    gyro_bias_rad_s: np.ndarray,
    request: CalibrationRequest,
):
    keyframes = build_keyframes(
        working_frames,
        min_translation_m=request.min_pair_translation_m,
        min_rotation_deg=request.min_pair_rotation_deg,
    )
    pair_set = build_motion_pairs(
        session_id=session_id,
        keyframes=list(keyframes.frames),
        keyframe_indices=keyframes.indices,
        imu=imu,
        delta_t_s=delta_t_s,
        gyro_bias_rad_s=gyro_bias_rad_s,
        min_rotation_deg=request.min_pair_rotation_deg,
        min_translation_m=request.min_pair_translation_m,
    )
    handeye = solve_rotation_handeye(pair_set.pairs)
    return keyframes, pair_set, handeye


def _session_details(
    session: SessionInput,
    request: CalibrationRequest,
    vehicle_config: dict[str, Any] | None,
) -> dict[str, Any]:
    imu = load_imu_samples(session.imu_source)
    frames = load_lidar_frames(session.lidar_source)

    ts = audit_timestamps(imu, frames)
    if not ts.ok:
        return {"ok": False, "stage": "timestamp_audit", "report": asdict(ts)}

    imu_report = audit_imu(imu)
    if not imu_report.ok:
        return {"ok": False, "stage": "imu_audit", "report": asdict(imu_report)}

    offset = estimate_time_offset(
        imu,
        frames,
        gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
        search_s=request.time_offset_search_s,
    )
    if not offset.ok:
        return {"ok": False, "stage": "time_offset", "report": asdict(offset)}

    working_frames = frames
    r_x = np.eye(3)
    handeye = None
    pair_set = None
    keyframes = None
    pairs_notes: list[str] = []
    pair_count = 0
    time_offset_notes = list(offset.notes)

    for iteration in range(max(1, request.max_iterations)):
        if iteration > 0:
            working_frames = deskew_lidar_frames(
                frames,
                imu,
                delta_t_s=offset.delta_t_s,
                R_IMU_lidar=r_x,
                gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
            )
        keyframes, pair_set, handeye = _build_pairs_and_handeye(
            session_id=session.session_id,
            working_frames=working_frames,
            imu=imu,
            delta_t_s=offset.delta_t_s,
            gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
            request=request,
        )
        pairs_notes = list(pair_set.notes)
        pair_count = len(pair_set.pairs)
        if handeye.pair_count < 3:
            return {
                "ok": False,
                "stage": "rotation_handeye",
                "iteration": iteration,
                "time_offset": asdict(offset),
                "imu_audit": asdict(imu_report),
                "timestamp_audit": asdict(ts),
                "keyframes": len(keyframes.indices),
                "pair_notes": pairs_notes,
                "handeye": asdict(handeye),
            }
        # Use candidate R even if RMS gate failed, so signed δt refine can still run.
        r_x = handeye.R_IMU_lidar

        # Phase-A: alternate signed δt refine with current R (up to 2 rounds).
        for _ in range(2):
            refined = refine_time_offset_signed(
                imu,
                frames,
                delta_t_s=offset.delta_t_s,
                R_IMU_lidar=r_x,
                gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
                search_s=min(0.12, max(0.04, 0.25 * request.time_offset_search_s)),
            )
            delta_shift = abs(refined.delta_t_s - offset.delta_t_s)
            offset = _merge_time_offset(offset, refined)
            time_offset_notes = list(offset.notes)
            if delta_shift < 1e-3:
                break
            keyframes, pair_set, handeye = _build_pairs_and_handeye(
                session_id=session.session_id,
                working_frames=working_frames,
                imu=imu,
                delta_t_s=offset.delta_t_s,
                gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
                request=request,
            )
            pairs_notes = list(pair_set.notes)
            pair_count = len(pair_set.pairs)
            if handeye.pair_count < 3:
                return {
                    "ok": False,
                    "stage": "rotation_handeye",
                    "iteration": iteration,
                    "time_offset": asdict(offset),
                    "imu_audit": asdict(imu_report),
                    "timestamp_audit": asdict(ts),
                    "keyframes": len(keyframes.indices),
                    "pair_notes": pairs_notes,
                    "handeye": asdict(handeye),
                }
            r_x = handeye.R_IMU_lidar

        if not handeye.ok:
            return {
                "ok": False,
                "stage": "rotation_handeye",
                "iteration": iteration,
                "time_offset": asdict(offset),
                "imu_audit": asdict(imu_report),
                "timestamp_audit": asdict(ts),
                "keyframes": len(keyframes.indices),
                "pair_notes": pairs_notes,
                "handeye": asdict(handeye),
            }

    assert handeye is not None and pair_set is not None and keyframes is not None
    force_rotation_only = request.requested_mode == CalibrationMode.ROTATION_ONLY
    # Specific force opposing measured specific force ≈ −g in the static IMU frame.
    acc_mean = np.asarray(imu_report.static_acc_mean_m_s2, dtype=float).reshape(3)
    acc_n = float(np.linalg.norm(acc_mean))
    if acc_n > 1e-6:
        gravity_init = -acc_mean * (9.80665 / acc_n)
    else:
        gravity_init = np.array([0.0, 0.0, -9.80665])
    joint = solve_joint_extrinsic(
        pair_set.pairs,
        r_x,
        force_rotation_only=force_rotation_only,
        imu=imu,
        delta_t_s=offset.delta_t_s,
        gyro_bias_rad_s=imu_report.gyro_bias_rad_s,
        gravity_init_m_s2=gravity_init,
        enable_phase_c=not force_rotation_only,
    )

    offset_payload = asdict(offset)

    return {
        "ok": True,
        "session_id": session.session_id,
        "vehicle_config_loaded": vehicle_config is not None,
        "timestamp_audit": asdict(ts),
        "imu_audit": {
            **asdict(imu_report),
            "gyro_bias_rad_s": imu_report.gyro_bias_rad_s.tolist(),
            "static_acc_mean_m_s2": imu_report.static_acc_mean_m_s2.tolist(),
        },
        "time_offset": offset_payload,
        "keyframes": len(keyframes.indices),
        "pair_count": pair_count,
        "pair_notes": pairs_notes,
        "handeye": {
            "residual_rms_deg": handeye.residual_rms_deg,
            "residual_median_deg": handeye.residual_median_deg,
            "pair_count": handeye.pair_count,
            "ok": handeye.ok,
            "notes": handeye.notes,
            "R_IMU_lidar": handeye.R_IMU_lidar.tolist(),
        },
        "joint": {
            "translation_accepted": joint.translation_accepted,
            "residual_rms_rot_deg": joint.residual_rms_rot_deg,
            "residual_rms_trans_m": joint.residual_rms_trans_m,
            "observability": asdict(joint.observability),
            "notes": joint.notes,
            "T_IMU_lidar": joint.T_IMU_lidar.tolist(),
            "gyro_bias_rad_s": None
            if joint.gyro_bias_rad_s is None
            else np.asarray(joint.gyro_bias_rad_s, dtype=float).tolist(),
            "accel_bias_m_s2": None
            if joint.accel_bias_m_s2 is None
            else np.asarray(joint.accel_bias_m_s2, dtype=float).tolist(),
            "gravity_m_s2": None
            if joint.gravity_m_s2 is None
            else np.asarray(joint.gravity_m_s2, dtype=float).tolist(),
        },
        "T_IMU_lidar": joint.T_IMU_lidar,
        "time_offset_s": offset.delta_t_s,
        "translation_accepted": joint.translation_accepted,
        "rotation_ok": handeye.ok and joint.observability.rotation_observable,
    }


def run_calibration(request: CalibrationRequest) -> CalibrationResult:
    """Run the V1 calibration pipeline for one or more sessions."""

    if not request.sessions:
        return finalize_result(
            status=CalibrationStatus.BLOCKED,
            message="no sessions provided",
            details={},
            output_directory=request.output_directory,
        )

    vehicle_config = None
    if request.vehicle_config is not None:
        try:
            vehicle_config = load_vehicle_config(request.vehicle_config)
        except Exception as exc:  # noqa: BLE001 - surface config problems as blocked
            return finalize_result(
                status=CalibrationStatus.BLOCKED,
                message=f"vehicle config failed: {exc}",
                details={},
                output_directory=request.output_directory,
            )

    session_results = []
    for session in request.sessions:
        session_results.append(_session_details(session, request, vehicle_config))

    primary = session_results[0]
    if not primary.get("ok"):
        return finalize_result(
            status=CalibrationStatus.BLOCKED,
            message=f"blocked at stage {primary.get('stage')}",
            details={"sessions": session_results},
            output_directory=request.output_directory,
        )

    T = np.asarray(primary["T_IMU_lidar"], dtype=float)
    delta_t = float(primary["time_offset_s"])
    if request.requested_mode == CalibrationMode.FULL_SE3:
        if primary.get("translation_accepted"):
            status = CalibrationStatus.FULL_SE3_ACCEPTED
            message = "full SE3 accepted"
        else:
            status = CalibrationStatus.FULL_SE3_REJECTED
            message = "rotation accepted; translation rejected by observability/residual gates"
    else:
        status = CalibrationStatus.ROTATION_ONLY_ACCEPTED
        message = "rotation-only calibration accepted"
        T = T.copy()
        T[:3, 3] = 0.0

    return finalize_result(
        status=status,
        message=message,
        details={"sessions": [_public_session(s) for s in session_results]},
        T_IMU_lidar=T,
        time_offset_s=delta_t,
        output_directory=request.output_directory,
    )


def _public_session(session_result: dict[str, Any]) -> dict[str, Any]:
    payload = dict(session_result)
    payload.pop("T_IMU_lidar", None)
    return payload