#!/usr/bin/env python3
"""Benchmark compression.zstd decompression (CPython gh-155460).

Measures ``zstd.decompress()`` with pyperf, on two workloads:

* one frame of a fixed decompressed size, for a range of sizes;
* a chunk-oriented workload: 1000 frames whose decompressed sizes are drawn
  from a normal distribution in log2(size), centred on 2 MiB and spanning
  512 KiB to 8 MiB at +-3 sigma (the endpoints are geometrically symmetric
  around the target).  One iteration decompresses all 1000 frames.

The corpus is deterministic (fixed seeds) and cached on disk next to this
script, so every pyperf worker process decompresses byte-identical input:
a 16 MiB blob of random 16-byte words drawn from a 200-word dictionary,
which compresses about 9:1 at level 1.

Run it once per interpreter, then compare:

    ./python bench_zstd_decompress.py --level 3 -o main.json
    ../other/python bench_zstd_decompress.py --level 3 -o patched.json
    ./python bench_zstd_decompress.py --compare main.json patched.json

The last command prints a table of mean +- standard deviation for both
runs.  ``python -m pyperf compare_to main.json patched.json --table`` works
too and adds pyperf's significance test.

Requires pyperf (``pip install pyperf``) for both interpreters.
"""

import math
import os
import random
import struct
import sys

from compression import zstd

CACHE_DIR = os.path.dirname(os.path.abspath(__file__))
BLOB_SEED = 42
CHUNK_SEED = 1234
NCHUNKS = 1000
MIN_CHUNK, TARGET_CHUNK, MAX_CHUNK = 512 << 10, 2 << 20, 8 << 20

SIZES = {
    "64K": 1 << 16, "256K": 1 << 18, "1M": 1 << 20, "2M": 2 << 20,
    "3M": 3 << 20, "4M": 4 << 20, "5M": 5 << 20, "6M": 6 << 20,
    "7M": 7 << 20, "8M": 8 << 20, "9M": 9 << 20, "10M": 10 << 20,
    "11M": 11 << 20, "12M": 12 << 20, "13M": 13 << 20, "14M": 14 << 20,
    "15M": 15 << 20, "16M": 16 << 20, "32M": 32 << 20, "64M": 64 << 20,
    "128M": 128 << 20, "256M": 256 << 20,
}
DEFAULT_SIZES = "64K,256K,1M,2M,3M,4M,5M,6M,8M,12M,16M,32M,64M,128M,256M"


def get_blob():
    """16 MiB of moderately compressible data (cached on disk)."""
    path = os.path.join(CACHE_DIR, "zstd_bench_blob.bin")
    if not os.path.exists(path):
        rng = random.Random(BLOB_SEED)
        words = [bytes(rng.randrange(256) for _ in range(16))
                 for _ in range(200)]
        blob = b"".join(rng.choice(words) for _ in range(1 << 20))
        _write_atomic(path, blob)
    with open(path, "rb") as fp:
        return fp.read()


def _write_atomic(path, data):
    with open(path + ".tmp", "wb") as fp:
        fp.write(data)
    os.replace(path + ".tmp", path)


def get_frame(label, level):
    """One frame of the given decompressed size (cached on disk).

    Cached so that a high compression level is not paid once per pyperf
    worker process, and so that all workers use byte-identical input.
    """
    path = os.path.join(CACHE_DIR, f"zstd_bench_frame_{label}_l{level}.bin")
    if not os.path.exists(path):
        size = SIZES[label]
        blob = get_blob()
        data = (blob * (size // len(blob) + 1))[:size]
        frame = zstd.compress(data, level)
        assert zstd.decompress(frame) == data, "round-trip failed"
        del data, blob
        _write_atomic(path, frame)
    with open(path, "rb") as fp:
        return fp.read()


def chunk_sizes():
    """1000 chunk sizes: normal in log2(size), 512 KiB..8 MiB, mode 2 MiB."""
    lo, mu, hi = (math.log2(MIN_CHUNK), math.log2(TARGET_CHUNK),
                  math.log2(MAX_CHUNK))
    sigma = (hi - mu) / 3.0
    rng = random.Random(CHUNK_SEED)
    sizes, offsets = [], []
    while len(sizes) < NCHUNKS:
        x = rng.gauss(mu, sigma)
        if lo <= x <= hi:
            sizes.append(int(2 ** x))
            offsets.append(rng.randrange(1 << 24))
    return sizes, offsets


def get_chunk_frames(level):
    """The 1000 compressed chunks (cached on disk, one file per level)."""
    path = os.path.join(CACHE_DIR, f"zstd_bench_chunks_l{level}.bin")
    if not os.path.exists(path):
        pool = get_blob() * 2
        sizes, offsets = chunk_sizes()
        buf = []
        for size, off in zip(sizes, offsets):
            frame = zstd.compress(pool[off:off + size], level)
            buf.append(struct.pack("<I", len(frame)))
            buf.append(frame)
        _write_atomic(path, b"".join(buf))
    frames = []
    with open(path, "rb") as fp:
        while header := fp.read(4):
            frames.append(fp.read(struct.unpack("<I", header)[0]))
    return frames


def decompress_all(frames):
    for frame in frames:
        out = zstd.decompress(frame)
    del out


def human_time(seconds):
    for unit, scale in (("s", 1), ("ms", 1e-3), ("us", 1e-6), ("ns", 1e-9)):
        if seconds >= scale or unit == "ns":
            return f"{seconds / scale:.3g} {unit}"


def compare(path_a, path_b):
    """Print a human-readable mean +- stdev comparison of two pyperf runs."""
    import pyperf

    def load(path):
        suite = pyperf.BenchmarkSuite.load(path)
        return {bench.get_name(): bench for bench in suite}

    a, b = load(path_a), load(path_b)
    names = [n for n in a if n in b]
    if not names:
        sys.exit(f"no benchmark names in common between {path_a} and {path_b}")

    label_a, label_b = os.path.basename(path_a), os.path.basename(path_b)
    width = max(len(n) for n in names)
    col = max(len(label_a), len(label_b), 22)
    print(f"{'benchmark':<{width}}  {label_a:>{col}}  {label_b:>{col}}  "
          f"{'speedup':>9}")
    print("-" * (width + 2 * col + 17))
    for name in names:
        mean_a, sd_a = a[name].mean(), a[name].stdev()
        mean_b, sd_b = b[name].mean(), b[name].stdev()
        cell_a = f"{human_time(mean_a)} +- {human_time(sd_a)}"
        cell_b = f"{human_time(mean_b)} +- {human_time(sd_b)}"
        speedup = (mean_a - mean_b) / mean_a * 100.0
        # flag differences smaller than the combined spread as inconclusive
        flag = " ~" if abs(mean_a - mean_b) < (sd_a + sd_b) else ""
        print(f"{name:<{width}}  {cell_a:>{col}}  {cell_b:>{col}}  "
              f"{speedup:+8.1f}%{flag}")
    print("\n'~' marks differences smaller than the combined standard "
          "deviations.\nFor pyperf's significance test, use: "
          f"python -m pyperf compare_to {path_a} {path_b} --table")


def add_cmdline_args(cmd, args):
    # pyperf spawns worker processes; custom arguments must be forwarded
    # explicitly, otherwise the workers silently use the defaults.
    cmd.extend(("--level", str(args.level)))
    cmd.extend(("--sizes", args.sizes))
    if args.skip_chunks:
        cmd.append("--skip-chunks")


def main():
    if "--compare" in sys.argv:
        pos = sys.argv.index("--compare")
        try:
            return compare(sys.argv[pos + 1], sys.argv[pos + 2])
        except IndexError:
            sys.exit("--compare needs two pyperf JSON files")

    import pyperf

    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
    runner.argparser.add_argument(
        "--level", type=int, default=zstd.COMPRESSION_LEVEL_DEFAULT,
        help="zstd compression level used to build the corpus "
             f"(default: {zstd.COMPRESSION_LEVEL_DEFAULT}, the "
             "compression.zstd default)")
    runner.argparser.add_argument(
        "--sizes", default=DEFAULT_SIZES,
        help=f"comma-separated decompressed sizes (default: {DEFAULT_SIZES})")
    runner.argparser.add_argument(
        "--skip-chunks", action="store_true",
        help="skip the 1000-chunk mixed-size workload")
    runner.argparser.add_argument(
        "--compare", nargs=2, metavar=("A.json", "B.json"),
        help="print a mean +- stdev comparison of two result files and exit")
    args = runner.parse_args()

    runner.metadata["zstd_version"] = zstd.zstd_version
    runner.metadata["zstd_level"] = str(args.level)

    for label in args.sizes.split(","):
        label = label.strip()
        if label not in SIZES:
            sys.exit(f"unknown size {label!r}, pick from {', '.join(SIZES)}")
        frame = get_frame(label, args.level)
        runner.bench_func(f"decompress {label} (level {args.level})",
                          zstd.decompress, frame)

    if not args.skip_chunks:
        frames = get_chunk_frames(args.level)
        runner.bench_func(
            f"decompress {len(frames)} chunks (level {args.level})",
            decompress_all, frames)


if __name__ == "__main__":
    main()