calibration/imu_lidar/cli.py

"""Command-line entry point for LiDAR–IMU calibration."""

from __future__ import annotations

import argparse
from pathlib import Path

from .contracts import CalibrationMode, CalibrationRequest, CalibrationStatus, SessionInput
from .pipeline import describe_pipeline, run_calibration


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="LiDAR–IMU extrinsic calibration (V1)")
    subcommands = parser.add_subparsers(dest="command", required=True)

    plan = subcommands.add_parser("plan", help="显示标定阶段，不读取数据")
    plan.add_argument("--vehicle-config", help="车辆配置路径（仅展示，plan 不读取）")
    plan.add_argument(
        "--mode",
        choices=[mode.value for mode in CalibrationMode],
        default=CalibrationMode.ROTATION_ONLY.value,
    )

    run = subcommands.add_parser(
        "run",
        help="执行 V1 标定流水线（可重复 --imu/--lidar/--session-id 做多会话联合）",
    )
    run.add_argument(
        "--session-id",
        action="append",
        default=None,
        help="会话 ID（可重复；与 --imu/--lidar 一一对应）",
    )
    run.add_argument(
        "--imu",
        action="append",
        required=True,
        help="IMU CSV/NPZ 路径（可重复）",
    )
    run.add_argument(
        "--lidar",
        action="append",
        required=True,
        help="LiDAR 会话目录（可重复）",
    )
    run.add_argument("--vehicle-config", required=True, help="车辆配置 YAML")
    run.add_argument("--output", required=True, help="输出目录")
    run.add_argument(
        "--mode",
        choices=[mode.value for mode in CalibrationMode],
        default=CalibrationMode.ROTATION_ONLY.value,
    )
    run.add_argument("--max-iterations", type=int, default=2)
    run.add_argument("--time-offset-search-s", type=float, default=1.0)
    run.add_argument("--min-pair-rotation-deg", type=float, default=3.0)
    run.add_argument("--min-pair-translation-m", type=float, default=0.3)
    return parser


def _build_sessions(args: argparse.Namespace) -> tuple[SessionInput, ...]:
    imus = [Path(p) for p in args.imu]
    lidars = [Path(p) for p in args.lidar]
    if len(imus) != len(lidars):
        raise SystemExit(f"--imu count ({len(imus)}) must match --lidar count ({len(lidars)})")
    if args.session_id is None:
        session_ids = [f"session{i}" for i in range(len(imus))]
    else:
        session_ids = list(args.session_id)
        if len(session_ids) != len(imus):
            raise SystemExit(
                f"--session-id count ({len(session_ids)}) must match --imu/--lidar ({len(imus)})"
            )
    return tuple(
        SessionInput(session_id=sid, imu_source=imu, lidar_source=lidar)
        for sid, imu, lidar in zip(session_ids, imus, lidars)
    )


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "plan":
        request = CalibrationRequest(
            vehicle_config=Path(args.vehicle_config) if args.vehicle_config else None,
            requested_mode=CalibrationMode(args.mode),
        )
        print("LiDAR–IMU calibration stages:")
        print(f"requested mode: {request.requested_mode.value}")
        for index, stage in enumerate(describe_pipeline(request), start=1):
            print(f"{index}. {stage.name}: {stage.responsibility}")
        return 0

    if args.command == "run":
        sessions = _build_sessions(args)
        request = CalibrationRequest(
            vehicle_config=Path(args.vehicle_config),
            sessions=sessions,
            requested_mode=CalibrationMode(args.mode),
            output_directory=Path(args.output),
            max_iterations=args.max_iterations,
            min_pair_rotation_deg=args.min_pair_rotation_deg,
            min_pair_translation_m=args.min_pair_translation_m,
            time_offset_search_s=args.time_offset_search_s,
        )
        result = run_calibration(request)
        print(f"status: {result.status.value}")
        print(f"message: {result.message}")
        if result.time_offset_s is not None:
            print(f"time_offset_s (first session; t_imu = t_lidar + dt): {result.time_offset_s:.6f}")
        joint = (result.details or {}).get("joint") or {}
        if joint:
            print(f"merged_pair_count: {joint.get('merged_pair_count')}")
            print(f"pair_counts_per_session: {joint.get('pair_counts_per_session')}")
        if result.T_IMU_lidar is not None:
            print("T_IMU_lidar:")
            print(result.T_IMU_lidar)
        print(f"report directory: {args.output}")
        return 0 if result.status != CalibrationStatus.BLOCKED else 2

    parser.error(f"unknown command {args.command}")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())