/* Computing the solver for heat equation in multithreads, then compare the 
 *  result with the solver got in a sequencial way.
 *
 *  Command line example:
 *  civl verify -inpuNPROCS=3 diffusion_1d.cvl 
 *  Notice: Verify this program may take much long time and much large memory space.
 **/ 

#include<civlc.h>
#include "mp_root.cvh"
#define FROM_LEFT 1           /* mpi tag which means the source is send or receive
				                 from left chunk*/
#define FROM_RIGHT 2          /* mpi tag which means the source is send or receive
				                 from right chunk*/
#define REDUCE 3              /* reduce tag */
#define Threshold 1           /* the threshold of difference */
#define K 0.3                 /* k = alpha^2 * dt/(dx^2) */
#define ELENUM 7              /* the length of the array */
#define RELEASE 0             /* sync signal */
#define HOLD 1                /* sync signal */
#define NSTEPS 2              /* time */

int reduce_count = 0;           /* used for sync */
short __release = HOLD;         /* used for sync */
double globle_u[ELENUM];        /* used for store the whole array */
double seq_u[ELENUM];           /* used for store the whole array in the sequential way */

/* return the absolute double value */
double doubleAbs(double v){
  if(v < 0)
    return -v;
  else
    return v;
}

/* return boundary values. Here we just use a constant value */
double boundaryValues(int x){
  return 4.0;
}

/* update the array using heat equation, 
   returns the max difference between the previous one and updated one*/
double update(double * u, double * u_new, int start, int end){
    int i;
    int chunk_length = end - start + 1;

    // chunk_length = (the length of u or u_new) - 2
    for(i=1; i<chunk_length+1; i++)
      u_new [i] = u[i] + K * (u[i-1] + u[i+1] - 2*u[i]);

    // compute the difference between u and u_new
    double max_diff = 0;
    double diff;

    for(i=1; i<chunk_length+1; i++){
      diff = doubleAbs(u_new[i] - u[i]);
      if(diff > max_diff)
	max_diff = diff;
      // update u
      u[i] = u_new[i];
    }
    return max_diff;
}

/* Communicate with left chunk and right chunk,
   send the first interior element to left chunk and receive for 
   the last interior element */
void communicate(double * u, int left, int right, int start, 
		 int end, int rank){
  #include "mp_proc_diffusion.cvh"

  int chunk_length = end - start + 1;
  // the most left or right chunks just need to exchange with one side 
    if(left == -1 && right == -1)
      return;
    else if(left == -1){ 
      send(&u[chunk_length], 1, right, FROM_LEFT, rank); 
      recv(&u[chunk_length + 1], 1, right, FROM_RIGHT, rank);
    }
    else if(right == -1){
      send(&u[1], 1, left, FROM_RIGHT, rank); 
      recv(&u[0], 1, left, FROM_LEFT, rank);
    }else{
      send(&u[chunk_length], 1, right, FROM_LEFT, rank);
      send(&u[1], 1, left, FROM_RIGHT, rank);
      recv(&u[chunk_length + 1], 1, right, FROM_RIGHT, rank);
      recv(&u[0], 1, left, FROM_LEFT, rank);
    }
}

/* compute the value of start, end, left chunk and right chunk
   return an array with those values in such order*/
void chunkProperties(int length, int nprocs, int rank, int* results){
  if(length < 3 || nprocs > (length - 2))
    return -1;
  
    int remainder = (length - 2) % nprocs;
    int chunk_length =(length - 2) / nprocs;

    int start = rank * chunk_length + 1; 

    // the last chunk takes the remainder
    if(rank == nprocs - 1){
      chunk_length = chunk_length + remainder;
    }

    int end = start + chunk_length - 1;

    int left_chunk = rank - 1;
    int right_chunk = rank + 1;
    if(right_chunk >= nprocs)
      right_chunk = -1;

    results[0] = start;
    results[1] = end;
    results[2] = left_chunk;
    results[3] = right_chunk;

  }

/* the numbers of elements in u and u_new are 2 more than chunk_length because of 
   ghost elements */
int initChunk(double * u, double * u_new, int start, int end,
         int left_chunk, int right_chunk){
    int i;
    int chunk_length = end - start + 1; // the number of interior elements

    /* initiate interior array */
    for(i=1; i< chunk_length+1; i++){
      u[i] = 0;
      u_new[i] = 0;
    }

  /* for the most left or right chunks, initiate the first ghost ele and 
     the last ghost ele */
    if(left_chunk == -1){
	u[0] = boundaryValues(0);
	u_new[0] = boundaryValues(0);
    }
    if(right_chunk == -1){
      i = chunk_length + 1;
      u[i] = boundaryValues(i);
      u_new[i] = boundaryValues(i);
    }

    return chunk_length;
}

/* combine all separate chunks from all processes */
void combineU(int start, int end, double * u, int rank, double * whole_u){
#include "mp_proc_diffusion.cvh"
  double receive_whole_u[ELENUM];
  whole_u[0] = boundaryValues(0);
  whole_u[ELENUM - 1] = boundaryValues(ELENUM - 1);
  int i = start;
  int j = 1;
  int k = 1;
  for(; k<ELENUM - 1; k++){
		whole_u[k] = 0;
  }
  for(; i < end+1; i++){
		whole_u[i] = u[j];
		j++;
  }
  if(rank != 0){
	reduce_send(whole_u, ELENUM, 0, REDUCE, rank);
  }else{
	for(i = 1; i<NPROCS; i++){
	  reduce_recv(receive_whole_u, ELENUM, i, REDUCE, rank);
	  for(j = 1; j<ELENUM-1; j++){
		  whole_u[j] += receive_whole_u[j];
	  }
	}
  }
}

/* synchronization */
void procsHold(int nprocs){
  $atomic{
    if(reduce_count < nprocs -1){
      reduce_count++;
    }else{
      reduce_count++;
      __release = RELEASE;
    }
  }
  $when (__release == RELEASE);
}
/* synchronization */
void procsRelease(){
  $atomic{
    if(reduce_count > 1){
      reduce_count--;
    }else{
      reduce_count--;
      __release = HOLD;
    }
  }
 $when (__release == HOLD);
}

/* computing the solver in a sequential way*/
void seqDiffusion1d(double * seq_u){
  double diff,max_diff = 0.0;
  double u[ELENUM], u_new[ELENUM];
  int i;
	int nsteps = NSTEPS;

  //Initiate                                                              
  u[0] = boundaryValues(0);
  u_new[0] = boundaryValues(0);
  u[ELENUM-1] = boundaryValues(ELENUM - 1);
  u_new[ELENUM-1] = boundaryValues(ELENUM - 1);
  for(i=1; i<ELENUM-1; i++){
    u[i] = 0;
    u_new[i] = 0;
  }
  //Jacobi Iteration             
  while(nsteps > 0){
    //update                                         
    for(i=1; i<ELENUM-1; i++){
      u_new[i] = u[i] + K * (u[i-1] + u[i+1] - 2*u[i]);
    }
    for(i=1; i<ELENUM-1; i++){
      diff = doubleAbs(u_new[i] - u[i]);
      if(diff > max_diff)
        max_diff = diff;
      // update u
      u[i] = u_new[i];
    }
    //termination                                    
    if(max_diff*3 <= Threshold)
      break;
    else
      max_diff = 0;
	  nsteps --;
  }
  for(i =0 ;i<ELENUM; i++){
    seq_u[i] = u[i];
  }
}

void MPI_Reduce(double * buf, int rank, int count){
#include "mp_proc_diffusion.cvh"
	int nprocs = NPROCS;
	int i=0;
	double receive_buf;
	
	if(rank != 0){
		reduce_send(buf, count, 0, REDUCE, rank);
		reduce_recv(buf, count, 0, REDUCE, rank);
	}
	else{
		for(i=1; i<nprocs; i++){
		reduce_recv(&receive_buf, count, i, REDUCE, rank);
			(*buf) += receive_buf;	
		}
		for(i=1; i<nprocs; i++){
			reduce_send(buf, count, i, REDUCE, rank);
		}
	}
}

void MPI_Process (int rank){
  
  $when (__start);
  double diff, myTotalDiff;        /* the max difference between previous 
				                      function and updated function*/
  double whole_u[ELENUM];
  int nprocs;                      /* number of processes*/
  int start, end;                  /* the index of the chunk*/
  int left_chunk, right_chunk;     /* the index of the left chunk 
                                      and right chunk*/
  int chunk_length;                /* number of elements in this chunk*/
  int temp[4];                     /*temp buffer for start, end, left_chunk and right_chunk */
  int nsteps = NSTEPS;
  nprocs = NPROCS;
	
  chunkProperties(ELENUM, nprocs, rank, temp);
  start = temp[0];
  end = temp[1];
  left_chunk = temp[2];
  right_chunk = temp[3];
  chunk_length = end - start + 1;
  double u[chunk_length + 2];
  double u_new[chunk_length + 2];

  initChunk(u, u_new, start, end, left_chunk, right_chunk);
  /* Jacobi Iterations*/
  while(nsteps > 0){
    communicate(u, left_chunk, right_chunk, start, end, rank);
    diff = update(u, u_new, start, end);
	MPI_Reduce(&diff,rank,1);
    //accumulate diff
//    procsHold(nprocs);
//    total_diff += diff;
//    procsRelease();
//    procsHold(nprocs);
//    myTotalDiff = total_diff;
//    procsRelease();
    if(diff <= Threshold){
      break;
    }
    else{
      diff = 0;
	  nsteps --;
    }
  }
//  procsHold(nprocs);
    combineU(start,end,u,rank, whole_u);
//  procsRelease();
  /* the process with rank 0 takes charge of comparing
     the results of sequential way and MPI way */
  if(rank == 0){
	int i;
	  
	seqDiffusion1d(seq_u);
	for(i=0; i<ELENUM; i++){
	  double test_dif = doubleAbs((seq_u[i] - whole_u[i]));
    }
  }
}