#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Summary of a device profile of ``profile_ops.py`` (CPU only, standard library): the OPT_BASELINE.md tables.

    python code/scripts/profile_summary.py <ops_perf_results_*.csv> <out.json> [--md out.md]

Reads the Tracy ops CSV (``python -m tracy -r ...``) and its sections ``eager`` .. ``eager_end`` and ``trace`` ..
``trace_end``. The eager section carries per-block (``m:<block>``) and per-conv (``c:<conv key>``) signposts; the
replay has the same op sequence (checked op code by op code), so every replayed op inherits the block / layer of the
eager op at the same position (``c:<key>|post`` marks the glue after a layer: SE gate multiply, residual add,
SceneContext transpose; ``c:<key>.se_sum`` the SE row sum). Per section: program count, device kernel sum, FW sum,
op-to-op gap sum and median, span (first FW start -> last FW end at the CSV's chip frequency), cores per op, math fidelities; tables per category,
block, op code and layer; the top single ops; the largest gaps; and the cost split of the layers
(:func:`precision_split`: compute, conv data movement, the activation split and the fp32 sum / bias / cast).
Input shapes of the CSV are those of the first op that built a cached program (``--no-op-info-cache`` in tracy
shows every op's own), so read shapes of later ops of the same program with care; durations are per op.
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import statistics
from collections import OrderedDict, defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

KERNEL = "DEVICE KERNEL DURATION [ns]"
FW = "DEVICE FW DURATION [ns]"
GAP = "OP TO OP LATENCY [ns]"
START = "DEVICE FW START CYCLE"
END = "DEVICE FW END CYCLE"

CATEGORIES = [  # (category, regex on the op code), first match wins
    ("conv2d / conv_transpose2d", r"Conv2d|conv2d|Conv2D|ConvTranspose"),
    ("halo", r"Halo|halo|UntilizeWithHalo"),
    ("matmul (1x1 conv)", r"Matmul|matmul|Linear"),
    ("slice / DRAM-slicing copies", r"Slice|slice|PaddedSlice"),
    ("binary (term sums, fp32 bias add, residual, SE gate)", r"Binary|binary|Where|where|Ternary"),
    ("typecast", r"Typecast|typecast"),
    ("unary (SiLU, GELU, sigmoid)", r"Unary|unary|Clamp|clamp|Sigmoid"),
    ("concat", r"Concat|concat"),
    ("upsample", r"Upsample|upsample"),
    ("tilize / untilize / layout", r"Tilize|tilize|Untilize|untilize|ToLayout|Layout"),
    ("reshape / transpose / permute", r"Reshape|reshape|Transpose|transpose|Permute|permute|View"),
    ("argmax", r"ArgMax|Argmax|argmax"),
    ("reduce (SE / context row sums)", r"Reduce|reduce|Sum"),
    ("move / copy / reshard", r"Move|move|Copy|copy|Shard|shard|Interleaved|interleaved|Reallocate"),
    ("pad / fill", r"Pad|pad|Fill|fill"),
]


def category(op: str) -> str:
    for name, pat in CATEGORIES:
        if re.search(pat, op):
            return name
    return "other"


def num(row: Dict[str, str], key: str) -> float:
    v = (row.get(key) or "").strip()
    try:
        return float(v) if v else 0.0
    except ValueError:
        return 0.0


def logical(v: str) -> str:
    """'32[16]' (padded[logical]) -> '16'; '7' -> '7'."""
    m = re.match(r"\s*([0-9]+)(?:\[([0-9]+)\])?", v or "")
    return (m.group(2) or m.group(1)) if m else "?"


def chip_freq_mhz(csv_path: Path) -> Optional[float]:
    for cand in [csv_path.parent / "profile_log_device.csv", csv_path.parent.parent / "profile_log_device.csv",
                 csv_path.parent / ".logs" / "profile_log_device.csv", csv_path.parent.parent / ".logs" / "profile_log_device.csv"]:
        try:
            head = cand.read_text(errors="ignore").splitlines()[0]
        except (OSError, IndexError):
            continue
        m = re.search(r"CHIP_FREQ\[MHz\]:\s*([0-9.]+)", head)
        if m:
            return float(m.group(1))
    return None


def sections(rows: List[Dict[str, str]], name_key: str) -> Dict[str, List[Tuple[Dict[str, str], str, str]]]:
    """{"eager": [(row, block, layer)], "trace": [...]} from the signposts."""
    out: Dict[str, List[Tuple[Dict[str, str], str, str]]] = {"eager": [], "trace": []}
    cur: Optional[str] = None
    block, layer = "", ""
    for r in rows:
        if (r.get("OP TYPE") or "") == "signpost":
            code = r.get(name_key, "")
            if code in ("eager", "trace"):
                cur, block, layer = code, "", ""
            elif code in ("eager_end", "trace_end"):
                cur = None
            elif code.startswith("m:"):
                block, layer = code[2:], ""
            elif code.startswith("c:"):
                layer = code[2:]
            continue
        if cur is not None:
            out[cur].append((r, block, layer))
    return out


COMPUTE = re.compile(r"Conv2d|Matmul|ConvTranspose")
GLUE = re.compile(r"BinaryNg|Typecast|Unary")


def precision_split(sec, kern, name_key) -> Dict[str, Any]:
    """The cost structure of the layers (``c:<key>`` signposts, ``|post`` glue excluded): per layer the conv /
    matmul compute kernels, the other ops between the first and the last compute op (conv data movement: halo,
    shard moves, DRAM slices, reshapes), the elementwise glue before the first compute op (the float32 activation
    split into bf16 terms) and after the last one (the fp32 sum of the products, the fp32 bias add, the activation,
    the cast). ``products`` counts the compute ops: weight terms x activation terms minus lo x lo
    (``SCENESEG_WEIGHT_TERMS`` / ``SCENESEG_ACT_TERMS``), times the DRAM slices of a sliced conv."""
    layers: Dict[str, List[int]] = OrderedDict()
    for i, (r, b, l) in enumerate(sec):
        if l and not l.endswith("|post"):
            layers.setdefault(l, []).append(i)
    tot = {"layers": 0, "compute_us": 0.0, "movement_us": 0.0, "split_us": 0.0, "finish_us": 0.0,
           "compute_ops": 0, "ops": 0}
    per: List[Dict[str, Any]] = []
    for l, idx in layers.items():
        codes = [sec[i][0].get(name_key, "") for i in idx]
        comp = [k for k, c in enumerate(codes) if COMPUTE.search(c)]
        if not comp:
            continue
        a, z = comp[0], comp[-1]
        d = {"layer": l, "ops": len(idx), "compute_ops": len(comp),
             "compute_us": sum(kern[idx[k]] for k in comp),
             "movement_us": sum(kern[idx[k]] for k in range(a, z + 1) if k not in comp),
             "split_us": sum(kern[idx[k]] for k in range(a)),
             "finish_us": sum(kern[idx[k]] for k in range(z + 1, len(idx)))}
        per.append({k: (round(v, 1) if isinstance(v, float) else v) for k, v in d.items()})
        tot["layers"] += 1
        for k in ("compute_us", "movement_us", "split_us", "finish_us", "compute_ops", "ops"):
            tot[k] += d[k]
    out = {k: (round(v, 1) if isinstance(v, float) else v) for k, v in tot.items()}
    out["per_layer"] = sorted(per, key=lambda x: -(x["compute_us"] + x["movement_us"] + x["split_us"] + x["finish_us"]))
    return out


def summarize(sec: List[Tuple[Dict[str, str], str, str]], name_key: str, freq: float, top: int = 10) -> Dict[str, Any]:
    rows = [r for r, _, _ in sec]
    kern = [num(r, KERNEL) / 1e3 for r in rows]
    fw = [num(r, FW) / 1e3 for r in rows]
    gaps = [num(r, GAP) / 1e3 for r in rows]
    if gaps:
        gaps[0] = 0.0            # the first op's gap reaches back across the section boundary
    starts = [num(r, START) for r in rows if num(r, START) > 0]
    ends = [num(r, END) for r in rows if num(r, END) > 0]
    span = (max(ends) - min(starts)) / freq if starts and ends else None
    total = sum(kern) or 1.0

    def group(keyfn) -> List[Dict[str, Any]]:
        agg: Dict[str, List[float]] = OrderedDict()
        for (r, b, l), k, g in zip(sec, kern, gaps):
            key = keyfn(r, b, l)
            a = agg.setdefault(key, [0, 0.0, 0.0])
            a[0] += 1
            a[1] += k
            a[2] += g
        return [{"key": key, "ops": int(c), "kernel_us": round(t, 1), "share": round(t / total, 4),
                 "gap_us": round(gs, 1)} for key, (c, t, gs) in sorted(agg.items(), key=lambda kv: -kv[1][1])]

    fid: Dict[str, int] = defaultdict(int)
    cores = []
    for r in rows:
        f = (r.get("MATH FIDELITY") or "").strip()
        if f:
            fid[f] += 1
        cores.append(num(r, "CORE COUNT"))
    order = sorted(range(len(rows)), key=lambda i: -kern[i])[:top]
    top_ops = []
    for i in order:
        r, b, l = sec[i]
        shape = "x".join(logical(r.get(f"INPUT_0_{d}_PAD[LOGICAL]") or r.get(f"INPUT_0_{d}") or "") for d in "WZYX")
        top_ops.append({"index": i, "op": r.get(name_key), "block": b, "layer": l, "kernel_us": round(kern[i], 1),
                        "fw_us": round(fw[i], 1), "cores": int(num(r, "CORE COUNT")),
                        "fidelity": (r.get("MATH FIDELITY") or "").strip(), "input0": shape,
                        "in0_mem": (r.get("INPUT_0_MEMORY") or "").strip(),
                        "out0_mem": (r.get("OUTPUT_0_MEMORY") or "").strip()})
    gap_order = sorted(range(1, len(rows)), key=lambda i: -gaps[i])[:top]
    top_gaps = [{"index": i, "op": rows[i].get(name_key), "after": rows[i - 1].get(name_key), "gap_us": round(gaps[i], 1),
                 "block": sec[i][1], "layer": sec[i][2]} for i in gap_order]
    nz_gaps = [g for g in gaps[1:] if g > 0]
    two_term = precision_split(sec, kern, name_key)
    return {
        "ops": len(rows), "kernel_sum_us": round(sum(kern), 1), "fw_sum_us": round(sum(fw), 1),
        "op2op_sum_us": round(sum(gaps), 1),
        "op2op_median_us": round(statistics.median(nz_gaps), 2) if nz_gaps else None,
        "span_us": round(span, 1) if span else None, "freq_mhz": freq,
        "cores_median": statistics.median(cores) if cores else None,
        "ops_below_120_cores": sum(1 for c in cores if 0 < c < 120),
        "fidelity": dict(fid),
        "by_op": group(lambda r, b, l: r.get(name_key, "?")),
        "by_category": group(lambda r, b, l: category(r.get(name_key, "?"))),
        "by_block": group(lambda r, b, l: b or "?"),
        "by_layer": group(lambda r, b, l: l or f"[{b or '?'} glue]"),
        "top_ops": top_ops, "top_gaps": top_gaps, "precision_split": two_term,
    }


def md_table(rows: List[Dict[str, Any]], n: int, label: str) -> str:
    out = [f"| {label} | ops | kernel us | share | op-to-op us |", "|---|---:|---:|---:|---:|"]
    for r in rows[:n]:
        out.append(f"| {r['key']} | {r['ops']} | {r['kernel_us']:.1f} | {100 * r['share']:.1f} % | {r['gap_us']:.1f} |")
    return "\n".join(out)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("csv")
    ap.add_argument("out_json")
    ap.add_argument("--md")
    ap.add_argument("--freq-mhz", type=float, default=None)
    a = ap.parse_args()
    path = Path(a.csv)
    with open(path, newline="") as f:
        rows = list(csv.DictReader(f))
    name_key = "OP CODE" if rows and "OP CODE" in rows[0] else "OP NAME"
    freq = a.freq_mhz or chip_freq_mhz(path) or 1350.0
    secs = sections(rows, name_key)
    eager, trace = secs["eager"], secs["trace"]
    # the replay inherits the block / layer of the eager op at the same position (identical op sequence)
    same = len(eager) == len(trace) and all(e[0].get(name_key) == t[0].get(name_key) for e, t in zip(eager, trace))
    if same:
        trace = [(r, e[1], e[2]) for (r, _, _), e in zip(trace, eager)]
    res = {"csv": str(path), "columns": list(rows[0].keys()) if rows else [], "freq_mhz": freq,
           "trace_matches_eager_op_sequence": same,
           "trace_ids": sorted({(r.get("METAL TRACE ID") or "").strip() for r, _, _ in secs["trace"]}),
           "replay_sessions": sorted({(r.get("METAL TRACE REPLAY SESSION ID") or "").strip() for r, _, _ in secs["trace"]}),
           "eager": summarize(eager, name_key, freq), "trace": summarize(trace, name_key, freq)}
    Path(a.out_json).write_text(json.dumps(res, indent=1) + "\n")
    t = res["trace"]
    lines = [f"csv: {path}", f"freq {freq} MHz; trace op sequence == eager: {same}; trace ids {res['trace_ids']}; "
             f"replay sessions {res['replay_sessions']}"]
    for s in ("eager", "trace"):
        x = res[s]
        lines.append(f"{s}: ops {x['ops']}  kernel {x['kernel_sum_us']:.1f} us  fw {x['fw_sum_us']:.1f} us  "
                     f"op2op {x['op2op_sum_us']:.1f} us (median {x['op2op_median_us']})  span {x['span_us']} us  "
                     f"cores median {x['cores_median']}  <120 cores: {x['ops_below_120_cores']}  fidelity {x['fidelity']}")
    lines += ["", md_table(t["by_category"], 20, "category (trace)"), "", md_table(t["by_block"], 20, "block (trace)"),
              "", md_table(t["by_op"], 25, "op code (trace)"), "", md_table(t["by_layer"], 25, "layer (trace)"), "",
              "| # | op | block | layer | kernel us | cores | fidelity | input0 | in mem | out mem |",
              "|---:|---|---|---|---:|---:|---|---|---|---|"]
    for k, o in enumerate(t["top_ops"], 1):
        lines.append(f"| {k} | {o['op']} | {o['block']} | {o['layer']} | {o['kernel_us']:.1f} | {o['cores']} | "
                     f"{o['fidelity']} | {o['input0']} | {o['in0_mem']} | {o['out0_mem']} |")
    ps = t["precision_split"]
    lines += ["", "layer cost split (trace): " + json.dumps({k: v for k, v in ps.items() if k != "per_layer"}),
              "", "| layer | ops | compute ops | compute us | movement us | split us | finish us |",
              "|---|---:|---:|---:|---:|---:|---:|"]
    for x in ps["per_layer"][:25]:
        lines.append(f"| {x['layer']} | {x['ops']} | {x['compute_ops']} | {x['compute_us']} | {x['movement_us']} | "
                     f"{x['split_us']} | {x['finish_us']} |")
    lines += ["", "largest op-to-op gaps (trace): " + json.dumps(t["top_gaps"])]
    text = "\n".join(lines)
    print(text)
    if a.md:
        Path(a.md).write_text(text + "\n")


if __name__ == "__main__":
    main()