# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "joblib",
#     "pyvista",
# ]
# ///
import concurrent
import concurrent.futures
import itertools
import time

import joblib
import pyvista as pv


def process(mesh: pv.PolyData, idx: int) -> float:
    """A simple function that accesses data from the mesh."""
    return mesh.points[idx, 0]


def main() -> None:
    # Create a large mesh object to demonstrate the issue.
    # A smaller `level` will reduce the mesh size and the performance gap.
    mesh: pv.PolyData = pv.Box(level=200)
    N_JOBS: int = 1000
    print("Mesh details:")
    print(mesh)
    print("-" * 20)

    # Benchmark joblib with automatic batching
    with joblib.parallel_config(n_jobs=8, verbose=1, prefer="processes"):
        time_start: float = time.perf_counter()
        parallel = joblib.Parallel(batch_size="auto")
        _ = parallel(joblib.delayed(process)(mesh, i) for i in range(N_JOBS))
        time_end: float = time.perf_counter()
        rate: float = N_JOBS / (time_end - time_start)
        print(f"joblib (auto batch size): {rate:.2f} it/sec")

        # Benchmark joblib with manual batching
        time_start: float = time.perf_counter()
        parallel = joblib.Parallel(batch_size=N_JOBS // 16)
        _ = parallel(joblib.delayed(process)(mesh, i) for i in range(N_JOBS))
        time_end: float = time.perf_counter()
        rate: float = N_JOBS / (time_end - time_start)
        print(f"joblib (manual batch size): {rate:.2f} it/sec")

    # Benchmark concurrent.futures.ProcessPoolExecutor
    time_start = time.perf_counter()
    with concurrent.futures.ProcessPoolExecutor(max_workers=8) as executor:
        _ = executor.map(
            process,
            itertools.repeat(mesh),
            range(N_JOBS),
            chunksize=N_JOBS // 16,
        )
    time_end = time.perf_counter()
    rate: float = N_JOBS / (time_end - time_start)
    print(f"concurrent.futures: {rate:.2f} it/sec")


if __name__ == "__main__":
    main()