"""Serialize / deserialize motion pairs for fast visualization."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np

from .contracts import MotionPair

SCHEMA_VERSION = 2

# Keep visualization fields plus the compact 3x3 rotation metadata needed to
# rerun Phase-A without repeating LiDAR registration. Full 9x9 Phase-C matrices
# remain excluded from this cache.
_METADATA_KEEP = frozenset(
    {
        "backend",
        "rotation_deg_A",
        "rotation_deg_B",
        "translation_m_B",
        "weight",
        "duration_s",
        "mean_gyro_norm",
        "preint_sigma_rad",
        "cov",
        "J_bg",
        "phase_a_metadata_rehydrated",
        "rehydrated_R_A_error_deg",
        "t_i_imu_s",
        "t_j_imu_s",
        "gyro_bias0_rad_s",
        "accel_bias0_m_s2",
        "time_offset_s",
        "keyframe_span",
        "is_consecutive",
        "modeling",
    }
)


def _to_list(value: Any) -> Any:
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, (np.floating, np.integer, np.bool_)):
        return value.item()
    return value


def pair_to_dict(pair: MotionPair) -> dict[str, Any]:
    meta = {
        str(k): _to_list(v)
        for k, v in (pair.metadata or {}).items()
        if str(k) in _METADATA_KEEP
    }
    return {
        "session_id": pair.session_id,
        "i": int(pair.i),
        "j": int(pair.j),
        "t_i_s": float(pair.t_i_s),
        "t_j_s": float(pair.t_j_s),
        "R_A": np.asarray(pair.R_A, dtype=float).reshape(3, 3).tolist(),
        "R_B": np.asarray(pair.R_B, dtype=float).reshape(3, 3).tolist(),
        "t_A_m": None if pair.t_A_m is None else np.asarray(pair.t_A_m, dtype=float).reshape(3).tolist(),
        "t_B_m": None if pair.t_B_m is None else np.asarray(pair.t_B_m, dtype=float).reshape(3).tolist(),
        "fitness": float(pair.fitness),
        "metadata": meta,
    }


def pair_from_dict(payload: dict[str, Any]) -> MotionPair:
    t_a = payload.get("t_A_m")
    t_b = payload.get("t_B_m")
    return MotionPair(
        session_id=str(payload.get("session_id", "")),
        i=int(payload["i"]),
        j=int(payload["j"]),
        t_i_s=float(payload["t_i_s"]),
        t_j_s=float(payload["t_j_s"]),
        R_A=np.asarray(payload["R_A"], dtype=float).reshape(3, 3),
        R_B=np.asarray(payload["R_B"], dtype=float).reshape(3, 3),
        t_A_m=None if t_a is None else np.asarray(t_a, dtype=float).reshape(3),
        t_B_m=None if t_b is None else np.asarray(t_b, dtype=float).reshape(3),
        fitness=float(payload.get("fitness", 0.0)),
        metadata=dict(payload.get("metadata") or {}),
    )


def build_motion_pairs_payload(
    *,
    prepared_sessions: list[dict[str, Any]],
) -> dict[str, Any]:
    """Build a JSON-serializable cache from pipeline ``prepared`` session dicts."""

    sessions_out: list[dict[str, Any]] = []
    for prep in prepared_sessions:
        pairs = prep.get("pairs") or ()
        sessions_out.append(
            {
                "session_id": prep.get("session_id"),
                "delta_t_s": float(prep.get("time_offset_s", 0.0)),
                "gyro_bias_rad_s": np.asarray(prep.get("gyro_bias_rad_s", np.zeros(3)), dtype=float)
                .reshape(3)
                .tolist(),
                "pair_count": len(pairs),
                "pairs": [pair_to_dict(pair) for pair in pairs],
            }
        )
    return {
        "schema_version": SCHEMA_VERSION,
        "sessions": sessions_out,
        "note": "Cached motion pairs for visualization; A=IMU preintegration, B=LiDAR registration",
    }


def save_motion_pairs(path: Path | str, payload: dict[str, Any]) -> Path:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return destination


def load_motion_pairs(path: Path | str) -> dict[str, Any]:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    version = int(payload.get("schema_version", 0))
    if version not in {1, SCHEMA_VERSION}:
        raise ValueError(
            f"unsupported motion_pairs schema_version={version}; "
            f"expected 1 or {SCHEMA_VERSION}"
        )
    return payload


def pairs_for_session(payload: dict[str, Any], session_id: str | None = None) -> list[MotionPair]:
    sessions = payload.get("sessions") or []
    if not sessions:
        return []
    if session_id is None:
        chosen = sessions[0]
    else:
        chosen = next((s for s in sessions if s.get("session_id") == session_id), None)
        if chosen is None:
            raise KeyError(f"session_id {session_id!r} not found in motion_pairs cache")
    return [pair_from_dict(item) for item in chosen.get("pairs") or []]


def resolve_motion_pairs_path(summary_path: Path | str) -> Path | None:
    """Return ``motion_pairs.json`` next to a summary if it exists."""

    summary = Path(summary_path)
    candidate = summary.parent / "motion_pairs.json"
    return candidate if candidate.is_file() else None