"""gelu(prefetch) -> cmm_linear(prefetch) (fc2 with its first weight pairs prefetched during the GELU) vs the
model's gelu -> library 2D linear.  usage: cmm_pf_test.py [pairs=P] [la=L] [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)
P, LA, KB = int(arg("pairs", "3")), int(arg("la", "1")), int(arg("kb", "16"))
dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20)
try:
    torch.manual_seed(0)
    L1 = ttnn.L1_MEMORY_CONFIG
    M, K, N = 192, 5120, 1280
    at = (torch.randn(1, 1, M, K) * 1.5).to(torch.bfloat16)
    wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16)
    bt = (torch.randn(N) * 0.1).to(torch.bfloat16)
    a = ttnn.from_torch(at, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    b = ttnn.from_torch(bt.reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
    omem = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK,
                                             orientation=ttnn.ShardOrientation.ROW_MAJOR)
    wpk = tk.CmmWeight(wt, dev, 10, 4, KB, ttnn.bfloat8_b)
    pf = tk.CmmPrefetch(dev, 10, 6, wpk, P) if P else None
    defs = ("CMM_FLUSH", f"CMM_LA={LA}", "CMM_DRAM0")
    def fused():
        g = tk.gelu(a, memory_config=L1, prefetch=pf)
        o = tk.cmm_linear(g, wpk, b, omem, kb=KB, defines=defs, prefetch=pf)
        ttnn.deallocate(g)
        return o
    ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True)
    pc2 = tuned_matmul_configs(dev.compute_with_storage_grid_size(), fp32_acc=True)["fc2"]
    def lib():
        g = tk.gelu(a, memory_config=L1)
        o = ttnn.linear(g, w, bias=b, program_config=pc2, memory_config=omem, compute_kernel_config=ck)
        ttnn.deallocate(g)
        return o
    print(f"[pf] pairs={P} la={LA} kb={KB}", flush=True)
    o = fused(); ttnn.synchronize_device(dev)
    ot = ttnn.to_torch(o).float()
    lt = ttnn.to_torch(lib()).float()
    ref = torch.nn.functional.gelu(at.float()).bfloat16().float().reshape(M, K) @ ttnn.to_torch(w).float() + bt.float()
    print(f"[pf] lib PCC {pcc(lt.reshape(M, N), ref):.7f}", flush=True)
    print(f"[pf] first call OK  PCC {pcc(ot.reshape(M, N), ref):.7f}  bit-identical to library pair: {bool(torch.equal(ot, lt))}  max|d| {float((ot - lt).abs().max()):.4f}", flush=True)
    if "time" in sys.argv:
        print(f"[pf] gelu+cmm(prefetch) {time_op(dev, fused):.2f} us", flush=True)
        print(f"[pf] gelu+library 2D    {time_op(dev, lib):.2f} us", flush=True)
        print(f"[pf] gelu(prefetch) alone {time_op(dev, lambda: tk.gelu(a, memory_config=L1, prefetch=pf)):.2f} us   gelu alone {time_op(dev, lambda: tk.gelu(a, memory_config=L1)):.2f} us", flush=True)
    n = int(arg("soak", "0"))
    if n:
        reps = 50
        tid = ttnn.begin_trace_capture(dev, cq_id=0)
        outs = []
        try:
            for _ in range(reps):
                outs.append(fused())
        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(x).float(), ot)) for x in outs)
            print(f"[pf] soak {n // reps * reps} calls OK, outputs differing from the first: {bad}/{reps}", flush=True)
        finally:
            ttnn.release_trace(dev, tid)
finally:
    ttnn.close_device(dev)