#!/usr/bin/env python3
"""
Portable Microbench (no external deps)
- CPU integer: Sieve of Eratosthenes
- CPU float: sin(sqrt) accumulation
- Hash: SHA-256 throughput
- Text/Regex: findall on synthesized text
- JSON/Gzip: serialize/deserialize & compress
- SQLite: insert & aggregate in-memory
- Parallel scaling: float workload with multiprocessing

Usage:
    python3 microbench.py            # run with defaults
    python3 microbench.py --secs 2   # target seconds per test (default 1.5)
    python3 microbench.py --quick    # faster (≈0.7s/test)
    python3 microbench.py --long     # slower (≈3s/test)
    python3 microbench.py --no-par   # skip parallel test
    python3 microbench.py --csv out.csv  # also save results as CSV
    python3 microbench.py --json out.json # also save results as JSON

Test notes:
- Designed to run on tiny SBCs up to beefier ARM boards.
- Pure-Python by default; uses stdlib only.
- Results are *relative* across devices; absolute values are for rough comparison.
"""

import argparse
import gzip
import hashlib
import io
import json
import math
import os
import platform
import random
import re
import sqlite3
import statistics
import sys
import time
from multiprocessing import Pool, cpu_count

# ---------- Utilities ----------

def human(n):
    # human-friendly number formatting
    if n is None:
        return "-"
    if n >= 1e9:
        return f"{n/1e9:.2f}G"
    if n >= 1e6:
        return f"{n/1e6:.2f}M"
    if n >= 1e3:
        return f"{n/1e3:.2f}k"
    return f"{n:.2f}"

def now():
    return time.perf_counter()

def calibrate(fn, make_args, target_secs, min_iters=1):
    """Find an iteration count so total runtime ~ target_secs."""
    iters = min_iters
    # warmup
    fn(*make_args(iters))
    start = now()
    fn(*make_args(iters))
    elapsed = now() - start
    if elapsed <= 0:
        elapsed = 1e-6
    # scale to hit target
    scale = max(target_secs / elapsed, 1.0)
    iters = max(int(iters * scale), min_iters)
    # Limit excessive iters
    iters = min(iters, 10**9)
    return iters

def run_bench(fn, make_args, target_secs, unit_label, higher_is_better=True):
    iters = calibrate(fn, make_args, target_secs)
    start = now()
    fn(*make_args(iters))
    elapsed = now() - start
    # throughput = iterations per second (or bytes/sec inside fn if returns custom)
    # If the function returns a numeric throughput, use that; else use iterations.
    throughput = None
    try:
        result = fn(*make_args(iters))
        if isinstance(result, (int, float)) and result > 0:
            throughput = result / elapsed
    except Exception:
        # Some tests might not be re-entrant; run once for timing only
        pass
    if throughput is None:
        throughput = iters / elapsed
    return {
        "iters": iters,
        "elapsed_s": elapsed,
        "throughput": throughput,
        "unit": unit_label,
        "higher_is_better": higher_is_better,
    }

# ---------- Workloads ----------

def sieve_count(n):
    """Return the count of primes up to n using a simple sieve."""
    sieve = bytearray(b"\x01") * (n + 1)
    sieve[0:2] = b"\x00\x00"
    m = int(n ** 0.5) + 1
    for p in range(2, m):
        if sieve[p]:
            step = p
            start = p*p
            sieve[start:n+1:step] = b"\x00" * (((n - start) // step) + 1)
    return int(sum(sieve))

def bench_sieve(iters, limit=200_000):
    # Count primes up to 'limit', repeated 'iters' times
    total = 0
    for _ in range(iters):
        total += sieve_count(limit)
    # Return iterations as measure (ops = runs)
    return iters

def bench_float(iters, n=400_000):
    # Floating-point heavy loop
    acc = 0.0
    for _ in range(iters):
        # sum sin(sqrt(i)) for i in range(n)
        a = 0.0
        for i in range(n):
            a += math.sin(math.sqrt(i + 0.123))
        acc += a
    # Return iterations as measure
    # Prevent dead-code elimination
    if acc == -1.0:
        print("unlikely")
    return iters

def bench_hash(iters, block_size=64, blocks=1_000_000):
    # Hash 'blocks' of pseudo-random data per iter; fixed seed
    rnd = random.Random(42)
    data = bytes(rnd.getrandbits(8) for _ in range(block_size))
    total_bytes = 0
    for _ in range(iters):
        h = hashlib.sha256()
        # stream multiple blocks
        for _b in range(blocks):
            h.update(data)
        _ = h.digest()
        total_bytes += block_size * blocks
    # Return bytes hashed as measure
    return total_bytes

_LOREM = " ".join([
    "lorem", "ipsum", "dolor", "sit", "amet,", "consectetur",
    "adipiscing", "elit.", "Sed", "do", "eiusmod", "tempor",
    "incididunt", "ut", "labore", "et", "dolore", "magna", "aliqua."
])
_REGEX = re.compile(r"\b(?:dolor|elit|tempor|aliqua)\b")

def bench_regex(iters, para_reps=5000):
    text = (" " + _LOREM) * para_reps
    count = 0
    for _ in range(iters):
        count += len(_REGEX.findall(text))
    if count == -1:
        print("never")
    return iters

def bench_json_gzip(iters, objs=15_000):
    payload = [
        {"id": i, "x": i % 7, "y": (i * 13) % 97, "name": f"name_{i}", "ok": (i % 2 == 0)}
        for i in range(objs)
    ]
    total_bytes = 0
    for _ in range(iters):
        s = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        # round-trip JSON
        _ = json.loads(s)
        # compress
        buf = io.BytesIO()
        with gzip.GzipFile(fileobj=buf, mode="wb", compresslevel=5) as f:
            f.write(s)
        total_bytes += len(s)
    return total_bytes

def bench_sqlite(iters, rows=40_000):
    total_rows = 0
    for _ in range(iters):
        con = sqlite3.connect(":memory:")
        cur = con.cursor()
        cur.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, a INT, b INT)")
        cur.executemany("INSERT INTO t (a,b) VALUES (?,?)",
                        ((i % 97, (i*7) % 101) for i in range(rows)))
        cur.execute("SELECT SUM(a*b) FROM t WHERE a BETWEEN 10 AND 60")
        _ = cur.fetchone()[0]
        con.close()
        total_rows += rows
    return total_rows

def _float_worker(args):
    # worker for parallel float test
    iters, n = args
    return bench_float(iters, n)

def bench_parallel_float(target_secs, workers=None, n=250_000):
    cores = max(1, cpu_count() or 1)
    if workers is None:
        workers = min(4, cores)
    # Calibrate per-process iters
    iters = calibrate(bench_float, lambda it: (it, n), target_secs / 1.2)  # a tad shorter
    # Run in parallel
    start = now()
    with Pool(processes=workers) as pool:
        results = pool.map(_float_worker, [(iters, n)] * workers)
    elapsed = now() - start
    total_iters = sum(results)
    thr = total_iters / elapsed
    return {
        "workers": workers,
        "iters_per_worker": iters,
        "elapsed_s": elapsed,
        "throughput": thr,
        "unit": "iters/s",
        "higher_is_better": True,
        "cores": cores,
    }

# ---------- Runner ----------

TESTS = [
    ("CPU int (sieve)", lambda secs: run_bench(bench_sieve, lambda it: (it, 200_000), secs, "runs/s")),
    ("CPU float (sin√)", lambda secs: run_bench(bench_float, lambda it: (it, 400_000), secs, "runs/s")),
    ("SHA-256",        lambda secs: run_bench(bench_hash,  lambda it: (it, 64, 200_000), secs, "bytes/s")),
    ("Regex",          lambda secs: run_bench(bench_regex, lambda it: (it, 5000), secs, "runs/s")),
    ("JSON+Gzip",      lambda secs: run_bench(bench_json_gzip, lambda it: (it, 10_000), secs, "bytes/s")),
    ("SQLite",         lambda secs: run_bench(bench_sqlite, lambda it: (it, 30_000), secs, "rows/s")),
]

def gather_sysinfo():
    info = {
        "python": sys.version.split()[0],
        "impl": platform.python_implementation(),
        "platform": platform.platform(),
        "machine": platform.machine(),
        "processor": platform.processor(),
        "cpu_count": cpu_count() or 1,
    }
    # Try to read /proc/cpuinfo & meminfo if present
    try:
        with open("/proc/cpuinfo", "r") as f:
            info["/proc/cpuinfo"] = f.read().strip()
    except Exception:
        pass
    try:
        with open("/proc/meminfo", "r") as f:
            info["/proc/meminfo"] = f.read().strip()
    except Exception:
        pass
    return info

def geometric_mean(nums):
    nums = [x for x in nums if x > 0]
    if not nums:
        return 0.0
    log_sum = sum(math.log(x) for x in nums)
    return math.exp(log_sum / len(nums))

def main():
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group()
    g.add_argument("--quick", action="store_true", help="~0.7s per test")
    g.add_argument("--long", action="store_true", help="~3s per test")
    ap.add_argument("--secs", type=float, default=None, help="target seconds per test (default 1.5)")
    ap.add_argument("--no-par", action="store_true", help="skip parallel test")
    ap.add_argument("--csv", type=str, default=None, help="also write CSV here")
    ap.add_argument("--json", type=str, default=None, help="also write JSON here")
    args = ap.parse_args()

    target = 1.5
    if args.quick:
        target = 0.7
    if args.long:
        target = 3.0
    if args.secs is not None:
        target = max(0.3, float(args.secs))

    print("Portable Microbench\n")
    sysinfo = gather_sysinfo()
    print(f"Python: {sysinfo['python']} ({sysinfo['impl']})")
    print(f"Platform: {sysinfo['platform']}")
    print(f"Machine: {sysinfo['machine']} | CPU count: {sysinfo['cpu_count']}")
    print()

    results = []
    for name, runner in TESTS:
        print(f"[{name}] target ≈ {target:.1f}s ...", flush=True)
        r = runner(target)
        r["name"] = name
        results.append(r)
        print(f"  {name:<18} {human(r['throughput'])} {r['unit']}  (elapsed {r['elapsed_s']:.2f}s)")

    par = None
    if not args.no_par:
        print(f"[Parallel float] target ≈ {target:.1f}s ...", flush=True)
        par = bench_parallel_float(target, None, 250_000)
        print(f"  Parallel float    {human(par['throughput'])} {par['unit']}  "
              f"({par['workers']} workers, elapsed {par['elapsed_s']:.2f}s)")

    # Compute a composite score using geometric mean of normalized throughputs.
    # We normalize each test by dividing by the median throughput of all tests on this device,
    # so the composite is roughly balanced across kinds of work.
    thr_values = [r["throughput"] for r in results if r["higher_is_better"]]
    median_thr = statistics.median(thr_values) if thr_values else 1.0
    normalized = [r["throughput"] / median_thr for r in results if r["higher_is_better"]]
    score = geometric_mean(normalized) * 100.0  # mean around 100 on each device

    print("\nSummary:")
    print(f"{'Test':<20} {'Throughput':>14}  {'Unit'}")
    print("-"*40)
    for r in results:
        print(f"{r['name']:<20} {human(r['throughput']):>14}  {r['unit']}")
    if par:
        print(f"{'Parallel float':<20} {human(par['throughput']):>14}  {par['unit']} ({par['workers']} workers)")
    print("-"*40)
    print(f"Composite score (≈relative on this device): {score:.1f}")

    out = {
        "sysinfo": sysinfo,
        "results": results,
        "parallel": par,
        "composite_score": score,
        "target_secs": target,
        "ts": time.time(),
    }

    if args.csv:
        try:
            with open(args.csv, "w") as f:
                f.write("name,throughput,unit,elapsed_s,iters\n")
                for r in results:
                    f.write(f"{r['name']},{r['throughput']},{r['unit']},{r['elapsed_s']},{r['iters']}\n")
                if par:
                    f.write(f"Parallel float,{par['throughput']},{par['unit']},{par['elapsed_s']},{par['iters_per_worker']*par['workers']}\n")
            print(f"\nWrote CSV: {args.csv}")
        except Exception as e:
            print(f"Failed to write CSV: {e}", file=sys.stderr)

    if args.json:
        try:
            with open(args.json, "w") as f:
                json.dump(out, f, indent=2)
            print(f"Wrote JSON: {args.json}")
        except Exception as e:
            print(f"Failed to write JSON: {e}", file=sys.stderr)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nInterrupted.")