/*
 Nattawut Khuadplod
 .
 .
 .
HeatTransfer using MPI Collective Communication
*/
#include<mpi.h>
#include<stdio.h>
#include<stdlib.h>


int main(int argc, char *argv[]){

int rank, size , rankCounter;
    MPI_Init(&argc,&argv);
    MPI_Comm_rank(MPI_COMM_WORLD,&rank);
    MPI_Comm_size(MPI_COMM_WORLD,&size);

    int row ;				// number of nodes in x direction
    int col ;				// number of nodes in y direction
    int iters;				// number of iterations
    int i, j, k, frac=0;
    int print_i,print_j;				// initial variables
    float *heatMatrix;
    float *buff_heatMatrix,*buff_temp;
    float sum;

    int partrow,frac_data;

    int *sendcounts = (int*)malloc(sizeof(int)*size);
    int *displs = (int*)malloc(sizeof(int)*size);

    FILE *fp;


    if(rank==0){
        fp = fopen(argv[1],"r");
        if(fp){
            fscanf(fp,"%d %d \n",&row,&col);
            fclose(fp);
        }
        iters = atoi(argv[3]);

        partrow = row/size;
        frac_data = row%size;

        heatMatrix = (float*)calloc(row*col,sizeof(float));

        //top
        for(i=0;i<col;i++){
          heatMatrix[i] = 255;
        }
        //left
        for(i=1;i<row;i++){
          heatMatrix[i*col] = 255;
        }
        //bottom
        for(i=0;i<col;i++){
          heatMatrix[(row-1)*col+i] = 255;
        }
        //right
        for(i=1;i<row;i++){
          heatMatrix[i*col+(col-1)] = 255;
        }
    }

    MPI_Bcast(&row,1,MPI_INT,0,MPI_COMM_WORLD);
    MPI_Bcast(&col,1,MPI_INT,0,MPI_COMM_WORLD);
    MPI_Bcast(&iters,1,MPI_INT,0,MPI_COMM_WORLD);
    MPI_Bcast(&frac_data,1,MPI_INT,0,MPI_COMM_WORLD);
    MPI_Barrier(MPI_COMM_WORLD);

    int *displs_top = (int*)calloc(size,sizeof(int));
    int *displs_bottom = (int*)calloc(size,sizeof(int));

    int sumDispls=0;
    // Calculate send counts and displacements
        for (int i = 0; i < size; i++) {
            sendcounts[i] = (row/size)*col;

            /*
            if (frac_data > 0 && i==0) {
                sendcounts[i]+=frac_data*col;
                frac_data=0;
            }
            */

            if (frac_data > 0 ) {
                sendcounts[i]+=col;
                frac_data--;
            }
            displs[i] = sumDispls;
            sumDispls += sendcounts[i];
        }

    buff_heatMatrix = (float*)calloc(row*col,sizeof(float));
    buff_temp = (float*)calloc(row*col,sizeof(float));

    // Scatter global data to all processes
    MPI_Scatterv(&heatMatrix[0], sendcounts, displs, MPI_FLOAT, buff_heatMatrix,row*col, MPI_FLOAT, 0, MPI_COMM_WORLD);

    for(k=0;k<iters;k++){
    if(rank==0){
        //Rank 0

        partrow = sendcounts[rank]/col;

        for(i=1;i<size;i++){
            if (i==(size-1)){
                displs_top[i]=displs[i]-col;
                MPI_Send(&heatMatrix[displs_top[i]],col, MPI_FLOAT, i, 0, MPI_COMM_WORLD);
            }else{
                displs_top[i]=displs[i]-col;
                displs_bottom[i]=displs[i+1];
                MPI_Send(&heatMatrix[displs_top[i]],col, MPI_FLOAT, i, 0, MPI_COMM_WORLD);
                MPI_Send(&heatMatrix[displs_bottom[i]],col, MPI_FLOAT, i, 1, MPI_COMM_WORLD);
            }
        }

        for(i=1; i<partrow; i++)
          {
            for(j=1; j<col-1; j++)
            {
              //left = heatMatrix[(i*col)+j-1];
              //right = heatMatrix[(i*col)+j+1];
              //top = heatMatrix[((i-1)*col)+j];
              //bot = heatMatrix[((i+1)*col)+j];
              buff_temp[(i*col)+j] = .25*(heatMatrix[(i*col)+j-1]+heatMatrix[(i*col)+j+1]+heatMatrix[((i-1)*col)+j]+heatMatrix[((i+1)*col)+j]);
            }
          }

        //Set result
        for(i=0;i<partrow;i++){
            for(j=0;j<col;j++){
                if (heatMatrix[(i*col)+j] >=255){
                     buff_heatMatrix[(i*col)+j] = 255;
                }
                else{
                     buff_heatMatrix[(i*col)+j] = buff_temp[(i*col)+j];
                }
            }
        }

    }else{
        //Other rank
        partrow = sendcounts[rank]/col;
        float *buff_top = (float*)calloc(col,sizeof(float));
        float *buff_bottom = (float*)calloc(col,sizeof(float));

        if(rank==size-1){
            MPI_Recv(&buff_top[0],col, MPI_FLOAT, 0, 0, MPI_COMM_WORLD,MPI_STATUS_IGNORE);

            for(i=0; i<partrow-1; i++){
                for(j=1; j<col-1; j++){

                    if(i==0){
                        buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_top[j]+buff_heatMatrix[((i+1)*col)+j]);

                    }else{
                        buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_heatMatrix[((i-1)*col)+j]+buff_heatMatrix[((i+1)*col)+j]);
                    }

               }
            }

            //Set result
            for(i=0;i<partrow;i++){
                for(j=0;j<col;j++){
                    if (buff_heatMatrix[(i*col)+j] >=255){
                        buff_heatMatrix[(i*col)+j] = 255;
                    }
                    else{
                        buff_heatMatrix[(i*col)+j] = buff_temp[(i*col)+j];
                    }
                }
            }
            MPI_Send(&buff_top[0],col, MPI_FLOAT, 0, 2, MPI_COMM_WORLD);
        }else{
            MPI_Recv(&buff_top[0],col, MPI_FLOAT, 0, 0, MPI_COMM_WORLD,MPI_STATUS_IGNORE);
            MPI_Recv(&buff_bottom[0],col, MPI_FLOAT, 0, 1, MPI_COMM_WORLD,MPI_STATUS_IGNORE);

            for(i=0; i<partrow; i++){
                for(j=1; j<col-1; j++){
                    if(i==0){
                        if(i==partrow-1){
                            buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_top[j]+buff_bottom[j]);
                        }
                        else{
                            buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_top[j]+buff_heatMatrix[((i+1)*col)+j]);
                        }
                    }else{
                        if(i==partrow-1){
                            buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_heatMatrix[((i-1)*col)+j]+buff_bottom[j]);
                        }
                        else{
                            buff_temp[(i*col)+j] = .25*(buff_heatMatrix[(i*col)+j-1]+buff_heatMatrix[(i*col)+j+1]+buff_heatMatrix[((i-1)*col)+j]+buff_heatMatrix[((i+1)*col)+j]);

                        }
                    }
                }
            }

            //Set result
            for(i=0;i<partrow;i++){
                for(j=0;j<col;j++){
                    if (buff_heatMatrix[(i*col)+j] >=255){
                        buff_heatMatrix[(i*col)+j] = 255;
                    }
                    else{
                        buff_heatMatrix[(i*col)+j] = buff_temp[(i*col)+j];
                    }
                }
            }
        }
            free(buff_top);
            free(buff_bottom);
    }

        MPI_Barrier(MPI_COMM_WORLD);
        MPI_Gatherv(buff_heatMatrix,sendcounts[rank],MPI_FLOAT,heatMatrix,sendcounts,displs,MPI_FLOAT,0,MPI_COMM_WORLD);

    }
    if(rank==0){
        fp = fopen(argv[2],"w");
        if(fp){
            fprintf(fp,"%d %d\n",row,col);
            for(i=0;i<row;i++){
                for(j=0;j<col;j++){
                    fprintf(fp, "%.0f ",heatMatrix[(i*col)+j]);
                }
                fprintf(fp, "\n");
            }
        fclose(fp);
        }
        printf("Successful! , Let's see result file.\n");

    free(heatMatrix);
    }

    MPI_Finalize();

    free(buff_heatMatrix);
    free(buff_temp);
    free(sendcounts);
    free(displs);

  return 0;
}