#
#  multiprocessing/vectorize.py
#  bxtorch
#
#  Created by Oliver Borchert on June 02, 2019.
#  Copyright (c) 2019 Oliver Borchert. All rights reserved.
#

import os
import numpy as np
import torch
import torch.multiprocessing as mp

class Vectorizer:
    """
    The Vectorizer class ought to be used in cases where a result tensor
    of size N is filled with values computed in some complex way. The
    computation of these N computations can then be parallelized over
    multiple processes.
    """

    # MARK: Initialization
    def __init__(self, worker_func, worker_init=None, callback_func=None,
                 num_workers=-1, **kwargs):
        """
        Initializes a new vectorizer.

        Parameters:
        -----------
        - worker_func: callable
            The function which receives as input an item of the input to process
            and outputs a value which ought to be returned.
        - worker_init: callable, default: None
            The function receives as input the rank of the worker (i.e. every
            time this function is called, it is called with a different integer
            as parameter). Its return values are passed as *last* parameters
            to ``worker_func`` upon every invocation.
        - callback_func: callable, default: None
            A function to call after every item has been processed. Must not
            need to be a free function as it is called on the main thread.
        - num_workers: int, default: -1
            The number of processes to use. If set to -1, it defaults to the
            number of available cores. If set to 0, everything is executed on
            the main thread.
        - kwargs: keyword arguments
            Additional arguments passed to the worker initialization function.
        """
        self.num_workers = os.cpu_count() if num_workers == -1 else num_workers
        self.worker_func = worker_func
        self.worker_init = worker_init
        self.callback_func = callback_func
        self.init_kwargs = kwargs

    # MARK: Instance Methods
    def process(self, items, *args):
        """
        Uses the vectorizer's worker function in order to process all items
        in parallel.

        Parameters:
        -----------
        - items: list of object or iterable of object
            The items to be processed by the workers. If given as iterable only
            (i.e. it does not support index access), the performance might
            suffer slightly due to an increased number of synchronizations.
        - args: variadic arguments
            Additional arguments passed directly to the worker function.

        Returns:
        --------
        - list of object
            The output generated by the worker function for each of the input
            items.
        """
        if self.num_workers == 0: # execute sequentially
            result = []
            init = _init_if_needed(self.worker_init, 0, **self.init_kwargs)
            all_args = _combine_args(args, init)
            for item in items:
                result.append(self.worker_func(item, *all_args))
            return result
        else:
            process_batches = hasattr(items, '__getitem__')
            
            if process_batches:
                return self._process_batches(items, *args)
            else:
                return self._process_consumers(items, *args)

    # MARK: Private Methods
    def _process_batches(self, items, *args):
        num_items = len(items)

        splits = np.array_split(np.arange(num_items), self.num_workers)
        splits = [0] + [a[-1] + 1 for a in splits]

        result = []

        processes = []
        queues = []
        done = mp.Event()
        if self.callback_func is not None:
            tick_queue = mp.Queue()
        else:
            tick_queue = None

        for i in range(self.num_workers):
            queue = mp.Queue()
            process = mp.Process(
                target=_batch_worker,
                args=(
                    queue, done, tick_queue, i, 
                    self.worker_func, self.worker_init,
                    self.init_kwargs, items[splits[i]:splits[i+1]],
                    *args
                )
            )
            process.start()
            processes.append(process)
            queues.append(queue)

        if self.callback_func is not None:
            for _ in range(num_items):
                tick_queue.get()
                self.callback_func()

        for i, q in enumerate(queues):
            result.extend(q.get())
            q.close()

        done.set()

        for p in processes:
            p.join()

        return result

    def _process_consumers(self, items, *args):
        result = []

        processes = []
        push_queue = mp.Queue()
        pull_queue = mp.Queue()

        for i in range(self.num_workers):
            process = mp.Process(
                target=_consumer_worker,
                args=(push_queue, pull_queue, i,
                      self.worker_func, self.worker_init, self.init_kwargs, 
                      *args)
            )
            process.start()
            processes.append(process)

        iterator = iter(items)
        index = 0
        expect = 0
        try:
            for _ in range(self.num_workers):
                item = next(iterator)
                expect += 1
                push_queue.cancel_join_thread()
                push_queue.put((index, item))
                index += 1

            while True:
                result.append(pull_queue.get())
                self.callback_func()
                expect -= 1
                item = next(iterator)
                expect += 1
                push_queue.cancel_join_thread()
                push_queue.put((index, item))
                index += 1

        except StopIteration:
            for _ in range(expect):
                result.append(pull_queue.get())
                self.callback_func()

            pull_queue.close()

            for _ in range(len(processes)):
                push_queue.cancel_join_thread()
                push_queue.put(None)

            push_queue.close()

            for p in processes:
                p.join()

        return [r[1] for r in sorted(result, key=lambda r:r[0])]


def _batch_worker(push_queue, done, tick_queue, rank,
                  worker_func, worker_init, init_kwargs, items, *args):

    init = _init_if_needed(worker_init, rank, **init_kwargs)
    all_args = _combine_args(args, init)

    result = []

    for i, item in enumerate(items):
        result.append(worker_func(item, *all_args))
        if tick_queue is not None:
            tick_queue.cancel_join_thread()
            tick_queue.put(None)

    push_queue.cancel_join_thread()
    push_queue.put(result)
    done.wait()


def _consumer_worker(pull_queue, push_queue, rank, 
                     worker_func, worker_init, init_kwargs, *args):

    init = _init_if_needed(worker_init, rank, **init_kwargs)
    all_args = _combine_args(args, init)

    while True:
        item = pull_queue.get()
        if item is None:
            break

        idx, item = item
        result = worker_func(item, *all_args)

        push_queue.cancel_join_thread()
        push_queue.put((idx, result))
    

def _init_if_needed(init, rank, **kwargs):
    if init is None:
        return None
    else:
        return init(rank, **kwargs)


def _combine_args(a, b):
    if b is None:
        return a
    elif isinstance(b, tuple):
        return a + b
    else:
        return a + (b,)
