"""Main module."""
from itertools import starmap 
from typing import TYPE_CHECKING 


if TYPE_CHECKING:
    from typing import Optional, Callable, Dict 
    from collections.abc import Iterable 


def _map_parallel_multiprocessing(
    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True
) -> list:
    from concurrent.futures import ProcessPoolExecutor 

    with ProcessPoolExecutor(max_workers=processes) as process_pool_executor:
        res = process_pool_executor.map(f, *args)
        if return_results:
            return list(res)
        else:
            return []


def _map_parallel_multithreading(
    f: Callable, 
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True
) -> list:
    from concurrent.futures import ThreadPoolExecutor

    with ThreadPoolExecutor(max_workers=processes) as thread_pool_executor:
        res = thread_pool_executor.map(f, *args)
        if return_results:
            return list(res)
        else:
            return []


def _map_parallel_dask(
    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True
) -> list:
    from dask.distributed import Client 
    from dask.distributed import LocalCluster 

    cluster = LocalCluster(max_workers=processes, dashboard=None)
    client = Client(cluster)
    if return_results:
        return [future.result() for future in client.map(f, *args)] 
    else:
        for future in client.map(f, *args):
            future.result()
        return []


def _map_parallel_mpi(
    f: Callable,
    args: Iterable[Iterable],
    return_results: bool = True,
    **kwargs
) -> list:
    from mpi4py.futures import MPIPoolExecutor

    with MPIPoolExecutor() as mpi_pool_executor:
        res = mpi_pool_executor.map(f, args)
        if return_results:
            return list(res)
        else:
            return []


def _starmap_parallel_mpi(
    f: Callable, 
    args: Iterable[Iterable],
    return_results: bool = True,
    **kwargs  
) -> list:
    from mpi4py.futures import MPIPoolExecutor

    with MPIPoolExecutor() as mpi_pool_executor:
        res = mpi_pool_executor.starmap(f, args) 
        if return_results:
            return list(res)
        else:
            return []


def _starmap_parallel_mpi_simple(
    f: Callable, 
    args: Iterable[Iterable],
    return_results: bool = True,
    **kwargs,
) -> list:
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()

    args_list = list(args)
    if args_list:
        n = len(args_list)
        start = (rank * n) // size 
        end = (( rank + 1) * n) // size 
        local_args = args_list[start:end]
        res = list(starmap(f, local_args))

        if return_results:
            res = comm.gather(res, root=0)
            if rank == 0:
                return sum(res, [])
    return []


_map_parallel_func: 'Dict[str, Callable]' = {
    'multiprocessing': _map_parallel_multiprocessing,
    'multithreading': _map_parallel_multithreading,
    'dask': _map_parallel_dask,
    'mpi': _map_parallel_mpi,
}


_starmap_parallel_func: 'Dict[str, Callable]' = {
    'mpi': _starmap_parallel_mpi,
    'mpi_simple': _starmap_parallel_mpi_simple,
}


def map_parallel(
    f: 'Callable',
    *args: 'Iterable',
    processes: 'Optional[int]' = None,
    mode: 'str' = 'multiprocessing',
    return_results: 'bool' = True,
) -> 'list':
    '''equiv to `map(f, *args)` but in parallel
    :param str mode: backend for parallelization
        - multiprocessing: using multiprocessing from standard library
        - multithreading: using multithreading from standard library
        - dask: using dask.distributed
        - mpi: using mpi4py.futures. May not work depending on your MPI vendor
        - mpi_simple: using mpi4py with simple scheduling that divides works into equal chunks
        - serial: using map
    :param int processes: no. of parallel processes
        (in the case of mpi, it is determined by mpiexec/mpirun args)
    :param bool return_results: (Only affects mode == 'mpi_simple') if True, return results in rank 0.
    '''
    if processes is None or processes > 1:
        if mode in _map_parallel_func:
            return _map_parallel_func[mode](f, *args, processes=processes, return_results=return_results)
        elif mode in _starmap_parallel_func:
            return _starmap_parallel_func[mode](f, zip(*args), processes=processes, return_results=return_results)
    return list(map(f, *args))


def starmap_parallel(
    f: 'Callable',
    args: 'Iterable[Iterable]',
    processes: 'Optional[int]' = None,
    mode: 'str' = 'multiprocessing',
    return_results: 'bool' = True,
) -> 'list':
    '''equiv to `starmap(f, args)` but in parallel
    See docstring from :func:`~map_parallel.map_parallel`
    '''
    if processes is None or processes > 1:
        if mode in _starmap_parallel_func:
            return _starmap_parallel_func[mode](f, args, processes=processes, return_results=return_results)
        elif mode in _map_parallel_func:
            return _map_parallel_func[mode](f, *zip(*args), processes=processes, return_results=return_results)
    return list(starmap(f, args))
