# -----------------------------------------------------------------------------

from array import array
from threading import Thread

from mpi4py import MPI


class Counter:
    def __init__(self, comm):
        # duplicate communicator
        assert not comm.Is_inter()
        self.comm = comm.Dup()
        # start counter thread
        self.thread = None
        rank = self.comm.Get_rank()
        if rank == 0:
            self.thread = Thread(target=self._counter_thread)
            self.thread.start()

    def _counter_thread(self):
        incr = array("i", [0])
        ival = array("i", [0])
        status = MPI.Status()
        while True:  # server loop
            self.comm.Recv(
                [incr, MPI.INT], MPI.ANY_SOURCE, MPI.ANY_TAG, status
            )
            if status.Get_tag() == 1:
                return
            self.comm.Ssend([ival, MPI.INT], status.Get_source(), 0)
            ival[0] += incr[0]

    def free(self):
        self.comm.Barrier()
        # stop counter thread
        rank = self.comm.Get_rank()
        if rank == 0:
            self.comm.Ssend([None, MPI.INT], 0, 1)
            self.thread.join()
        #
        self.comm.Free()

    def next(self):
        incr = array("i", [1])
        ival = array("i", [0])
        self.comm.Ssend([incr, MPI.INT], 0, 0)
        self.comm.Recv([ival, MPI.INT], 0, 0)
        nxtval = ival[0]
        return nxtval


# -----------------------------------------------------------------------------


def test_thread_level():
    import sys

    flag = MPI.Query_thread() == MPI.THREAD_MULTIPLE
    flag = MPI.COMM_WORLD.bcast(flag, root=0)
    if not flag:
        if MPI.COMM_WORLD.Get_rank() == 0:
            sys.stderr.write("MPI does not provide enough thread support\n")
        sys.exit(0)


def test():
    vals = []
    counter = Counter(MPI.COMM_WORLD)
    for _i in range(5):
        c = counter.next()
        vals.append(c)
    counter.free()

    vals = MPI.COMM_WORLD.allreduce(vals)
    assert sorted(vals) == list(range(len(vals)))


if __name__ == "__main__":
    test_thread_level()
    test()

# -----------------------------------------------------------------------------
