"""Round 15 (D5 fused): the MANO-head q GEMV with the attention in the same program (hgemv(attn=...): 16 GEMV cores signal
8 hattn2 cores through a semaphore) vs the two separate programs.  Expect bit-identical q and context.
usage: hgattn_test.py [time] [soak=N]"""
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, K, N = 8, 192, 64, 1024, 512
    L1 = ttnn.L1_MEMORY_CONFIG
    x = ttnn.from_torch(torch.randn(1, 1, 32, K), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    w = tk.CmmWeight((torch.randn(K, N) * K ** -0.5).to(torch.bfloat16), dev, N // 32, 1, K // 32, ttnn.bfloat16)
    b = ttnn.from_torch(torch.zeros(1, 1, 32, N), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    tt_ = lambda t: ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    Kt, Vt = tt_(torch.randn(1, H, S, D).to(torch.bfloat16)), tt_(torch.randn(1, H, S, D).to(torch.bfloat16))
    gcores = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 3))})
    acores = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(8, 8), ttnn.CoreCoord(11, 9))})

    def sep():
        q = tk.hgemv(x, w, b, ln_eps=1e-5, out_dtype=ttnn.bfloat16, out_heads=H, cores=gcores)
        return q, tk.hattn2(q, Kt, Vt, D ** -0.5, cores=acores)

    def fused():
        return tk.hgemv(x, w, b, ln_eps=1e-5, out_dtype=ttnn.bfloat16, out_heads=H, cores=gcores,
                        attn=(Kt, Vt, D ** -0.5, acores))
    q1, c1 = fused()
    ttnn.synchronize_device(dev)
    print("[hga] fused first call OK", flush=True)
    q0, c0 = sep()
    Q0, C0 = ttnn.to_torch(q0).float()[0, :, 0], ttnn.to_torch(c0).float()[0, :, 0]
    Q1, C1 = ttnn.to_torch(q1).float()[0, :, 0], ttnn.to_torch(c1).float()[0, :, 0]
    print(f"[hga] q equal {bool(torch.equal(Q0, Q1))}  ctx equal {bool(torch.equal(C0, C1))} (max|d| {float((C0 - C1).abs().max()):.3e})", flush=True)
    if "time" in sys.argv:
        print(f"[hga] separate {time_op(dev, sep):.2f} us", flush=True)
        print(f"[hga] fused    {time_op(dev, fused):.2f} us", flush=True)
    n = int(arg("soak", "0"))
    if n:
        bad = 0
        for _ in range(n):
            q_, c_ = fused()
            bad += int(not torch.equal(ttnn.to_torch(c_).float()[0, :, 0], C1))
            ttnn.deallocate(q_); ttnn.deallocate(c_)
        print(f"[hga] soak {n}: differing {bad}", flush=True)
finally:
    ttnn.close_device(dev)