#!/usr/bin/env python3
"""Read-only artificial GNHPR dropout validation for the engineering R0 bridge.

For each requested gap, Q4 HPR samples inside an otherwise 0.1 s contiguous
run are withheld.  Their true baseline directions are compared with normalized
linear interpolation and short-horizon HI13 gyro propagation.  No lever-arm
solve, prior, bootstrap, sensitivity run, or threshold update is performed.
"""
from __future__ import annotations

import argparse
import json
import math
import sys
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 _all_hpr, _baseline_angle_deg, _interpolate_baseline, _world_rtk
from rtk_imu.rtk_imu_multisource import load_unified_sessions

GAPS_S = (0.2, 0.3, 0.4, 0.5, 0.6, 0.8)
R2G_RPY_DEG = (0.4543066225, -0.0026392019, 0.0122384129)
BASELINE_FACTOR_SIGMA_DEG = float(np.degrees(0.035))


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 isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (tuple, list)):
        return [_jsonable(v) for v in value]
    return value


def _contiguous_runs(times: np.ndarray, valid: np.ndarray) -> list[np.ndarray]:
    indices = np.flatnonzero(valid)
    if indices.size == 0:
        return []
    runs: list[list[int]] = [[int(indices[0])]]
    for previous, index in zip(indices[:-1], indices[1:]):
        dt = float(times[index] - times[previous])
        if 0.075 <= dt <= 0.125:
            runs[-1].append(int(index))
        else:
            runs.append([int(index)])
    return [np.asarray(run, dtype=int) for run in runs if len(run) >= 3]


def _propagate_baseline(session, R_WI_start: np.ndarray, baseline_I: np.ndarray,
                        start_s: float, targets_s: np.ndarray) -> list[np.ndarray]:
    """Propagate with vectorized gyro interpolation over the short withheld gap."""
    targets = np.asarray(targets_s, dtype=float)
    if targets.size == 0:
        return []
    grid = np.unique(np.concatenate(([start_s], session.imu.t_s[
        (session.imu.t_s > start_s) & (session.imu.t_s < targets[-1])], targets)))
    gyro = np.column_stack([
        np.interp(grid, session.imu.t_s, session.imu.gyro_rad_s[:, axis]) for axis in range(3)
    ])
    R_WI = np.asarray(R_WI_start, dtype=float).copy()
    predicted: list[np.ndarray] = []
    target_index = 0
    for index, (left, right) in enumerate(zip(grid[:-1], grid[1:])):
        omega = 0.5 * (gyro[index] + gyro[index + 1])
        R_WI = R_WI @ Rotation.from_rotvec(omega * float(right - left)).as_matrix()
        while target_index < targets.size and abs(targets[target_index] - right) < 1e-9:
            predicted.append(R_WI @ baseline_I)
            target_index += 1
    return predicted

def _session_validation(session, rotation: np.ndarray, max_windows: int) -> dict[str, object]:
    hpr = _all_hpr(session)
    runs = _contiguous_runs(hpr.t_s, hpr.valid)
    baseline_I = rotation.T[:, 0]
    result: dict[str, object] = {"session_id": session.session_id, "q4_runs": len(runs), "gaps": {}}
    for requested_gap in GAPS_S:
        interpolation_errors: list[float] = []
        propagation_errors: list[float] = []
        actual_gaps: list[float] = []
        windows = 0
        for run in runs:
            nominal_dt = float(np.median(np.diff(hpr.t_s[run])))
            step_count = max(2, int(round(requested_gap / nominal_dt)))
            # The endpoints remain observed, every internal Q4 point is withheld.
            for start in range(0, run.size - step_count, max(1, step_count)):
                subset = run[start:start + step_count + 1]
                if subset.size != step_count + 1:
                    continue
                left, right = int(subset[0]), int(subset[-1])
                actual_gap = float(hpr.t_s[right] - hpr.t_s[left])
                if abs(actual_gap - requested_gap) > 0.08:
                    continue
                withheld = subset[1:-1]
                if withheld.size == 0:
                    continue
                left_b, right_b = hpr.baseline_enu[left], hpr.baseline_enu[right]
                targets = hpr.t_s[withheld]
                predicted_prop = _propagate_baseline(
                    session, _world_rtk(left_b) @ rotation, baseline_I, float(hpr.t_s[left]), targets
                )
                for index, predicted in zip(withheld, predicted_prop):
                    fraction = float((hpr.t_s[index] - hpr.t_s[left]) / actual_gap)
                    interpolation_errors.append(_baseline_angle_deg(
                        _interpolate_baseline(left_b, right_b, fraction), hpr.baseline_enu[index]
                    ))
                    propagation_errors.append(_baseline_angle_deg(predicted, hpr.baseline_enu[index]))
                actual_gaps.append(actual_gap)
                windows += 1
                if windows >= max_windows:
                    break
            if windows >= max_windows:
                break
        def summary(errors: list[float]) -> dict[str, float | int | None]:
            finite = np.asarray([e for e in errors if np.isfinite(e)], dtype=float)
            if finite.size == 0:
                return {"count": 0, "p50_deg": None, "p95_deg": None, "max_deg": None}
            return {"count": int(finite.size), "p50_deg": float(np.percentile(finite, 50)),
                    "p95_deg": float(np.percentile(finite, 95)), "max_deg": float(np.max(finite))}
        interp, prop = summary(interpolation_errors), summary(propagation_errors)
        choice = "linear_baseline_interpolation" if (interp["p95_deg"] or np.inf) <= (prop["p95_deg"] or np.inf) else "imu_gyro_propagation"
        selected_p95 = interp["p95_deg"] if choice.startswith("linear") else prop["p95_deg"]
        result["gaps"][f"{requested_gap:.1f}"] = {
            "requested_gap_s": requested_gap,
            "actual_gap_p50_s": float(np.median(actual_gaps)) if actual_gaps else None,
            "window_count": windows,
            "linear_baseline_interpolation": interp,
            "imu_gyro_propagation": prop,
            "preferred_method_by_p95": choice,
            "selected_p95_deg": selected_p95,
            "within_baseline_factor_sigma": bool(selected_p95 is not None and selected_p95 <= BASELINE_FACTOR_SIGMA_DEG),
        }
    return result


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("--session", action="append")
    parser.add_argument("--max-windows", type=int, default=500)
    args = parser.parse_args()
    sessions = load_unified_sessions(args.manifest, selected_session_ids=None if args.session is None else set(args.session))
    rotation = Rotation.from_euler("xyz", R2G_RPY_DEG, degrees=True).as_matrix()
    session_results = [_session_validation(session, rotation, args.max_windows) for session in sessions]
    recommendation: dict[str, object] = {}
    for gap in GAPS_S:
        entries = [item["gaps"][f"{gap:.1f}"] for item in session_results if item["gaps"].get(f"{gap:.1f}")]
        p95 = [entry["selected_p95_deg"] for entry in entries if entry["selected_p95_deg"] is not None]
        recommendation[f"{gap:.1f}"] = {
            "sessions_with_samples": len(p95),
            "worst_session_selected_p95_deg": float(max(p95)) if p95 else None,
            "passes_all_sessions_factor_sigma": bool(p95 and max(p95) <= BASELINE_FACTOR_SIGMA_DEG),
        }
    passing = [float(key) for key, value in recommendation.items() if value["passes_all_sessions_factor_sigma"]]
    payload = {
        "scope": "artificial HPR dropout validation only; no solve/prior/bootstrap/sensitivity/threshold change",
        "r2g_rotation_rpy_deg": R2G_RPY_DEG,
        "baseline_factor_sigma_deg": BASELINE_FACTOR_SIGMA_DEG,
        "sessions": session_results,
        "bridge_recommendation": recommendation,
        "largest_gap_passing_all_sessions_s": max(passing) if passing else None,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(_jsonable(payload), ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"largest_gap_passing_all_sessions_s": payload["largest_gap_passing_all_sessions_s"],
                      "bridge_recommendation": recommendation}, ensure_ascii=False, indent=2))
    return 0


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