/*
 *    It evolves the equation:
 *                            u,t + u,x + u,y = 0
 *    Using a Lax scheme.
 *    The initial data is a cruddy gaussian.
 *    Boundaries are flat: copying the value of the neighbour
 */
 
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>

#define NX 100
#define NY 100
#define LX 2.0
#define LY 2.0

#define sigma 0.1
#define tmax 100



/* conversions from discrete to real coordinates: can be better */
float ix2x(int ix){
    return ix*LX/NX-LX/2; 
}


/*
 * THIS CHANGES: the every processor has a different offset
 * I pass also the proc_y and I do not use proc_me directy because
 * saving the file I use this parameter to evaluate the real y coordinate
 * for all the data of other processors that are saved by 0 however.
 */

float iy2y(int iy, int NLY, int proc_y){ 
    return iy*LY/NY-LY/2 + proc_y*NLY*LY/NY;  
}

/* 
 * initialize the system with a gaussian temperature distribution 
 */

int init_transport(float *temp, int NLY, int proc_me){
    
    int ix,iy;
    float x,y;

    for(iy=1;iy<NLY+1;++iy)
	for(ix=1;ix<NX+1;++ix){	
	    x=ix2x(ix);
	    y=iy2y(iy, NLY, proc_me);
	    temp[((NX+2)*iy)+ix] = tmax*exp((-((x*x)+(y*y)))/(2.0*(sigma*sigma)));
	}
    return 0;
}

/*
 * save the temperature distribution
 * the ascii format is suitable for splot gnuplot function
 */

int save_gnuplot(char *filename, float *temp, int NLY, int proc_me, int nprocs){

    float *buf;
    int ix, iy, iproc=0;
    FILE *fp;

    MPI_Status status;

    if(proc_me == 0){

	buf = (float *) malloc ((NX+2)*(NLY+2)*sizeof(float));
	fp = fopen(filename, "w");

	for(iy=1;iy<=NLY;++iy){		
	    for(ix=1;ix<=NX;++ix){
		fprintf(fp, "\t%f\t%f\t%g\n", ix2x(ix), iy2y(iy, NLY, proc_me), temp[((NX+2)*iy)+ix]);
	    }
	    fprintf(fp, "\n");
	}

	while(++iproc<nprocs){
	    MPI_Recv(buf, (NX+2)*(NLY+2), MPI_REAL, iproc, 0, MPI_COMM_WORLD, &status);
	    for(iy=1;iy<=NLY;++iy){		
		for(ix=1;ix<=NX;++ix){
		    fprintf(fp, "\t%f\t%f\t%g\n", ix2x(ix), iy2y(iy, NLY, iproc), buf[((NX+2)*iy)+ix]);
		}
		fprintf(fp, "\n");
	    }
	}
	fclose(fp);
    }
    else{
	MPI_Send(temp, (NX+2)*(NLY+2), MPI_REAL, 0, 0, MPI_COMM_WORLD);
    }

    return 0;
}

int update_boundaries_FLAT(float *temp, int NLY, int nprocs, int proc_me, int proc_up, int proc_down){

    MPI_Status status, status1;
    int iy=0, ix=0;
    
    for(iy=1;iy<=NLY;++iy){
	temp[(NX+2)*iy] = temp[((NX+2)*iy)+1];
	temp[((NX+2)*iy)+(NX+1)] = temp[((NX+2)*iy)+NX];
    }

/*  
 * only the lowest has the lower boundary condition  
 */
    if (proc_me==0) for(ix=0;ix<=NX+1;++ix) temp[ix] = temp[(NX+2)+ix];
    
/*
 * only the highest has the upper boundary condition  
 */
    if (proc_me==nprocs-1) for(ix=0;ix<=NX+1;++ix) temp[((NX+2)*(NLY+1))+ix] = temp[((NX+2)*(NLY))+ix];


/*  
 *communicate the ghost-cells  
 *
 *  
 * lower-down  
 */
    MPI_Sendrecv(&temp[(NX+2)+1], NX, MPI_REAL, proc_down, 0, &temp[((NX+2)*(NLY+1))+1], NX, MPI_REAL, proc_up, 0, MPI_COMM_WORLD, &status);


/*
 * higher-up  
 */
    MPI_Sendrecv(&temp[((NX+2)*(NLY))+1], NX, MPI_INTEGER, proc_up, 0, &temp[1], NX, MPI_INTEGER, proc_down, 0, MPI_COMM_WORLD, &status1);

    return 0;
}


int evolve(float dtfact, float *temp, float *temp_new, int NLY){
    
    float dx, dt;
    int ix, iy;
    float temp0;
    
    dx = 2*LX/NX;
    dt = dtfact*dx/sqrt(3.0);
    
    for(iy=1;iy<=NLY;++iy)
	for(ix=1;ix<=NX;++ix){
	    temp0 = temp[((NX+2)*iy)+ix];
	    temp_new[((NX+2)*iy)+ix] = temp0-0.5*dt*(temp[((NX+2)*(iy+1))+ix]-temp[((NX+2)*(iy-1))+ix]+temp[((NX+2)*iy)+(ix+1)]-temp[((NX+2)*iy)+(ix-1)])/dx;
	}
    
    for(iy=0;iy<NLY+2;++iy)
	for(ix=0;ix<NX+2;++ix)
	    temp[((NX+2)*iy)+ix] = temp_new[((NX+2)*iy)+ix];
    
    return 0;
}


int main(int argc, char* argv[]){

    int nprocs, proc_me, proc_up, proc_down;

    int i=0, NLY;
    float *temp, *temp_new;


    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &proc_me);

/*  
 *  all the communications from/to MPI_PROC_NULL do nothing
 */

    proc_down = proc_me-1;
    proc_up = proc_me+1;

    if(proc_down < 0)
	proc_down = MPI_PROC_NULL;
    
    if(proc_up == nprocs) proc_up = MPI_PROC_NULL;

    NLY = NY/nprocs;
    
    temp = (float *) malloc ((NX+2)*(NLY+2)*sizeof(float));
    temp_new = (float *) malloc ((NX+2)*(NLY+2)*sizeof(float));

    init_transport(temp, NLY, proc_me);
    update_boundaries_FLAT(temp, NLY, nprocs, proc_me, proc_up, proc_down);

    save_gnuplot("transport.dat", temp, NLY, proc_me, nprocs);   

    for(i=1;i<=500;++i){
	evolve(0.1, temp, temp_new, NLY);
	update_boundaries_FLAT(temp, NLY, nprocs, proc_me, proc_up, proc_down);
    }
    
    save_gnuplot("transport_end.dat", temp, NLY, proc_me, nprocs);   

    MPI_Finalize();

/*    return 0; */
}
