"""Round 14: why the server's model call (serve_gate model p50) is slower than bench's back-to-back served call.
Times TtHamer.__call__ (fresh u8 crop each call) after different host-side gaps: none (bench), a busy-wait gap, a
sleeping gap, and the server's own work in between (decode + native preprocess of a media frame).  Run with the serve env
(eth dispatch, 12x10)."""
import json, os, statistics, sys, time, io
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
os.environ["HAMER_DISPATCH"] = "eth"
import torch
from PIL import Image
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
from hamer.preprocess import preprocess_pil

hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
ref = Hamer().eval(); load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
tt = build_from_reference(ref, device_id=0, strict=True)
media = _REPO.parent / "media"; boxes = json.loads((media / "bboxes.json").read_text())
reqs = [((media / n).read_bytes(), e["bbox"], bool(e.get("is_right", 1))) for n, es in sorted(boxes.items()) for e in es]
xs = [preprocess_pil(Image.open(io.BytesIO(r[0])), r[1], is_right=r[2], raw=True)[0] for r in reqs]
with torch.inference_mode():
    tt(xs[0]); tt(xs[1]); assert tt.trace_active
N = int(os.environ.get("N", "200"))

def run(mode, gap_ms=0.0):
    ts = []
    with torch.inference_mode():
        for i in range(N):
            if mode == "busy":
                e = time.perf_counter() + gap_ms / 1e3
                while time.perf_counter() < e: pass
            elif mode == "sleep":
                time.sleep(gap_ms / 1e3)
            elif mode == "work":
                r = reqs[i % len(reqs)]
                im = Image.open(io.BytesIO(r[0])); im.load()
                x = preprocess_pil(im, r[1], is_right=r[2], raw=True)[0]
            x = xs[i % len(xs)] if mode != "work" else x
            t0 = time.perf_counter(); tt(x); ts.append((time.perf_counter() - t0) * 1e3)
    ts = ts[10:]
    print(f"[gap] {mode:5s} gap {gap_ms:4.1f} ms: call median {statistics.median(ts):.3f} min {min(ts):.3f} p90 {sorted(ts)[int(.9*len(ts))]:.3f}", flush=True)

for rep in range(2):
    run("none"); run("busy", 1.0); run("busy", 2.0); run("busy", 5.0); run("sleep", 2.0); run("work")
tt.close()