# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model-card Python snippet on the demo frame (media/sample1.jpg, both hands).

    pip install -e code/                       # once, in an environment with ttnn / tt-metal
    python examples/quickstart.py [--device-id 0] [--mano /path/to/MANO_RIGHT.pkl]

Writes examples/out/sample1_hands.json (the predictions) and examples/out/sample1_hands.png (the boxes, and with
MANO the projected 21 joints).
"""
import argparse
import json
from pathlib import Path

from PIL import Image, ImageDraw

from hamer import HamerModel

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "examples" / "out"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--device-id", type=int, default=0)
    ap.add_argument("--mano", default=None, help="MANO_RIGHT.pkl for vertices / joints / 2-D keypoints")
    a = ap.parse_args()

    # --- the model-card snippet ---
    with HamerModel.from_pretrained(device_id=a.device_id, mano_path=a.mano) as model:
        hands = model(ROOT / "media/sample1.jpg",
                      bbox=[[20, 170, 295, 260], [130, 340, 320, 510]],   # [x1, y1, x2, y2] per hand
                      is_right=[True, False])
        for h in hands:
            print(h.is_right, h.rotmats.shape, h.betas.shape, h.cam_t_full.round(3), f"{h.timing_ms['model']:.2f} ms")
    # ---

    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "sample1_hands.json").write_text(json.dumps([h.to_dict() for h in hands], indent=1))
    img = Image.open(ROOT / "media/sample1.jpg").convert("RGB")
    draw = ImageDraw.Draw(img)
    for h in hands:
        color = (40, 200, 255) if h.is_right else (255, 120, 40)
        draw.rectangle(h.bbox, outline=color, width=3)
        draw.text((h.bbox[0] + 4, h.bbox[1] + 4), f"{'R' if h.is_right else 'L'} tz={h.cam_t_full[2]:.2f}", fill=color)
        if h.keypoints_2d is not None:
            for x, y in h.keypoints_2d:
                draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=color)
    img.save(OUT / "sample1_hands.png")
    print(f"wrote {OUT / 'sample1_hands.json'} and {OUT / 'sample1_hands.png'}")


if __name__ == "__main__":
    main()