#!/usr/bin/env python3
# MIKE — AI / ML dataset benchmark.  Reproduces the table on
# https://www.nohathacker.com/mike/  — every file is CRC-verified after a
# full compress + decompress round-trip.  No trust required: run it yourself.
#
#   python3 bench.py            # uses ./datasets and the `mike` on your PATH
#   MIKE=/path/to/mike python3 bench.py
#
# Competitors are optional — any that aren't installed are simply skipped.
import os, subprocess, time, zlib, tempfile, shutil, shutil as _sh

MIKE = os.environ.get("MIKE", "mike")
DS   = os.path.join(os.path.dirname(os.path.abspath(__file__)), "datasets")

FILES = [
    ("HTRU2 pulsar features (float32 tensor)", "htru2-features.f32"),
    ("California housing (numeric CSV)",       "california-housing.csv"),
    ("HTRU2 dataset (numeric CSV)",            "htru2.csv"),
    ("Wine quality (numeric CSV)",             "winequality-red.csv"),
    ("MNIST test images (uint8 tensor)",       "mnist-t10k-images.uint8"),
    ("bert-tiny weights (float32)",            "bert-tiny-weights.f32.bin"),
]

def have(x): return _sh.which(x) is not None

def crc(p):
    c = 0
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            c = zlib.crc32(b, c)
    return c & 0xffffffff

def run(cmd):
    t = time.perf_counter()
    r = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return time.perf_counter() - t, r.returncode

def bench_mike(src, wd):
    arc, out = os.path.join(wd, "a.mike"), os.path.join(wd, "a.out")
    ct, rc = run([MIKE, "c", src, arc]); csz = os.path.getsize(arc) if rc == 0 else -1
    dt, rd = run([MIKE, "d", arc, out])
    return csz, ct, dt, (rd == 0 and crc(src) == crc(out))

def bench_cmd(src, wd, comp, decomp, ext):
    arc, out = os.path.join(wd, "a." + ext), os.path.join(wd, "out.bin")
    ct, rc = run(comp(src, arc)); csz = os.path.getsize(arc) if rc == 0 and os.path.exists(arc) else -1
    dt, rd = run(decomp(arc, out))
    return csz, ct, dt, (rd == 0 and os.path.exists(out) and crc(src) == crc(out))

TOOLS = [("MIKE (SAFE)", None)]
if have("zip"):  TOOLS.append(("zip -9",   (lambda s,a:["zip","-9","-q","-j",a,s], lambda a,o:["sh","-c",f"unzip -p '{a}' > '{o}'"], "zip")))
if have("rar"):  TOOLS.append(("rar -m5",  (lambda s,a:["rar","a","-m5","-ep","-idq",a,s], lambda a,o:["sh","-c",f"unrar p -idq '{a}' > '{o}'"], "rar")))
if have("7z"):   TOOLS.append(("7z (LZMA2)",(lambda s,a:["7z","a","-mx=9","-bd",a,s], lambda a,o:["sh","-c",f"7z e -so '{a}' > '{o}'"], "7z")))
if have("xz"):   TOOLS.append(("xz -9e",   (lambda s,a:["sh","-c",f"xz -9e -c '{s}' > '{a}'"], lambda a,o:["sh","-c",f"xz -dc '{a}' > '{o}'"], "xz")))
if have("zstd"): TOOLS.append(("zstd -19", (lambda s,a:["sh","-c",f"zstd -19 -q -c '{s}' > '{a}'"], lambda a,o:["sh","-c",f"zstd -dq -c '{a}' > '{o}'"], "zst")))

wins = 0
for label, fn in FILES:
    p = os.path.join(DS, fn)
    if not os.path.exists(p): print(f"skip (missing) {fn}"); continue
    osz = os.path.getsize(p)
    print(f"\n### {label}  —  {osz:,} bytes ({osz/1e6:.2f} MB)")
    print(f"{'tool':<14}{'out bytes':>13}{'ratio':>8}{'saved%':>8}{'comp s':>9}{'decomp s':>10}  crc")
    rows = []
    for tool, spec in TOOLS:
        wd = tempfile.mkdtemp()
        try:
            csz, ct, dt, ok = bench_mike(p, wd) if spec is None else bench_cmd(p, wd, *spec)
        finally:
            shutil.rmtree(wd, ignore_errors=True)
        if csz < 0: print(f"{tool:<14}{'ERR':>13}"); continue
        rows.append((tool, csz))
        print(f"{tool:<14}{csz:>13,}{osz/csz:>7.2f}x{100*(1-csz/osz):>7.1f}%{ct:>9.2f}{dt:>10.2f}  {'OK' if ok else 'FAIL'}")
    if rows:
        best = min(rows, key=lambda r: r[1])
        if best[0].startswith("MIKE"): wins += 1
        print(f"  -> smallest: {best[0]} ({best[1]:,} b)")
print(f"\nMIKE produced the smallest file on {wins}/{len(FILES)} datasets.")
