/*

h = 0.0333
dz = 0.000333
s[2] = 0.00425
s[1] = 0.00190
s[0] = 0.00048
s[-1]= 0.00000
s[-2]= 0.00000

dzdzdzS = (0.00425 - .0038)/2dzdzdz = .0002 / dzdzdz = large !

*/

/***

todo : STEPS_R >> STEPS_Z  so that dr = dz roughly
		(i.e. don't make them constants)

make max_megs the cli arg

***/

/* } * { */

/***
 *
 * note : all numbers are in cm, sec. and gm.
 *
 ***/

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

#define smartfree(ptr) if ( ptr == NULL ) ; else { free(ptr); ptr = NULL; }
typedef unsigned long ulong;

#ifndef PI
#define PI ((double)3.14159265358979323846)
#endif

/** constants : **/
#define PAD					(2)
#define GRAVITY				(980.0)
#define CHART_CNT			(20)
#define NUM_STEPS_X			(200)
#define NUM_STEPS_T			(100)

/* } * { */

/** the world: **/
/*	h is height[r]
 *	s is stream function of [r][z]
 *	p is pressure of [r][z]
**/
double *h,**s,**p;		//new
double *h1=NULL,**s1=NULL,**p1=NULL;	//allocs
double *h2=NULL,**s2=NULL,**p2=NULL;
double base_dr,base_dz,dt;

/* } * { */

/** parameters: **/

const double flux = 20; // cm*cm*cm/s
const double jetr = 0.2; // cm
const double visc = 0.01; // water, gm*cm*cm/s

/** computed secondary parameters: **/
double Re,Fr;
double jetv,rjump;
double rmin,rmax,hseed;

/* } * { */

/** protos **/

void conditionBoundary(void);
void conditionTimeZero(void);
void solveNextTmaster(void);
void solveNextT(void);
void saveInfoHeader(FILE *fp,char *name,int length,int timestep);
void saveVector(char *name,double *vector,int length,int timestep);
void saveMatrix(char *name,double **matrix,int length,int step,int timestep);
void saveVelocities(int time);
double *vector(int size);
double **matrix(int sizex,int sizey);
void freevector(double *v);
void freematrix(double **m);
void cleanup(char *mess,int ret);

/*** do it ***/

/* } * { */

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

	/** flux,jetr,visc are the seeds; find other init-values **/

	jetv = flux/(PI*jetr*jetr);

	Re = jetv * jetr / visc;
	Fr = jetv / sqrt(GRAVITY*jetr);

	printf("Reynolds : %f , Froude: %f\n",Re,Fr);

	rjump = pow(Re,1.0/3.0);
	rjump *= jetr; // prediction from trivial Godwin

	rmin = jetr * 3.0;
	rmax = rjump * 3.0;

	hseed = flux / (2*PI*rmin*jetv);

	printf("Godwin rjump : %f\n",rjump);

	printf("rmin: %f , rmax: %f\n",rmin,rmax);

	base_dr = (rmax - rmin) / NUM_STEPS_X;
	base_dz = hseed * 2.0 / NUM_STEPS_X;

	dt = (base_dr / jetv) / 10.0;
		/** do (dr/jetv) to make a dt in the right units
			then take 1/10 cuz dt << dr
		**/

	h1 = vector(NUM_STEPS_X);
	s1 = matrix(NUM_STEPS_X,NUM_STEPS_X);
	p1 = matrix(NUM_STEPS_X,NUM_STEPS_X);
	h2 = vector(NUM_STEPS_X);
	s2 = matrix(NUM_STEPS_X,NUM_STEPS_X);
	p2 = matrix(NUM_STEPS_X,NUM_STEPS_X);

	 h = h1;	 s = s1;	 p = p1;

	conditionTimeZero();

	solveNextTmaster();

	cleanup("exit",0);
	return 0;
}

void cleanup(char *mess,int ret)
{
	freevector(h1);
	freevector(h2);
	freematrix(s1);
	freematrix(s2);
	freematrix(p1);
	freematrix(p2);

	fprintf(stderr,"%s\n",mess);

	exit(ret);
}

/* } * { */

void solveNextTmaster(void)
{
int i;

	i = 0;
	for(;;) {

		saveVector("height",h,NUM_STEPS_X,i);
		saveMatrix("pressure",p,NUM_STEPS_X,10,i);
		saveMatrix("streamF",s,NUM_STEPS_X,10,i);
		saveVelocities(i);

		printf("time step # %d\n",i);

		if ( ++i == NUM_STEPS_T )
			break;

		conditionBoundary();
		solveNextT();
	}

}

/* } * { */

/***********
 *
 * solve next t
 *
 * starts with h,p,s
 *
 * returns h,p,s at new t
 *
*************/

#define dzS(r,z)	((s[r][z+1] - s[r][z-1])/(2.0*dz))
#define drS(r,z)	((s[r+1][z] - s[r-1][z])/(2.0*dr))
#define drdzS(r,z)	( (s[r+1][z+1] - s[r+1][z-1] - s[r-1][z+1] + s[r-1][z-1])/(2.0*dr*dz) )
#define dzdzS(r,z)	( (s[r][z+1]+s[r][z-1]-2*s[r][z])/(dz*dz) )
#define drdrS(r,z)	( (s[r+1][z]+s[r-1][z]-2*s[r][z])/(dr*dr) )
#define drP(r,z)	((p[r+1][z] - p[r-1][z])/(2.0*dr))
#define lapS(r,z)	( drdrS(r,z) + dzdzS(r,z) )
#define drdzdzS(r,z)	( (s[r+1][z+1] - s[r+1][z-1] - s[r-1][z+1] + s[r-1][z-1])/(2.0*dr*dz) )
#define dzdzdzS(r,z) (( s[r][z+2] - s[r][z-2] - 2*(s[r][z+1] - s[r][z-1]))/(2.0*dz*dz*dz))
#define drdrdrS(r,z) (( s[r+2][z] - s[r-2][z] - 2*(s[r+1][z] - s[r-1][z]))/(2.0*dr*dr*dr))

/* } * { */

void solveNextT(void)
{
int ri,zi,zmax;
double r,z,rest;
double *nh,**ns,**np;
double delta,temp;
double drLapS,dzLapS,drdtS,dzdtS;
double dr,dz;

	if ( h == h1 ) {
		nh = h2; ns = s2; np = p2;
	} else {
		nh = h1; ns = s1; np = p1;
	}

	/** first find s at new t **/

	for(ri=0;ri<NUM_STEPS_X;ri++) {
		r = (ri * base_dr) + rmin;
		zmax = h[ri] / base_dz;
	
		temp = r + base_dr;
		dr = temp - r;

		if ( zmax < 0 ) {
			cleanup("negative height found!",0);
		} else if ( zmax >= NUM_STEPS_X ) {
			cleanup("double initial height found!",0);
		}

		for(zi=0;zi<=zmax;zi++) {
			z = zi * base_dz;

			temp = z + base_dz;
			dz = temp - z;

				{
			double dzdzdzS,dzdrdrS;
			dzdzdzS = dzdzdzS(ri,zi);
			dzdrdrS = (drdrS(ri,zi+1) - drdrS(ri,zi-1))/(dz+dz);
			dzLapS =  dzdzdzS + dzdrdrS;
			}

			delta = (1/r)*(	dzS(ri,zi)*dzS(ri,zi)/r
						- dzS(ri,zi)*drdzS(ri,zi) - drS(ri,zi)*dzdzS(ri,zi))
						- drP(ri,zi)
						+ (visc/r)*dzLapS;
			dzdtS = s[ri][zi+1] - s[ri][zi-1] + 2*dz*dt * delta;
			ns[ri][zi+1] = ns[ri][zi-1] + dzdtS;
		}
		rest = ns[ri][zmax];
		for(zi=zmax+1;zi<NUM_STEPS_X;zi++) {
			ns[ri][zi] = rest;
		}
	}

	/** now find p at new t **/

	for(ri=0;ri<NUM_STEPS_X;ri++) {
		r = (ri * base_dr) + rmin;
		zmax = h[ri] / base_dz;

		temp = r + base_dr;
		dr = temp - r;

		for(zi=0;zi<=zmax ;zi++) {
			z = zi * base_dz;

			temp = z + base_dz;
			dz = temp - z;

			drdtS = (ns[ri+1][zi] - ns[ri-1][zi] - s[ri+1][zi] +s[ri-1][zi])/(2*dr*dt);

				{
			double drdzdzS,drdrdrS;
			drdzdzS = (dzdzS(ri+1,zi) - dzdzS(ri-1,zi))/(dr+dr);
			drdrdrS = drdrdrS(ri,zi);
			drLapS =  drdzdzS + drdrdrS;
			}

			delta = ( (dzS(ri,zi)*drS(ri,zi)/(r*r) - 
				(1/r)*dzS(ri,zi)*drdrS(ri,zi) + (1/r)*drS(ri,zi)*drdzS(ri,zi))
				+ GRAVITY + (visc/r)*( drLapS + drS(ri,zi)/(r*r) )
				- drdtS );

			np[ri][zi+1] = np[ri][zi-1] - 2*dz*delta;

		}
		rest = np[ri][zmax];
		for(zi=zmax+1;zi<NUM_STEPS_X;zi++) {
			np[ri][zi] = rest;
		}
	}

	/** now find h at new t **/

	for(ri=0;ri<NUM_STEPS_X;ri++) {
		r = (ri * base_dr) + rmin;
		z = h[ri];
		zi = z / base_dz;

		temp = r + base_dr;
		dr = temp - r;
		temp = z + base_dz;
		dz = temp - z;

			{
		double drS,dzS,drH,delta;
		drS = drS(ri,zi);
		dzS = (s[ri,zi] - s[ri,zi-1])/dz;
		drH = ( h[ri+1] - h[ri-1] ) /(dr+dr);
		delta = (drS + dzS * drH)/r;
		nh[ri] = h[ri] - dt*delta;
		}
	}

	/** swap the pointers **/

	if ( h == h1 ) {
		h = h2; s = s2; p = p2;
	} else {
		h = h1; s = s1; p = p1;
	}

}

/* } * { */

void saveInfoHeader(FILE *fp,char *name,int length,int timestep)
{
fprintf(fp,"# hydraulic jump %s profile\n",name);
fprintf(fp,"#\n");
fprintf(fp,"# time step : %d\n",timestep);
fprintf(fp,"#\n");
fprintf(fp,"# all units are in seconds, cm, and grams.\n");
fprintf(fp,"#\n");
fprintf(fp,"# Reynolds : %0.9f\n",Re);
fprintf(fp,"# Froude : %0.9f\n",Fr);
fprintf(fp,"# flux : %0.9f\n",flux);
fprintf(fp,"# viscocity (Kinematic) : %0.9f\n",visc);
fprintf(fp,"# jetr : %0.9f\n",jetr);
fprintf(fp,"# jetv : %0.9f\n",jetv);
fprintf(fp,"# Godwin jump R : %0.9f\n",rjump);
fprintf(fp,"# hseed = h(rmin) : %0.9f\n",hseed);
fprintf(fp,"# rmin : %0.9f\n",rmin);
fprintf(fp,"# rmax : %0.9f\n",rmax);
fprintf(fp,"# dr : %0.9f\n",base_dr);
fprintf(fp,"# dz : %0.9f\n",base_dz);
fprintf(fp,"# dt : %0.9f\n",dt);
fprintf(fp,"# num_steps : %ld\n",length);
fprintf(fp,"#\n");
}

void saveVector(char *name,double *vector,int length,int timestep)
{
FILE *fp;
char fname[100];
double dr = base_dr;

	sprintf(fname,"dat\\%s.%03d",name,timestep);

	fp = fopen(fname,"wb");
	if ( ! fp ) {
		puts("fopen failed!");
	} else {
		int i;

		saveInfoHeader(fp,name,length,timestep);

		for(i=0;i<length;i++) {
			fprintf(fp,"%7f\t%7f\n",rmin+dr*i,vector[i]);
		}

		fclose(fp);
	}
}

void saveMatrix(char *name,double **matrix,int length,int step,int timestep)
{
FILE *fp;
char fname[100];
double dr = base_dr,dz = base_dz,r;

	sprintf(fname,"dat\\%s.%03d",name,timestep);

	fp = fopen(fname,"wb");
	if ( ! fp ) {
		puts("fopen failed!");
	} else {
		int ri,zi,zmax;

		saveInfoHeader(fp,name,length,timestep);

		for(ri=0;ri<length;ri+=step) {
			r = rmin+dr*ri;
			zmax = h[ri] / dz;
			if ( zmax < 0 ) zmax = 0;
			if ( zmax > NUM_STEPS_X ) zmax = NUM_STEPS_X;
			for(zi=0;zi<zmax;zi+=step) {
				fprintf(fp,"%7f %7f %7f\n",r,dz*zi,matrix[ri][zi]);
			}
		}

		fclose(fp);
	}
}

/*

saveMatrix("u",(1/r)dzS);
saveMatrix("v",(-1/r)drS);

*/
void saveVelocities(int timestep)
{
FILE *fp;
char fname[100];
double dr = base_dr,dz = base_dz;
double u,r,w;
int length = NUM_STEPS_X;
int step = 10,zmax;

	sprintf(fname,"dat\\u.%03d",timestep);
	fp = fopen(fname,"wb");
	if ( ! fp ) {
		puts("fopen failed!");
	} else {
		int ri,zi;

		saveInfoHeader(fp,"u",length,timestep);

		for(ri=0;ri<length;ri+=step) {
			r = rmin+dr*ri;
			zmax = h[ri] / dz;
			if ( zmax < 0 ) zmax = 0;
			if ( zmax > NUM_STEPS_X ) zmax = NUM_STEPS_X;
			for(zi=0;zi<zmax;zi+=step) {
				u = dzS(ri,zi)/r;
				fprintf(fp,"%7f %7f %7f\n",r,dz*zi,u);
			}
		}

		fclose(fp);
	}

	sprintf(fname,"dat\\w.%03d",timestep);
	fp = fopen(fname,"wb");
	if ( ! fp ) {
		puts("fopen failed!");
	} else {
		int ri,zi;

		saveInfoHeader(fp,"w",length,timestep);

		for(ri=0;ri<length;ri+=step) {
			r = rmin+dr*ri;
			zmax = h[ri] / dz;
			if ( zmax < 0 ) zmax = 0;
			if ( zmax > NUM_STEPS_X ) zmax = NUM_STEPS_X;
			for(zi=0;zi<zmax;zi+=step) {
				w = -drS(ri,zi)/r;
				fprintf(fp,"%7f %7f %7f\n",r,dz*zi,w);
			}
		}

		fclose(fp);
	}

}

/* } * { **/

void conditionBoundary(void)
{
double *nh,**ns,**np;
int ri,zi;

	if ( h == h1 ) {
		nh = h2; ns = s2; np = p2;
	} else {
		nh = h1; ns = s1; np = p1;
	}

	ri = -PAD;
	for(;;) {
		if ( ri == 1 )
			ri = NUM_STEPS_X - 1;
		else if ( ri == (NUM_STEPS_X+PAD) )
			break;

		nh[ri] = h[ri];

		zi = -PAD;
		for(;;) {
			if ( zi == 1 )
				zi = NUM_STEPS_X - 1;
			else if ( zi == (NUM_STEPS_X+PAD) )
				break;

			np[ri][zi] = p[ri][zi];
			ns[ri][zi] = s[ri][zi];

			zi++;
		}

		ri++;
	}

}

void conditionTimeZero(void)
{
int ri,zi;
double r,z,height,similar;
double s_pre,h_pre;
double dr = base_dr,dz = base_dz;

/***

fill h,s,p with time=0 conditions

we'll use the zero-viscocity solution

h(r) = jetr*jetr/ 2 * r

S(r,z) = (U/aa)*((rz)^2)*(3-(z/h))

p(r,z) = g*(h(r) - z)

***/

	h_pre = jetr*jetr/2.0;
	s_pre = flux/(4.0*PI);

	for(ri=-PAD;ri<(NUM_STEPS_X+PAD);ri++) {
		r = (ri * dr) + rmin;

		height = h_pre / r;

		h[ri] = height;

		for(zi=-PAD;zi<(NUM_STEPS_X+PAD);zi++) {
			z = zi * dz;
		
			p[ri][zi] = GRAVITY * ( height - z );

			similar = z / height;

			s[ri][zi] = s_pre * similar * similar * ( 3 - similar );
		}
		for(zi=-PAD;zi<=0;zi++) {
			s[ri][zi] = 0;
		}
	}

}

/* } * { */

// vectors and matrices have [-1] PAD extra spaces

double *vector(int size)
{
double *ret;
	if ( (ret = malloc(sizeof(double)*(size+2*PAD))) == NULL ) {
		cleanup("out of memory. exiting",10);
	}
	ret+=PAD;
return ret;
}

double **matrix(int sizex,int sizey)
{
double **ret;
int i;
	if ( (ret = malloc(sizeof(double *)*(sizex+1+2*PAD))) == NULL ) {
		cleanup("out of memory. exiting",10);
	}
	memset(ret,0,sizeof(double *)*(sizex+1+2*PAD));
	ret[0] = (double *)sizex;
	ret++;
	for(i=0;i<(sizex+2*PAD);i++) {
		ret[i] = vector(sizey);
	}
	ret+= PAD;
return ret;
}

void freevector(double *v)
{
if ( v == NULL ) return;
v -= PAD;
free( v );
}

void freematrix(double **m)
{
int i;
int sizex;
if ( m == NULL ) return;
	m -= PAD;
	sizex = (int)m[-1];
	for(i=0;i<(sizex + PAD * 2);i++) {
		freevector(m[i]);
	}
	m --;
	free( m );
}

/* } */