# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Render every annotated hand's predicted MANO mesh on the original image.

Bbox annotations are a JSON of the form::

    {
      "image1.jpg": [
          {"bbox": [x1, y1, x2, y2], "is_right": 1},
          {"bbox": [x1, y1, x2, y2], "is_right": 0}
      ],
      ...
    }

Mirrors the upstream HaMeR demo:

  * Each hand bbox is processed independently (HaMeR is single-hand).
  * For ``is_right=0`` (left hand) the crop is mirrored horizontally
    before the model, then the predicted ``cam.x`` and ``verts.x`` are
    flipped back — see ``hamer/demo.py`` lines 136-138 / 178-179.
  * Mesh is projected onto the original frame via ``cam_crop_to_full``
    with ``focal = 5000 / 256 · max(W, H)`` (no spurious x-flip — that
    only exists in upstream because ``pyrender`` places the camera at
    ``cam_t``; here we translate the mesh in front of a fixed camera).

Run from the repo root:

    python scripts/visualize.py \
        --ckpt ... --mean ... --mano ... \
        --images media --bboxes media/bboxes.json --out media
"""
from __future__ import annotations

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

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image

_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
from hamer.rendering import overlay_mesh, project_mesh  # noqa: E402


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, help="MANO_RIGHT.pkl")
    p.add_argument("--images", type=Path, required=True, help="dir of images")
    p.add_argument("--bboxes", type=Path, required=True,
                   help="JSON: {image_name: [{bbox, is_right}, ...]}")
    p.add_argument("--out", type=Path, default=Path("media"))
    p.add_argument("--device-id", type=int, default=int(os.environ.get("HAMER_DEVICE", "0")))
    return p.parse_args()


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 normalize_entry(e):
    """Accept either ``[x1,y1,x2,y2]`` (legacy single-hand) or
    ``{bbox: [...], is_right: 0/1}``.  Returns dict with ``bbox`` +
    ``is_right``."""
    if isinstance(e, dict):
        return {"bbox": tuple(e["bbox"]), "is_right": int(e.get("is_right", 1))}
    return {"bbox": tuple(e), "is_right": 1}


def main() -> int:
    args = parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    raw_bboxes = json.loads(args.bboxes.read_text())
    bboxes = {
        name: [normalize_entry(e) for e in (entries if isinstance(entries, list) else [entries])]
        for name, entries in raw_bboxes.items()
    }

    print("Loading reference + checkpoint…")
    torch.manual_seed(0)
    ref = Hamer(); ref.eval()
    load_hamer_checkpoint(ref, args.ckpt, args.mean)
    print("Loading MANO model…")
    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)

    # Distinct colour per detection so multi-hand overlays read clearly.
    REF_COLORS = [(0.20, 0.85, 1.00), (1.00, 0.85, 0.20), (0.30, 1.00, 0.50), (1.00, 0.50, 0.85)]
    TT_COLORS = [(1.00, 0.55, 0.20), (0.55, 0.85, 1.00), (1.00, 0.20, 0.40), (0.40, 0.95, 0.65)]

    for name, hands in bboxes.items():
        img_path = args.images / name
        if not img_path.exists():
            print(f"  skip {name} — file not found")
            continue
        original = np.asarray(Image.open(img_path).convert("RGB"))
        h, w = original.shape[:2]

        ref_projs = []
        tt_projs = []
        for hand in hands:
            x, bbox_meta = preprocess(img_path, bbox_xyxy=hand["bbox"], is_right=bool(hand["is_right"]))
            cx, cy, bbox_size = bbox_meta
            with torch.no_grad():
                ref_out = ref(x).detach()
                tt_out = tt_model(x).detach()
            rot_ref, betas_ref, cam_ref = split_output(ref_out)
            rot_tt, betas_tt, cam_tt = split_output(tt_out)
            with torch.no_grad():
                verts_ref = mano(rot_ref, betas_ref)[0]
                verts_tt = mano(rot_tt, betas_tt)[0]
            ref_projs.append(project_mesh(verts_ref, cam_ref[0], (cx, cy), bbox_size, w, h, is_right=bool(hand["is_right"])))
            tt_projs.append(project_mesh(verts_tt, cam_tt[0], (cx, cy), bbox_size, w, h, is_right=bool(hand["is_right"])))

        fig, axes = plt.subplots(1, 3, figsize=(3 * w / 90, h / 90))
        axes[0].imshow(original)
        axes[0].set_title(f"input — {name}", fontsize=11)
        axes[0].axis("off")

        # Ref panel — every hand in cyan family
        axes[1].imshow(original)
        axes[1].set_xlim(0, w); axes[1].set_ylim(h, 0)
        axes[1].set_aspect("equal"); axes[1].axis("off")
        axes[1].set_title("torch CPU reference meshes", fontsize=11)
        for i, (hand, (v2d, vz)) in enumerate(zip(hands, ref_projs)):
            overlay_mesh(axes[1], original, v2d, vz, mano.faces, color=REF_COLORS[i % len(REF_COLORS)], is_right=bool(hand["is_right"]))

        # TT panel — every hand in orange family
        axes[2].imshow(original)
        axes[2].set_xlim(0, w); axes[2].set_ylim(h, 0)
        axes[2].set_aspect("equal"); axes[2].axis("off")
        axes[2].set_title("tt-nn port meshes", fontsize=11)
        for i, (hand, (v2d, vz)) in enumerate(zip(hands, tt_projs)):
            overlay_mesh(axes[2], original, v2d, vz, mano.faces, color=TT_COLORS[i % len(TT_COLORS)], is_right=bool(hand["is_right"]))

        plt.tight_layout(pad=0.5)
        out_path = args.out / f"{Path(name).stem}_result.png"
        plt.savefig(out_path, dpi=110, bbox_inches="tight")
        plt.close(fig)
        print(f"  → {out_path}")

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


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