#!/usr/bin/env python
"""
For basic applications, MPI is as easy to use as any other 
message-passing system. The sample code below contains the complete 
communications skeleton for a dynamically load balanced master/slave 
application. Following the code is a description of the few functions 
necessary to write typical parallel applications.

important parameters
--------------------

status = MPI.Status()               # where all info is stored

# Receive results from a slave
result = comm.recv(                 # message buffer
         source=MPI.ANY_SOURCE,     # receive from any sender (-1)
         tag=MPI.ANY_TAG,           # any type of message (-1)
         status=status)             # info about the received msg (class)

# Send the slave a new work unit
comm.send(work,                     # message buffer
         dest=status.Get_source(),  # to whom we just received from
         tag=WORKTAG)               # user chosen message tag

"""

from mpi4py import MPI
import numpy as np

WORKTAG = 1
DIETAG = 0

class Work(object):
    def __init__(self):
	
        n_to_iterate = [[x] for x in range(100)]
    
        #CREATE QUEUE OF DATA TO SEND TO SLAVES
        self.work = n_to_iterate

    def get_next(self):
        if len(self.work) == 0:
            return None
        return self.work.pop(0)
 

def do_work(n):
    #arr = range(50000000)
    arr = np.arange(50000000)
    #mult = arr*n[0]
    print 'done %s'%(n[0])
    return


def master(comm):
    num_procs = comm.Get_size()
    status = MPI.Status()
    
    # generate work queue
    wq = Work()

    # Seed the slaves, send one unit of work to each slave (rank)
    for rank in xrange(1, num_procs):
        work = wq.get_next()
        print 'master sending work'
        if not work: break
        comm.send(work, dest=rank, tag=WORKTAG)
    
    # Loop over getting new work requests until there is no more work to be done
    while True:
        work = wq.get_next()
        if not work: break
    
        # Receive results from a slave
        result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        #process_result(result)

        # Send the slave a new work unit
        comm.send(work, dest=status.Get_source(), tag=WORKTAG)
    
    # No more work to be done, receive all outstanding results from slaves
    for rank in xrange(1, num_procs): 
        result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        del result
        #process_result(result)

    # Tell all the slaves to exit by sending an empty message with DIETAG
    for rank in xrange(1, num_procs):
        comm.send(0, dest=rank, tag=DIETAG)
        
    return


def slave(comm):
    my_rank = comm.Get_rank()
    status = MPI.Status()
    
    print 'slave %s online'%(my_rank)

    while True:
        # Receive a message from the master (the work to do)
        work = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
        print 'slave %s has work'%(my_rank)
        
        # Check the tag of the received message
        if status.Get_tag() == DIETAG: break 

        # Do the work and send it back
        result = do_work(work)
        comm.send(result, dest=0)
        

def main():
    comm = MPI.COMM_WORLD
    my_rank = comm.Get_rank()
    my_name = MPI.Get_processor_name()
    #comm.Barrier()
    #start = MPI.Wtime()
    
    if my_rank == 0:
        master(comm)
    else:
        slave(comm)
    
    #comm.Barrier()
    #end = MPI.Wtime()
    #print 'time:', end - start


if __name__ == '__main__':
    main()
