"""Device-profiler zones of cmm_linear (CMM_PROF). usage: TT_METAL_DEVICE_PROFILER=1 cmm_prof.py kb=16 slots=3 [defs=..]"""
import sys, os, glob, csv
from pathlib import Path
from collections import defaultdict
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)
kb, s = int(arg("kb", "16")), int(arg("slots", "3"))
defs = ("CMM_PROF",) + tuple(d for d in arg("defs", "").split(",") if d)
lay = arg("layout", "out")
dev = open_dev() if "parse" not in sys.argv else None
try:
  if dev is not None:
      GX, GY, NT, KT = 10, 6, 4, 160
      M, K, N = 192, 5120, 1280
      a = ttnn.from_torch(torch.randn(1, 1, M, K), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
      wt = torch.randn(K, N) * 0.01
      w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
      if "packed" in sys.argv:
          w = tk.CmmWeight(wt, dev, GX, NT, kb, ttnn.bfloat8_b)
      b = ttnn.from_torch(torch.zeros(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)
      for _ in range(4):
          o = tk.cmm_linear(a, w, 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); ttnn.synchronize_device(dev); ttnn.deallocate(o)
      ttnn.ReadDeviceProfiler(dev)
finally:
    if dev is not None:
        ttnn.close_device(dev)
f = os.environ["TT_METAL_PROFILER_DIR"] + "/.logs/profile_log_device.csv"
rows = list(csv.reader(open(f)))[2:]
ev = defaultdict(list)  # (run, core, risc) -> [(cyc, name, typ)]
for r in rows:
    try:
        cx, cy, risc, cyc, name, typ, run = int(r[1]), int(r[2]), r[3].strip(), int(r[5]), r[10].strip(), r[11].strip(), r[7].strip()
    except Exception:
        continue
    ev[(run, cx, cy, risc)].append((cyc, name, typ))
runs = sorted({k[0] for k in ev}, key=lambda v: int(v) if v.isdigit() else 0)
last = runs[-1]
t0 = min(c for k, v in ev.items() if k[0] == last for c, _, _ in v)
dur = defaultdict(float); cnt = defaultdict(int); kend = []
for (run, cx, cy, risc), v in ev.items():
    if run != last:
        continue
    st = {}
    for cyc, name, typ in sorted(v):
        if "START" in typ: st[name] = cyc
        elif "END" in typ and name in st:
            dur[(risc, name)] += (cyc - st.pop(name)) / 1.35e3; cnt[(risc, name)] += 1
    kend.append((max(c for c, _, _ in v) - t0) / 1.35e3)
ncore = len({(k[1], k[2]) for k in ev if k[0] == last})
print(f"run {last}: cores {ncore}, last event at {max(kend):.2f} us after first")
for k in sorted(dur):
    print(f"  {k[0]:8s} {k[1]:24s} n={cnt[k]:5d} total/core={dur[k] / ncore:7.2f} us")
# kernel start skew
ks = sorted(((min(c for c, n, t in v if "KERNEL" in n) - t0) / 1.35e3, k[1:]) for k, v in ev.items() if k[0] == last and any("KERNEL" in n for _, n, _ in v))
print("kernel start skew: min %.2f med %.2f max %.2f us" % (ks[0][0], ks[len(ks) // 2][0], ks[-1][0]))
if "full" in sys.argv:
    for cy in range(2, 8):
        for risc in ("NCRISC", "BRISC"):
            key = (last, 1, cy, risc)
            if key in ev:
                print(risc, (1, cy), " ".join(f"{n[:5]}{'+' if 'START' in t else '-'}{(c - t0) / 1.35e3:.1f}" for c, n, t in sorted(ev[key])))
# timeline of the NCRISC (weights) of column 0
for cy in sorted({k[2] for k in ev if k[0] == last})[:3]:
    for key in [k for k in ev if k[0] == last and k[2] == cy and k[3] == "NCRISC"][:1]:
        evs = sorted(ev[key])
        print("core", key[1:3], " ".join(f"{n}{'+' if 'START' in t else '-'}{(c - t0) / 1.35e3:.1f}" for c, n, t in evs if n in ("WRDY", "RECV", "XFER", "RSV")))