/*  This c code is a finite-difference solver for the one way wave equation 

    p_t + c p_x = 0 ; p(x,0) = A exp(-(x-x0)^2)   -E < x < E
    
*/

#include <stdio.h>
#include <math.h>

#include <mpi.h>

/*
	This is the Master proc. It reads all the input from stdin and outs 
	the files to times and filexxxx based
	on the input given.
*/

void RunMaster(int size)
{

	float c,A,E,dx,dt;
	int numTimes,numGrids;
	char fileName[256]="file",buf[256]="";
	float *p1,*p2,*temp;
	FILE *fout;
	int i,j,k,length,pos,*lengths;
	int result;
	MPI_Status status,*statuses;
	MPI_Request request,*requests;
	int count;
	
//read in the input values

	printf("Give c : ");
	scanf("%f",&c);
	printf("Give A : ");
	scanf("%f",&A);
	printf("Give E : ");
	scanf("%f",&E);
	printf("Give the grid size dx : ");
	scanf("%f",&dx);
	printf("Give the time step dt : ");
	scanf("%f",&dt);
	printf("Give the number of time steps : ");
	scanf("%d",&numTimes);
	printf("Save matlab data as : ");fflush(stdout);
	//we must flush or we will read the newline just outputed
	scanf("%s",fileName);

	numGrids=2*E/dx;		//calc num of grids
//	printf("The number of grids is %d \n ",numGrids);
	
// Build the time file
	
	fout=fopen("times","w");	//output the times file
	if (fout==NULL)
	{
		perror("times");
		return;
	}
	for(i=0;i<=numTimes;i++)
	{
		fprintf(fout,"%f\n",i*dt);	//output the times into the file
	}
	fclose(fout);
	
// Set up the memory for the pressure
	
	p1=(float *)malloc(sizeof(float)*(numGrids+2));	//create p1,p2
	if (p1==NULL)
	{
		perror("memory");
		return;
	}
	p2=(float *)malloc(sizeof(float)*(numGrids+2));
	if (p2==NULL)
	{
		perror("memory");
		free((void *)p1);
		return;
	}
	
// Broadcast the global values to the nodes

	result=MPI_Bcast((void *)&c,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&A,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&E,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&dx,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&dt,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&numTimes,1,MPI_INT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&numGrids,1,MPI_INT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;

	
	snprintf(buf,255,"%s%08d",fileName,0);
	
	//create filename for the initial output file
	
	fout=fopen(buf,"w");
	if (fout==NULL)
	{
		perror(buf);
		free((void *)p1);
		free((void *)p2);
		return;
	}
	
// The initial conditions

	i=0;	//set up i
	pos=0;
	
	//calc the initial values for p1 array
	
	requests=(MPI_Request *)malloc(sizeof(MPI_Request)*3*(size-1));
	statuses=(MPI_Status *)malloc(sizeof(MPI_Status)*3*(size-1));
	
	lengths=(int *)malloc(sizeof(int)*(size-1));
	
	j=1;
	
// Send grids to each processor (node).

	for(k=1;k<size;k++)
	{
	
		result=MPI_Isend((void *)&j,1,MPI_INT,k,100,MPI_COMM_WORLD,
		&(requests[(k-1)*3+0]));
		if (result!=MPI_SUCCESS) return;
		result=MPI_Isend((void *)&i,1,MPI_INT,k,200,MPI_COMM_WORLD,
		&(requests[(k-1)*3+1]));
		if (result!=MPI_SUCCESS) return;
		
		lengths[k-1]=(numGrids+2)/(size-1);

/*		
		if (k==(size-1) && (k>1)){
			lengths[k-1]=(numGrids+2)%(size-1);
			printf("The number of grids is %d \n ",lengths[k-1]);
		}
*/
		
		result=MPI_Isend((void *)&(lengths[k-1]),1,MPI_INT,k,201,
		MPI_COMM_WORLD,&(requests[(k-1)*3+2]));
		
		if (result!=MPI_SUCCESS) return;
		i+=lengths[k-1];

	}
	
	//MPI_Waitall((size-1)*3,requests,statuses);

// Send p to each processor (node).
	
	for(k=1;k<size;k++)
	{
	
		length=(numGrids+2)/(size-1);

/*		
		if (k==(size-1) && (k>1)){
			length=(numGrids+2)%(size-1);
		}
*/
		
		result=MPI_Recv((void *)&(p1[pos]),length,MPI_FLOAT,k,202,
		MPI_COMM_WORLD,&status);
		if (result!=MPI_SUCCESS) return;
		result=MPI_Get_count(&status,MPI_FLOAT,&count);
		if (result!=MPI_SUCCESS) return;
		
		if (count!=length) return;
		pos+=length;
		
	}
	
	fwrite(p1,sizeof(float),numGrids+2,fout);
	
//	for(i=0;i<(numGrids+2);i++) fprintf(fout,"%f\n",p1[i]);
	fclose(fout);

// Go through the time steps. Send and receive data from the processors (nodes).
	
	for(i=1;i<=numTimes;i++)
	{	//run for the num of time steps
		 
		sprintf(buf,"%s%08d",fileName,i);	//create filename
		fout=fopen(buf,"w");
		if (fout==NULL)
		{
			perror(buf);
			free((void *)p1);
			free((void *)p2);
			return;
		}

// send p
		
		pos=0;
		for(j=1;j<size;j++)
		{
		
			k=2;
			result=MPI_Bsend((void *)&k,1,MPI_INT,j,100,MPI_COMM_WORLD);
			if (result!=MPI_SUCCESS) return;
			
			length=(numGrids+2)/(size-1)+1;
			
/*
			if (j==(size-1)){
			
				length=numGrids+2-pos;
				
			}
*/
			
			result=MPI_Bsend((void *)&length,1,MPI_INT,j,300,
			MPI_COMM_WORLD);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Bsend((void *)&(p1[pos]),length,MPI_FLOAT,j,301,
			MPI_COMM_WORLD);
			if (result!=MPI_SUCCESS) return;
			pos+=length-1;
			
		}
		
		pos=0;

// receive p
		
	        for(j=1;j<size;j++)
		{
	       
	       		length=(numGrids+2)/(size-1);
		
/*	
			if (j==(size-1)){
			
				length=numGrids+2-pos-1;
				
			}
*/
			
	    		result=MPI_Recv((void *)&(p1[pos]),length,MPI_FLOAT,j,302,
			MPI_COMM_WORLD,&status);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Get_count(&status,MPI_FLOAT,&count);
			if (result!=MPI_SUCCESS) return;
			
			if (count!=length) return;
			pos+=length;
	       }

// output p

               fwrite(p1,sizeof(float),numGrids+2,fout);
	       
/*		for(j=0;j<(numGrids+2);j++)
		{
			fprintf(fout,"%f\n",p1[j]);
		}
*/
		
		fclose(fout);	//close the file
	}   //here we go again
	
}

/* This routine is run on processes with rank 1 through n-1. It  

        1: receives the global values
 	1: calculates the initial value
	2: calculates the updated values
	3: exits the slave processes
*/

void RunSlave(void)
{

	MPI_Status status;
	int result;
	float value;
	int method,count;
	float c,A,E,dx,dt;
	int numTimes,numGrids;
	int i,j,length;
	float *p,*p1;

//Get all the global values that are broadcasted

	result=MPI_Bcast((void *)&c,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&A,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&E,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&dx,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&dt,1,MPI_FLOAT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&numTimes,1,MPI_INT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	result=MPI_Bcast((void *)&numGrids,1,MPI_INT,0,MPI_COMM_WORLD);
	if (result!=MPI_SUCCESS) return;
	
//now repeat forever

	while(1){
	
		//get what thing we want to do
		
		result=MPI_Recv((void *)&method,1,MPI_INT,0,100,MPI_COMM_WORLD,
		&status);
		if (result!=MPI_SUCCESS) return;
		result=MPI_Get_count(&status,MPI_INT,&count);
		if (result!=MPI_SUCCESS) return;
		if (count!=1) return;
		
		switch(method)
		{
		case 1:
			//calc initial value
			//get the value of the loop
				
			result=MPI_Recv((void *)&i,1,MPI_INT,0,200,MPI_COMM_WORLD,
			&status);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Get_count(&status,MPI_INT,&count);
			if (result!=MPI_SUCCESS) return;
			if (count!=1) return;
			result=MPI_Recv((void *)&length,1,MPI_INT,0,201,
			MPI_COMM_WORLD,&status);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Get_count(&status,MPI_INT,&count);
			if (result!=MPI_SUCCESS) return;
			if (count!=1) return;
			p1=(float *)malloc(sizeof(float)*length);
			if (p1==NULL) return;
			
// Assign the initial values

			for(j=0;j<length;i++,j++)
			{
			
				p1[j]=A*exp(-1*(pow((i*dx-E),2))); 
			
			}
		
// send back to the host
	
			result=MPI_Bsend((void *)p1,length,MPI_FLOAT,0,202,
			                  MPI_COMM_WORLD);
					  
				//return it to root proc
				
			if (result!=MPI_SUCCESS) return;
			free((void *)p1);
			p1=NULL;
			break;
		case 2:
			//calc the updated value
			//first get the p[j] and p[j+1] values
			
			result=MPI_Recv((void *)&length,1,MPI_FLOAT,0,300,
			MPI_COMM_WORLD,&status);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Get_count(&status,MPI_INT,&count);
			if (result!=MPI_SUCCESS) return;
			if (count!=1) return;
			
			p=(float *)malloc(sizeof(float)*length);
			if (p==NULL) return;

// receive p from the host
		
			result=MPI_Recv((void *)p,length,MPI_FLOAT,0,301,
			MPI_COMM_WORLD,&status);
			if (result!=MPI_SUCCESS) return;
			result=MPI_Get_count(&status,MPI_INT,&count);
			if (result!=MPI_SUCCESS) return;
			if (count!=length) return;
			length--;
			p1=(float *)malloc(sizeof(float)*length);
			if (p1==NULL) return;

// determine the values for p
			
			for(i=0;i<length;i++)
			{
				p1[i]=p[i] + (c*c*dt/dx)*(p[i+1]-p[i]);	
				
			}
			
			free((void *)p);
			
// send back to the host
	
			result=MPI_Bsend((void *)p1,length,MPI_FLOAT,0,302,
			MPI_COMM_WORLD);
			
				//return it
				
			if (result!=MPI_SUCCESS) return;
			free((void *)p1);
			break;
			
		case 3:	//exit if 3
			return;
			break;
			
		default:	//otherwise abort the program
			MPI_Abort(MPI_COMM_WORLD,1);
			return;
			break;

		}		

	}


}


/* Main entry point for the program. It ignores all arguesments passed.*/
int main(int argc,char *argv[])
{

	int size,rank;
	int i,j,result;
	void *ptr;

	MPI_Init(&argc,&argv);			//init MPI
	MPI_Comm_size(MPI_COMM_WORLD,&size);	//get num of procs in session
	MPI_Comm_rank(MPI_COMM_WORLD,&rank);	//get our process rank
	
	ptr=(void *)malloc(1024*1024*10);
	
	MPI_Buffer_attach(ptr,1024*1024*10);
	
	if (rank==0)
	{
				//if we are the root node, run Master Process
		RunMaster(size);
		j=3;			//pass 3 to terminate the child process
		for(i=1;i<size;i++)
		{	//for each slave node
			
		  result=MPI_Send((void *)&j,1,MPI_INT,i,100,MPI_COMM_WORLD);	
		        //send msg
			//if (result!=MPI_SUCCESS) return;
			//ignore errors
		}
	}
	else RunSlave();	//otherwise we are a slave node
	MPI_Finalize();		//cleanup MPI
	return 0;		//exit with 0
}