#!/usr/bin/env python3
"""A/B two release builds of a benchmark and say whether the difference is real.

Every performance number in this repository used to come from a hand-rolled
shell loop, and the loops kept producing wrong answers - see the log in
`docs/PERFORMANCE.md`. This machine is shared, its load average moves by a
factor of eight inside an afternoon, and a benchmark that reads one figure from
one build and one from another is measuring the afternoon.

What this does instead:

1. Runs the two binaries **alternately**, swapping which goes first each pair,
   so neither build owns the quiet half of the run.
2. Reads two figures out of **each** run: the `--metric` under test, and a
   `--control` the change is not supposed to touch.
3. Normalises **within the run** - metric / control - so the machine state that
   both figures shared cancels before anything is compared across builds.
4. Compares the medians of those per-run ratios, with a bootstrap interval.

And it refuses to report at all when the control's own before/after ratio has
moved away from 1.0, because that means the machine moved under the measurement
in a way the normalisation did not absorb, and no number from that run is
readable. The refusal is the point of the tool.

    scripts/ab.py --before /tmp/perf-main --after /tmp/perf-branch --pairs 15 \
        --metric 'score 160 candidates' \
        --control 'parse one observation'

A selector is a regular expression. Capture a group and that group is the
number; capture nothing and it picks lines, and the last number on a matching
line is taken - which is where `crates/sds-core/examples/perf.rs` puts the
figure, so a bare row name works and `score 160 candidates` reads its `us`
rather than the 160 in its name. A selector matching two rows is an error.

`--json` keeps every run's whole output, so `--from-json` can ask a different
row of a measurement already taken rather than re-running it on a machine that
is no longer the machine it was measured on.

No third-party modules: this decides whether a result is real, so it runs
anywhere `python3` does.
"""

import argparse
import json
import os
import random
import re
import resource
import statistics
import subprocess
import sys
from dataclasses import asdict, dataclass, field

# Enough resamples that the interval is stable to the digit we print, cheap
# enough that it is not worth a flag.
BOOTSTRAP_RESAMPLES = 2000
# Fixed, so re-running the tool on the same samples prints the same interval.
BOOTSTRAP_SEED = 20260828

NUMBER = r"[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?"


class AbError(Exception):
    """Something about the measurement is wrong, not the thing measured."""


class Selector:
    """A regular expression that names one number in a run's output.

    With a capture group, that group is the number. Without one, the pattern
    picks lines and the **last** number on a matching line is taken - which is
    where `crates/sds-core/examples/perf.rs` puts it, so a bare row name works
    and `score 160 candidates` does not come back as 160.
    """

    def __init__(self, pattern: str):
        self.pattern = pattern
        try:
            self.regex = re.compile(pattern)
        except re.error as exc:
            raise AbError(f"selector {pattern!r} is not a regular expression: {exc}") from exc
        if self.regex.groups > 1:
            raise AbError(
                f"selector {pattern!r} captures {self.regex.groups} groups; "
                "it must capture one number"
            )
        self.explicit = self.regex.groups == 1

    def findall(self, text: str) -> list[str]:
        if self.explicit:
            return self.regex.findall(text)
        found = []
        for line in text.splitlines():
            if not self.regex.search(line):
                continue
            numbers = re.findall(NUMBER, line)
            if not numbers:
                raise AbError(
                    f"selector {self.pattern!r} matched a line with no number: {line.strip()!r}"
                )
            found.append(numbers[-1])
        return found


def compile_selector(pattern: str) -> Selector:
    return Selector(pattern)


def scrape(text: str, selector: Selector) -> float:
    """Pull the one number a selector names out of a run's output.

    Ambiguity is an error rather than a choice. A selector that matches two
    rows would otherwise silently measure whichever the benchmark printed
    first, and that is the failure this tool exists to stop.
    """
    found = selector.findall(text)
    if not found:
        raise AbError(f"selector {selector.pattern!r} matched nothing in the output")
    if len(found) > 1:
        raise AbError(
            f"selector {selector.pattern!r} matched {len(found)} times ({', '.join(found[:4])}); "
            "narrow it so it names one row"
        )
    try:
        return float(found[0])
    except ValueError as exc:
        raise AbError(f"selector {selector.pattern!r} captured {found[0]!r}, not a number") from exc


@dataclass
class Run:
    """One execution of one build."""

    build: str
    pair: int
    metric: float
    control: float
    cpu: float = 0.0
    # The run's whole output, kept so `--json` is an audit trail: a figure
    # somebody doubts can be re-scraped with `--from-json` instead of re-run on
    # a machine that is no longer the machine it was measured on.
    output: str = ""

    @property
    def ratio(self) -> float:
        return self.metric / self.control


@dataclass
class Summary:
    """What the two sets of runs say, and whether it may be believed."""

    pairs: int = 0
    before_metric: float = 0.0
    after_metric: float = 0.0
    before_control: float = 0.0
    after_control: float = 0.0
    before_ratio: float = 0.0
    after_ratio: float = 0.0
    # after/before of the normalised ratio: below 1.0 is faster.
    normalised: float = 0.0
    # after/before of the raw metric, ignoring the control. Printed so the
    # difference between the naive reading and the controlled one is visible.
    raw: float = 0.0
    # after/before of the control alone. This is the refusal.
    control_drift: float = 0.0
    # How far the two statistics disagree, in points of change. Two readings of
    # the same seven samples came out at 7% and 24%; the gap is the signal.
    divergence: float = 0.0
    # Robust dispersion of the control within each build, as a fraction.
    control_spread: float = 0.0
    interval: tuple[float, float] = (0.0, 0.0)
    readable: bool = True
    refusal: str | None = None
    warnings: list[str] = field(default_factory=list)


def spread(values: list[float]) -> float:
    """Interquartile range over the median: how much this sample wandered.

    Robust rather than a standard deviation, for the same reason the report is
    medians: one run that landed while somebody's build ran should not set the
    figure.
    """
    if len(values) < 2:
        return 0.0
    ordered = sorted(values)
    middle = statistics.median(ordered)
    if middle == 0:
        return 0.0
    lower = statistics.median(ordered[: len(ordered) // 2])
    upper = statistics.median(ordered[(len(ordered) + 1) // 2 :])
    return (upper - lower) / middle


def bootstrap_interval(
    before: list[float],
    after: list[float],
    resamples: int = BOOTSTRAP_RESAMPLES,
    seed: int = BOOTSTRAP_SEED,
) -> tuple[float, float]:
    """A 95% percentile interval on after-median / before-median.

    Resampled rather than assumed: these are ratios of timings on a loaded
    machine, and nothing about them is normal.
    """
    if len(before) < 2 or len(after) < 2:
        return (float("nan"), float("nan"))
    rng = random.Random(seed)
    draws = []
    for _ in range(resamples):
        b = statistics.median(rng.choices(before, k=len(before)))
        a = statistics.median(rng.choices(after, k=len(after)))
        if b == 0:
            continue
        draws.append(a / b)
    if not draws:
        return (float("nan"), float("nan"))
    draws.sort()
    lo = draws[int(0.025 * (len(draws) - 1))]
    hi = draws[int(0.975 * (len(draws) - 1))]
    return (lo, hi)


def summarise(
    runs: list[Run],
    tolerance: float = 0.10,
    min_pairs: int = 8,
    divergence: float = 0.05,
) -> Summary:
    """Reduce the runs to a verdict.

    `tolerance` is how far the control's own before/after ratio may sit from
    1.0 before the whole result is refused.
    """
    before = [r for r in runs if r.build == "before"]
    after = [r for r in runs if r.build == "after"]
    if not before or not after:
        raise AbError("need at least one run of each build")

    summary = Summary(pairs=min(len(before), len(after)))
    summary.before_metric = statistics.median([r.metric for r in before])
    summary.after_metric = statistics.median([r.metric for r in after])
    summary.before_control = statistics.median([r.control for r in before])
    summary.after_control = statistics.median([r.control for r in after])
    summary.before_ratio = statistics.median([r.ratio for r in before])
    summary.after_ratio = statistics.median([r.ratio for r in after])
    summary.normalised = summary.after_ratio / summary.before_ratio
    summary.raw = summary.after_metric / summary.before_metric
    summary.control_drift = summary.after_control / summary.before_control
    summary.control_spread = max(
        spread([r.control for r in before]), spread([r.control for r in after])
    )
    summary.interval = bootstrap_interval([r.ratio for r in before], [r.ratio for r in after])

    summary.divergence = abs(summary.raw - summary.normalised)

    drift = abs(summary.control_drift - 1.0)
    if drift > tolerance:
        summary.readable = False
        summary.refusal = (
            f"the control moved {summary.control_drift:.3f}x between builds "
            f"({drift * 100:.1f}% from 1.0, tolerance {tolerance * 100:.1f}%). "
            "The change does not touch the control, so this is the machine moving under "
            "the measurement. No figure from this run is readable."
        )
    if summary.pairs < min_pairs:
        summary.warnings.append(
            f"only {summary.pairs} pairs; {min_pairs} is the floor for a quotable figure"
        )
    if summary.control_spread > 0.25:
        summary.warnings.append(
            f"the control's own spread is {summary.control_spread * 100:.0f}% of its median; "
            "the machine was noisy even within a build"
        )
    if summary.divergence > divergence:
        summary.warnings.append(
            f"the two statistics disagree by {summary.divergence * 100:.1f} points "
            f"(raw {percent(summary.raw)}, normalised {percent(summary.normalised)}). "
            "Either the sample is too small or the machine moved; add pairs before "
            "quoting either figure."
        )
    return summary


def child_cpu() -> float:
    used = resource.getrusage(resource.RUSAGE_CHILDREN)
    return used.ru_utime + used.ru_stime


def execute(command: list[str], timeout: float | None) -> tuple[str, float]:
    """Run one build once and hand back its output and the CPU it took."""
    start = child_cpu()
    try:
        done = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except FileNotFoundError as exc:
        raise AbError(f"cannot run {command[0]!r}: {exc}") from exc
    except subprocess.TimeoutExpired as exc:
        raise AbError(f"{command[0]!r} did not finish inside {timeout}s") from exc
    if done.returncode != 0:
        raise AbError(f"{command[0]!r} exited {done.returncode}\n{done.stderr.strip()[-2000:]}")
    return done.stdout + done.stderr, child_cpu() - start


def collect(args, metric: Selector, control: Selector) -> list[Run]:
    """Run the two builds alternately and scrape both figures from each run."""
    builds = {"before": [args.before, *args.rest], "after": [args.after, *args.rest]}

    for _ in range(args.warmup):
        for command in builds.values():
            execute(command, args.timeout)

    runs: list[Run] = []
    for pair in range(args.pairs):
        # Swap which build goes first each pair. A monotone drift inside a pair
        # then lands on the two builds equally instead of always on the second.
        order = ("before", "after") if pair % 2 == 0 else ("after", "before")
        for build in order:
            output, cpu = execute(builds[build], args.timeout)
            runs.append(
                Run(
                    build=build,
                    pair=pair,
                    metric=scrape(output, metric),
                    control=scrape(output, control),
                    cpu=cpu,
                    output=output,
                )
            )
            if not args.quiet:
                last = runs[-1]
                print(
                    f"  pair {pair + 1:>2}/{args.pairs} {build:<6} "
                    f"metric {last.metric:>12.4g}  control {last.control:>12.4g}  "
                    f"ratio {last.ratio:>8.4f}",
                    file=sys.stderr,
                )
    return runs


def percent(ratio: float) -> str:
    """A ratio as the change a reader wants: negative is faster."""
    return f"{(ratio - 1.0) * 100:+.1f}%"


def report(args, runs: list[Run], summary: Summary, load_start, load_end) -> str:
    lines = []
    lines.append("")
    lines.append(f"before   {args.before}")
    lines.append(f"after    {args.after}")
    if args.from_json:
        lines.append(f"re-read  {args.from_json} (no binary was run)")
    lines.append(f"metric   {args.metric}")
    lines.append(f"control  {args.control}")
    lines.append(
        f"pairs    {summary.pairs}   load {load_start[0]:.2f} at start, {load_end[0]:.2f} at end"
    )
    lines.append("")
    lines.append("  pair  build   metric        control       metric/control   cpu s")
    for run in runs:
        lines.append(
            f"  {run.pair + 1:>4}  {run.build:<6}  {run.metric:>12.4g}  "
            f"{run.control:>12.4g}  {run.ratio:>14.4f}  {run.cpu:>6.2f}"
        )
    lines.append("")
    lines.append("  medians          before        after         change")
    lines.append(
        f"  metric      {summary.before_metric:>12.4g}  {summary.after_metric:>12.4g}  "
        f"{percent(summary.raw):>13}   (uncontrolled - do not quote this)"
    )
    lines.append(
        f"  control     {summary.before_control:>12.4g}  {summary.after_control:>12.4g}  "
        f"{percent(summary.control_drift):>13}   (should be ~0%)"
    )
    lines.append(
        f"  normalised  {summary.before_ratio:>12.4f}  {summary.after_ratio:>12.4f}  "
        f"{percent(summary.normalised):>13}"
    )
    lines.append("")
    lo, hi = summary.interval
    lines.append(f"  control spread within a build: {summary.control_spread * 100:.0f}%")
    if summary.readable:
        lines.append(
            f"  RESULT  {percent(summary.normalised)} normalised, "
            f"{percent(summary.raw)} raw  "
            f"(95% interval {percent(lo)} to {percent(hi)}, {summary.pairs} pairs, "
            f"load {load_start[0]:.2f}-{load_end[0]:.2f})"
        )
    for warning in summary.warnings:
        lines.append(f"  WARNING  {warning}")
    if not summary.readable:
        lines.append("")
        lines.append("  REFUSED TO REPORT")
        lines.append(f"  {summary.refusal}")
        lines.append("  Wait for the machine to quiet down and run it again.")
    lines.append("")
    return "\n".join(lines)


def parse_args(argv: list[str] | None = None):
    parser = argparse.ArgumentParser(
        prog="scripts/ab.py",
        description="A/B two benchmark binaries with a control, and refuse to lie.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Build the two binaries first, one per worktree:\n"
            "  cargo build -j 2 --release -p sds-core --example perf\n"
            "  cp target/release/examples/perf /tmp/perf-<branch>\n"
        ),
    )
    parser.add_argument("--before", help="the baseline binary")
    parser.add_argument("--after", help="the binary with the change in it")
    parser.add_argument("--metric", required=True, help="selector for the figure under test")
    parser.add_argument(
        "--control",
        required=True,
        help="selector for a figure the change does not touch",
    )
    parser.add_argument(
        "--pairs", type=int, default=15, help="alternating pairs to run (default 15)"
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=1,
        help="runs of each build to throw away first (default 1)",
    )
    parser.add_argument(
        "--tolerance",
        type=float,
        default=0.10,
        help="how far the control's before/after ratio may sit from 1.0 (default 0.10)",
    )
    parser.add_argument(
        "--max-divergence",
        type=float,
        default=0.05,
        help=(
            "how far the raw and normalised readings may disagree, in points "
            "of change, before it is called out (default 0.05)"
        ),
    )
    parser.add_argument(
        "--min-pairs",
        type=int,
        default=8,
        help="fewer pairs than this warns (default 8)",
    )
    parser.add_argument("--timeout", type=float, default=900.0, help="seconds per run")
    parser.add_argument("--json", help="also write the runs, their output and the summary here")
    parser.add_argument(
        "--from-json",
        help="re-scrape a saved run with these selectors instead of running anything",
    )
    parser.add_argument("--quiet", action="store_true", help="no per-run progress")
    parser.add_argument(
        "rest",
        nargs=argparse.REMAINDER,
        help="arguments after -- are passed to both binaries",
    )
    args = parser.parse_args(argv)
    if args.rest and args.rest[0] == "--":
        args.rest = args.rest[1:]
    if args.pairs < 1:
        parser.error("--pairs must be at least 1")
    if not args.from_json and not (args.before and args.after):
        parser.error("--before and --after are required unless --from-json is given")
    return args


def reanalyse(path: str, metric: Selector, control: Selector):
    """Re-scrape a saved run with different selectors.

    The runs are the expensive part and the machine they happened on is gone.
    Asking a different question of the same output is free, and it is the only
    honest way to answer "what did that build do to the other row".
    """
    with open(path) as handle:
        saved = json.load(handle)
    runs = []
    for entry in saved["runs"]:
        if not entry.get("output"):
            raise AbError(f"{path} was written before outputs were kept; re-run instead")
        runs.append(
            Run(
                build=entry["build"],
                pair=entry["pair"],
                metric=scrape(entry["output"], metric),
                control=scrape(entry["output"], control),
                cpu=entry.get("cpu", 0.0),
                output=entry["output"],
            )
        )
    load = saved.get("load", {})
    return (
        runs,
        tuple(load.get("start", (0.0, 0.0, 0.0))),
        tuple(load.get("end", (0.0, 0.0, 0.0))),
        (saved.get("before", "?"), saved.get("after", "?")),
    )


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    try:
        metric = compile_selector(args.metric)
        control = compile_selector(args.control)
        if args.from_json:
            runs, load_start, load_end, names = reanalyse(args.from_json, metric, control)
            args.before = args.before or names[0]
            args.after = args.after or names[1]
        else:
            load_start = os.getloadavg()
            runs = collect(args, metric, control)
            load_end = os.getloadavg()
        summary = summarise(
            runs,
            tolerance=args.tolerance,
            min_pairs=args.min_pairs,
            divergence=args.max_divergence,
        )
    except AbError as exc:
        print(f"ab: {exc}", file=sys.stderr)
        return 1

    print(report(args, runs, summary, load_start, load_end))
    if args.json:
        with open(args.json, "w") as handle:
            json.dump(
                {
                    "before": args.before,
                    "after": args.after,
                    "metric": args.metric,
                    "control": args.control,
                    "load": {"start": load_start, "end": load_end},
                    "runs": [asdict(r) for r in runs],
                    "summary": asdict(summary),
                },
                handle,
                indent=2,
            )
    return 0 if summary.readable else 2


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