#!/usr/bin/env python3
"""
bench.py — the measurement harness you use in every project of this journey.

Why this file exists
--------------------
Every project in this track has an experiment section, and every experiment
section demands the same four things:

  1. A *distribution* of latencies, not a mean. A mean hides the tail, and the
     tail is where systems fail. p99 is the number your users feel.
  2. Warmup separated from measurement. First iterations pay for page faults,
     JIT/interpreter caches, cold branch predictors and cold file cache. If you
     average them in, you are measuring startup, not steady state.
  3. An uncertainty estimate. "A is 4% faster than B" is meaningless unless you
     know the run-to-run spread. We bootstrap a confidence interval on the
     median so you can say "no measurable difference" honestly.
  4. Recorded environment. A number without its machine, CPU governor, Python
     version and load state is not reproducible.

Usage as a library
------------------
    from bench import measure, report, compare

    r = measure(lambda: my_index.search(q, k=10), n=2000, warmup=200)
    report("hnsw ef=64", r)

    compare("brute", measure(brute, n=200),
            "hnsw",  measure(hnsw,  n=200))

Usage as a CLI (self-demo; also a smoke test that the harness works)
--------------------------------------------------------------------
    python3 bench.py demo

No third-party dependencies. `numpy` is used only by the demo, and only if
importable.
"""

from __future__ import annotations

import json
import math
import os
import platform
import statistics
import sys
import time
from dataclasses import dataclass, field, asdict
from typing import Callable, Sequence

# ---------------------------------------------------------------------------
# Clock choice
# ---------------------------------------------------------------------------
# perf_counter_ns is monotonic, has nanosecond resolution, and is not affected
# by NTP steps or wall-clock adjustment. Never use time.time() for latency:
# it can move backwards. Never use time.process_time() for a system benchmark:
# it excludes the time you blocked on I/O, which is usually the thing you care
# about.
CLOCK = time.perf_counter_ns


@dataclass
class Result:
    """A latency sample plus the environment it was taken in."""

    label: str
    samples_ns: list[int] = field(default_factory=list)
    env: dict = field(default_factory=dict)

    # -- summary statistics -------------------------------------------------
    @property
    def n(self) -> int:
        return len(self.samples_ns)

    def pct(self, q: float) -> float:
        """Nearest-rank percentile in milliseconds.

        Nearest-rank (not interpolated) is the right choice for latency: it
        returns a value that actually occurred, so p99 is always a real
        observation. With n samples the p-th percentile is the ceil(p*n)-th
        smallest, 1-indexed.
        """
        if not self.samples_ns:
            raise ValueError("no samples")
        s = sorted(self.samples_ns)
        rank = max(1, math.ceil(q / 100.0 * len(s)))
        return s[rank - 1] / 1e6

    @property
    def mean_ms(self) -> float:
        return statistics.fmean(self.samples_ns) / 1e6

    @property
    def stdev_ms(self) -> float:
        if self.n < 2:
            return 0.0
        return statistics.stdev(self.samples_ns) / 1e6

    @property
    def ops_per_sec(self) -> float:
        """Throughput implied by the *median*, not the mean.

        This is single-threaded, one-op-at-a-time throughput. It is NOT the
        throughput of a concurrent server: with C concurrent workers you get
        somewhere between 1x and Cx this number depending on contention, and
        finding out where is the whole point of a scaling experiment.
        """
        med = self.pct(50)
        return 0.0 if med == 0 else 1000.0 / med

    def ci_median_ms(self, iters: int = 2000, conf: float = 95.0,
                     seed: int = 0) -> tuple[float, float]:
        """Bootstrap confidence interval on the median.

        We resample the observed latencies with replacement `iters` times,
        take the median of each resample, and report the empirical quantiles
        of those medians. This makes no normality assumption -- which matters,
        because latency distributions are right-skewed and multi-modal, so a
        t-interval on the mean would be wrong.
        """
        import random

        rng = random.Random(seed)
        s = self.samples_ns
        if len(s) < 2:
            return (self.pct(50), self.pct(50))
        meds = []
        for _ in range(iters):
            resample = [s[rng.randrange(len(s))] for _ in range(len(s))]
            meds.append(statistics.median(resample))
        meds.sort()
        lo_q = (100.0 - conf) / 2.0
        hi_q = 100.0 - lo_q
        lo = meds[max(0, math.ceil(lo_q / 100 * len(meds)) - 1)]
        hi = meds[max(0, math.ceil(hi_q / 100 * len(meds)) - 1)]
        return (lo / 1e6, hi / 1e6)


def environment() -> dict:
    """Everything needed to know whether a number is comparable to another."""
    return {
        "python": sys.version.split()[0],
        "impl": platform.python_implementation(),
        "platform": platform.platform(),
        "machine": platform.machine(),
        "processor": platform.processor() or "unknown",
        "cpu_count": os.cpu_count(),
        # loadavg tells you whether something else was competing for the CPU.
        # A benchmark taken at load 8 on an 8-core box is not a benchmark.
        "loadavg_1m": round(os.getloadavg()[0], 2) if hasattr(os, "getloadavg") else None,
    }


def measure(fn: Callable[[], object], n: int = 1000, warmup: int = 100,
            label: str = "") -> Result:
    """Run `fn` `warmup` times untimed, then `n` times timed.

    One timed call per sample. Do NOT batch k calls into one sample and divide:
    that averages away the tail you are trying to see. If a single call is too
    fast to time (< ~1 microsecond, i.e. within a few clock reads), make the
    unit of work bigger instead -- benchmark a batch of 1000 lookups as one
    operation and report per-batch latency, saying so explicitly.
    """
    for _ in range(warmup):
        fn()
    samples: list[int] = []
    for _ in range(n):
        t0 = CLOCK()
        fn()
        t1 = CLOCK()
        samples.append(t1 - t0)
    return Result(label=label, samples_ns=samples, env=environment())


def report(label: str, r: Result, out=None) -> None:
    # out=None rather than out=sys.stdout: a default argument is evaluated once at
    # import time, so `out=sys.stdout` would capture whatever stdout was THEN and
    # ignore any later redirection (pytest's capsys, contextlib.redirect_stdout, a
    # logging harness). Late binding is the fix. Found by writing tests/test_tools.py.
    out = sys.stdout if out is None else out
    r.label = label or r.label
    lo, hi = r.ci_median_ms()
    print(
        f"{r.label:<28} n={r.n:<6} "
        f"p50={r.pct(50):9.4f}ms  p95={r.pct(95):9.4f}ms  "
        f"p99={r.pct(99):9.4f}ms  max={r.pct(100):9.4f}ms  "
        f"mean={r.mean_ms:9.4f}ms  {r.ops_per_sec:10.1f} op/s  "
        f"[median 95% CI {lo:.4f}–{hi:.4f}]",
        file=out,
    )


def compare(label_a: str, a: Result, label_b: str, b: Result,
            out=None) -> None:
    """Report the ratio of medians with an honest 'no difference' verdict.

    If the bootstrap CIs on the two medians overlap, you have not measured a
    difference. Say that. Reporting "3% faster" from overlapping intervals is
    the single most common benchmarking lie.
    """
    report(label_a, a, out)
    report(label_b, b, out)
    ma, mb = a.pct(50), b.pct(50)
    alo, ahi = a.ci_median_ms()
    blo, bhi = b.ci_median_ms()
    overlap = not (ahi < blo or bhi < alo)
    ratio = mb / ma if ma else float("inf")
    verdict = (
        "OVERLAPPING CIs — no measurable difference"
        if overlap
        else f"{label_b} is {ratio:.2f}x the median latency of {label_a}"
    )
    print(f"  -> {verdict}", file=out)


def save(path: str, results: Sequence[Result]) -> None:
    """Persist raw samples. Never keep only the summary.

    You will want to re-plot, re-percentile, or check for bimodality months
    later, and you cannot recover a distribution from a p50.
    """
    with open(path, "w") as f:
        json.dump([asdict(r) for r in results], f)


# ---------------------------------------------------------------------------
# demo / smoke test
# ---------------------------------------------------------------------------
def _demo() -> int:
    print("environment:", json.dumps(environment(), indent=2))
    print()

    # 1. Two functions that differ enormously: pure-Python triple-loop matmul
    #    vs. the same arithmetic in a BLAS call. This is the arithmetic-
    #    intensity lesson of Project 14 in miniature.
    N = 64
    A = [[(i * 31 + j) % 7 - 3 for j in range(N)] for i in range(N)]
    B = [[(i * 17 + j) % 5 - 2 for j in range(N)] for i in range(N)]

    def naive():
        C = [[0] * N for _ in range(N)]
        for i in range(N):
            Ai = A[i]
            Ci = C[i]
            for k in range(N):
                aik = Ai[k]
                if aik == 0:
                    continue
                Bk = B[k]
                for j in range(N):
                    Ci[j] += aik * Bk[j]
        return C

    r_naive = measure(naive, n=30, warmup=3)
    report(f"python matmul {N}x{N}", r_naive)

    try:
        import numpy as np

        An = np.array(A, dtype=np.float32)
        Bn = np.array(B, dtype=np.float32)
        r_np = measure(lambda: An @ Bn, n=2000, warmup=200)
        report(f"numpy  matmul {N}x{N}", r_np)
        flops = 2 * N ** 3
        print(f"\n  {N}x{N} matmul = 2*N^3 = {flops:,} FLOPs")
        print(f"  python: {flops / (r_naive.pct(50) / 1e3) / 1e6:10.2f} MFLOP/s")
        print(f"  numpy : {flops / (r_np.pct(50) / 1e3) / 1e9:10.2f} GFLOP/s")
        print(f"  ratio : {r_naive.pct(50) / r_np.pct(50):10.1f}x")
    except ImportError:
        print("(numpy not available; skipping BLAS comparison)")

    # 2. Two functions that are actually the same. The harness must be able to
    #    say "no difference" -- if it cannot, it will manufacture wins for you.
    print()
    data = list(range(5000))
    compare(
        "sum(list) A", measure(lambda: sum(data), n=1500, warmup=200),
        "sum(list) B", measure(lambda: sum(data), n=1500, warmup=200),
    )
    return 0


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "demo":
        raise SystemExit(_demo())
    print(__doc__)