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

from __future__ import annotations

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

import numpy as np

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Select one static LiDAR frame per exported station and rebuild yaw-only RTK body poses."
    )
    parser.add_argument("--export-root", type=Path, required=True, help="Directory containing exported station directories.")
    parser.add_argument("--output", type=Path, required=True, help="Output prepared directory.")
    parser.add_argument("--pose-name", default="rear_gga_raw_rear_to_front", help="Suffix of body_poses_<name>.csv.")
    parser.add_argument("--heading-offset-deg", type=float, required=True, help="Added to raw_heading_deg before ENU yaw conversion.")
    parser.add_argument("--antenna-lever", type=float, nargs=3, metavar=("X", "Y", "Z"), required=True, help="Antenna position in body coordinates [m].")
    parser.add_argument("--expected-stations", type=int, default=0, help="Require exactly this many usable stations; 0 disables.")
    parser.add_argument("--min-stations", type=int, default=30, help="Fail below this many usable stations.")
    parser.add_argument("--groups", nargs="*", default=None, help="Optional explicit station directory names.")
    parser.add_argument("--accepted-fixes", type=int, nargs="+", default=[4, 5], help="Accepted GGA fix values.")
    parser.add_argument("--heading-std-limit-deg", type=float, default=float("inf"), help="Reject a station above this heading circular stddev.")
    parser.add_argument("--overwrite", action="store_true", help="Allow replacing generated files in a non-empty output directory.")
    return parser.parse_args()


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


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 read_csv(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8-sig", newline="") as stream:
        return list(csv.DictReader(stream))


def require_columns(rows: list[dict[str, str]], path: Path, columns: list[str]) -> None:
    if not rows:
        raise ValueError(f"{path}: no rows")
    missing = [column for column in columns if column not in rows[0]]
    if missing:
        raise ValueError(f"{path}: missing required columns: {', '.join(missing)}")


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


def circular_std_deg(values: np.ndarray) -> float:
    radians = np.deg2rad(values)
    resultant = max(math.hypot(float(np.mean(np.cos(radians))), float(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_ecef: np.ndarray, origin_lat_deg: float, origin_lon_deg: float) -> np.ndarray:
    lat, lon = math.radians(origin_lat_deg), math.radians(origin_lon_deg)
    sin_lat, cos_lat, sin_lon, cos_lon = math.sin(lat), math.cos(lat), math.sin(lon), math.cos(lon)
    rotation = np.array([[-sin_lon, cos_lon, 0.0], [-sin_lat * cos_lon, -sin_lat * sin_lon, cos_lat], [cos_lat * cos_lon, cos_lat * sin_lon, sin_lat]], dtype=float)
    return rotation @ (ecef - origin_ecef)


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]], dtype=float)


def select_frame(station: Path) -> tuple[dict[str, str], int, int]:
    path = station / "reports" / "manifest.csv"
    rows = read_csv(path)
    require_columns(rows, path, ["status", "unix_time_ns", "output_file"])
    exported = [row for row in rows if row["status"] in {"exported", "resumed"}]
    if not exported:
        raise ValueError(f"{station.name}: no exported/resumed LiDAR frames")
    valid = [row for row in exported if truth(row.get("rtk_position_valid")) and truth(row.get("rtk_heading_valid"))]
    candidates = valid or exported
    candidates.sort(key=lambda row: int(row["unix_time_ns"]))
    return candidates[len(candidates) // 2], len(exported), len(valid)


def validation_checks(station: Path) -> dict[str, bool]:
    path = station / "reports" / "validation_report.json"
    if not path.exists():
        return {}
    document = json.loads(path.read_text(encoding="utf-8"))
    return {str(item.get("name")): bool(item.get("ok")) for item in document.get("checks", [])}


def process_station(station: Path, args: argparse.Namespace) -> tuple[dict[str, Any], dict[str, Any]]:
    rtk_path = station / "rtk" / "gps_post_z.csv"
    rows = read_csv(rtk_path)
    columns = ["lat", "lon", "alt_m", "raw_heading_deg", "fix", "position_valid", "heading_valid", "unix_time_ns"]
    require_columns(rows, rtk_path, columns)
    fixes = set(args.accepted_fixes)
    good = [row for row in rows if truth(row["position_valid"]) and truth(row["heading_valid"]) and int(row["fix"]) in fixes]
    if not good:
        raise ValueError(f"{station.name}: no RTK samples pass position/heading/fix filters")
    lat = np.asarray([float(row["lat"]) for row in good])
    lon = np.asarray([float(row["lon"]) for row in good])
    alt = np.asarray([float(row["alt_m"]) for row in good])
    heading = np.asarray([float(row["raw_heading_deg"]) for row in good])
    times = np.asarray([int(row["unix_time_ns"]) for row in good], dtype=np.int64)
    heading_std = circular_std_deg(heading)
    if heading_std > args.heading_std_limit_deg:
        raise ValueError(f"{station.name}: heading circular stddev {heading_std:.3f} deg exceeds {args.heading_std_limit_deg:.3f} deg")
    frame, frame_count, valid_frame_count = select_frame(station)
    source = station / frame["output_file"]
    if not source.is_file():
        raise FileNotFoundError(f"{station.name}: selected frame is absent: {source}")
    checks = validation_checks(station)
    selected = {"station": station.name, "source_frame": source, "frame_time": int(frame["unix_time_ns"]) / 1e9, "lat_deg": float(np.mean(lat)), "lon_deg": float(np.mean(lon)), "alt_m": float(np.mean(alt)), "raw_heading_deg": circular_mean_deg(heading), "raw_heading_std_deg": heading_std}
    summary = {"station": station.name, "selected_frame": frame["output_file"], "exported_frames": frame_count, "frames_with_valid_rtk": valid_frame_count, "valid_rtk_samples": len(good), "raw_heading_mean_deg": selected["raw_heading_deg"], "raw_heading_std_deg": heading_std, "alt_std_m": float(np.std(alt)), "rtk_span_sec": float((times.max() - times.min()) / 1e9) if len(times) > 1 else 0.0, "payload_ok": checks.get("lidar_payload_length", ""), "rtk_parse_ok": checks.get("gps_post_z_parse", ""), "time_check_3s_ok": checks.get("rtk_time_alignment", "")}
    return selected, summary


def main() -> int:
    args = parse_args()
    if args.min_stations < 2 or args.expected_stations < 0:
        raise ValueError("--min-stations must be >=2 and --expected-stations must be >=0")
    if not args.export_root.is_dir():
        raise NotADirectoryError(args.export_root)
    names = args.groups or [path.name for path in args.export_root.iterdir() if path.is_dir()]
    stations = [args.export_root / name for name in sorted(names, key=natural_key)]
    missing = [str(path) for path in stations if not path.is_dir()]
    if missing:
        raise FileNotFoundError("station directories do not exist: " + ", ".join(missing))
    selected, summaries, rejected = [], [], []
    for station in stations:
        try:
            item, summary = process_station(station, args)
            selected.append(item); summaries.append(summary)
        except (ValueError, FileNotFoundError) as error:
            rejected.append({"station": station.name, "reason": str(error)})
    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}")
    output, frames_dir = args.output, args.output / "frames_all"
    if output.exists() and any(output.iterdir()) and not args.overwrite:
        raise FileExistsError(f"{output} is non-empty; pass --overwrite to replace generated files")
    frames_dir.mkdir(parents=True, exist_ok=True)
    origin = selected[0]
    origin_ecef = geodetic_to_ecef(origin["lat_deg"], origin["lon_deg"], origin["alt_m"])
    lever = np.asarray(args.antenna_lever, dtype=float)
    pose_rows = []
    for index, item in enumerate(selected, 1):
        destination = frames_dir / f"station_{index:02d}.npz"
        shutil.copy2(item["source_frame"], destination)
        antenna_enu = ecef_to_enu(geodetic_to_ecef(item["lat_deg"], item["lon_deg"], item["alt_m"]), origin_ecef, origin["lat_deg"], origin["lon_deg"])
        corrected_heading = (item["raw_heading_deg"] + args.heading_offset_deg) % 360.0
        yaw = math.radians(90.0 - corrected_heading)
        body_position = antenna_enu - yaw_rotation(yaw) @ lever
        pose_rows.append(dict(zip(POSE_FIELDS, [item["frame_time"], *body_position, 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 = 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)
    fields = sorted({key for row in summaries for key in row})
    with (output / "station_summary.csv").open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields); writer.writeheader(); writer.writerows(summaries)
    manifest = {"source_export_root": str(args.export_root.resolve()), "station_count": len(selected), "pose_csv": pose_path.name, "frame_directory": "frames_all", "selection_policy": "middle frame among exported frames with valid RTK; falls back to middle exported frame", "rtk_filter": {"position_valid": True, "heading_valid": True, "accepted_fixes": args.accepted_fixes, "heading_std_limit_deg": args.heading_std_limit_deg}, "body_pose_configuration": {"raw_heading_offset_deg": args.heading_offset_deg, "antenna_lever_body_m": args.antenna_lever, "body_axes": "x forward, y left, z up; origin must match the supplied lever", "orientation_model": "yaw-only from raw_heading_deg; this tool does not reconstruct RTK pitch or roll"}, "enu_origin": {"lat_deg": origin["lat_deg"], "lon_deg": origin["lon_deg"], "alt_m": origin["alt_m"]}, "stations": [{"sequence": index + 1, "source_station": item["station"], "source_frame": str(item["source_frame"]), "prepared_frame": f"station_{index + 1:02d}.npz"} for index, item in enumerate(selected)]}
    (output / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"prepared": str(output.resolve()), "stations": len(selected), "rejected": rejected, "pose_csv": pose_path.name}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise