#Brian Refsdal's parallel_map, from astropython.org
#Not sure what license this is released under, but until I know better:
#
#Copyright (c) 2010, Brian Refsdal
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are
#met:
#
#1. Redistributions of source code must retain the above copyright
#notice, this list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright
#notice, this list of conditions and the following disclaimer in the
#documentation and/or other materials provided with the distribution.
#
#3. The name of the author may not be used to endorse or promote
#products derived from this software without specific prior written
#permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
#A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import platform

import numpy

_multi=False
_ncpus=1

try:
  # May raise ImportError
  import multiprocessing
  _multi=True

  # May raise NotImplementedError
  _ncpus = multiprocessing.cpu_count()
except:
  pass

try:
    import tqdm
    _TQDM_LOADED= True
except ImportError: #pragma: no cover
    _TQDM_LOADED= False

__all__ = ('parallel_map',)


def worker(f, ii, chunk, out_q, err_q, lock,
           progressbar, iter_elapsed, tot_iter, pbar_proc):
  """
  A worker function that maps an input function over a
  slice of the input iterable.

  :param f  : callable function that accepts argument from iterable
  :param ii  : process ID
  :param chunk: slice of input iterable
  :param out_q: thread-safe output queue
  :param err_q: thread-safe queue to populate on exception
  :param lock : thread-safe lock to protect a resource
         ( useful in extending parallel_map() )
  :param progressbar: if True, display a progress bar
  :param iter_elapsed: shared-memory Value to track progress
  :param tot_iter: total number of iterations (for progressbar)
  :param pbar_proc: process to use to display the progressbar
  """
  vals = []

  progressbar*= _TQDM_LOADED
  if progressbar and ii == pbar_proc:
    pbar= tqdm.tqdm(total=tot_iter,leave=False)

  # iterate over slice
  for val in chunk:
    try:
      result = f(val)
    except Exception as e:
      err_q.put(e)
      return
    # Update progress bar; only update in first process, accumulate in others
    if progressbar:
      if ii == pbar_proc:
        pbar.update(iter_elapsed.value+1)
        iter_elapsed.value= 0
      else:
        iter_elapsed.value+= 1

    vals.append(result)

  if progressbar and ii == pbar_proc:
    pbar.close()

  # output the result and task ID to output queue
  out_q.put( (ii, vals) )


def run_tasks(procs, err_q, out_q, num):
  """
  A function that executes populated processes and processes
  the resultant array. Checks error queue for any exceptions.

  :param procs: list of Process objects
  :param out_q: thread-safe output queue
  :param err_q: thread-safe queue to populate on exception
  :param num : length of resultant array

  """
  # function to terminate processes that are still running.
  die = (lambda vals : [val.terminate() for val in vals
             if val.exitcode is None])

  try:
    for proc in procs:
      proc.start()

    for proc in procs:
      proc.join()

  except Exception as e:
    # kill all slave processes on ctrl-C
    try:
      die(procs)
    finally:
      raise e

  if not err_q.empty():
    # kill all on any exception from any one slave
    try:
      die(procs)
    finally:
      raise err_q.get()

  # Processes finish in arbitrary order. Process IDs double
  # as index in the resultant array.
  results=[None]*num;
  while not out_q.empty():
    idx, result = out_q.get()
    results[idx] = result

  # Remove extra dimension added by array_split
  return list(numpy.concatenate(results))


def parallel_map(function, sequence, numcores=None, progressbar=False):
  """
  A parallelized version of the native Python map function that
  utilizes the Python multiprocessing module to divide and
  conquer sequence.

  parallel_map does not yet support multiple argument sequences.

  :param function: callable function that accepts argument from iterable
  :param sequence: iterable sequence
  :param numcores: number of cores to use
  :param progressbar: if True, display a progressbar using tqdm
  """
  if not callable(function):
    raise TypeError("input function '%s' is not callable" %
              repr(function))

  if not numpy.iterable(sequence):
    raise TypeError("input '%s' is not iterable" %
              repr(sequence))

  size = len(sequence)

  if not _multi or size == 1:
    return map(function, sequence)

  if numcores is None:
    numcores = _ncpus

  if platform.system() == 'Windows': # JB: don't think this works on Win
    return list(map(function,sequence))

  # Use fork-based parallelism (because spawn fails with pickling issues, #457)
  ctx= multiprocessing.get_context('fork')

  # Returns a started SyncManager object which can be used for sharing
  # objects between processes. The returned manager object corresponds
  # to a spawned child process and has methods which will create shared
  # objects and return corresponding proxies.
  manager = ctx.Manager()

  # Create FIFO queue and lock shared objects and return proxies to them.
  # The managers handles a server process that manages shared objects that
  # each slave process has access to. Bottom line -- thread-safe.
  out_q = manager.Queue()
  err_q = manager.Queue()
  lock = manager.Lock()

  # if sequence is less than numcores, only use len sequence number of
  # processes
  if size < numcores:
    numcores = size

  # group sequence into numcores-worth of chunks
  sequence = numpy.array_split(sequence, numcores)

  # For progressbar: shared-memory variable to track iterations elapsed
  # in non-displaying processes
  iter_elapsed= multiprocessing.Value('i',0)

  procs = [ctx.Process(target=worker,
           args=(function, ii, chunk, out_q, err_q, lock,
                 progressbar,iter_elapsed,size,numcores-1))
         for ii, chunk in enumerate(sequence)]

  return run_tasks(procs, err_q, out_q, numcores)


if __name__ == "__main__":
  """
  Unit test of parallel_map()

  Create an arbitrary length list of references to a single
  matrix containing random floats and compute the eigenvals
  in serial and parallel. Compare the results and timings.
  """

  import time

  numtasks = 5
  #size = (1024,1024)
  size = (512,512)

  vals = numpy.random.rand(*size)
  f = numpy.linalg.eigvals

  iterable = [vals]*numtasks

  print ('Running numpy.linalg.eigvals %iX on matrix size [%i,%i]' %
      (numtasks,size[0],size[1]))

  tt = time.time()
  presult = parallel_map(f, iterable)
  print('parallel map in %g secs' % (time.time()-tt))

  tt = time.time()
  result = map(f, iterable)
  print('serial map in %g secs' % (time.time()-tt))

  assert (numpy.asarray(result) == numpy.asarray(presult)).all()
