#!/usr/bin/env python3
"""
9_evaluateMarker.py
===================
Compare reconstructed markers against ground-truth (render_*.json).

Reports, per link and overall:
  * 3D position error (mm)
  * orientation error (deg) between measured and GT normal — only meaningful
    when the detected file carries a MEASURED normal (aruco_marker_poses.json).

Backwards compatible: the original positional CLI
    python 9_evaluateMarker.py detected.json original.json
still works and prints the familiar summary; --out adds a JSON report.
"""
from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from typing import Any, Dict, List, Optional


def dist3(a, b) -> float:
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


def angle_deg(a, b) -> Optional[float]:
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    if na < 1e-9 or nb < 1e-9:
        return None
    c = sum(x * y for x, y in zip(a, b)) / (na * nb)
    c = max(-1.0, min(1.0, c))
    ang = math.degrees(math.acos(c))
    return min(ang, 180.0 - ang)  # flip-invariant


def _pct(values: List[float], q: float) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    idx = max(0, min(len(s) - 1, int(q * len(s)) - 1))
    return s[idx]


def analyze(detected_file: str, original_file: str) -> Dict[str, Any]:
    detected = json.load(open(detected_file, "r", encoding="utf-8"))
    original = json.load(open(original_file, "r", encoding="utf-8"))

    det_markers = detected.get("markers", detected if isinstance(detected, list) else [])
    det_by_id = {int(m.get("marker_id", m.get("id", -1))): m for m in det_markers}
    orig_by_id = {int(m["id"]): m for m in original}

    det_ids = set(det_by_id) - {-1}
    orig_ids = set(orig_by_id)
    recognized = det_ids & orig_ids
    missing = orig_ids - det_ids

    per_link_pos: Dict[str, List[float]] = defaultdict(list)
    per_link_ang: Dict[str, List[float]] = defaultdict(list)
    rows = []
    for mid in sorted(recognized):
        o = orig_by_id[mid]
        d = det_by_id[mid]
        link = o.get("link", d.get("link", "?"))
        pos_err = dist3(o["position_m"], d["position_m"]) * 1000.0  # mm
        per_link_pos[link].append(pos_err)
        ang_err = None
        if o.get("normal") is not None and d.get("normal") is not None:
            ang_err = angle_deg(o["normal"], d["normal"])
            if ang_err is not None:
                per_link_ang[link].append(ang_err)
        rows.append({"marker_id": mid, "link": link, "pos_err_mm": pos_err, "normal_err_deg": ang_err})

    all_pos = [r["pos_err_mm"] for r in rows]
    all_ang = [r["normal_err_deg"] for r in rows if r["normal_err_deg"] is not None]

    return {
        "n_original": len(orig_ids),
        "n_recognized": len(recognized),
        "n_missing": len(missing),
        "recognition_rate": (len(recognized) / len(orig_ids) * 100.0) if orig_ids else 0.0,
        "per_link": {
            ln: {
                "n": len(per_link_pos[ln]),
                "pos_mean_mm": sum(per_link_pos[ln]) / len(per_link_pos[ln]),
                "pos_max_mm": max(per_link_pos[ln]),
                "normal_mean_deg": (sum(per_link_ang[ln]) / len(per_link_ang[ln])) if per_link_ang.get(ln) else None,
                "normal_max_deg": max(per_link_ang[ln]) if per_link_ang.get(ln) else None,
            } for ln in sorted(per_link_pos, key=lambda k: -len(per_link_pos[k]))
        },
        "overall": {
            "pos_mean_mm": (sum(all_pos) / len(all_pos)) if all_pos else 0.0,
            "pos_p90_mm": _pct(all_pos, 0.9),
            "pos_max_mm": max(all_pos) if all_pos else 0.0,
            "normal_mean_deg": (sum(all_ang) / len(all_ang)) if all_ang else None,
            "normal_p90_deg": _pct(all_ang, 0.9) if all_ang else None,
            "normal_max_deg": max(all_ang) if all_ang else None,
        },
        "rows": rows,
    }


def print_report(r: Dict[str, Any]) -> None:
    # familiar summary (backwards compatible wording)
    print(f"Erkannte Marker: {r['n_recognized']}")
    print(f"Nicht erkannte Marker: {r['n_missing']}")
    print(f"Gesamtzahl der Original-Marker: {r['n_original']}")
    print(f"Erkennungsrate: {r['recognition_rate']:.2f}%")
    o = r["overall"]
    print(f"Gemittelter 3D-Abstand: {o['pos_mean_mm']/1000.0:.4f}m")
    print(f"90%-Radius: {o['pos_p90_mm']/1000.0:.4f}m")
    print(f"Schlechtester Abstand: {o['pos_max_mm']/1000.0:.4f}m")
    if o["normal_mean_deg"] is not None:
        print(f"Normalen-Fehler: mean {o['normal_mean_deg']:.2f}deg / p90 {o['normal_p90_deg']:.2f}deg / max {o['normal_max_deg']:.2f}deg")

    print("\nPro Glied:")
    print(f"  {'link':>10} | {'n':>3} | {'pos mean/max [mm]':>18} | {'normal mean/max [deg]':>21}")
    print("  " + "-" * 60)
    for ln, s in r["per_link"].items():
        nd = f"{s['normal_mean_deg']:6.2f} /{s['normal_max_deg']:6.2f}" if s["normal_mean_deg"] is not None else "      -          "
        print(f"  {ln:>10} | {s['n']:>3} | {s['pos_mean_mm']:7.2f} /{s['pos_max_mm']:7.2f}     | {nd:>21}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Analysiert die Markererkennung (Position + Orientierung).")
    ap.add_argument("detected_file", help="aruco_marker_poses.json oder aruco_positions_*.json")
    ap.add_argument("original_file", help="Ground-Truth render_*.json")
    ap.add_argument("--out", default=None, help="optional JSON-Report")
    args = ap.parse_args()

    r = analyze(args.detected_file, args.original_file)
    if r["n_recognized"] == 0:
        print("Keine gemeinsamen Marker gefunden, um die Genauigkeit zu bewerten.")
        return
    print_report(r)
    if args.out:
        json.dump(r, open(args.out, "w", encoding="utf-8"), indent=2)
        print(f"\n[INFO] wrote {args.out}")


if __name__ == "__main__":
    main()