/*

the whole velocity field is blowing up, and then the flux is falling off in a curve
something's blowing out at zi = 0 

----

there's a wiggle in h(r) with an r = 0.2 period
	the period seems to grow linearly with r
	-> it's the way zimax(r) has big steps in it.

*/

/* {

hydr2 : approximate model.

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(drdz(w) - dzdz(u)) - g dr(h)

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 = 100;
long steps_t = 1001;
long max_memory = 5; // 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,dz,dt,otwodr,otwodz,odzdz,odrdr;
double *fluxr;

/*}*{ 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,zmax,hseed;
long steps_r,steps_z;

/*}*{ protos **/

void dbf(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);

/*}*{* do it ***/

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

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

	rmin = rjump * 0.5;
	rmax = rjump * 20.0;
	zmax = (jetr * jetr)/rmin; // twice z at r = rmin

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

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

	dr = (rmax - rmin) / 10000;
	for(;;) {
		dz = dr / 100.0; // is this ratio ok?
		steps_r = (rmax - rmin)/dr;
		steps_z = (zmax/dz);
		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);
	}

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

	dt = dz*dz / (300*visc);

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

	otwodr = 0.5 / dr;
	otwodz = 0.5 / dz;
	odzdz = 1.0 / (dz*dz);
	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,ri;

	time = 0;
	timestep = save_steps_t;
	do {

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

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

			/*
			 *  check to make sure input data won't crash the machine
			 *
			\***/

			for(ri=0;ri<steps_r;ri++) {
				int zi;
				zi = h[ri] / dz;
				if ( zi < 0 ) {
					saveState(time);
					printf("h(%f) = %f \n",(ri*dr + rmin),h[ri]);
					dbf(); cleanup("h[ri] less than zero",10);
				} else if ( (zi+5) > steps_z ) {
					saveState(time);
					printf("h(%f) = %f \n",(ri*dr + rmin),h[ri]);
					dbf(); cleanup("h[ri] too large",10);
				}
			}

			solveNextT();

			++time;
		}

		timestep = 0;

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

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

void solveNextT(void)
{
int ri,zi,zimax;
double *nh,**nu,**nw;

	// no need to "make dr" , we never do (r+dr)

	/*{

	we are given pointer to h,u,w

	set "new" pointers.  We will make h,u,w point to these
		new values after we will them in.

	**/

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

	/*}*{

		fill u[-1] with half-parabolic jet flux

	*/


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

		zoverh = 2.0*rmin*zi*dz/(jetr*jetr);

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

		u[-1][zi] = Ubdry;
		w[-1][zi] = w[0][zi];
		// w[-1][zi] = - Ubdry*zi*dz / rmin;  , not necessary
		//						we may violate incompressibility near here

		u[steps_r][zi] = u[steps_r-1][zi];
		w[steps_r][zi] = w[steps_r-1][zi];

		h[-1] = h[0];
		h[steps_r] = h[steps_r - 1];
	}

	/*}*{

		find u at new time.  The equation is:

		dt(u) = - (udr + wdz)u -v(drdz(w) - dzdz(u)) - g dr(h)

		*******
		t=0 ; 
			dzU = 3jetv/h( (1-n/2) -n/2) = 3jetv*z/hh
			dzdzU = 3jetv/hh
			dzdzU = 3*159.15/h^2
			h = a^2/2r
			h^2 = a^4/4r^2
			dzdzU = 1193625*r^2 = 
			dzdzU = 3jetv/hh

			dtU = dzdzU causes a smoothing of the function
		***********

	*/

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

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

		r = ri*dr + rmin;

		zmax = (int)( h[ri]/dz + 1.5 );
		//zmax = steps_z;

		for(zi=0;zi<zmax;zi++) {
			double dtU,drU,dzU,drdrU,dzdzU,Uval,Wval,laprU;
			double momentum,viscosity,pressure;

			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-1) < 0 || (zi+1) >= zmax ) {	// we're at a z-boundary

				if ( zi == 0 ) {
					dzU   = ( u[ri][1] - u[ri][0] )/dz;
					dzdzU = ( u[ri][2] + u[ri][0] - 2*u[ri][1])/(dz*dz);
				} else {
					dzU   = ( u[ri][zi] - u[ri][zi-1] )/dz;
					dzdzU = ( u[ri][zi] + u[ri][zi-2] - 2*u[ri][zi-1])/(dz*dz);
				}

			} 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);
			pressure = - gravity*drH;

			dtU = pressure + viscosity - momentum;

			dtU *= dt;

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

			if ( (dtU / Uval) > 0.5 && (zi*dz) < h[ri] )
				printf(" differential too big!\n");

		}
	}

	u = nu;
	w = nw; // use nu, not u

	/*}*{

		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;

		r = (ri * dr) + rmin;

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

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

			UoverR = nu[ri][zi] / r;

			if ( ri == 0 )
				drU = ( nu[1][zi] - nu[0][zi] )/dr;
			else if ( ri == (steps_r - 1) )
				drU = ( nu[ri][zi] - nu[ri-1][zi] )/dr;
			else
				drU = ( nu[ri+1][zi] - nu[ri-1][zi] )*otwodr;
		
			dzW = - ( UoverR + drU );

			if ( zi == 0 )
				nw[ri][1] = dz * dzW;
/**
			else if ( zi == 1 )
				nw[ri][2] = nw[ri][1] + dz * dzW;
**/
			else
				nw[ri][zi+1] = nw[ri][zi-1] + 2*dz * dzW;
		}
	}


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

	for(ri=0;ri<steps_r;ri++) {
		double drH,dtH,r,height;
		int newzi;

		r  = ri * dr + rmin;
		height = h[ri];
		zi = (int)height/dz; // + 0.5 ?

		// find drH, taking care at bdry

		if ( ri == 0 )
			drH = (h[1] - h[0])/dr;
		else if ( (ri == steps_r-1) )
			drH = (h[ri] - h[ri-1])/dr;
		else
			drH = ( h[ri+1] - h[ri-1] )*otwodr;

		dtH = nw[ri][zi] - nu[ri][zi] * drH;

		nh[ri] = height + dt * dtH;

		newzi = nh[ri]/dz;

		if ( newzi < 0 || newzi >= steps_z ) {
			saveState(-1);
			printf("nh(%f) = %f ; oldh = %f ; ri = %d, zi = %d, newzi = %d\n",r,nh[ri],height,ri,zi,newzi);
			printf("drH = %f, u = %f, w = %f, dtH = %f, dt*dtH = %f\n",drH,nu[ri][zi],nw[ri][zi],dtH,dt*dtH);
			dbf(); cleanup("nh[ri] blown!",10);
		}

		if ( newzi > zi ) {
			int i;
			// we've opend up a new square(s):
			// copy the heighest value on up
			for(i=zi+1;i<=newzi;i++) {
				u[ri][i] = u[ri][zi];
				w[ri][i] = w[ri][zi];
			}
		}
	}

	/*}* swap the pointers **/

	h = nh;
	u = nu;
	w = nw;
}

/* } * { */

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

	for(ri=0;ri<steps_r;ri++) {
		zimax = (int)(((double)h[ri])/dz + 0.5);
		integral = 0;
		for(zi=0;zi<=zimax;zi++) {
			integral += f[ri][zi];
		}

		integral *= h[ri] / zimax; // 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,"# zmax = h(jetr) : %0.9f\n",zmax);
fprintf(fp,"# rmin : %0.9f\n",rmin);
fprintf(fp,"# rmax : %0.9f\n",rmax);
fprintf(fp,"# dr : %0.9f\n",dr);
fprintf(fp,"# dz : %0.9f\n",dz);
fprintf(fp,"# dt : %0.9f\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,"#\n");
}

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

	sprintf(fname,"dat\\%s.%03d",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.%03d",name,timestep);

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

		saveInfoHeader(fp,name,timestep);

		writestr_init(fp);

		for(ri=0;ri<steps_r;ri+=10) {
			r = rmin+dr*ri;
			zimax = h[ri] / dz;
			if ( zimax < 1 ) zimax = 1;          
			if ( zimax > steps_z ) zimax = steps_z;
			for(zi=0;zi<zimax;zi++) {
				sprintf(str,"%7f %7f %7f\n",r,(double)(dz*zi),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.%03d",name,timestep);

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

		saveInfoHeader(fp,name,timestep);

		writestr_init(fp);

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

		writestr_flush();

		fclose(fp);
	}
}

/* } * { **/

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 * dz;
		
			zoverh = z/height;

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

}

/* } * { */

// 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++;
	}
}

/* } */