"""Round 8: add+LN zone TIMELINE (start / end relative to the kernel start, median over cores) for the three TRISCs and
BRISC.  usage: ADDLN_EXP=prof TT_METAL_DEVICE_PROFILER=1 addln_tl.py"""
import sys, os, 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
dev = open_dev()
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)
x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
y = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
C = tk.addln_constants(torch.ones(1280), torch.zeros(1280), dev)
G = dev.compute_with_storage_grid_size()
sg = ttnn.num_cores_to_corerangeset(10, G, row_wise=True)
hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
                         ttnn.ShardSpec(sg, [192, 128], ttnn.ShardOrientation.ROW_MAJOR))
for _ in range(4):
    s, h = tk.add_layernorm(x, y, C, 1e-6, out_memory_config=hmem, out_dtype=ttnn.bfloat8_b)
    ttnn.synchronize_device(dev); ttnn.deallocate(s); ttnn.deallocate(h)
ttnn.ReadDeviceProfiler(dev)
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)
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)].append((cyc, risc, name, typ))
last = sorted({k[0] for k in ev}, key=lambda v: int(v) if v.isdigit() else 0)[-1]
tl = defaultdict(lambda: defaultdict(list))
for (run, cx, cy), v in ev.items():
    if run != last:
        continue
    t0 = min(c for c, rk, n, t in v if "KERNEL" in n)
    st = {}
    for c, rk, n, t in sorted(v):
        if "START" in t:
            st[(rk, n)] = c
        elif "END" in t and (rk, n) in st:
            tl[rk][n].append(((st.pop((rk, n)) - t0) / 1.35e3, (c - t0) / 1.35e3))
for rk in sorted(tl):
    for n, v in sorted(tl[rk].items(), key=lambda kv: sorted(kv[1])[len(kv[1]) // 2][0]):
        s_ = sorted(a for a, b in v)[len(v) // 2]; e_ = sorted(b for a, b in v)[len(v) // 2]
        print(f"{rk:8s} {n:22s} start {s_:6.2f} end {e_:6.2f} us")