calibration/imu_lidar/pipeline.py

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

from __future__ import annotations

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

import numpy as np

from .contracts import (
    CalibrationMode,
    CalibrationRequest,
    CalibrationResult,
    CalibrationStatus,
    MotionPair,
    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 .motion_pairs_io import build_motion_pairs_payload
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, prior_enabled

# Remap keyframe indices so multi-session Phase-C graphs do not collide.
_SESSION_INDEX_OFFSET = 1_000_000


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"),
    PipelineStage("lidar_motion", "各会话关键帧、可选去畸变与 LiDAR 相对运动"),
    PipelineStage("motion_pairs", "各会话构造运动对，再合并"),
    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,
    R_prior: np.ndarray | None = None,
    prior_sigma_deg: float | None = None,
):
    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,
        R_prior=R_prior,
        prior_sigma_deg=prior_sigma_deg,
    )
    return keyframes, pair_set, handeye


def _translation_prior_from_config(
    vehicle_config: dict[str, Any] | None,
) -> tuple[np.ndarray | None, np.ndarray | float | None]:
    if vehicle_config is None or not prior_enabled(vehicle_config, "translation_prior"):
        return None, None
    init_cfg = vehicle_config.get("initialization") or {}
    tp = init_cfg.get("translation_prior") or {}
    if tp.get("t_IMU_lidar_m") is None:
        return None, None
    return np.asarray(tp["t_IMU_lidar_m"], dtype=float).reshape(3), tp.get("sigma_m", [0.05, 0.05, 0.05])


def _rotation_prior_from_config(
    vehicle_config: dict[str, Any] | None,
) -> tuple[np.ndarray | None, float | None]:
    if vehicle_config is None or not prior_enabled(vehicle_config, "rotation_prior"):
        return None, None
    init_cfg = vehicle_config.get("initialization") or {}
    rp = init_cfg.get("rotation_prior") or {}
    if rp.get("R_IMU_lidar") is None:
        return None, None
    return np.asarray(rp["R_IMU_lidar"], dtype=float).reshape(3, 3), float(rp.get("sigma_deg", 15.0))


def _prepare_session_pairs(
    session: SessionInput,
    request: CalibrationRequest,
    *,
    R_prior: np.ndarray | None = None,
    prior_sigma_deg: float | None = None,
) -> dict[str, Any]:
    """Per-session: audit, δt, keyframes/pairs. No joint extrinsic yet."""

    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", "session_id": session.session_id, "report": asdict(ts)}

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

    if request.fixed_time_offset_s is not None:
        offset = TimeOffsetResult(
            delta_t_s=float(request.fixed_time_offset_s),
            correlation_peak=1.0,
            search_s=0.0,
            notes=(
                f"fixed_time_offset_s={float(request.fixed_time_offset_s):.6f} "
                "(skip |ω| search; intended for host-UTC-bridged sessions)",
            ),
            ok=True,
        )
    else:
        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", "session_id": session.session_id, "report": asdict(offset)}

    coarse_delta_t = float(offset.delta_t_s)
    working_frames = frames
    r_x = np.eye(3) if R_prior is None else np.asarray(R_prior, dtype=float).reshape(3, 3)
    handeye = None
    pair_set = None
    keyframes = None
    pairs_notes: list[str] = []
    pair_count = 0

    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,
            R_prior=R_prior,
            prior_sigma_deg=prior_sigma_deg,
        )
        pairs_notes = list(pair_set.notes)
        pair_count = len(pair_set.pairs)
        if pair_count < 3:
            return {
                "ok": False,
                "stage": "motion_pairs",
                "session_id": session.session_id,
                "iteration": iteration,
                "time_offset": asdict(offset),
                "imu_audit": asdict(imu_report),
                "timestamp_audit": asdict(ts),
                "keyframes": 0 if keyframes is None else len(keyframes.indices),
                "pair_notes": pairs_notes,
                "handeye": asdict(handeye),
            }
        r_x = handeye.R_IMU_lidar

        if not request.enable_signed_time_refine:
            continue

        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)),
                max_shift_s=request.max_signed_refine_shift_s,
            )
            # Also bound total walk away from the original coarse estimate.
            if abs(refined.delta_t_s - coarse_delta_t) > request.max_signed_refine_shift_s:
                refined = TimeOffsetResult(
                    delta_t_s=float(offset.delta_t_s),
                    correlation_peak=refined.correlation_peak,
                    search_s=refined.search_s,
                    notes=tuple(
                        list(refined.notes)
                        + [
                            f"signed refine clamped: |δt-coarse| would exceed "
                            f"{request.max_signed_refine_shift_s:.3f}s"
                        ]
                    ),
                    ok=True,
                )
            delta_shift = abs(refined.delta_t_s - offset.delta_t_s)
            offset = _merge_time_offset(offset, refined)
            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,
                R_prior=R_prior,
                prior_sigma_deg=prior_sigma_deg,
            )
            pairs_notes = list(pair_set.notes)
            pair_count = len(pair_set.pairs)
            if pair_count < 3:
                return {
                    "ok": False,
                    "stage": "motion_pairs",
                    "session_id": session.session_id,
                    "iteration": iteration,
                    "time_offset": asdict(offset),
                    "imu_audit": asdict(imu_report),
                    "timestamp_audit": asdict(ts),
                    "keyframes": 0 if keyframes is None else len(keyframes.indices),
                    "pair_notes": pairs_notes,
                    "handeye": asdict(handeye),
                }
            r_x = handeye.R_IMU_lidar

    assert handeye is not None and pair_set is not None and keyframes is not None
    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])

    return {
        "ok": True,
        "session_id": session.session_id,
        "pairs": tuple(pair_set.pairs),
        "gyro_bias_rad_s": np.asarray(imu_report.gyro_bias_rad_s, dtype=float).reshape(3),
        "gravity_init_m_s2": gravity_init,
        "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": asdict(offset),
        "time_offset_s": float(offset.delta_t_s),
        "keyframes": len(keyframes.indices),
        "pair_count": pair_count,
        "pair_notes": pairs_notes,
        "handeye_local": {
            "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(),
        },
    }


def _remap_pairs_for_joint(prepared: list[dict[str, Any]]) -> list[MotionPair]:
    merged: list[MotionPair] = []
    for index, prep in enumerate(prepared):
        id_offset = (index + 1) * _SESSION_INDEX_OFFSET
        for pair in prep["pairs"]:
            merged.append(
                replace(
                    pair,
                    i=int(pair.i) + id_offset,
                    j=int(pair.j) + id_offset,
                )
            )
    return merged


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

    Multi-session: each session estimates its own δt and builds motion pairs;
    rotation hand-eye and joint SE3 are solved once on the merged pair set.
    """

    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,
            )

    r_prior, prior_sigma_deg = _rotation_prior_from_config(vehicle_config)

    prepared: list[dict[str, Any]] = []
    for session in request.sessions:
        prep = _prepare_session_pairs(
            session,
            request,
            R_prior=r_prior,
            prior_sigma_deg=prior_sigma_deg,
        )
        if not prep.get("ok"):
            return finalize_result(
                status=CalibrationStatus.BLOCKED,
                message=f"blocked at stage {prep.get('stage')} ({prep.get('session_id')})",
                details={"sessions": [prep]},
                output_directory=request.output_directory,
            )
        prepared.append(prep)

    all_pairs = _remap_pairs_for_joint(prepared)
    handeye = solve_rotation_handeye(
        all_pairs,
        R_prior=r_prior,
        prior_sigma_deg=prior_sigma_deg,
    )
    if not handeye.ok:
        return finalize_result(
            status=CalibrationStatus.BLOCKED,
            message="blocked at stage rotation_handeye (joint)",
            details={
                "sessions": [_public_session(p) for p in prepared],
                "joint_handeye": asdict(handeye),
                "merged_pair_count": len(all_pairs),
            },
            output_directory=request.output_directory,
        )

    force_rotation_only = request.requested_mode == CalibrationMode.ROTATION_ONLY
    t_prior, t_prior_sigma = _translation_prior_from_config(vehicle_config)
    gyro_bias = np.mean(np.stack([p["gyro_bias_rad_s"] for p in prepared], axis=0), axis=0)
    gravity_init = np.mean(np.stack([p["gravity_init_m_s2"] for p in prepared], axis=0), axis=0)
    g_n = float(np.linalg.norm(gravity_init))
    if g_n > 1e-6:
        gravity_init = gravity_init * (9.80665 / g_n)

    joint = solve_joint_extrinsic(
        all_pairs,
        handeye.R_IMU_lidar,
        force_rotation_only=force_rotation_only,
        imu=None,
        delta_t_s=0.0,
        gyro_bias_rad_s=gyro_bias,
        gravity_init_m_s2=gravity_init,
        enable_phase_c=not force_rotation_only,
        t_init_m=t_prior,
        t_prior_m=t_prior,
        t_prior_sigma_m=t_prior_sigma,
    )

    session_results = []
    for prep in prepared:
        session_results.append(
            {
                **_public_session(prep),
                "vehicle_config_loaded": vehicle_config is not None,
                "handeye": {
                    "residual_rms_deg": handeye.residual_rms_deg,
                    "residual_median_deg": handeye.residual_median_deg,
                    "pair_count": handeye.pair_count,
                    "ok": handeye.ok,
                    "notes": tuple(list(handeye.notes) + [f"joint over {len(request.sessions)} sessions"]),
                    "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(),
                },
                "translation_accepted": joint.translation_accepted,
                "rotation_ok": handeye.ok and joint.observability.rotation_observable,
            }
        )

    T = np.asarray(joint.T_IMU_lidar, dtype=float)
    # Report per-session δt list; keep first as scalar for backward-compatible field.
    delta_t = float(prepared[0]["time_offset_s"])
    if request.requested_mode == CalibrationMode.FULL_SE3:
        if joint.translation_accepted:
            status = CalibrationStatus.FULL_SE3_ACCEPTED
            message = f"full SE3 accepted (joint {len(prepared)} sessions, {len(all_pairs)} pairs)"
        else:
            status = CalibrationStatus.FULL_SE3_REJECTED
            message = (
                f"rotation accepted jointly ({len(prepared)} sessions); "
                "translation rejected by observability/residual gates"
            )
    else:
        status = CalibrationStatus.ROTATION_ONLY_ACCEPTED
        message = f"rotation-only calibration accepted (joint {len(prepared)} sessions, {len(all_pairs)} pairs)"
        T = T.copy()
        T[:3, 3] = 0.0

    return finalize_result(
        status=status,
        message=message,
        details={
            "sessions": session_results,
            "joint": {
                "session_count": len(prepared),
                "merged_pair_count": len(all_pairs),
                "pair_counts_per_session": {p["session_id"]: p["pair_count"] for p in prepared},
                "time_offset_s_per_session": {p["session_id"]: p["time_offset_s"] for p in prepared},
                "handeye_rms_deg": handeye.residual_rms_deg,
                "translation_accepted": joint.translation_accepted,
            },
        },
        T_IMU_lidar=T,
        time_offset_s=delta_t,
        output_directory=request.output_directory,
        motion_pairs_payload=build_motion_pairs_payload(prepared_sessions=prepared),
    )


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