#!/usr/bin/env python3
"""Combine independent RTK-direct hand-eye batches for a shared extrinsic.

Each batch contributes only its within-batch A/B motion pairs and LiDAR ground
planes.  No cross-batch motion pair is created, so different ENU origins and
capture locations are valid as long as every batch uses the same RTK-direct
frame definition and unchanged physical sensor installation.
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import numpy as np


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--batch-name", action="append", required=True)
    parser.add_argument("--pairs", action="append", required=True, type=Path)
    parser.add_argument("--ground-planes", action="append", required=True, type=Path)
    parser.add_argument("--output-pairs", required=True, type=Path)
    parser.add_argument("--output-ground-planes", required=True, type=Path)
    parser.add_argument("--summary", required=True, type=Path)
    return parser.parse_args()


def load_planes(path: Path, batch_name: str) -> list[dict[str, str]]:
    with path.open(encoding="utf-8-sig", newline="") as stream:
        rows = list(csv.DictReader(stream))
    if not rows:
        raise ValueError(f"no ground planes in {path}")
    for row in rows:
        for key in ("nx", "ny", "nz", "d"):
            if key not in row or row[key] in (None, ""):
                raise ValueError(f"missing {key} in {path}")
        row["source_batch"] = batch_name
    return rows


def main() -> int:
    args = parse_args()
    count = len(args.batch_name)
    if count < 2 or len(args.pairs) != count or len(args.ground_planes) != count:
        raise ValueError("provide the same number of --batch-name, --pairs, and --ground-planes (at least two)")

    pair_parts: list[dict[str, np.ndarray]] = []
    plane_rows: list[dict[str, str]] = []
    batch_summaries: list[dict[str, object]] = []
    for index, (name, pairs_path, planes_path) in enumerate(zip(args.batch_name, args.pairs, args.ground_planes)):
        with np.load(pairs_path, allow_pickle=False) as source:
            required = ("A", "B", "meta", "station_times", "rtk_nearest_dt_s")
            missing = [key for key in required if key not in source]
            if missing:
                raise ValueError(f"{pairs_path} missing {missing}")
            a = np.asarray(source["A"], float)
            b = np.asarray(source["B"], float)
            meta = np.asarray(source["meta"], float)
            times = np.asarray(source["station_times"], float)
            rtk_dt = np.asarray(source["rtk_nearest_dt_s"], float)
        if len(a) == 0 or len(a) != len(b) or len(a) != len(meta):
            raise ValueError(f"invalid A/B/meta sizes in {pairs_path}")
        pair_parts.append({"A": a, "B": b, "meta": meta, "station_times": times, "rtk_dt": rtk_dt})
        rows = load_planes(planes_path, name)
        plane_rows.extend(rows)
        batch_summaries.append({
            "name": name,
            "pairs_path": str(pairs_path.resolve()),
            "ground_planes_path": str(planes_path.resolve()),
            "pairs": len(a),
            "stations": len(times),
            "ground_planes": len(rows),
            "pair_offset": sum(item["A"].shape[0] for item in pair_parts[:-1]),
        })

    output_pairs = args.output_pairs
    output_pairs.parent.mkdir(parents=True, exist_ok=True)
    batch_index = np.concatenate([np.full(len(part["A"]), index, np.int32) for index, part in enumerate(pair_parts)])
    np.savez_compressed(
        output_pairs,
        A=np.concatenate([part["A"] for part in pair_parts]),
        B=np.concatenate([part["B"] for part in pair_parts]),
        meta=np.concatenate([part["meta"] for part in pair_parts]),
        station_times=np.concatenate([part["station_times"] for part in pair_parts]),
        rtk_nearest_dt_s=np.concatenate([part["rtk_dt"] for part in pair_parts]),
        batch_index=batch_index,
        batch_names=np.asarray(args.batch_name),
        backend=np.asarray("independent_batch_consensus"),
    )

    output_planes = args.output_ground_planes
    output_planes.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = ["nx", "ny", "nz", "d", "source_batch"]
    with output_planes.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fieldnames)
        writer.writeheader()
        for row in plane_rows:
            writer.writerow({key: row[key] for key in fieldnames})

    summary = {
        "schema_version": 1,
        "convention": "Shared T_RTK_lidar; only within-batch A_ij and B_ij are combined.",
        "batches": batch_summaries,
        "total_pairs": int(len(batch_index)),
        "total_ground_planes": len(plane_rows),
        "output_pairs": str(output_pairs.resolve()),
        "output_ground_planes": str(output_planes.resolve()),
    }
    args.summary.parent.mkdir(parents=True, exist_ok=True)
    args.summary.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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