"""signed-error / rounding-bias check of the custom GELU vs ttnn.gelu on fc1-like data."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels
dev = open_dev()
try:
    L1 = ttnn.L1_MEMORY_CONFIG
    torch.manual_seed(0)
    xt = torch.randn(1, 1, 192, 5120) * 2.0 - 0.5
    x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    xb = xt.bfloat16().double()
    ref = torch.nn.functional.gelu(xb)
    ulp = lambda r: torch.where(r == 0, torch.full_like(r, 1e-30), 2.0 ** (torch.floor(torch.log2(r.abs().clamp_min(1e-30))) - 7))
    u = ulp(ref)
    for name, t in [("custom", ttnn.to_torch(tt_kernels.gelu(x, memory_config=L1)).double()),
                    ("custom_bf16dst", ttnn.to_torch(tt_kernels.gelu(x, memory_config=L1, fp32_dest=False)).double()),
                    ("ttnn", ttnn.to_torch(ttnn.gelu(x, memory_config=L1)).double()),
                    ("rne(ref)", ref.bfloat16().double())]:
        e = (t - ref) / u
        m = ref.abs() > 1e-3
        print(f"{name:15s} signed err mean {float(e[m].mean()):+.4f} ulp  |err| mean {float(e[m].abs().mean()):.4f} ulp  max {float(e[m].abs().max()):.3f}  sum(t-ref)/sum|ref| {float((t-ref).sum()/ref.abs().sum()):+.3e}", flush=True)
finally:
    ttnn.close_device(dev)