#!/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 frontlidar_dlog_export.py and parsed
RTK/IMU JSONL files from parse_rtk_imu_v2.py.  Raw .rscap files remain the
traceability source; this script never modifies them.
"""

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


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("--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 100 Hz 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 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), ("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 main() -> int:
    args = parse_args()
    if args.out.exists() and any(args.out.iterdir()) and not args.overwrite:
        raise FileExistsError(f"{args.out} is non-empty; pass --overwrite")
    frames_out = args.out / "frames"
    frames_out.mkdir(parents=True, exist_ok=True)

    rtk_rows = load_jsonl(args.rtk)
    gga = sorted(
        [row for row in rtk_rows if row.get("type") == "GGA" and row.get("checksum_valid") and row.get("lat_deg") is not None],
        key=lambda row: int(row["host_receive_utc_ns"]),
    )
    heading = sorted(
        [row for row in rtk_rows if row.get("type") == "UNIHEADINGA" and row.get("checksum_valid") and row.get("heading_valid")],
        key=lambda row: int(row["host_receive_utc_ns"]),
    )
    imu = estimate_imu_times(load_jsonl(args.imu))
    gga_times = np.asarray([int(row["host_receive_utc_ns"]) for row in gga], dtype=np.int64)
    heading_times = np.asarray([int(row["host_receive_utc_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(args.rtk_max_dt_ms * 1_000_000)
    before_ns = int(args.imu_before_ms * 1_000_000)
    after_ns = int(args.imu_after_ms * 1_000_000)

    for segment_name, frame_dir in args.lidar:
        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_time_ns = int(scalar(values["unix_time_ns"]))

            gga_index = nearest_index(gga_times, lidar_time_ns)
            heading_index = nearest_index(heading_times, lidar_time_ns)
            gga_row = gga[gga_index] if gga_index >= 0 else None
            heading_row = heading[heading_index] if heading_index >= 0 else None
            gga_dt = int(gga_times[gga_index]) - lidar_time_ns if gga_index >= 0 else None
            heading_dt = int(heading_times[heading_index]) - lidar_time_ns if heading_index >= 0 else None
            gga_ok = gga_row is not None and abs(gga_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", gga_row if gga_ok else None, gga_dt)
            add_rtk(values, "rtk_heading", heading_row if heading_ok else None, heading_dt)
            initialize_rtk_measurements(values)

            if gga_ok and gga_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([gga_row.get(key, default)], dtype=dtype)
                values["rtk_gga_satellites"] = np.asarray([gga_row.get("satellites", -1)], dtype=np.int32)
                values["rtk_fixed"] = np.asarray([int(gga_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), ("heading_stddev_deg", np.float64, np.nan),
                    ("pitch_stddev_deg", np.float64, np.nan),
                    ("solution_satellites", np.int32, -1),
                ):
                    values[f"rtk_{key}"] = np.asarray([heading_row.get(key, default)], 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, args.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 args.imu], 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(args.out)),
                "source_lidar": str(source.resolve()),
                "lidar_time_ns": lidar_time_ns,
                "rtk_gga_dt_ns": gga_dt,
                "rtk_heading_dt_ns": heading_dt,
                "rtk_valid": gga_ok,
                "heading_valid": heading_ok,
                "rtk_fix_quality": gga_row.get("fix_quality") if gga_ok and gga_row else None,
                "rtk_fixed": bool(gga_ok and gga_row and int(gga_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 (args.out / "manifest.csv").open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(manifest)
    summary = {
        "frames": len(manifest),
        "segments": {name: sum(row["segment"] == name for row in manifest) for name, _ in args.lidar},
        "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": args.rtk_max_dt_ms,
        "imu_window_ms": [-args.imu_before_ms, args.imu_after_ms],
        "time_basis": "LiDAR and serial host UTC; RTK GNSS time and IMU device time are retained for clock-model refinement",
        "imu_orientation_warning": "IMU values are in the raw IMU sensor frame; no LiDAR/body extrinsic is applied",
    }
    (args.out / "dataset_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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