from typing import Any, Callable, Iterable, Optional

import numpy as np
from joblib import Parallel, delayed
from numpy.typing import NDArray

from .typing import FilePath

__all__ = ["execute_parallel", "reduce_parallel"]


def execute_parallel(
    function: Callable[[FilePath, FilePath], None],
    srcs: Iterable[FilePath],
    dsts: Iterable[FilePath],
    *,
    n_jobs: Optional[int] = None,
    **kwargs: Any,
) -> None:
    """Execute function in parallel."""

    delayed_function = delayed(function)
    with Parallel(n_jobs=n_jobs, **kwargs) as parallel:
        parallel(delayed_function(src, dst) for src, dst in zip(srcs, dsts))


def reduce_parallel(
    function: Callable[[FilePath], NDArray],
    srcs: Iterable[FilePath],
    *,
    n_jobs: Optional[int] = None,
    **kwargs: Any,
) -> NDArray:
    """Reduce function in parallel."""

    delayed_function = delayed(function)
    with Parallel(n_jobs=n_jobs, **kwargs) as parallel:
        result = parallel(delayed_function(src) for src in srcs)
    return np.stack(result)
