"""How does the packer round fp32 DST -> bfp8_b (bfp8_pack_precise)?  Copy-only GELU kernel (HG_NOP) with a bfp8 output
vs host round-to-nearest / truncation of the same bf16 values (round 8)."""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels

def bfp8(x, mode):
    s = x.shape
    g = x.reshape(-1, 16)  # tile rows of 16 inside a face: TILE layout handled by comparing in torch layout
    return g, s

dev = open_dev()
try:
    torch.manual_seed(0)
    xt = (torch.randn(1, 1, 64, 64) * torch.exp(torch.randn(1, 1, 64, 1))).bfloat16()
    x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
    o = ttnn.to_torch(tt_kernels.gelu(x, defines=("HG_NOP",), out_dtype=ttnn.bfloat8_b)).float()
    h = ttnn.to_torch(ttnn.from_torch(xt.float(), dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float()  # host conversion
    xf = xt.float()
    # host reference: groups of 16 consecutive columns per row (face rows)
    g = xf.reshape(64, 4, 16)
    mx = g.abs().amax(-1, keepdim=True).clamp_min(1e-30)
    step = torch.pow(2.0, torch.floor(torch.log2(mx)) - 6)
    rn = (torch.round(g / step).clamp(-127, 127) * step).reshape(1, 1, 64, 64)
    tr = (torch.trunc(g / step).clamp(-127, 127) * step).reshape(1, 1, 64, 64)
    for name, ref in (("round-nearest", rn), ("truncate", tr), ("host ttnn.from_torch", h)):
        print(f"[pack] device vs {name:22s}: equal {float((o == ref).float().mean()):.4f}  max|d| {float((o - ref).abs().max()):.3e}", flush=True)
    print(f"[pack] mean signed error device-x {float((o - xf).mean()):+.3e}  rn-x {float((rn - xf).mean()):+.3e}  tr-x {float((tr - xf).mean()):+.3e}")
finally:
    ttnn.close_device(dev)