/************************************************************************************************
*												*
*	rb.c											*
*												*
*												*
*       Solves the rayleigh-benard equation in bousinesq approximation.				*
*       Pseudospectral scheme for spatial derivatives,						*
*       RK4 for temporal integration.								*
*       Volume penalization for no-slip boundary conditions.					*
*												*
*												*
*	compile with										*
*		gcc -std=c99 -O3 -Wall -o rb rb.c -I(FFTW3_PATH)/include			*
*			-L(FFTW3_PATH)/lib -lfftw3 -I(DISLIN_PATH) -L(DISLIN_PATH) -ldislin -lm	*
*												*
*												*
*	author: Johannes Luelff	(johannes.luelff@uni-muenster.de)				*
*												*
************************************************************************************************/


#include <stdio.h>
#include <stdlib.h>
#include <complex.h>
#include <math.h>
#include <time.h>
#include "fftw3.h"
#include "dislin.h"

#define USE_DISLIN

//constants etc.
#define PI 3.14159265358979323846264338327950288419716939937510582097494459230781
const int N = 256;		//Grid resolution
const double L = 2*PI;		//Physical length

//Rayleigh-Benard parameters
const double nu = 1e-2;
const double kappa = 1e-2;
const double alpha = 1;
const double beta = 1;

//time(step)
double t;
double dt=2e-3;

const int MAX_ITER = 40*1000;
const int EVERY_DISPLAY = 40;

//qvec is the statevector; qvec[0] is the vorticity field, qvec[1] the temperature field
complex *qvec[2];
double *qvec_real[2];
complex *ux, *uy;
complex *nonlin[2], *Nx[2], *Ny[2];
complex *rk1[2], *rk2[2], *rk3[2], *rk4[2], *rktmp[2];
double *kvec[3];
float *dis_real, *dis_spec;

//penalization terms
double *wampe;
const double wampe_eps = 1e-2;

//fftw plans
fftw_plan plan_r2c, plan_c2r;


//two routines to handle the FT
void fftR2C(complex *arr) {
	fftw_execute_dft_r2c(plan_r2c, (double*)arr, (complex*)arr);
	for(int i=0; i<N*(N/2+1); i++)
		arr[i] /= sqrt(N*N);
}
void fftC2R(complex *arr) {
	fftw_execute_dft_c2r(plan_c2r, (complex*)arr, (double*)arr);
	for(int i=0; i<N*(N/2+1); i++)
		arr[i] /= sqrt(N*N);
}

//some routines for 'array arithmetic'
void assignArrays(complex *out, complex *in) {
	for(int i=0; i<N*(N/2+1); i++)
		out[i] = in[i];
}
void add2Arrays2(complex *out[2], complex *in1[2], double s, complex *in2[2]) {
	for(int i=0; i<N*(N/2+1); i++) {
		out[0][i] = in1[0][i] + s*in2[0][i];
		out[1][i] = in1[1][i] + s*in2[1][i];
	}
}
void add2Arrays(complex *out, complex *in1, double s, complex *in2) {
	for(int i=0; i<N*(N/2+1); i++)
		out[i] = in1[i] + s*in2[i];
}

//dealiases the field using Orszag's 2/3-rule
void dealiase(complex *arr) {
	for(int i=0; i<N*(N/2+1); i++)
		if(kvec[2][i] > pow(2.0/3.0 * (N/2+1), 2))
			arr[i] = 0;
}

//calculates nonlinear part of our equations
void calcNonlin(complex *in[2], complex *out[2]) {
	//calc velocity
	ux[0] = 0;
	uy[0] = 0;
	for(int i=1; i<N*(N/2+1); i++) {
		ux[i] =  I*kvec[1][i]*in[0][i] / kvec[2][i];
		uy[i] = -I*kvec[0][i]*in[0][i] / kvec[2][i];
	}
	fftC2R(ux);
	fftC2R(uy);

	//calc vorticity part in qvec[0]
	assignArrays((complex*)qvec_real[0], in[0]);
	fftC2R((complex*)qvec_real[0]);
	for(int i=0; i<N*2*(N/2+1); i++) {
		((double*)Nx[0])[i] =  ((double*)uy)[i] * qvec_real[0][i] - 1.0/wampe_eps * wampe[i] * ((double*)ux)[i];
		((double*)Ny[0])[i] = -((double*)ux)[i] * qvec_real[0][i] - 1.0/wampe_eps * wampe[i] * ((double*)uy)[i];
	}
	fftR2C(Nx[0]);
	fftR2C(Ny[0]);
	for(int i=0; i<N*(N/2+1); i++)
		out[0][i] = I * (kvec[0][i]*Ny[0][i] - kvec[1][i]*Nx[0][i]);
	dealiase(out[0]);


	//calc temperature T
	assignArrays((complex*)qvec_real[1], in[1]);
	fftC2R((complex*)qvec_real[1]);
	for(int i=0; i<N*2*(N/2+1); i++) {
		((double*)Nx[1])[i] =  ((double*)uy)[i] * qvec_real[1][i];
		((double*)Ny[1])[i] = -((double*)ux)[i] * qvec_real[1][i];
	}
	fftR2C(Nx[1]);
	fftR2C(Ny[1]);
	for(int i=0; i<N*(N/2+1); i++)
		out[1][i] = I * (kvec[0][i]*Ny[1][i] - kvec[1][i]*Nx[1][i]);

	//wampe for temperature T
	for(int i=0; i<N*2*(N/2+1); i++)
		qvec_real[1][i] = 1.0/wampe_eps * wampe[i] * qvec_real[1][i];
	fftR2C((complex*)qvec_real[1]);
	for(int i=0; i<N*(N/2+1); i++)
		out[1][i] += -((complex*)qvec_real[1])[i];
	dealiase(out[1]);
}


//solves the whole RHS of the two equation for temperature and velocity field
void solveRHS(complex *in[2], complex *out[2], double t) {
        //nonlinearity
	calcNonlin(in, nonlin);

        //linear part
	for(int i=0; i<N*(N/2+1); i++) {
                //calculate y-component of velocity field
		complex uy;
		if(kvec[2][i]!=0)
			uy = (-I*kvec[0][i]*in[0][i] / kvec[2][i]);
		else
			uy = 0;

		out[0][i] = nonlin[0][i] -    nu * kvec[2][i] * in[0][i] + alpha*I*kvec[0][i]*in[1][i];
 		out[1][i] = nonlin[1][i] - kappa * kvec[2][i] * in[1][i] + beta * uy;
	}
}


//calculate next time step via RK4 scheme
void doTimeStep(complex *qvec[2]) {
	//k1
	solveRHS(qvec, rk1, t);
	//k2
	add2Arrays2(rktmp, qvec, 0.5*dt, rk1);
	solveRHS(rktmp, rk2, t+0.5*dt);
	//k3
	add2Arrays2(rktmp, qvec, 0.5*dt, rk2);
	solveRHS(rktmp, rk3, t+0.5*dt);
	//k4
	add2Arrays2(rktmp, qvec, dt, rk3);
	solveRHS(rktmp, rk4, t+dt);
	//all together
	for(int i=0; i<N*(N/2+1); i++) {
		qvec[0][i] += dt/6.0 * (rk1[0][i] + 2*rk2[0][i] + 2*rk3[0][i] + rk4[0][i]);
		qvec[1][i] += dt/6.0 * (rk1[1][i] + 2*rk2[1][i] + 2*rk3[1][i] + rk4[1][i]);
	}

	t += dt;
}


void display(complex *in1, complex *in2) {
	double min1 = HUGE_VAL;
	double max1 = -HUGE_VAL;
	double min2 = HUGE_VAL;
	double max2 = -HUGE_VAL;

	complex *ctmp = rktmp[0];

	assignArrays(ctmp, in1);
	fftC2R(ctmp);
	for(int i=0; i<N; i++) {
		for(int j=0; j<N; j++) {
			dis_real[i*N+j] = ((double*)ctmp)[i*2*(N/2+1)+j];
			if(dis_real[i*N+j]<min1)
				min1 = dis_real[i*N+j];
			if(dis_real[i*N+j]>max1)
				max1 = dis_real[i*N+j];
		}
	}

	assignArrays(ctmp, in2);
	fftC2R(ctmp);
	for(int i=0; i<N; i++) {
		for(int j=0; j<N; j++) {
			dis_spec[i*N+j] = ((double*)ctmp)[i*2*(N/2+1)+j];
			if(dis_spec[i*N+j]<min2)
				min2 = dis_spec[i*N+j];
			if(dis_spec[i*N+j]>max2)
				max2 = dis_spec[i*N+j];
		}
	}

	if(min1==max1) {
		min1 -= 1e-3;
		max1 += 1e-3;
	}
	erase();
	axspos(2970*0.1f,2970*0.85f);
	autres(N,N);
	titlin("omega", 1);
	char s[256];
	sprintf(s, "t=%f", t);
	titlin(s, 3);
	graf3(0,L,0,L/2, 0,L,0,L/2, min1, max1, min1, (max1-min1)/2);
	title();
	crvmat(dis_real, N, N, ((512/N)<=0)?1:(512/N), ((512/N)<=0)?1:(512/N));
	endgrf();

	if(min2==max2) {
		min2 -= 1e-3;
		max2 += 1e-3;
	}
	axspos(2970*1.1f,2970*0.85f);
	autres(N,N);
	titlin("Theta", 1);
	sprintf(s, "dt=%f", dt);
	titlin(s, 3);
	graf3(0,L,0,L/2, 0,L,0,L/2, min2, max2, min2, (max2-min2)/2);
	title();
	crvmat(dis_spec, N, N, ((512/N)<=0)?1:(512/N), ((512/N)<=0)?1:(512/N));
	endgrf();
}


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

        //allocate memory
	for(int i=0; i<2; i++) {
		qvec[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		qvec_real[i] = (double*)fftw_malloc(N*2*(N/2+1) * sizeof(double));
		nonlin[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		Nx[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		Ny[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		rk1[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		rk2[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		rk3[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		rk4[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
		rktmp[i] = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
	}
	ux = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
	uy = (complex*)fftw_malloc(N*(N/2+1) * sizeof(complex));
	for(int i=0; i<3; i++)
		kvec[i] = (double*)fftw_malloc(N*(N/2+1) * sizeof(double));
	dis_real = (float*)fftw_malloc(N*N*sizeof(float));
	dis_spec = (float*)fftw_malloc(N*N*sizeof(float));

	wampe = (double*)fftw_malloc(N*2*(N/2+1)*sizeof(double));

        //create plans
	plan_c2r = fftw_plan_dft_c2r_2d(N, N, qvec[0], (double*)qvec[0], FFTW_ESTIMATE);
	plan_r2c = fftw_plan_dft_r2c_2d(N, N, (double*)qvec[0], qvec[0], FFTW_ESTIMATE);


	//init kvec
	for(int i=0; i<N; i++) {
		for(int j=0; j<N/2+1; j++) {
			kvec[0][i*(N/2+1)+j] = 2*PI*((i<N/2+1) ? (i) : (i-N)) / L;
			kvec[1][i*(N/2+1)+j] = 2*PI*j / L;
			kvec[2][i*(N/2+1)+j] = kvec[0][i*(N/2+1)+j]*kvec[0][i*(N/2+1)+j] + kvec[1][i*(N/2+1)+j]*kvec[1][i*(N/2+1)+j];
		}
	}

	//init wampe
	for(int i=0; i<N; i++) {
		for(int j=0; j<N; j++) {
			wampe[i*2*(N/2+1)+j] = ((j>1*N/8)&&(j<7*N/8)) ? 0.0 : 1.0;	//bottom / top
		}
	}

	//init field
	for(int i=0; i<N; i++) {
		for(int j=0; j<N; j++) {
			if(wampe[i*2*(N/2+1)+j] == 0.0) {
				//white noise
				//((double*)qvec[0])[i*2*(N/2+1)+j] =  0;
				//((double*)qvec[1])[i*2*(N/2+1)+j] =  10 + 1*(1-2*rand()/(double)RAND_MAX);

				//raylor-tayleigh instability
				((double*)qvec[0])[i*2*(N/2+1)+j] =  0;
				((double*)qvec[1])[i*2*(N/2+1)+j] = ((j<N/2)?(10):(-10)) + 0.01*(1-2*rand()/(double)RAND_MAX);

				//linear; theta=linear means T=const
				//((double*)qvec[0])[i*2*(N/2+1)+j] =  0;
				//((double*)qvec[1])[i*2*(N/2+1)+j] = beta*(j/(double)N-0.5) + 0.01*(1-2*rand()/(double)RAND_MAX);
			} else {
				((double*)qvec[0])[i*2*(N/2+1)+j] = 0;
				((double*)qvec[1])[i*2*(N/2+1)+j] = 0;
			}
		}
	}
	fftR2C(qvec[0]);
	dealiase(qvec[0]);
	qvec[0][0] = 0;
	fftR2C(qvec[1]);
	dealiase(qvec[1]);
	qvec[1][0] = 0;


#ifdef USE_DISLIN
	//init dislin
	metafl("XWIN");
	scrmod("REVERS");
	winsiz(1600,800);
	page(2970*2,2970);
	disini();
	ax3len(2970*0.75f, 2970*0.75f, 2970*0.75f);
#endif //USE_DISLIN


        //main loop
	for(int iter=0; iter<=MAX_ITER; iter++) {

		if(iter%EVERY_DISPLAY == 0) {
			display(qvec[0], qvec[1]);
			printf("t = %f , dt = %f\n", t, dt);
		}

		doTimeStep(qvec);
	}

	//clean up
	fftw_destroy_plan(plan_c2r);
	fftw_destroy_plan(plan_r2c);

	for(int i=0; i<2; i++) {
		fftw_free(qvec[i]);
		fftw_free(qvec_real[i]);
		fftw_free(nonlin[i]);
		fftw_free(Nx[i]);
		fftw_free(Ny[i]);
		fftw_free(rk1[i]);
		fftw_free(rk2[i]);
		fftw_free(rk3[i]);
		fftw_free(rk4[i]);
		fftw_free(rktmp[i]);
	}
	fftw_free(ux);
	fftw_free(uy);
	for(int i=0; i<3; i++)
		fftw_free(kvec[i]);
	fftw_free(dis_real);
	fftw_free(dis_spec);
	fftw_free(wampe);

#ifdef USE_DISLIN
	//finalize dislin
	disfin();
#endif //USE_DISLIN

	printf("READY\n");

}