"""Per (prev op, op) pair: count, mean op-to-op gap, mean kernel duration of op; last forward of scripts/prof.py.
usage: python scripts/prof_pairs.py <ops_perf.csv> [<ops_perf B.csv>]   (B: printed side by side by pair key)"""
import csv, sys
from collections import defaultdict


def load(p):
    rows = [r for r in csv.DictReader(open(p)) if r.get("DEVICE KERNEL DURATION [ns]")]
    st = [i for i, r in enumerate(rows[:-1]) if r["OP CODE"].startswith("GenericOp")
          and rows[i + 1]["OP CODE"].startswith("MinimalMatmul")]
    last = rows[st[-1]:]
    g = defaultdict(lambda: [0, 0.0, 0.0])
    for i in range(1, len(last)):
        k = (last[i - 1]["OP CODE"][:24], last[i]["OP CODE"][:24])
        s = g[k]; s[0] += 1
        s[1] += float(last[i]["OP TO OP LATENCY [ns]"] or 0) / 1e3
        s[2] += float(last[i]["DEVICE KERNEL DURATION [ns]"]) / 1e3
    kern = sum(float(r["DEVICE KERNEL DURATION [ns]"]) for r in last) / 1e3
    gap = sum(float(r["OP TO OP LATENCY [ns]"] or 0) for r in last[1:]) / 1e3
    return g, len(last), kern, gap


A = load(sys.argv[1]); B = load(sys.argv[2]) if len(sys.argv) > 2 else None
print(f"A: ops={A[1]} kernel={A[2]:.1f} us gaps={A[3]:.1f} us" + (f"   B: ops={B[1]} kernel={B[2]:.1f} us gaps={B[3]:.1f} us" if B else ""))
keys = set(A[0]) | (set(B[0]) if B else set())
def tot(k, X):
    s = X[0].get(k); return (s[1] + s[2]) if s else 0.0
for k in sorted(keys, key=lambda k: -max(tot(k, A), tot(k, B) if B else 0)):
    line = f"{k[0]:24s} -> {k[1]:24s}"
    for X in ([A, B] if B else [A]):
        s = X[0].get(k)
        line += f" | n={s[0]:3d} gap={s[1]/s[0]:6.2f} dur={s[2]/s[0]:6.2f}" if s else " | " + " " * 31
    print(line)