#!/usr/bin/env python3
"""Prepare one static LiDAR frame and one yaw-only RTK body pose per NPZ segment."""

from __future__ import annotations

import argparse
import csv
import json
import math
import re
import shutil
from pathlib import Path
from typing import Any

import numpy as np

POSE_FIELDS = ["time", "x", "y", "z", "qx", "qy", "qz", "qw"]


def natural_key(value: str) -> list[Any]:
    return [int(part) if part.isdigit() else part.lower() for part in re.split(r"(\d+)", value)]


def truth(value: Any) -> bool:
    return str(value).strip().lower() in {"1", "true", "yes", "y"}


def circular_mean_deg(values: np.ndarray) -> float:
    radians = np.deg2rad(values)
    return float(np.rad2deg(math.atan2(np.mean(np.sin(radians)), np.mean(np.cos(radians)))) % 360.0)


def circular_std_deg(values: np.ndarray) -> float:
    radians = np.deg2rad(values)
    resultant = max(math.hypot(np.mean(np.cos(radians)), np.mean(np.sin(radians))), 1e-12)
    return float(np.rad2deg(math.sqrt(-2.0 * math.log(resultant))))


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)
    sin_lat, cos_lat, sin_lon, cos_lon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon)
    n = a / math.sqrt(1.0 - e2 * sin_lat * sin_lat)
    return np.array([(n + height_m) * cos_lat * cos_lon, (n + height_m) * cos_lat * sin_lon,
                     (n * (1.0 - e2) + height_m) * sin_lat], 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)
    rotation = np.array([[-slon, clon, 0.0], [-slat * clon, -slat * slon, clat],
                         [clat * clon, clat * slon, slat]], dtype=float)
    return rotation @ (ecef - origin)


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


def scalar(data: np.lib.npyio.NpzFile, name: str) -> float:
    return float(np.asarray(data[name]).reshape(-1)[0])


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--combined-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--pose-name", default="rear_gga_raw_rear_to_front")
    parser.add_argument("--heading-offset-deg", type=float, required=True)
    parser.add_argument("--antenna-lever", type=float, nargs=3, required=True, metavar=("X", "Y", "Z"))
    parser.add_argument("--accepted-fixes", type=int, nargs="+", default=[4, 5])
    parser.add_argument("--heading-std-limit-deg", type=float, default=0.5)
    parser.add_argument("--min-stations", type=int, default=30)
    parser.add_argument("--expected-stations", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    manifest_path = args.combined_root / "manifest.csv"
    with manifest_path.open("r", encoding="utf-8-sig", newline="") as stream:
        rows = list(csv.DictReader(stream))
    required = {"segment", "output", "lidar_time_ns", "rtk_valid", "heading_valid", "rtk_fix_quality"}
    if not rows or not required.issubset(rows[0]):
        raise ValueError(f"{manifest_path} is empty or lacks {sorted(required)}")
    groups: dict[str, list[dict[str, str]]] = {}
    for row in rows:
        groups.setdefault(row["segment"], []).append(row)

    selected, summaries, rejected = [], [], []
    accepted_fixes = set(args.accepted_fixes)
    for segment in sorted(groups, key=natural_key):
        group = sorted(groups[segment], key=lambda row: int(row["lidar_time_ns"]))
        good = [row for row in group if truth(row["rtk_valid"]) and truth(row["heading_valid"])
                and int(row["rtk_fix_quality"]) in accepted_fixes]
        if not good:
            rejected.append({"station": segment, "reason": "no associated fixed RTK position and valid heading"})
            continue
        samples = []
        for row in good:
            path = args.combined_root / Path(row["output"])
            with np.load(path, allow_pickle=False) as data:
                samples.append((scalar(data, "rtk_lat_deg"), scalar(data, "rtk_lon_deg"),
                                scalar(data, "rtk_altitude_m"), scalar(data, "rtk_raw_heading_deg"),
                                scalar(data, "rtk_pitch_deg"), scalar(data, "rtk_heading_stddev_deg")))
        values = np.asarray(samples, dtype=float)
        heading_std = circular_std_deg(values[:, 3])
        if heading_std > args.heading_std_limit_deg:
            rejected.append({"station": segment, "reason": f"heading std {heading_std:.4f} deg exceeds limit"})
            continue
        frame = good[len(good) // 2]
        source = args.combined_root / Path(frame["output"])
        selected.append({"station": segment, "source": source, "time": int(frame["lidar_time_ns"]) / 1e9,
                         "lat": float(np.mean(values[:, 0])), "lon": float(np.mean(values[:, 1])),
                         "alt": float(np.mean(values[:, 2])), "heading": circular_mean_deg(values[:, 3])})
        summaries.append({"station": segment, "frames": len(group), "valid_fixed_frames": len(good),
                          "heading_mean_deg": circular_mean_deg(values[:, 3]),
                          "heading_circular_std_deg": heading_std, "rtk_pitch_mean_deg": float(np.mean(values[:, 4])),
                          "reported_heading_std_mean_deg": float(np.nanmean(values[:, 5])),
                          "altitude_std_m": float(np.std(values[:, 2])), "selected_source": str(source)})

    if args.expected_stations and len(selected) != args.expected_stations:
        raise RuntimeError(f"expected {args.expected_stations} usable stations, got {len(selected)}; rejected={rejected}")
    if len(selected) < args.min_stations:
        raise RuntimeError(f"need at least {args.min_stations} usable stations, got {len(selected)}; rejected={rejected}")
    if args.output.exists() and any(args.output.iterdir()) and not args.overwrite:
        raise FileExistsError(f"{args.output} is non-empty; pass --overwrite")
    frames = args.output / "frames_all"
    frames.mkdir(parents=True, exist_ok=True)
    origin = selected[0]
    origin_ecef = geodetic_to_ecef(origin["lat"], origin["lon"], origin["alt"])
    lever = np.asarray(args.antenna_lever, dtype=float)
    pose_rows = []
    for index, item in enumerate(selected, 1):
        destination = frames / f"station_{index:02d}.npz"
        shutil.copy2(item["source"], destination)
        antenna = ecef_to_enu(geodetic_to_ecef(item["lat"], item["lon"], item["alt"]), origin_ecef,
                              origin["lat"], origin["lon"])
        corrected_heading = (item["heading"] + args.heading_offset_deg) % 360.0
        yaw = math.radians(90.0 - corrected_heading)
        body = antenna - yaw_rotation(yaw) @ lever
        pose_rows.append(dict(zip(POSE_FIELDS, [item["time"], *body, 0.0, 0.0,
                                                math.sin(yaw / 2.0), math.cos(yaw / 2.0)])))
        summaries[index - 1].update({"sequence": index, "prepared_frame": destination.name,
                                     "corrected_heading_deg": corrected_heading})
    pose_path = args.output / f"body_poses_{args.pose_name}.csv"
    with pose_path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=POSE_FIELDS); writer.writeheader(); writer.writerows(pose_rows)
    with (args.output / "station_summary.csv").open("w", encoding="utf-8", newline="") as stream:
        fields = sorted({key for row in summaries for key in row})
        writer = csv.DictWriter(stream, fieldnames=fields); writer.writeheader(); writer.writerows(summaries)
    document = {"source_combined_root": str(args.combined_root.resolve()), "station_count": len(selected),
                "rejected": rejected, "pose_csv": pose_path.name,
                "selection_policy": "middle LiDAR frame among fixed-position and valid-heading associations",
                "body_pose_configuration": {"raw_heading_offset_deg": args.heading_offset_deg,
                    "antenna_lever_body_m": args.antenna_lever,
                    "orientation_model": "yaw-only, identical to the previous calibration workflow"},
                "stations": [{"sequence": i + 1, "source_station": item["station"],
                              "source_frame": str(item["source"]), "prepared_frame": f"station_{i + 1:02d}.npz"}
                             for i, item in enumerate(selected)]}
    (args.output / "manifest.json").write_text(json.dumps(document, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"prepared": str(args.output.resolve()), "stations": len(selected),
                      "rejected": rejected, "pose_csv": pose_path.name}, ensure_ascii=False, indent=2))
    return 0


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