#!/usr/bin/env python3
"""
eval_pose.py
============
Compare estimated joint angles (robot_state.json) against ground truth
(simulation/SceneX/pose.json -> "position").

Per-joint error:
  revolute (y,z,a,b,c): angular error in degrees, wrap-aware (179 vs -179 = 2deg)
  linear   (x,e):       error in millimetres

Prints a table and optionally writes a JSON summary. Returns nonzero if any
observable joint exceeds a tolerance (for scripted regression checks).
"""
from __future__ import annotations

import argparse
import json
import sys
from typing import Any, Dict

LINEAR = {"x", "e"}
JOINTS = ["x", "y", "z", "a", "b", "c", "e"]


def load_estimate(path: str) -> Dict[str, Dict[str, Any]]:
    d = json.load(open(path, "r", encoding="utf-8"))
    mv = d.get("movements", {}) or {}
    out: Dict[str, Dict[str, Any]] = {}
    for v in JOINTS:
        e = mv.get(v, {})
        # tolerate several historical schemas
        val = e.get("value", e.get("value_mm", e.get("value_deg")))
        out[v] = {
            "value": float(val) if val is not None else 0.0,
            "observable": bool(e.get("observable", True)),
            "n_markers": int(e.get("n_markers", -1)),
        }
    return out


def load_gt(path: str) -> Dict[str, float]:
    d = json.load(open(path, "r", encoding="utf-8"))
    pos = d.get("position", d)
    return {v: float(pos[v]) for v in JOINTS if v in pos}


def joint_error(v: str, est: float, gt: float) -> float:
    if v in LINEAR:
        return abs(est - gt)
    return abs(((est - gt + 180.0) % 360.0) - 180.0)


def evaluate(estimate_path: str, gt_path: str) -> Dict[str, Any]:
    est = load_estimate(estimate_path)
    gt = load_gt(gt_path)

    rows = []
    ang_errs, lin_errs = [], []
    for v in JOINTS:
        if v not in gt:
            continue
        e = est.get(v, {"value": 0.0, "observable": False, "n_markers": -1})
        err = joint_error(v, e["value"], gt[v])
        unit = "mm" if v in LINEAR else "deg"
        rows.append({"joint": v, "estimate": e["value"], "gt": gt[v], "error": err,
                     "unit": unit, "observable": e["observable"], "n_markers": e["n_markers"]})
        if e["observable"]:
            (lin_errs if v in LINEAR else ang_errs).append(err)

    summary = {
        "n_joints": len(rows),
        "mean_abs_deg": (sum(ang_errs) / len(ang_errs)) if ang_errs else None,
        "max_abs_deg": max(ang_errs) if ang_errs else None,
        "mean_abs_mm": (sum(lin_errs) / len(lin_errs)) if lin_errs else None,
        "max_abs_mm": max(lin_errs) if lin_errs else None,
    }
    return {"rows": rows, "summary": summary}


def main() -> int:
    ap = argparse.ArgumentParser(description="Evaluate estimated joint angles vs ground truth")
    ap.add_argument("estimate", help="robot_state.json")
    ap.add_argument("gt", help="simulation/SceneX/pose.json")
    ap.add_argument("--out", default=None)
    ap.add_argument("--tolDeg", type=float, default=2.0)
    ap.add_argument("--tolMm", type=float, default=3.0)
    args = ap.parse_args()

    res = evaluate(args.estimate, args.gt)
    print(f"{'joint':>6} | {'est':>9} | {'gt':>9} | {'error':>9} | obs | nMk")
    print("-" * 58)
    worst = 0.0
    for r in res["rows"]:
        flag = " " if r["observable"] else "U"
        print(f"{r['joint']:>6} | {r['estimate']:9.2f} | {r['gt']:9.2f} | "
              f"{r['error']:7.2f}{r['unit']:>2} | {flag:>3} | {r['n_markers']:>3}")
    s = res["summary"]
    print("-" * 58)
    md = f"{s['mean_abs_deg']:.2f}" if s["mean_abs_deg"] is not None else "-"
    xd = f"{s['max_abs_deg']:.2f}" if s["max_abs_deg"] is not None else "-"
    mm = f"{s['mean_abs_mm']:.2f}" if s["mean_abs_mm"] is not None else "-"
    xm = f"{s['max_abs_mm']:.2f}" if s["max_abs_mm"] is not None else "-"
    print(f"angles: mean {md}deg / max {xd}deg   |   linear: mean {mm}mm / max {xm}mm")

    if args.out:
        json.dump(res, open(args.out, "w", encoding="utf-8"), indent=2)
        print(f"[INFO] wrote {args.out}")

    over = [r for r in res["rows"] if r["observable"] and
            r["error"] > (args.tolMm if r["joint"] in LINEAR else args.tolDeg)]
    return 1 if over else 0


if __name__ == "__main__":
    sys.exit(main())