# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""First-call latency of the Python API, one call variant per FRESH process.

    python code/scripts/first_call.py --variant right [--mano PATH] [--out .work/first_call] [-- from_pretrained kwargs as JSON]

It times ``HamerModel.from_pretrained(device_id=0)``, then 30 real calls of the variant.  Every call has a NEW real input
(the 8 demo hands of ``media/``, the box shifted by ``k`` px from call 9 on), never the warm-up input.  It prints and
writes ``<out>/<variant>.json``: boot seconds, the 1st / 2nd / 10th call, the median of calls 11..30, the
first-call penalty (1st - steady median), the median of the later calls on the same demo hand as call 1 (calls 9, 17, 25:
the crop size, and so the host time, depends on the hand), and a control: one more new call after a 0.5 s pause in the
warm process (the cost of the first call after ANY idle gap: cold CPU caches; warm-up cannot remove it).  The raw outputs of every call go to ``<out>/<variant>.npz`` so two runs
(before / after a change) can be compared bit for bit.

Variants (the device graph has ONE shape: a batch-1 256x192 crop; every variant below is a host-side code path):
  right        model(path, bbox, is_right=True)                 single right hand, JPEG file path (the quickstart)
  left         model(path, bbox, is_right=False)                single left hand (mirrored crop)
  multi        model(path, [box_a, box_b], is_right=[R, L])     two hands of one frame
  array        model(np.ndarray, bbox)                          decoded uint8 RGB frame
  tensor       model(torch.Tensor (3,H,W) uint8, bbox)
  center       model(path, bbox=None)                           centre-square crop
  mano         model(path, bbox, is_right=R/L) with MANO        mesh, joints, 2D keypoints (needs --mano)
  crops        model.forward_crops(uint8 (2,3,256,256))         upstream-style batch, B=2
  crops_float  model.forward_crops(normalized float (2,3,256,256))
  run_crop     model.run_crop(uint8 (1,3,256,192))
  big          model(path, bbox) on 1920x1252 JPEG frames (the demo frames upscaled; written to <out>/big_frames)
"""
from __future__ import annotations

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

import numpy as np

REPO = Path(__file__).resolve().parents[2]
MEDIA = REPO / "media"
VARIANTS = ("right", "left", "multi", "array", "tensor", "center", "mano", "crops", "crops_float", "run_crop", "big")
BIG_SCALE = 3.75  # "big": the demo frames upscaled 512x334 -> 1920x1252 (larger than the 640x480 warm-up frame)
N_CALLS = 30


def hands():
    boxes = json.loads((MEDIA / "bboxes.json").read_text())
    return [(MEDIA / name, e["bbox"], bool(e.get("is_right", 1))) for name in sorted(boxes) for e in boxes[name]]


def inputs(variant: str, k: int, out: Path):
    """The k-th real input (k = 0..N_CALLS-1); distinct for every k."""
    hs = hands()
    if variant == "big":
        from PIL import Image

        path, box, right = hs[k % len(hs)]
        d = k // len(hs)
        big = out / "big_frames" / path.name
        if not big.exists():
            big.parent.mkdir(parents=True, exist_ok=True)
            im = Image.open(path).convert("RGB")
            im.resize((round(im.width * BIG_SCALE), round(im.height * BIG_SCALE)), Image.BICUBIC).save(big, quality=92)
        return (big,), dict(bbox=[(v + d) * BIG_SCALE for v in box], is_right=right)
    path, box, right = hs[k % len(hs)]
    d = k // len(hs)  # shift the box by d px on later rounds: a new crop every call
    box = [box[0] + d, box[1] + d, box[2] + d, box[3] + d]
    if variant == "right":
        return (path,), dict(bbox=box, is_right=True)
    if variant == "left":
        return (path,), dict(bbox=box, is_right=False)
    if variant == "mano":
        return (path,), dict(bbox=box, is_right=right)
    if variant == "multi":
        p = sorted({h[0] for h in hs})[k % 4]
        bs = [[b[0] + d, b[1] + d, b[2] + d, b[3] + d] for q, b, _ in hs if q == p]
        return (p,), dict(bbox=bs, is_right=[r for q, _, r in hs if q == p])
    if variant == "center":
        return (hs[(k + d) % len(hs)][0],), dict(bbox=None) if k < 4 else dict(bbox=None, rescale_factor=2.0 - 0.01 * k)
    from PIL import Image

    a = np.asarray(Image.open(path).convert("RGB"))
    if variant == "array":
        return (a,), dict(bbox=box, is_right=right)
    if variant == "tensor":
        import torch

        return (torch.from_numpy(a.copy()).permute(2, 0, 1),), dict(bbox=box, is_right=right)
    from hamer.preprocess import crop_bbox

    import torch

    def crop256(i):
        p, b, r = hs[(k * 2 + i) % len(hs)]
        b = [b[0] + d, b[1] + d, b[2] + d, b[3] + d]
        c, _ = crop_bbox(Image.open(p).convert("RGB"), b)
        if not r:
            c = c.transpose(Image.FLIP_LEFT_RIGHT)
        return torch.from_numpy(np.asarray(c).copy()).permute(2, 0, 1)

    if variant == "run_crop":
        return (crop256(0)[None, :, :, 32:224].contiguous(),), {}
    u = torch.stack([crop256(0), crop256(1)])
    if variant == "crops":
        return (u,), {}
    if variant == "crops_float":
        from hamer.preprocess import normalize_raw

        return (normalize_raw(u),), {}
    raise ValueError(variant)


def flat(out) -> np.ndarray:
    """All numbers of an output, for a bit-exact comparison."""
    import torch

    from hamer import HandPrediction

    if isinstance(out, list):
        return np.concatenate([flat(o) for o in out])
    if isinstance(out, HandPrediction):
        parts = [out.rotmats, out.betas, out.cam, out.cam_t_full, out.bbox_center, [out.bbox_size, out.focal_length]]
        parts += [p for p in (out.vertices, out.joints, out.keypoints_2d) if p is not None]
        return np.concatenate([np.asarray(p, np.float64).ravel() for p in parts])
    if isinstance(out, dict):
        return np.concatenate([flat(out[k]) for k in sorted(out)])
    if isinstance(out, torch.Tensor):
        return out.detach().double().numpy().ravel()
    raise TypeError(type(out))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", required=True, choices=VARIANTS)
    ap.add_argument("--mano", default=None)
    ap.add_argument("--out", default=str(REPO / ".work" / "first_call"))
    ap.add_argument("--kwargs", default="{}", help="extra from_pretrained kwargs as JSON (e.g. warmup_variants)")
    a = ap.parse_args()
    if a.variant == "mano" and not a.mano:
        sys.exit("--variant mano needs --mano")
    real = [inputs(a.variant, k, Path(a.out)) for k in range(N_CALLS + 1)]  # host-side input preparation is not part of the timing

    t0 = time.perf_counter()
    import torch  # noqa: F401  (timed as part of a user's start-up)

    from hamer import HamerModel

    t_imp = time.perf_counter() - t0
    t0 = time.perf_counter()
    model = HamerModel.from_pretrained(device_id=0, mano_path=a.mano, **json.loads(a.kwargs))
    t_boot = time.perf_counter() - t0
    fn = {"crops": model.forward_crops, "crops_float": model.forward_crops, "run_crop": model.run_crop}.get(a.variant, model)
    ts, outs = [], []
    for args, kw in real[:N_CALLS]:
        t = time.perf_counter()
        out = fn(*args, **kw)
        ts.append((time.perf_counter() - t) * 1e3)
        outs.append(flat(out))
    time.sleep(0.5)
    t = time.perf_counter()
    fn(*real[N_CALLS][0], **real[N_CALLS][1])
    idle_ms = (time.perf_counter() - t) * 1e3
    model.close()
    steady = statistics.median(ts[10:])
    per = 4 if a.variant in ("multi", "crops", "crops_float") else 8  # inputs repeat (box shifted) with this period
    same = statistics.median(ts[per::per][:3])
    res = {"variant": a.variant, "import_s": round(t_imp, 2), "from_pretrained_s": round(t_boot, 2),
           "call1_ms": round(ts[0], 2), "call2_ms": round(ts[1], 2), "call10_ms": round(ts[9], 2),
           "steady_median_ms": round(steady, 2), "first_call_penalty_ms": round(ts[0] - steady, 2),
           "call1_over_steady": round(ts[0] / steady, 3), "same_input_median_ms": round(same, 2),
           "call1_over_same_input": round(ts[0] / same, 3), "after_idle_ms": round(idle_ms, 2),
           "after_idle_over_same_input": round(idle_ms / statistics.median(ts[N_CALLS % per + per::per]), 3),
           "warmup": model.info.get("warmup"),
           "warmup_variants": model.info.get("warmup_variants"), "chip": os.environ.get("TT_CHIP"),
           "host_load": [round(v, 1) for v in os.getloadavg()], "all_ms": [round(v, 3) for v in ts]}
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    (out / f"{a.variant}.json").write_text(json.dumps(res, indent=1))
    np.savez(out / f"{a.variant}.npz", *outs)
    print("[first_call] " + json.dumps({k: v for k, v in res.items() if k not in ("all_ms", "warmup")}), flush=True)


if __name__ == "__main__":
    main()