#!/usr/bin/env python3
"""Audit strict RTK--IMU segments for lever-arm excitation and marginal information."""

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 _fit_segments, _fit_summary, _segments
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, (tuple, list)):
        return [_jsonable(item) for item in value]
    return value


def _orientation_spans_deg(segment, rotation_rtk_imu: np.ndarray) -> np.ndarray:
    matrices = [segment.R_WRTK_initial @ rotation_rtk_imu]
    for pre in segment.preintegrations:
        matrices.append(matrices[-1] @ pre.delta_R)
    euler = Rotation.from_matrix(np.asarray(matrices)).as_euler("xyz", degrees=False)
    return np.degrees(np.ptp(np.unwrap(euler, axis=0), axis=0))


def _audit_segment(segment, rotation_rtk_imu: np.ndarray, max_nfev: int) -> dict[str, object]:
    gyro = np.asarray([node.gyro_rad_s for node in segment.nodes])
    gyro_norm = np.linalg.norm(gyro, axis=1)
    best_count = sum(node.source == "BESTNAVA" for node in segment.nodes)
    doppler_count = sum(
        node.source == "BESTNAVA" and node.velocity_enu_m_s is not None
        for node in segment.nodes
    )
    fit, residual, detail = _fit_segments([segment], rotation_rtk_imu, max_nfev=max_nfev)
    summary = _fit_summary(fit, residual, detail)
    item: dict[str, object] = {
        "segment_id": segment.segment_id,
        "session_id": segment.session_id,
        "node_count": len(segment.nodes),
        "duration_s": float(segment.nodes[-1].t_s - segment.nodes[0].t_s),
        "yaw_pitch_roll_span_deg": _orientation_spans_deg(segment, rotation_rtk_imu)[[2, 1, 0]],
        "gyro_rms_deg_s": np.degrees(np.sqrt(np.mean(gyro ** 2, axis=0))),
        "gyro_peak_deg_s": np.degrees(np.max(np.abs(gyro), axis=0)),
        "gyro_norm_rms_deg_s": float(np.degrees(np.sqrt(np.mean(gyro_norm ** 2)))),
        "gyro_norm_peak_deg_s": float(np.degrees(np.max(gyro_norm))),
        "bestnava_count": best_count,
        "doppler_count": doppler_count,
        "fit": None,
    }
    if summary is not None:
        item["fit"] = {
            "optimizer_converged": summary.optimizer_converged,
            "lever_information_singular_values": summary.lever_information_singular_values,
            "lever_information_condition_number": summary.lever_information_condition_number,
            "lever_precision_rank": summary.lever_precision_rank,
            "weakest_lever_direction_I": summary.weakest_lever_direction_I,
            "lever_std_m": summary.l_I_std_m,
        }
        singular = summary.lever_information_singular_values
        item["information_score"] = float(singular[-1]) if summary.lever_precision_rank == 3 else 0.0
    else:
        item["information_score"] = 0.0
    spans = np.asarray(item["yaw_pitch_roll_span_deg"])
    # Short strict runs rarely accumulate a full vehicle turn; retain clearly non-straight motion.
    item["turn_or_slope"] = bool(spans[0] >= 3.0 or abs(spans[1]) >= 0.5 or abs(spans[2]) >= 0.5)
    return item


def _recommended(items: list[dict[str, object]]) -> list[str]:
    candidates = [
        item for item in items
        if item["turn_or_slope"] and int(item["bestnava_count"]) >= 6
        and int(item["doppler_count"]) >= 6
        and item["fit"] is not None
        and int(item["fit"]["lever_precision_rank"]) == 3
    ]
    candidates.sort(key=lambda item: float(item["information_score"]), reverse=True)
    selected: list[str] = []
    per_session: dict[str, int] = {}
    for item in candidates:
        session_id = str(item["session_id"])
        if per_session.get(session_id, 0) >= 2:
            continue
        selected.append(str(item["segment_id"]))
        per_session[session_id] = per_session.get(session_id, 0) + 1
    return selected


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("--sample-period-s", type=float, default=1.0)
    parser.add_argument("--rotation-rpy-deg", nargs=3, type=float,
                        default=[0.4543066225, -0.0026392019, 0.0122384129])
    parser.add_argument("--max-nfev", type=int, default=80)
    args = parser.parse_args(argv)

    sessions = load_unified_sessions(args.manifest)
    rotation = Rotation.from_euler("xyz", args.rotation_rpy_deg, degrees=True).as_matrix()
    items = [
        _audit_segment(segment, rotation, args.max_nfev)
        for segment in _segments(sessions, args.sample_period_s)
    ]
    payload = {
        "session_count": len(sessions),
        "strict_segment_count": len(items),
        "sample_period_s": args.sample_period_s,
        "rotation_rpy_deg": args.rotation_rpy_deg,
        "recommended_segment_ids": _recommended(items),
        "segments": items,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(_jsonable(payload), ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    print(json.dumps({
        "strict_segment_count": len(items),
        "recommended_segment_ids": payload["recommended_segment_ids"],
    }, ensure_ascii=False, indent=2))
    return 0


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