#!/usr/bin/env python3
"""
Compare two multiprocessing.Pool patterns for a large read-only NumPy array.

Naive pattern:
    The large array is included in every submitted task.

Optimized pattern:
    The large array is installed as persistent worker state by a Pool
    initializer. Individual tasks then contain only small (start, length)
    tuples.

The benchmark deliberately uses chunksize=1 so that the communication
pattern is easy to interpret.

Examples
--------
Use Python's default multiprocessing start method:

    python pool_large_data_demo.py

Choose a start method explicitly:

    python pool_large_data_demo.py --start-method forkserver
    python pool_large_data_demo.py --start-method fork
    python pool_large_data_demo.py --start-method spawn
"""

import argparse
import multiprocessing as mp
from time import perf_counter

import numpy as np


# Persistent state inside each worker of the optimized Pool.
# This is one of the relatively few cases where a module-level global is useful:
# process-local, read-only worker state initialized once by the Pool.
_WORKER_READONLY_DATA = None


def reduction_kernel(data, job):
    """Do the same numerical work in both benchmark variants."""
    start, length = job
    x = data[start:start + length]
    return np.mean(x * x)


def task_naive(args):
    """Task receives the large array together with the small job tuple."""
    data, job = args
    return reduction_kernel(data, job)


def init_worker(data):
    """Install the large read-only array once as worker-local global state."""
    global _WORKER_READONLY_DATA
    _WORKER_READONLY_DATA = data


def task_optimized(job):
    """Task receives only the small job tuple."""
    return reduction_kernel(_WORKER_READONLY_DATA, job)


def get_context(name):
    """Return the requested multiprocessing context."""
    if name == "default":
        return mp.get_context()
    return mp.get_context(name)


def run_naive(ctx, data, jobs, workers):
    # This list does NOT make independent copies of data in the parent:
    # every tuple initially contains a reference to the same NumPy array.
    task_args = [(data, job) for job in jobs]

    t0 = perf_counter()
    with ctx.Pool(processes=workers) as pool:
        result = pool.map(task_naive, task_args, chunksize=1)
    elapsed = perf_counter() - t0

    return result, elapsed, task_args


def run_optimized(ctx, data, jobs, workers):
    t0 = perf_counter()
    with ctx.Pool(
        processes=workers,
        initializer=init_worker,
        initargs=(data,),
    ) as pool:
        result = pool.map(task_optimized, jobs, chunksize=1)
    elapsed = perf_counter() - t0

    return result, elapsed


def parse_args():
    parser = argparse.ArgumentParser(
        description="Naive versus optimized multiprocessing.Pool data passing"
    )
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--array-mib", type=int, default=32)
    parser.add_argument("--tasks", type=int, default=80)
    parser.add_argument("--slice-length", type=int, default=200_000)
    parser.add_argument(
        "--start-method",
        choices=["default", "fork", "forkserver", "spawn"],
        default="default",
        help="multiprocessing start method (default: Python/platform default)",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    ctx = get_context(args.start_method)
    method = ctx.get_start_method()

    # float64 = 8 bytes
    n = args.array_mib * 1024**2 // np.dtype(np.float64).itemsize
    if args.slice_length > n:
        raise ValueError("slice length must not exceed array length")

    rng = np.random.default_rng(12345)
    data = rng.random(n)

    max_start = n - args.slice_length
    starts = rng.integers(0, max_start + 1, size=args.tasks)
    jobs = [(int(start), args.slice_length) for start in starts]

    nominal_payload_gib = data.nbytes * args.tasks / 1024**3

    print("=" * 72)
    print("LARGE READ-ONLY ARRAY: NAIVE POOL VS OPTIMIZED POOL")
    print("=" * 72)
    print(f"Python start method       : {method}")
    print(f"workers                   : {args.workers}")
    print(f"tasks                     : {args.tasks}")
    print(f"array size                : {data.nbytes / 1024**2:.1f} MiB")
    print(f"slice length              : {args.slice_length:,}")
    print(f"slice size                : {args.slice_length * data.itemsize / 1024**2:.3f} MiB")
    print(f"naive nominal array payload: {nominal_payload_gib:.2f} GiB")
    print()

    r_naive, t_naive, task_args = run_naive(
        ctx, data, jobs, args.workers
    )

    # Demonstrate that the parent-side argument list contains references,
    # not independent NumPy copies.
    if len(task_args) >= 2:
        print(
            "same parent array object in task 0 and task 1:",
            task_args[0][0] is task_args[1][0],
        )

    r_opt, t_opt = run_optimized(
        ctx, data, jobs, args.workers
    )

    print()
    print("Timings")
    print("-------")
    print(f"naive Pool     : {t_naive:9.3f} s")
    print(f"optimized Pool : {t_opt:9.3f} s")
    print(f"speedup        : {t_naive / t_opt:9.2f} x")
    print()
    print("Validation")
    print("----------")
    print("same numerical results:", np.allclose(r_naive, r_opt))
    print()
    print("Interpretation")
    print("--------------")
    print(
        "Naive Pool: the large array is part of every submitted task, so it is "
        "repeatedly serialized/transferred through multiprocessing IPC."
    )
    print(
        "Optimized Pool: the large array is installed as persistent worker state; "
        "each queued task contains only a small (start, length) tuple."
    )
    if method == "fork":
        print(
            "With fork, read-only inherited data can additionally benefit from "
            "copy-on-write."
        )
    else:
        print(
            f"With {method}, the large object must be made available to each worker, "
            "but only once per worker rather than once per task."
        )


if __name__ == "__main__":
    main()