"""
Test the parallel module.
"""

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# Copyright (c) 2010-2011 Gael Varoquaux
# License: BSD Style, 3 clauses.

import time
import sys
import os
from math import sqrt
import threading
from multiprocessing import TimeoutError
from time import sleep

from joblib import parallel

from joblib.test.common import np, with_numpy
from joblib.test.common import with_multiprocessing
from joblib.testing import raises, check_subprocess_call, SkipTest, warns
from joblib._compat import PY3_OR_LATER

try:
    import cPickle as pickle
    PickleError = TypeError
except ImportError:
    import pickle
    PickleError = pickle.PicklingError


if PY3_OR_LATER:
    PickleError = pickle.PicklingError

try:
    # Python 2/Python 3 compat
    unicode('str')
except NameError:
    unicode = lambda s: s

try:
    from queue import Queue
except ImportError:
    # Backward compat
    from Queue import Queue

try:
    import posix
except ImportError:
    posix = None

from joblib._parallel_backends import SequentialBackend
from joblib._parallel_backends import ThreadingBackend
from joblib._parallel_backends import MultiprocessingBackend
from joblib._parallel_backends import SafeFunction
from joblib._parallel_backends import WorkerInterrupt

from joblib.parallel import Parallel, delayed
from joblib.parallel import register_parallel_backend, parallel_backend

from joblib.parallel import mp, cpu_count, BACKENDS, effective_n_jobs
from joblib.my_exceptions import JoblibException


ALL_VALID_BACKENDS = [None] + sorted(BACKENDS.keys())
# Add instances of backend classes deriving from ParallelBackendBase
ALL_VALID_BACKENDS += [BACKENDS[backend_str]() for backend_str in BACKENDS]

if hasattr(mp, 'get_context'):
    # Custom multiprocessing context in Python 3.4+
    ALL_VALID_BACKENDS.append(mp.get_context('spawn'))


def division(x, y):
    return x / y


def square(x):
    return x ** 2


class MyExceptionWithFinickyInit(Exception):
    """An exception class with non trivial __init__
    """
    def __init__(self, a, b, c, d):
        pass


def exception_raiser(x, custom_exception=False):
    if x == 7:
        raise (MyExceptionWithFinickyInit('a', 'b', 'c', 'd')
               if custom_exception else ValueError)
    return x


def interrupt_raiser(x):
    time.sleep(.05)
    raise KeyboardInterrupt


def f(x, y=0, z=0):
    """ A module-level function so that it can be spawn with
    multiprocessing.
    """
    return x ** 2 + y + z


def _active_backend_type():
    return type(parallel.get_active_backend()[0])


def parallel_func(inner_n_jobs):
    return Parallel(n_jobs=inner_n_jobs)(delayed(square)(i) for i in range(3))


###############################################################################
def test_cpu_count():
    assert cpu_count() > 0


def test_effective_n_jobs():
    assert effective_n_jobs() > 0


###############################################################################
# Test parallel
def check_simple_parallel(backend):
    X = range(5)
    for n_jobs in (1, 2, -1, -2):
        assert ([square(x) for x in X] ==
                Parallel(n_jobs=n_jobs, backend=backend)(
                    delayed(square)(x) for x in X))
    for verbose in (2, 11, 100):
        Parallel(n_jobs=-1, verbose=verbose, backend=backend)(
            delayed(square)(x) for x in X)
        Parallel(n_jobs=1, verbose=verbose, backend=backend)(
            delayed(square)(x) for x in X)
        Parallel(n_jobs=2, verbose=verbose, pre_dispatch=2,
                 backend=backend)(
            delayed(square)(x) for x in X)
        Parallel(n_jobs=2, verbose=verbose, backend=backend)(
            delayed(square)(x) for x in X)


def test_simple_parallel():
    for backend in ALL_VALID_BACKENDS:
        check_simple_parallel(backend)


def check_main_thread_renamed_no_warning(backend):
    with warns(None) as warninfo:
        results = Parallel(n_jobs=2, backend=backend)(
            delayed(square)(x) for x in range(3))
        assert results == [0, 1, 4]
    # The multiprocessing backend will raise a warning when detecting that is
    # started from the non-main thread. Let's check that there is no false
    # positive because of the name change.
    assert len(warninfo) == 0


def test_main_thread_renamed_no_warning():
    # Check that no default backend relies on the name of the main thread:
    # https://github.com/joblib/joblib/issues/180#issuecomment-253266247
    # Some programs use a different name for the main thread. This is the case
    # for uWSGI apps for instance.
    main_thread = threading.current_thread()
    original_name = main_thread.name
    try:
        main_thread.name = "some_new_name_for_the_main_thread"
        for backend in ALL_VALID_BACKENDS:
            check_main_thread_renamed_no_warning(backend)
    finally:
        main_thread.name = original_name


def nested_loop(backend):
    Parallel(n_jobs=2, backend=backend)(
        delayed(square)(.01) for _ in range(2))


def check_nested_loop(parent_backend, child_backend):
    Parallel(n_jobs=2, backend=parent_backend)(
        delayed(nested_loop)(child_backend) for _ in range(2))


def test_nested_loop():
    for parent_backend in BACKENDS:
        for child_backend in BACKENDS:
            check_nested_loop(parent_backend, child_backend)


def test_mutate_input_with_threads():
    """Input is mutable when using the threading backend"""
    q = Queue(maxsize=5)
    Parallel(n_jobs=2, backend="threading")(
        delayed(q.put, check_pickle=False)(1) for _ in range(5))
    assert q.full()


def test_parallel_kwargs():
    """Check the keyword argument processing of pmap."""
    lst = range(10)
    for n_jobs in (1, 4):
        assert ([f(x, y=1) for x in lst] ==
                Parallel(n_jobs=n_jobs)(delayed(f)(x, y=1) for x in lst))


def check_parallel_as_context_manager(backend):
    lst = range(10)
    expected = [f(x, y=1) for x in lst]
    with Parallel(n_jobs=4, backend=backend) as p:
        # Internally a pool instance has been eagerly created and is managed
        # via the context manager protocol
        managed_backend = p._backend
        if mp is not None:
            assert managed_backend is not None
            assert managed_backend._pool is not None

        # We make call with the managed parallel object several times inside
        # the managed block:
        assert expected == p(delayed(f)(x, y=1) for x in lst)
        assert expected == p(delayed(f)(x, y=1) for x in lst)

        # Those calls have all used the same pool instance:
        if mp is not None:
            assert managed_backend._pool is p._backend._pool

    # As soon as we exit the context manager block, the pool is terminated and
    # no longer referenced from the parallel object:
    if mp is not None:
        assert p._backend._pool is None

    # It's still possible to use the parallel instance in non-managed mode:
    assert expected == p(delayed(f)(x, y=1) for x in lst)
    if mp is not None:
        assert p._backend._pool is None


def test_parallel_context_manager():
    for backend in ['multiprocessing', 'threading']:
        check_parallel_as_context_manager(backend)


def test_parallel_pickling():
    """ Check that pmap captures the errors when it is passed an object
        that cannot be pickled.
    """
    def g(x):
        return x ** 2

    try:
        # pickling a local function always fail but the exception
        # raised is a PickleError for python <= 3.4 and AttributeError
        # for python >= 3.5
        pickle.dumps(g)
    except Exception as exc:
        exception_class = exc.__class__

    with raises(exception_class):
        Parallel()(delayed(g)(x) for x in range(10))


def test_parallel_timeout_success():
    # Check that timeout isn't thrown when function is fast enough
    for backend in ['multiprocessing', 'threading']:
        assert len(Parallel(n_jobs=2, backend=backend, timeout=10)(
            delayed(sleep)(0.001) for x in range(10))) == 10


@with_multiprocessing
def test_parallel_timeout_fail():
    # Check that timeout properly fails when function is too slow
    for backend in ['multiprocessing', 'threading']:
        with raises(TimeoutError):
            Parallel(n_jobs=2, backend=backend, timeout=0.01)(
                delayed(sleep)(10) for x in range(10))


def test_error_capture():
    # Check that error are captured, and that correct exceptions
    # are raised.
    if mp is not None:
        # A JoblibException will be raised only if there is indeed
        # multiprocessing
        with raises(JoblibException):
            Parallel(n_jobs=2)(
                [delayed(division)(x, y)
                    for x, y in zip((0, 1), (1, 0))])
        with raises(WorkerInterrupt):
            Parallel(n_jobs=2)(
                [delayed(interrupt_raiser)(x) for x in (1, 0)])

        # Try again with the context manager API
        with Parallel(n_jobs=2) as parallel:
            assert parallel._backend._pool is not None
            original_pool = parallel._backend._pool

            with raises(JoblibException):
                parallel([delayed(division)(x, y)
                          for x, y in zip((0, 1), (1, 0))])

            # The managed pool should still be available and be in a working
            # state despite the previously raised (and caught) exception
            assert parallel._backend._pool is not None

            # The pool should have been interrupted and restarted:
            assert parallel._backend._pool is not original_pool

            assert ([f(x, y=1) for x in range(10)] ==
                    parallel(delayed(f)(x, y=1) for x in range(10)))

            original_pool = parallel._backend._pool
            with raises(WorkerInterrupt):
                parallel([delayed(interrupt_raiser)(x) for x in (1, 0)])

            # The pool should still be available despite the exception
            assert parallel._backend._pool is not None

            # The pool should have been interrupted and restarted:
            assert parallel._backend._pool is not original_pool

            assert ([f(x, y=1) for x in range(10)] ==
                    parallel(delayed(f)(x, y=1) for x in range(10)))

        # Check that the inner pool has been terminated when exiting the
        # context manager
        assert parallel._backend._pool is None
    else:
        with raises(KeyboardInterrupt):
            Parallel(n_jobs=2)(
                [delayed(interrupt_raiser)(x) for x in (1, 0)])

    # wrapped exceptions should inherit from the class of the original
    # exception to make it easy to catch them
    with raises(ZeroDivisionError):
        Parallel(n_jobs=2)(
            [delayed(division)(x, y) for x, y in zip((0, 1), (1, 0))])

    with raises(MyExceptionWithFinickyInit):
        Parallel(n_jobs=2, verbose=0)(
            (delayed(exception_raiser)(i, custom_exception=True)
             for i in range(30)))

    try:
        # JoblibException wrapping is disabled in sequential mode:
        ex = JoblibException()
        Parallel(n_jobs=1)(
            delayed(division)(x, y) for x, y in zip((0, 1), (1, 0)))
    except Exception as ex:
        assert not isinstance(ex, JoblibException)


class Counter(object):
    def __init__(self, list1, list2):
        self.list1 = list1
        self.list2 = list2

    def __call__(self, i):
        self.list1.append(i)
        assert len(self.list1) == len(self.list2)


def consumer(queue, item):
    queue.append('Consumed %s' % item)


def check_dispatch_one_job(backend):
    """ Test that with only one job, Parallel does act as a iterator.
    """
    queue = list()

    def producer():
        for i in range(6):
            queue.append('Produced %i' % i)
            yield i

    # disable batching
    Parallel(n_jobs=1, batch_size=1, backend=backend)(
        delayed(consumer)(queue, x) for x in producer())
    assert queue == [
        'Produced 0', 'Consumed 0',
        'Produced 1', 'Consumed 1',
        'Produced 2', 'Consumed 2',
        'Produced 3', 'Consumed 3',
        'Produced 4', 'Consumed 4',
        'Produced 5', 'Consumed 5',
    ]
    assert len(queue) == 12

    # empty the queue for the next check
    queue[:] = []

    # enable batching
    Parallel(n_jobs=1, batch_size=4, backend=backend)(
        delayed(consumer)(queue, x) for x in producer())
    assert queue == [
        # First batch
        'Produced 0', 'Produced 1', 'Produced 2', 'Produced 3',
        'Consumed 0', 'Consumed 1', 'Consumed 2', 'Consumed 3',

        # Second batch
        'Produced 4', 'Produced 5', 'Consumed 4', 'Consumed 5',
    ]
    assert len(queue) == 12


def test_dispatch_one_job():
    for backend in BACKENDS:
        check_dispatch_one_job(backend)


def check_dispatch_multiprocessing(backend):
    """ Check that using pre_dispatch Parallel does indeed dispatch items
        lazily.
    """
    if mp is None:
        raise SkipTest()
    manager = mp.Manager()
    queue = manager.list()

    def producer():
        for i in range(6):
            queue.append('Produced %i' % i)
            yield i

    Parallel(n_jobs=2, batch_size=1, pre_dispatch=3, backend=backend)(
        delayed(consumer)(queue, 'any') for _ in producer())

    # Only 3 tasks are dispatched out of 6. The 4th task is dispatched only
    # after any of the first 3 jobs have completed.
    first_four = list(queue)[:4]
    # The the first consumption event can sometimes happen before the end of
    # the dispatching, hence, pop it before introspecting the "Produced" events
    first_four.remove('Consumed any')
    assert first_four == ['Produced 0', 'Produced 1', 'Produced 2']
    assert len(queue) == 12


def test_dispatch_multiprocessing():
    for backend in BACKENDS:
        if backend != "sequential":
            check_dispatch_multiprocessing(backend)


def test_batching_auto_threading():
    # batching='auto' with the threading backend leaves the effective batch
    # size to 1 (no batching) as it has been found to never be beneficial with
    # this low-overhead backend.

    with Parallel(n_jobs=2, batch_size='auto', backend='threading') as p:
        p(delayed(id)(i) for i in range(5000))  # many very fast tasks
        assert p._backend.compute_batch_size() == 1


def test_batching_auto_multiprocessing():
    with Parallel(n_jobs=2, batch_size='auto', backend='multiprocessing') as p:
        p(delayed(id)(i) for i in range(5000))  # many very fast tasks

        # It should be strictly larger than 1 but as we don't want heisen
        # failures on clogged CI worker environment be safe and only check that
        # it's a strictly positive number.
        assert p._backend.compute_batch_size() > 0


def test_exception_dispatch():
    "Make sure that exception raised during dispatch are indeed captured"
    with raises(ValueError):
        Parallel(n_jobs=2, pre_dispatch=16, verbose=0)(
            delayed(exception_raiser)(i) for i in range(30))


def test_nested_exception_dispatch():
    # Ensure TransportableException objects for nested joblib cases gets
    # propagated.
    with raises(JoblibException):
        Parallel(n_jobs=2, pre_dispatch=16, verbose=0)(
            delayed(SafeFunction(exception_raiser))(i) for i in range(30))


def _reload_joblib():
    # Retrieve the path of the parallel module in a robust way
    joblib_path = Parallel.__module__.split(os.sep)
    joblib_path = joblib_path[:1]
    joblib_path.append('parallel.py')
    joblib_path = '/'.join(joblib_path)
    module = __import__(joblib_path)
    # Reload the module. This should trigger a fail
    reload(module)


def test_multiple_spawning():
    # Test that attempting to launch a new Python after spawned
    # subprocesses will raise an error, to avoid infinite loops on
    # systems that do not support fork
    if not int(os.environ.get('JOBLIB_MULTIPROCESSING', 1)):
        raise SkipTest()
    with raises(ImportError):
        Parallel(n_jobs=2, pre_dispatch='all')(
            [delayed(_reload_joblib)() for i in range(10)])


class FakeParallelBackend(SequentialBackend):
    """Pretends to run concurrently while running sequentially."""

    def configure(self, n_jobs=1, parallel=None, **backend_args):
        self.n_jobs = self.effective_n_jobs(n_jobs)
        self.parallel = parallel
        return n_jobs

    def effective_n_jobs(self, n_jobs=1):
        if n_jobs < 0:
            n_jobs = max(mp.cpu_count() + 1 + n_jobs, 1)
        return n_jobs


def test_invalid_backend():
    with raises(ValueError):
        Parallel(backend='unit-testing')


def test_register_parallel_backend():
    try:
        register_parallel_backend("test_backend", FakeParallelBackend)
        assert "test_backend" in BACKENDS
        assert BACKENDS["test_backend"] == FakeParallelBackend
    finally:
        del BACKENDS["test_backend"]


def test_overwrite_default_backend():
    assert _active_backend_type() == MultiprocessingBackend
    try:
        register_parallel_backend("threading", BACKENDS["threading"],
                                  make_default=True)
        assert _active_backend_type() == ThreadingBackend
    finally:
        # Restore the global default manually
        parallel.DEFAULT_BACKEND = 'multiprocessing'
    assert _active_backend_type() == MultiprocessingBackend


def check_backend_context_manager(backend_name):
    with parallel_backend(backend_name, n_jobs=3):
        active_backend, active_n_jobs = parallel.get_active_backend()
        assert active_n_jobs == 3
        assert effective_n_jobs(3) == 3
        p = Parallel()
        assert p.n_jobs == 3
        if backend_name == 'multiprocessing':
            assert type(active_backend) == MultiprocessingBackend
            assert type(p._backend) == MultiprocessingBackend
        elif backend_name == 'threading':
            assert type(active_backend) == ThreadingBackend
            assert type(p._backend) == ThreadingBackend
        elif backend_name.startswith('test_'):
            assert type(active_backend) == FakeParallelBackend
            assert type(p._backend) == FakeParallelBackend


@with_multiprocessing
def test_backend_context_manager():
    all_test_backends = ['test_backend_%d' % i for i in range(3)]
    for test_backend in all_test_backends:
        register_parallel_backend(test_backend, FakeParallelBackend)
    all_backends = ['multiprocessing', 'threading'] + all_test_backends

    try:
        assert _active_backend_type() == MultiprocessingBackend
        # check that this possible to switch parallel backends sequentially
        for test_backend in all_backends: