from mpi4py import MPI
import math


def f(x):
  return x*x
# end def


def trapezoidal(a, b, n, h):

  s = 0.0
  s += h * f(a)
  for i in range(1, n):
    s += 2.0 * h * f(a + i*h)
  # end for
  s += h * f(b)
  return (s/2.)
# end def


# Main section
comm = MPI.COMM_WORLD
my_rank = comm.Get_rank()

p = comm.Get_size()

if (my_rank == 0):
  print("Number of ranks = ",p)
# end if
print('my_rank = ',my_rank)


# Integral parameters
a=0.0        # Left endpoint
b=2.0        # Right endpoint
n=1024       # Number of trapezoids
dest=0       # Destination for messages (rank 0 process)
total=-1.0   # Initialize to negative number

h = (b-a)/n         #  h = Trapezoid Base Length - the same for all processes
local_n = int(n/p)  #  Number of trapezoids for each process

#   Length of each process' interval of
#   integration = local_n*h.
local_a = a + my_rank*local_n*h   # local a (specific to each process)
local_b = local_a + local_n*h     # local b (specific to each process)

# Each rank performs its own integration
integral = trapezoidal(local_a, local_b, local_n, h)

# Add up the integrals calculated by each process
if (my_rank == 0):
  total=integral
  for source in range(1,p):
    integral=comm.recv(source=source)
    print("PE ",my_rank,"<-",source,",",integral,"\n")
    total=total+integral
  # end for
else:
  print("PE ",my_rank,"->",dest,",",integral,"\n")
  comm.send(integral, dest=0)
# end if


#  Print the result
if (my_rank==0):
  print("**With n=",n,", trapezoids,")
  print("** Final integral from",a,"to",b,"=",total,"\n")
# end if

MPI.Finalize()