#include <stdlib.h>
#include <stdio.h>
#include <iostream.h>
#include <time.h>
#include <sys/time.h>
//Uncomment if these need to be undefined
//#undef SEEK_SET
//#undef SEEK_CUR
//#undef SEEK_END
#include "mpi.h"

void PrintSurfaceMatrix(double** M, int LD, int D);

int main(int argc, char* argv[]){
  
  MPI_Init(&argc,&argv);
  int NCPUS, MY_PE;
  MPI_Comm_size(MPI_COMM_WORLD, &NCPUS);
  MPI_Comm_rank(MPI_COMM_WORLD, &MY_PE);

  // Read Input Parameters
  if(argc < 4) {
    cerr << "\nToo many input paramters.. should be four\n";
    exit(2);
  }

  int Dimension = atoi(argv[1]);
  int Iterations = atoi(argv[2]);
  int RowPeek = atoi(argv[3]);
  int ColPeek = atoi(argv[4]);
   
  // cout <<"\nInputs = " << Dimension << " "<<Iterations << " " 
  //     << RowPeek << " " << ColPeek << "\n";

  if((RowPeek > Dimension) || (ColPeek > Dimension)){
    cerr << "Cannot Peek a matrix element outside of the surface";
    cerr << "Arguments 3 and 4 must be smaller than " << Dimension;
    exit(3);
  }

  // Initialize Matrix
  // So on each processor, we need only enough memory to hold our rows
  // and the two other rows
  double **SurfaceMatrix;
  int *Offsets = (int*)malloc(sizeof(int)*NCPUS);
  int *LocalDims = (int*)malloc(sizeof(int)*NCPUS);
  int LocalDim = Dimension/NCPUS;
  for(int ip = 0; ip < NCPUS; ip++){
    LocalDims[ip] = LocalDim;
    if(ip < (Dimension % NCPUS)) LocalDims[ip]++;
  }
  Offsets[0] = 0;
  for(int ip =1; ip < NCPUS; ip++)
    Offsets[ip] = Offsets[ip-1] + LocalDims[ip-1];
  
  SurfaceMatrix = (double**)malloc(sizeof(double*)*(LocalDims[MY_PE]+2));
  for(int i=0;i<LocalDims[MY_PE]+2;i++)
    SurfaceMatrix[i] = (double*)malloc(sizeof(double)*(Dimension+2));
  
  double **SurfaceMatrix_t;
  SurfaceMatrix_t = (double**)malloc(sizeof(double*)*(LocalDims[MY_PE]+2));
  for(int i=0;i<LocalDims[MY_PE]+2;i++){
    SurfaceMatrix_t[i] = (double*)malloc(sizeof(double)*(Dimension+2));
    for(int j=0;j< Dimension+2;j++){
      SurfaceMatrix_t[i][j] = 0.0;
      SurfaceMatrix[i][j] = 0.0;
    }
  }
  
  for(int i=1;i<=LocalDims[MY_PE];i++)
    for(int j=1;j<=Dimension;j++)
      SurfaceMatrix[i][j] = 0.5;
 
  // Set up 1s Panel
  // Now it gets tricky
  double increment = 100.0 / (Dimension+1);
  int count = Offsets[MY_PE];

  for(int i=1;i<LocalDims[MY_PE]+1;i++){
    SurfaceMatrix[i][0] = (count+i)*increment;
    SurfaceMatrix_t[i][0] = (count+i)*increment;
    SurfaceMatrix[i][Dimension+1] = 0.0;
    SurfaceMatrix_t[i][Dimension+1] = 0.0;
  }

  // Top Row Boudary
  if(MY_PE==0){
    for(int i = 1; i <= Dimension+1;i++){
      SurfaceMatrix[0][i] = 0.0;
      SurfaceMatrix_t[0][i] = 0.0;
    }
  }
  // Bottom Row Boundary
  if(MY_PE==NCPUS-1){
    for(int i=1;i<=Dimension+1;i++){
      SurfaceMatrix[LocalDims[MY_PE]+1][Dimension+1-i]=i*increment;
      SurfaceMatrix_t[LocalDims[MY_PE]+1][Dimension+1-i]=i*increment;
    }
  }
  
  // Iterate
  double TimeStart = MPI_Wtime();
  //PrintSurfaceMatrix(SurfaceMatrix,LocalDims[MY_PE],Dimension);
  MPI_Barrier(MPI_COMM_WORLD);
  fflush(stdout);
  for(int iCount = 1; iCount <=Iterations; iCount++){
    MPI_Request req1,req2,req3,req4;
    MPI_Status stat;
    if(MY_PE != NCPUS-1){
       MPI_Isend(SurfaceMatrix[LocalDims[MY_PE]],Dimension+2,MPI_DOUBLE,
		 MY_PE+1,0,MPI_COMM_WORLD,&req1);
     }
     if(MY_PE !=0){
       MPI_Irecv(SurfaceMatrix[0],Dimension+2,MPI_DOUBLE,MY_PE-1,
		 0,MPI_COMM_WORLD,&req2);
     }
   
     if(MY_PE !=0){
       MPI_Isend(SurfaceMatrix[1],Dimension+2,MPI_DOUBLE,
		 MY_PE-1,0,MPI_COMM_WORLD,&req3);
     }
     if(MY_PE != NCPUS-1){
       MPI_Irecv(SurfaceMatrix[LocalDims[MY_PE]+1],Dimension+2,MPI_DOUBLE,
		MY_PE+1,0,MPI_COMM_WORLD,&req4);
     }
     
     if(MY_PE != 0) {
       MPI_Wait(&req2,&stat);
       MPI_Wait(&req3,&stat);
     }
     if(MY_PE != NCPUS-1){
       MPI_Wait(&req4,&stat);
       MPI_Wait(&req1,&stat);
     }
     
    //This will be a row dominant program.
    for(int i=1;i<=LocalDims[MY_PE];i++)
      for(int j=1;j<=Dimension;j++)
	SurfaceMatrix_t[i][j] = (0.25)*(SurfaceMatrix[i-1][j] +
				       SurfaceMatrix[i][j+1] +
				       SurfaceMatrix[i+1][j] +
				       SurfaceMatrix[i][j-1]);
      
    double ** tmp;
    tmp = SurfaceMatrix;
    SurfaceMatrix = SurfaceMatrix_t;
    SurfaceMatrix_t = tmp;
  }
  double TimeEnd = MPI_Wtime();
  
  if(MY_PE == 0) 
    cout <<"\n Time Iterations = "<<TimeEnd-TimeStart << " seconds";
  
  // Need to find= the proper matrix element to print out.
  int LocalRowPeek, IPRowPeek;
  for(int ip = 0;ip < NCPUS; ip++){
    if(RowPeek > Offsets[ip]){
      IPRowPeek = ip;
      LocalRowPeek = RowPeek - Offsets[ip];
    }
  }
  if(MY_PE == IPRowPeek){
    cout << "\n Result SurfaceMatrix["<<RowPeek<<"]["<<ColPeek<<"]="
	 <<SurfaceMatrix[LocalRowPeek+1][ColPeek+1]<<"\n";
  }
  for(int i=0;i<LocalDims[MY_PE]+2;i++){
    free(SurfaceMatrix[i]);
    free(SurfaceMatrix_t[i]);
  }
  free(SurfaceMatrix);
  free(SurfaceMatrix_t);
  
  MPI_Finalize();
  return 0;
}

void PrintSurfaceMatrix(double**SurfaceMatrix, int LocalDim, int Dimension){
  
  cout << "\t";
  int mype,ncpus;
  MPI_Comm_rank(MPI_COMM_WORLD,&mype);
  MPI_Comm_size(MPI_COMM_WORLD,&ncpus);
  MPI_Barrier(MPI_COMM_WORLD);
  for(int ip=0;ip<ncpus;ip++){
    if(mype == ip){
      cout << "\nMatrix for CPU "<< ip << "\n";
      for(int i=0;i<Dimension+2;i++){
	cout << i << "\t";
      }

      cout << "\n";
  
      for(int i=0;i<LocalDim+2;i++){
	cout << i << "\t";  
	for(int j=0;j<Dimension+2;j++){
	  cout << SurfaceMatrix[i][j] << "\t";
	}
	cout << "\n";
      }
    }
    MPI_Barrier(MPI_COMM_WORLD);
  }
}