# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Real-checkpoint, real-image PCC + numeric-drift evaluation.

Run from the repo root:

    python scripts/eval.py \
        --ckpt /path/to/_DATA/hamer_ckpts/checkpoints/hamer.ckpt \
        --mean /path/to/_DATA/data/mano_mean_params.npz \
        --images /path/to/hamer/example_data \
        --use-tt
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import List

import torch

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

import json  # noqa: E402

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


def pcc(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.detach().reshape(-1).float()
    b = b.detach().reshape(-1).float()
    a = a - a.mean()
    b = b - b.mean()
    denom = (a.norm() * b.norm()).clamp_min(1e-12)
    return float((a @ b) / denom)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--ckpt", type=Path, required=True, help="hamer.ckpt from upstream demo data")
    p.add_argument("--mean", type=Path, required=True, help="mano_mean_params.npz")
    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]}")
    p.add_argument("--use-tt", action="store_true", help="Also run the tt-nn port")
    p.add_argument("--device-id", type=int, default=int(os.environ.get("HAMER_DEVICE", "0")))
    return p.parse_args()


def main() -> int:
    args = parse_args()
    bboxes_map = json.loads(args.bboxes.read_text())
    image_paths: List[Path] = sorted(
        p for p in args.images.iterdir() if p.suffix.lower() in (".jpg", ".jpeg", ".png")
    )
    if not image_paths:
        print(f"ERROR: no images in {args.images}")
        return 2

    print("=== Loading reference HaMeR checkpoint ===")
    torch.manual_seed(0)
    ref = Hamer()
    ref.eval()
    report = load_hamer_checkpoint(ref, args.ckpt, args.mean)
    print(f"  {report['loaded']} tensors mapped, {report['skipped_upstream']} upstream keys ignored")
    if report["missing_in_ref"] or report["unexpected_in_ref"]:
        print(f"  WARN: missing={report['missing_in_ref']} unexpected={report['unexpected_in_ref']}")

    tt_model = None
    if args.use_tt:
        os.environ["HAMER_USE_TT"] = "1"
        os.environ["HAMER_FORCE_TT"] = "1"
        from hamer import build_from_reference
        print(f"=== Opening tt-nn device {args.device_id} ===")
        tt_model = build_from_reference(ref, device_id=args.device_id)

    print()
    print(f"{'image':<24} {'ref|cam|':>10} {'pcc(tt,ref)':>12} {'maxabs_pose':>12} {'maxabs_betas':>13} {'maxabs_cam':>11}")
    print("-" * 90)

    pccs = []
    for img_path in image_paths:
        entries = bboxes_map.get(img_path.name)
        if entries is None:
            continue
        # Take the first hand for each image (full multi-hand sweep is in mpjpe.py).
        if isinstance(entries, dict):
            bbox, is_right = entries["bbox"], int(entries.get("is_right", 1))
        elif isinstance(entries, list) and entries and isinstance(entries[0], dict):
            bbox, is_right = entries[0]["bbox"], int(entries[0].get("is_right", 1))
        else:
            bbox, is_right = entries, 1
        x, _ = preprocess(img_path, bbox_xyxy=tuple(bbox), is_right=bool(is_right))
        with torch.no_grad():
            ref_out = ref(x).detach()
        cam_norm = float(ref_out[:, -3:].norm())
        if tt_model is not None:
            with torch.no_grad():
                tt_out = tt_model(x).detach()
            p = pcc(tt_out, ref_out) * 100.0
            diff = (tt_out - ref_out).abs()
            mx_pose = float(diff[:, :144].max())
            mx_betas = float(diff[:, 144:154].max())
            mx_cam = float(diff[:, 154:157].max())
            pccs.append(p)
            print(f"{img_path.name:<24} {cam_norm:>10.4f} {p:>11.4f}% {mx_pose:>12.4e} {mx_betas:>13.4e} {mx_cam:>11.4e}")
        else:
            print(f"{img_path.name:<24} {cam_norm:>10.4f}")

    if pccs:
        print()
        print(f"mean PCC across {len(pccs)} images: {sum(pccs)/len(pccs):.4f}%")
        print(f"min  PCC: {min(pccs):.4f}%")

    if tt_model is not None and hasattr(tt_model, "close"):
        tt_model.close()
    return 0


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