#!/usr/bin/env python3
"""Keep common A/B edges on which Open3D and small_gicp agree, without using X."""
import argparse
import json
from pathlib import Path

import numpy as np
from scipy.spatial.transform import Rotation


def key(meta):
    return int(meta[0]), int(meta[1])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--open3d-pairs", required=True)
    parser.add_argument("--small-pairs", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--audit")
    parser.add_argument("--max-translation", type=float, default=0.05)
    parser.add_argument("--max-rotation", type=float, default=0.50)
    parser.add_argument("--min-pairs", type=int, default=25)
    args = parser.parse_args()

    with np.load(args.open3d_pairs, allow_pickle=False) as source:
        open_a = np.asarray(source["A"], float)
        open_b = np.asarray(source["B"], float)
        open_meta = np.asarray(source["meta"], float)
        station_times = np.asarray(source["station_times"])
        rtk_dt = np.asarray(source["rtk_nearest_dt_s"])
    with np.load(args.small_pairs, allow_pickle=False) as source:
        small = {key(meta): np.asarray(b, float)
                 for meta, b in zip(source["meta"], source["B"])}

    keep, audit = [], []
    for meta, b_open in zip(open_meta, open_b):
        edge = key(meta)
        if edge not in small:
            audit.append({"i": edge[0], "j": edge[1], "accepted": False,
                          "reason": "not_in_small_gicp_refined"})
            keep.append(False)
            continue
        delta = np.linalg.inv(b_open) @ small[edge]
        translation = float(np.linalg.norm(delta[:3, 3]))
        rotation = float(np.rad2deg(Rotation.from_matrix(delta[:3, :3]).magnitude()))
        accepted = translation <= args.max_translation and rotation <= args.max_rotation
        keep.append(accepted)
        audit.append({
            "i": edge[0], "j": edge[1],
            "open3d_small_translation_m": translation,
            "open3d_small_rotation_deg": rotation,
            "accepted": accepted,
            "reason": "" if accepted else "backend_disagreement",
        })
    keep = np.asarray(keep, bool)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        output, A=open_a[keep], B=open_b[keep], meta=open_meta[keep],
        station_times=station_times, rtk_nearest_dt_s=rtk_dt,
        backend=np.asarray("open3d_gicp_cross_backend_consensus"),
    )
    audit_path = Path(args.audit or output.with_suffix(".consensus.json"))
    audit_path.write_text(json.dumps({
        "selection_is_X_independent": True,
        "B_source": "Open3D; small_gicp is used only as an agreement gate",
        "max_translation_m": args.max_translation,
        "max_rotation_deg": args.max_rotation,
        "input_open3d_pairs": len(open_b),
        "accepted_pairs": int(np.count_nonzero(keep)),
        "pairs": audit,
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    if np.count_nonzero(keep) < args.min_pairs:
        raise RuntimeError(f"only {np.count_nonzero(keep)} consensus pairs")
    print(json.dumps({"accepted_pairs": int(np.count_nonzero(keep)),
                      "output": str(output.resolve()), "audit": str(audit_path.resolve())}, indent=2))


if __name__ == "__main__":
    main()
