"""Cached Phase-A replay: rehydrate gyro factors, compare variants, write reports."""

from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path
from typing import Any

import numpy as np

from .imu_io import load_imu_samples
from .motion_pairs_io import (
    build_motion_pairs_payload,
    load_motion_pairs,
    pair_from_dict,
    save_motion_pairs,
)
from .phase_a import (
    ProgressCallback,
    phase_a_comparison_to_dict,
    phase_a_metadata_complete,
    rehydrate_phase_a_pairs,
    solve_phase_a_comparison,
)
from .vehicle_config import load_vehicle_config, prior_enabled


def _rotation_prior(
    vehicle_config_path: Path,
) -> tuple[np.ndarray | None, float]:
    config = load_vehicle_config(vehicle_config_path)
    if not prior_enabled(config, "rotation_prior"):
        return None, 15.0
    prior = (config.get("initialization") or {}).get("rotation_prior") or {}
    matrix = prior.get("R_IMU_lidar")
    if matrix is None:
        return None, float(prior.get("sigma_deg", 15.0))
    return (
        np.asarray(matrix, dtype=float).reshape(3, 3),
        float(prior.get("sigma_deg", 15.0)),
    )


def _sanitize_json(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): _sanitize_json(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_sanitize_json(item) for item in value]
    if isinstance(value, np.ndarray):
        return _sanitize_json(value.tolist())
    if isinstance(value, (np.floating, float)):
        number = float(value)
        return number if np.isfinite(number) else None
    if isinstance(value, (np.integer, np.bool_)):
        return value.item()
    return value


def _write_json(path: Path, payload: Any) -> None:
    path.write_text(
        json.dumps(_sanitize_json(payload), indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def _load_cached_sessions(
    motion_pairs_path: Path,
) -> tuple[
    dict[str, Any],
    list,
    dict[str, np.ndarray],
    dict[str, float],
]:
    payload = load_motion_pairs(motion_pairs_path)
    pairs = []
    biases: dict[str, np.ndarray] = {}
    offsets: dict[str, float] = {}
    for session in payload.get("sessions") or []:
        session_id = str(session["session_id"])
        biases[session_id] = np.asarray(
            session.get("gyro_bias_rad_s", np.zeros(3)),
            dtype=float,
        ).reshape(3)
        offsets[session_id] = float(session.get("delta_t_s", 0.0))
        pairs.extend(
            pair_from_dict(item)
            for item in session.get("pairs") or []
        )
    if not pairs:
        raise ValueError(f"motion-pair cache is empty: {motion_pairs_path}")
    return payload, pairs, biases, offsets


def run_phase_a_replay(
    *,
    motion_pairs_path: Path,
    vehicle_config_path: Path,
    output_directory: Path,
    imu_paths_by_session: dict[str, Path] | None = None,
    excluded_sessions: set[str] | None = None,
    strong_rotation_min_deg: float = 1.0,
    decorrelation_block_s: float = 3.0,
    max_pairs_per_block: int = 1,
    bias_prior_sigma_rad_s: float = 0.002,
    yaw_std_max_deg: float = 0.5,
    leave_one_out_yaw_range_max_deg: float = 1.0,
    data_prior_difference_max_deg: float = 1.0,
    max_nfev: int = 200,
    progress_callback: ProgressCallback | None = None,
) -> dict[str, Any]:
    """Run Phase-A only. Existing LiDAR relative motions are never recomputed."""

    output_directory.mkdir(parents=True, exist_ok=True)
    source_payload, pairs, bias0, offsets = _load_cached_sessions(
        motion_pairs_path
    )
    session_ids = sorted(bias0)
    if progress_callback is not None:
        progress_callback(
            "cache_loaded",
            {
                "schema_version": source_payload.get("schema_version"),
                "sessions": len(session_ids),
                "pairs": len(pairs),
            },
        )

    rehydration_report: dict[str, Any] = {
        "required": not phase_a_metadata_complete(pairs),
        "pair_count": len(pairs),
    }
    if not phase_a_metadata_complete(pairs):
        supplied_paths = {} if imu_paths_by_session is None else imu_paths_by_session
        missing = [sid for sid in session_ids if sid not in supplied_paths]
        if missing:
            raise ValueError(
                "v1 cache lacks J_bg/cov; provide --session-imu for: "
                + ", ".join(missing)
            )
        imu_by_session = {
            sid: load_imu_samples(supplied_paths[sid])
            for sid in session_ids
        }
        pairs, details = rehydrate_phase_a_pairs(
            pairs,
            imu_by_session=imu_by_session,
            bias0_by_session=bias0,
            progress_callback=progress_callback,
        )
        rehydration_report.update(details)
        if float(details["max_R_A_error_deg"]) > 0.05:
            raise ValueError(
                "rehydrated IMU rotations do not match cached R_A: "
                f"max error={details['max_R_A_error_deg']:.6f} deg; "
                "check session-to-IMU path mapping"
            )

    grouped: dict[str, list] = defaultdict(list)
    for pair in pairs:
        grouped[pair.session_id].append(pair)
    enriched_payload = build_motion_pairs_payload(
        prepared_sessions=[
            {
                "session_id": sid,
                "time_offset_s": offsets[sid],
                "gyro_bias_rad_s": bias0[sid],
                "pairs": tuple(grouped[sid]),
            }
            for sid in session_ids
        ]
    )
    enriched_cache_path = save_motion_pairs(
        output_directory / "motion_pairs_phase_a_v2.json",
        enriched_payload,
    )

    rotation_prior, rotation_prior_sigma_deg = _rotation_prior(
        vehicle_config_path
    )
    comparison = solve_phase_a_comparison(
        pairs,
        gyro_bias_rad_s_by_session=bias0,
        rotation_prior=rotation_prior,
        rotation_prior_sigma_deg=rotation_prior_sigma_deg,
        preexcluded_session_ids=excluded_sessions,
        strong_rotation_min_deg=strong_rotation_min_deg,
        decorrelation_block_s=decorrelation_block_s,
        max_pairs_per_block=max_pairs_per_block,
        bias_prior_sigma_rad_s=bias_prior_sigma_rad_s,
        yaw_std_max_deg=yaw_std_max_deg,
        leave_one_out_yaw_range_max_deg=(
            leave_one_out_yaw_range_max_deg
        ),
        data_prior_difference_max_deg=data_prior_difference_max_deg,
        run_leave_one_out=True,
        max_nfev=max_nfev,
        progress_callback=progress_callback,
    )
    full = phase_a_comparison_to_dict(comparison)
    full["input"] = {
        "motion_pairs": str(motion_pairs_path),
        "source_schema_version": source_payload.get("schema_version"),
        "vehicle_config": str(vehicle_config_path),
        "session_imu_paths": {
            sid: str(path)
            for sid, path in (imu_paths_by_session or {}).items()
        },
        "excluded_sessions": sorted(excluded_sessions or set()),
    }
    full["rehydration"] = rehydration_report
    full["enriched_cache"] = str(enriched_cache_path)
    full["parameters"] = {
        "strong_rotation_min_deg": strong_rotation_min_deg,
        "decorrelation_block_s": decorrelation_block_s,
        "max_pairs_per_block": max_pairs_per_block,
        "bias_prior_sigma_rad_s": bias_prior_sigma_rad_s,
        "rotation_prior_sigma_deg": rotation_prior_sigma_deg,
        "yaw_std_max_deg": yaw_std_max_deg,
        "leave_one_out_yaw_range_max_deg": (
            leave_one_out_yaw_range_max_deg
        ),
        "data_prior_difference_max_deg": (
            data_prior_difference_max_deg
        ),
        "max_nfev": max_nfev,
    }

    variants = full["variants"]
    summary = {
        "status": comparison.solution_status,
        "accepted": comparison.accepted,
        "partial_accepted": comparison.partial_accepted,
        "acceptance_checks": comparison.acceptance_checks,
        "primary_result": comparison.recommended_result,
        "variants": {
            name: {
                "rpy_deg_xyz": item["rpy_deg_xyz"],
                "R_IMU_lidar": item["R_IMU_lidar"],
                "residual_rms_deg": item["residual_rms_deg"],
                "residual_p95_deg": item["residual_p95_deg"],
                "accepted": item["accepted"],
                "gyro_bias_rad_s_per_session": item[
                    "gyro_bias_rad_s_per_session"
                ],
            }
            for name, item in variants.items()
            if item is not None
        },
        "marginal_observability_A1": full[
            "marginal_observability_A1"
        ],
        "data_vs_prior_yaw_diff_deg": (
            comparison.data_vs_prior_yaw_diff_deg
        ),
        "data_vs_prior_geodesic_deg": (
            comparison.data_vs_prior_geodesic_deg
        ),
        "leave_one_out_yaw_range_deg": (
            comparison.leave_one_out_yaw_range_deg
        ),
        "leave_one_out_observable_max_deg": (
            comparison.leave_one_out_observable_max_deg
        ),
        "strong_pair_candidate_count": (
            comparison.strong_pair_candidate_count
        ),
        "decorrelated_pair_count": comparison.decorrelated_pair_count,
        "strong_pair_counts_per_session": (
            comparison.strong_pair_counts_per_session
        ),
        "excluded_sessions": list(comparison.excluded_sessions),
        "rehydration": rehydration_report,
        "comparison_file": "phase_a_comparison.json",
        "observability_file": "phase_a_observability.json",
        "leave_one_out_file": "phase_a_leave_one_out.json",
        "enriched_cache_file": enriched_cache_path.name,
    }

    _write_json(output_directory / "phase_a_comparison.json", full)
    _write_json(
        output_directory / "phase_a_observability.json",
        full["marginal_observability_A1"],
    )
    _write_json(
        output_directory / "phase_a_leave_one_out.json",
        full["leave_one_out"],
    )
    _write_json(output_directory / "phase_a_summary.json", summary)
    return summary