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 process2(mesh, idx):
    return process(mesh.result(), idx)

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()
        shelved_mesh = joblib.shelve(mesh)
        parallel = joblib.Parallel()
        _ = parallel(joblib.delayed(process2)(shelved_mesh, i) for i in range(N_JOBS))
        shelved_mesh.clear()
        time_end: float = time.perf_counter()
        rate: float = N_JOBS / (time_end - time_start)
        print(f"joblib (with shelving): {rate:.2f} it/sec")

        time_start: float = time.perf_counter()
        parallel = joblib.Parallel()
        _ = 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: {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()