"""hattn (kernels/hattn) vs torch fp64 and vs head_split + SDPA + head_concat timing.  usage: hattn_test.py [layer=0] [time=1]"""
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
a = dict(s.split("=") for s in sys.argv[1:])
l = int(a.get("layer", 0))
dev = open_dev()
try:
    torch.manual_seed(0)
    H, L = 8, 6
    qv = torch.randn(H * 64) * 2
    kvt = torch.randn(192, L * 1024) * 1.5
    rowf = lambda v: v.reshape(1, 1, 1, -1).expand(1, 1, 32, v.numel()).contiguous()
    q = ttnn.from_torch(rowf(qv * 0.125).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
                        memory_config=ttnn.L1_MEMORY_CONFIG)
    kv = ttnn.from_torch(kvt.reshape(1, 1, 192, -1).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT,
                         device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
    kb = kvt.to(torch.bfloat16).double()
    qb = (qv * 0.125).to(torch.bfloat16).double()
    ref = []
    for h in range(H):
        K = kb[:, l * 1024 + h * 64: l * 1024 + h * 64 + 64]
        V = kb[:, l * 1024 + 512 + h * 64: l * 1024 + 512 + h * 64 + 64]
        p = torch.softmax(K @ qb[h * 64:(h + 1) * 64], 0)
        ref.append(p @ V)
    ref = torch.cat(ref)
    o = tk.hattn(q, kv, l)
    ttnn.synchronize_device(dev)
    got = ttnn.to_torch(o).double()[0, 0, 0]
    print(f"[ha] layer {l}: pcc {pcc(got, ref):.7f} max|d| {float((got - ref).abs().max()):.5f}", flush=True)
    if a.get("time", "1") == "1":
        print(f"[ha] hattn {time_op(dev, lambda: tk.hattn(q, kv, l)):.2f} us/call", flush=True)
finally:
    from open_device_12x10 import close_device
    close_device(dev)