"""CPU-bound pure-Python work: free-threaded threads vs. processes.

Every task runs the same loop and touches no shared Python object, so threads
should scale like processes. --interleave starts an idle (detached) thread
after every worker, so no two busy threads get neighbouring QSBR slots.
"""
import argparse
import multiprocessing
import os
import queue
import sys
import threading
import time
from concurrent.futures import ProcessPoolExecutor

LOOP = 200_000
TASKS = 256


def work(_=None):
    acc = 0
    for i in range(LOOP):
        acc += i & 7
    return acc


def run_threads(workers, interleave):
    tasks = queue.SimpleQueue()
    for _ in range(TASKS):
        tasks.put(None)
    go = threading.Event()
    stop = threading.Event()

    def worker():
        go.wait()
        while True:
            try:
                tasks.get_nowait()
            except queue.Empty:
                return
            work()

    busy, idle = [], []
    for _ in range(workers):
        t = threading.Thread(target=worker)
        t.start()
        busy.append(t)
        if interleave:
            t = threading.Thread(target=stop.wait)
            t.start()
            idle.append(t)
    time.sleep(0.05)  # let every thread reach its wait
    t0 = time.perf_counter()
    go.set()
    for t in busy:
        t.join()
    elapsed = time.perf_counter() - t0
    stop.set()
    for t in idle:
        t.join()
    return elapsed


def run_processes(workers):
    ctx = multiprocessing.get_context("spawn")
    with ProcessPoolExecutor(workers, mp_context=ctx) as pool:
        list(pool.map(work, range(2 * workers)))  # start + warm up
        t0 = time.perf_counter()
        list(pool.map(work, range(TASKS), chunksize=1))
        return time.perf_counter() - t0


def best(fn, *args, repeat=3):
    return min(fn(*args) for _ in range(repeat))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", default=None,
                        help="comma-separated worker counts (default: 1,4,8,...,cpu_count)")
    parser.add_argument("--interleave", action="store_true",
                        help="start an idle thread between the worker threads")
    parser.add_argument("--no-processes", action="store_true")
    args = parser.parse_args()

    ncpu = os.cpu_count()
    counts = ([int(n) for n in args.workers.split(",")] if args.workers
              else sorted({1, 4, *range(8, ncpu + 1, 2), ncpu}))
    print(sys.version)
    print("GIL enabled:", sys._is_gil_enabled() if hasattr(sys, "_is_gil_enabled") else True,
          "| cpu_count:", ncpu, "| tasks:", TASKS, "x", LOOP, "iterations")

    rows = [("threads" + (" (interleaved)" if args.interleave else ""),
             lambda n: best(run_threads, n, args.interleave))]
    if not args.no_processes:
        rows.append(("processes", lambda n: best(run_processes, n)))
    for name, fn in rows:
        t1 = fn(1)
        speedups = " ".join(f"{n}:{t1 / fn(n):4.1f}x" for n in counts)
        print(f"{name:<24} 1 worker {t1 * 1000:6.0f} ms   speedup  {speedups}", flush=True)


if __name__ == "__main__":
    main()