#!/usr/bin/env python3
"""Build one LiDAR-centric NPZ per frame with matched RTK and an IMU window.

Inputs are LiDAR frame NPZ files from ``export_h32_rscap_station.py`` (or legacy
``frontlidar_dlog_export.py``) and parsed RTK/IMU JSONL from
``parse_rtk_imu_v2.py``.  Raw ``.rscap`` files remain the traceability source;
this script never modifies them.

Position rows may be NMEA ``GGA`` or G90 ``PVTSLNA`` (both expose ``lat_deg`` /
``lon_deg`` / ``altitude_m``).  Default time basis is LiDAR device time vs GNSS
week/TOW; ``--time-basis host`` keeps the legacy host-receive nearest-neighbour
association for old dlog datasets.
"""

from __future__ import annotations

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

import numpy as np


GPS_EPOCH_UNIX_NS = 315964800 * 1_000_000_000
POSITION_TYPES = {"GGA", "PVTSLNA"}
HEADING_TYPES = {"UNIHEADINGA", "GNHPR"}


def parse_named_path(text: str) -> tuple[str, Path]:
    if "=" not in text:
        raise argparse.ArgumentTypeError("expected NAME=PATH")
    name, raw_path = text.split("=", 1)
    if not name.strip():
        raise argparse.ArgumentTypeError("segment name is empty")
    return name.strip(), Path(raw_path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--lidar",
        type=parse_named_path,
        action="append",
        required=True,
        metavar="NAME=FRAMES_DIR",
        help="Repeat for each LiDAR segment; directory contains exported *.npz frames.",
    )
    parser.add_argument("--rtk", type=Path, action="append", required=True, help="Parsed rtk.jsonl; repeat per session.")
    parser.add_argument("--imu", type=Path, action="append", required=True, help="Parsed imu.jsonl; repeat per session.")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--rtk-max-dt-ms", type=float, default=150.0)
    parser.add_argument("--imu-before-ms", type=float, default=100.0)
    parser.add_argument("--imu-after-ms", type=float, default=100.0)
    parser.add_argument("--gps-utc-leap-seconds", type=int, default=18)
    parser.add_argument(
        "--time-basis",
        choices=("device_gnss", "host"),
        default="device_gnss",
        help="device_gnss: LiDAR unix_time_ns ↔ GNSS week/TOW; host: legacy host-receive association.",
    )
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def load_jsonl(paths: list[Path]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for source_index, path in enumerate(paths):
        source_file = str(path.resolve())
        with path.open("r", encoding="utf-8") as stream:
            for line_number, line in enumerate(stream, start=1):
                if not line.strip():
                    continue
                row = json.loads(line)
                row["_source_file"] = source_file
                row["_source_index"] = source_index
                row["_source_line"] = line_number
                rows.append(row)
    return rows


def utf8_array(value: Any) -> np.ndarray:
    return np.frombuffer(str(value if value is not None else "").encode("utf-8"), dtype=np.uint8)


def scalar(array: np.ndarray) -> Any:
    return array.reshape(-1)[0].item()


def nearest_index(times: np.ndarray, target: int) -> int:
    if not len(times):
        return -1
    right = int(np.searchsorted(times, target, side="left"))
    candidates = [index for index in (right - 1, right) if 0 <= index < len(times)]
    return min(candidates, key=lambda index: abs(int(times[index]) - target))


def estimate_imu_times(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Recover timing inside each serial chunk from device timestamps.

    A capture chunk has one host receive timestamp but may contain several IMU
    frames.  The last frame is anchored to the chunk receive time and earlier
    frames are moved backwards by their device timestamp difference.
    """
    groups: dict[tuple[int, int], list[dict[str, Any]]] = {}
    for row in rows:
        if not row.get("crc_valid") or row.get("device_timestamp_ms") is None:
            continue
        key = (int(row["_source_index"]), int(row.get("source_chunk_sequence_last", -1)))
        groups.setdefault(key, []).append(row)
    result: list[dict[str, Any]] = []
    for group in groups.values():
        group.sort(key=lambda row: (int(row["device_timestamp_ms"]), int(row["_source_line"])))
        last_device = int(group[-1]["device_timestamp_ms"])
        host_ns = int(group[-1]["host_receive_utc_ns"])
        for row in group:
            delta_ms = (last_device - int(row["device_timestamp_ms"])) & 0xFFFFFFFF
            if delta_ms > 60_000:
                delta_ms = 0
            copied = dict(row)
            copied["estimated_time_ns"] = host_ns - delta_ms * 1_000_000
            result.append(copied)
    result.sort(key=lambda row: int(row["estimated_time_ns"]))
    return result


def gnss_utc_ns(row: dict[str, Any], leap_seconds: int) -> int | None:
    week, tow_ms = row.get("gnss_week"), row.get("gnss_tow_ms")
    if week is None or tow_ms is None:
        return None
    seconds = int(week) * 604800 + float(tow_ms) / 1000.0 - leap_seconds
    return GPS_EPOCH_UNIX_NS + int(round(seconds * 1_000_000_000))


def association_time_ns(row: dict[str, Any], time_basis: str, leap_seconds: int) -> int | None:
    if time_basis == "host":
        host = row.get("host_receive_utc_ns")
        return int(host) if host is not None else None
    device = gnss_utc_ns(row, leap_seconds)
    if device is not None:
        return device
    host = row.get("host_receive_utc_ns")
    return int(host) if host is not None else None


def numeric_array(rows: list[dict[str, Any]], key: str, dtype: Any, default: Any) -> np.ndarray:
    return np.asarray([row.get(key, default) if row.get(key) is not None else default for row in rows], dtype=dtype)


def raw_frame_matrix(rows: list[dict[str, Any]]) -> tuple[np.ndarray, np.ndarray]:
    frames = [bytes.fromhex(str(row.get("raw_frame_hex", ""))) for row in rows]
    lengths = np.asarray([len(frame) for frame in frames], dtype=np.int32)
    width = max(lengths, default=0)
    matrix = np.zeros((len(frames), width), dtype=np.uint8)
    for index, frame in enumerate(frames):
        matrix[index, : len(frame)] = np.frombuffer(frame, dtype=np.uint8)
    return matrix, lengths


def add_rtk(values: dict[str, np.ndarray], prefix: str, row: dict[str, Any] | None, dt_ns: int | None) -> None:
    values[f"{prefix}_valid"] = np.asarray([row is not None], dtype=np.uint8)
    values[f"{prefix}_dt_ns"] = np.asarray([dt_ns or 0], dtype=np.int64)
    values[f"{prefix}_host_receive_utc_ns"] = np.asarray([0], dtype=np.int64)
    values[f"{prefix}_raw_utf8"] = utf8_array("")
    values[f"{prefix}_source_file_utf8"] = utf8_array("")
    values[f"{prefix}_source_raw_file_offset"] = np.asarray([-1], dtype=np.int64)
    values[f"{prefix}_source_raw_byte_length"] = np.asarray([0], dtype=np.int32)
    if row is None:
        return
    values[f"{prefix}_host_receive_utc_ns"] = np.asarray([row.get("host_receive_utc_ns", 0)], dtype=np.int64)
    values[f"{prefix}_raw_utf8"] = utf8_array(row.get("raw_line", ""))
    values[f"{prefix}_source_file_utf8"] = utf8_array(row.get("_source_file", ""))
    values[f"{prefix}_source_raw_file_offset"] = np.asarray([row.get("source_raw_file_offset", -1)], dtype=np.int64)
    values[f"{prefix}_source_raw_byte_length"] = np.asarray([row.get("source_raw_byte_length", 0)], dtype=np.int32)


def initialize_rtk_measurements(values: dict[str, np.ndarray]) -> None:
    for key, dtype, default in (
        ("lat_deg", np.float64, np.nan), ("lon_deg", np.float64, np.nan),
        ("altitude_m", np.float64, np.nan), ("hdop", np.float64, np.nan),
        ("fix_quality", np.int32, -1), ("gga_satellites", np.int32, -1),
        ("differential_age_s", np.float64, np.nan),
        ("gnss_week", np.int32, -1), ("gnss_tow_ms", np.int64, -1),
        ("baseline_length_m", np.float64, np.nan), ("raw_heading_deg", np.float64, np.nan),
        ("pitch_deg", np.float64, np.nan), ("roll_deg", np.float64, np.nan),
        ("heading_stddev_deg", np.float64, np.nan),
        ("pitch_stddev_deg", np.float64, np.nan), ("heading_satellites", np.int32, -1),
        ("solution_satellites", np.int32, -1),
    ):
        values[f"rtk_{key}"] = np.asarray([default], dtype=dtype)
    values["rtk_fixed"] = np.asarray([0], dtype=np.uint8)
    values["rtk_heading_solution_utf8"] = utf8_array("")
    values["rtk_heading_gnss_utc_ns"] = np.asarray([0], dtype=np.int64)
    values["rtk_heading_host_minus_gnss_ns"] = np.asarray([0], dtype=np.int64)


def build_combined(
    lidar_segments: list[tuple[str, Path]],
    rtk_paths: list[Path],
    imu_paths: list[Path],
    out: Path,
    *,
    rtk_max_dt_ms: float = 150.0,
    imu_before_ms: float = 100.0,
    imu_after_ms: float = 100.0,
    gps_utc_leap_seconds: int = 18,
    time_basis: str = "device_gnss",
    overwrite: bool = False,
) -> dict[str, Any]:
    """Associate LiDAR frames with RTK/IMU and write ``out/`` combined package."""

    if out.exists() and any(out.iterdir()) and not overwrite:
        raise FileExistsError(f"{out} is non-empty; pass overwrite=True")
    frames_out = out / "frames"
    frames_out.mkdir(parents=True, exist_ok=True)

    rtk_rows = load_jsonl(rtk_paths)
    positions = []
    for row in rtk_rows:
        if row.get("type") not in POSITION_TYPES or not row.get("checksum_valid"):
            continue
        if row.get("lat_deg") is None or row.get("lon_deg") is None:
            continue
        assoc = association_time_ns(row, time_basis, gps_utc_leap_seconds)
        if assoc is None:
            continue
        copied = dict(row)
        copied["_assoc_time_ns"] = assoc
        positions.append(copied)
    positions.sort(key=lambda row: int(row["_assoc_time_ns"]))

    heading = []
    for row in rtk_rows:
        if row.get("type") not in HEADING_TYPES or not row.get("checksum_valid") or not row.get("heading_valid"):
            continue
        assoc = association_time_ns(row, time_basis, gps_utc_leap_seconds)
        if assoc is None:
            continue
        copied = dict(row)
        copied["_assoc_time_ns"] = assoc
        heading.append(copied)
    heading.sort(key=lambda row: int(row["_assoc_time_ns"]))

    imu = estimate_imu_times(load_jsonl(imu_paths))
    position_times = np.asarray([int(row["_assoc_time_ns"]) for row in positions], dtype=np.int64)
    heading_times = np.asarray([int(row["_assoc_time_ns"]) for row in heading], dtype=np.int64)
    imu_times = np.asarray([int(row["estimated_time_ns"]) for row in imu], dtype=np.int64)

    manifest: list[dict[str, Any]] = []
    global_index = 0
    max_rtk_ns = int(rtk_max_dt_ms * 1_000_000)
    before_ns = int(imu_before_ms * 1_000_000)
    after_ns = int(imu_after_ms * 1_000_000)

    for segment_name, frame_dir in lidar_segments:
        frame_paths = sorted(frame_dir.glob("*.npz"))
        if not frame_paths:
            raise FileNotFoundError(f"no NPZ frames under {frame_dir}")
        for segment_index, source in enumerate(frame_paths):
            with np.load(source, allow_pickle=False) as frame:
                values = {key: np.asarray(frame[key]) for key in frame.files}
            lidar_device_time_ns = int(scalar(values["unix_time_ns"]))
            if time_basis == "host":
                lidar_time_ns = int(scalar(values["host_receive_utc_ns"]))
                if lidar_time_ns <= 0:
                    raise ValueError(f"host time requested but missing in {source}")
            else:
                lidar_time_ns = lidar_device_time_ns
            values["lidar_association_time_ns"] = np.asarray([lidar_time_ns], dtype=np.int64)

            position_index = nearest_index(position_times, lidar_time_ns)
            heading_index = nearest_index(heading_times, lidar_time_ns)
            position_row = positions[position_index] if position_index >= 0 else None
            heading_row = heading[heading_index] if heading_index >= 0 else None
            position_dt = int(position_times[position_index]) - lidar_time_ns if position_index >= 0 else None
            heading_dt = int(heading_times[heading_index]) - lidar_time_ns if heading_index >= 0 else None
            position_ok = position_row is not None and abs(position_dt or 0) <= max_rtk_ns
            heading_ok = heading_row is not None and abs(heading_dt or 0) <= max_rtk_ns
            add_rtk(values, "rtk_gga", position_row if position_ok else None, position_dt)
            add_rtk(values, "rtk_heading", heading_row if heading_ok else None, heading_dt)
            initialize_rtk_measurements(values)

            if position_ok and position_row:
                for key, dtype, default in (
                    ("lat_deg", np.float64, np.nan), ("lon_deg", np.float64, np.nan),
                    ("altitude_m", np.float64, np.nan), ("hdop", np.float64, np.nan),
                    ("fix_quality", np.int32, -1), ("gga_satellites", np.int32, -1),
                    ("differential_age_s", np.float64, np.nan),
                ):
                    values[f"rtk_{key}"] = np.asarray([position_row.get(key, default)], dtype=dtype)
                values["rtk_gga_satellites"] = np.asarray([position_row.get("satellites", -1)], dtype=np.int32)
                if position_row.get("gnss_week") is not None:
                    values["rtk_gnss_week"] = np.asarray([position_row.get("gnss_week", -1)], dtype=np.int32)
                    values["rtk_gnss_tow_ms"] = np.asarray([position_row.get("gnss_tow_ms", -1)], dtype=np.int64)
                values["rtk_fixed"] = np.asarray([int(position_row.get("fix_quality", -1)) in {4, 5}], dtype=np.uint8)
            if heading_ok and heading_row:
                for key, dtype, default in (
                    ("gnss_week", np.int32, -1), ("gnss_tow_ms", np.int64, -1),
                    ("baseline_length_m", np.float64, np.nan), ("raw_heading_deg", np.float64, np.nan),
                    ("pitch_deg", np.float64, np.nan), ("roll_deg", np.float64, np.nan),
                    ("heading_stddev_deg", np.float64, np.nan),
                    ("pitch_stddev_deg", np.float64, np.nan),
                    ("solution_satellites", np.int32, -1),
                ):
                    value = heading_row.get(key, default)
                    if key == "roll_deg" and value is None:
                        value = 0.0
                    values[f"rtk_{key}"] = np.asarray([value], dtype=dtype)
                values["rtk_heading_satellites"] = np.asarray([heading_row.get("satellites", -1)], dtype=np.int32)
                values["rtk_heading_solution_utf8"] = utf8_array(heading_row.get("heading_solution", ""))
                device_ns = gnss_utc_ns(heading_row, gps_utc_leap_seconds)
                values["rtk_heading_gnss_utc_ns"] = np.asarray([device_ns or 0], dtype=np.int64)
                values["rtk_heading_host_minus_gnss_ns"] = np.asarray(
                    [int(heading_row["host_receive_utc_ns"]) - device_ns if device_ns is not None else 0], dtype=np.int64
                )

            left = int(np.searchsorted(imu_times, lidar_time_ns - before_ns, side="left"))
            right = int(np.searchsorted(imu_times, lidar_time_ns + after_ns, side="right"))
            window = imu[left:right]
            values["imu_window_count"] = np.asarray([len(window)], dtype=np.int32)
            values["imu_valid"] = np.asarray([bool(window)], dtype=np.uint8)
            values["imu_time_ns"] = numeric_array(window, "estimated_time_ns", np.int64, 0)
            values["imu_host_receive_utc_ns"] = numeric_array(window, "host_receive_utc_ns", np.int64, 0)
            for key in ("device_timestamp_ms", "pps_sync_stamp_ms", "tag"):
                values[f"imu_{key}"] = numeric_array(window, key, np.int64, -1)
            for key in (
                "temperature_c", "air_pressure_pa", "accel_x_mps2", "accel_y_mps2", "accel_z_mps2",
                "gyro_x_radps", "gyro_y_radps", "gyro_z_radps", "mag_x_ut", "mag_y_ut", "mag_z_ut",
                "roll_deg", "pitch_deg", "yaw_deg", "quaternion_w", "quaternion_x", "quaternion_y", "quaternion_z",
            ):
                values[f"imu_{key}"] = numeric_array(window, key, np.float64, np.nan)
            values["imu_source_index"] = numeric_array(window, "_source_index", np.int32, -1)
            values["imu_source_raw_file_offset"] = numeric_array(window, "source_raw_file_offset", np.int64, -1)
            raw_matrix, raw_lengths = raw_frame_matrix(window)
            values["imu_raw_frame_bytes"] = raw_matrix
            values["imu_raw_frame_length"] = raw_lengths
            values["imu_source_files_json_utf8"] = utf8_array(
                json.dumps([str(path.resolve()) for path in imu_paths], ensure_ascii=False)
            )
            values["source_lidar_file_utf8"] = utf8_array(source.resolve())
            values["segment_name_utf8"] = utf8_array(segment_name)

            output = frames_out / f"{segment_name}_{segment_index:06d}.npz"
            np.savez_compressed(output, **values)
            manifest.append({
                "global_index": global_index,
                "segment": segment_name,
                "segment_index": segment_index,
                "output": str(output.relative_to(out)),
                "source_lidar": str(source.resolve()),
                "lidar_time_ns": lidar_time_ns,
                "lidar_device_time_ns": lidar_device_time_ns,
                "rtk_gga_dt_ns": position_dt,
                "rtk_heading_dt_ns": heading_dt,
                "rtk_valid": position_ok,
                "heading_valid": heading_ok,
                "rtk_fix_quality": position_row.get("fix_quality") if position_ok and position_row else None,
                "rtk_fixed": bool(position_ok and position_row and int(position_row.get("fix_quality", -1)) in {4, 5}),
                "imu_window_count": len(window),
            })
            global_index += 1

    fields = sorted({key for row in manifest for key in row})
    with (out / "manifest.csv").open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(manifest)
    if time_basis == "device_gnss":
        time_basis_text = (
            "LiDAR MSOP/device unix_time_ns ↔ RTK GNSS week/TOW (fallback host receive); "
            "IMU still windowed on host-anchored device deltas"
        )
    else:
        time_basis_text = (
            "LiDAR and serial host UTC; RTK GNSS time and IMU device time are retained for clock-model refinement"
        )
    summary = {
        "frames": len(manifest),
        "segments": {name: sum(row["segment"] == name for row in manifest) for name, _ in lidar_segments},
        "rtk_valid": sum(bool(row["rtk_valid"]) for row in manifest),
        "heading_valid": sum(bool(row["heading_valid"]) for row in manifest),
        "rtk_fixed": sum(bool(row["rtk_fixed"]) for row in manifest),
        "imu_window_nonempty": sum(int(row["imu_window_count"]) > 0 for row in manifest),
        "rtk_max_dt_ms": rtk_max_dt_ms,
        "imu_window_ms": [-imu_before_ms, imu_after_ms],
        "time_basis": time_basis_text,
        "time_basis_mode": time_basis,
        "position_message_types": sorted(POSITION_TYPES),
        "imu_orientation_warning": "IMU values are in the raw IMU sensor frame; no LiDAR/body extrinsic is applied",
    }
    (out / "dataset_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    return summary


def main() -> int:
    args = parse_args()
    summary = build_combined(
        args.lidar,
        args.rtk,
        args.imu,
        args.out,
        rtk_max_dt_ms=args.rtk_max_dt_ms,
        imu_before_ms=args.imu_before_ms,
        imu_after_ms=args.imu_after_ms,
        gps_utc_leap_seconds=args.gps_utc_leap_seconds,
        time_basis=args.time_basis,
        overwrite=args.overwrite,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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