"""Model-level soak: N served calls (fresh input upload + trace replay + readback) cycling the 4 sample crops;
every output must be bit-identical to the first output for the same crop.  usage: soak.py [N] (HAMER_DISPATCH)"""
import os, sys, time, json
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 5000
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"; boxes = json.loads((media / "bboxes.json").read_text())
xs = []
for name in sorted(boxes):
    e = boxes[name]
    bbox = e[0]["bbox"] if isinstance(e, list) and e and isinstance(e[0], dict) else (e["bbox"] if isinstance(e, dict) else e)
    xs.append(preprocess(media / name, bbox_xyxy=tuple(bbox), is_right=True)[0])
tt = build_from_reference(ref, device_id=0, strict=True)
with torch.no_grad():
    tt(xs[0]); tt(xs[1])
    first = [tt(x).clone() for x in xs]
    bad = 0; t0 = time.perf_counter()
    for i in range(N):
        o = tt(xs[i % 4])
        if not torch.equal(o, first[i % 4]):
            bad += 1
    dt = time.perf_counter() - t0
print(f"SOAK {N} calls in {dt:.1f} s, outputs differing from the first: {bad}")
tt.close()