import numpy as np
from mpi4py import MPI
import time

def sequential_matrix_multiplication(A, B):
    return np.dot(A, B)

def print_matrix(matrix):
    for row in matrix:
        for element in row:
            print("{:.2f}".format(element), end="\t")
        print()

def generate_random_matrix(size):
    return np.random.rand(size, size)

if __name__ == "__main__":
    # Step 1: Matrix Generation
    matrix_sizes = [100, 200, 400, 800]
    A_matrices = []
    B_matrices = []

    for size in matrix_sizes:
        A = generate_random_matrix(size)
        B = generate_random_matrix(size)
        A_matrices.append(A)
        B_matrices.append(B)

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

    for i in range(len(matrix_sizes)):
        N = matrix_sizes[i]
        A = A_matrices[i]
        B = B_matrices[i]

        if rank == 0:
            print(f"Matrix size: {N} x {N}")

        comm.Barrier()

        if rank == 0:
            # Step 2: Sequential Matrix Multiplication
            start_time = time.time()
            C_seq = sequential_matrix_multiplication(A, B)
            end_time = time.time()

            if rank == 0:
                print("Sequential Matrix Multiplication:")
                print_matrix(C_seq)
                print(f"Execution Time: {end_time - start_time} seconds")
                print()

        comm.Barrier()

        if rank == 0:
            print(f"Parallel Matrix Multiplication using {size} processes:")

        comm.Barrier()

        # Step 3: MPI Scatter
        local_A = np.zeros((N // size, N), dtype=np.float64)
        comm.Scatter(A, local_A, root=0)

        if rank == 0:
            print(f"Process {rank} - Local A:")
            print_matrix(local_A)
            print()

        local_B = np.zeros((N, N // size), dtype=np.float64)
        comm.Scatter(B, local_B, root=0)

        if rank == 0:
            print(f"Process {rank} - Local B:")
            print_matrix(local_B)
            print()

        comm.Barrier()

        # Step 4: Matrix Multiplication
        local_C = np.dot(local_A, local_B)

        if rank == 0:
            print(f"Process {rank} - Local C:")
            print_matrix(local_C)
            print()

        comm.Barrier()

        # Step 5: MPI Reduce
        C = np.zeros((N, N), dtype=np.float64)
        comm.Reduce(local_C, C, op=MPI.SUM, root=0)

        if rank == 0:
            print("Parallel Matrix Multiplication:")
            print_matrix(C)
            print()

        comm.Barrier()

    if rank == 0:
        # Step 6: Performance Analysis
        num_processes = [2, 4, 6, 8]

        for i in range(len(matrix_sizes)):
            N = matrix_sizes[i]
            A = A_matrices[i]
            B = B_matrices[i]

            print(f"Matrix size: {N} x {N}")

            for num_procs in num_processes:
                if num_procs > size:
                    continue

                comm.Barrier()

                print(f"Parallel Matrix Multiplication using {num_procs} processes:")

                comm.Barrier()

                start_time = time.time()
                C_parallel = parallel_matrix_multiplication(A, B, num_procs, comm)
                end_time = time.time()

                if rank == 0:
                    print("Parallel Matrix Multiplication:")
                    print_matrix(C_parallel)
                    print(f"Execution Time: {end_time - start_time} seconds")
                    print()

                comm.Barrier()