# SPDX-License-Identifier: Apache-2.0
"""Baseline perf breakdown of the served HaMeR path on one Blackhole chip.

Measures, on the REAL checkpoint and the media/ sample crops (fresh input every request,
no same-input hash cache):

  device_trace_ms   execute_trace back-to-back (non-blocking, one sync at the end) / N
                    = pure device time of the traced ViT + MANO head graph
  per-request stage timings of the served path (median of N):
      im2col        host input prep: F.unfold -> fp32 [1,1,192,768] (patch=fp32w), or just the
                    fp32 contiguous crop (patch=device: the im2col runs on device)
      h2d           ttnn.from_torch(host) + copy_host_to_device_tensor (590 KB fp32)
      exec          execute_trace(blocking=True)
      d2h           ttnn.to_torch(dec [1,1,1,109])
      finalize      host init-param add + rot6d -> rotmat
      total         sum (what one served request costs on the device side)
  PCC vs the torch CPU reference per sample (trace replay with fresh input).

Grid / dispatch selection (must be decided before ttnn creates its runtime options):
  --dispatch eth|worker   ETH dispatch (default, patched tree, 1 CQ) or Tensix-column dispatch (opt-in)
  --grid    XxY (<=12x10) cap the compute grid via TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE
                          (11x10 = p150 stock, 12x10 = p150 + ETH dispatch / Galaxy stock)
"""
from __future__ import annotations

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

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


def parse():
    p = argparse.ArgumentParser()
    p.add_argument("--dispatch", default="eth", choices=["eth", "worker"],
                   help="eth (default: p150-equivalent, 1 CQ, 12x10) | worker (Galaxy-only opt-in; 11x10 on a p150)")
    p.add_argument("--grid", default="12x10", help="12x10 (policy grid) or smaller, e.g. 11x10 = stock p150")
    p.add_argument("--iters", type=int, default=50)
    p.add_argument("--ckpt", default=None)
    p.add_argument("--mean", default=None)
    p.add_argument("--json", default=None, help="write results JSON here")
    return p.parse_args()


def main():
    a = parse()
    # POLICY: tune/report on the 12x10 grid only (13x10 Galaxy-only grid is forbidden).
    x, y = (int(v) for v in a.grid.lower().split("x"))
    if x > 12 or y > 10:
        raise SystemExit("grid larger than 12x10 is not allowed (p150 12x10 policy)")
    os.environ["TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"] = f"{x - 1},{y - 1}"
    os.environ.setdefault("TT_FUSED_HASH", "0")  # served config: fresh image every request

    import torch
    import ttnn
    from hamer import Hamer, build_from_reference, load_hamer_checkpoint
    from hamer import model as hm
    from hamer import mano_head as mh
    from hamer.fused import host_im2col
    from hamer.preprocess import preprocess

    hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
    ckpt = Path(a.ckpt) if a.ckpt else hf / "hamer.ckpt"
    mean = Path(a.mean) if a.mean else hf / "mano_mean_params.npz"
    torch.manual_seed(0)
    ref = Hamer().eval()
    load_hamer_checkpoint(ref, ckpt, mean)

    # dispatch core type: the model's own HAMER_DISPATCH switch (worker | eth)
    os.environ["HAMER_DISPATCH"] = a.dispatch

    tt = build_from_reference(ref, device_id=0, strict=True)
    g = tt.device.compute_with_storage_grid_size()
    print(f"[bench] dispatch={tt.dispatch or a.dispatch} cqs=1 grid={g.x}x{g.y} fused={tt.fused_opts.describe()}")

    media = _REPO.parent / "media"
    boxes = json.loads((media / "bboxes.json").read_text())
    xs = []
    xs_raw = []  # uint8 crops: what the server hands a TT_FUSED_INPUT=u8 model
    for name in sorted(boxes):
        e = boxes[name]
        if isinstance(e, list) and e and isinstance(e[0], dict):
            bbox, right = e[0]["bbox"], bool(e[0].get("is_right", 1))
        elif isinstance(e, dict):
            bbox, right = e["bbox"], bool(e.get("is_right", 1))
        else:
            bbox, right = e, True
        x, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right)
        xr, _ = preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=right, raw=True)
        xs.append((name, x))
        xs_raw.append(xr)

    def pcc(u, v):
        u = u.reshape(-1).float(); v = v.reshape(-1).float()
        u = u - u.mean(); v = v - v.mean()
        return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12))

    with torch.no_grad():
        refs = [ref(x) for _, x in xs]
        tt(xs[0][1])  # eager (JIT)
        tt(xs[1][1])  # trace capture + replay
        assert tt.trace_active
        pccs = {n: pcc(tt(x), r) for (n, x), r in zip(xs, refs)}
    print("[bench] PCC (trace replay, fresh input):", {k: round(v, 6) for k, v in pccs.items()})

    trace_id, dec_buf = tt._trace
    dev = tt.device
    # pure device time: back-to-back replays
    for _ in range(5):
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
    ttnn.synchronize_device(dev)
    t0 = time.perf_counter()
    for _ in range(a.iters):
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
    ttnn.synchronize_device(dev)
    dev_ms = (time.perf_counter() - t0) / a.iters * 1e3
    # 10 chunks of back-to-back replays -> median / min of the chunk means (noise on the shared host)
    chunks = []
    for _ in range(10):
        t0 = time.perf_counter()
        for _ in range(20):
            ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
        ttnn.synchronize_device(dev)
        chunks.append((time.perf_counter() - t0) / 20 * 1e3)
    dev_med, dev_min = statistics.median(chunks), min(chunks)
    print(f"[bench] device_trace_ms={dev_ms:.3f} chunk median={dev_med:.3f} min={dev_min:.3f}")
    if not tt.fused_opts.device_patch_embed:  # stage loop below assumes the fused im2col input
        print("RESULT " + json.dumps({"dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms, "pcc": pccs}))
        tt.close()
        return

    # served path, stage by stage, fresh input each request
    u8 = tt.fused_opts.input == "u8"  # served input = the uint8 crop (server: preprocess_pil(raw=True))
    hp = tt._head_params
    io = tt._io
    if io is not None:
        # round 13 fast host I/O (TtHamer._forward_fast, stage by stage): prep = centring into the persistent host
        # buffer, h2d = copy_host_to_device_tensor, exec = replay + readback into the persistent host tensor + the
        # wait (spin on a CQ event, or a blocking read with HAMER_SPIN_WAIT=0), finalize = HostFinalizer
        st = {k: [] for k in ("prep", "h2d", "exec", "finalize", "total")}
        for i in range(a.iters):
            x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
            c0 = time.perf_counter()
            tt._write_input(io, x)  # centring (round 14: + the tiled im2col with TT_FUSED_TILE_IN)
            c1 = time.perf_counter()
            ttnn.copy_host_to_device_tensor(io["host_in"], tt._tokens_dev)
            c2 = time.perf_counter()
            ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=False)
            if tt.spin_wait:
                ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=False)
                ttnn.event_synchronize(ttnn.record_event(dev, 0))
            else:
                ttnn.copy_device_to_host_tensor(dec_buf, io["host_out"], blocking=True)
            c3 = time.perf_counter()
            out = io["fin"](io["out_np"])
            c4 = time.perf_counter()
            for k, (s_, e_) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c0, c4)]):
                st[k].append((e_ - s_) * 1e3)
    else:
      st = {k: [] for k in ("im2col", "h2d", "exec", "d2h", "finalize", "total")}
      for i in range(a.iters):
        x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
        c0 = time.perf_counter()
        host, hdt, lay = tt._host_input(x)
        c1 = time.perf_counter()
        h = ttnn.from_torch(host, device=None, layout=lay, dtype=hdt,  # round 15: the input layout (TILE with TT_FUSED_TILE_IN)
                            mesh_mapper=ttnn.ReplicateTensorToMesh(dev))
        ttnn.copy_host_to_device_tensor(h, tt._tokens_dev)
        c2 = time.perf_counter()
        ttnn.execute_trace(dev, trace_id, cq_id=0, blocking=True)
        c3 = time.perf_counter()
        dec_h = tt._dec_host(ttnn.to_torch(dec_buf), 1)
        c4 = time.perf_counter()
        out = mh.host_finalize(dec_h, hp["dec_split"], tt._init_pose, tt._init_betas, tt._init_cam, init=tt._fin_init)
        c5 = time.perf_counter()
        for k, (s, e) in zip(st, [(c0, c1), (c1, c2), (c2, c3), (c3, c4), (c4, c5), (c0, c5)]):
            st[k].append((e - s) * 1e3)
    med = {k: statistics.median(v) for k, v in st.items()}
    # last output sanity vs reference
    last = (a.iters - 1) % len(xs)
    med_pcc_last = pcc(out, refs[last])
    # full TtHamer.__call__ (what the server calls), fresh input
    call = []
    with torch.no_grad():
        for i in range(a.iters):
            x = xs_raw[i % len(xs)] if u8 else xs[i % len(xs)][1]
            c0 = time.perf_counter(); tt(x); call.append((time.perf_counter() - c0) * 1e3)
    res = {
        "dispatch": a.dispatch, "grid": f"{g.x}x{g.y}", "device_trace_ms": dev_ms,
        "device_trace_chunk_median_ms": dev_med, "device_trace_chunk_min_ms": dev_min,
        "stages_ms": med, "ttHamer_call_ms": statistics.median(call), "ttHamer_call_min_ms": min(call),
        "pcc": pccs, "pcc_bench_last": med_pcc_last,
    }
    print(f"[bench] device_trace_ms={dev_ms:.3f}  (med {dev_med:.3f} min {dev_min:.3f})  served-call={res['ttHamer_call_ms']:.3f} ms (min {res['ttHamer_call_min_ms']:.3f})")
    print("[bench] stages (median ms): " + "  ".join(f"{k}={v:.3f}" for k, v in med.items()))
    print("RESULT " + json.dumps(res))
    if a.json:
        Path(a.json).write_text(json.dumps(res, indent=2))
    tt.close()


if __name__ == "__main__":
    main()