#include <mpi.h>
#define MPE_GRAPHICS
#include "mpe.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <math.h>

int main(int argc, char *argv[])
{
  char *displayname = ":0";
  int Nx, Ny, iters ;
  int rank, size ;
  int next, prev ;
  int i, j, k;
  int my_xsize, my_ysize, nremaining;
  int my_xoffset=0, my_yoffset=0;
  int width = 600, height = 600;
  int xloc, yloc, xwid, ywid;
  int ierr;
  double **matrix, **temp, **addr, **trees;
  double factor, sum, tree_min, tree_max, tree_mult;
  double tree_min_local, tree_max_local;
  double slavetime, totaltime, starttime ;
  MPE_XGraph graph;
  MPE_Color *color_array;

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

  /* Set neighbors */
  if (rank == 0)
    prev = MPI_PROC_NULL;
  else
    prev = rank-1;
  if (rank == size - 1)
    next = MPI_PROC_NULL;
  else
    next = rank+1;
    
  /* If I'm process 0, determine the matrix size and # of iterations */
  
  if ( rank == 0 ) {
    printf("X Size : ") ;
    scanf("%d",&Nx) ;
    printf("Y Size : ") ;
    scanf("%d",&Ny) ;
    printf("Iterations : ") ;
    scanf("%d",&iters) ;
  }

  /* Broadcast the size and # of iterations to all processes */
  MPI_Bcast(&Nx, 1, MPI_INT, 0, MPI_COMM_WORLD) ;
  MPI_Bcast(&Ny, 1, MPI_INT, 0, MPI_COMM_WORLD) ;
  MPI_Bcast(&iters, 1, MPI_INT, 0, MPI_COMM_WORLD) ;

  /* Determine my part of the matrix for y dimension */
  my_xsize = Nx/size;
  /* Take care of leftover cells */
  /* % is modulo or remainder operator */
  nremaining = Nx%size;
  if (rank < nremaining) my_xsize++;

  /* Determine offset for my part of matrix */
  my_xoffset = rank * (Nx/size);
  /* Take care of extra cells */  
  if (rank > nremaining) {
  	 my_xoffset += nremaining;
  }
  else {
  	 my_xoffset += rank;
  }

  /* My x size is the same for all processors */
  my_ysize = Ny;
  my_yoffset = 0;

  /* allocate the memory dynamically for the matrix */
  /* the extra 2 is for the ghost boundaries on each end */
  matrix = (double **)malloc(sizeof(double *)*(my_xsize+2)) ;
  temp = (double **)malloc(sizeof(double *)*(my_xsize+2)) ;
  trees = (double **)malloc(sizeof(double *)*(my_xsize+2)) ;
  for (i = 0; i < my_xsize+2; i++) {
    matrix[i] = (double *)malloc(sizeof(double)*(my_ysize+2)) ;
    temp[i] = (double *)malloc(sizeof(double)*(my_ysize+2)) ;
    trees[i] = (double *)malloc(sizeof(double)*(my_ysize+2)) ;
  }

  /* Initialize the life matrix */
  srand48((long)(1000^(my_xsize)));
  for (i = 0; i <= my_xsize+1; i++){
    for (j = 0; j<= my_ysize+1; j++){
      factor = drand48()/2.0+0.75;
      matrix[i][j] = 0.0;
      temp[i][j] = 0.0;
      trees[i][j] = factor*(0.2+1.6*((double)j/(double)(Ny+1)));
    }
  }
  if ((Nx+1)/2-my_xoffset >=0 && (Nx+1)/2-my_xoffset <= my_xsize+1){
    matrix[(Nx+1)/2-my_xoffset][(Ny+1)/2] = 8100.0;
  }
  
  /* Open the graphics display */
  MPE_Open_graphics( &graph, MPI_COMM_WORLD, displayname,
                   -1, -1, width, height, 0 );
  
  color_array = (MPE_Color *) malloc(sizeof(MPE_Color)*65000);

  ierr = MPE_Make_color_array(graph, sizeof(color_array), color_array);
  if (ierr && rank == 0) printf("Error(Make_color_array): ierr is %d\n",ierr);

  /* Do the calculation for the given number of iterations */
  starttime = MPI_Wtime() ;
  for (k = 0; k < iters; k++) {
    MPI_Request      req[4];
    MPI_Status       status[4];
    
    for (j=0; j<=my_ysize+1; j++){
      matrix[0][j] = matrix[1][j];
      matrix[my_xsize+1][j] = matrix[my_xsize][j];
    }

    for (i=0; i<=my_xsize+1; i++){
      matrix[i][0] = matrix[i][1];
      matrix[i][my_ysize+1] = matrix[i][my_ysize];
    }

    /* Send and receive boundary information */
    MPI_Isend(&matrix[1][0],my_ysize+2,MPI_DOUBLE,prev,1,MPI_COMM_WORLD,req);
    MPI_Irecv(&matrix[0][0],my_ysize+2,MPI_DOUBLE,prev,0,MPI_COMM_WORLD,req+1);
    MPI_Isend(&matrix[my_xsize][0],my_ysize+2,MPI_DOUBLE,next,0,MPI_COMM_WORLD,req+2);
    MPI_Irecv(&matrix[my_xsize+1][0],my_ysize+2,MPI_DOUBLE,next,1,MPI_COMM_WORLD,req+3);
    MPI_Waitall(4, req, status);

    /* For each element of the matrix ... */  
    for (i = 0; i <= my_xsize+1; i++) {
      for (j = 0; j<= my_ysize+1; j++) {
      	matrix[i][j] *= trees[i][j];
      }
    }

    /* printf("i j xloc yloc xwid ywid value\n");*/
    for (i = 1; i <= my_xsize; i++) {
      xloc = ((my_xoffset + i - 1) * width)  / Nx;
      xwid = ((my_xoffset + i) * width)  / Nx - xloc;
      for (j = 1; j <= my_ysize; j++) {
        yloc = ((my_yoffset + j - 1) * height) / Ny;
        ywid = ((my_yoffset + j) * height) / Ny - yloc;

        /* find out the value of the current cell */
        sum = matrix[i][j]     + matrix[i-1][j-1] + matrix[i-1][j] + 
              matrix[i-1][j+1] + matrix[i][j-1]   + matrix[i][j+1] +
              matrix[i+1][j-1] + matrix[i+1][j]   + matrix[i+1][j+1];
        temp[i][j] = sum/9.0;
        trees[i][j] -= 2.0*temp[i][j]/8100.0;
        if (trees[i][j] < 0.0) trees[i][j]=0.0;
        if ((int)temp[i][j] > sizeof(color_array)-2) temp[i][j] = (double)(sizeof(color_array)-2);

/*        printf("DEBUG %d %d %d %d\n",rank, i, j, sizeof(color_array)-(int)temp[i][j]);*/
        MPE_Fill_rectangle( graph, xloc, yloc, xwid, ywid, sizeof(color_array)-(int)temp[i][j]);
        /*printf("%d %d %d %d %d %d %d %lf\n",rank,i, j, xloc, yloc, xwid, ywid, temp[i][j])*/;
      }
    }
    
    MPE_Update( graph );

 
    /* Swap the matrices */
    addr = matrix ;
    matrix = temp ;
    temp = addr ;
  }

  /* Return the average time taken/processor */
  slavetime = MPI_Wtime() - starttime;
  MPI_Reduce (&slavetime, &totaltime, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
  
  /* Print the total time taken */
  if (rank == 0){
    printf("[%d] Beetlemania finished in %lf seconds\n",rank,totaltime/(double)size/100);
  }
    
  sleep(10);
  
  if (rank == 0){
    printf("\n\n[%d] Now Plotting Tree Distribution\n",rank);
  }

  tree_max_local = -1.0E34;
  tree_min_local =  1.0E34;
  for (i = 0; i <= my_xsize+1; i++) {
    for (j = 0; j<= my_ysize+1; j++) {
      /*printf("%d %d %d %lf\n",rank,i, j, trees[i][j]);*/
      if (trees[i][j] < tree_min_local) tree_min_local = trees[i][j];
      if (trees[i][j] > tree_max_local) tree_max_local = trees[i][j];
    }
  }
  MPI_Reduce (&tree_min_local, &tree_min, 1, MPI_DOUBLE, MPI_MIN, 0, MPI_COMM_WORLD);
  MPI_Reduce (&tree_max_local, &tree_max, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
  tree_mult = (double)(sizeof(color_array-2))/(tree_max - tree_min);

  for (i = 0; i <= my_xsize+1; i++) {
    for (j = 0; j<= my_ysize+1; j++) {
      trees[i][j] = trees[i][j] * 0.70* tree_mult - tree_min;
    }
  }

  for (i = 1; i <= my_xsize; i++) {
    xloc = ((my_xoffset + i - 1) * width)  / Nx;
    xwid = ((my_xoffset + i) * width)  / Nx - xloc;
    for (j = 1; j <= my_ysize; j++) {
      yloc = ((my_yoffset + j - 1) * height) / Ny;
      ywid = ((my_yoffset + j) * height) / Ny - yloc;

      MPE_Fill_rectangle( graph, xloc, yloc, xwid, ywid, sizeof(color_array)-(int)trees[i][j]);
    }
  }
    
  MPE_Update( graph );

  sleep(10);
  MPE_Close_graphics(&graph);
  MPI_Finalize();
  exit(0);
}