from mpi4py import MPI
import sys
import time
import numpy

comm_to_function_wrapper = MPI.Comm.Get_parent()
rank = comm_to_function_wrapper.Get_rank()
size = comm_to_function_wrapper.Get_size()

x_length = int(sys.argv[1])

# Initialize and receive the starting point from function_wrapper.py
x = numpy.array(numpy.zeros(x_length))
comm_to_function_wrapper.Bcast(x,root=0) 

print x
print "Starting! The function evaluation of x=%s was given %d ranks. I am rank %d." % (str(x),size,rank)
local_y = (sum(x)**2)/size # Sum the numbers in x and divide by the number of ranks
print local_y
comm_to_function_wrapper.Reduce(local_y, None, op=MPI.SUM, root=0) # Give the value back to function_wrapper.py
print "Ending! The function evaluation of x=%s was given %d ranks. I am rank %d." % (str(x),size,rank)

comm_to_function_wrapper.Disconnect()