#!/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

With --robot, additionally computes FK-based position errors (mm) at the points
in TRACK_POINTS — the euclidean distance between estimated and true world
position of a point on the robot:
  wrist_error_mm  : Hand origin   (depends only on arm joints x,y,z,a)
  finger_error_mm : FingerA tip   (depends on the full chain x..e)
A point whose chain contains an UNOBSERVABLE joint yields None (the true
position is simply unknown) rather than a misleading value. n_unobservable
counts how many of the 7 joints were unobservable.

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 pathlib import Path
from typing import Any, Dict, Optional

import numpy as np

sys.path.insert(0, str(Path(__file__).parent.parent / "pipeline"))
from robot_fk import RobotFK  # noqa: E402

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

# Messpunkte entlang der Kette: name -> (link, lokaler Offset in mm).
# wrist  = Hand-Ursprung (hängt nur von den Armgelenken x,y,z,a ab)
# finger = FingerA-Skelettspitze (hängt von der ganzen Kette ab)
TRACK_POINTS = {
    "wrist":  ("Hand",    [0.0, 0.0, 0.0]),
    "finger": ("FingerA", [0.0, -60.0, 0.0]),
}


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 _point_world(fk: RobotFK, vals: Dict[str, float],
                 link: str, local: list) -> np.ndarray:
    """Weltposition eines lokalen Punktes auf einem Link (mm)."""
    T = fk.compute(vals)
    h = np.array([local[0], local[1], local[2], 1.0])
    return (T[link] @ h)[:3]


def _dependent_joints(fk: RobotFK, gt_vals: Dict[str, float],
                      link: str, local: list,
                      eps: float = 5.0, thresh: float = 1e-6) -> set:
    """Gelenke, die diesen Punkt bewegen — numerisch per Perturbation am GT-Zustand.

    Robust für beliebige Robotermodelle: ein Gelenk zählt als abhängig, wenn
    eine kleine Auslenkung den Punkt messbar verschiebt.
    """
    base = _point_world(fk, gt_vals, link, local)
    deps = set()
    for j in JOINTS:
        v = dict(gt_vals)
        v[j] = v[j] + eps
        if float(np.linalg.norm(_point_world(fk, v, link, local) - base)) > thresh:
            deps.add(j)
    return deps


def point_error_mm(fk: RobotFK, est_vals: Dict[str, float], gt_vals: Dict[str, float],
                   observable: Dict[str, bool], link: str, local: list) -> Optional[float]:
    """Euklidischer Positionsfehler eines Punktes (mm).

    Gibt None zurück, wenn irgendein Gelenk, von dem der Punkt abhängt,
    unbeobachtbar ist — dann ist die wahre Position schlicht unbekannt.
    """
    deps = _dependent_joints(fk, gt_vals, link, local)
    if any(not observable.get(j, False) for j in deps):
        return None
    p_est = _point_world(fk, est_vals, link, local)
    p_gt  = _point_world(fk, gt_vals, link, local)
    return float(np.linalg.norm(p_est - p_gt))


def evaluate(estimate_path: str, gt_path: str,
             robot_path: Optional[str] = None) -> 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,
        "n_unobservable": sum(1 for v in JOINTS if not est.get(v, {}).get("observable", False)),
    }

    # FK-basierte Positionsfehler (mm) je Messpunkt — nur wenn robot.json gegeben
    if robot_path:
        fk = RobotFK.from_file(robot_path)
        est_vals = {v: est[v]["value"] for v in JOINTS}
        gt_vals  = {v: gt.get(v, 0.0) for v in JOINTS}
        observable = {v: bool(est.get(v, {}).get("observable", False)) for v in JOINTS}
        for name, (link, local) in TRACK_POINTS.items():
            summary[f"{name}_error_mm"] = point_error_mm(
                fk, est_vals, gt_vals, observable, link, local)

    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("--robot", default=None,
                    help="robot.json — aktiviert FK-basierten Fingerfehler in mm")
    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, robot_path=args.robot)
    print(f"{'joint':>6} | {'est':>9} | {'gt':>9} | {'error':>9} | obs | nMk")
    print("-" * 58)
    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")
    print(f"unobservable joints: {s.get('n_unobservable', 0)}")
    for name in TRACK_POINTS:
        pe = s.get(f"{name}_error_mm")
        txt = f"{pe:.2f} mm" if pe is not None else "n/a (Gelenk unbeobachtbar)"
        print(f"{name:>6} position error: {txt}")

    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())