"""Round 10: MANO-head cross-attention (1 query token, 8 heads x 64, 192 keys): the library SDPA as used now
([1, 8, 32, 64] query, row 0 meaningful) vs flash-decode SDPA ([1, 1, 8->32, 64] query, heads as rows).
usage: head_sdpa_decode.py"""
import sys, os
from pathlib import Path
from common import *
dev = open_dev()
try:
    torch.manual_seed(0)
    H, S, D = 8, 192, 64
    q = torch.randn(H, D)
    k = torch.randn(H, S, D)
    v = torch.randn(H, S, D)
    ref = torch.softmax((q[:, None, :] @ k.transpose(-1, -2)) * D ** -0.5, -1) @ v   # [H, 1, D]
    ref = ref[:, 0]
    L1 = ttnn.L1_MEMORY_CONFIG
    G = dev.compute_with_storage_grid_size()
    tk_ = lambda t: ttnn.from_torch(t.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    K, V = tk_(k[None]), tk_(v[None])
    qh = torch.zeros(1, H, 32, D); qh[0, :, 0] = q
    Q = tk_(qh)
    pc = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=G, 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)
    f0 = 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)
    o0 = ttnn.to_torch(f0()).float()[0, :, 0]
    print(f"[sdpa] current: PCC {pcc(o0, ref):.7f} max|d| {float((o0 - ref).abs().max()):.4f}  {time_op(dev, f0):.2f} us", flush=True)
    qd = torch.zeros(1, 1, 32, D); qd[0, 0, :H] = q
    DR = ttnn.DRAM_MEMORY_CONFIG
    tkd = lambda t: ttnn.from_torch(t.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR)
    QD, K, V = tkd(qd), tkd(k[None]), tkd(v[None])
    for kc in (32, 64, 128):
        pcd = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=G, q_chunk_size=32, k_chunk_size=kc, exp_approx_mode=False)
        f1 = lambda pcd=pcd: ttnn.transformer.scaled_dot_product_attention_decode(
            QD, K, V, is_causal=True, cur_pos=[S - 1], scale=D ** -0.5, program_config=pcd, memory_config=DR,
            compute_kernel_config=cc)
        try:
            o1 = ttnn.to_torch(f1()).float()[0, 0, :H]
            print(f"[sdpa] decode k_chunk {kc}: PCC {pcc(o1, ref):.7f} max|d| {float((o1 - ref).abs().max()):.4f}  "
                  f"{time_op(dev, f1):.2f} us", flush=True)
        except Exception as e:  # noqa: BLE001
            print(f"[sdpa] decode k_chunk {kc} FAIL {str(e).splitlines()[0][:200]}", flush=True)
finally:
    ttnn.close_device(dev)