from mpi4py import MPI
import random
import numpy as np
import time
N = 2000
comm = MPI.COMM_WORLD
rank = comm.Get_rank()

def pingpong(size):
    if rank is 0:
        array = [ [1] for i in range(size)]
        comm.barrier()
        start = time.time()
        for i in range(1,N):
            comm.send(array, dest=1, tag=i)
        end = time.time()
    if rank is 1:
        comm.barrier()
        start = time.time()
        for i in range(1,N):
            array = comm.recv(source=0, tag=i)

        end = time.time()
    return ((end-start)*1e6/N)

def pingpong_np1(size):
    if rank is 0:
        array = np.random.randint(0,size, size=size, dtype='int')
        comm.barrier()
        start = time.time()
        for i in range(1,N):
            comm.send(array, dest=1, tag=i)
        end = time.time()
    if rank is 1:
        comm.barrier()
        start = time.time()
        for i in range(1,N):
            array = comm.recv(source=0, tag=i)

        end = time.time()
    return ((end-start)*1e6/N)


def pingpong_np(size):
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    if rank is 0:
        comm.barrier()
        start = time.time()
        array = np.random.randint(0,size, size=size, dtype='int')
        for i in range(1,N):
            comm.Send([array, MPI.INT], dest=1, tag=i)
        end = time.time()
    if rank is 1:
        comm.barrier()
        start = time.time()
        array = np.zeros(size, dtype='int')
        for i in range(1,N):
            comm.Recv([array,MPI.INT], source=0, tag=i)

        end = time.time()
    return ((end-start)*1e6/N)



sizes = (2**exp for exp in range(0, 15))

for n in  sizes:
    lat = np.zeros(3, dtype="double")
    lat_sum = np.zeros(3, dtype="double")
    lat[0] = pingpong(n)
    lat[1] = pingpong_np1(n)
    lat[2] =pingpong_np(n)

    comm.Reduce([lat, MPI.DOUBLE], [lat_sum, MPI.DOUBLE],
            op=MPI.SUM, root=0)
    lat=lat_sum/2
    if rank is 0:
#        print ('%4d %4.2f %4.2f  %4.2f' % n, py_lat, np1_lat, np_lat)
        print('{0:6}\t{1:7.2f}\t{2:4.2f}\t{3:4.2f}'. format(n, lat[0], lat[1], lat[2]))
