"""Custom fused add+LayerNorm (tt_kernels.add_layernorm) vs torch fp32 and vs the library path
(ttnn.add + sharded ttnn.layer_norm + sharded_to_interleaved).  usage: addln_test.py [eth|worker] [soak N]"""
import sys
from pathlib import Path
from common import *
from open_device_12x10 import open_device
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
disp = sys.argv[1] if len(sys.argv) > 1 else "eth"
soak = int(sys.argv[3]) if len(sys.argv) > 3 and sys.argv[2] == "soak" else 0
timing = "time" in sys.argv
GM = int(next((a.split("=")[1] for a in sys.argv if a.startswith("mode=")), "2"))
NOY = "noy" in sys.argv  # bisect: plain LN only (HAS_Y=0)
dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20)
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK,
                                        orientation=ttnn.ShardOrientation.ROW_MAJOR)
xt = (torch.randn(1, 1, 192, 1280) * 2 + torch.randn(1280) * 1.5).to(torch.bfloat16)
yt = (torch.randn(1, 1, 192, 1280) * 0.7).to(torch.bfloat16)
gt = 1 + 0.3 * torch.randn(1280); bt = 0.2 * torch.randn(1280)
eps = 1e-6
x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
y = ttnn.from_torch(yt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
C = tk.addln_constants(gt, bt, dev)
if NOY:
    _, h0 = tk.add_layernorm(x, None, C, eps)
    ttnn.synchronize_device(dev)
    h0r = torch.nn.functional.layer_norm(xt.float(), (1280,), gt, bt, eps)
    hh = ttnn.to_torch(h0).float()
    print(f"mode={GM} noy first call OK; h PCC {pcc(hh, h0r):.7f} nan={int(hh.isnan().sum())} inf={int(hh.isinf().sum())}", flush=True)
    if GM == 0:  # stub: every core's stats are 10x its own 128-col partials
        xb = xt.float().reshape(192, 10, 128)
        mu = xb.sum(-1, keepdim=True) * 10 / 1280
        var = ((xb - mu) ** 2).sum(-1, keepdim=True) * 10 / 1280
        sr = (((xb - mu) / (var + eps).sqrt()).reshape(1, 1, 192, 1280) * gt + bt)
        print(f"stub ref PCC {pcc(hh.nan_to_num(), sr):.7f}; per-core-col nan counts", hh.reshape(192, 10, 128).isnan().sum((0, 2)).tolist(),
              "per row-tile", hh.reshape(6, 32, 1280).isnan().sum((1, 2)).tolist())
        print("sample", hh[0, 0, 0, :6].tolist(), sr[0, 0, 0, :6].tolist())
        z = ((hh - bt.to(torch.bfloat16).float()) / gt.to(torch.bfloat16).float()).reshape(192, 10, 128).double()
        xx = xb.double()
        for r in (0, 1, 31, 40):
            for c in (0, 5):
                A = torch.stack([xx[r, c], torch.ones(128, dtype=torch.float64)], 1)
                sol = torch.linalg.lstsq(A, z[r, c].unsqueeze(1)).solution.squeeze()
                rs, mu_dev = float(sol[0]), float(-sol[1] / sol[0])
                print(f"row {r} col {c}: dev rstd {rs:.4g} mean {mu_dev:.4g} | expect rstd {float((var[r, c] + eps).rsqrt()):.4g} mean {float(mu[r, c]):.4g} | var {float(var[r,c]):.4g}")
    ttnn.close_device(dev); sys.exit(0)
s, h = tk.add_layernorm(x, y, C, eps)
ttnn.synchronize_device(dev)
if "wshard" in sys.argv:
    G = dev.compute_with_storage_grid_size()
    sg = ttnn.num_cores_to_corerangeset(8, ttnn.CoreCoord(G.x, G.y), row_wise=True)
    wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
                            ttnn.ShardSpec(sg, [192, 160], ttnn.ShardOrientation.ROW_MAJOR))
    s3, h3 = tk.add_layernorm(x, y, C, eps, out_memory_config=wmc)
    ttnn.synchronize_device(dev)
    print("wshard h == interleaved h:", bool(torch.equal(ttnn.to_torch(h3).float(), ttnn.to_torch(h).float())),
          "s equal:", bool(torch.equal(ttnn.to_torch(s3).float(), ttnn.to_torch(s).float())), flush=True)
    if "time" in sys.argv:
        print(f"wshard addln {time_op(dev, lambda: list(tk.add_layernorm(x, y, C, eps, out_memory_config=wmc))):.2f} us", flush=True)
print(f"mode={GM} first call OK", flush=True)
st = ttnn.to_torch(s).float(); ht = ttnn.to_torch(h).float()
sref = (xt.float() + yt.float()).to(torch.bfloat16).float()
href = torch.nn.functional.layer_norm(sref, (1280,), gt, bt, eps)
print("s exact:", bool(torch.equal(st, sref)), "max|ds|", float((st - sref).abs().max()))
print(f"h PCC {pcc(ht, href):.7f} max|dh| {float((ht - href).abs().max()):.4f} mean|dh| {float((ht - href).abs().mean()):.5f}")
# library path
ga = ttnn.from_torch(gt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
be = ttnn.from_torch(bt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
lnpc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(10, 6), subblock_w=4, block_h=1, block_w=4, inplace=False)
def lib():
    s2 = ttnn.add(x, y, memory_config=res)
    hh = ttnn.layer_norm(s2, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc)
    hi = ttnn.sharded_to_interleaved(hh, L1)
    ttnn.deallocate(hh)
    return [s2, hi]
s2, hl = lib()
s2t = ttnn.to_torch(s2).float()
print("lib s exact:", bool(torch.equal(s2t, sref)), "max|ds|", float((s2t - sref).abs().max()), "custom==lib s:", bool(torch.equal(st, s2t)),
      "n diff", int((st != sref).sum()))
hl = ttnn.to_torch(hl).float()
href2 = torch.nn.functional.layer_norm(sref, (1280,), gt.to(torch.bfloat16).float(), bt.to(torch.bfloat16).float(), eps)
print(f"lib h PCC {pcc(hl, href):.7f} max|dh| {float((hl - href).abs().max()):.4f} mean|dh| {float((hl - href).abs().mean()):.5f}")
print(f"custom vs bf16-gamma ref PCC {pcc(ht, href2):.7f} mean|dh| {float((ht - href2).abs().mean()):.5f}; lib {float((hl - href2).abs().mean()):.5f}")
# plain LN variant (no y)
_, h0 = tk.add_layernorm(x, None, C, eps)
h0r = torch.nn.functional.layer_norm(xt.float(), (1280,), gt, bt, eps)
print(f"noy h PCC {pcc(ttnn.to_torch(h0).float(), h0r):.7f}")
if timing:
    print(f"custom addln {time_op(dev, lambda: list(tk.add_layernorm(x, y, C, eps))):.2f} us; library add+ln+s2i {time_op(dev, lib):.2f} us", flush=True)
if soak:
    tid = ttnn.begin_trace_capture(dev, cq_id=0)
    try:
        outs = [tk.add_layernorm(x, y, C, eps) for _ in range(10)]
    finally:
        ttnn.end_trace_capture(dev, tid, cq_id=0)
    for i in range(soak // 10):
        ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
        if i % 50 == 0:
            ttnn.synchronize_device(dev)
    ttnn.synchronize_device(dev)
    bad = sum(int(not torch.equal(ttnn.to_torch(o[1]).float(), ht)) for o in outs)
    print(f"soak {soak} replays OK, outputs differing from first call: {bad}/10")
    ttnn.release_trace(dev, tid)
ttnn.close_device(dev)