# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""3-D joint MPJPE evaluation for the tt-nn port.

Two modes:

* **default** — torch CPU reference treated as GT.  Reports the MPJPE
  between tt-nn-predicted joints and reference-predicted joints in mm,
  matching the convention used by InterHand2.6M / FreiHAND papers.  The
  number is the joint-space analogue of the regression-vector PCC reported
  by ``scripts/eval.py``.

* **--gt-json /path/to/interhand_test.json** — when an InterHand2.6M-style
  annotation file is available, treats those 3-D joints as GT and reports
  the absolute MPJPE / PA-MPJPE per sample.  The JSON is expected to look
  like::

      [
        {"image": "...", "bbox": [x1, y1, x2, y2], "joints_3d_mm": [[...]],
         "is_right": 1},
        ...
      ]

Run from the repo root:

    python scripts/mpjpe.py \
        --ckpt /path/to/_DATA/hamer_ckpts/checkpoints/hamer.ckpt \
        --mean /path/to/_DATA/data/mano_mean_params.npz \
        --mano /path/to/_DATA/data/mano/MANO_RIGHT.pkl \
        --images /path/to/hamer/example_data
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np
import torch

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

from hamer import Hamer, load_hamer_checkpoint  # noqa: E402
from hamer.preprocess import preprocess  # noqa: E402
from hamer.mano import MANO  # noqa: E402


def split_output(out: torch.Tensor):
    rotmats = out[:, :144].reshape(-1, 16, 3, 3)
    betas = out[:, 144:154]
    cam = out[:, 154:157]
    return rotmats, betas, cam


def root_relative(joints_mm: np.ndarray) -> np.ndarray:
    """Subtract the wrist joint (joint 0) — root-relative MPJPE matches the
    reporting convention of InterHand2.6M and FreiHAND."""
    return joints_mm - joints_mm[..., :1, :]


def mpjpe(pred: np.ndarray, gt: np.ndarray) -> float:
    """Mean per-joint position error in mm — both inputs ``(J, 3)``."""
    return float(np.linalg.norm(pred - gt, axis=-1).mean())


def procrustes_align(pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
    """Similarity-align ``pred`` onto ``gt`` (Procrustes / PA-MPJPE pre-step).
    Returns the aligned ``pred``."""
    mu_p, mu_g = pred.mean(0), gt.mean(0)
    p, g = pred - mu_p, gt - mu_g
    s_p = np.linalg.norm(p)
    s_g = np.linalg.norm(g)
    p_n = p / (s_p + 1e-12)
    g_n = g / (s_g + 1e-12)
    U, _, Vt = np.linalg.svd(p_n.T @ g_n)
    R = (U @ Vt).T
    if np.linalg.det(R) < 0:
        Vt[-1] *= -1
        R = (U @ Vt).T
    s = (s_g / (s_p + 1e-12)) * np.trace((p_n.T @ g_n) @ R.T)
    return s * (pred - mu_p) @ R.T + mu_g


def parse_args():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--ckpt", type=Path, required=True)
    p.add_argument("--mean", type=Path, required=True)
    p.add_argument("--mano", type=Path, required=True)
    p.add_argument("--images", type=Path, required=True, help="dir of input images")
    p.add_argument("--bboxes", type=Path, required=True,
                   help="JSON {image_name: [x1,y1,x2,y2]} for the synthetic-GT mode")
    p.add_argument("--gt-json", type=Path, default=None,
                   help="Optional InterHand2.6M-style GT (image/bbox/joints_3d_mm). "
                        "Without this, the torch CPU reference is treated as GT.")
    p.add_argument("--device-id", type=int, default=int(os.environ.get("HAMER_DEVICE", "0")))
    return p.parse_args()


def load_gt_samples(gt_json: Path):
    return json.loads(gt_json.read_text())


def collect_demo_samples(images_dir: Path, bboxes_path: Path):
    """Build a sample list from a directory + multi-hand bboxes JSON.  GT
    joints come from the torch CPU reference at run time (the
    synthetic-GT mode)."""
    bboxes_map = json.loads(bboxes_path.read_text())
    out = []
    for name, entries in bboxes_map.items():
        p = images_dir / name
        if not p.exists():
            continue
        if isinstance(entries, dict):
            entries = [entries]
        elif entries and not isinstance(entries[0], (dict, list)):
            entries = [entries]
        for i, e in enumerate(entries):
            if isinstance(e, dict):
                bbox, is_right = e["bbox"], int(e.get("is_right", 1))
            else:
                bbox, is_right = e, 1
            out.append({
                "image": str(p),
                "bbox": list(bbox),
                "is_right": is_right,
                "joints_3d_mm": None,
                "label": f"{name}#{i}",
            })
    return out


def main() -> int:
    args = parse_args()
    print("Loading reference + checkpoint…")
    torch.manual_seed(0)
    ref = Hamer(); ref.eval()
    load_hamer_checkpoint(ref, args.ckpt, args.mean)
    mano = MANO(args.mano)

    os.environ["HAMER_USE_TT"] = "1"
    os.environ["HAMER_FORCE_TT"] = "1"
    from hamer import build_from_reference
    print("Opening tt-nn device…")
    tt_model = build_from_reference(ref, device_id=args.device_id)

    samples = load_gt_samples(args.gt_json) if args.gt_json else collect_demo_samples(args.images, args.bboxes)
    if not samples:
        print("ERROR: no samples")
        return 2

    rows = []
    for s in samples:
        x, _ = preprocess(s["image"], bbox_xyxy=tuple(s["bbox"]), is_right=bool(s.get("is_right", 1)))
        with torch.no_grad():
            ref_out = ref(x)
            tt_out = tt_model(x)
        rot_ref, betas_ref, _ = split_output(ref_out)
        rot_tt, betas_tt, _ = split_output(tt_out)
        with torch.no_grad():
            j_ref = mano.joints_21(rot_ref, betas_ref)[0].cpu().numpy() * 1000.0  # → mm
            j_tt = mano.joints_21(rot_tt, betas_tt)[0].cpu().numpy() * 1000.0

        gt = j_ref if s["joints_3d_mm"] is None else np.asarray(s["joints_3d_mm"])
        gt_rr, ref_rr, tt_rr = root_relative(gt), root_relative(j_ref), root_relative(j_tt)
        m_ref = mpjpe(ref_rr, gt_rr)                                    # 0 in synthetic-GT mode
        m_tt = mpjpe(tt_rr, gt_rr)
        m_tt_pa = mpjpe(procrustes_align(tt_rr, gt_rr), gt_rr)
        label = s.get("label") or Path(s["image"]).name
        rows.append((label, m_ref, m_tt, m_tt_pa))

    print()
    print(f"{'image':<24} {'MPJPE(ref)':>11} {'MPJPE(tt)':>10} {'PA-MPJPE(tt)':>13}    [mm]")
    print("-" * 76)
    for name, m_r, m_t, m_t_pa in rows:
        print(f"{name:<24} {m_r:>10.3f}  {m_t:>9.3f}  {m_t_pa:>12.3f}")
    print("-" * 76)
    arr = np.array([r[1:] for r in rows])
    print(f"{'mean':<24} {arr[:, 0].mean():>10.3f}  {arr[:, 1].mean():>9.3f}  {arr[:, 2].mean():>12.3f}")

    if hasattr(tt_model, "close"):
        tt_model.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())