"""
This function provides functions for parallel execution of loops and function
mappings, using the builtin Python module multiprocessing.
"""
__all__ = ['parfor', 'parallel_map', 'serial_map']

from scipy import array
import multiprocessing
from functools import partial
import os
import sys
import signal
import qutip.settings as qset
from qutip.ui.progressbar import BaseProgressBar, TextProgressBar


if sys.platform == 'darwin':
    Pool = multiprocessing.get_context('fork').Pool
else:
    Pool = multiprocessing.Pool


def _task_wrapper(args):
    try:
        return args[0](*args[1])
    except KeyboardInterrupt:
        os.kill(args[2], signal.SIGINT)
        sys.exit(1)


def _task_wrapper_with_args(args, user_args):
    try:
        return args[0](*args[1], **user_args)
    except KeyboardInterrupt:
        os.kill(args[2], signal.SIGINT)
        sys.exit(1)


def parfor(func, *args, **kwargs):
    """Executes a multi-variable function in parallel on the local machine.

    Parallel execution of a for-loop over function `func` for multiple input
    arguments and keyword arguments.

    .. note::

        From QuTiP 3.1, we recommend to use :func:`qutip.parallel.parallel_map`
        instead of this function.

    Parameters
    ----------
    func : function_type
        A function to run in parallel on the local machine. The function 'func'
        accepts a series of arguments that are passed to the function as
        variables. In general, the function can have multiple input variables,
        and these arguments must be passed in the same order as they are
        defined in the function definition.  In addition, the user can pass
        multiple keyword arguments to the function.

    The following keyword argument is reserved:

    num_cpus : int
        Number of CPU's to use.  Default uses maximum number of CPU's.
        Performance degrades if num_cpus is larger than the physical CPU
        count of your machine.

    Returns
    -------
    result : list
        A ``list`` with length equal to number of input parameters
        containing the output from `func`.

    """
    os.environ['QUTIP_IN_PARALLEL'] = 'TRUE'
    kw = _default_kwargs()
    if 'num_cpus' in kwargs.keys():
        kw['num_cpus'] = kwargs['num_cpus']
        del kwargs['num_cpus']
    if len(kwargs) != 0:
        task_func = partial(_task_wrapper_with_args, user_args=kwargs)
    else:
        task_func = _task_wrapper

    if kw['num_cpus'] > qset.num_cpus:
        print("Requested number of CPUs (%s) " % kw['num_cpus'] +
              "is larger than physical number (%s)." % qset.num_cpus)
        print("Reduce 'num_cpus' for greater performance.")

    pool = Pool(processes=kw['num_cpus'])
    args = [list(arg) for arg in args]
    var = [[args[j][i] for j in range(len(args))]
           for i in range(len(list(args[0])))]
    try:
        map_args = ((func, v, os.getpid()) for v in var)
        par_return = list(pool.map(task_func, map_args))

        pool.terminate()
        pool.join()
        os.environ['QUTIP_IN_PARALLEL'] = 'FALSE'
        if isinstance(par_return[0], tuple):
            par_return = [elem for elem in par_return]
            num_elems = len(par_return[0])
            dt = [type(ii) for ii in par_return[0]]
            return [array([elem[ii] for elem in par_return], dtype=dt[ii])
                    for ii in range(num_elems)]
        else:
            return par_return

    except KeyboardInterrupt:
        os.environ['QUTIP_IN_PARALLEL'] = 'FALSE'
        pool.terminate()


def serial_map(task, values, task_args=tuple(), task_kwargs={}, **kwargs):
    """
    Serial mapping function with the same call signature as parallel_map, for
    easy switching between serial and parallel execution. This
    is functionally equivalent to::

        result = [task(value, *task_args, **task_kwargs) for value in values]

    This function work as a drop-in replacement of
    :func:`qutip.parallel.parallel_map`.

    Parameters
    ----------
    task : a Python function
        The function that is to be called for each value in ``task_vec``.
    values : array / list
        The list or array of values for which the ``task`` function is to be
        evaluated.
    task_args : list / dictionary
        The optional additional argument to the ``task`` function.
    task_kwargs : list / dictionary
        The optional additional keyword argument to the ``task`` function.
    progress_bar : ProgressBar
        Progress bar class instance for showing progress.

    Returns
    --------
    result : list
        The result list contains the value of
        ``task(value, *task_args, **task_kwargs)`` for each
        value in ``values``.

    """
    try:
        progress_bar = kwargs['progress_bar']
        if progress_bar is True:
            progress_bar = TextProgressBar()
    except:
        progress_bar = BaseProgressBar()

    progress_bar.start(len(values))
    results = []
    for n, value in enumerate(values):
        progress_bar.update(n)
        result = task(value, *task_args, **task_kwargs)
        results.append(result)
    progress_bar.finished()

    return results


def parallel_map(task, values, task_args=tuple(), task_kwargs={}, **kwargs):
    """
    Parallel execution of a mapping of `values` to the function `task`. This
    is functionally equivalent to::

        result = [task(value, *task_args, **task_kwargs) for value in values]

    Parameters
    ----------
    task : a Python function
        The function that is to be called for each value in ``task_vec``.
    values : array / list
        The list or array of values for which the ``task`` function is to be
        evaluated.
    task_args : list / dictionary
        The optional additional argument to the ``task`` function.
    task_kwargs : list / dictionary
        The optional additional keyword argument to the ``task`` function.
    progress_bar : ProgressBar
        Progress bar class instance for showing progress.

    Returns
    --------
    result : list
        The result list contains the value of
        ``task(value, *task_args, **task_kwargs)`` for
        each value in ``values``.

    """
    os.environ['QUTIP_IN_PARALLEL'] = 'TRUE'
    kw = _default_kwargs()
    if 'num_cpus' in kwargs:
        kw['num_cpus'] = kwargs['num_cpus']

    try:
        progress_bar = kwargs['progress_bar']
        if progress_bar is True:
            progress_bar = TextProgressBar()
    except:
        progress_bar = BaseProgressBar()

    progress_bar.start(len(values))
    nfinished = [0]

    def _update_progress_bar(x):
        nfinished[0] += 1
        progress_bar.update(nfinished[0])

    try:
        pool = Pool(processes=kw['num_cpus'])

        async_res = [pool.apply_async(task, (value,) + task_args, task_kwargs,
                                      _update_progress_bar)
                     for value in values]

        while not all([ar.ready() for ar in async_res]):
            for ar in async_res:
                ar.wait(timeout=0.1)

        pool.terminate()
        pool.join()

    except KeyboardInterrupt as e:
        os.environ['QUTIP_IN_PARALLEL'] = 'FALSE'
        pool.terminate()
        pool.join()
        raise e

    progress_bar.finished()
    os.environ['QUTIP_IN_PARALLEL'] = 'FALSE'
    return [ar.get() for ar in async_res]


def _default_kwargs():
    settings = {'num_cpus': qset.num_cpus}
    return settings
