#!/usr/bin/env python3
"""Export paired G90/HI13 captures without replacing sensor time by host time."""

from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from collections import Counter
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.rscap_v2.hi13_imu import Hi13Sample, iter_hi13_imu_samples
from tools.time_alignment import AffineClockModel, fit_affine_clock

GPS_EPOCH_UNIX_S = datetime(1980, 1, 6, tzinfo=timezone.utc).timestamp()
DEFAULT_SOURCES = (
    ("0808", Path("D:/data/raw_serial_capture_v2"), "20260808"),
    ("0815", Path("D:/data/0815/raw_serial_capture_v2"), None),
    ("0819", Path("D:/data/0819/raw_serial_capture_v2"), None),
)
RTK_FIELDS = [
    "message_type", "t_device_s", "measurement_utc_s", "gnss_time_s",
    "gnss_week", "gnss_tow_ms", "leap_seconds", "host_receive_utc_s",
    "host_minus_measurement_s", "receive_utc_ticks", "checksum_valid",
    "position_time_utc", "lat_deg", "lon_deg", "altitude_m",
    "position_status", "position_type", "position_fixed", "fix_quality",
    "satellites", "solution_satellites", "hdop", "undulation_m",
    "lat_std_m", "lon_std_m", "altitude_std_m", "differential_age_s",
    "solution_age_s", "station_id", "heading_deg", "pitch_deg", "roll_deg",
    "heading_quality", "heading_satellites", "heading_age_s",
    "heading_station_id", "heading_valid", "baseline_length_m",
    "heading_type", "heading_solution_satellites", "velocity_status",
    "velocity_type", "doppler_velocity_valid", "velocity_latency_s",
    "velocity_age_s", "horizontal_speed_m_s", "track_ground_deg",
    "velocity_east_m_s", "velocity_north_m_s", "vertical_speed_m_s",
    "horizontal_speed_std_m_s", "vertical_speed_std_m_s", "raw_line",
]


def _capture_time(path: Path) -> datetime:
    match = re.search(r"_(\d{8}-\d{6}\.\d+)_", path.name)
    if not match:
        raise ValueError(f"capture filename has no timestamp: {path}")
    return datetime.strptime(match.group(1), "%Y%m%d-%H%M%S.%f")


def discover_pairs(
    sources: tuple[tuple[str, Path, str | None], ...],
    *,
    max_start_delta_s: float = 1.5,
) -> tuple[list[dict], list[dict]]:
    pairs: list[dict] = []
    unmatched: list[dict] = []
    for batch, folder, date_prefix in sources:
        rtk = sorted(folder.glob("wheeltec-g90*.rscap"))
        imu = sorted(folder.glob("hi13*.rscap"))
        if date_prefix:
            rtk = [path for path in rtk if date_prefix in path.name]
            imu = [path for path in imu if date_prefix in path.name]
        available = set(imu)
        for rtk_path in rtk:
            candidates = sorted(
                (
                    (abs((_capture_time(path) - _capture_time(rtk_path)).total_seconds()), path)
                    for path in available
                ),
                key=lambda item: item[0],
            )
            if not candidates or candidates[0][0] > max_start_delta_s:
                unmatched.append({"batch": batch, "rtk_rscap": str(rtk_path)})
                continue
            delta, imu_path = candidates[0]
            available.remove(imu_path)
            stamp = _capture_time(rtk_path).strftime("%Y%m%d_%H%M%S")
            pairs.append(
                {
                    "session_id": f"{batch}_{stamp}",
                    "batch_id": batch,
                    "rtk_rscap": rtk_path,
                    "imu_rscap": imu_path,
                    "capture_start_delta_s": delta,
                }
            )
    return pairs, unmatched


def _fit_clock(samples: list[Hi13Sample]) -> AffineClockModel:
    stride = max(1, len(samples) // 20000)
    selected = samples[::stride]
    device = np.asarray([sample.t_s for sample in selected], dtype=float)
    host = np.asarray(
        [utc_dotnet_ticks_to_unix_s(sample.host_receive_utc_ticks) for sample in selected],
        dtype=float,
    )
    return fit_affine_clock(device, host)


def _nmea_utc_to_unix_s(value: str, host_s: float) -> float:
    packed = float(value)
    hour = int(packed // 10000)
    minute = int((packed - hour * 10000) // 100)
    second = packed - hour * 10000 - minute * 100
    receive = datetime.fromtimestamp(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 - host_s))


def _measurement_time(row: RtkSentence) -> tuple[float, float | None]:
    fields = row.fields
    week = fields.get("gnss_week")
    tow_ms = fields.get("gnss_tow_ms")
    leap = fields.get("leap_seconds")
    if week is not None and tow_ms is not None:
        gnss_s = float(week) * 604800.0 + float(tow_ms) * 1e-3
        utc_s = GPS_EPOCH_UNIX_S + gnss_s - float(leap or 0)
        return utc_s, gnss_s
    value = fields.get("position_time_utc")
    if value is None:
        raise ValueError(f"{row.sentence_type} has no sensor measurement time")
    host_s = utc_dotnet_ticks_to_unix_s(row.receive_utc_ticks)
    return _nmea_utc_to_unix_s(str(value), host_s), None


def _format(value: object) -> object:
    if value is None:
        return ""
    if isinstance(value, bool):
        return int(value)
    return value


def _write_imu_npz(path: Path, samples: list[Hi13Sample]) -> None:
    np.savez_compressed(
        path,
        system_time_s=np.asarray([row.t_s for row in samples], dtype=np.float64),
        system_time_ms=np.asarray([row.system_time_ms for row in samples], dtype=np.uint32),
        host_receive_utc_s=np.asarray(
            [utc_dotnet_ticks_to_unix_s(row.host_receive_utc_ticks) for row in samples],
            dtype=np.float64,
        ),
        gyro_rad_s=np.asarray([row.gyro_rad_s for row in samples], dtype=np.float64),
        accel_m_s2=np.asarray([row.accel_m_s2 for row in samples], dtype=np.float64),
        rpy_deg=np.asarray([row.rpy_deg for row in samples], dtype=np.float64),
        quaternion_wxyz=np.asarray([row.quaternion_wxyz for row in samples], dtype=np.float64),
        mag_ut=np.asarray([row.mag_ut for row in samples], dtype=np.float64),
        pps_sync_stamp_ms=np.asarray([row.pps_sync_stamp_ms for row in samples], dtype=np.uint16),
        temperature_c=np.asarray([row.temperature_c for row in samples], dtype=np.int16),
        air_pressure_pa=np.asarray([row.air_pressure_pa for row in samples], dtype=np.float64),
        frame_tag=np.asarray([row.frame_tag for row in samples], dtype=np.uint8),
    )


def _write_rtk_csv(
    path: Path,
    rows: list[RtkSentence],
    clock: AffineClockModel,
) -> Counter:
    counts: Counter = Counter()
    with path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=RTK_FIELDS)
        writer.writeheader()
        for row in rows:
            host_s = utc_dotnet_ticks_to_unix_s(row.receive_utc_ticks)
            try:
                measurement_s, gnss_s = _measurement_time(row)
                t_device_s = clock.inverse(measurement_s)
                host_minus_measurement_s = host_s - measurement_s
            except (TypeError, ValueError):
                measurement_s = None
                gnss_s = None
                t_device_s = None
                host_minus_measurement_s = None
                counts[f"{row.sentence_type}_invalid_measurement_time"] += 1
            fields = dict(row.fields)
            fields.update(
                {
                    "message_type": row.sentence_type,
                    "t_device_s": t_device_s,
                    "measurement_utc_s": measurement_s,
                    "gnss_time_s": gnss_s,
                    "host_receive_utc_s": host_s,
                    "host_minus_measurement_s": host_minus_measurement_s,
                    "receive_utc_ticks": row.receive_utc_ticks,
                    "checksum_valid": row.checksum_valid,
                    "raw_line": row.raw_line,
                }
            )
            writer.writerow({name: _format(fields.get(name)) for name in RTK_FIELDS})
            counts[row.sentence_type] += 1
    return counts


def export_pair(pair: dict, output_root: Path, *, overwrite: bool) -> dict:
    destination = output_root / str(pair["session_id"])
    if destination.exists() and not overwrite:
        raise FileExistsError(f"{destination} exists; pass --overwrite")
    destination.mkdir(parents=True, exist_ok=True)
    imu_capture = read_capture(pair["imu_rscap"])
    rtk_capture = read_capture(pair["rtk_rscap"])
    imu_rows = iter_hi13_imu_samples(imu_capture)
    if len(imu_rows) < 2:
        raise ValueError(f"not enough CRC-valid HI91 rows: {pair['imu_rscap']}")
    if np.any(np.diff(np.asarray([row.t_s for row in imu_rows])) <= 0):
        raise ValueError(f"HI13 system_time is not strictly increasing: {pair['imu_rscap']}")
    clock = _fit_clock(imu_rows)
    rtk_rows = iter_g90_sentences(rtk_capture)
    _write_imu_npz(destination / "imu.npz", imu_rows)
    counts = _write_rtk_csv(destination / "rtk.csv", rtk_rows, clock)
    quaternion = np.asarray([row.quaternion_wxyz for row in imu_rows], dtype=float)
    quaternion_norm = np.linalg.norm(quaternion, axis=1)
    summary = {
        **{key: str(value) if isinstance(value, Path) else value for key, value in pair.items()},
        "time_policy": {
            "master": "HI13 system_time; RTK GNSS measurement time mapped into that clock",
            "host_receive_time": "diagnostic and affine cross-clock bridge only",
        },
        "imu_capture": file_summary(imu_capture),
        "rtk_capture": file_summary(rtk_capture),
        "imu_valid_rows": len(imu_rows),
        "imu_time_span_s": float(imu_rows[-1].t_s - imu_rows[0].t_s),
        "quaternion_norm_p01_p50_p99": [
            float(np.percentile(quaternion_norm, percentile)) for percentile in (1, 50, 99)
        ],
        "rtk_records": dict(counts),
        "clock_model_device_to_host": clock.to_dict(),
    }
    (destination / "export_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return summary


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--session", action="append")
    args = parser.parse_args(argv)
    pairs, unmatched = discover_pairs(DEFAULT_SOURCES)
    if args.session:
        selected = set(args.session)
        pairs = [pair for pair in pairs if pair["session_id"] in selected]
    args.output_root.mkdir(parents=True, exist_ok=True)
    summaries = []
    for index, pair in enumerate(pairs, 1):
        print(f"[{index}/{len(pairs)}] {pair['session_id']}", flush=True)
        summary_path = args.output_root / pair["session_id"] / "export_summary.json"
        if args.resume and summary_path.is_file():
            summaries.append(json.loads(summary_path.read_text(encoding="utf-8")))
            continue
        summaries.append(export_pair(pair, args.output_root, overwrite=args.overwrite))
    manifest = {
        "schema_version": 3,
        "session_count": len(summaries),
        "unmatched_rtk": unmatched,
        "sessions": [
            {
                "session_id": item["session_id"],
                "batch_id": item["batch_id"],
                "directory": str((args.output_root / item["session_id"]).resolve()),
                "rtk_records": item["rtk_records"],
                "imu_valid_rows": item["imu_valid_rows"],
            }
            for item in summaries
        ],
    }
    (args.output_root / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps({"output_root": str(args.output_root), **manifest}, ensure_ascii=False, indent=2))
    return 0


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