#include<stdlib.h>
#include<mpi.h>

#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 
double kappa;

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

int myrank;
int nprocs;
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 */

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(myrank);  
	nxl = countForProc(myrank);
	if (first == 0 || nxl == 0)
		left = MPI_PROC_NULL;
	else
		left = owner(first-1);
	if (first+nxl >= nx || nxl == 0)
		right = MPI_PROC_NULL;
	else
		right = owner(first+nxl);  
	bufSize = nx/nprocs;
	/* if (nx%nprocs != 0) */ bufSize++;
	buf = (double *) malloc(bufSize * sizeof(double));
	v = (double *) malloc((nxl + 2) * sizeof(double));
	v_new = (double *) malloc((nxl + 2) * sizeof(double));
	for (i=1; i<=nxl; i++) v[i] = u0[first+i-1];
	owner = owner(0);
	if (myrank == owner) start=2; else start=1;
	owner = owner(nx-1);
	if (myrank == 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 (myrank != 0) {
		MPI_Send(v+1, nxl, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
	} else { 
		j=0;
		for (source=0; source<nprocs; source++) {
			count = countForProc(source);
			if (source != 0)
				MPI_Recv(buf, count, MPI_DOUBLE, source, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
			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() {
    MPI_Send(v+1, 1, MPI_DOUBLE, left, 0, MPI_COMM_WORLD);
    MPI_Recv(v+(nxl+1), 1, MPI_DOUBLE, right, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    MPI_Send(v+nxl, 1, MPI_DOUBLE, right, 0, MPI_COMM_WORLD);
    MPI_Recv(v, 1, MPI_DOUBLE, left, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
}

/* 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 argc;
	char **argv;
	int iter;
	
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
	MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
	init();
	write(0);
	for (iter = 1; iter <= nsteps; iter++) {
		exchange_ghost_cells();
		update();
		write(iter);
	}
	cleanup();
	MPI_Finalize();
}
