#!/usr/bin/env python3
"""Simple throughput benchmark for Parallel Python.

Website: https://www.parallelpython.com
License: Apache-2.0 (see LICENSE and NOTICE)

Runs the same CPU-bound workload with 1 worker (serial) and with N workers,
then prints elapsed times and the measured speedup.  Modeled after the
original pp speed_tests (sum_of_short / sum_of_long), but fully local and
self-contained.

Usage: python benchmark.py [ncpus] [repeats]
"""

import sys
import time

import pp


def _sum_primes(n):
    """Sum of primes below n (CPU-bound, ~0.2-0.3s for n=200k)."""
    def isprime(x):
        if x < 2:
            return False
        i = 2
        while i * i <= x:
            if x % i == 0:
                return False
            i += 1
        return True

    return sum(x for x in range(2, n) if isprime(x))


def _run(server, n_jobs):
    jobs = [server.submit(_sum_primes, (200_000,)) for _ in range(n_jobs)]
    t0 = time.monotonic()
    for j in jobs:
        j()
    return time.monotonic() - t0


def main():
    ncpus = int(sys.argv[1]) if len(sys.argv) > 1 else 4
    repeats = int(sys.argv[2]) if len(sys.argv) > 2 else 8

    ppservers = ()

    serial = pp.Server(1, ppservers=ppservers)
    print(f"Serial (1 worker), {repeats} jobs: ", end="", flush=True)
    t_serial = _run(serial, repeats)
    serial.destroy()
    print(f"{t_serial:.2f}s")

    parallel = pp.Server(ncpus, ppservers=ppservers)
    print(f"Parallel ({ncpus} workers), {repeats} jobs: ", end="", flush=True)
    t_parallel = _run(parallel, repeats)
    parallel.destroy()
    print(f"{t_parallel:.2f}s")

    print(f"Speedup: {t_serial / t_parallel:.2f}x")


if __name__ == "__main__":
    main()