import numpy as np

import time

from mpi4py import MPI

def generate_matrices(N):

    A = np.random.randint(0, 10, (N, N))

    B = np.random.randint(0, 10, (N, N))

    return A, B

def sequential_matrix_multiply(A, B):

    return np.dot(A, B)

def parallel_matrix_multiply(A, B):

    comm = MPI.COMM_WORLD

    rank = comm.Get_rank()

    size = comm.Get_size()

    N = A.shape[0]

    # Scatter matrix A to all other processes

    local_A = np.zeros((N // size, N), dtype=np.int)

    local_A = np.array_split(A, size)[rank]

    comm.Scatter(A, local_A, root=0)

    # Scatter matrix B to all processes

    local_B = np.zeros((N, N), dtype=np.int)

    comm.Scatter(B, local_B, root=0)

    # Perform local matrix multiplication

    local_C = np.dot(local_A, local_B)

    # Gather the results from all processes

    global_C = np.zeros((N, N), dtype=np.int)

    comm.Allgather(local_C, global_C)

    return global_C

# Step 1: Matrix Generation

N = 100  # Change the matrix size as needed

A, B = generate_matrices(N)

# Step 2: Sequential Matrix Multiplication

start_time = time.time()

sequential_result = sequential_matrix_multiply(A, B)

sequential_time = time.time() - start_time

# Step 3: MPI Initialization

comm = MPI.COMM_WORLD

rank = comm.Get_rank()

size = comm.Get_size()

# Print the received portion of matrix A on each process

print(f"Process {rank} received matrix A:\n{A}")

# Print the received portion of matrix B on each process

print(f"Process {rank} received matrix B:\n{B}")

# Step 4: Matrix Multiplication

start_time = time.time()

parallel_result = parallel_matrix_multiply(A, B)

parallel_time = time.time() - start_time

# Step 5: MPI Reduce

comm.Reduce(sequential_result, parallel_result, op=MPI.SUM, root=0)

# Step 6: Performance Analysis

if rank == 0:

    print(f"Sequential time: {sequential_time} seconds")

    print(f"Parallel time: {parallel_time} seconds")

    print("Sequential result:")

    print(sequential_result)

    print("Parallel result:")

    print(parallel_result)