#!/usr/bin/env python3
"""Run the conditional engineering RTK--IMU 6DoF lever-arm branch."""

from __future__ import annotations

import argparse
import json
import math
import sys
from dataclasses import asdict
from pathlib import Path

import numpy as np
from scipy.spatial.transform import Rotation

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

from rtk_imu.rtk_imu_engineering import (
    DEFAULT_MANUAL_L_I_M,
    DEFAULT_MANUAL_L_I_COVARIANCE_M2,
    solve_engineering_6dof,
)
from rtk_imu.rtk_imu_multisource import load_unified_sessions


def _jsonable(value):
    if isinstance(value, np.ndarray):
        return _jsonable(value.tolist())
    if isinstance(value, np.generic):
        return _jsonable(value.item())
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if hasattr(value, "__dataclass_fields__"):
        return {key: _jsonable(item) for key, item in asdict(value).items()}
    if isinstance(value, dict):
        return {str(key): _jsonable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_jsonable(item) for item in value]
    return value


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--session", action="append", required=True,
        help="Dynamic/static sessions intentionally admitted to this engineering run.",
    )
    parser.add_argument(
        "--rotation-rpy-deg", nargs=3, type=float,
        default=[0.4543066225, -0.0026392019, 0.0122384129],
    )
    parser.add_argument(
        "--manual-l-i-m", nargs=3, type=float, default=DEFAULT_MANUAL_L_I_M.tolist(),
        help="Mechanical ANT1 phase-centre lever arm p_ANT1^I in metres.",
    )
    prior_group = parser.add_mutually_exclusive_group()
    prior_group.add_argument(
        "--manual-l-i-std-m", nargs=3, type=float,
        default=np.sqrt(np.diag(DEFAULT_MANUAL_L_I_COVARIANCE_M2)).tolist(),
        help="Mechanical 1-sigma prior standard deviation for lx, ly, lz in metres.",
    )
    prior_group.add_argument(
        "--manual-l-i-covariance-m2", nargs=9, type=float,
        help="Row-major 3x3 mechanical prior covariance in m^2.",
    )
    parser.add_argument(
        "--no-manual-prior", action="store_true",
        help="Run a free solve only; retain the mechanical reference in no optimisation factor.",
    )
    parser.add_argument(
        "--segment-id", action="append",
        help="Strict-continuity segment id from the excitation audit; may be repeated.",
    )
    parser.add_argument("--sample-period-s", type=float, default=0.5)
    parser.add_argument("--skip-loo", action="store_true")
    parser.add_argument("--run-bootstrap", action="store_true")
    parser.add_argument("--bootstrap-repetitions", type=int, default=40)
    parser.add_argument("--bootstrap-seed", type=int, default=0)
    parser.add_argument("--run-rotation-sensitivity", action="store_true")
    args = parser.parse_args(argv)

    manual_l_i_m = None if args.no_manual_prior else args.manual_l_i_m
    prior_covariance = None
    if not args.no_manual_prior:
        if args.manual_l_i_std_m is not None:
            prior_covariance = np.diag(np.square(args.manual_l_i_std_m))
        elif args.manual_l_i_covariance_m2 is not None:
            prior_covariance = np.asarray(args.manual_l_i_covariance_m2, dtype=float).reshape(3, 3)
        else:
            parser.error("mechanical prior requires --manual-l-i-std-m or --manual-l-i-covariance-m2")

    sessions = load_unified_sessions(args.manifest, selected_session_ids=set(args.session))
    rotation = Rotation.from_euler("xyz", args.rotation_rpy_deg, degrees=True).as_matrix()
    result = solve_engineering_6dof(
        sessions,
        R_RTK_IMU=rotation,
        manual_l_I_m=manual_l_i_m,
        manual_l_I_covariance_m2=prior_covariance,
        sample_period_s=args.sample_period_s,
        run_loo=not args.skip_loo,
        run_bootstrap=args.run_bootstrap,
        bootstrap_repetitions=args.bootstrap_repetitions,
        bootstrap_seed=args.bootstrap_seed,
        run_rotation_sensitivity=args.run_rotation_sensitivity,
        selected_segment_ids=None if args.segment_id is None else set(args.segment_id),
    )
    payload = {
        "data_only_6dof_accepted": False,
        "data_only_translation_accepted": result.data_only_translation_accepted,
        "engineering_6dof": _jsonable(result),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8"
    )
    print(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False))
    return 0


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