"""Round 15 (D5): tt_kernels.hattn2 (dense-score 1-query attention, one core per head) vs the library SDPA as the MANO head
uses it (8x1 grid, q chunk 32, k chunk 192, HiFi4 + fp32) and vs torch fp32.  usage: hattn2_test.py [time] [soak=N] [big]
  big: score magnitudes up to ~60 (stress the max subtraction / exp range)"""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
dev = open_dev()
try:
    torch.manual_seed(0)
    H, S, D = 8, 192, 64
    sc = 3.0 if "big" in sys.argv else 1.0
    q = (torch.randn(H, D) * sc).to(torch.bfloat16)
    k = (torch.randn(H, S, D) * sc).to(torch.bfloat16)
    v = torch.randn(H, S, D).to(torch.bfloat16)
    ref = torch.softmax((q.float()[:, None, :] @ k.float().transpose(-1, -2)) * D ** -0.5, -1) @ v.float()  # [H, 1, D]
    ref = ref[:, 0]
    L1 = ttnn.L1_MEMORY_CONFIG
    tt_ = lambda t: ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    qh = torch.zeros(1, H, 32, D, dtype=torch.bfloat16)
    qh[0, :, 0] = q
    Q, K, V = tt_(qh), tt_(k[None]), tt_(v[None])
    pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=ttnn.CoreCoord(8, 1), q_chunk_size=32, k_chunk_size=192,
                                exp_approx_mode=False)
    cc = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
                                                math_approx_mode=False)
    lib = lambda: ttnn.transformer.scaled_dot_product_attention(Q, K, V, is_causal=False, scale=D ** -0.5, program_config=pc,
                                                                memory_config=L1, compute_kernel_config=cc)
    new = lambda: tk.hattn2(Q, K, V, D ** -0.5)
    o1 = new()
    ttnn.synchronize_device(dev)
    O1 = ttnn.to_torch(o1).float()[0, :, 0]
    o0 = lib()
    O0 = ttnn.to_torch(o0).float()[0, :, 0]
    print(f"[ha2] first call OK; finite {bool(torch.isfinite(O1).all())}", flush=True)
    print(f"[ha2] vs fp32 ref: hattn2 PCC {pcc(O1, ref):.7f} max|d| {float((O1 - ref).abs().max()):.3e} mean|d| "
          f"{float((O1 - ref).abs().mean()):.3e} | lib PCC {pcc(O0, ref):.7f} max|d| {float((O0 - ref).abs().max()):.3e} "
          f"mean|d| {float((O0 - ref).abs().mean()):.3e}", flush=True)
    print(f"[ha2] hattn2 vs lib: PCC {pcc(O1, O0):.7f} max|d| {float((O1 - O0).abs().max()):.3e}", flush=True)
    if "time" in sys.argv:
        print(f"[ha2] lib SDPA {time_op(dev, lib):.2f} us", flush=True)
        print(f"[ha2] hattn2   {time_op(dev, new):.2f} us", flush=True)
    if "prof" in sys.argv:  # HA2_PROF=1 TT_METAL_DEVICE_PROFILER=1: zone timeline of the last hattn2 call
        for _ in range(4):
            ttnn.deallocate(new())
            ttnn.synchronize_device(dev)
        ttnn.ReadDeviceProfiler(dev)
    n = int(arg("soak", "0"))
    if n:
        bad = 0
        for _ in range(n):
            o = new()
            bad += int(not torch.equal(ttnn.to_torch(o).float()[0, :, 0], O1))
            ttnn.deallocate(o)
        print(f"[ha2] soak {n}: differing {bad}", flush=True)
finally:
    ttnn.close_device(dev)
if "prof" in sys.argv:
    tl_report(risks=("NCRISC", "BRISC", "TRISC_0", "TRISC_1", "TRISC_2"))