"""add_layernorm on the MANO-head token layout: [1,1,32,1024] WIDTH-sharded on 8 cores (one tile row)."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
dev = open_dev()
try:
    torch.manual_seed(0)
    cg = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 0))})
    wm = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
                           ttnn.ShardSpec(cg, [32, 128], ttnn.ShardOrientation.ROW_MAJOR))
    xt = torch.randn(1, 1, 32, 1024) * 2; yt = torch.randn(1, 1, 32, 1024)
    g = torch.randn(1024) * 0.2 + 1; b = torch.randn(1024) * 0.1
    x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wm)
    y = ttnn.from_torch(yt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wm)
    C = tk.addln_constants(g, b, dev)
    s, h = tk.add_layernorm(x, y, C, 1e-5, out_memory_config=wm)
    ttnn.synchronize_device(dev)
    sr = xt.bfloat16().float() + yt.bfloat16().float()
    hr = torch.nn.functional.layer_norm(sr.bfloat16().float(), (1024,), g.bfloat16().float(), b.bfloat16().float(), 1e-5)
    st, ht = ttnn.to_torch(s).float(), ttnn.to_torch(h).float()
    print(f"[ahead] first call OK  s pcc {pcc(st, sr):.6f}  h pcc {pcc(ht, hr):.6f}  max|dh| {float((ht - hr).abs().max()):.4f}", flush=True)
    ln = lambda: ttnn.layer_norm(ttnn.add(x, y, memory_config=wm), weight=None, bias=None, epsilon=1e-5, memory_config=wm)
    if "time" in sys.argv:
        print(f"[ahead] addln {time_op(dev, lambda: tk.add_layernorm(x, y, C, 1e-5, out_memory_config=wm)):.2f} us", flush=True)
finally:
    ttnn.close_device(dev)