"""
* Experimental *

Like the map function, but can use a pool of threads.

Really easy to use threads.  eg.  tmap(f, alist)

If you know how to use the map function, you can use threads.
"""

__author__ = "Rene Dudfield"
__version__ = "0.3.0"
__license__ = "Python license"

from queue import Queue, Empty
import threading


Thread = threading.Thread

STOP = object()
FINISH = object()

# DONE_ONE = object()
# DONE_TWO = object()

# a default worker queue.
_wq = None

# if we are using threads or not.  This is the number of workers.
_use_workers = 0

# Set this to the maximum for the amount of Cores/CPUs
#    Note, that the tests early out.
#    So it should only test the best number of workers +2
MAX_WORKERS_TO_TEST = 64


def init(number_of_workers=0):
    """Does a little test to see if threading is worth it.
      Sets up a global worker queue if it's worth it.

    Calling init() is not required, but is generally better to do.
    """
    global _wq, _use_workers

    if number_of_workers:
        _use_workers = number_of_workers
    else:
        _use_workers = benchmark_workers()

    # if it is best to use zero workers, then use that.
    _wq = WorkerQueue(_use_workers)


def quit():
    """cleans up everything."""
    global _wq, _use_workers
    _wq.stop()
    _wq = None
    _use_workers = False


def benchmark_workers(a_bench_func=None, the_data=None):
    """does a little test to see if workers are at all faster.
    Returns the number of workers which works best.
    Takes a little bit of time to run, so you should only really call
      it once.
    You can pass in benchmark data, and functions if you want.
    a_bench_func - f(data)
    the_data - data to work on.
    """
    # TODO: try and make this scale better with slower/faster cpus.
    #  first find some variables so that using 0 workers takes about 1.0 seconds.
    #  then go from there.

    # note, this will only work with pygame 1.8rc3+
    # replace the doit() and the_data with something that releases the GIL

    import pygame
    import pygame.transform
    import time

    if not a_bench_func:

        def doit(x):
            return pygame.transform.scale(x, (544, 576))

    else:
        doit = a_bench_func

    if not the_data:
        thedata = [pygame.Surface((155, 155), 0, 32) for x in range(10)]
    else:
        thedata = the_data

    best = time.time() + 100000000
    best_number = 0
    # last_best = -1

    for num_workers in range(0, MAX_WORKERS_TO_TEST):
        wq = WorkerQueue(num_workers)
        t1 = time.time()
        for _ in range(20):
            print(f"active count:{threading.active_count()}")
            tmap(doit, thedata, worker_queue=wq)
        t2 = time.time()

        wq.stop()

        total_time = t2 - t1
        print(f"total time num_workers:{num_workers}: time:{total_time}:")

        if total_time < best:
            # last_best = best_number
            best_number = num_workers
            best = total_time

        if num_workers - best_number > 1:
            # We tried to add more, but it didn't like it.
            #   so we stop with testing at this number.
            break

    return best_number


class WorkerQueue:
    def __init__(self, num_workers=20):
        self.queue = Queue()
        self.pool = []
        self._setup_workers(num_workers)

    def _setup_workers(self, num_workers):
        """Sets up the worker threads
        NOTE: undefined behaviour if you call this again.
        """
        self.pool = []

        for _ in range(num_workers):
            self.pool.append(Thread(target=self.threadloop))

        for a_thread in self.pool:
            a_thread.setDaemon(True)
            a_thread.start()

    def do(self, f, *args, **kwArgs):
        """puts a function on a queue for running later."""
        self.queue.put((f, args, kwArgs))

    def stop(self):
        """Stops the WorkerQueue, waits for all of the threads to finish up."""
        self.queue.put(STOP)
        for thread in self.pool:
            thread.join()

    def threadloop(self):  # , finish=False):
        """Loops until all of the tasks are finished."""
        while True:
            args = self.queue.get()
            if args is STOP:
                self.queue.put(STOP)
                self.queue.task_done()
                break
            try:
                args[0](*args[1], **args[2])
            finally:
                # clean up the queue, raise the exception.
                self.queue.task_done()
                # raise

    def wait(self):
        """waits until all tasks are complete."""
        self.queue.join()


class FuncResult:
    """Used for wrapping up a function call so that the results are stored
    inside the instances result attribute.
    """

    def __init__(self, f, callback=None, errback=None):
        """f - is the function we that we call
        callback(result) - this is called when the function(f) returns
        errback(exception) - this is called when the function(f) raises
                               an exception.
        """
        self.f = f
        self.exception = None
        self.result = None
        self.callback = callback
        self.errback = errback

    def __call__(self, *args, **kwargs):
        # we try to call the function here.  If it fails we store the exception.
        try:
            self.result = self.f(*args, **kwargs)
            if self.callback:
                self.callback(self.result)
        except Exception as e:
            self.exception = e
            if self.errback:
                self.errback(self.exception)


def tmap(f, seq_args, num_workers=20, worker_queue=None, wait=True, stop_on_error=True):
    """like map, but uses a thread pool to execute.
    num_workers - the number of worker threads that will be used.  If pool
                    is passed in, then the num_workers arg is ignored.
    worker_queue - you can optionally pass in an existing WorkerQueue.
    wait - True means that the results are returned when everything is finished.
           False means that we return the [worker_queue, results] right away instead.
           results, is returned as a list of FuncResult instances.
    stop_on_error -
    """

    if worker_queue:
        wq = worker_queue
    else:
        # see if we have a global queue to work with.
        if _wq:
            wq = _wq
        else:
            if num_workers == 0:
                return map(f, seq_args)

            wq = WorkerQueue(num_workers)

    # we short cut it here if the number of workers is 0.
    # normal map should be faster in this case.
    if len(wq.pool) == 0:
        return map(f, seq_args)

    # print("queue size:%s" % wq.queue.qsize())

    # TODO: divide the data (seq_args) into even chunks and
    #       then pass each thread a map(f, equal_part(seq_args))
    #      That way there should be less locking, and overhead.

    results = []
    for sa in seq_args:
        results.append(FuncResult(f))
        wq.do(results[-1], sa)

    # wq.stop()

    if wait:
        # print("wait")
        wq.wait()
        # print("after wait")
        # print("queue size:%s" % wq.queue.qsize())
        if wq.queue.qsize():
            raise RuntimeError("buggy threadmap")
        # if we created a worker queue, we need to stop it.
        if not worker_queue and not _wq:
            # print("stopping")
            wq.stop()
            if wq.queue.qsize():
                um = wq.queue.get()
                if um is not STOP:
                    raise RuntimeError("buggy threadmap")

        # see if there were any errors.  If so raise the first one.  This matches map behaviour.
        # TODO: the traceback doesn't show up nicely.
        # NOTE: TODO: we might want to return the results anyway?  This should be an option.
        if stop_on_error:
            error_ones = list(filter(lambda x: x.exception, results))
            if error_ones:
                raise error_ones[0].exception

        return (x.result for x in results)
    return [wq, results]