from mpi4py import MPI
import numpy as np
from time import time

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

Start = time()

'''
    Approximation of the function f(x) using the extended trapezoid method, distributed on more processes.
    It is required to send the parameter of intervals to number of intervals + 1 because we require the rank 0 process
    to distribute the data and intervals to other processes using point to point blocking communication.


    Parameters
    ----------
    f : function
        Vectorized function of a single variable
    a , b : numbers
        Interval of integration [a,b]
    N : integer
        Number of sub intervals of [a,b]
    x : Array
        Return evenly spaced numbers over a specified interval.
        numpy.linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0)
        N+1 points make N sub intervals
    arrayX : Array
        Converts array x into a dtype=np.float64.
    dx : float
        Lambda, calculates the required length of each array.
    k : Array
        Used to temporary store received date from other processes.
    T : Array
        Approximation of the function using sums of each result from the other processes.
    y : Array
        Inputs f into array y.
    a : Array
        Temporary array used to store intervals and their data.
    y_left : Array
        Left endpoints.
    y_right: Array
        Right endpoints.
    G : float
        Function used to calculate the approximation using the trapezoid method.
    r : Dictionary
        set of key(i): value pairs(receive massage)
    s : Dictionary
        set of key(i): value pairs(sent massage)
    r1 : Dictionary
        set of key(rank): value pairs(receive massage)
    s1 : Dictionary
        set of key(rank): value pairs(sent massage)
    Start: float
        at the start of program function time() returns the time in seconds since the epoch as a floating point number.
    End: float
        at the end of program function time() returns the time in seconds since the epoch as a floating point number.

    Returns
    -------
    Array
        Approximation of the function of f(x) from a to b using the
        trapezoid rule with N sub intervals of equal length.
'''
a = 0
b = 6
N = size - 1
f = np.square
x = np.linspace(a, b, N + 1)
arrayX = np.array(x, dtype=np.float64)
dx = (b - a) / N

r = {}
s = {}

r1 = {}
s1 = {}

#ts = {}
#tr = {}

if rank == 0:
    print(f(arrayX))
    k = np.zeros(1, dtype=np.float64)
    T = np.zeros(1, dtype=np.float64)
    for i in range(0, size - 1):
        y = f(arrayX[i:i + 2])
        print(f(arrayX[i:i + 2]))
        s[i] = comm.Isend(y, dest=i + 1)
        #s[i].Wait()
    for i in range(0, size -1):
        r[i] = comm.irecv(k, source=i + 1)
        r[i].Wait()
        T = T + k

    End = time()
    print("Result: ", T[0], "Time: ", End - Start)
    #print("Timing starts at: ", Start, "and ends at: ", End))

elif rank != 0:
    #ts[rank] = time()
    a = np.zeros(2, dtype=np.float64)
    r[rank] = comm.irecv(a, source=0)
    r[rank].Wait()
    y_right = a[0:1]  # right endpoint
    y_left = a[1:2]  # left endpoint
    G = (dx / 2) * np.sum(y_right + y_left)
    s[rank] = comm.Isend(G, dest=0)
    #s[rank].Wait()
    #tr[rank] = time()
    #print("Rank ", rank, " Start: ", ts[rank], " End: ", tr[rank], " Time: ", tr[rank]-ts[rank])
else:
    exit()