#!/usr/bin/env python3
"""Interactive 3D inspection of LiDAR–IMU motion-pair registration.

Similar to the RTK–LiDAR ``visualize_pair_3d`` viewer, but A comes from IMU
preintegration and X is ``T_IMU_lidar``.

Modes (keyboard):
  1          raw source (no transform)
  2          IMU prediction with X=I  (B_pred = A)
  3          LiDAR registration B     (reference)
  4          calibrated prediction    B_pred = X^{-1} A X
             (rotation_only runs default to R conjug + t_B so bad IMU Δp
              does not dominate the overlay)
  N / ]      next motion pair
  P / [      previous motion pair
  Q / Esc    exit

Blue = target keyframe i; orange = source keyframe j after the selected transform.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from imu_lidar.geometry import (
    inverse_transform,
    make_transform,
    rotation_angle_deg,
    rpy_deg_xyz,
    transform_points,
)
from imu_lidar.imu_io import load_imu_samples
from imu_lidar.imu_preintegration import preintegrate_imu
from imu_lidar.keyframes import build_keyframes
from imu_lidar.lidar_io import load_lidar_frame_at, load_lidar_frames
from imu_lidar.motion_pairs import build_motion_pairs
from imu_lidar.motion_pairs_io import (
    load_motion_pairs,
    pairs_for_session,
    resolve_motion_pairs_path,
)
from imu_lidar.registration import register_lidar_pair
from imu_lidar.time_offset import lidar_time_to_imu_time


class _LazyFrameStore:
    """Load NPZ frames on demand; indices match mid-time-sorted ``load_lidar_frames``."""

    def __init__(self, lidar_dir: Path, *, max_cached: int = 16):
        self.lidar_dir = Path(lidar_dir)
        self.max_cached = max_cached
        self._cache: dict[int, object] = {}
        self._order: list[int] = []

    def __getitem__(self, index: int):
        index = int(index)
        if index in self._cache:
            return self._cache[index]
        frame = load_lidar_frame_at(self.lidar_dir, index)
        self._cache[index] = frame
        self._order.append(index)
        while len(self._order) > self.max_cached:
            old = self._order.pop(0)
            self._cache.pop(old, None)
        return frame


COLORS = {
    "target": [0.10, 0.65, 1.00],
    "source": [1.00, 0.35, 0.05],
}

MODE_NAMES = (
    "1 raw",
    "2 IMU (X=I)",
    "3 LiDAR B",
    "4 calibrated X^-1 A X",
)


def _load_extrinsic(summary_path: Path) -> tuple[np.ndarray, float, np.ndarray, dict[str, Any]]:
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    t_block = summary.get("T_IMU_lidar")
    meta: dict[str, Any] = {
        "rotation_only": False,
        "translation_accepted": False,
        "status": str(summary.get("status") or ""),
    }
    if t_block is None:
        matrix = summary.get("matrix")
        if matrix is not None:
            return np.asarray(matrix, dtype=float), 0.0, np.zeros(3), meta
        raise ValueError(f"no T_IMU_lidar in {summary_path}")
    t_mat = np.asarray(t_block["matrix"], dtype=float)
    dt = float(summary.get("time_offset_s") or 0.0)
    session = (summary.get("details") or {}).get("sessions", [{}])[0]
    joint = session.get("joint") or {}
    bias = np.asarray(
        (session.get("imu_audit") or {}).get("gyro_bias_rad_s")
        or joint.get("gyro_bias_rad_s")
        or [0.0, 0.0, 0.0],
        dtype=float,
    ).reshape(3)
    status = str(summary.get("status") or "")
    translation_accepted = bool(
        joint.get("translation_accepted")
        or (summary.get("details") or {}).get("joint", {}).get("translation_accepted")
    )
    rotation_only = ("rotation_only" in status) or (
        not translation_accepted and float(np.linalg.norm(t_mat[:3, 3])) < 1e-9
    )
    meta.update(
        {
            "rotation_only": rotation_only,
            "translation_accepted": translation_accepted,
            "status": status,
        }
    )
    return t_mat, dt, bias, meta


def _delta_components(reference: np.ndarray, candidate: np.ndarray) -> dict:
    delta = inverse_transform(reference) @ candidate
    translation = np.asarray(delta[:3, 3], dtype=float)
    return {
        "translation_xyz_cm": (translation * 100.0).tolist(),
        "translation_norm_cm": float(np.linalg.norm(translation) * 100.0),
        "rotation_rpy_deg_xyz": rpy_deg_xyz(delta[:3, :3]).tolist(),
        "rotation_angle_deg": rotation_angle_deg(delta[:3, :3]),
    }


def _print_delta(name: str, reference: np.ndarray, candidate: np.ndarray) -> dict:
    item = _delta_components(reference, candidate)
    tx, ty, tz = item["translation_xyz_cm"]
    roll, pitch, yaw = item["rotation_rpy_deg_xyz"]
    print(
        f"{name}: B^-1*motion  "
        f"t_xyz=[{tx:+.3f}, {ty:+.3f}, {tz:+.3f}] cm  "
        f"rpy=[{roll:+.3f}, {pitch:+.3f}, {yaw:+.3f}] deg  "
        f"|t|={item['translation_norm_cm']:.3f} cm  "
        f"|R|={item['rotation_angle_deg']:.4f} deg"
    )
    return item


def _cloud(o3d, points: np.ndarray, color, voxel: float):
    item = o3d.geometry.PointCloud()
    item.points = o3d.utility.Vector3dVector(points)
    if voxel > 0:
        item = item.voxel_down_sample(voxel)
    item.paint_uniform_color(color)
    return item


def _set_cloud_points(cloud, points: np.ndarray, color, voxel: float, o3d) -> None:
    tmp = _cloud(o3d, points, color, voxel)
    cloud.points = tmp.points
    cloud.colors = tmp.colors


def _build_pair_list(
    *,
    lidar_dir: Path,
    imu_path: Path,
    delta_t_s: float,
    gyro_bias: np.ndarray,
    min_rotation_deg: float,
    min_translation_m: float,
):
    frames = load_lidar_frames(lidar_dir)
    imu = load_imu_samples(imu_path)
    keyframes = build_keyframes(
        frames,
        min_translation_m=min_translation_m,
        min_rotation_deg=min_rotation_deg,
    )
    pair_set = build_motion_pairs(
        session_id="viz",
        keyframes=list(keyframes.frames),
        keyframe_indices=keyframes.indices,
        imu=imu,
        delta_t_s=delta_t_s,
        gyro_bias_rad_s=gyro_bias,
        min_rotation_deg=min_rotation_deg,
        min_translation_m=min_translation_m,
    )
    return frames, imu, keyframes, pair_set


def _pair_from_indices(
    frames,
    imu,
    *,
    i: int,
    j: int,
    delta_t_s: float,
    gyro_bias: np.ndarray,
):
    frame_i = frames[i]
    frame_j = frames[j]
    reg = register_lidar_pair(frame_j.points_xyz, frame_i.points_xyz)
    t_i = lidar_time_to_imu_time(frame_i.t_mid_s, delta_t_s)
    t_j = lidar_time_to_imu_time(frame_j.t_mid_s, delta_t_s)
    preint = preintegrate_imu(
        imu.t_s,
        imu.gyro_rad_s,
        imu.acc_m_s2,
        t_i,
        t_j,
        gyro_bias,
        np.zeros(3),
    )
    a = make_transform(preint.delta_p, preint.delta_R)
    return frame_i, frame_j, a, reg.transform


def _transforms_for_pair(
    x: np.ndarray,
    a_ij: np.ndarray,
    b_gicp: np.ndarray,
    *,
    mode4_translation: str = "imu",
) -> dict[str, np.ndarray]:
    """Build overlay transforms.

    ``mode4_translation``:
      - ``imu``: full SE3 conjug ``X^{-1} A X`` (needs trustworthy IMU Δp)
      - ``gicp``: rotation conjug only; translation taken from LiDAR B
        (correct check for rotation_only calibrations)
    """

    calibrated = inverse_transform(x) @ a_ij @ x
    if mode4_translation == "gicp":
        calibrated = make_transform(b_gicp[:3, 3], calibrated[:3, :3])
    elif mode4_translation != "imu":
        raise ValueError(f"unknown mode4_translation={mode4_translation!r}")
    return {
        MODE_NAMES[0]: np.eye(4),
        MODE_NAMES[1]: a_ij.copy(),
        MODE_NAMES[2]: b_gicp.copy(),
        MODE_NAMES[3]: calibrated,
    }


def _resolve_pair(
    frames,
    pairs,
    pair_index: int,
    x: np.ndarray,
    *,
    mode4_translation: str = "imu",
):
    pair = pairs[pair_index]
    frame_i = frames[pair.i]
    frame_j = frames[pair.j]
    a_ij = make_transform(
        pair.t_A_m if pair.t_A_m is not None else np.zeros(3),
        pair.R_A,
    )
    b_gicp = make_transform(
        pair.t_B_m if pair.t_B_m is not None else np.zeros(3),
        pair.R_B,
    )
    transforms = _transforms_for_pair(
        x, a_ij, b_gicp, mode4_translation=mode4_translation
    )
    label = (
        f"pair {pair_index + 1}/{len(pairs)}  "
        f"frames {pair.i} <- {pair.j}  "
        f"rotB={rotation_angle_deg(pair.R_B):.2f} deg  "
        f"|tB|={0.0 if pair.t_B_m is None else float(np.linalg.norm(pair.t_B_m)):.3f} m"
    )
    return frame_i, frame_j, a_ij, b_gicp, transforms, label


def _print_pair_header(label: str, b_gicp: np.ndarray, transforms: dict[str, np.ndarray]) -> None:
    print("-" * 72)
    print(label)
    print("blue=target i | orange=source j")
    print("1-4: overlay mode | N/]: next pair | P/[: prev pair | Q/Esc: exit")
    baseline = _print_delta("mode4 minus mode3", b_gicp, transforms[MODE_NAMES[3]])
    roll, pitch, yaw = np.abs(baseline["rotation_rpy_deg_xyz"])
    if max(roll, pitch) > max(0.10, 2.0 * yaw):
        print("note: roll/pitch dominate yaw on this pair.")
    tx, ty, tz = np.abs(baseline["translation_xyz_cm"])
    if tz > max(tx, ty):
        print("note: largest translation component is Z for this pair.")


def _save_topdown_png(
    path: Path,
    target: np.ndarray,
    source: np.ndarray,
    transforms: dict[str, np.ndarray],
) -> None:
    import matplotlib.pyplot as plt

    def downsample(points: np.ndarray) -> np.ndarray:
        if points.shape[0] <= 8000:
            return points
        idx = np.linspace(0, points.shape[0] - 1, 8000).astype(int)
        return points[idx]

    names = list(transforms.keys())
    fig, axes = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)
    tgt = downsample(target)
    for ax, name in zip(axes.ravel(), names):
        src = downsample(transform_points(source, transforms[name]))
        ax.scatter(tgt[:, 0], tgt[:, 1], s=1, c="tab:blue", alpha=0.35, label="target i")
        ax.scatter(src[:, 0], src[:, 1], s=1, c="tab:orange", alpha=0.35, label="source j")
        ax.set_title(name)
        ax.set_aspect("equal", adjustable="box")
        ax.grid(alpha=0.3)
    axes[0, 0].legend(loc="upper right", markerscale=4)
    fig.suptitle("LiDAR–IMU pair overlay (XY top-down)")
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=140)
    plt.close(fig)
    print(f"wrote {path}")


def _run_gui(
    *,
    frames,
    pairs,
    x: np.ndarray,
    start_index: int,
    voxel: float,
    fixed_single_pair: tuple | None,
    mode4_translation: str = "imu",
) -> None:
    import open3d as o3d

    if fixed_single_pair is not None:
        frame_i, frame_j, a_ij, b_gicp = fixed_single_pair
        transforms = _transforms_for_pair(
            x, a_ij, b_gicp, mode4_translation=mode4_translation
        )
        label = f"fixed frames (no pair switching)"
        pair_index = 0
        n_pairs = 1
        use_list = False
    else:
        pair_index = int(np.clip(start_index, 0, len(pairs) - 1))
        n_pairs = len(pairs)
        use_list = True
        frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, pair_index, x, mode4_translation=mode4_translation
        )

    viewer = o3d.visualization.VisualizerWithKeyCallback()
    viewer.create_window("LiDAR–IMU registration inspection", 1400, 900)
    target_cloud = _cloud(o3d, frame_i.points_xyz, COLORS["target"], voxel)
    source_cloud = _cloud(o3d, frame_j.points_xyz, COLORS["source"], voxel)
    viewer.add_geometry(target_cloud)
    viewer.add_geometry(source_cloud)
    viewer.add_geometry(o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0))
    viewer.get_render_option().background_color = np.array([0.02, 0.02, 0.02])
    viewer.get_render_option().point_size = 2.0

    state = {
        "pair_index": pair_index,
        "mode_name": MODE_NAMES[3],
        "current": np.eye(4),
        "transforms": transforms,
        "b_gicp": b_gicp,
        "frame_i": frame_i,
        "frame_j": frame_j,
    }

    def apply_mode(vis, mode_name: str, *, announce: bool = True) -> None:
        desired = state["transforms"][mode_name]
        source_cloud.transform(desired @ inverse_transform(state["current"]))
        state["current"] = desired
        state["mode_name"] = mode_name
        vis.update_geometry(source_cloud)
        if announce:
            if mode_name == MODE_NAMES[2]:
                print(f"{mode_name}: registration reference; delta = 0")
            else:
                _print_delta(mode_name + " minus mode3", state["b_gicp"], desired)

    def load_pair(vis, new_index: int) -> None:
        if not use_list:
            print("pair switching disabled in --frame-i/--frame-j mode")
            return
        new_index = int(new_index) % n_pairs
        frame_i, frame_j, _a, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, new_index, x, mode4_translation=mode4_translation
        )
        state["pair_index"] = new_index
        state["transforms"] = transforms
        state["b_gicp"] = b_gicp
        state["frame_i"] = frame_i
        state["frame_j"] = frame_j
        state["current"] = np.eye(4)
        _set_cloud_points(target_cloud, frame_i.points_xyz, COLORS["target"], voxel, o3d)
        _set_cloud_points(source_cloud, frame_j.points_xyz, COLORS["source"], voxel, o3d)
        vis.update_geometry(target_cloud)
        vis.update_geometry(source_cloud)
        _print_pair_header(label, b_gicp, transforms)
        apply_mode(vis, state["mode_name"], announce=True)

    def make_mode_cb(mode_name: str):
        def callback(vis):
            apply_mode(vis, mode_name, announce=True)
            return False

        return callback

    def next_pair(vis):
        load_pair(vis, state["pair_index"] + 1)
        return False

    def prev_pair(vis):
        load_pair(vis, state["pair_index"] - 1)
        return False

    _print_pair_header(label, b_gicp, transforms)
    for key, name in zip((ord("1"), ord("2"), ord("3"), ord("4")), MODE_NAMES):
        viewer.register_key_callback(key, make_mode_cb(name))
    for key in (ord("N"), ord("n"), ord("]")):
        viewer.register_key_callback(key, next_pair)
    for key in (ord("P"), ord("p"), ord("[")):
        viewer.register_key_callback(key, prev_pair)

    apply_mode(viewer, MODE_NAMES[3], announce=False)
    viewer.run()
    viewer.destroy_window()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--lidar", required=True, type=Path, help="LiDAR session directory")
    parser.add_argument(
        "--imu",
        type=Path,
        default=None,
        help="IMU CSV (only needed when rebuilding pairs without motion_pairs.json)",
    )
    parser.add_argument(
        "--summary",
        required=True,
        type=Path,
        help="summary.json (or T_IMU_lidar.json) from a calibration run",
    )
    parser.add_argument(
        "--motion-pairs",
        type=Path,
        default=None,
        help="Cached motion_pairs.json (default: next to --summary)",
    )
    parser.add_argument(
        "--session-id",
        default=None,
        help="Session id inside multi-session motion_pairs.json",
    )
    parser.add_argument(
        "--rebuild-pairs",
        action="store_true",
        help="Ignore cache and rebuild pairs from IMU/LiDAR (slow)",
    )
    parser.add_argument("--pair-index", type=int, default=0, help="Starting motion-pair index")
    parser.add_argument("--frame-i", type=int, default=None, help="Optional explicit frame index i")
    parser.add_argument("--frame-j", type=int, default=None, help="Optional explicit frame index j")
    parser.add_argument("--voxel", type=float, default=0.12)
    parser.add_argument("--min-pair-rotation-deg", type=float, default=3.0)
    parser.add_argument("--min-pair-translation-m", type=float, default=0.3)
    parser.add_argument(
        "--save-png",
        type=Path,
        default=None,
        help="Write a 2x2 XY top-down comparison PNG for the starting pair",
    )
    parser.add_argument(
        "--no-gui",
        action="store_true",
        help="Skip Open3D window (use with --save-png)",
    )
    parser.add_argument(
        "--mode4-translation",
        choices=("auto", "gicp", "imu"),
        default="auto",
        help=(
            "Mode-4 translation source: gicp=R conjug + t_B (rotation check); "
            "imu=full X^-1 A X; auto=gicp for rotation_only summaries"
        ),
    )
    args = parser.parse_args(argv)

    x, delta_t_s, gyro_bias, extr_meta = _load_extrinsic(args.summary)
    if args.mode4_translation == "auto":
        mode4_translation = "gicp" if extr_meta.get("rotation_only") else "imu"
    else:
        mode4_translation = args.mode4_translation
    cache_path = args.motion_pairs or resolve_motion_pairs_path(args.summary)
    use_cache = (not args.rebuild_pairs) and cache_path is not None and args.frame_i is None

    frames = None
    pairs: tuple = ()
    fixed_single_pair = None
    label = ""
    b_gicp = np.eye(4)
    transforms: dict[str, np.ndarray] = {}
    frame_i = frame_j = None

    if use_cache:
        payload = load_motion_pairs(cache_path)
        pair_list = pairs_for_session(payload, args.session_id)
        if not pair_list:
            raise SystemExit(f"no pairs in cache: {cache_path}")
        if not 0 <= args.pair_index < len(pair_list):
            raise SystemExit(
                f"pair-index {args.pair_index} outside [0, {len(pair_list) - 1}] "
                f"({len(pair_list)} pairs in cache)"
            )
        frames = _LazyFrameStore(args.lidar)
        pairs = tuple(pair_list)
        frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, args.pair_index, x, mode4_translation=mode4_translation
        )
        print(f"loaded {len(pairs)} cached pairs from {cache_path}")
    else:
        if args.imu is None:
            raise SystemExit("--imu is required when motion_pairs.json is missing (or use --rebuild-pairs with --imu)")
        frames, imu, keyframes, pair_set = _build_pair_list(
            lidar_dir=args.lidar,
            imu_path=args.imu,
            delta_t_s=delta_t_s,
            gyro_bias=gyro_bias,
            min_rotation_deg=args.min_pair_rotation_deg,
            min_translation_m=args.min_pair_translation_m,
        )
        if args.frame_i is not None and args.frame_j is not None:
            frame_i, frame_j, a_ij, b_gicp = _pair_from_indices(
                frames,
                imu,
                i=args.frame_i,
                j=args.frame_j,
                delta_t_s=delta_t_s,
                gyro_bias=gyro_bias,
            )
            transforms = _transforms_for_pair(
                x, a_ij, b_gicp, mode4_translation=mode4_translation
            )
            label = f"frames {args.frame_i} <- {args.frame_j}"
            fixed_single_pair = (frame_i, frame_j, a_ij, b_gicp)
            pairs = ()
        else:
            if not pair_set.pairs:
                raise SystemExit("no motion pairs rebuilt; loosen min-pair thresholds or check data")
            if not 0 <= args.pair_index < len(pair_set.pairs):
                raise SystemExit(
                    f"pair-index {args.pair_index} outside [0, {len(pair_set.pairs) - 1}] "
                    f"({len(pair_set.pairs)} pairs available)"
                )
            pairs = pair_set.pairs
            frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
                frames, pairs, args.pair_index, x, mode4_translation=mode4_translation
            )
            print(f"rebuilt {len(pairs)} pairs from {len(keyframes.indices)} keyframes")

    print(
        f"mode4 translation={mode4_translation}  "
        f"(status={extr_meta.get('status') or 'n/a'}, "
        f"rotation_only={bool(extr_meta.get('rotation_only'))})"
    )
    if mode4_translation == "gicp":
        print(
            "note: mode4 uses R conjug + t_B; IMU Δp is ignored "
            "(typical for rotation_only — raw Δp often has large Z drift)."
        )
    if mode4_translation == "imu":
        print(
            "note: mode4 uses full X^-1 A X. If clouds stack vertically, "
            "IMU Δp is likely bad; retry with --mode4-translation gicp."
        )

    if args.save_png is not None:
        _print_pair_header(label, b_gicp, transforms)
        _save_topdown_png(args.save_png, frame_i.points_xyz, frame_j.points_xyz, transforms)

    if args.no_gui:
        return 0

    try:
        import open3d  # noqa: F401
    except ImportError as exc:
        raise SystemExit(
            "Open3D is required for interactive view. "
            "Install with: python -m pip install -e \".[open3d]\" "
            "or use --no-gui --save-png out.png"
        ) from exc

    _run_gui(
        frames=frames,
        pairs=pairs,
        x=x,
        start_index=args.pair_index,
        voxel=args.voxel,
        fixed_single_pair=fixed_single_pair,
        mode4_translation=mode4_translation,
    )
    return 0


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