"""cmm_linear (kb, slots) sweep on one shape, PCC-checked.  usage: cmm_sweep.py [eth|worker] shape=fc2 cfgs=8x3,8x6,16x3 [defs=...]"""
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
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
disp = sys.argv[1] if len(sys.argv) > 1 and sys.argv[1] in ("eth", "worker") else "eth"
shape = arg("shape", "fc2")
GX, GY, NT, KT = {"fc2": (10, 6, 4, 160), "proj": (10, 6, 4, 48), "qkv": (12, 6, 12, 40), "fc1": (10, 6, 16, 40)}[shape]
cfgs = [tuple(int(v) for v in c.split("x")) for c in arg("cfgs", "8x3").split(",")]
lay = arg("layout", "out")
defs = tuple(d for d in arg("defs", "").split(",") if d)
dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20)
try:
    torch.manual_seed(0)
    M, K, N = 32 * GY, 32 * KT, 32 * NT * GX
    at = (torch.randn(1, 1, M, K) * 0.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=ttnn.L1_MEMORY_CONFIG)
    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=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK,
                                             orientation=ttnn.ShardOrientation.ROW_MAJOR)
    aref = torch.nn.functional.gelu(at.float()).to(torch.bfloat16).float() if "gelu" in sys.argv else at.float()
    ref = aref.reshape(M, K) @ ttnn.to_torch(w).float() + bt.float()
    for kb, s in cfgs:
        wf = tk.CmmWeight(wt, dev, GX, NT, kb, ttnn.bfloat8_b) if "packed" in sys.argv else w
        f = lambda: tk.cmm_linear(a, wf, b, omem, kb=kb, slots=s, defines=defs, layout=lay, rotate="rot" in sys.argv, gelu="gelu" in sys.argv, v1="v1" in sys.argv)
        o = f(); ttnn.synchronize_device(dev)
        p = pcc(ttnn.to_torch(o).float().reshape(M, N), ref)
        print(f"[sweep] {shape} {lay} packed={'packed' in sys.argv} rot={'rot' in sys.argv} kb={kb} slots={s} defs={defs} PCC {p:.7f}  {time_op(dev, f):.2f} us", flush=True)
    if "lib" in sys.argv:
        from hamer.vit import tuned_matmul_configs
        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)[shape]
        def libf():
            g = tk.gelu(a, memory_config=ttnn.L1_MEMORY_CONFIG) if "gelu" in sys.argv else a
            o = ttnn.linear(g, w, bias=b, program_config=pc2, memory_config=omem, compute_kernel_config=ck)
            if g is not a:
                ttnn.deallocate(g)
            return o
        o = libf(); ttnn.synchronize_device(dev)
        p = pcc(ttnn.to_torch(o).float().reshape(M, N), ref)
        print(f"[sweep] {shape} LIBRARY (gelu kernel={'gelu' in sys.argv} + 2D linear) PCC {p:.7f}  {time_op(dev, libf):.2f} us", flush=True)
finally:
    ttnn.close_device(dev)