#!/usr/bin/env python3
"""Cut G90 captures into per-window RTK CSV files.

NMEA GGA/GNHPR UTC is the measurement time. Host receive UTC is retained only
for diagnostics. The measurement UTC is mapped onto the IMU device clock with
one robust affine clock model per window.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from tools.h32_dlog.timeutil import utc_dotnet_ticks_to_unix_s
from tools.rscap_v2.capture_format_v2 import file_summary, read_capture
from tools.rscap_v2.g90_rtk import RtkSentence, iter_g90_sentences
from tools.time_alignment import AffineClockModel, fit_affine_clock

LOCAL_TZ = timezone(timedelta(hours=8))
HPR_MATCH_S = 0.08
CSV_FIELDS = [
    "t",
    "t_measurement_utc_s",
    "t_host_utc_s",
    "receive_delay_s",
    "t_local",
    "receive_utc_ticks",
    "lat_deg",
    "lon_deg",
    "altitude_m",
    "fix_quality",
    "satellites",
    "hdop",
    "heading_deg",
    "pitch_deg",
    "roll_deg",
    "heading_quality",
    "heading_satellites",
    "heading_age_s",
    "heading_station_id",
    "heading_valid",
    "gga_utc",
    "hpr_utc",
    "hpr_measurement_utc_s",
    "checksum_valid",
]


def _fmt(value) -> str:
    if value is None:
        return ""
    if isinstance(value, bool):
        return "1" if value else "0"
    if isinstance(value, float):
        if math.isnan(value):
            return ""
        return f"{value:.12g}"
    return str(value)


def load_imu_clock_model(
    imu_csv: Path,
) -> tuple[float, float, float, float, AffineClockModel]:
    """Return device/host spans and a robust ``IMU device -> host UTC`` model."""

    t_device: list[float] = []
    t_host: list[float] = []
    with imu_csv.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            t_device.append(float(row["t"]))
            t_host.append(float(row["t_host_utc_s"]))
    if not t_host:
        raise ValueError(f"empty IMU csv: {imu_csv}")
    model = fit_affine_clock(np.asarray(t_device), np.asarray(t_host))
    return min(t_device), max(t_device), min(t_host), max(t_host), model


def sentence_host_s(row: RtkSentence) -> float:
    return utc_dotnet_ticks_to_unix_s(row.receive_utc_ticks)


def nmea_utc_to_unix_s(value: str | None, receive_host_s: float) -> float:
    """Resolve NMEA ``hhmmss.s`` to the UTC day nearest host receive time."""

    if value is None or not str(value).strip():
        raise ValueError("missing NMEA UTC time")
    packed = float(value)
    hour = int(packed // 10000)
    minute = int((packed - hour * 10000) // 100)
    second = packed - hour * 10000 - minute * 100
    if not (0 <= hour < 24 and 0 <= minute < 60 and 0.0 <= second < 60.0):
        raise ValueError(f"invalid NMEA UTC time: {value!r}")
    receive = datetime.fromtimestamp(receive_host_s, tz=timezone.utc)
    midnight = datetime(
        receive.year,
        receive.month,
        receive.day,
        tzinfo=timezone.utc,
    ).timestamp()
    same_day = midnight + hour * 3600 + minute * 60 + second
    return min(
        (same_day - 86400.0, same_day, same_day + 86400.0),
        key=lambda candidate: abs(candidate - receive_host_s),
    )


def sentence_measurement_utc_s(row: RtkSentence) -> float:
    return nmea_utc_to_unix_s(
        row.fields.get("position_time_utc"),
        sentence_host_s(row),
    )


def nearest_hpr(gga: RtkSentence, hpr_rows: list[RtkSentence]) -> RtkSentence | None:
    if not hpr_rows:
        return None
    lo, hi = 0, len(hpr_rows) - 1
    target = sentence_measurement_utc_s(gga)
    while lo < hi:
        mid = (lo + hi) // 2
        if sentence_measurement_utc_s(hpr_rows[mid]) < target:
            lo = mid + 1
        else:
            hi = mid
    best = hpr_rows[lo]
    if lo > 0 and abs(sentence_measurement_utc_s(hpr_rows[lo - 1]) - target) < abs(
        sentence_measurement_utc_s(best) - target
    ):
        best = hpr_rows[lo - 1]
    if abs(sentence_measurement_utc_s(best) - target) > HPR_MATCH_S:
        return None
    return best


def merged_row(
    gga: RtkSentence,
    hpr: RtkSentence | None,
    imu_to_host: AffineClockModel,
) -> dict:
    t_host = sentence_host_s(gga)
    t_measurement = sentence_measurement_utc_s(gga)
    t_hpr = None if hpr is None else sentence_measurement_utc_s(hpr)
    local = datetime.fromtimestamp(t_measurement, tz=timezone.utc).astimezone(LOCAL_TZ)
    fields = gga.fields
    hpr_fields = hpr.fields if hpr is not None else {}
    checksum = gga.checksum_valid and (hpr is None or hpr.checksum_valid)
    return {
        "t": imu_to_host.inverse(t_measurement),
        "t_measurement_utc_s": t_measurement,
        "t_host_utc_s": t_host,
        "receive_delay_s": t_host - t_measurement,
        "t_local": local.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3],
        "receive_utc_ticks": gga.receive_utc_ticks,
        "lat_deg": fields.get("lat_deg"),
        "lon_deg": fields.get("lon_deg"),
        "altitude_m": fields.get("altitude_m"),
        "fix_quality": fields.get("fix_quality"),
        "satellites": fields.get("satellites"),
        "hdop": fields.get("hdop"),
        "heading_deg": hpr_fields.get("heading_deg"),
        "pitch_deg": hpr_fields.get("pitch_deg"),
        "roll_deg": hpr_fields.get("roll_deg"),
        "heading_quality": hpr_fields.get("heading_quality"),
        "heading_satellites": hpr_fields.get("heading_satellites"),
        "heading_age_s": hpr_fields.get("heading_age_s"),
        "heading_station_id": hpr_fields.get("heading_station_id"),
        "heading_valid": hpr_fields.get("heading_valid"),
        "gga_utc": fields.get("position_time_utc"),
        "hpr_utc": hpr_fields.get("position_time_utc"),
        "hpr_measurement_utc_s": t_hpr,
        "checksum_valid": checksum,
    }


def write_rtk_csv(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=CSV_FIELDS)
        writer.writeheader()
        for row in rows:
            writer.writerow({key: _fmt(row[key]) for key in CSV_FIELDS})


def discover_windows(sessions_root: Path) -> list[Path]:
    windows = sorted(
        path
        for path in sessions_root.iterdir()
        if path.is_dir() and (path / "imu.csv").is_file()
    )
    if not windows:
        raise SystemExit(f"no session dirs with imu.csv under {sessions_root}")
    return windows


def find_default_rscap(sessions_root: Path) -> Path:
    parents = [sessions_root, sessions_root.parent]
    matches: list[Path] = []
    for folder in parents:
        matches.extend(sorted(folder.glob("wheeltec-g90*.rscap")))
        matches.extend(sorted(folder.glob("*g90*.rscap")))
    if not matches:
        raise SystemExit(f"no G90 .rscap next to {sessions_root}")
    return matches[0]


def _delay_summary(rows: list[dict]) -> dict[str, float] | None:
    if not rows:
        return None
    values = np.asarray([row["receive_delay_s"] for row in rows], dtype=np.float64)
    return {
        "median_s": float(np.median(values)),
        "p05_s": float(np.percentile(values, 5.0)),
        "p95_s": float(np.percentile(values, 95.0)),
        "std_s": float(np.std(values)),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sessions-root", type=Path, required=True)
    parser.add_argument("--rtk-rscap", type=Path, action="append")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args(argv)
    sessions_root = args.sessions_root.resolve()
    rscaps = [path.resolve() for path in (args.rtk_rscap or [find_default_rscap(sessions_root)])]
    for rscap in rscaps:
        if not rscap.is_file():
            raise SystemExit(f"missing RTK capture: {rscap}")

    sentences: list[RtkSentence] = []
    captures_meta = []
    for rscap in rscaps:
        print(f"reading {rscap}", flush=True)
        capture = read_capture(rscap)
        print(f"chunks={len(capture.chunks)}", flush=True)
        sentences.extend(iter_g90_sentences(capture))
        captures_meta.append(file_summary(capture))
    sentences.sort(key=lambda row: row.receive_utc_ticks)
    gga_all = [row for row in sentences if row.sentence_type == "GGA"]
    hpr_all = [row for row in sentences if row.sentence_type == "GNHPR"]
    hpr_all.sort(key=sentence_measurement_utc_s)
    print(
        f"parsed sentences={len(sentences)} GGA={len(gga_all)} GNHPR={len(hpr_all)}",
        flush=True,
    )

    summaries = {
        "rtk_rscap": [str(path) for path in rscaps],
        "captures": captures_meta,
        "parsed_sentences": len(sentences),
        "gga": len(gga_all),
        "gnhpr": len(hpr_all),
        "time_source": "NMEA measurement UTC mapped through IMU device->host affine clock",
        "windows": [],
    }
    for window in discover_windows(sessions_root):
        out_csv = window / "rtk.csv"
        if out_csv.exists() and not args.overwrite:
            raise SystemExit(f"{out_csv} exists; pass --overwrite")
        dev_min, dev_max, host_min, host_max, imu_to_host = load_imu_clock_model(
            window / "imu.csv"
        )
        gga = [
            row
            for row in gga_all
            if dev_min <= imu_to_host.inverse(sentence_measurement_utc_s(row)) <= dev_max
        ]
        hpr = [
            row
            for row in hpr_all
            if dev_min - HPR_MATCH_S
            <= imu_to_host.inverse(sentence_measurement_utc_s(row))
            <= dev_max + HPR_MATCH_S
        ]
        merged = [merged_row(row, nearest_hpr(row, hpr), imu_to_host) for row in gga]
        write_rtk_csv(out_csv, merged)
        brief = {
            "window": window.name,
            "imu_host_span_s": [host_min, host_max],
            "imu_device_span_s": [dev_min, dev_max],
            "imu_device_to_host_clock": imu_to_host.to_dict(),
            "rtk_receive_delay": _delay_summary(merged),
            "gga": len(gga),
            "gnhpr_in_window": len(hpr),
            "rows_written": len(merged),
            "heading_matched": sum(1 for row in merged if row["heading_deg"] is not None),
            "fix_quality_4_or_5": sum(
                1 for row in merged if row["fix_quality"] in {4, 5}
            ),
            "out": str(out_csv),
        }
        summaries["windows"].append(brief)
        print(json.dumps(brief, ensure_ascii=False), flush=True)

    manifest = sessions_root / "rtk_export_summary.json"
    manifest.write_text(json.dumps(summaries, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"summary: {manifest}")
    return 0


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