#!/usr/bin/env python3
"""聚合一次批量跑（run_batch.py 产物目录）的结果为 summary.csv。

数据源是各算子 run 目录里的 run.log（run_pytest.sh 已去 ANSI 色）：
- 逐表解析 `Operator: <op_name> (dtype=..., mode=..., level=...)` 标题
  （兼容 _pretty_report_plugin 与上游原生两种格式）；
- 逐行解析 SUCCESS/FAILED 表行：状态 + 前三个数值列固定为
  torch_lat / gems_lat / speedup（tflops/gbps 等追加列收进 extra_metrics），
  余下为 size_detail；
- 汇入 run_batch.py 写的 ops_status.csv（每算子总状态 / 备注）。

也可单独使用：python batch_summary.py runs/batch_xxx [-o summary.csv]
"""

from __future__ import annotations

import argparse
import csv
import itertools
import os
import re
import statistics
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# 兼容两种标题：
#   pretty:   Operator: softmax  (dtype=torch.float16, mode=kernel, level=core)
#   upstream: Operator: softmax  Performance Test (dtype=torch.float16, mode=kernel,level=core)
OPERATOR_HEADER_RE = re.compile(
    r"^Operator:\s+(?P<name>\S+)\s+(?:Performance Test\s*)?"
    r"\(dtype=(?P<dtype>[^,]+),\s*mode=(?P<mode>[^,]+),\s*level=(?P<level>[^)]+)\)"
)
ROW_RE = re.compile(r"^(?P<status>SUCCESS|FAILED)\s+(?P<rest>\S.*)$")
NUM_TOKEN_RE = re.compile(
    r"^(?:N/A|nan|-?inf"
    r"|-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)$"
)
PYTEST_SKIPPED_RE = re.compile(r"(?P<n>\d+)\s+(?:skipped|xfailed)\b")
DIRECT_NO_CUDAGRAPH_RE = re.compile(r"BENCHMARK_DIRECT_NO_CUDAGRAPH\b")
REPLAY_FALLBACK_RE = re.compile(r"AUTOTUNE_REPLAY_FALLBACK\b")

SUMMARY_FIELDS = [
    "repo", "op", "backend", "device", "timer",
    "fg_op_name", "dtype", "case_id", "row_status",
    "torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail",
    "op_status", "notes",
]


def _identity_from_status(status_row: Dict[str, str], op_dir: Path) -> Dict[str, str]:
    """Backend/device/timer from ops_status.csv, falling back to run_identity.json."""
    backend = (status_row.get("backend") or "").strip()
    device = (status_row.get("device") or status_row.get("gpu") or "").strip()
    timer = (status_row.get("timer") or "").strip()
    ident_path = op_dir / "run_identity.json"
    if ident_path.is_file() and not (backend and device and timer):
        try:
            from _run_identity import load_identity
            ident = load_identity(str(op_dir))
            backend = backend or str(ident.get("backend") or "")
            device = device or str(ident.get("device") or ident.get("gpu") or "")
            timer = timer or str(ident.get("timer") or "")
        except Exception:
            pass
    return {"backend": backend, "device": device, "timer": timer}


def parse_run_log(text: str) -> List[Dict[str, str]]:
    """把 run.log 里的全部结果表解析为行列表（跨多个 dtype 表）。"""
    rows: List[Dict[str, str]] = []
    cur_op, cur_dtype = "", ""
    for line in text.splitlines():
        stripped = line.strip()
        header = OPERATOR_HEADER_RE.match(stripped)
        if header:
            cur_op = header.group("name")
            cur_dtype = header.group("dtype").strip()
            continue
        m = ROW_RE.match(stripped)
        if not m:
            continue
        tokens = m.group("rest").split()
        metrics: List[str] = []
        while tokens and len(metrics) < 7 and NUM_TOKEN_RE.match(tokens[0]):
            metrics.append(tokens.pop(0))
        # 前三个数值列固定：torch_lat, gems_lat, speedup（上游列序承诺不变）
        while len(metrics) < 3:
            metrics.append("")
        rows.append({
            "fg_op_name": cur_op,
            "dtype": cur_dtype,
            "row_status": m.group("status"),
            "torch_lat": metrics[0],
            "gems_lat": metrics[1],
            "speedup": metrics[2],
            "extra_metrics": " ".join(metrics[3:]),
            "size_detail": " ".join(tokens),
        })
    return rows


def pytest_skipped_count(text: str) -> int:
    """runtime skipif 与上游 xfail 标记都算"无测量信号"，供 SKIP 判定。"""
    count = 0
    for m in PYTEST_SKIPPED_RE.finditer(text[-4000:]):
        count += int(m.group("n"))
    return count


def marker_notes(text: str) -> str:
    """cudagraph 回退 / replay 兜底标记计数，供备注列（口径解读用）。"""
    notes = []
    n = len(DIRECT_NO_CUDAGRAPH_RE.findall(text))
    if n:
        notes.append(f"no_cudagraph_rows={n}")
    n = len(REPLAY_FALLBACK_RE.findall(text))
    if n:
        notes.append(f"replay_fallback={n}")
    return ";".join(notes)


def classify(rc: int, rows: List[Dict[str, str]], log_text: str) -> str:
    """算子级状态：FAIL / SKIP / PASS（TIMEOUT 由 run_batch 在外层判）。"""
    if rc != 0:
        return "FAIL"
    if not rows:
        return "SKIP" if pytest_skipped_count(log_text) else "FAIL"
    if any(r["row_status"] != "SUCCESS" for r in rows):
        return "FAIL"
    return "PASS"


def _load_ops_status(batch_dir: Path) -> Dict[str, Dict[str, str]]:
    path = batch_dir / "ops_status.csv"
    if not path.is_file():
        return {}
    with path.open(newline="") as f:
        return {f"{r['repo']}/{r['run_name']}": r for r in csv.DictReader(f)}


def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
    batch_dir = batch_dir.resolve()
    out_path = out_path or (batch_dir / "summary.csv")
    status_by_dir = _load_ops_status(batch_dir)

    all_rows: List[Dict[str, str]] = []
    for repo_dir in sorted(p for p in batch_dir.iterdir() if p.is_dir()):
        repo = repo_dir.name
        for op_dir in sorted(p for p in repo_dir.iterdir() if p.is_dir()):
            log = op_dir / "run.log"
            if not log.is_file():
                continue
            text = log.read_text(errors="replace")
            status_row = status_by_dir.get(f"{repo}/{op_dir.name}", {})
            op = status_row.get("op") or op_dir.name
            ident = _identity_from_status(status_row, op_dir)
            table_rows = parse_run_log(text)
            if not table_rows:
                all_rows.append({
                    "repo": repo, "op": op, **ident, "fg_op_name": "", "dtype": "",
                    "case_id": "", "row_status": "", "torch_lat": "",
                    "gems_lat": "", "speedup": "", "extra_metrics": "",
                    "size_detail": "",
                    "op_status": status_row.get("status", ""),
                    "notes": status_row.get("notes", ""),
                })
                continue
            case_counter: Dict[Tuple[str, str], int] = {}
            for row in table_rows:
                key = (row["fg_op_name"], row["dtype"])
                idx = case_counter.get(key, 0)
                case_counter[key] = idx + 1
                all_rows.append({
                    "repo": repo, "op": op, **ident,
                    "case_id": f"case_{idx:03d}",
                    "op_status": status_row.get("status", ""),
                    "notes": status_row.get("notes", ""),
                    **row,
                })

    tmp = out_path.with_name(out_path.name + ".tmp")
    with tmp.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=SUMMARY_FIELDS)
        writer.writeheader()
        writer.writerows(all_rows)
    os.replace(tmp, out_path)
    return out_path


def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]:
    """把两次 run.log 的测量行按 (op, dtype, size_detail) 配对。

    speedup 列是 gems 对 torch，不是前后对比，所以这里用 gems_lat 自己算
    before/after - 1。只有两侧都 SUCCESS 且延迟为正的行才计入；对不上的行
    单独返回，避免用不完整的配对报一个看起来完整的区间。
    """
    def index(text: str) -> Dict[Tuple[str, str, str], float]:
        out: Dict[Tuple[str, str, str], float] = {}
        for row in parse_run_log(text):
            if row.get("row_status") != "SUCCESS":
                continue
            try:
                lat = float(row.get("gems_lat", ""))
            except ValueError:
                continue
            if lat > 0:
                out[(row.get("fg_op_name", ""), row.get("dtype", ""),
                     row.get("size_detail", ""))] = lat
        return out

    before, after = index(before_text), index(after_text)
    paired: List[Dict[str, str]] = []
    for key in sorted(before.keys() & after.keys()):
        op, dtype, size = key
        b, a = before[key], after[key]
        paired.append({
            "fg_op_name": op, "dtype": dtype, "size_detail": size,
            "before_lat": f"{b:.6f}", "after_lat": f"{a:.6f}",
            "speedup_percent": f"{(b / a - 1.0) * 100.0:.2f}",
        })
    return paired


def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]:
    """打印一个算子的前后对比：逐行 speedup 与整体区间。"""
    before_text = before.read_text(encoding="utf-8", errors="replace")
    after_text = after.read_text(encoding="utf-8", errors="replace")
    rows = ab_rows(before_text, after_text)
    name = label or (rows[0]["fg_op_name"] if rows else "?")
    if not rows:
        print(f"[ab] {name}: 没有可配对的测量行"
              f"（before={before} after={after}）")
        return rows
    values = [float(r["speedup_percent"]) for r in rows]
    for r in rows:
        print(f"  {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}%  "
              f"{r['before_lat']} -> {r['after_lat']}  {r['size_detail']}")
    fallback = ""
    if DIRECT_NO_CUDAGRAPH_RE.search(before_text) or \
            DIRECT_NO_CUDAGRAPH_RE.search(after_text):
        # 回退行用的是 do_bench 口径，与 cudagraph 行不同表可比性
        fallback = "  [含 BENCHMARK_DIRECT_NO_CUDAGRAPH 回退行]"
    print(f"[ab] {name}: {len(rows)} 行配对，speedup "
          f"{min(values):+.2f}% ~ {max(values):+.2f}%{fallback}")
    return rows


def _lat_index(text: str) -> Dict[Tuple[str, str, int], float]:
    """(dtype, size_detail, occurrence) -> gems latency, SUCCESS rows only.

    Occurrence disambiguates a shape a benchmark measures twice in one table,
    which BLAS does at --level comprehensive.
    """
    out: Dict[Tuple[str, str, int], float] = {}
    seen: Dict[Tuple[str, str], int] = {}
    for row in parse_run_log(text):
        if row.get("row_status") != "SUCCESS":
            continue
        try:
            lat = float(row.get("gems_lat", ""))
        except ValueError:
            continue
        if lat <= 0:
            continue
        base = (row.get("dtype", ""), row.get("size_detail", ""))
        i = seen.get(base, 0)
        seen[base] = i + 1
        out[base + (i,)] = lat
    return out


def ab_round_rows(root: Path) -> Tuple[List[Dict[str, str]], List[float], int]:
    """Aggregate an N-round A/B laid out as <root>/r<k>_<A|B>/run.log.

    Returns the per-shape rows, the per-round deltas, and the round count.

    A single pair cannot say whether its own verdict is stable, which is the
    question that matters once a delta lands near zero: measured on an S5000,
    topk eliminated 32 converts yet 17 of its 42 shapes changed the SIGN of
    their delta between two equally valid runs. Two dispersions answer different
    halves of that and both are reported:

      side spread  max/min-1 within one side's samples -> how noisy the machine
                   was while measuring
      round spread min~max of the per-round deltas -> whether the VERDICT holds,
                   which is what a single pair cannot show
    """
    rounds: List[Tuple[Dict, Dict]] = []
    for k in itertools.count(1):
        a, b = root / f"r{k}_A" / "run.log", root / f"r{k}_B" / "run.log"
        if not (a.is_file() and b.is_file()):
            break
        rounds.append((
            _lat_index(a.read_text(encoding="utf-8", errors="replace")),
            _lat_index(b.read_text(encoding="utf-8", errors="replace")),
        ))
    if not rounds:
        return [], [], 0

    keys = set(rounds[0][0]) & set(rounds[0][1])
    for a, b in rounds[1:]:
        keys &= set(a) & set(b)

    per_round = [
        statistics.median((a[k] / b[k] - 1.0) * 100.0 for k in keys)
        for a, b in rounds
    ] if keys else []

    rows: List[Dict[str, str]] = []
    for key in sorted(keys):
        av = [a[key] for a, _ in rounds]
        bv = [b[key] for _, b in rounds]
        am, bm = statistics.median(av), statistics.median(bv)
        deltas = [(a / b - 1.0) * 100.0 for a, b in zip(av, bv)]
        rows.append({
            "dtype": key[0], "size_detail": key[1],
            "before_lat": f"{am:.6f}", "after_lat": f"{bm:.6f}",
            "speedup_percent": f"{(am / bm - 1.0) * 100.0:.2f}",
            "delta_min": f"{min(deltas):.2f}", "delta_max": f"{max(deltas):.2f}",
            "sign_stable": "yes" if all(d > 0 for d in deltas)
                           or all(d <= 0 for d in deltas) else "NO",
            "before_spread": f"{(max(av) / min(av) - 1.0) * 100.0:.2f}",
            "after_spread": f"{(max(bv) / min(bv) - 1.0) * 100.0:.2f}",
        })
    return rows, per_round, len(rounds)


def report_ab_rounds(root: Path, label: str = "") -> List[Dict[str, str]]:
    """Print the multi-round A/B: per-shape medians plus both dispersions."""
    rows, per_round, n = ab_round_rows(root)
    name = label or root.name
    if n < 2:
        print(f"[ab-rounds] {name}: 需要至少 2 轮 <root>/r<k>_<A|B>/run.log"
              f"（找到 {n} 轮）")
        return rows
    if not rows:
        print(f"[ab-rounds] {name}: {n} 轮，但没有全轮次都成功的配对行")
        return rows
    for r in rows:
        flag = "" if r["sign_stable"] == "yes" else "  符号不稳"
        print(f"  {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}%  "
              f"[{float(r['delta_min']):+.2f}%~{float(r['delta_max']):+.2f}%]  "
              f"{r['before_lat']} -> {r['after_lat']}  "
              f"{r['size_detail']}{flag}")
    med = [float(r["speedup_percent"]) for r in rows]
    spread = [float(r["before_spread"]) for r in rows] + \
             [float(r["after_spread"]) for r in rows]
    unstable = sum(1 for r in rows if r["sign_stable"] != "yes")
    print(f"[ab-rounds] {name}: {n} 轮 × {len(rows)} shape，"
          f"speedup 中位 {statistics.median(med):+.2f}%，"
          f"逐轮 {min(per_round):+.2f}% ~ {max(per_round):+.2f}%，"
          f"单侧噪声中位 {statistics.median(spread):.2f}%")
    if unstable:
        # The verdict, not the machine, is what these rows fail to pin down.
        print(f"[ab-rounds] {name}: {unstable}/{len(rows)} 个 shape 的 delta "
              "在轮次间改变符号——这些行的方向不可结论，与计时口径无关")
    return rows


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("batch_dir", type=Path, nargs="?",
                   help="run_batch.py 的产物目录")
    p.add_argument("-o", "--output", type=Path, default=None,
                   help="输出 CSV（默认 <batch_dir>/summary.csv）")
    p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"),
                   help="两份 run.log，按 shape 配对输出该算子的 speedup 区间")
    p.add_argument("--ab-rounds", type=Path, default=None, metavar="ROOT",
                   help="多轮 A/B 目录，含 r<k>_A/run.log 与 r<k>_B/run.log；"
                        "逐 shape 取轮次中位，并报逐轮区间与符号稳定性")
    p.add_argument("--label", default="", help="--ab / --ab-rounds 报告里的算子名")
    args = p.parse_args()
    if args.ab_rounds:
        if not args.ab_rounds.is_dir():
            raise SystemExit(f"目录不存在: {args.ab_rounds}")
        report_ab_rounds(args.ab_rounds, args.label)
        return
    if args.ab:
        for path in args.ab:
            if not path.is_file():
                raise SystemExit(f"run.log 不存在: {path}")
        report_ab(args.ab[0], args.ab[1], args.label)
        return
    if args.batch_dir is None:
        raise SystemExit("需要 batch_dir，或用 --ab BEFORE AFTER")
    if not args.batch_dir.is_dir():
        raise SystemExit(f"目录不存在: {args.batch_dir}")
    out = summarize(args.batch_dir, args.output)
    with out.open(newline="") as f:
        rows = list(csv.DictReader(f))
    n_ok = sum(1 for r in rows if r["row_status"] == "SUCCESS")
    n_bad = sum(1 for r in rows if r["row_status"] == "FAILED")
    print(f"[summary] {out}: {len(rows)} 行（SUCCESS={n_ok} FAILED={n_bad}）",
          file=sys.stderr)


if __name__ == "__main__":
    main()