"""Route boards under several connection orders and compare variants fairly.

One routing run says little: which connection routes first moves a result by
ten connections either way on a crowded board. This runs every board with
every variant under the same seeds (``--seed``), so variant A and variant B
face the same order on seed 3, and compares them seed by seed.

    uv run python bench/bench.py BOARDS -v base= -v hist1="--history-cost 1" \\
        --seeds 6 --jobs 4 --out bench/results/hist.jsonl

BOARDS is a text file, one board per line: a name, a DSN path and any flags
that board always needs, e.g.

    sense  /path/to/charge-sense.dsn  --edge-clearance 0.5
    # comments and blank lines are ignored

A variant is ``NAME=FLAGS``, or ``NAME@SRC=FLAGS`` to run the code under SRC
(a worktree's ``src`` directory) rather than the installed package. The first
variant is the baseline. Results are appended to --out as JSON lines, and runs
already in it are skipped, so an interrupted benchmark picks up where it left
off. Exits non-zero if any run crashed or timed out.
"""

from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import json
import math
import os
from pathlib import Path
import re
import shlex
import subprocess
import sys
import tempfile
import threading
import time

NOT_ROUTED = re.compile(r"rpmroute: (\d+) connection\(s\) not routed")
TO_ZONE = re.compile(r"rpmroute: (\d+) more connection\(s\) left to a zone fill")


def parse_boards(path: Path) -> list[tuple[str, Path, list[str]]]:
    boards = []
    for line in path.read_text().splitlines():
        words = shlex.split(line, comments=True)
        if not words:
            continue
        if len(words) < 2:
            raise SystemExit(f"{path}: want 'name path [flags]', got {line!r}")
        dsn = Path(words[1]).expanduser()
        if not dsn.is_file():
            raise SystemExit(f"{path}: no such DSN {dsn}")
        boards.append((words[0], dsn, words[2:]))
    return boards


def parse_variant(spec: str) -> tuple[str, str | None, list[str]]:
    head, sep, flags = spec.partition("=")
    if not sep or not head:
        raise SystemExit(f"variant {spec!r}: want NAME=FLAGS or NAME@SRC=FLAGS")
    name, _, src = head.partition("@")
    if src and not (Path(src) / "rpmroute").is_dir():
        raise SystemExit(f"variant {name}: {src} has no rpmroute package")
    return name, src or None, shlex.split(flags)


def run_one(board, variant, seed: int, timeout: float) -> dict:
    name, dsn, board_flags = board
    vname, src, flags = variant
    env = dict(os.environ)
    if src:
        env["PYTHONPATH"] = src + os.pathsep + env.get("PYTHONPATH", "")
    with tempfile.TemporaryDirectory() as tmp:
        cmd = [
            sys.executable, "-m", "rpmroute.cli", str(dsn), "-do", f"{tmp}/out.ses",
            *board_flags, *flags, "--seed", str(seed),
        ]
        start = time.monotonic()
        proc = subprocess.Popen(
            cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True
        )
        expired = threading.Event()

        def expire() -> None:
            expired.set()
            proc.kill()

        timer = threading.Timer(timeout, expire)
        timer.start()
        err = proc.stderr.read()
        # reap it here rather than through Popen, to keep its peak memory
        _, wstatus, usage = os.wait4(proc.pid, 0)
        timer.cancel()
        proc.returncode = os.waitstatus_to_exitcode(wstatus)
        proc.stderr.close()
        status = "timeout" if expired.is_set() else (
            "ok" if proc.returncode == 0 else f"exit {proc.returncode}"
        )
        rss_mb = round(usage.ru_maxrss / 1024)
        secs = round(time.monotonic() - start, 1)
        try:
            with open(f"{tmp}/out.ses", encoding="utf-8") as f:
                vias = f.read().count("(via ")
        except OSError:
            vias = None
    missing = NOT_ROUTED.search(err)
    zone = TO_ZONE.search(err)
    return {
        "board": name, "variant": vname, "seed": seed, "status": status,
        "unrouted": int(missing.group(1)) if missing else 0,
        "to_zone": int(zone.group(1)) if zone else 0,
        "secs": secs, "rss_mb": rss_mb, "vias": vias,
        "tail": err.strip().splitlines()[-1] if status != "ok" and err.strip() else "",
    }


def mean_sd(xs: list[float]) -> tuple[float, float]:
    m = sum(xs) / len(xs)
    sd = math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1)) if len(xs) > 1 else 0.0
    return m, sd


def report(rows: list[dict], boards, variants) -> None:
    ok = {(r["board"], r["variant"], r["seed"]): r for r in rows if r["status"] == "ok"}
    base = variants[0][0]
    print(f"\nunrouted per board (lower is better), baseline {base!r}\n")
    print(f"{'board':14} {'variant':14} {'n':>3} {'mean':>6} {'sd':>5} {'min':>4} {'max':>4}"
          f" {'secs':>6}   vs baseline, seed by seed")
    for bname, _, _ in boards:
        for vname, _, _ in variants:
            runs = sorted((k[2], r) for k, r in ok.items() if k[:2] == (bname, vname))
            if not runs:
                print(f"{bname:14} {vname:14}   no finished runs")
                continue
            xs = [r["unrouted"] for _, r in runs]
            m, sd = mean_sd(xs)
            t, _ = mean_sd([r["secs"] for _, r in runs])
            line = (f"{bname:14} {vname:14} {len(xs):>3} {m:>6.1f} {sd:>5.1f} {min(xs):>4}"
                    f" {max(xs):>4} {t:>6.0f}")
            if vname != base:
                diffs = [r["unrouted"] - ok[(bname, base, s)]["unrouted"]
                         for s, r in runs if (bname, base, s) in ok]
                if diffs:
                    dm, dsd = mean_sd(diffs)
                    se = dsd / math.sqrt(len(diffs)) if len(diffs) > 1 else float("inf")
                    better = sum(d < 0 for d in diffs)
                    worse = sum(d > 0 for d in diffs)
                    # a difference inside two standard errors is not a finding
                    verdict = ("better" if dm < 0 else "worse") if abs(dm) > 2 * se else "no clear change"
                    line += (f"   {dm:+.1f} ± {se:.1f}  ({better} better, {worse} worse,"
                             f" {len(diffs) - better - worse} same)  {verdict}")
            print(line)
    # vias, where the runs recorded them (older result files did not)
    print("\nvias per variant, summed over boards (mean over seeds per board)\n")
    for vname, _, _ in variants:
        total, boards_counted = 0.0, 0
        for bname, _, _ in boards:
            vs = [r["vias"] for k, r in ok.items() if k[:2] == (bname, vname) and r.get("vias") is not None]
            if vs:
                total += sum(vs) / len(vs)
                boards_counted += 1
        print(f"{vname:14} {total:8.1f} vias over {boards_counted} boards")
    failed = [r for r in rows if r["status"] != "ok"]
    for r in failed:
        print(f"FAILED {r['board']} {r['variant']} seed {r['seed']}: {r['status']} {r['tail']}")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("boards", type=Path, help="board list: name, DSN path, flags")
    ap.add_argument("-v", "--variant", action="append", required=True,
                    help="NAME=FLAGS or NAME@SRC=FLAGS; the first is the baseline")
    ap.add_argument("--seeds", type=int, default=6, help="orders per board and variant")
    ap.add_argument("--first-seed", type=int, default=0)
    ap.add_argument("--jobs", type=int, default=4, help="runs at once")
    ap.add_argument("--timeout", type=float, default=1800, help="seconds per run")
    ap.add_argument("--out", type=Path, required=True, help="JSON lines, appended")
    ap.add_argument("--report-only", action="store_true", help="summarise --out, run nothing")
    args = ap.parse_args(argv)

    boards = parse_boards(args.boards)
    variants = [parse_variant(v) for v in args.variant]
    if len({v[0] for v in variants}) != len(variants):
        raise SystemExit("variant names must differ")
    seeds = range(args.first_seed, args.first_seed + args.seeds)

    rows: list[dict] = []
    if args.out.exists():
        rows = [json.loads(line) for line in args.out.read_text().splitlines() if line.strip()]
    done = {(r["board"], r["variant"], r["seed"]) for r in rows if r["status"] == "ok"}
    todo = [(b, v, s) for b in boards for v in variants for s in seeds
            if (b[0], v[0], s) not in done]

    if not args.report_only and todo:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        print(f"{len(todo)} runs, {args.jobs} at a time", file=sys.stderr)
        # retried runs replace their earlier failure
        rows = [r for r in rows if (r["board"], r["variant"], r["seed"]) in done]
        with ThreadPoolExecutor(args.jobs) as pool, args.out.open("w") as out:
            for r in rows:
                out.write(json.dumps(r) + "\n")
            futures = [pool.submit(run_one, b, v, s, args.timeout) for b, v, s in todo]
            for n, fut in enumerate(as_completed(futures), 1):
                r = fut.result()
                rows.append(r)
                out.write(json.dumps(r) + "\n")
                out.flush()
                print(f"[{n}/{len(todo)}] {r['board']} {r['variant']} seed {r['seed']}:"
                      f" {r['status']} unrouted {r['unrouted']} ({r['secs']}s,"
                      f" {r['rss_mb']} MB)", file=sys.stderr)

    wanted = {(b[0], v[0], s) for b in boards for v in variants for s in seeds}
    rows = [r for r in rows if (r["board"], r["variant"], r["seed"]) in wanted]
    report(rows, boards, variants)
    return 1 if any(r["status"] != "ok" for r in rows) or len(rows) < len(wanted) else 0


if __name__ == "__main__":
    raise SystemExit(main())