#!/usr/bin/env python3
"""Check the PERF_PLAN Phase-0 acceptance criterion.

Runs bench_pipeline several times and reports, per row, how far the passes
spread around their median. The plan's gate is: the same configuration run 7x
lands within +/-5% on every row. Until that holds, no measured difference
between KPN and TBB is worth acting on.

Exits non-zero if any row exceeds the tolerance, so it can gate a session of
performance work rather than merely inform one.

A full sweep is hours, so the run is observable and restartable rather than
opaque: rows stream to --out-dir as each pass produces them, and a progress
bar tracks rows within the pass. Killing the run keeps everything already
written; --resume picks up from the completed passes on disk.

Usage:
    scripts/bench_repro_check.py ./build_bench/benchmarks/bench_pipeline \\
        --passes 7 --tolerance 5 -- --work=10 --topos=chain,wide --reps=5
"""

import argparse
import datetime
import os
import pathlib
import statistics
import subprocess
import sys

KEY_COLS = ("topology", "size", "work_us", "threads")

# bench_pipeline reports throughput, bench_dispatch reports per-dispatch cost.
# Either is a valid thing to demand reproducibility of; deviation from the
# median is symmetric, so it does not matter which direction is "better".
METRIC_COLS = ("items_per_sec", "ns_per_dispatch")


def _progress(total, desc):
    """A tqdm bar if tqdm is installed, else a minimal stderr fallback.

    The fallback exists because this script gates a benchmark run; refusing to
    start over a missing progress dependency would be the wrong trade.
    """
    try:
        from tqdm import tqdm
    except ImportError:
        class Fallback:
            def __init__(self):
                self.n = 0

            def update(self, k=1):
                self.n += k
                end = "\n" if (total and self.n >= total) else "\r"
                print(f"  {desc}: {self.n}/{total or '?'} rows",
                      file=sys.stderr, end=end, flush=True)

            def close(self):
                pass

            def __enter__(self):
                return self

            def __exit__(self, *exc):
                self.close()

        return Fallback()

    return tqdm(total=total, desc=desc, unit="row", leave=False,
                bar_format="  {desc}: {n_fmt}/{total_fmt} rows "
                           "|{bar}| {elapsed}<{remaining}",
                file=sys.stderr)


def parse_csv_lines(lines, metric=None):
    """Return ({(topology, size, work_us, threads): value}, metric_name)."""
    rows = {}
    header = None
    for line in lines:
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        fields = line.split(",")
        if header is None:
            if fields[0] != "topology":
                continue
            header = fields
            if metric is None:
                for cand in METRIC_COLS:
                    if cand in header:
                        metric = cand
                        break
                else:
                    sys.exit(f"no metric column found in header: {header}")
            elif metric not in header:
                sys.exit(f"metric {metric!r} not in header: {header}")
            continue
        rec = dict(zip(header, fields))
        try:
            key = tuple(rec[c] for c in KEY_COLS)
            rows[key] = float(rec[metric])
        except (KeyError, ValueError):
            continue
    return rows, metric


def parse_csv(text, metric=None):
    return parse_csv_lines(text.splitlines(), metric)


def count_rows(binary, extra):
    """Enumerate the sweep cheaply, so the progress bar has a real total.

    Asks the binary itself rather than reimplementing the sweep in Python,
    which would silently drift from the C++ defaults. Returns None if the
    probe fails -- an unknown total degrades the bar, it does not stop the run.
    """
    probe = [binary] + extra + ["--reps=0", "--warmup=0"]
    try:
        proc = subprocess.run(probe, capture_output=True, text=True,
                              timeout=600)
    except (subprocess.SubprocessError, OSError):
        return None
    if proc.returncode != 0:
        return None
    rows, _ = parse_csv(proc.stdout)
    return len(rows) or None


def run_pass(binary, extra, total, desc, sink, metric):
    """Run one pass, streaming rows to `sink` and the bar as they arrive.

    capture_output would withhold every row until the pass ended, which for a
    multi-hour sweep means no way to tell a slow run from a wedged one.
    """
    lines = []
    bar = _progress(total, desc)
    proc = subprocess.Popen([binary] + extra, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE, text=True, bufsize=1)
    try:
        for line in proc.stdout:
            lines.append(line)
            if sink:
                sink.write(line)
                sink.flush()          # a killed run keeps its rows
            stripped = line.strip()
            if (stripped and not stripped.startswith("#")
                    and "," in stripped
                    and not stripped.startswith("topology,")):
                bar.update(1)
    finally:
        bar.close()
        proc.stdout.close()
        stderr = proc.stderr.read()
        proc.stderr.close()
        rc = proc.wait()

    if rc != 0:
        print(stderr, file=sys.stderr)
        sys.exit(f"{binary} failed with {rc}")
    return parse_csv_lines(lines, metric)


def report(passes, metric, tolerance, npasses):
    keys = set(passes[0])
    for p in passes[1:]:
        keys &= set(p)
    if not keys:
        sys.exit("no rows common to every pass")

    print(f"\n{'row':<34} {'median ' + metric:>20} {'worst dev':>10}  verdict")
    print("-" * 72)

    failures = 0
    for key in sorted(keys):
        values = [p[key] for p in passes]
        med = statistics.median(values)
        worst = max(abs(v - med) / med * 100 for v in values) if med else 0.0
        ok = worst <= tolerance
        failures += not ok
        label = "{}-{} w={} s={}".format(*key)
        print(f"{label:<34} {med:>20.1f} {worst:>9.1f}%  {'ok' if ok else 'NOISY'}")

    print("-" * 72)
    if failures:
        print(f"{failures}/{len(keys)} rows exceed +/-{tolerance:g}% — "
              f"the Phase-0 gate is not met.")
        return 1
    print(f"all {len(keys)} rows within +/-{tolerance:g}% "
          f"over {npasses} passes — Phase-0 gate met.")
    return 0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("binary", help="path to bench_pipeline")
    ap.add_argument("--passes", type=int, default=7)
    ap.add_argument("--tolerance", type=float, default=5.0,
                    help="max allowed deviation from the median, percent")
    ap.add_argument("--metric", default=None, choices=METRIC_COLS,
                    help="column to check (default: whichever the CSV carries)")
    ap.add_argument("--out-dir", default=None,
                    help="write pass-NN.csv as rows arrive "
                         "(default: bench_runs/<timestamp>)")
    ap.add_argument("--resume", action="store_true",
                    help="reuse complete pass-NN.csv files in --out-dir")

    # Everything after a standalone `--` goes to bench_pipeline verbatim.
    # argparse.REMAINDER would swallow this script's own flags instead.
    argv = sys.argv[1:]
    extra = []
    if "--" in argv:
        cut = argv.index("--")
        argv, extra = argv[:cut], argv[cut + 1:]
    args = ap.parse_args(argv)

    out_dir = args.out_dir
    if out_dir is None:
        if args.resume:
            sys.exit("--resume needs an explicit --out-dir")
        stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        out_dir = os.path.join("bench_runs", stamp)
    out = pathlib.Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    total = count_rows(args.binary, extra)
    print(f"writing to {out}/", file=sys.stderr)
    if total:
        print(f"{total} rows per pass, {args.passes} passes", file=sys.stderr)

    passes = []
    metric = args.metric
    for i in range(args.passes):
        path = out / f"pass-{i + 1:02d}.csv"

        if args.resume and path.exists():
            rows, metric = parse_csv(path.read_text(), metric)
            # A partial file from a killed run must not be silently averaged
            # in as if it were a whole pass.
            if total and len(rows) < total:
                print(f"pass {i + 1}/{args.passes}: {path.name} has "
                      f"{len(rows)}/{total} rows — rerunning", file=sys.stderr)
            else:
                print(f"pass {i + 1}/{args.passes}: reusing {path.name} "
                      f"({len(rows)} rows)", file=sys.stderr)
                passes.append(rows)
                continue

        print(f"pass {i + 1}/{args.passes} ...", file=sys.stderr, flush=True)
        with open(path, "w") as sink:
            rows, metric = run_pass(args.binary, extra, total,
                                    f"pass {i + 1}/{args.passes}", sink, metric)
        passes.append(rows)

    return report(passes, metric, args.tolerance, args.passes)


if __name__ == "__main__":
    sys.exit(main())