from mpi4py import MPI
import numpy as np
import math

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

chunk = 0

if rank == 0:
    N = 100000
    x = np.random.uniform(0,1,(N,1))

    chunksize = math.floor(N / size)
    for i in range(1,size):
        from_index = 0 + chunksize * i
        to_index = (chunksize * (i + 1)) - 1
        chunk = x[from_index:to_index]
        comm.send(chunk, dest = i)
else:
    chunk = comm.recv(source = 0)

chunk_sum = 0.0
for i in range(len(chunk)):
    chunk_sum += chunk[i] * chunk[i]

comm.Barrier()

sum = comm.reduce(chunk_sum, op = MPI.SUM, root = 0)

if rank == 0:
    print(f"The squared euclidian norm of x is {sum[0]:.2f}.")