# SPDX-License-Identifier: Apache-2.0
"""Summarize a tracy ops_perf_results CSV of scripts/prof.py: the LAST forward (segmented on the
ingest tilize op), op count, kernel-sum, op-to-op gap sum, time per op code and per ViT/head region.
usage: python scripts/prof_summary.py <ops_perf_results.csv> [--out trimmed.csv]"""
import csv
import sys
from collections import OrderedDict, defaultdict

path = sys.argv[1]
rows = list(csv.DictReader(open(path)))
rows = [r for r in rows if r.get("DEVICE KERNEL DURATION [ns]")]
starts = [i for i, r in enumerate(rows) if r["OP CODE"].startswith("Tilize")]
if not starts:  # patch=device: the forward starts with the input copy feeding the im2col generic_op
    # (the im2col generic_op immediately followed by the patch-embed minimal_matmul)
    starts = [i for i, r in enumerate(rows[:-1]) if r["OP CODE"].startswith("GenericOp")
              and rows[i + 1]["OP CODE"].startswith("MinimalMatmul")]
last = rows[starts[-1]:]
kern = sum(float(r["DEVICE KERNEL DURATION [ns]"]) for r in last) / 1e3
gaps = [float(r["OP TO OP LATENCY [ns]"] or 0) for r in last][1:]
print(f"forwards found: {len(starts)}  ops in last forward: {len(last)}")
print(f"kernel sum {kern:.1f} us   op-to-op gaps (excl. first) {sum(gaps)/1e3:.1f} us")
by = OrderedDict()
for r in last:
    k = r["OP CODE"]
    n, t = by.get(k, (0, 0.0))
    by[k] = (n + 1, t + float(r["DEVICE KERNEL DURATION [ns]"]) / 1e3)
for k, (n, t) in sorted(by.items(), key=lambda kv: -kv[1][1]):
    print(f"  {k:55s} n={n:4d}  {t:9.1f} us")
if "--out" in sys.argv:
    o = sys.argv[sys.argv.index("--out") + 1]
    keep = ["OP CODE", "CORE COUNT", "MATH FIDELITY", "DEVICE KERNEL DURATION [ns]", "OP TO OP LATENCY [ns]",
            "INPUT_0_Y_PAD[LOGICAL]", "INPUT_0_X_PAD[LOGICAL]", "INPUT_0_DATATYPE", "INPUT_0_MEMORY",
            "INPUT_1_Y_PAD[LOGICAL]", "INPUT_1_X_PAD[LOGICAL]", "INPUT_1_DATATYPE", "OUTPUT_0_MEMORY"]
    with open(o, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["idx"] + keep)
        for i, r in enumerate(last):
            w.writerow([i] + [r.get(k, "") for k in keep])
    print("wrote", o)