#include <iostream>
#include <cstdlib>
#include <iomanip>
#include <mpi.h>

using namespace std;

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    const int NRA = 4, NCA = 3, NCB = 2; // Example dimensions of matrices
    const int MASTER = 0;
    double start_time;
    start_time = MPI_Wtime();

    double *a = nullptr;
    double *b = new double[NCA * NCB];
    double *c = new double[NRA * NCB];

    if (rank == MASTER) {
        cout << "Number of processes: " << size << endl;
        cout << "Matrix A" << endl;
        a = new double[NRA * NCA];
        for (int i = 0; i < NRA; i++) {
            for (int j = 0; j < NCA; j++) {
                a[i * NCA + j] = 1 + i + j; // Example initialization
                cout << setw(5) << a[i * NCA + j] << " ";
            }
            cout << endl;
        }

        cout << "Matrix B" << endl;
        for (int i = 0; i < NCA; i++) {
            for (int j = 0; j < NCB; j++) {
                b[i * NCB + j] = 1 + i * j; // Example initialization
                cout << setw(5) << b[i * NCB + j] << " ";
            }
            cout << endl;
        }
    }

    // Broadcast Matrix B to all processes
    MPI_Bcast(b, NCA * NCB, MPI_DOUBLE, MASTER, MPI_COMM_WORLD);

    // Allocate memory for Matrix A on all processes
    if (rank != MASTER) {
        a = new double[NRA * NCA];
    }

    // Broadcast Matrix A to all processes
    MPI_Bcast(a, NRA * NCA, MPI_DOUBLE, MASTER, MPI_COMM_WORLD);

    // Perform matrix multiplication
    for (int i = 0; i < NRA; i++) {
        for (int j = 0; j < NCB; j++) {
            c[i * NCB + j] = 0;
            for (int k = 0; k < NCA; k++) {
                c[i * NCB + j] += a[i * NCA + k] * b[k * NCB + j];
            }
        }
    }

    // Master process prints the resulting matrix C
    if (rank == MASTER) {
        cout << "Matrix C" << endl;
        for (int i = 0; i < NRA; i++) {
            for (int j = 0; j < NCB; j++) {
                cout << setw(5) << c[i * NCB + j] << " ";
            }
            cout << endl;
        }
    }

    // Free allocated memory
    delete[] b;
    delete[] c;
    if (rank == MASTER) {
        delete[] a;
    }

    MPI_Finalize();
    return 0;
}