from mpi4py import MPI
import sys
import timeit

from matrix_mpi import MatrixMPI as Matrix
from vector_mpi import VectorMPI as Vector

# from matrix import Matrix
# from vector import Vector

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

n = int(sys.argv[1])

m_data = [(i / k) for i in range(1, n+1) for k in range(1, n+1)]
A = Matrix(m_data, (n, n))
v = Vector(list(range(1, n+1)))

time = timeit.timeit(lambda: A * v, number=1)

if rank == 0:
    print(f"ranks = {size}: time = {time}s")