#!/usr/bin/env python3
"""Export IMU + H32 LiDAR captures to Lidar-IMU V1 intermediate format.

IMU sources:
- ``--imu-kind hi13`` (HI13R4 / HI91) or ``n300`` or ``auto``
- one or more ``--imu-rscap`` files (concatenated)

LiDAR sources (exactly one):
- ``--lidar-dlog``: Medulla dlog dir **or recovered zip** (MSOP + DIFOP)
- ``--lidar-rscap``: legacy H32 MSOP V2 ``.rscap``

Optional host-time window (local wall clock, DateTime.Now.Ticks convention):
- ``--host-start`` / ``--host-end`` e.g. ``2026-08-08T17:40:05``

Output under ``--out``:

    imu.csv
    lidar/
      frames_index.csv
      frames/frame_XXXXX.npz
    export_summary.json

Device times stay in ``t`` / ``t_start``/``t_end``. Host UTC receive times are
also written so LiDAR–IMU alignment can bridge clocks without forcing first-frame
device coincidence.
"""

from __future__ import annotations

import argparse
import csv
import json
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 tools.h32_dlog.load_session import load_h32_dlog_lidar
from tools.h32_dlog.timeutil import (
    local_wall_to_dotnet_ticks,
    local_wall_to_utc_dotnet_ticks,
    utc_dotnet_ticks_to_unix_s,
)
from tools.rscap_v2.capture_format_v2 import file_summary, read_capture
from tools.rscap_v2.h32_msop import iter_h32_frames, iter_h32_frames_from_packets
from tools.rscap_v2.hi13_imu import iter_hi13_imu_samples
from tools.rscap_v2.n300_imu import ImuSample, iter_n300_imu_samples, samples_to_arrays


def write_imu_csv(path: Path, samples: list[ImuSample]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(
            [
                "t",
                "gx",
                "gy",
                "gz",
                "ax",
                "ay",
                "az",
                "t_host_utc_s",
                "receive_utc_ticks",
            ]
        )
        for sample in samples:
            ticks = int(sample.host_receive_utc_ticks)
            t_host = utc_dotnet_ticks_to_unix_s(ticks) if ticks > 0 else float("nan")
            writer.writerow(
                [
                    f"{sample.t_s:.9f}",
                    f"{sample.gyro_rad_s[0]:.12g}",
                    f"{sample.gyro_rad_s[1]:.12g}",
                    f"{sample.gyro_rad_s[2]:.12g}",
                    f"{sample.accel_m_s2[0]:.12g}",
                    f"{sample.accel_m_s2[1]:.12g}",
                    f"{sample.accel_m_s2[2]:.12g}",
                    f"{t_host:.9f}" if ticks > 0 else "",
                    ticks,
                ]
            )


def write_lidar_session(root: Path, frames) -> dict:
    frames_dir = root / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)
    index_path = root / "frames_index.csv"
    with index_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(
            [
                "frame_id",
                "filename",
                "t_start",
                "t_end",
                "host_receive_utc_ticks",
                "t_host_utc_s",
                "host_receive_utc_end_ticks",
                "t_host_utc_end_s",
            ]
        )
        point_counts = []
        host_ok = 0
        for index, frame in enumerate(frames):
            rel = f"frames/frame_{index:05d}.npz"
            np.savez_compressed(root / rel, points=np.asarray(frame.points_xyz, dtype=np.float32))
            h0 = int(getattr(frame, "host_receive_utc_ticks_start", 0) or 0)
            h1 = int(getattr(frame, "host_receive_utc_ticks_end", 0) or 0)
            t_host0 = utc_dotnet_ticks_to_unix_s(h0) if h0 > 0 else float("nan")
            t_host1 = utc_dotnet_ticks_to_unix_s(h1) if h1 > 0 else float("nan")
            if h0 > 0:
                host_ok += 1
            writer.writerow(
                [
                    index,
                    rel,
                    f"{frame.t_start_s:.9f}",
                    f"{frame.t_end_s:.9f}",
                    h0,
                    f"{t_host0:.9f}" if h0 > 0 else "",
                    h1,
                    f"{t_host1:.9f}" if h1 > 0 else "",
                ]
            )
            point_counts.append(int(frame.points_xyz.shape[0]))
    return {
        "frames": len(frames),
        "frames_with_host_utc": host_ok,
        "points_min": int(min(point_counts)) if point_counts else 0,
        "points_max": int(max(point_counts)) if point_counts else 0,
        "points_mean": float(np.mean(point_counts)) if point_counts else 0.0,
        "t_start": float(frames[0].t_start_s) if frames else None,
        "t_end": float(frames[-1].t_end_s) if frames else None,
    }


def detect_imu_kind(paths: list[Path], explicit: str) -> str:
    if explicit != "auto":
        return explicit
    joined = " ".join(path.name.lower() for path in paths)
    if "hi13" in joined or "hipnuc" in joined:
        return "hi13"
    if "n300" in joined or "wheeltec" in joined:
        return "n300"
    return "hi13"


def load_imu_samples(
    paths: list[Path],
    *,
    kind: str,
    host_ticks_min: int | None,
    host_ticks_max: int | None,
) -> tuple[list[ImuSample], list[dict], str]:
    samples: list[ImuSample] = []
    captures_meta: list[dict] = []
    for path in paths:
        capture = read_capture(path)
        captures_meta.append(file_summary(capture))
        if kind == "hi13":
            part = iter_hi13_imu_samples(
                capture,
                host_utc_ticks_min=host_ticks_min,
                host_utc_ticks_max=host_ticks_max,
            )
        elif kind == "n300":
            part = iter_n300_imu_samples(capture)
            if host_ticks_min is not None or host_ticks_max is not None:
                part = [
                    sample
                    for sample in part
                    if (host_ticks_min is None or sample.host_receive_utc_ticks >= host_ticks_min)
                    and (host_ticks_max is None or sample.host_receive_utc_ticks <= host_ticks_max)
                ]
        else:
            raise ValueError(f"unsupported imu kind: {kind}")
        samples.extend(part)
    samples.sort(key=lambda sample: (sample.t_s, sample.device_timestamp_us))
    return samples, captures_meta, kind


def export_session(
    *,
    imu_rscap: list[Path] | Path,
    out: Path,
    lidar_rscap: Path | None = None,
    lidar_dlog: Path | None = None,
    imu_kind: str = "auto",
    msop_object: str = "frontlidar-msop-raw",
    difop_object: str = "frontlidar-difop-raw",
    require_difop: bool = False,
    host_start: str | None = None,
    host_end: str | None = None,
    frame_stride: int = 1,
    max_points_per_frame: int | None = 80000,
    min_range_m: float = 0.3,
    max_range_m: float = 120.0,
    min_frame_points: int = 100,
) -> dict:
    if (lidar_rscap is None) == (lidar_dlog is None):
        raise ValueError("provide exactly one of lidar_rscap or lidar_dlog")

    imu_paths = [imu_rscap] if isinstance(imu_rscap, Path) else list(imu_rscap)
    if not imu_paths:
        raise ValueError("at least one --imu-rscap is required")

    # LiDAR DObject tic uses DateTime.Now; IMU/MSOP host fields use UTC.
    lidar_ticks_min = local_wall_to_dotnet_ticks(host_start) if host_start else None
    lidar_ticks_max = local_wall_to_dotnet_ticks(host_end) if host_end else None
    imu_ticks_min = local_wall_to_utc_dotnet_ticks(host_start) if host_start else None
    imu_ticks_max = local_wall_to_utc_dotnet_ticks(host_end) if host_end else None
    kind = detect_imu_kind(imu_paths, imu_kind)

    out.mkdir(parents=True, exist_ok=True)
    samples, imu_captures, kind = load_imu_samples(
        imu_paths,
        kind=kind,
        host_ticks_min=imu_ticks_min,
        host_ticks_max=imu_ticks_max,
    )
    t, _gyro, _accel = samples_to_arrays(samples)
    imu_csv = out / "imu.csv"
    write_imu_csv(imu_csv, samples)
    imu_host_ok = sum(1 for sample in samples if sample.host_receive_utc_ticks > 0)

    if lidar_dlog is not None:
        session = load_h32_dlog_lidar(
            lidar_dlog,
            msop_object=msop_object,
            difop_object=difop_object,
            require_difop=require_difop,
            host_ticks_min=lidar_ticks_min,
            host_ticks_max=lidar_ticks_max,
        )
        frames = iter_h32_frames_from_packets(
            session.msop_packets,
            host_utc_ticks=session.msop_host_utc_ticks,
            min_frame_points=min_frame_points,
            frame_stride=frame_stride,
            min_range_m=min_range_m,
            max_range_m=max_range_m,
            max_points_per_frame=max_points_per_frame,
            vertical_deg=session.vertical_deg,
            horizontal_deg=session.horizontal_deg,
        )
        lidar_meta = {
            "source": "dlog",
            "lidar_dlog": session.dlog_root,
            "msop_object": session.msop_object,
            "difop_object": session.difop_object,
            "msop_packets": len(session.msop_packets),
            "msop_packets_with_host_utc": sum(1 for ticks in session.msop_host_utc_ticks if ticks > 0),
            "msop_batches": session.msop_batch_count,
            "difop_records": session.difop_record_count,
            "session_id": session.session_id,
            "lidar_ip": session.lidar_ip,
            "angle_source": session.angle_source,
            "timestamp_note": (
                "device: h32_msop_device_timestamp -> seconds; "
                "host: MSOP HostReceiveUtcTicks -> unix seconds"
            ),
        }
    else:
        assert lidar_rscap is not None
        lidar_capture = read_capture(lidar_rscap)
        frames = iter_h32_frames(
            lidar_capture,
            min_frame_points=min_frame_points,
            frame_stride=frame_stride,
            min_range_m=min_range_m,
            max_range_m=max_range_m,
            max_points_per_frame=max_points_per_frame,
        )
        lidar_meta = {
            "source": "rscap_v2",
            "lidar_rscap": str(lidar_rscap),
            "capture": file_summary(lidar_capture),
            "angle_source": "default_msop_only_vertical_-16_to_16_deg",
            "timestamp_note": (
                "device: h32_msop_device_timestamp_ms -> seconds; "
                "host: rscap receive_utc_ticks -> unix seconds"
            ),
        }

    lidar_dir = out / "lidar"
    lidar_stats = write_lidar_session(lidar_dir, frames)
    imu_time_note = (
        "hi13_device_timestamp_ms -> seconds"
        if kind == "hi13"
        else "n300_device_timestamp_us -> seconds"
    )

    summary = {
        "imu_rscap": [str(path) for path in imu_paths],
        "imu_kind": kind,
        "out": str(out),
        "host_window": {
            "host_start": host_start,
            "host_end": host_end,
            "lidar_ticks_min": lidar_ticks_min,
            "lidar_ticks_max": lidar_ticks_max,
            "imu_ticks_min": imu_ticks_min,
            "imu_ticks_max": imu_ticks_max,
            "note": "local wall cut; lidar DObject tic=DateTime.Now, IMU/MSOP host=UTC",
        },
        "timestamp_policy": {
            "imu_device": imu_time_note,
            "imu_host": "rscap receive_utc_ticks -> t_host_utc_s",
            "lidar_device": "MSOP device timestamp -> t_start/t_end",
            "lidar_host": "MSOP HostReceiveUtcTicks -> t_host_utc_s",
            "calibration_align": "bridge via host UTC; do not force first device samples to coincide",
        },
        "imu": {
            "samples": int(t.shape[0]),
            "samples_with_host_utc": imu_host_ok,
            "t_start": float(t[0]) if t.size else None,
            "t_end": float(t[-1]) if t.size else None,
            "captures": imu_captures,
        },
        "lidar": {
            **lidar_stats,
            "frame_stride": int(frame_stride),
            "max_points_per_frame": max_points_per_frame,
            **lidar_meta,
        },
        "outputs": {
            "imu_csv": str(imu_csv),
            "lidar_session": str(lidar_dir),
        },
    }
    (out / "export_summary.json").write_text(
        json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    return summary


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--imu-rscap",
        type=Path,
        action="append",
        required=True,
        help="IMU V2 .rscap (repeatable)",
    )
    parser.add_argument(
        "--imu-kind",
        choices=("auto", "hi13", "n300"),
        default="auto",
        help="IMU decoder (default: auto from filename)",
    )
    lidar = parser.add_mutually_exclusive_group(required=True)
    lidar.add_argument(
        "--lidar-dlog",
        type=Path,
        help="H32 dlog directory or recovered zip (indices.log + data.bin)",
    )
    lidar.add_argument(
        "--lidar-rscap",
        type=Path,
        help="Legacy H32 MSOP V2 .rscap",
    )
    parser.add_argument("--msop-object", default="frontlidar-msop-raw")
    parser.add_argument("--difop-object", default="frontlidar-difop-raw")
    parser.add_argument("--require-difop", action="store_true")
    parser.add_argument("--host-start", type=str, default=None, help="Local wall start, e.g. 2026-08-08T17:40:05")
    parser.add_argument("--host-end", type=str, default=None, help="Local wall end, e.g. 2026-08-08T17:45:15")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--frame-stride", type=int, default=1)
    parser.add_argument("--max-points-per-frame", type=int, default=80000)
    parser.add_argument("--min-range-m", type=float, default=0.3)
    parser.add_argument("--max-range-m", type=float, default=120.0)
    parser.add_argument("--min-frame-points", type=int, default=100)
    args = parser.parse_args()
    max_points = None if args.max_points_per_frame <= 0 else args.max_points_per_frame
    summary = export_session(
        imu_rscap=args.imu_rscap,
        lidar_rscap=args.lidar_rscap,
        lidar_dlog=args.lidar_dlog,
        imu_kind=args.imu_kind,
        msop_object=args.msop_object,
        difop_object=args.difop_object,
        require_difop=args.require_difop,
        host_start=args.host_start,
        host_end=args.host_end,
        out=args.out,
        frame_stride=args.frame_stride,
        max_points_per_frame=max_points,
        min_range_m=args.min_range_m,
        max_range_m=args.max_range_m,
        min_frame_points=args.min_frame_points,
    )
    print(
        json.dumps(
            {
                "imu_kind": summary["imu_kind"],
                "imu_samples": summary["imu"]["samples"],
                "imu_host_utc": summary["imu"]["samples_with_host_utc"],
                "lidar_frames": summary["lidar"]["frames"],
                "lidar_host_utc": summary["lidar"]["frames_with_host_utc"],
                "lidar_source": summary["lidar"]["source"],
                "angle_source": summary["lidar"]["angle_source"],
                "host_window": summary["host_window"],
                "imu_csv": summary["outputs"]["imu_csv"],
                "lidar_session": summary["outputs"]["lidar_session"],
                "export_summary": str(Path(args.out) / "export_summary.json"),
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    if summary["imu"]["samples"] == 0:
        raise SystemExit("no valid IMU samples decoded in window")
    if summary["lidar"]["frames"] == 0:
        raise SystemExit("no valid H32 frames decoded in window")
    if summary["lidar"]["frames_with_host_utc"] == 0:
        raise SystemExit("no LiDAR frames with MSOP HostReceiveUtcTicks; cannot host-bridge align")
    if summary["imu"]["samples_with_host_utc"] == 0:
        raise SystemExit("no IMU samples with host receive UTC; cannot host-bridge align")
    return 0


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