/* This is the block-row partitioned version */
/* I already wrote some parts of the code */
#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include<iostream>

using namespace std;

#define SIZE1 64
#define SIZE2 128
#define SIZE3 64

double a[SIZE1][SIZE2];
double b[SIZE2][SIZE3];

double c[SIZE1][SIZE3]; // c = a * b

int main(int argc, char *argv[]){
    int i,j,k;
    int rank; // process rank
    int size; // number of processes
    FILE *myfile1, *myfile2, *myfile3;

    // read array a from file a.txt
    myfile1 = fopen("a.txt", "r");
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if(rank==0){
        // read array a from a.txt file
        myfile1 = fopen("a.txt", "r");
        for(i=0; i<SIZE1; i++)
            for(j=0; j<SIZE2; j++)
                fscanf(myfile1,"%lf",&a[i][j]);
        fclose(myfile1);

        // read array b from b.txt file
        myfile2 = fopen("b.txt", "r");
        for(i=0; i<SIZE2; i++)
            for(j=0; j<SIZE3; j++)
                fscanf(myfile2,"%lf",&b[i][j]);
        fclose(myfile2);
    }

    //SIZE1 -> a rows // SIZE2 -> a columns b rows // SIZE3 -> b columns
    int chunk = SIZE1/size;     //size =2 -> chunk = 2 -> tmp[2][8] * b[8][4] ->  tmp[2][4];

    double tmp[chunk][SIZE3];

    if(rank==0) {
        //MPI_Scatter (sendbuf, sendcnt, sendtype,recvbuf, recvcnt, recvtype, root, comm)
        MPI_Scatter(a, chunk * SIZE2, MPI_DOUBLE, tmp, chunk * SIZE2, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    }

    for(i=0; i<chunk; i++){
        for(j=0; j<chunk; j++){
            tmp[i][j] = 0;
            for(k=0; k<SIZE2; k++)
                tmp[i][j] += a[i][k] * b[k][j];
        }
    }

    //MPI_Gather (sendbuf, sendcnt, sendtype,recvbuf, recvcnt, recvtype, root, comm)
    MPI_Gather (tmp, chunk * SIZE2, MPI_DOUBLE, c, chunk * SIZE2, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    if(rank==0){
        //write array c to c.txt file
        myfile3 = fopen("c.txt","wb");
        for(i=0; i<SIZE1; i++){
            for(j=0; j<SIZE3; j++)
                fprintf(myfile3, "%lf ",c[i][j]);
            fprintf(myfile3, "\n");
        }
        fclose(myfile3);
    }

    MPI_Finalize();
}