#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);

	// dynamic allocations
	double **U0 = (double**)malloc((Nx[0]+2) * sizeof(double*));
	for(i = 0; i < Nx[0]+2; ++i)
		U0[i] = (double*)malloc((Nx[1]+2) * sizeof(double));
	
	double **U1 = (double**)malloc((Nx[0]+2) * sizeof(double*));	
	for(i = 0; i < Nx[0]+2; ++i)
		U1[i] = (double*)malloc((Nx[1]+2) * sizeof(double));

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

	rank = 0;

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

	// print inpout
	FILE *file = NULL;
	file = fopen(fileName, "w");
	/*
	for(i = 1; i < Nx[0]+1; i += 10)
		for(j = 1; j < Nx[1]+1; j += 10)
			fprintf(file, "%d\t%d\t%lf\n", i, j, U0[i][j]);
	*/
	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]);
	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(i = 1; i < Nx[0]+1; ++i)
			for(j = 1; j < Nx[1]+1; ++j) 
				U1[i][j] = U0[i][j] + DtAlpha * ((U0[i+1][j  ] -2.0*U0[i][j] + U0[i-1][j  ])*InvDxDx[0] + 
				                                 (U0[i  ][j+1] -2.0*U0[i][j] + U0[i  ][j-1])*InvDxDx[1]);
#pragma omp for schedule(runtime)
		for(i = 1; i < Nx[0]+1; ++i)
			for(j = 1; j < Nx[1]+1; ++j) 
				U0[i][j] = U1[i][j] + DtAlpha * ((U1[i+1][j  ] -2.0*U1[i][j] + U1[i-1][j  ])*InvDxDx[0] + 
				                                 (U1[i  ][j+1] -2.0*U1[i][j] + U1[i  ][j-1])*InvDxDx[1]);
	}

	// print output
	sprintf(fileName, "OUT.%d.dat", rank);
	file = fopen(fileName, "w");
	/*
	for(i = 1; i < Nx[0]+1; i += 10)
		for(j = 1; j < Nx[1]+1; j += 10)
			fprintf(file, "%d\t%d\t%lf\n", i, j, U0[i][j]);
	*/
	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]);
	fclose(file);

	// free buffers
	for(i = 0; i < Nx[0]+2; ++i)
		free(U0[i]);
	free(U0);
	for(i = 0; i < Nx[0]+2; ++i)
		free(U1[i]);
	free(U1);

	return 0;
}