'''Module provides functions for the user to turn on and off parallel processing,
and for other modules to access a pool of worker processes when parrallel processing
is turned on. Parallel processing is disabled by default but can be turned on using
enable_multiprocessing().  

When enabled, parallel processing will be automatically used by model fitting and
permutation testing functions. As parralel processing  incurs overheads, speedup will
be limited or even negative on machines with few threads and for small jobs. 
Multiprocessing is always used at outermost loop at which it is supported.

Each parallel process has different sequences generated by numpy.random.'''

from multiprocessing import pool, cpu_count
import numpy as np
import itertools
import builtins

global enabled, _mp_pool, _pool_available, _n_proc
_n_proc    = None
_mp_pool   = None
_pool_available = False
enabled = False


def enable_multiprocessing(n_processes=None):
    global enabled, _mp_pool, _pool_available, _n_proc
    if n_processes == None: n_processes = cpu_count()-1
    if _mp_pool: _mp_pool.close()
    enabled = True
    _mp_pool = pool.Pool(processes=n_processes)
    _pool_available = True
    _n_proc = n_processes
        

def disable_multiprocessing():
    global enabled, _mp_pool, _pool_available, _n_proc
    if _mp_pool: _mp_pool.close()
    enabled = False
    _mp_pool = None
    _pool_available = False
    _n_proc = None


class _rand_seed_wrap:
    '''Callable object which first seeds the numpy random number generator then behaves 
    as func, used to avoid parallel processed functions all using the same random sequences.'''
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        np.random.seed()
        return self.func(*args, **kwargs)


def map(function, iterable):
    '''Map function over iterable, return list.'''
    global _mp_pool, _pool_available, _n_proc
    if _pool_available:
        _pool_available = False
        out = _mp_pool.map(_rand_seed_wrap(function), iterable)
        #out = _mp_pool.map(function, iterable)
        _pool_available = True
        return out
    else:
        return list(builtins.map(function, iterable))


def starmap(function, iterable):
    '''Like map() except that the elements of the iterable are expected to be 
    iterables that are unpacked as arguments.'''
    global _mp_pool, _pool_available, _n_proc
    if _pool_available:
        _pool_available = False
        out = _mp_pool.starmap(_rand_seed_wrap(function), iterable)
        _pool_available = True
        return out
    else:
        return list(itertools.starmap(function, iterable))
        

def imap(function, iterable, ordered = True):
    '''Map function over iterable, return iterator.  If ordered is set to False, order of
    returned results is not preserved and parallel processing will be faster.'''
    global enabled, _mp_pool, _pool_available
    if _pool_available: 
        _pool_available = False
        if ordered:
            gen = _mp_pool.imap(_rand_seed_wrap(function), iterable)
        else:
            gen = _mp_pool.imap_unordered(_rand_seed_wrap(function), iterable)
    else:
        gen = builtins.map(function, iterable)
    for i in gen:
        yield i
    if enabled: _pool_available = True 


