#!/usr/bin/env python3
"""Run the independent RTK--IMU calibration against the project inventory."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from rtk_imu.rtk_imu_replay import (
    DEFAULT_RTK_FRAME_DEFINITION,
    DEFAULT_RTK_REFERENCE_POINT,
    load_inventory,
    load_sessions,
    run_calibration,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--inventory",
        type=Path,
        default=ROOT / "artifacts" / "rtk_imu_inventory_v1" / "rtk_session_inventory.csv",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=ROOT / "artifacts" / "rtk_imu_calibration_v1",
    )
    parser.add_argument("--session", action="append", help="session id to include; repeatable")
    parser.add_argument("--batch", action="append", help="batch id to include; repeatable")
    parser.add_argument("--rotation-only", action="store_true")
    parser.add_argument("--no-loo", action="store_true")
    parser.add_argument("--knot-step-s", type=float, default=2.0)
    parser.add_argument("--per-batch", action="store_true")
    parser.add_argument("--rtk-frame-definition", default=DEFAULT_RTK_FRAME_DEFINITION)
    parser.add_argument("--rtk-reference-point", default=DEFAULT_RTK_REFERENCE_POINT)
    args = parser.parse_args(argv)

    entries = load_inventory(args.inventory)
    if args.session:
        selected = set(args.session)
        entries = [entry for entry in entries if entry.session_id in selected]
    if args.batch:
        selected_batches = set(args.batch)
        entries = [entry for entry in entries if entry.batch_id in selected_batches]
    if not entries:
        raise SystemExit("no inventory rows match the requested selection")
    sessions = load_sessions(entries)
    rotation, translation = run_calibration(
        sessions,
        args.output_dir,
        rotation_only=args.rotation_only,
        compute_loo=not args.no_loo,
        knot_step_s=args.knot_step_s,
        rtk_frame_definition=args.rtk_frame_definition,
        rtk_reference_point=args.rtk_reference_point,
    )
    print(
        json.dumps(
            {
                "output": str(args.output_dir.resolve()),
                "sessions": [session.session_id for session in sessions],
                "rotation_ok": rotation.ok,
                "rotation_rpy_deg": rotation.rpy_deg.tolist(),
                "rotation_rms_deg": rotation.residual_rms_deg,
                "translation_ok": None if translation is None else translation.ok,
                "translation_m": None if translation is None else translation.t_RTK_IMU_m.tolist(),
            },
            ensure_ascii=False,
            indent=2,
        )
    )

    if args.per_batch:
        for batch in sorted({entry.batch_id for entry in entries}):
            batch_entries = [entry for entry in entries if entry.batch_id == batch]
            if len(batch_entries) < 2:
                continue
            run_calibration(
                load_sessions(batch_entries),
                args.output_dir / "per_batch" / batch,
                rotation_only=args.rotation_only,
                compute_loo=False,
                knot_step_s=args.knot_step_s,
                rtk_frame_definition=args.rtk_frame_definition,
                rtk_reference_point=args.rtk_reference_point,
            )
    return 0


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