calibration/tools/export_raw_to_combined.py

#!/usr/bin/env python3
"""One-shot export: raw H32/G90/N300 captures → RTK–LiDAR ``combined/`` package.

Analogous to Lidar-IMU ``tools/export_rscap_to_v1.py``: raw ``.rscap`` in,
calibration-ready intermediate out. Downstream prepare/solve consume ``combined/``
only (``manifest.csv`` + associated frame NPZs).

Expected raw layout:

    stations/
      001/h32.rscap
      002/h32.rscap
      ...
    captures/   (paths passed explicitly)
      rtk.rscap   # G90: #PVTSLNA + #UNIHEADINGA
      imu.rscap   # N300 (associated only; not used in AX=XB)

Output under ``--out``:

    export/<station>/frames/*.npz   # internal LiDAR frames
    parsed/rtk.jsonl, imu.jsonl
    combined/frames/*.npz + manifest.csv + dataset_summary.json
    export_summary.json

Legacy dlog stations (``dobject`` + ``dobject_recording``) are still accepted;
use ``--time-basis host`` for those datasets.

Raw ``.rscap`` / dlog files are never modified.
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "rscap_v2"))

from build_multisensor_npz import build_combined  # noqa: E402
from capture_format_v2 import file_summary, read_capture  # noqa: E402
from export_h32_rscap_station import export_station_h32, resolve_lidar_rscap  # noqa: E402
from pipeline_common_corrected import (  # noqa: E402
    parse_imu_capture,
    parse_rtk_capture,
    write_json,
    write_jsonl,
)


def is_h32_station(station: Path, capture_name: str) -> bool:
    try:
        resolve_lidar_rscap(station, capture_name)
        return True
    except FileNotFoundError:
        return False


def is_dlog_station(station: Path) -> bool:
    return (station / "dobject").is_dir() and (station / "dobject_recording").is_dir()


def discover_stations(stations_root: Path, names: list[str], capture_name: str) -> list[Path]:
    if names:
        stations = [stations_root / name for name in names]
        missing = [str(path) for path in stations if not path.is_dir()]
        if missing:
            raise FileNotFoundError(f"station directories missing: {missing}")
        return stations
    stations = sorted(
        [
            path
            for path in stations_root.iterdir()
            if path.is_dir() and (is_h32_station(path, capture_name) or is_dlog_station(path))
        ],
        key=lambda path: path.name,
    )
    if not stations:
        raise FileNotFoundError(
            f"no station with {capture_name}/lidar.rscap or dobject+dobject_recording under {stations_root}"
        )
    return stations


def export_legacy_dlog_station(
    station: Path,
    out: Path,
    *,
    lidar_object: str,
    timezone: str,
    stride: int,
) -> None:
    exporter = ROOT / "frontlidar_dlog_export.py"
    command = [
        sys.executable,
        str(exporter),
        "--dlog",
        str(station),
        "--out",
        str(out),
        "--object",
        lidar_object,
        "--format",
        "npz",
        "--timezone",
        timezone,
        "--stride",
        str(stride),
        "--compress",
        "--skip-rtk",
        "--write-reports",
        "--resume",
    ]
    completed = subprocess.run(command, check=False)
    if completed.returncode != 0:
        raise RuntimeError(f"legacy dlog export failed for {station} (exit {completed.returncode})")


def parse_serial(rtk_rscap: Path, imu_rscap: Path, parsed_root: Path) -> dict[str, Any]:
    parsed_root.mkdir(parents=True, exist_ok=True)
    rtk_capture = read_capture(rtk_rscap)
    imu_capture = read_capture(imu_rscap)
    rtk_rows = parse_rtk_capture(rtk_capture)
    imu_rows = parse_imu_capture(imu_capture)
    write_jsonl(parsed_root / "rtk.jsonl", rtk_rows)
    write_jsonl(parsed_root / "imu.jsonl", imu_rows)
    summary = {
        "rtk_capture": file_summary(rtk_capture),
        "imu_capture": file_summary(imu_capture),
        "rtk_records": len(rtk_rows),
        "rtk_checksum_valid": sum(bool(row.get("checksum_valid")) for row in rtk_rows),
        "rtk_pvtslna": sum(row.get("type") == "PVTSLNA" and row.get("checksum_valid") for row in rtk_rows),
        "rtk_gga": sum(row.get("type") == "GGA" and row.get("checksum_valid") for row in rtk_rows),
        "rtk_heading_valid": sum(row.get("type") == "UNIHEADINGA" and row.get("heading_valid") for row in rtk_rows),
        "imu_frames": len(imu_rows),
        "imu_crc_valid": sum(bool(row.get("crc_valid")) for row in imu_rows),
        "imu_types": sorted({str(row.get("type")) for row in imu_rows}),
    }
    write_json(parsed_root / "parse_summary.json", summary)
    return summary


def export_raw_to_combined(
    *,
    stations_root: Path,
    rtk_rscap: Path,
    imu_rscap: Path,
    out: Path,
    station_names: list[str] | None = None,
    lidar_capture_name: str = "h32.rscap",
    lidar_object: str = "frontlidar",
    timezone: str = "+08:00",
    stride: int = 1,
    rtk_max_dt_ms: float = 150.0,
    imu_before_ms: float = 100.0,
    imu_after_ms: float = 100.0,
    time_basis: str = "device_gnss",
    overwrite: bool = False,
) -> dict[str, Any]:
    """Full raw → combined export. Returns ``export_summary`` dict."""

    if not stations_root.is_dir():
        raise FileNotFoundError(f"stations root does not exist: {stations_root}")
    if not rtk_rscap.is_file():
        raise FileNotFoundError(f"RTK capture missing: {rtk_rscap}")
    if not imu_rscap.is_file():
        raise FileNotFoundError(f"IMU capture missing: {imu_rscap}")
    if out.exists() and any(out.iterdir()) and not overwrite:
        raise FileExistsError(f"{out} is non-empty; pass --overwrite")
    if overwrite and out.exists():
        # Keep out root but clear known children so rebuild is deterministic.
        for child in ("export", "parsed", "combined", "export_summary.json", "capture_audit.json"):
            target = out / child
            if target.is_dir():
                shutil.rmtree(target)
            elif target.is_file():
                target.unlink()

    out.mkdir(parents=True, exist_ok=True)
    export_root = out / "export"
    parsed_root = out / "parsed"
    combined_root = out / "combined"

    stations = discover_stations(stations_root, station_names or [], lidar_capture_name)
    parse_summary = parse_serial(rtk_rscap, imu_rscap, parsed_root)

    station_meta: list[dict[str, Any]] = []
    lidar_segments: list[tuple[str, Path]] = []
    saw_dlog = False
    for station in stations:
        station_out = export_root / station.name
        if is_h32_station(station, lidar_capture_name):
            meta = export_station_h32(
                station,
                station_out,
                capture_name=lidar_capture_name,
                stride=stride,
                write_reports=True,
                resume=False,
            )
            kind = "h32_rscap"
        elif is_dlog_station(station):
            saw_dlog = True
            export_legacy_dlog_station(
                station,
                station_out,
                lidar_object=lidar_object,
                timezone=timezone,
                stride=stride,
            )
            meta = {"source": str(station.resolve()), "kind": "legacy_dlog"}
            kind = "legacy_dlog"
        else:
            raise RuntimeError(f"station {station.name} has neither H32 .rscap nor dlog layout")
        frames_dir = station_out / "frames"
        if not frames_dir.is_dir() or not any(frames_dir.glob("*.npz")):
            raise RuntimeError(f"no exported frames for station {station.name}: {frames_dir}")
        lidar_segments.append((station.name, frames_dir))
        station_meta.append({"station": station.name, "kind": kind, "frames_dir": str(frames_dir), **meta})

    if saw_dlog and time_basis == "device_gnss":
        print(
            "[warn] legacy dlog stations use host/DObject time; prefer --time-basis host",
            file=sys.stderr,
        )

    combined_summary = build_combined(
        lidar_segments,
        [parsed_root / "rtk.jsonl"],
        [parsed_root / "imu.jsonl"],
        combined_root,
        rtk_max_dt_ms=rtk_max_dt_ms,
        imu_before_ms=imu_before_ms,
        imu_after_ms=imu_after_ms,
        time_basis=time_basis,
        overwrite=True,
    )

    summary = {
        "role": "RTK-LiDAR one-shot raw export (like Lidar-IMU export_rscap_to_v1)",
        "stations_root": str(stations_root.resolve()),
        "rtk_rscap": str(rtk_rscap.resolve()),
        "imu_rscap": str(imu_rscap.resolve()),
        "out": str(out.resolve()),
        "station_count": len(stations),
        "stations": station_meta,
        "parsed": parse_summary,
        "combined": combined_summary,
        "outputs": {
            "combined": str(combined_root.resolve()),
            "manifest": str((combined_root / "manifest.csv").resolve()),
            "parsed": str(parsed_root.resolve()),
            "export": str(export_root.resolve()),
        },
        "timestamp_policy": {
            "default_time_basis": time_basis,
            "lidar_h32": "MSOP device timestamp → unix_time_ns",
            "rtk": "GNSS week/TOW when time_basis=device_gnss; else host_receive_utc_ns",
            "imu": "associated only; host-anchored device deltas in combined window",
            "host_utc": "kept for audit; not the default calibration timeline for new captures",
        },
        "next_step": "run/run_direct_rtk_lidar.ps1 -CombinedRoot <out>/combined ...",
    }
    (out / "export_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    return summary


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--stations-root", type=Path, required=True, help="Directory of per-station folders")
    parser.add_argument("--rtk-rscap", type=Path, required=True, help="Continuous G90/RTK V2 .rscap")
    parser.add_argument("--imu-rscap", type=Path, required=True, help="Continuous N300/IMU V2 .rscap")
    parser.add_argument("--out", type=Path, required=True, help="Output package root (contains combined/)")
    parser.add_argument("--station", action="append", default=[], help="Optional station name filter; repeatable")
    parser.add_argument("--lidar-capture-name", default="h32.rscap")
    parser.add_argument("--lidar-object", default="frontlidar", help="Legacy dlog DObject name")
    parser.add_argument("--timezone", default="+08:00", help="Legacy dlog tick timezone")
    parser.add_argument("--stride", type=int, default=1)
    parser.add_argument("--rtk-max-dt-ms", type=float, default=150.0)
    parser.add_argument("--imu-before-ms", type=float, default=100.0)
    parser.add_argument("--imu-after-ms", type=float, default=100.0)
    parser.add_argument("--time-basis", choices=("device_gnss", "host"), default="device_gnss")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.stride < 1:
        raise SystemExit("stride must be >= 1")
    summary = export_raw_to_combined(
        stations_root=args.stations_root,
        rtk_rscap=args.rtk_rscap,
        imu_rscap=args.imu_rscap,
        out=args.out,
        station_names=args.station,
        lidar_capture_name=args.lidar_capture_name,
        lidar_object=args.lidar_object,
        timezone=args.timezone,
        stride=args.stride,
        rtk_max_dt_ms=args.rtk_max_dt_ms,
        imu_before_ms=args.imu_before_ms,
        imu_after_ms=args.imu_after_ms,
        time_basis=args.time_basis,
        overwrite=args.overwrite,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"\nCombined package ready: {summary['outputs']['combined']}")
    return 0


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