#include <iostream>
#include <mpi.h>
#include <cmath>

using namespace std;

const int COLUMN_TAG{20};

const int PRINT_TAG{10};
const int DONE_TAG{11};

const int CELL_WIDTH{4};
const int HEIGHT{5};

//                       space for the borders
using MapT = int[HEIGHT][CELL_WIDTH+2];

void PrintMySlice(const MapT myMap, int rank);
void PrintMap(const MapT myMaps,  int rank, int size);
void InitMyMap(MapT myMap, int rank);
void TradeBorders(MapT mySlice, int rank, int size);
void UpdateSlice(MapT mySlice);
void BuildFinalMap(MapT mySlice, int rank, int size);
void BlockingMessage(string msg, int rank);

static MPI_Datatype colType;

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

    // set up the data type
    //      the data will be   b d d d d ... d b
    //      there will be CELL_WIDTH d's
    //      so the stride will be CELL_WIDTH+2
    
    MPI_Type_vector(HEIGHT, 1, CELL_WIDTH+2, MPI_INT, &colType);
    MPI_Type_commit(&colType);

    // I want two, because is a real game of life, one would hold the 
    // current state and one would hold the next state
    MapT mySlices[2];

    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    InitMyMap(mySlices[0], rank);

    BlockingMessage("Inital Data", rank);
    PrintMap(mySlices[0], rank, size);

    TradeBorders(mySlices[0], rank, size);
    
    BlockingMessage("After Border Swap", rank);
    PrintMap(mySlices[0], rank, size);

    BlockingMessage("The map", rank);
    BuildFinalMap(mySlices[0], rank, size);

    UpdateSlice(mySlices[0]);
    TradeBorders(mySlices[0], rank, size);

    BlockingMessage("After Update", rank);
    PrintMap(mySlices[0], rank, size);

    BlockingMessage("Map After Update", rank);
    BuildFinalMap(mySlices[0], rank, size);

    MPI_Type_free(&colType);
    MPI_Finalize();
    return 0;
}

void BlockingMessage(string msg, int rank) {
    MPI_Barrier(MPI_COMM_WORLD) ;
    if (rank == 0) {
       cout << msg << endl;
    }
}

void UpdateSlice(MapT mySlice){
   for(int i = 0 ; i < HEIGHT; ++i) {
      for(int j = 1; j <= CELL_WIDTH ;++j) {
         mySlice[i][j] += 1000;
      }
   }
}

void PrintFinalMap(int * picture, int rows, int cols){
   for(int i = 0; i < rows; ++i) {
      for(int j = 0; j < cols; ++j) {
         cout << picture[i*cols + j] << " " ;
      }
      cout << endl;
   }
   cout << endl;
}

void BuildFinalMap(MapT mySlice, int rank, int size){
   int * picture;
   int blockSize = CELL_WIDTH*size;

   if (rank == 0) {
       picture = new int[HEIGHT*blockSize] ;
   }

   for(int row = 0; row < HEIGHT; ++row) {
       MPI_Gather(&(mySlice[row][1]), CELL_WIDTH, MPI_INT,
                  &picture[row*blockSize], CELL_WIDTH, MPI_INT, 
                  0, MPI_COMM_WORLD);
   }

   if (rank == 0) {
       PrintFinalMap(picture, HEIGHT, blockSize);
       delete[] picture;
   }
}

void TradeBorders(MapT mySlice, int rank, int size){

    int dest = rank-1;
    // send left
    if (rank == 0) {
        dest = size-1;
    } 
    int src = (rank + 1) % size;

    MPI_Status status;

    MPI_Sendrecv(&mySlice[0][1], 1, colType, dest, COLUMN_TAG, 
                 &mySlice[0][CELL_WIDTH+1], 1, colType, src, COLUMN_TAG, 
                 MPI_COMM_WORLD, &status);

    swap(src, dest);

    // send right
    MPI_Sendrecv(&mySlice[0][CELL_WIDTH], 1, colType, dest, COLUMN_TAG, 
                 &mySlice[0][0], 1, colType, src, COLUMN_TAG, 
                 MPI_COMM_WORLD, &status);

}

void InitMyMap(MapT  myMap, int rank){
    
        for(int i =0; i < HEIGHT; ++i) {
           myMap[i][0] =  rank;
           myMap[i][CELL_WIDTH+1] = rank;

           for(int j = 1; j <= CELL_WIDTH; ++j) {
              myMap[i][j] = (rank+1)*100+i*10 + j;  
           }
        }
}

void PrintMap(const MapT myMap, int rank, int size){

    MPI_Status status;

    if (rank == 0) {
        PrintMySlice(myMap, rank);

        for(int i = 1; i < size; ++i) {
           MPI_Send( nullptr, 0, MPI_INT, i, PRINT_TAG, MPI_COMM_WORLD);
           MPI_Recv(nullptr, 0, MPI_INT, i, DONE_TAG, MPI_COMM_WORLD, &status);
        }

        cout << endl << endl;
    } else {
        MPI_Recv(nullptr, 0, MPI_INT, 0, PRINT_TAG, MPI_COMM_WORLD, &status);

        PrintMySlice(myMap, rank);

        MPI_Send( nullptr, 0, MPI_INT, 0, DONE_TAG, MPI_COMM_WORLD);
    }
}

void PrintMySlice(const MapT  myMap, int rank){
    cout << "Slice for process " << rank << endl;
    for(int i = 0; i < HEIGHT; ++i) {
       for(int j = 0; j < CELL_WIDTH+2; ++j) {
           cout << myMap[i][j] << " " ;
       }
       cout << endl;
    }
    cout << endl;
}