#!/usr/bin/env python3
"""Build LiDAR GT/quality tables for a continuous LiDAR + dual-RTK + IMU run."""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import json
import math
from pathlib import Path
from typing import Any

import numpy as np


def args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--lidar-manifest", type=Path, required=True)
    p.add_argument("--rtk-jsonl", type=Path, required=True)
    p.add_argument("--imu-jsonl", type=Path, required=True)
    p.add_argument("--extrinsic", type=Path, required=True)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--max-bracket-ms", type=float, default=150.0)
    p.add_argument("--heading-std-limit-deg", type=float, default=0.5)
    return p.parse_args()


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    with path.open(encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def geodetic_to_ecef(lat_deg: float, lon_deg: float, height_m: float) -> np.ndarray:
    a, e2 = 6378137.0, 6.69437999014e-3
    lat, lon = math.radians(lat_deg), math.radians(lon_deg)
    slat, clat, slon, clon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon)
    n = a / math.sqrt(1.0 - e2 * slat * slat)
    return np.array([(n + height_m) * clat * clon,
                     (n + height_m) * clat * slon,
                     (n * (1.0 - e2) + height_m) * slat], dtype=float)


def ecef_to_enu(ecef: np.ndarray, origin: np.ndarray, lat_deg: float, lon_deg: float) -> np.ndarray:
    lat, lon = math.radians(lat_deg), math.radians(lon_deg)
    slat, clat, slon, clon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon)
    r = np.array([[-slon, clon, 0.0],
                  [-slat * clon, -slat * slon, clat],
                  [clat * clon, clat * slon, slat]], dtype=float)
    return r @ (ecef - origin)


def yaw_matrix(yaw_rad: float) -> np.ndarray:
    c, s = math.cos(yaw_rad), math.sin(yaw_rad)
    return np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=float)


def matrix_to_quat_xyzw(r: np.ndarray) -> np.ndarray:
    # Stable branch-based conversion; output convention is x,y,z,w.
    tr = float(np.trace(r))
    if tr > 0.0:
        s = math.sqrt(tr + 1.0) * 2.0
        q = np.array([(r[2, 1] - r[1, 2]) / s,
                      (r[0, 2] - r[2, 0]) / s,
                      (r[1, 0] - r[0, 1]) / s, 0.25 * s])
    else:
        i = int(np.argmax(np.diag(r)))
        if i == 0:
            s = math.sqrt(1.0 + r[0, 0] - r[1, 1] - r[2, 2]) * 2.0
            q = np.array([0.25 * s, (r[0, 1] + r[1, 0]) / s,
                          (r[0, 2] + r[2, 0]) / s, (r[2, 1] - r[1, 2]) / s])
        elif i == 1:
            s = math.sqrt(1.0 + r[1, 1] - r[0, 0] - r[2, 2]) * 2.0
            q = np.array([(r[0, 1] + r[1, 0]) / s, 0.25 * s,
                          (r[1, 2] + r[2, 1]) / s, (r[0, 2] - r[2, 0]) / s])
        else:
            s = math.sqrt(1.0 + r[2, 2] - r[0, 0] - r[1, 1]) * 2.0
            q = np.array([(r[0, 2] + r[2, 0]) / s,
                          (r[1, 2] + r[2, 1]) / s, 0.25 * s,
                          (r[1, 0] - r[0, 1]) / s])
    if q[3] < 0.0:
        q = -q
    return q / np.linalg.norm(q)


def bracket(rows: list[dict[str, Any]], times: np.ndarray, t: int,
            max_ns: int) -> tuple[dict[str, Any], dict[str, Any], float] | None:
    right = int(np.searchsorted(times, t, side="left"))
    if right == 0 or right >= len(times):
        return None
    left = right - 1
    t0, t1 = int(times[left]), int(times[right])
    if t1 <= t0 or t - t0 > max_ns or t1 - t > max_ns:
        return None
    return rows[left], rows[right], (t - t0) / (t1 - t0)


def circular_lerp_deg(a: float, b: float, u: float) -> float:
    delta = (b - a + 180.0) % 360.0 - 180.0
    return (a + u * delta) % 360.0


def iso_utc(ns: int) -> str:
    return dt.datetime.fromtimestamp(ns / 1e9, dt.timezone.utc).isoformat(timespec="microseconds")


def write_imu_csv(rows: list[dict[str, Any]], path: Path) -> None:
    fields = [
        "host_receive_utc_ns", "device_timestamp_ms", "pps_sync_stamp_ms", "crc_valid",
        "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", "temperature_c", "air_pressure_pa",
        "roll_deg", "pitch_deg", "yaw_deg",
        "quaternion_x", "quaternion_y", "quaternion_z", "quaternion_w",
        "source_chunk_sequence_first", "source_raw_file_offset",
    ]
    with path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for row in rows:
            w.writerow({key: row.get(key) for key in fields})


def main() -> int:
    a = args()
    a.out.mkdir(parents=True, exist_ok=True)
    with a.lidar_manifest.open(encoding="utf-8-sig", newline="") as f:
        lidar = [row for row in csv.DictReader(f) if not row.get("error")]
    rtk = read_jsonl(a.rtk_jsonl)
    imu = [row for row in read_jsonl(a.imu_jsonl) if row.get("crc_valid")]
    gga = sorted([r for r in rtk if r.get("type") == "GGA" and r.get("checksum_valid")
                  and r.get("lat_deg") is not None], key=lambda r: int(r["host_receive_utc_ns"]))
    heading = sorted([r for r in rtk if r.get("type") == "UNIHEADINGA" and r.get("checksum_valid")
                      and r.get("heading_valid") and r.get("raw_heading_deg") is not None],
                     key=lambda r: int(r["host_receive_utc_ns"]))
    if not lidar or len(gga) < 2 or len(heading) < 2:
        raise RuntimeError("insufficient LiDAR/GGA/heading data")

    ext = json.loads(a.extrinsic.read_text(encoding="utf-8"))
    t_r_l = np.asarray(ext["matrix_4x4"], dtype=float)
    if t_r_l.shape != (4, 4):
        raise ValueError("extrinsic matrix_4x4 must be 4x4")

    gga_times = np.asarray([int(r["host_receive_utc_ns"]) for r in gga], dtype=np.int64)
    heading_times = np.asarray([int(r["host_receive_utc_ns"]) for r in heading], dtype=np.int64)
    origin_row = next(r for r in gga if int(r.get("fix_quality", -1)) == 4)
    origin_lat, origin_lon, origin_alt = (float(origin_row[k]) for k in ("lat_deg", "lon_deg", "altitude_m"))
    origin_ecef = geodetic_to_ecef(origin_lat, origin_lon, origin_alt)
    max_ns = int(a.max_bracket_ms * 1_000_000)
    pose_rows: list[dict[str, Any]] = []

    for index, frame in enumerate(lidar):
        t = int(frame["unix_time_ns"])
        gb, hb = bracket(gga, gga_times, t, max_ns), bracket(heading, heading_times, t, max_ns)
        reasons: list[str] = []
        available = gb is not None and hb is not None
        row: dict[str, Any] = {
            "frame_index": index, "lidar_time_ns": t, "lidar_time_utc": iso_utc(t),
            "lidar_file": frame["output_file"], "point_count": frame["point_count"],
            "pose_available": int(available), "gt_valid": 0, "invalid_reason": "",
        }
        if not available:
            if gb is None: reasons.append("GGA_NOT_BRACKETED")
            if hb is None: reasons.append("HEADING_NOT_BRACKETED")
            row.update({k: "" for k in ("x_m", "y_m", "z_m", "qx", "qy", "qz", "qw",
                                                 "rtk_x_m", "rtk_y_m", "rtk_z_m", "raw_heading_deg")})
            row["invalid_reason"] = ";".join(reasons)
            pose_rows.append(row)
            continue

        g0, g1, gu = gb
        h0, h1, hu = hb
        p0 = geodetic_to_ecef(float(g0["lat_deg"]), float(g0["lon_deg"]), float(g0["altitude_m"]))
        p1 = geodetic_to_ecef(float(g1["lat_deg"]), float(g1["lon_deg"]), float(g1["altitude_m"]))
        p_rtk = ecef_to_enu((1.0 - gu) * p0 + gu * p1, origin_ecef, origin_lat, origin_lon)
        raw_heading = circular_lerp_deg(float(h0["raw_heading_deg"]), float(h1["raw_heading_deg"]), hu)
        yaw = math.radians(90.0 - raw_heading)
        t_w_r = np.eye(4)
        t_w_r[:3, :3] = yaw_matrix(yaw)
        t_w_r[:3, 3] = p_rtk
        t_w_l = t_w_r @ t_r_l
        q = matrix_to_quat_xyzw(t_w_l[:3, :3])

        fix0, fix1 = int(g0.get("fix_quality", -1)), int(g1.get("fix_quality", -1))
        sol0, sol1 = str(h0.get("heading_solution", "")), str(h1.get("heading_solution", ""))
        std0 = float(h0.get("heading_stddev_deg") or math.inf)
        std1 = float(h1.get("heading_stddev_deg") or math.inf)
        if fix0 != 4 or fix1 != 4: reasons.append("RTK_POSITION_NOT_FIXED")
        if sol0 != "NARROW_INT" or sol1 != "NARROW_INT": reasons.append("HEADING_NOT_NARROW_INT")
        if max(std0, std1) > a.heading_std_limit_deg: reasons.append("HEADING_STD_EXCEEDED")
        row.update({
            "gt_valid": int(not reasons), "invalid_reason": ";".join(reasons),
            "x_m": t_w_l[0, 3], "y_m": t_w_l[1, 3], "z_m": t_w_l[2, 3],
            "qx": q[0], "qy": q[1], "qz": q[2], "qw": q[3],
            "rtk_x_m": p_rtk[0], "rtk_y_m": p_rtk[1], "rtk_z_m": p_rtk[2],
            "raw_heading_deg": raw_heading, "yaw_enu_deg": math.degrees(yaw),
            "gga_fix_before": fix0, "gga_fix_after": fix1,
            "heading_solution_before": sol0, "heading_solution_after": sol1,
            "heading_std_max_deg": max(std0, std1),
            "gga_before_dt_ms": (t - int(g0["host_receive_utc_ns"])) / 1e6,
            "gga_after_dt_ms": (int(g1["host_receive_utc_ns"]) - t) / 1e6,
            "heading_before_dt_ms": (t - int(h0["host_receive_utc_ns"])) / 1e6,
            "heading_after_dt_ms": (int(h1["host_receive_utc_ns"]) - t) / 1e6,
        })
        pose_rows.append(row)

    fields = list(dict.fromkeys(k for row in pose_rows for k in row))
    pose_path = a.out / "lidar_gt_pose_enu.csv"
    with pose_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader(); w.writerows(pose_rows)
    write_imu_csv(imu, a.out / "imu_parsed.csv")

    summary = {
        "coordinate_convention": "T_W_L maps raw LiDAR points to local ENU; T_W_L = T_W_RTK @ T_RTK_lidar",
        "world_frame": "local ENU, origin is the first RTK FIX GGA sample",
        "rtk_frame": "x is rawHeading baseline direction projected horizontally, y left, z up",
        "orientation_model": "RTK pose is yaw-only; IMU orientation is not fused",
        "time_basis": "LiDAR and serial host UTC; no jointly estimated clock offset/drift",
        "lidar_frames": len(pose_rows),
        "pose_available_frames": sum(int(r["pose_available"]) for r in pose_rows),
        "gt_valid_frames": sum(int(r["gt_valid"]) for r in pose_rows),
        "gt_invalid_frames": sum(not int(r["gt_valid"]) for r in pose_rows),
        "imu_frames": len(imu),
        "enu_origin": {"lat_deg": origin_lat, "lon_deg": origin_lon, "altitude_m": origin_alt},
        "quality_rule": "GGA endpoints fix_quality=4, heading endpoints NARROW_INT, heading std <= limit, both streams bracket LiDAR time",
        "heading_std_limit_deg": a.heading_std_limit_deg,
        "max_bracket_ms": a.max_bracket_ms,
        "warning": "gt_valid is a quality gate, not independent proof of +/-3 cm absolute accuracy",
    }
    (a.out / "delivery_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())