#pragma TASS input 
int NX_BOUND;
#pragma TASS input 
int NSTEPS_BOUND;
#pragma TASS input {nx>=2 && nx<=NX_BOUND} 
int nx;
#pragma TASS input 
double u0[nx];
#pragma TASS input {nsteps>=0 && nsteps<=NSTEPS_BOUND} 
int nsteps;
#pragma TASS input  {dt>0}
double dt;
#pragma TASS input 
double kappa;

#pragma TASS output 
double data[nsteps+1][nx];


int left;      /* rank of left neighbor */
int right;     /* rank of right neighbor on torus */
int nxl;       /* number cells on this process (excluding ghost cells) */
int first;     /* global index for local index 0 */
int start;     /* first local index to update */
int stop;      /* last local index to update */
double* buf;     /* temp. buffer used on proc 0 only */ 
double* v;       /* temperature function */
double* v_new;   /* temp. used to update u */


void MPI_Send(double* buffer, int count, int destination, int tag) {
  double tempBuffer[count];
  int i;
  
  for (i=0; i<count; i++) tempBuffer[i] = buffer[i];
  send(tempBuffer, destination, tag);
}

void MPI_Recv(double* buffer, int count, int source, int tag) {
  double tempBuffer[count];
  int i;
  
  recv(tempBuffer, source, tag);
  for (i=0; i<count; i++) buffer[i] = tempBuffer[i];
}

int owner(int index) { return (NPROCS*(index+1)-1)/nx; }

int firstForProc(int rank) { return (rank*nx)/NPROCS; }

int countForProc(int rank) {
  int a;
  int b;
  
  a = firstForProc(rank+1);
  b = firstForProc(rank);
  return a-b;
}

/* init: initializes global variables. */
void init() { 
  int i;
  int owner;
  int bufSize;
  
  for (i=0; i<nx; i++) ;
  first = firstForProc(PID);  
  nxl = countForProc(PID);
  if (first == 0 || nxl == 0) left = -1; else left = owner(first-1);
  if (first+nxl >= nx || nxl == 0) right = -1; else right = owner(first+nxl);  
  bufSize = nx/NPROCS;
  /* if (nx%NPROCS != 0) */ bufSize++;
  buf = malloc(bufSize * sizeof(double));
  v = malloc((nxl + 2) * sizeof(double));
  v_new = malloc((nxl + 2) * sizeof(double));
  for (i=1; i<=nxl; i++) v[i] = u0[first+i-1];
  owner = owner(0);
  if (PID == owner) start=2; else start=1;
  owner = owner(nx-1);
  if (PID == owner) stop=nxl-1; else stop=nxl;
}

void cleanup() {
  free(buf);
  free(v);
  free(v_new);
}

void write(int time) {
  int source;
  int i;
  int j;
  int count;
  
  if (PID != 0) {
    MPI_Send(v+1, nxl, 0, 0);
  } else { 
    j=0;
    for (source=0; source<NPROCS; source++) {
      count = countForProc(source);
      if (source != 0)
        MPI_Recv(buf, count, source, 0);
      else
        for (i=0; i<count; i++) buf[i] = v[i+1];
      for (i=0; i<count; i++) {
        data[time][j] = buf[i];
        j++;
      }
    }
  }
}

/* exchange_ghost_cells: updates ghost cells using MPI communication */
void exchange_ghost_cells() {
  if (left >= 0) MPI_Send(v+1, 1, left, 0);
  if (right >= 0) {
    MPI_Recv(v+(nxl+1), 1, right, 0);
    MPI_Send(v+nxl, 1, right, 0);
  }
  if (left >= 0) MPI_Recv(v, 1, left, 0);
}

/* update: updates v.  Uses ghost cells.  Purely local operation.  Do not update the
 * first and last global cells: they are the constant boundary cells. */
void update() {
  int i;
  
  for (i=start; i<=stop; i++) v_new[i] = v[i] + kappa*(v[i+1] + v[i-1] - 2*v[i]);
  for (i=start; i<=stop; i++) v[i] = v_new[i];
}

/* main: executes simulation, creates one output file for each time step */
void main() {
  int iter;

  init();
  write(0);
  for (iter = 1; iter <= nsteps; iter++) {
    exchange_ghost_cells();
    update();
    write(iter);
  }
  cleanup();
}
