/*{
*/

#define DT_SCALE (1)		// (dt*(du/dt)/u) = 1/DT_SCALE
#define HDT_SCALE (10)		// hdt = dt * HDT_SCALE , h[new] = n[old] + hdt * (dh/dt)

#define FORCE_FLUX_H			// scale h to conserve Q
//#define FORCE_FLUX_U			// scale u to conserve Q
//#define WALK_H				// else use dh/dt

/*

four ways: (ff = FORCE_FLUX_U, wh = walk_height)

ff = 0 wh = 0
ff = 1 wh = 0
ff = 0 wh = 1
ff = 1 wh = 1	-> bad!

they give slightly different results, but all seem to do very little after 100,000 steps

next todo : feed a closer fit of h(r)

walk_h seems to be ugly

********

# dt : 0.0000000000040074

this is 249538354045 steps in a second

TOO SMALL!

todo : solve dt(h) by r-integration instead

---

hydr3 : like hydr2, but on a rectangular domain

z is rescaled to y = z/h , 0<y<1

we approximate h as nearly constant on the scale of dr, so that d/dr is still a difference
	at the same y (it should be a difference at constant z = y*h, so slightly different y
	values).

we also approximate : after a dh/dt changes the scaling of y, we don't change u(y)

hydr2 :

the approximation is :

	set P = pg(h-z)

	this overspecifies u and w, so now we drop the Navier-Stokes in w.

	note that w is NOT zero or "directly" connected to u now.

The equations of motion are:

	incompressibility : dz(w) = (-1/r)*dr(ru)

	navier-stokes : (dt + udr + wdz)u = v(lapr(u) + dzdz(u)) - g dr(h)
		
		lapr(u) = drdr(u) + drU/r - u/r^2 = dr(1/r(dr(ru)))

	boundary evolution : dt(h) = w@h - u@h dr(h)

**/

/***
 *
 * 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	1
const double gravity = 980.0;

/** CLI options **/

long save_steps_t = 10000;
long steps_t = 1<<30;
long max_memory = 2; // megs

/* the world: **/
/*	h is height[r]
 *	u and w are velocites [r][z]
**/
double *h,**u,**w;
double *h1=NULL,**u1=NULL,**w1=NULL;	//allocs
double *h2=NULL,**u2=NULL,**w2=NULL;
double dr,dt,hdt,otwodr,otwodz,odzdz,odrdr;
double *fluxr;
double maxdhdt,maxdudt;

/*}{ parameters: **/

const double flux = 10; // cm*cm*cm/s
const double jetr = 0.1; // cm
const double visc = 0.01; // water, gm*cm*cm/s

/** computed secondary parameters: **/
double Re,Fr;
double jetv,rjump,hjump;
double rmin,rmax,hrmin,hrmax;
long steps_r,steps_z,topz;

/* protos **/

void dbf(void) { };

void solveU(double **nu);
void forceFluxU(double **u);
void forceFluxH(double **u);
void solveW(double **nw);
void solveH_dhdt(double *nh);
void solveH_dhdr(double *nh);
void conditionBoundary(void);
void conditionTimeZero(void);
void solveNextTmaster(void);
void solveNextT(void);
void saveState(int time);
void saveInfoHeader(FILE *fp,char *name,int timestep);
void saveVector(char *name,double *vector,int length,int timestep);
void saveMatrix(char *name,double **matrix,int timestep);
void saveMatrices(char *name,double **matrix1,double **matrix2,int timestep);
void zintegrate(double *integralptr,double **f);
double *vector(int size);
double **matrix(int sizex,int sizey);
void freevector(double *v);
void freematrix(double **m);
void cleanup(char *mess,int ret);

void writestr_init(FILE *fp);
void writestr_flush(void);
void writestr(char *str);

/*
 *
 * solve next t
 *
 * starts with h,u,w
 *
 * returns h,u,w at new t
 *
*************/

/*}{*/

void solveNextT(void)
{
double *nh,**nu,**nw;

	if ( h == h1 ) {
		nh = h2; nu = u2; nw = w2;
	} else {
		nh = h1; nu = u1; nw = w1;
	}

	conditionBoundary();

	solveU(nu); u = nu;

#ifdef FORCE_FLUX_U
	forceFluxU(u);
#else
#ifdef FORCE_FLUX_H
	forceFluxH(u);
#endif
#endif

	solveW(nw);	w = nw;

#ifdef WALK_H
	solveH_dhdr(nh); h = nh;
#else
	solveH_dhdt(nh); h = nh;
#endif

}

/*}{*/

void solveU(double **nu)
{
int ri,zi;
int diffstoobig = 0;

	/*
		find u at new time.  The equation is:

		dt(u) = - (udr + wdz)u -v(lapr(u) + dzdz(u)) - g dr(h)
	*/

	maxdudt = 0;

	for(ri=0;ri<steps_r;ri++) {
		double drH,r,dz,otwodz,odzdz;
		double pressure;

		r = ri*dr + rmin;
		dz = h[ri]/(double)topz;
		otwodz = 0.5/dz;
		odzdz = 1.0/(dz*dz);

		drH = ( h[ri+1] - h[ri-1] )*otwodr;
		pressure = - gravity*drH;

		for(zi=1;zi<steps_z;zi++) {
			double dtU,drU,dzU,drdrU,dzdzU,Uval,Wval,laprU,z;
			double momentum,viscosity;

			Uval = u[ri][zi];
			Wval = w[ri][zi];

			drU   = ( u[ri+1][zi] - u[ri-1][zi] )*otwodr;
			drdrU = ( u[ri+1][zi] + u[ri-1][zi] - (Uval+Uval) )*odrdr;

			if ( zi == topz ) { // we're at a z-boundary

				dzU   = ( u[ri][zi] - u[ri][zi-1] )/dz;
				dzdzU = ( u[ri][zi] + u[ri][zi-2] - 2*u[ri][zi-1])*odzdz;

			} else {  // not at bdry:

				dzU   = ( u[ri][zi+1] - u[ri][zi-1] )*otwodz;
				dzdzU = ( u[ri][zi+1] + u[ri][zi-1] - (Uval+Uval) )*odzdz;

			}

			laprU = drdrU + drU/r - Uval/(r*r);

			momentum = Uval*drU + Wval*dzU;
			viscosity = visc*(dzdzU + laprU);

			dtU = pressure + viscosity - momentum;

			if ( abs(dtU)/Uval > maxdudt )
				maxdudt = dtU/Uval;

			dtU *= dt;

			nu[ri][zi] = Uval + dtU;

			if ( (dtU / Uval) > 0.2 )
				diffstoobig++;

		}
	}

if ( diffstoobig > 0 )
	printf("%d differentials too big!\n",diffstoobig);

}

void forceFluxU(double **u)
{
int ri,zi;
double curflux,fixflux;

	for(ri=0;ri<steps_r;ri++) {
		curflux = 0;
		for(zi=0;zi<steps_z;zi++) {
			curflux += u[ri][zi];
		}

		curflux *= h[ri] / (double)topz; // correction factor

		curflux *= 2*PI*(ri*dr + rmin);

		if ( curflux == 0 )
			cleanup("found a flux of zero",10);

		fixflux = flux / curflux;

		for(zi=0;zi<steps_z;zi++) {
			u[ri][zi] *= fixflux;
		}
	}
}

void forceFluxH(double **u)
{
int ri,zi;
double curflux;

	for(ri=0;ri<steps_r;ri++) {
		curflux = 0;
		for(zi=0;zi<steps_z;zi++) {
			curflux += u[ri][zi];
		}

		curflux *= h[ri] / (double)topz; // correction factor

		curflux *= 2*PI*(ri*dr + rmin);

		if ( curflux == 0 )
			cleanup("found a flux of zero",10);

		h[ri] *= flux / curflux;
	}
}

/*}{*/

void solveW(double **w)
{
int ri,zi;

	/*
		we have u at new time. now find w.
		we use incompressibility:

		dz(w)	= (-1/r)*dr(ru)
				= - (u/r + dru)
		w(z=0) = 0
	*/

	for(ri=0;ri<steps_r;ri++) {
		double r,dz;

		r = (ri * dr) + rmin;

		dz = h[ri]/(double)topz;

		// this is the bdry condition:
		w[ri][-1] = 0;
		w[ri][ 0] = 0;

		for(zi=0;zi<topz;zi++) { // this will set nw[1] through nw[zimax]
			double UoverR,drU,dzW;

			UoverR = u[ri][zi] / r;
			drU = ( u[ri+1][zi] - u[ri-1][zi] )*otwodr;
			dzW = - ( UoverR + drU );

			w[ri][zi+1] = w[ri][zi-1] + 2*dz*dzW;

		/***
			if ( zi == 0 )
				w[ri][1] = dz * dzW;
			else if ( zi == 1 )
				w[ri][2] = w[ri][1] + dz * dzW; //w[ri][2] = 2*dz*dzW;  is the no-special-case
			else
		***/

		}
	}
}

/*}{*/

void solveH_dhdt(double *nh)
{
int ri;

	/* now find h at new t **/
	/** use incomp. at surface (from shallow water) */

	maxdhdt = 0;

	h[-1] = (jetr*jetr/2)/(rmin-dr);
	h[0] = (jetr*jetr/2)/(rmin);
	h[steps_r] = 2*h[steps_r - 1] - h[steps_r - 2];

	for(ri=0;ri<steps_r;ri++) {
		double drH,dtH;

		drH = ( h[ri+1] - h[ri-1] )*otwodr;

		dtH = w[ri][topz] - u[ri][topz] * drH;

		if ( abs(dtH)/h[ri] > maxdhdt )
			maxdhdt = dtH/h[ri];

		nh[ri] = h[ri] + hdt*dtH;
	}
}

void solveH_dhdr(double *h)
{
int ri;
double drH,uv,wv,lastdrH;

	/* now find h at new t **/
	/** use incomp. at surface (from shallow water) */

	h[-1] = (jetr*jetr/2)/(rmin-dr);
	h[0] = (jetr*jetr/2)/(rmin);

	lastdrH = 0;

	for(ri=0;ri<steps_r;ri++) {

		uv = u[ri][topz];
		wv = w[ri][topz];

		if ( uv == 0.0 ) {
			drH = lastdrH;
		} else {
			lastdrH = drH = wv / uv;
		}

		//h[ri+1] = h[ri-1] + 2*dr*drH;
		h[ri] = h[ri-1] + dr*drH;
	}
}

void conditionBoundary(void)
{
int zi,ri;
double hr_lo,hr_hi;

	/*
	 *	 enfore no-slip etc.
	 */

	for(ri=0;ri<steps_r;ri++) {
		u[ri][-1] = u[ri][0] = 0;	// no slip
		u[ri][steps_z] = u[ri][topz]; // no stress

		w[ri][-1] = w[ri][0] = 0;	// no flux through wall
		//w[ri][steps_z] = w[ri][topz]; // no stress ?
	}

	/*
	 *	put parabolic time-0 u into u[-1] to make sure the jet keeps flowing
	 */

	hr_lo = h[-1] = (jetr*jetr)/(2*(rmin-dr));
	hr_hi = h[steps_r] = h[steps_r-1];
	hr_lo /= rmin;
	hr_hi /= rmax;

	for(zi=0;zi<steps_z;zi++) {
		double Ubdry,zoverh;

		zoverh = (double)zi/(double)topz;

		Ubdry = 3.0 * jetv * zoverh  * ( 1 - (zoverh/2.0) );

		u[-1][zi] = Ubdry;
		w[-1][zi] = - Ubdry*hr_lo*zoverh;

		u[steps_r][zi] = Ubdry;
		w[steps_r][zi] = - Ubdry*hr_hi*zoverh;
	}
}

/*}{*/

void conditionTimeZero(void)
{
int ri,zi;
double r,z,height,zoverh;
double h_pre;

/***

fill h,u,w with time=0 conditions

we'll use the zero-viscocity solution

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

u(r,z) = 3*jetv*(z/h)(1-(z/h)/2)
w(r,z) = -u*z/r

***/

	h_pre = jetr*jetr/2.0;

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

		height = h_pre / r;

		h[ri] = height;

		for(zi=-PAD;zi<(steps_z+PAD);zi++) {
			z = zi * height / steps_z;
		
			zoverh = z/height;

			u[ri][zi] = 3.0 * jetv * zoverh * ( 1 - (zoverh/2.0) );
			
			w[ri][zi] = - u[ri][zi] * z / r;
		}
	}

}

void zintegrate(double *integralptr,double **f)
{
int ri,zi;
double integral;

	for(ri=0;ri<steps_r;ri++) {
		integral = 0;
		for(zi=0;zi<steps_z;zi++) {
			integral += f[ri][zi];
		}

		integral *= h[ri] / (double)topz; // correction factor <>

		integralptr[ri] = integral;
	}

}

void saveInfoHeader(FILE *fp,char *name,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,"# Godwin jump H : %0.9f\n",hjump);
fprintf(fp,"# rmin : %0.9f\n",rmin);
fprintf(fp,"# rmax : %0.9f\n",rmax);
fprintf(fp,"# dr : %0.16f\n",dr);
fprintf(fp,"# dt : %0.16f\n",dt);
fprintf(fp,"# steps_r : %ld\n",steps_r);
fprintf(fp,"# steps_z : %ld\n",steps_z);
fprintf(fp,"# steps_t : %ld\n",steps_t);
fprintf(fp,"# max[ (du/dt) / u ] : %f\n",maxdudt);
fprintf(fp,"# max[ (dh/dt) / u ] : %f\n",maxdhdt);
fprintf(fp,"#\n");
}

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

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

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

		saveInfoHeader(fp,name,timestep);

		writestr_init(fp);

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

		writestr_flush();

		fclose(fp);
	}
}

void saveMatrix(char *name,double **matrix,int timestep)
{
FILE *fp;
char fname[100];
char str[100];

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

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

		saveInfoHeader(fp,name,timestep);

		writestr_init(fp);

		for(ri=0;ri<steps_r;ri+=10) {
			r = rmin+dr*ri;
			for(zi=0;zi<steps_z;zi++) {
				z = zi*h[ri]/steps_z;
				sprintf(str,"%7f %7f %7f\n",r,z,matrix[ri][zi]);
				writestr(str);
			}
		}

		writestr_flush();

		fclose(fp);
	}
}

void saveMatrices(char *name,double **matrix1,double **matrix2,int timestep)
{
FILE *fp;
char fname[100];
char str[100];

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

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

		saveInfoHeader(fp,name,timestep);

		writestr_init(fp);

		for(ri=0;ri<steps_r;ri+=10) {
			r = rmin+dr*ri;
			for(zi=0;zi<steps_z;zi++) {
				z = zi*h[ri]/steps_z;
				sprintf(str,"%7f %7f %7f %7f\n",r,z,matrix1[ri][zi],matrix2[ri][zi]);
				writestr(str);
			}
		}

		writestr_flush();

		fclose(fp);
	}
}

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

double *vector(int size)
{
double *ret;
	if ( (ret = malloc(sizeof(double)*(size+2*PAD))) == NULL ) {
		dbf(); 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 ) {
		dbf(); 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 );
}


#define WRITESTR_LEN 65536
static FILE *writestr_file = NULL;
static char writestr_buf[WRITESTR_LEN];
static char *writestr_ptr = NULL,*writestr_ptrend = NULL;

void writestr_init(FILE *fp)
{
writestr_file = fp;
writestr_ptr = writestr_buf;
writestr_ptrend = writestr_buf + WRITESTR_LEN;
}

void writestr_flush(void)
{

if ( writestr_ptr == writestr_buf ) return;

fwrite(writestr_buf,(writestr_ptr - writestr_buf),1,writestr_file);

writestr_ptr = writestr_buf;

}

void writestr(char *str)
{
	while(*str) {
		if ( writestr_ptr == writestr_ptrend ) {
			writestr_flush();
		}
		*writestr_ptr++ = *str++;
	}
}

/*}{*/

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

	/** get cli options **/

	while(--argc) {
		char *arg;
		arg = *(++argv);
		while(*arg == '-' || *arg == '/') arg++;
		switch(*arg++) {
			case 'm':
			case 'M':
				while(*arg == '=' || *arg == ':') arg++;
				max_memory = atol(arg);
				fprintf(stderr,"got max_mem = %d megs\n",max_memory);
				break;
			case 's':
			case 'S':
				while(*arg == '=' || *arg == ':') arg++;
				save_steps_t = atol(arg);
				fprintf(stderr,"got save_time_steps = %d\n",save_steps_t);
				break;
			case 't':
			case 'T':
				while(*arg == '=' || *arg == ':') arg++;
				steps_t = atol(arg);
				fprintf(stderr,"got time_steps = %d\n",steps_t);
				break;
			case '?':
			case 'h':
			case 'H':
				fprintf(stderr,"usage : hydr [-,/] <tag:value>\n");
				fprintf(stderr,"tags:\n");
				fprintf(stderr,"\tt\tnumber of time steps\n");
				fprintf(stderr,"\tm\tmaximum megs of memory to use\n");
				fprintf(stderr,"\ts\ttime steps between saves\n");
				break;
			default:
				fprintf(stderr,"warning: unknown tag %s\n",arg-1);
				break;
		}
	}

	/** 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);

	// prediction from trivial Godwin
	rjump = jetr * pow(Re,0.33333333333333333);
	hjump = (jetr*jetr)/(2*rjump);

	rmin = rjump / 2;
	rmax = rjump * 5;

	printf("Godwin rjump : %f, rmin: %f, rmax: %f\n",rjump,rmin,rmax);
	
		{
	long memory;

	dr = (rmax - rmin) / 10000;
	for(;;) {
		steps_r = (rmax - rmin)/dr;
		steps_z = steps_r/2;
		memory = (steps_r*steps_z*sizeof(double)*4) >> 20; //count megs
		if ( memory < max_memory )
			break;
		dr *= 2.0;
	}

	printf("steps_r : %d steps_z : %d memory used: %d megs\n",steps_r,steps_z,memory);
	}

	topz = steps_z - 1;

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

	/***
		at time 0:
			dzdzU = 3jetv/hh
		choose a dt such that::
			(dt * visc * dzdzU) = U/10
		dt = U/(10*visc*dzdzU)

		approx h = dz
	***/

	hrmin = (jetr*jetr)/(2*rmin);
	hrmax = (jetr*jetr)/(2*rmax);
	dz = hrmax/(double)topz;
	dt = dz*dz / (3*visc);

	dt /= DT_SCALE;

	hdt = dt * HDT_SCALE;

	printf("dr: %f dz : %f dt : %f\n",dr,dz,dt);

	otwodr = 0.5 / dr;
	odrdr = 1.0 / (dr*dr);

	h1 = vector(steps_r);
	u1 = matrix(steps_r,steps_z);
	w1 = matrix(steps_r,steps_z);
	h2 = vector(steps_r);
	u2 = matrix(steps_r,steps_z);
	w2 = matrix(steps_r,steps_z);
	fluxr = vector(steps_r);

	h = h1;
	u = u1;
	w = w1;

	conditionTimeZero();

	solveNextTmaster();

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

void cleanup(char *mess,int ret)
{

	freevector(fluxr);
	freevector(h1);
	freevector(h2);
	freematrix(u1);
	freematrix(u2);
	freematrix(w1);
	freematrix(w2);

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

	exit(ret);
}

/* } * { */

void solveNextTmaster(void)
{
int time,timestep;

	time = 0;
	timestep = save_steps_t;
	do {

		for(;timestep<save_steps_t;timestep++) {

			solveNextT();
			printf("time step # %d, dtH = %f %%, dtU = %f %%\n",time,maxdhdt,maxdudt);

			++time;
		}
		timestep = 0;

		printf("saving, time = %f seconds\n",(time*dt));
		saveState(time);

	} while ( time < steps_t );

}

void saveState(int time)
{
int ri;

saveVector("h",h,steps_r,time);
//saveMatrix("u",u,time);
//saveMatrix("w",w,time);
saveMatrices("uw",u,w,time);

zintegrate(fluxr,u);
for(ri=0;ri<steps_r;ri++) fluxr[ri] *= 2*PI*(ri*dr + rmin);
saveVector("fluxr",fluxr,steps_r,time);
}

/* } */