#!/usr/bin/env python3
"""X-independent second-stage filter for stationary A/B pairs."""
import argparse
import json
from pathlib import Path

import numpy as np

from rigorous_calibration import read_pairs, rotation_angle_deg


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pairs", required=True)
    parser.add_argument("--quality-json", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--audit")
    parser.add_argument("--min-pairs", type=int, default=25)
    parser.add_argument("--min-inlier-ratio", type=float, default=0.70)
    parser.add_argument("--max-inlier-rmse", type=float, default=0.13)
    parser.add_argument("--max-rotation-invariant-error", type=float, default=0.75)
    parser.add_argument("--reverse-translation-tolerance", type=float, default=0.05)
    parser.add_argument("--reverse-rotation-tolerance", type=float, default=0.50)
    args = parser.parse_args()

    a_array, b_array, meta, _ = read_pairs(args.pairs)
    quality = json.loads(Path(args.quality_json).read_text(encoding="utf-8-sig"))
    reports = {(int(item["i"]), int(item["j"])): item for item in quality["pairs"]}
    keep, audit = [], []
    for a_ij, b_ij, item_meta in zip(a_array, b_array, meta):
        key = (int(item_meta[0]), int(item_meta[1]))
        report = reports[key]
        heldout = report["heldout_symmetric"]
        reverse = report["forward_reverse"]
        invariant = abs(rotation_angle_deg(a_ij[:3, :3]) - rotation_angle_deg(b_ij[:3, :3]))
        reasons = []
        if heldout["inlier_ratio"] < args.min_inlier_ratio:
            reasons.append("overlap_ratio")
        if heldout["inlier_rmse_m"] is None or heldout["inlier_rmse_m"] > args.max_inlier_rmse:
            reasons.append("heldout_rmse")
        if invariant > args.max_rotation_invariant_error:
            reasons.append("rotation_conjugacy_invariant")
        if reverse["translation_m"] > args.reverse_translation_tolerance:
            reasons.append("forward_reverse_translation")
        if reverse["rotation_deg"] > args.reverse_rotation_tolerance:
            reasons.append("forward_reverse_rotation")
        accepted = not reasons
        keep.append(accepted)
        audit.append({
            "i": key[0], "j": key[1], "heldout_inlier_ratio": heldout["inlier_ratio"],
            "heldout_inlier_rmse_m": heldout["inlier_rmse_m"],
            "rotation_invariant_error_deg": invariant,
            "reverse_translation_m": reverse["translation_m"],
            "reverse_rotation_deg": reverse["rotation_deg"],
            "accepted": accepted, "rejection_reasons": reasons,
        })
    keep = np.asarray(keep, bool)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with np.load(args.pairs, allow_pickle=False) as source:
        np.savez_compressed(
            output, A=a_array[keep], B=b_array[keep], meta=meta[keep],
            station_times=np.asarray(source["station_times"]),
            rtk_nearest_dt_s=np.asarray(source["rtk_nearest_dt_s"]),
            backend=np.asarray(source["backend"]),
        )
    audit_path = Path(args.audit or output.with_suffix(".refinement.json"))
    audit_path.write_text(json.dumps({
        "selection_is_X_independent": True,
        "criteria": {
            "min_inlier_ratio": args.min_inlier_ratio,
            "max_inlier_rmse_m": args.max_inlier_rmse,
            "max_rotation_invariant_error_deg": args.max_rotation_invariant_error,
            "reverse_translation_tolerance_m": args.reverse_translation_tolerance,
            "reverse_rotation_tolerance_deg": args.reverse_rotation_tolerance,
        },
        "input_pairs": len(keep), "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)} refined pairs; need {args.min_pairs}")
    print(json.dumps({"input_pairs": len(keep), "accepted_pairs": int(np.count_nonzero(keep)),
                      "output": str(output.resolve()), "audit": str(audit_path.resolve())}, indent=2))


if __name__ == "__main__":
    main()