#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

int main(int argc, char** argv) 
{
	int niter,i,j,t,rank;
	double alpha,totaltime,Dt,Stab,DtAlpha;
	int Nx[2];
	double Lx[2], Dx[2], DxDx[2], InvDxDx[2];

	// parameters
	niter     = 8000;
	Nx[0]     = 700; 
	Nx[1]     = 700;
	alpha     = 0.4;
	totaltime = 0.005;

	// initializations computation
	Dt = totaltime/((niter-1)*1.0);
	Lx[0] = 1.0; 
	Lx[1] = 1.0;

	Dx[0] = Lx[0]/((Nx[0]-1)*1.0); 
	Dx[1] = Lx[1]/((Nx[1]-1)*1.0);

	InvDxDx[0] = 1.0/(Dx[0]*Dx[0]); 
	InvDxDx[1] = 1.0/(Dx[1]*Dx[1]);

	DxDx[0] = Dx[0]*Dx[0]; 
	DxDx[1] = Dx[1]*Dx[1];

	Stab    = alpha*Dt*(InvDxDx[0]) + alpha*Dt*(InvDxDx[1]);
	DtAlpha = Dt*alpha;
	
	// verify stability condition
	printf("Stability factor (need < 0.5): %lf\n", Stab);
	if(Stab > 0.5) 
		exit(0);

	double *U0 = malloc((Nx[0]+2) * (Nx[1]+2) * sizeof(double));
	double *U1 = malloc((Nx[0]+2) * (Nx[1]+2) * sizeof(double));

	// init grid
	for(j = 0; j < Nx[1]+2; ++j)
		for(i = 0; i < Nx[0]+2; ++i) 
		{ 
			U0[i+j*Nx[0]] = 1.0;
			U1[i+j*Nx[0]] = 1.0;
			if(i >= 100 && i <= 150 && j>=200 && j<=250) { U0[i+j*Nx[0]] = -10.0; U1[i+j*Nx[0]] = -10.0;}
			if(i >= 300 && i <= 350 && j>=400 && j<=450) { U0[i+j*Nx[0]] =  12.0; U1[i+j*Nx[0]] =  12.0;}
		}

	rank = 0;

	char fileName[2048];
	sprintf(fileName, "IN.%d.dat", rank);

	FILE *file = NULL;
	file = fopen(fileName, "w");
	for(j = 1; j < Nx[0]+1; j += 10)
		for(i = 1; i < Nx[1]+1; i += 10)
			fprintf(file, "%d\t%d\t%lf\n", i, j, U0[i+j*Nx[0]]);
	fclose(file);

	printf("Calculating, please wait...\n");
#pragma omp parallel default(shared) private(t,i,j)
	for(t = 0; t < niter; ++t) 
	{
#pragma omp for schedule(runtime)
		for(j = 1; j < Nx[1]+1; ++j)
			for(i = 1; i < Nx[0]+1; ++i) 
				U1[i+j*Nx[0]] = U0[i+j*Nx[0]] + DtAlpha * ((U0[i+1+j*Nx[0]]   -2.0*U0[i+j*Nx[0]] + U0[i-1+j*Nx[0]]  )*InvDxDx[0] + 
				                                           (U0[i+(j+1)*Nx[0]] -2.0*U0[i+j*Nx[0]] + U0[i+(j-1)*Nx[0]])*InvDxDx[1]);
#pragma omp for schedule(runtime)
		for(j = 1; j < Nx[1]+1; ++j)
			for(i = 1; i < Nx[0]+1; ++i) 
				U0[i+j*Nx[0]] = U1[i+j*Nx[0]] + DtAlpha * ((U1[i+1+j*Nx[0]]   -2.0*U1[i+j*Nx[0]] + U1[i-1+j*Nx[0]]  )*InvDxDx[0] + 
				                                           (U1[i+(j+1)*Nx[0]] -2.0*U1[i+j*Nx[0]] + U1[i+(j-1)*Nx[0]])*InvDxDx[1] );
	}

	sprintf(fileName, "OUT.%d.dat", rank);
	file = fopen(fileName, "w");
	for(j = 1; j < Nx[1]+1; j += 10)
		for(i = 1; i < Nx[0]+1; i += 10)
			fprintf(file, "%d\t%d\t%lf\n", i, j, U0[i+j*Nx[0]]);
	fclose(file);

	free(U0);
	free(U1);

	return 0;
}