"""custom GELU (tt_kernels.gelu) soak: one eager call, then N trace replays of 50 calls on the fc1-output shape;
every output must stay bit-identical.  usage: gelu_soak.py N"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels
n = int(sys.argv[1]) if len(sys.argv) > 1 else 1000
dev = open_dev()
try:
    L1 = ttnn.L1_MEMORY_CONFIG
    torch.manual_seed(0)
    xt = torch.randn(1, 1, 192, 5120) * 3.0
    x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    o0 = ttnn.to_torch(tt_kernels.gelu(x, memory_config=L1)).float()
    ref = torch.nn.functional.gelu(xt.bfloat16().float())
    print(f"[gsoak] first call OK  pcc {pcc(o0, ref):.8f} max|err| {float((o0 - ref).abs().max()):.3e}", flush=True)
    reps = 50
    tid = ttnn.begin_trace_capture(dev, cq_id=0)
    outs = []
    try:
        for _ in range(reps):
            outs.append(tt_kernels.gelu(x, memory_config=L1))
    finally:
        ttnn.end_trace_capture(dev, tid, cq_id=0)
    try:
        for i in range(max(1, n // reps)):
            ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
            if i % 40 == 39:
                ttnn.synchronize_device(dev)
        ttnn.synchronize_device(dev)
        bad = sum(int(not torch.equal(ttnn.to_torch(t).float(), o0)) for t in outs)
        print(f"[gsoak] {max(1, n // reps) * reps} calls OK, outputs differing from first: {bad}/{reps}", flush=True)
    finally:
        ttnn.release_trace(dev, tid)
finally:
    ttnn.close_device(dev)