"""In-process A/B of the served call: blocking vs non-blocking trace replay (TtHamer.nonblocking_replay),
interleaved in blocks so host-load drift hits both equally.  usage: ab_replay.py [N]"""
import os, sys, time, json, statistics
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
os.environ.setdefault("TT_FUSED_HASH", "0")
import torch
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
from hamer.preprocess import preprocess
N = int(sys.argv[1]) if len(sys.argv) > 1 else 400
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")
media = _REPO.parent / "media"
x, _ = preprocess(media / "sample1.jpg", bbox_xyxy=(20, 170, 295, 260), is_right=True)
tt = build_from_reference(ref, device_id=0, strict=True)
res = {True: [], False: []}
with torch.no_grad():
    tt(x); tt(x)
    o = {m: None for m in res}
    for blk in range(N // 20):
        for mode in (True, False) if blk % 2 == 0 else (False, True):
            tt.nonblocking_replay = mode
            for _ in range(10):
                t0 = time.perf_counter(); out = tt(x); res[mode].append((time.perf_counter() - t0) * 1e3)
                o[mode] = out
print("identical outputs:", torch.equal(o[True], o[False]))
for m, v in res.items():
    print(f"AB nonblocking={m}: median {statistics.median(v):.3f} ms  min {min(v):.3f}  p90 {sorted(v)[int(0.9 * len(v))]:.3f}  n={len(v)}")
tt.close()