#include <mpi.h>
#include <iostream>
using namespace std;

#include <vector>
#include <cstdlib>

using namespace std;

// Function to multiply matrices
void matrixMultiply(const vector<vector<int>>& A, const vector<vector<int>>& B, vector<vector<int>>& C, int startRow, int endRow, int N, int P) {
    for (int i = startRow; i < endRow; i++) {
        for (int j = 0; j < P; j++) {
            C[i][j] = 0;
            for (int k = 0; k < N; k++) {
                C[i][j] += A[i][k] * B[k][j];
            }
        }
    }
}

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);

    int M = 4, N = 4, P = 4; // Dimensions of matrices

    vector<vector<int>> A(M, vector<int>(N, 0));
    vector<vector<int>> B(N, vector<int>(P, 0));
    vector<vector<int>> C(M, vector<int>(P, 0));

    // Master initializes matrices
    if (rank == 0) {
        for (int i = 0; i < M; i++) {
            for (int j = 0; j < N; j++) {
                A[i][j] = rand() % 10; // Random values between 0 and 9
            }
        }
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < P; j++) {
                B[i][j] = rand() % 10;
            }
        }

        cout << "Matrix A:\n";
        for (const auto& row : A) {
            for (int val : row) cout << val << " ";
            cout << "\n";
        }

        cout << "\nMatrix B:\n";
        for (const auto& row : B) {
            for (int val : row) cout << val << " ";
            cout << "\n";
        }
        cout << "\n";
    }

    // Broadcast matrix B to all processes
    for (int i = 0; i < N; i++) {
        MPI_Bcast(B[i].data(), P, MPI_INT, 0, MPI_COMM_WORLD);
    }

    int rows_per_proc = M / size;
    vector<vector<int>> local_A(rows_per_proc, vector<int>(N));
    vector<vector<int>> local_C(rows_per_proc, vector<int>(P, 0));

    for (int i = 0; i < rows_per_proc; i++) {
        MPI_Scatter(A[i].data(), N, MPI_INT, local_A[i].data(), N, MPI_INT, 0, MPI_COMM_WORLD);
    }

    int startRow = rank * rows_per_proc;
    int endRow = startRow + rows_per_proc;
    matrixMultiply(A, B, C, startRow, endRow, N, P);

    for (int i = 0; i < rows_per_proc; i++) {
        MPI_Gather(local_C[i].data(), P, MPI_INT, C[startRow].data(), P, MPI_INT, 0, MPI_COMM_WORLD);
    }

    if (rank == 0) {
        cout << "Result Matrix C:\n";
        for (const auto& row : C) {
            for (int val : row) cout << val << " ";
            cout << "\n";
        }
    }

    MPI_Finalize();
    return 0;
}