#!/usr/bin/env python3
"""Align HI13/H32 via host-UTC bridge, then run rotation_only.

Device clocks (HI13 boot ms vs H32 absolute) must NOT be forced to share a
first-sample epoch. Instead map each LiDAR frame onto the IMU device timeline
by interpolating IMU device time at the frame's MSOP HostReceiveUtcTicks.
Optional |ω| correlation then refines residual host/path delay.
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
import sys
from pathlib import Path

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.cli import main as cli_main
from imu_lidar.geometry import so3_log
from imu_lidar.imu_io import load_imu_samples
from imu_lidar.lidar_io import load_lidar_frames
from imu_lidar.registration import estimate_frame_rotations
from imu_lidar.time_offset import _correlate_offset, _magnitude_series

SESSIONS = [
    "priority_174005_174515",
    "priority_174905_175450",
    "priority_175910_180530",
]


def _read_imu_host_table(imu_csv: Path) -> tuple[np.ndarray, np.ndarray]:
    rows = list(csv.DictReader(imu_csv.open(encoding="utf-8")))
    if not rows:
        raise RuntimeError(f"empty IMU csv: {imu_csv}")
    if "t_host_utc_s" not in rows[0] or not rows[0].get("t_host_utc_s"):
        raise RuntimeError(
            f"{imu_csv} missing t_host_utc_s; re-export with HostReceiveUtcTicks support"
        )
    t_dev = np.asarray([float(row["t"]) for row in rows], dtype=np.float64)
    t_host = np.asarray([float(row["t_host_utc_s"]) for row in rows], dtype=np.float64)
    order = np.argsort(t_host)
    return t_host[order], t_dev[order]


def _imu_device_at_host(t_host_query: np.ndarray, imu_host: np.ndarray, imu_dev: np.ndarray) -> np.ndarray:
    """Map host UTC seconds → IMU device seconds (linear interp, edge clamp)."""

    return np.interp(t_host_query, imu_host, imu_dev)


def rewrite_lidar_index_host_bridge(
    src_index: Path,
    dst_index: Path,
    *,
    imu_host: np.ndarray,
    imu_dev: np.ndarray,
    residual_delta_s: float = 0.0,
) -> dict:
    """Rewrite LiDAR times onto IMU device clock via host UTC bridge.

    For each frame:
      t_host_mid = mid of MSOP host receive window
      t_imu_mid  = interp(IMU device @ t_host_mid) + residual_delta
      keep device duration: t_start/t_end centered on t_imu_mid
    """

    rows = list(csv.DictReader(src_index.open(encoding="utf-8")))
    if not rows:
        raise RuntimeError(f"empty frames_index: {src_index}")
    if "t_host_utc_s" not in rows[0]:
        raise RuntimeError(
            f"{src_index} missing t_host_utc_s; re-export DLog with MSOP HostReceiveUtcTicks"
        )

    dst_index.parent.mkdir(parents=True, exist_ok=True)
    offsets: list[float] = []
    with dst_index.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["frame_id", "filename", "t_start", "t_end"])
        for row in rows:
            t0 = float(row["t_start"])
            t1 = float(row["t_end"])
            host0 = row.get("t_host_utc_s") or ""
            host1 = row.get("t_host_utc_end_s") or ""
            if not host0:
                raise RuntimeError(f"frame {row.get('frame_id')} missing t_host_utc_s")
            h0 = float(host0)
            h1 = float(host1) if host1 else h0
            host_mid = 0.5 * (h0 + h1)
            imu_mid = float(_imu_device_at_host(np.asarray([host_mid]), imu_host, imu_dev)[0])
            imu_mid += residual_delta_s
            duration = max(t1 - t0, 1e-3)
            new0 = imu_mid - 0.5 * duration
            new1 = imu_mid + 0.5 * duration
            offsets.append(imu_mid - 0.5 * (t0 + t1))
            writer.writerow(
                [
                    row["frame_id"],
                    row["filename"],
                    f"{new0:.9f}",
                    f"{new1:.9f}",
                ]
            )
    arr = np.asarray(offsets, dtype=np.float64)
    return {
        "frames": len(offsets),
        "bridge_offset_median_s": float(np.median(arr)),
        "bridge_offset_mean_s": float(np.mean(arr)),
        "bridge_offset_std_s": float(np.std(arr)),
        "bridge_offset_min_s": float(np.min(arr)),
        "bridge_offset_max_s": float(np.max(arr)),
        "residual_delta_s": float(residual_delta_s),
    }


def estimate_residual_delta(session_dir: Path, *, search_s: float = 5.0) -> tuple[float, float]:
    imu = load_imu_samples(session_dir / "imu.csv")
    frames = load_lidar_frames(session_dir / "lidar")
    # Short pairs only — large stride anti-correlates with IMU |gyro|.
    stride = 1 if len(frames) < 80 else 2
    rotations, pair_times = estimate_frame_rotations(frames, stride=stride)
    if len(rotations) < 8:
        rotations, pair_times = estimate_frame_rotations(frames, stride=1)
    lidar_t = []
    lidar_w = []
    for (t_a, t_b), rotation in zip(pair_times, rotations):
        dt_pair = max(t_b - t_a, 1e-3)
        omega = so3_log(rotation) / dt_pair
        lidar_t.append(0.5 * (t_a + t_b))
        lidar_w.append(omega)
    imu_t, imu_mag = _magnitude_series(imu.t_s, imu.gyro_rad_s)
    lidar_t_arr, lidar_mag = _magnitude_series(np.asarray(lidar_t), np.asarray(lidar_w))
    delta, peak = _correlate_offset(
        imu_t,
        imu_mag,
        lidar_t_arr,
        lidar_mag,
        search_s=search_s,
        sample_hz=20.0,
    )
    return float(delta), float(peak)


def align_session(src: Path, dst: Path, *, residual_search_s: float = 5.0) -> dict:
    if dst.exists():
        shutil.rmtree(dst)
    dst.mkdir(parents=True)
    shutil.copy2(src / "imu.csv", dst / "imu.csv")
    shutil.copytree(src / "lidar" / "frames", dst / "lidar" / "frames")

    imu_host, imu_dev = _read_imu_host_table(src / "imu.csv")
    bridge = rewrite_lidar_index_host_bridge(
        src / "lidar" / "frames_index.csv",
        dst / "lidar" / "frames_index.csv",
        imu_host=imu_host,
        imu_dev=imu_dev,
        residual_delta_s=0.0,
    )

    residual_delta, residual_peak = estimate_residual_delta(dst, search_s=residual_search_s)
    # Only apply residual when correlation is clearly positive; otherwise the
    # host-UTC bridge alone is the trusted alignment (weak peaks are noise).
    apply_residual = residual_peak >= 0.5 and abs(residual_delta) <= residual_search_s
    applied = float(residual_delta) if apply_residual else 0.0
    if apply_residual:
        bridge = rewrite_lidar_index_host_bridge(
            src / "lidar" / "frames_index.csv",
            dst / "lidar" / "frames_index.csv",
            imu_host=imu_host,
            imu_dev=imu_dev,
            residual_delta_s=applied,
        )

    meta = {
        "source": str(src),
        "aligned": str(dst),
        "method": "host_utc_bridge",
        "imu_host_span_s": [float(imu_host[0]), float(imu_host[-1])],
        "imu_device_span_s": [float(imu_dev[0]), float(imu_dev[-1])],
        "bridge": bridge,
        "residual_delta_s": residual_delta,
        "residual_peak": residual_peak,
        "residual_applied_s": applied,
        "residual_applied": apply_residual,
        "note": (
            "LiDAR t_* rewritten onto IMU device clock via MSOP/IMU HostReceiveUtc; "
            "not first-device-sample coincidence. residual |omega| shift applied only if peak>=0.5."
        ),
    }
    (dst / "align_meta.json").write_text(
        json.dumps(meta, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    return meta


def run_one(session_dir: Path, vehicle: Path, search_s: float) -> dict:
    output = session_dir / "out"
    if output.exists():
        shutil.rmtree(output)
    argv = [
        "run",
        "--session-id",
        session_dir.name,
        "--imu",
        str(session_dir / "imu.csv"),
        "--lidar",
        str(session_dir / "lidar"),
        "--vehicle-config",
        str(vehicle),
        "--output",
        str(output),
        "--mode",
        "rotation_only",
        "--time-offset-search-s",
        str(search_s),
        "--min-pair-rotation-deg",
        "2.0",
    ]
    code = cli_main(argv)
    summary_path = output / "summary.json"
    summary = {}
    if summary_path.is_file():
        summary = json.loads(summary_path.read_text(encoding="utf-8"))
    t_block = summary.get("T_IMU_lidar") or {}
    return {
        "session": session_dir.name,
        "exit_code": code,
        "status": summary.get("status"),
        "message": summary.get("message"),
        "time_offset_s": summary.get("time_offset_s"),
        "rotation_deg": t_block.get("rotation_deg") if isinstance(t_block, dict) else None,
        "summary": str(summary_path) if summary_path.is_file() else None,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--sessions-root",
        type=Path,
        default=Path(r"D:\data\calibration_usable_20260808\sessions_v1"),
    )
    parser.add_argument(
        "--aligned-root",
        type=Path,
        default=Path(r"D:\data\calibration_usable_20260808\sessions_v1_host_aligned"),
    )
    parser.add_argument(
        "--vehicle-config",
        type=Path,
        default=ROOT / "config" / "vehicle_hi13_h32_20260808.yaml",
    )
    parser.add_argument(
        "--residual-search-s",
        type=float,
        default=5.0,
        help="|ω| residual search after host bridge (seconds)",
    )
    parser.add_argument("--time-offset-search-s", type=float, default=1.0)
    args = parser.parse_args()

    results = []
    for name in SESSIONS:
        src = args.sessions_root / name
        if not src.is_dir():
            raise SystemExit(f"missing session: {src}")
        aligned = args.aligned_root / name
        print(f"=== align {name} ===", flush=True)
        meta = align_session(src, aligned, residual_search_s=args.residual_search_s)
        print(json.dumps(meta, ensure_ascii=False, indent=2), flush=True)
        print(f"=== calibrate {name} ===", flush=True)
        result = run_one(aligned, args.vehicle_config, args.time_offset_search_s)
        results.append({"align": meta, **result})
        print(json.dumps(result, ensure_ascii=False, indent=2), flush=True)

    manifest = args.aligned_root / "calibration_manifest.json"
    args.aligned_root.mkdir(parents=True, exist_ok=True)
    manifest.write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"manifest: {manifest}")
    return 0


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