"""Block-tail emulation: addln1(prefetch proj) -> proj = cmm(SDPA-layout ctx, head_concat folded) -> addln2(prefetch fc2)
-> gelu -> fc2 = cmm(prefetched pairs) vs the model's library path (head_concat + 2D linear, gelu + 2D linear).
usage: cmm_block_test.py [pf2=P] [time] [soak=N]"""
import sys
from pathlib import Path
from common import *
from open_device_12x10 import open_device
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
from hamer.vit import tuned_matmul_configs
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
P2 = int(arg("pf2", "3")); S2 = int(arg("s2", "3")); S1 = int(arg("s1", "3"))
dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20)
try:
    torch.manual_seed(0)
    L1 = ttnn.L1_MEMORY_CONFIG
    res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK,
                                            orientation=ttnn.ShardOrientation.ROW_MAJOR)
    T = lambda t, mem=L1, dt=ttnn.bfloat16: ttnn.from_torch(t, dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem)
    x = T(torch.randn(1, 1, 192, 1280), res); y = T(torch.randn(1, 1, 192, 1280), res)
    ctx_t = torch.randn(1, 16, 192, 96) * 0.5
    ctx = T(ctx_t)
    h1 = T(torch.randn(1, 1, 192, 5120) * 1.5)
    pw_t = (torch.randn(1536, 1280) * 1536 ** -0.5).bfloat16(); pb_t = (torch.randn(1280) * 0.1).bfloat16()
    fw_t = (torch.randn(5120, 1280) * 5120 ** -0.5).bfloat16(); fb_t = (torch.randn(1280) * 0.1).bfloat16()
    pw = T(pw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b); fw = T(fw_t, ttnn.DRAM_MEMORY_CONFIG, ttnn.bfloat8_b)
    pb = T(pb_t.reshape(1, 1, 1, -1)); fb = T(fb_t.reshape(1, 1, 1, -1))
    C1 = tk.addln_constants(torch.ones(1280), torch.zeros(1280), dev); C2 = tk.addln_constants(torch.ones(1280) * 0.9, torch.zeros(1280), dev)
    pwk = tk.CmmWeight(pw_t, dev, 10, 4, 8, ttnn.bfloat8_b); fwk = tk.CmmWeight(fw_t, dev, 10, 4, 16, ttnn.bfloat8_b)
    V1 = "v1" in sys.argv
    P1 = int(arg("pf1", "3"))
    pfp = tk.CmmPrefetch(dev, 10, 6, pwk, P1, single=V1) if P1 else None
    pff = tk.CmmPrefetch(dev, 10, 6, fwk, P2, single=V1) if P2 else None
    defs = ("CMM_FLUSH", f"CMM_LA={arg('la', '1')}") + (() if V1 else ("CMM_DRAM0",))
    ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True)
    pcs = tuned_matmul_configs(dev.compute_with_storage_grid_size(), fp32_acc=True)

    def cmm_path():
        s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1, prefetch=pfp)
        ttnn.deallocate(h)
        d = tk.cmm_linear(ctx, pwk, pb, res, kb=8, slots=S1, defines=defs, prefetch=pfp, a_heads=(16, 96), v1=V1)
        s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1, prefetch=pff)
        ttnn.deallocate(h2)
        g = tk.gelu(h1, memory_config=L1)
        d2 = tk.cmm_linear(g, fwk, fb, res, kb=16, slots=S2, defines=defs, prefetch=pff, v1=V1)
        for t in (s, d, g):
            ttnn.deallocate(t)
        return s2, d2

    def lib_path():
        s, h = tk.add_layernorm(x, y, C1, 1e-6, out_memory_config=L1)
        ttnn.deallocate(h)
        cm = tk.head_concat(ctx, memory_config=L1)
        d = ttnn.linear(cm, pw, bias=pb, program_config=pcs["proj"], memory_config=res, compute_kernel_config=ck)
        s2, h2 = tk.add_layernorm(s, d, C2, 1e-6, out_memory_config=L1)
        ttnn.deallocate(h2)
        g = tk.gelu(h1, memory_config=L1)
        d2 = ttnn.linear(g, fw, bias=fb, program_config=pcs["fc2"], memory_config=res, compute_kernel_config=ck)
        for t in (s, cm, d, g):
            ttnn.deallocate(t)
        return s2, d2

    a2, ad = cmm_path(); ttnn.synchronize_device(dev)
    b2, bd = lib_path(); ttnn.synchronize_device(dev)
    A2, AD, B2, BD = (ttnn.to_torch(t).float() for t in (a2, ad, b2, bd))
    print(f"[blk] first calls OK  s2 pcc {pcc(A2, B2):.7f} max|d| {float((A2 - B2).abs().max()):.4f}   fc2 pcc {pcc(AD, BD):.7f} max|d| {float((AD - BD).abs().max()):.4f}", flush=True)
    if "time" in sys.argv:
        print(f"[blk] cmm path {time_op(dev, cmm_path):.2f} us   library path {time_op(dev, lib_path):.2f} us", flush=True)
    n = int(arg("soak", "0"))
    if n:
        reps = 20
        tid = ttnn.begin_trace_capture(dev, cq_id=0)
        outs = []
        try:
            for _ in range(reps):
                outs.append(cmm_path())
        finally:
            ttnn.end_trace_capture(dev, tid, cq_id=0)
        try:
            for i in range(n // reps):
                ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
                if i % 40 == 39:
                    ttnn.synchronize_device(dev)
            ttnn.synchronize_device(dev)
            bad = sum(int(not (torch.equal(ttnn.to_torch(p).float(), A2) and torch.equal(ttnn.to_torch(q).float(), AD))) for p, q in outs)
            print(f"[blk] soak {n // reps * reps} block-tails OK, outputs differing from the first: {bad}/{reps}", flush=True)
        finally:
            ttnn.release_trace(dev, tid)
finally:
    ttnn.close_device(dev)