#include"tutorial.h"
#include<mpi.h>

double deltat;
int rank, size;

//Read from standard input
void load_stdin(double *p, double *v, double *m, int n)
{
	int i,j;
	double *pbuf = (double *)malloc(3*sizeof(double)*n/size);
	double *vbuf = (double *)malloc(3*sizeof(double)*n/size);
	double *mbuf = (double *)malloc(sizeof(double)*n/size);
	
	//Process 0 does this.	
	if(rank == 0)
	{
		// Read from file. The first n/size chunk should stay on process 0. Do not send it to 
		// another process.

		for(i = 0; i < n/size; i++)
		{
			scanf(INPUT_BODY,&pbuf[i*3],&pbuf[i*3+1],&pbuf[i*3+2],&vbuf[3*i],&vbuf[3*i+1],&vbuf[3*i+2],&mbuf[i]);
		}
		for(i = 0; i < 3*n/size; i++)
		{
			p[i] = pbuf[i];
			v[i] = vbuf[i];
		}

		for(i = 0; i < n/size; i++)
		{
			m[i] = mbuf[i];
		}
		
		// The remaining chunks are sent to processes 1 to size-1.
		for(j = 1; j < size ; j++)
		{	
			for(i = 0; i < n/size; i++)
			{
				scanf(INPUT_BODY,&pbuf[i*3],&pbuf[i*3+1],&pbuf[i*3+2],&vbuf[3*i],&vbuf[3*i+1],&vbuf[3*i+2],&mbuf[i]);
			}
			
			// Individual sends for position, velocity and mass.	
			MPI_Send(pbuf,3*n/size,MPI_DOUBLE,j,0,MPI_COMM_WORLD);
			MPI_Send(vbuf,3*n/size,MPI_DOUBLE,j,0,MPI_COMM_WORLD);
			MPI_Send(mbuf,n/size,MPI_DOUBLE,j,0,MPI_COMM_WORLD);
		}
	}
	// Other processes (not 0) do this. Corresponding "receives" for the above "sends".
	else
	{
		MPI_Status status;
		MPI_Recv(p,3*n/size,MPI_DOUBLE,0,0,MPI_COMM_WORLD,&status);
		MPI_Recv(v,3*n/size,MPI_DOUBLE,0,0,MPI_COMM_WORLD,&status);
		MPI_Recv(m,n/size,MPI_DOUBLE,0,0,MPI_COMM_WORLD,&status);
	}
	
	// Every processes now has it's data.

}

// Read from a file pointer. I did not talk about this during the tutorial. Instead of reading from
// standard input, we read from a specific file using file pointers. The advantage of this is that
// it leaves the stdin free for other inputs. The file is opened using fopen, which returns you 
// a file pointer. Then, instead of using scanf/printf functions (which read/write from/to standard input/output)
// we use corresponding fscanf and fprintf. 

// WARNING - I wrote these (file pointer) functions to demonstrate the alternative way of reading, but
// they are not complete. You can see that they do not involve the vector distribution to different 
// processes. You could complete the functions for practice. :)

void load_file(double *p, double *v, double *m, int n)
{
	int i;
	FILE *fp = fopen("input.txt","r");

	for(i = 0; i < n; i++)
	{
		fscanf(fp,INPUT_BODY,&p[i*3],&p[i*3+1],&p[i*3+2],&v[3*i],&v[3*i+1],&v[3*i+2],&m[i]);
	}
}

//Write to the standard output. You will see this on the screen
void print_stdout(double *p, double *v, double *m, int n)
{
	int i,j;
	
	// Process 0 does this
	if(rank == 0)
	{
		// Print out the first n/size chunk.
		for(i = 0; i < n/size; i++)	
		{
			printf(OUTPUT_BODY,p[i*3],p[i*3+1],p[i*3+2],v[3*i],v[3*i+1],v[3*i+2],m[i]);
		}
		
		MPI_Status status;

		
		double *pbuf = (double *)malloc(3*sizeof(double)*n/size);
		
		// Receive the remaining chunks from other processes. And then print them.
		for(j = 1; j < size ; j++)
		{	

			MPI_Recv(pbuf,3*n/size,MPI_DOUBLE,j,0,MPI_COMM_WORLD,&status);

			for(i = 0; i < n/size; i++)	
			{
				printf(OUTPUT_BODY,pbuf[i*3],pbuf[i*3+1],pbuf[i*3+2],v[3*i],v[3*i+1],v[3*i+2],m[i]);
			}
			
		}
	}
	else
	{
		MPI_Send(p,3*n/size,MPI_DOUBLE,0,0,MPI_COMM_WORLD);
	}

}

//Write into a file
void print_file(double *p, double *v, double *m, int n)
{
	
	//When opened in the write mode, the file is made blank
	FILE *fp = fopen("output.txt","w");
	int i;

	for(i = 0; i < n; i++)	
	{
		fprintf(fp,OUTPUT_BODY,p[i*3],p[i*3+1],p[i*3+2],v[3*i],v[3*i+1],v[3*i+2],m[i]);
	}

}

// Just a simple function which updates the body positions based on velocities. 
void update_position(double *p, double *v, double *m, int n)
{
	int i;
	
	for(i = 0; i < n; i++)
	{
		p[i] += v[i]*deltat;
	}
}

int main(int argc, char **argv)
{
	// Reading command line arguments
	int i;
	int n;
	
	n = atoi(argv[1]);
	int sw = atoi(argv[2]);
	
	// Initialization
	MPI_Init(&argc,&argv);

	// Determine which process is running the code, and the total number of processes spawned. 
	MPI_Comm_rank(MPI_COMM_WORLD,&rank);
	MPI_Comm_size(MPI_COMM_WORLD,&size);

	
	printf("I am process %d from among %d processes\n",rank,size);


	deltat = 0.1;
	// Allocate arrays for position, velocity and mass
	if((n%size))
	{	
		printf("Number of bodies is not a multiple of the number of processors\n");
		exit(-1);
	}

	double *p = (double *)malloc(3*sizeof(double)*n/size);
	double *v = (double *)malloc(3*sizeof(double)*n/size);
	double *m = (double *)malloc(sizeof(double)*n/size);

	
	// Switch = 0; Read from standard input	and write to standard output
	if(sw == 0)
	{
		load_stdin(p,v,m,n);
		//Do some computations on p and v
		update_position(p,v,m,3*n/size);
		print_stdout(p,v,m,n);
	}
	// Switch = 1; Read from a file and write to another 
	else
	{
		load_file(p,v,m,n);
		//Do some computations on p and v
		print_file(p,v,m,n);
	}
	
	// Dont forget the finalize function when you are done. 
	MPI_Finalize();
	return 0;

}