"""hgemv (kernels/hgemv, tt_kernels.hgemv) bring-up: correctness vs torch fp32 and timing.
usage: hgemv_test.py [K=1024] [N=1024] [ln=1] [gelu=0] [res=1] [xf32=1] [time=1]"""
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
a = dict(s.split("=") for s in sys.argv[1:])
K, N = int(a.get("K", 1024)), int(a.get("N", 1024))
ln, gelu, res, xf32 = a.get("ln", "1") == "1", a.get("gelu", "0") == "1", a.get("res", "0") == "1", a.get("xf32", "1") == "1"
dev = open_dev()
try:
    torch.manual_seed(0)
    xv = torch.randn(K) * 2 + 0.5
    W = torch.randn(K, N) / K ** 0.5
    bv = torch.randn(N) * 0.1
    rv = torch.randn(N)
    xdt = ttnn.float32 if xf32 else ttnn.bfloat16
    rowf = lambda v: v.reshape(1, 1, 1, -1).expand(1, 1, 32, v.numel()).contiguous()
    x = ttnn.from_torch(rowf(xv), dtype=xdt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
    D = int(a.get("din", 0))
    if D:  # dense replica input: D tiles, each = x reshaped 32x32 row-major
        x = ttnn.from_torch(tk.hg_dense_layout(xv).repeat(D, 1).reshape(1, 1, 32 * D, 32), dtype=xdt, layout=ttnn.TILE_LAYOUT,
                            device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
    b = ttnn.from_torch(rowf(bv).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
                        memory_config=ttnn.L1_MEMORY_CONFIG)  # as in the model (_hg_rowf)
    r = ttnn.from_torch(rowf(rv), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) if res else None
    w = tk.CmmWeight(W, dev, N // 32, 1, K // 32, ttnn.bfloat16)
    xin = xv.to(torch.float32 if xf32 else torch.bfloat16).double()
    h = torch.nn.functional.layer_norm(xin, (K,), eps=1e-5) if ln else xin
    ref = h @ W.to(torch.bfloat16).double() + bv.to(torch.bfloat16).double()
    if gelu:
        ref = torch.nn.functional.gelu(ref)
    if res:
        ref = ref + rv.double()
    DO = int(a.get("dout", 0))
    o = tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
    ttnn.synchronize_device(dev)
    if DO:
        o, od = o
        dd = ttnn.to_torch(od).double().reshape(DO, 32, 32)
        got0 = ttnn.to_torch(o).double()[0, 0, 0]
        print(f"[hg] dense out: replicas equal row-form: {all(torch.equal(tk.hg_dense_vector(dd[c]), got0) for c in range(DO))}", flush=True)
    got = ttnn.to_torch(o).double()[0, 0, 0]
    print(f"[hg] K{K} N{N} ln{int(ln)} gelu{int(gelu)} res{int(res)} xf32{int(xf32)}: pcc {pcc(got, ref):.7f} "
          f"max|d| {float((got - ref).abs().max()):.5f} rel {float((got - ref).abs().max() / ref.abs().max()):.2e}", flush=True)
    if a.get("pf", "0") == "1":  # round 10: cross-op weight prefetch -- GEMV 1 on rect A prefetches GEMV 2's weights
        rA = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 3))})
        rB = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 4), ttnn.CoreCoord(7, 7))})
        bufB = tk.hg_weight_buffer(dev, rB)
        kw = dict(ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
        def pair(pf):
            o1 = tk.hgemv(x, w, b, cores=rA, prefetch=(w, bufB, rB) if pf else None, **kw)
            o2 = tk.hgemv(x, w, b, cores=rB, w_l1=bufB if pf else None, **kw)
            return o1, o2
        p1, p2 = pair(True)
        ttnn.synchronize_device(dev)
        g = lambda o: ttnn.to_torch(o[0] if DO else o).double()[0, 0, 0]
        print(f"[hg] prefetch pair: out1 == out2 {bool(torch.equal(g(p1), g(p2)))}  == plain {bool(torch.equal(g(p2), got))}", flush=True)
        if a.get("time", "1") == "1":
            bufA = tk.hg_weight_buffer(dev, rA)
            def chain(pf, n=8):  # steady state: every program prefetches the next one's weights
                rs, bs, outs = (rA, rB), (bufA, bufB), []
                for i in range(n):
                    nxt = (w, bs[(i + 1) % 2], rs[(i + 1) % 2]) if pf and i + 1 < n else None
                    outs.append(tk.hgemv(x, w, b, cores=rs[i % 2], w_l1=bs[i % 2] if pf and i > 0 else None,
                                         prefetch=nxt, **kw))
                return outs
            print(f"[hg] chain of 8: plain {time_op(dev, lambda: chain(False), reps=4) / 8:.2f} us/gemv  prefetched "
                  f"{time_op(dev, lambda: chain(True), reps=4) / 8:.2f} us/gemv", flush=True)
            print(f"[hg] pair plain {time_op(dev, lambda: pair(False)):.2f} us  prefetched {time_op(dev, lambda: pair(True)):.2f} us", flush=True)
    if a.get("time", "1") == "1":
        print(f"[hg] {time_op(dev, lambda: tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)):.2f} us/call", flush=True)
    if a.get("prof", "0") == "1":  # round 10: zone timeline (HG_PROF=1 TT_METAL_DEVICE_PROFILER=1)
        for _ in range(4):
            o2 = tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
            ttnn.synchronize_device(dev)
        ttnn.ReadDeviceProfiler(dev)
finally:
    from open_device_12x10 import close_device
    close_device(dev)
if a.get("prof", "0") == "1":
    tl_report(risks=("NCRISC", "BRISC", "TRISC_0", "TRISC_1", "TRISC_2"))