#include <stdlib.h>

#include <stdio.h>

#include <math.h>

#include "mpi.h"

#include "block.h"

void indexFind(int num,int t[]) {

  int tmp;

  tmp = sqrt(num);

  while (num % tmp != 0)

    tmp --;

  t[0] = tmp;

  t[1] = num /tmp;

}

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

  int myRank,nbProc;

  int indice[2]; // indice de blocs

  int tmp,i,j;

  int counter = 0;

  int stride = BLOCK_SIZE+2;

  double *newBlock,*oldBlock;

  double *block1 = allocBlock();

  double *block2 = allocBlock();

  MPI_Datatype mpi_lin,mpi_col;

  MPI_Request req1[8],req2[8]; // order : sendUp,sendDown,sendLeft,sendRight,recvUp,recvDown,recvLeft,recvRight

  MPI_Request * req;

  MPI_Status status[8];

  int rankUp,rankDown,rankLeft,rankRight;

  int flag = 1;

  ZeroBlock(block1);

  ZeroBlock(block2);

  // new MPI type definition

  MPI_Type_vector(BLOCK_SIZE,1,1,MPI_DOUBLE,&mpi_lin);

  MPI_Type_vector(BLOCK_SIZE,1,BLOCK_SIZE+2,MPI_DOUBLE,&mpi_col);

  MPI_Type_commit(&mpi_lin);

  MPI_Type_commit(&mpi_col);

  MPI_Init(&argc,&argv);

  MPI_Comm_rank(MPI_COMM_WORLD,&myRank);

  MPI_Comm_size(MPI_COMM_WORLD,&nbProc);

  indexFind(nbProc,indice);

  // creation of persistent communication requests

  rankUp = (myRank < indice[0]) ? myRank : (myRank-indice[0]); // first line of the grille

  rankDown = (myRank >= nbProc-indice[0]) ? myRank : (myRank+indice[0]);// last line of the grille

  rankLeft = (myRank/indice[0] == 0) ? myRank : (myRank-1); // first column of the grille

  rankRight = ((myRank+1)/indice[0] == 0) ? myRank : (myRank+1); // last column of the grille

  // from block1 to block2

  MPI_Send_init(&(block1[1]),1,mpi_lin,rankUp,0,MPI_COMM_WORLD,&req1[0]);

  MPI_Recv_init(&(block2[1]),1,mpi_lin,rankUp,0,MPI_COMM_WORLD,&req1[4]);

  MPI_Send_init(&(block1[stride*(BLOCK_SIZE+1)+1]),1,mpi_lin,rankDown,0,MPI_COMM_WORLD,&req1[1]);

  MPI_Recv_init(&(block2[stride*(BLOCK_SIZE+1)+1]),1,mpi_lin,rankDown,0,MPI_COMM_WORLD,&req1[5]);

  MPI_Send_init(&(block1[stride]),1,mpi_col,rankLeft,0,MPI_COMM_WORLD,&req1[2]);

  MPI_Recv_init(&(block2[stride]),1,mpi_col,rankLeft,0,MPI_COMM_WORLD,&req1[6]);

  MPI_Send_init(&(block1[2*stride-1]),1,mpi_col,rankRight,0,MPI_COMM_WORLD,&req1[3]);

  MPI_Recv_init(&(block2[2*stride-1]),1,mpi_col,rankRight,0,MPI_COMM_WORLD,&req1[7]);

  // from block2 to block1

  MPI_Send_init(&(block2[1]),1,mpi_lin,rankUp,0,MPI_COMM_WORLD,&req2[0]);

  MPI_Recv_init(&(block1[1]),1,mpi_lin,rankUp,0,MPI_COMM_WORLD,&req2[4]);

  MPI_Send_init(&(block2[stride*(BLOCK_SIZE+1)+1]),1,mpi_lin,rankDown,0,MPI_COMM_WORLD,&req2[1]);

  MPI_Recv_init(&(block1[stride*(BLOCK_SIZE+1)+1]),1,mpi_lin,rankDown,0,MPI_COMM_WORLD,&req2[5]);

  MPI_Send_init(&(block2[stride]),1,mpi_col,rankLeft,0,MPI_COMM_WORLD,&req2[2]);

  MPI_Recv_init(&(block1[stride]),1,mpi_col,rankLeft,0,MPI_COMM_WORLD,&req2[6]);

  MPI_Send_init(&(block2[2*stride-1]),1,mpi_col,rankRight,0,MPI_COMM_WORLD,&req2[3]);

  MPI_Recv_init(&(block1[2*stride-1]),1,mpi_col,rankRight,0,MPI_COMM_WORLD,&req2[7]);

  // loop of calculation

  while (flag) {

    tmp = (counter%2==0);

    newBlock = tmp ? block1 : block2;

    oldBlock = tmp ? block2 : block1;

    req = tmp ? req1 : req2;

    counter++;

    ZeroBlock(newBlock);

    for (i=1;i<=BLOCK_SIZE;i++) {     // TODO : add thread

      for (j=1;j<=BLOCK_SIZE;j++) {

    tmp = oldBlock[i*stride+j]/5;

    newBlock[i*stride+j] += tmp;

    newBlock[i*stride+j-1] += tmp;

    newBlock[i*stride+j+1] += tmp;

    newBlock[(i+1)*stride+j] += tmp;

    newBlock[(i-1)*stride+j] += tmp;

      }

    }

    // communication (newBlock vers oldBlock)

    MPI_Startall(8,req);

    MPI_Waitall(8,req,status);

    // add to the result

    for (i=1;i<=BLOCK_SIZE;i++)

      newBlock[i] += oldBlock[i];

    for (i=stride*(BLOCK_SIZE+1)+1;i<(stride*stride-1);i++)

      newBlock[i] += oldBlock[i];

    for (i=stride;i<stride*(BLOCK_SIZE+1);i+=stride)

      newBlock[i] += oldBlock[i];

    for (i=2*stride-1;i<stride*(BLOCK_SIZE+1);i+=stride)

      newBlock[i] += oldBlock[i];

    // TODO Loop Condition : change flag to 0 if changes are no sensitive. need MPI_Gather

  }

  freeBlock(oldBlock);

  freeBlock(newBlock);

  MPI_Type_free(&mpi_lin);

  MPI_Type_free(&mpi_col);

  MPI_Finalize();

  return EXIT_SUCCESS;

}