#!/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 = "",
              strict: bool = True) -> List[Dict[str, str]]:
    """打印一个算子的前后对比：逐行 speedup 与整体区间。

    先核对两侧口径；不一致时默认拒绝出报告，因为不可比的两次运行会产出
    一份和有效报告长得一模一样的表。
    """
    if not _report_caliber([("before", before.parent), ("after", after.parent)],
                           label or "ab", strict):
        raise CaliberMismatch()
    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


# Fields whose disagreement makes two runs incomparable. Measured magnitudes,
# so the list is not precautionary: an unpinned autotune config moved
# router_gemm by 8pp; the timer moved replication_pad3d from +43% to +124.5%;
# a shared tuner cache made mm look like 945 compiled variants had become 33;
# --level core sees 5 mm shapes where comprehensive sees 29.
CALIBER_FIELDS = (
    "backend", "device", "timer", "bench_level", "shape_mode",
    "tuner_cache", "direct_repeat",
)


def caliber_of(run_dir: Path) -> Optional[Dict[str, str]]:
    """Caliber fields from a run's run_identity.json, or None if unreadable."""
    if not (run_dir / "run_identity.json").is_file():
        return None
    try:
        from _run_identity import load_identity
        ident = load_identity(str(run_dir))
    except Exception:
        return None
    return {k: str(ident.get(k, "")) for k in CALIBER_FIELDS} | {
        "autotune_mode": str(ident.get("autotune_mode", ""))
    }


def caliber_complaints(runs: List[Tuple[str, Path]]) -> List[str]:
    """Everything about this set of runs that makes their comparison unsound.

    Reported rather than silently tolerated because the failure is invisible in
    the numbers: two runs on different calibers produce a report that looks
    exactly like a valid one. This repository has already published three
    figures that were an artifact of one such difference.
    """
    out: List[str] = []
    cals = [(name, caliber_of(d)) for name, d in runs]
    missing = [n for n, c in cals if c is None]
    if missing:
        out.append(f"口径不可核实：{', '.join(missing)} 缺 run_identity.json"
                   "（早于该字段的产物）")
    known = [(n, c) for n, c in cals if c is not None]
    if len(known) < 2:
        return out

    for field in CALIBER_FIELDS:
        seen = {c[field] for _, c in known}
        if len(seen) > 1:
            vals = ", ".join(f"{n}={c[field] or '<空>'}" for n, c in known)
            out.append(f"{field} 不一致：{vals}")

    # A/B needs one side to pin the configs the other replays. Two record runs
    # each sweep independently, and the winner can differ; the README puts that
    # at the 10% level and one-directional per shape, which is enough to swamp
    # what is being measured.
    modes = [(n, c["autotune_mode"]) for n, c in known]
    if all(m.startswith("record") for _, m in modes):
        out.append("两侧都是 record 模式：config 未锁定，各自 sweep 可能选出不同 "
                   "winner。按 README，只有 record 数据支撑的结论按未验证处理")
    return out


def caliber_hazards(runs: List[Tuple[str, Path]]) -> List[str]:
    """Not mismatches — shared state that makes a matched caliber less solid.

    Warned about rather than refused, because refusing would reject the
    default configuration every A/B in this repo currently uses.
    """
    cals = [(n, c) for n, c in ((n, caliber_of(d)) for n, d in runs) if c]
    recording = [n for n, c in cals if c["autotune_mode"].startswith("record")]
    shared = [n for n, c in cals if c["tuner_cache"].startswith("shared")]
    if recording and shared:
        # 一次 record 跑在共享 libtuner cache 上，sweep 的起点是别的运行留下的
        # 结果。曾把 mm 的 945 个变体读成 33 个，误报为编译期收益。A/B 两侧都
        # replay 同一份 config 时结论仍成立，绝对值不成立。
        return [f"{', '.join(recording)} 在共享 tuner cache 上 record："
                "被 pin 的 config 取决于此前跑过什么。相对结论不受影响，"
                "绝对延迟/变体数不可跨机器引用（ZL_ISOLATE_TUNER_CACHE=1 可消除）"]
    return []


class CaliberMismatch(SystemExit):
    """Strict-mode refusal. Exit 4 so a driver script can tell it from a crash."""

    def __init__(self) -> None:
        super().__init__(4)


def _report_caliber(runs: List[Tuple[str, Path]], label: str, strict: bool) -> bool:
    """Print caliber problems. Returns False when strict mode should stop."""
    for h in caliber_hazards(runs):
        print(f"警告 [caliber] {label}: {h}")
    problems = caliber_complaints(runs)
    if not problems:
        return True
    head = "!!!" if strict else "警告"
    for p in problems:
        print(f"{head} [caliber] {label}: {p}")
    if strict:
        print(f"!!! [caliber] {label}: 拒绝出报告。确认口径确实可比后加 "
              "--allow-caliber-mismatch")
        return False
    return True


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 = "",
                     strict: bool = True) -> List[Dict[str, str]]:
    """Print the multi-round A/B: per-shape medians plus both dispersions."""
    name = label or root.name
    cells = sorted(d for d in root.glob("r*_[AB]") if (d / "run.log").is_file())
    if cells and not _report_caliber([(d.name, d) for d in cells], name, strict):
        raise CaliberMismatch()
    rows, per_round, n = ab_round_rows(root)
    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 报告里的算子名")
    p.add_argument("--caliber", type=Path, default=None, metavar="RUN_DIR",
                   help="打印该 run 的口径字段（TSV 一行），供筛选表复用")
    p.add_argument("--allow-caliber-mismatch", action="store_true",
                   help="口径不一致时仍出报告（默认拒绝）。只在确认两次运行"
                        "确实可比时使用")
    args = p.parse_args()
    if args.caliber is not None:
        # 口径的定义只有 CALIBER_FIELDS 一处；筛选表等复用它，不要在别处重列
        cal = caliber_of(args.caliber) or {}
        print("	".join(str(cal.get(k, "")) for k in CALIBER_FIELDS))
        return
    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,
                         strict=not args.allow_caliber_mismatch)
        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,
                  strict=not args.allow_caliber_mismatch)
        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()