"""Benchmark learners against standard compressors on three 1 MB streams.

  search    Python source the search optimised on
  held-out  Python source never used in the search
  english   English prose (Python's changelog): a kind of text the search never saw

  python benchmark.py --run runs/first
  python benchmark.py --run runs/first --out runs/first/benchmark.md

All numbers are bits per byte; lower is better. For the learners, the cost of their own
source code is included; the standard compressors are not charged for their program.
"""
from __future__ import annotations

import argparse
import bz2
import lzma
import time
import zlib
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from compsearch.evaluate import evaluate
from compsearch.search import load_results, ranked
from probe import read_streams

STREAMS = {"search": "data/search.bin", "held-out": "data/holdout.bin", "english": "data/english.bin"}


def standard_compressors() -> dict:
    tools = {
        "zlib -9 (LZ77, as in zip/gzip)": lambda data: zlib.compress(data, 9),
        "bz2 -9 (block sorting)": lambda data: bz2.compress(data, 9),
        "xz -9e (LZMA)": lambda data: lzma.compress(data, preset=9 | lzma.PRESET_EXTREME),
    }
    try:
        import brotli
        tools["brotli -11 (LZ77 + context modelling)"] = lambda data: brotli.compress(data, quality=11, lgwin=24)
    except ImportError:
        pass
    return tools


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--run", default="runs/first")
    parser.add_argument("--seconds", type=float, default=600.0)
    parser.add_argument("--mem-mb", type=int, default=1024)
    parser.add_argument("--out", help="also write the table to this Markdown file")
    args = parser.parse_args()

    read_streams()  # makes data/english.bin if it is missing
    data = {name: Path(path).read_bytes() for name, path in STREAMS.items()}
    rows = []

    for name, compress in standard_compressors().items():
        cells, started = [], time.monotonic()
        for stream in data.values():
            cells.append(8 * len(compress(stream)) / len(stream))
        speed = sum(len(s) for s in data.values()) / 1024 / (time.monotonic() - started)
        rows.append((name, cells, speed))

    run_dir = Path(args.run)
    best = ranked(load_results(run_dir))[0]
    learners = {
        "order-0 seed (byte frequencies)": Path("seeds/order0.py"),
        "sponge seed (hand-written start)": Path("seeds/sponge.py"),
        f"best searched learner (#{best['id']})": run_dir / "cands" / f"{best['id']:05d}" / "learner.py",
    }
    for name, path in learners.items():
        with ThreadPoolExecutor(max_workers=len(STREAMS)) as pool:
            results = list(pool.map(lambda p: evaluate(path, p, args.seconds, args.mem_mb), STREAMS.values()))
        cells = [r["bpb"] if r.get("status") == "ok" else float("nan") for r in results]
        seconds = [r["seconds"] for r in results if r.get("status") == "ok"]
        speed = sum(len(s) for s in data.values()) / 1024 / sum(seconds) if seconds else float("nan")
        rows.append((name, cells, speed))

    lines = ["| Model | " + " | ".join(STREAMS) + " | Speed (KB/s) |", "|---|" + "---|" * (len(STREAMS) + 1)]
    for name, cells, speed in rows:
        lines.append(f"| {name} | " + " | ".join(f"{cell:.3f}" for cell in cells) + f" | {speed:,.0f} |")
    table = "\n".join(lines)
    print(table)
    if args.out:
        Path(args.out).write_text(table + "\n", encoding="utf-8")
    return 0


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