#!/usr/bin/env python3
"""Run one prior-free engineering base fit on high-excitation qualified segments.

This diagnostic never adds a mechanical lever factor, never profiles the
mechanical reference, and never runs LOO/bootstrap/rotation sensitivity.
"""
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 (
    MIN_SEGMENT_DURATION_S,
    MIN_SEGMENT_NODE_COUNT,
    NUISANCE_DOF_PER_SEGMENT,
    _Segment,
    _fit_segments,
    _fit_summary,
    _height_reference,
    _initial_parameters,
    _marginal_lever_information,
    _nodes,
    _residual,
    _segment_residual_size,
    _world_rtk,
)
from imu_lidar.imu_preintegration import preintegrate_imu
from rtk_imu.rtk_imu_multisource import load_unified_sessions

MANUAL_REFERENCE_L_I_M = np.array([-0.45072, -0.25682, 0.73208], dtype=float)


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 _r0_segment_candidates(sessions, period_s: float) -> list[tuple[object, tuple, str]]:
    """Split R0 first; defer expensive preintegration until a candidate is selected."""
    reference = _height_reference(sessions)
    if reference is None:
        return []
    candidates: list[tuple[object, tuple, str]] = []
    for session in sessions:
        nodes = _nodes(session, reference, period_s)
        start = 0
        qualifying_index = 0
        for end in range(1, len(nodes) + 1):
            if end != len(nodes) and nodes[end].continuity_id == nodes[end - 1].continuity_id:
                continue
            run = tuple(nodes[start:end])
            start = end
            if len(run) < MIN_SEGMENT_NODE_COUNT or run[-1].t_s - run[0].t_s < MIN_SEGMENT_DURATION_S:
                continue
            if not any(node.hpr_factor_valid for node in run):
                continue
            candidates.append((session, run, f"{session.session_id}:{qualifying_index:02d}"))
            qualifying_index += 1
    return candidates


def _build_segment(session, run: tuple, segment_id: str) -> _Segment | None:
    pre = tuple(preintegrate_imu(
        session.imu.t_s, session.imu.gyro_rad_s, session.imu.acc_m_s2, left.t_s, right.t_s,
    ) for left, right in zip(run[:-1], run[1:]))
    if any(item.duration_s <= 0.0 for item in pre):
        return None
    initial_hpr = next((node for node in run if node.hpr_factor_valid), None)
    if initial_hpr is None:
        return None
    return _Segment(segment_id, session.session_id, run, pre, _world_rtk(initial_hpr.baseline_enu))

def _gyro_abs_rotation_deg(session, segment) -> np.ndarray:
    start_s, end_s = segment.nodes[0].t_s, segment.nodes[-1].t_s
    mask = (session.imu.t_s >= start_s) & (session.imu.t_s <= end_s)
    t_s = session.imu.t_s[mask]
    gyro = session.imu.gyro_rad_s[mask]
    if t_s.size < 2:
        return np.zeros(3)
    return np.degrees(np.trapezoid(np.abs(gyro), t_s, axis=0))


def _category_score(category: str, gyro_abs_deg: np.ndarray) -> float:
    if category in {"circle", "left_right"}:
        return float(gyro_abs_deg[2])
    return float(np.hypot(gyro_abs_deg[0], gyro_abs_deg[1]))


def _marginal_for_segment_indices(jacobian: np.ndarray, residual: np.ndarray,
                                  segments, indices: list[int]) -> tuple[np.ndarray, np.ndarray, float, int, np.ndarray]:
    row = 0
    rows: list[np.ndarray] = []
    for index, segment in enumerate(segments):
        count = _segment_residual_size(segment)
        if index in indices:
            rows.append(np.arange(row, row + count))
        row += count
    selected_rows = np.concatenate(rows) if rows else np.empty(0, dtype=int)
    columns = [0, 1, 2]
    for index in indices:
        offset = 3 + NUISANCE_DOF_PER_SEGMENT * index
        columns.extend(range(offset, offset + NUISANCE_DOF_PER_SEGMENT))
    marginal, singular, condition, rank, weakest, _ = _marginal_lever_information(
        jacobian[np.ix_(selected_rows, columns)], residual[selected_rows]
    )
    return marginal, singular, condition, rank, weakest


def main() -> 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("--circle-session", required=True)
    parser.add_argument("--left-right-session", required=True)
    parser.add_argument("--slope-session", required=True)
    parser.add_argument("--top-per-category", type=int, default=10)
    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])
    args = parser.parse_args()
    category_by_session = {
        args.circle_session: "circle",
        args.left_right_session: "left_right",
        args.slope_session: "slope",
    }
    sessions = load_unified_sessions(args.manifest, selected_session_ids=set(category_by_session))
    candidates: dict[str, list[tuple[float, object, tuple, str, np.ndarray]]] = {
        key: [] for key in category_by_session.values()
    }
    all_candidates = _r0_segment_candidates(sessions, args.sample_period_s)
    for session, run, segment_id in all_candidates:
        category = category_by_session[session.session_id]
        gyro_abs = _gyro_abs_rotation_deg(session, type("Run", (), {"nodes": run})())
        candidates[category].append((_category_score(category, gyro_abs), session, run, segment_id, gyro_abs))
    selected: list[object] = []
    selected_entries: list[dict[str, object]] = []
    category_indices: dict[str, list[int]] = {}
    for category, entries in candidates.items():
        selected_before = len(selected)
        for score, session, run, segment_id, gyro_abs in sorted(entries, key=lambda item: item[0], reverse=True):
            segment = _build_segment(session, run, segment_id)
            if segment is None:
                continue
            selected.append(segment)
            selected_entries.append({
                "category": category, "segment_id": segment.segment_id,
                "session_id": segment.session_id, "duration_s": segment.nodes[-1].t_s - segment.nodes[0].t_s,
                "node_count": len(segment.nodes), "excitation_score_deg": score,
                "cumulative_absolute_gyro_rotation_xyz_deg": gyro_abs,
            })
            if len(selected) - selected_before >= args.top_per_category:
                break
        category_indices[category] = list(range(selected_before, len(selected)))
    rotation = Rotation.from_euler("xyz", args.rotation_rpy_deg, degrees=True).as_matrix()
    initial_parameters = _initial_parameters(selected)
    initial_residual = _residual(initial_parameters, selected, rotation)
    fit, residual, detail = _fit_segments(selected, rotation)
    summary = _fit_summary(fit, residual, detail)
    if fit is None or summary is None:
        raise RuntimeError("no selected qualified segment could be fit")
    contributions: dict[str, dict[str, object]] = {}
    contribution_sum = np.zeros((3, 3))
    for category, indices in category_indices.items():
        marginal, singular, condition, rank, weakest = _marginal_for_segment_indices(
            fit.jac, residual, selected, indices
        )
        contribution_sum += marginal
        contributions[category] = {
            "selected_segment_count": len(indices),
            "lever_marginal_information": marginal,
            "lever_information_singular_values": singular,
            "condition_number": condition,
            "precision_rank": rank,
            "weakest_direction_I": weakest,
            "axis_information_diagonal_I": np.diag(marginal),
        }
    total_diag = np.diag(summary.lever_marginal_information)
    for category, item in contributions.items():
        item["axis_information_fraction_of_total_I"] = np.divide(
            item["axis_information_diagonal_I"], total_diag,
            out=np.full(3, np.nan), where=np.abs(total_diag) > 1e-12,
        )
    payload = {
        'solver_diagnostics': {
            'initial_cost': 0.5 * float(np.dot(initial_residual, initial_residual)),
            'final_cost': 0.5 * float(np.dot(residual, residual)),
            'cost_reduction': 0.5 * float(
                np.dot(initial_residual, initial_residual) - np.dot(residual, residual)
            ),
            'cost_definition': '0.5 * unmodified residual squared norm, comparable initial/final',
            'scipy_final_huber_cost': float(fit.cost),
            'nfev': int(fit.nfev), 'optimality': float(fit.optimality),
            'gradient_norm': float(np.linalg.norm(fit.grad)),
            'initial_l_I_m': initial_parameters[:3], 'final_l_I_m': fit.x[:3],
            'l_step_norm_m': float(np.linalg.norm(fit.x[:3] - initial_parameters[:3])),
        },
        "scope": "prior-free free base fit only; no mechanical factor/LOO/bootstrap/rotation sensitivity",
        "rotation_source": "R2G_gravity_level_prior",
        "translation_conditional_on_rotation": True,
        "manual_reference_comparison_only": {
            "manual_l_I_m": MANUAL_REFERENCE_L_I_M,
            "free_minus_manual_l_I_m": summary.l_I_m - MANUAL_REFERENCE_L_I_M,
            "euclidean_delta_m": float(np.linalg.norm(summary.l_I_m - MANUAL_REFERENCE_L_I_M)),
        },
        "sample_period_s": args.sample_period_s,
        "all_qualified_segment_count": len(all_candidates),
        "selected_segment_count": len(selected),
        "selected_segments": selected_entries,
        "free_solution": _jsonable(summary),
        "motion_category_information_contributions": _jsonable(contributions),
        "axis_information_total_I": np.diag(summary.lever_marginal_information),
        "marginal_additivity_error_fro": float(np.linalg.norm(
            contribution_sum - summary.lever_marginal_information, ord="fro"
        )),
    }
    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({
        "selected_segment_count": len(selected),
        "free_l_I_m": _jsonable(summary.l_I_m),
        "free_l_I_std_m": _jsonable(summary.l_I_std_m),
        "lever_information_singular_values": _jsonable(summary.lever_information_singular_values),
        "condition_number": summary.lever_information_condition_number,
        "precision_rank": summary.lever_precision_rank,
        "bestnava_xyz_vector_rms_p95_m": [summary.bestnava_xyz_residual.vector_rms, summary.bestnava_xyz_residual.vector_p95],
        "doppler_vector_rms_p95_m_s": [summary.doppler_velocity_residual.vector_rms, summary.doppler_velocity_residual.vector_p95],
    }, ensure_ascii=False, indent=2))
    return 0


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