from mpi4py import MPI
import numpy as np
import math
import sys

def dot_product(a, x):
    result = 0
    for i in range(len(a)):
        result += a[i] * x[i]
    return result

def row(i):
    row = []
    for j in range(1, n+1):
        row.append(i / j)
    return row

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

n = int(sys.argv[1])

x = list(range(1, n+1))

# each rank should compute almost the same amount of matrix.row * x
# for this split the list [1,...,n+1] in sublists containing the rownumbers every rank has to compute
# if n is a multiple of size all ranks have the same amount to compute, if not, the first (n % size) ranks compute each one more
chunks = np.array_split(list(range(1,n+1)), size) 

chunk = chunks[rank]
sub_b = []
for i in chunk:
    sub_b.append(dot_product(row(i), x)) # every rank computes its delegated rows times x

comm.barrier()

b = comm.gather(sub_b)

if rank == 0:
    b = np.concatenate(b).tolist() # b is a list of lists, np.concatenate 'flattens' the list of lists into an np.ndarray and tolist() to get an python list