from timeit import default_timer as timer
import numpy
from mpi4py import MPI 


iterations = 10
m = 1024
n= 1024
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

my_m = numpy.array(m/size, dtype='int')
if rank is 0:
    total_grid = numpy.zeros((m+2,n+2), dtype='double')
    total_grid[1,1:n+1] = list(range(n))
    total_grid[1:m+1,1] = list(range(m))

grid1 = numpy.zeros((my_m+2,n+2), dtype='double')
grid2 = numpy.zeros((my_m+2,n+2), dtype='double')

t00=timer()
##distribute init Filed
if rank is 0:
    comm.Scatter( [total_grid[1:m+1][:], MPI.DOUBLE], [grid1[1:][:], MPI.DOUBLE], root=0)
else:
    comm.Scatter( None, [grid1[1:][:], MPI.DOUBLE], root=0)

top = rank -1
bottom = rank+1

#proc NULL mean communication will be not executed (for corners)
if rank is 0:
    top = MPI.PROC_NULL
if rank is size-1:
    bottom = MPI.PROC_NULL
req = {}

t1=timer()
scatter_time = t1-t00
#print grid
#exit()
for k in range(iterations+1):
#
# Add here code to exchange the halo before every iteration
#

    b = 0 
#    # start timer after a warmup iteration
    if k<1:
        t0 = timer()

    for j in range(1,my_m+1):
        for i in range(1,n+1):
            grid2[j][i] =0.5*grid1[j][i] + 0.5*( grid1[j-1][i]+
                                             grid1[j+1][i]+
                                             grid1[j][i-1]+
                                             grid1[j][i+1])
    tmp=grid2
    grid2=grid1
    grid1=tmp

t1 = timer()
pipeline_time = t1 - t0

t0 = timer()
## add here code to collect the result with Gater on proc 0

t1=timer()

gather_time = t1-t0
total_time = t1-t00

if rank is 0:
    avgtime = pipeline_time/iterations
    print('Rate (MFlops/s): ',1.e-6*2*(m-1)*(n-1)/avgtime,' Avg time (s): ',avgtime)
    print("Comm time", scatter_time, "Comm time 2", gather_time);
    print ("Total time ", total_time)
