"""Face-granular head split (unpadded 80-wide heads -> 96 padded) / concat (96 -> 80): exactness + time."""
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()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
S, H, Dh, Dp = 192, 16, 80, 96
for trial in range(2):
    qkv_t = torch.randn(1, 1, S, 3 * H * Dh)
    qkv = ttnn.from_torch(qkv_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    ref = qkv_t.reshape(1, S, 3, H, Dh).permute(2, 0, 3, 1, 4)  # 3,1,H,S,Dh
    ref = torch.nn.functional.pad(ref, (0, Dp - Dh))
    outs = tk.head_split_faces(qkv, H, Dh, Dp)
    ok = all(torch.equal(ttnn.to_torch(o).float(), ref[i].to(torch.bfloat16).float()) for i, o in enumerate(outs))
    print(f"trial{trial} split exact={ok}", flush=True)
    ctx_t = torch.randn(1, H, S, Dp)
    ctx = ttnn.from_torch(ctx_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    c = tk.head_concat_faces(ctx, Dh)
    okc = torch.equal(ttnn.to_torch(c).float(), ctx_t[..., :Dh].permute(0, 2, 1, 3).reshape(1, 1, S, H * Dh).to(torch.bfloat16).float())
    print(f"trial{trial} concat exact={okc}", flush=True)
print("face split us", time_op(dev, lambda: tk.head_split_faces(qkv, H, Dh, Dp)))
print("face concat us", time_op(dev, lambda: tk.head_concat_faces(ctx, Dh)))
ttnn.close_device(dev)