from builtins import range
from builtins import object
__contributor__="Lisandro Dalcín"
""" MPIPool wrapped using mpi4py """
#import mpi4py
#mpi4py.rc.threaded = False
from mpi4py import MPI

class MPIPool(object):

    def __init__(self, comm=None, master=0):
        self.comm = MPI.COMM_WORLD if comm is None else comm
        self.master = master
        self.workers = set(range(self.comm.size))
        self.workers.discard(self.master)

    def is_master(self):
        return self.master == self.comm.rank

    def is_worker(self):
        return self.comm.rank in self.workers

    def map(self, function, iterable):
        assert self.is_master()

        comm = self.comm
        workerset = self.workers.copy()
        tasklist = [(tid, (function, arg)) for tid, arg in enumerate(iterable)]
        resultlist = [None] * len(tasklist)
        pending = len(tasklist)

        while pending:

            if workerset and tasklist:
                worker = workerset.pop()
                taskid, task = tasklist.pop()
                comm.send(task, dest=worker, tag=taskid)

            if tasklist:
                flag = comm.Iprobe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG)
                if not flag: continue
            else:
                comm.Probe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG)

            status = MPI.Status()
            result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            worker = status.source
            workerset.add(worker)
            taskid = status.tag
            resultlist[taskid] = result
            pending -= 1

        return resultlist

    def start(self):
        if not self.is_worker(): return
        comm = self.comm
        master = self.master
        status = MPI.Status()
        while True:
            task = comm.recv(source=master, tag=MPI.ANY_TAG, status=status)
            if task is None: break
            function, arg = task
            result = function(arg)
            comm.ssend(result, master, status.tag)

    def close(self):
        if not self.is_master(): return
        for worker in self.workers:
            self.comm.send(None, worker, 0)

