/* img_parker.c
 *
 * Take an img file of topography and use Parker's multi-term
 * FT of a Taylor series to estimate the gravity effect of the
 * topography.  Based on "img_filter_99", which does:
 * Filter an img file, using cosine transform in Y direction and
 * cosine blending of windows.  Use fftpack routines cost and rfft.
 #
 # 23 Feb 2009 WHFS:  changed to use check_imgfile_size to 
 # automatically get dimensions, and to pick strip width (576
 # or 2*576), eliminating -M and -N flags.
 *
 */
 
#include "img_predict.h"


main (argc, argv)
int	argc;
char	**argv;

{
	double	circum = 40075.01229;	/* Equatorial circumference, km */
	double	f = 1.0/298.257;
	double	depth = 0.0;		/* Depth for downward continuing  */
	double	flex_lambda = 135.0;	/* Wavelength of flexure for W1 filter  */
	double	w2a = 9500.0;		/* Constant in W2 filter  */
	double 	crust = 7.0;		/* Crust thickness for W1 filter  */
	double	esq, d2r, radius, phi, width, height, sinphi, wf, parkscale = -1.0;
	float	*blender, *z, *wrkx, *wrky, *tpr, *kx, *ky, *transp, *zsave;
	short int	*s;
	int	tracks = 0, verbose = 0, error = 0, do_flex = 0, do_w2 = 0;
	int	nximg, nyimg, nyfft = 0;
	int	i, j, k, iparker, iparkfact, iswath, nswaths, nyimghalf, nyffthalf, nyfftp1, nz, nparker = 0;
	size_t	kbyte;
	int	kbytesum;
	char	*infile, *outfile;
	void	load_wavenumbers(), costi_(), rffti_(), load_data(), blend_and_write();
	void	rfft2dg1c2(), filter();
	FILE	*fpin, *fpout;
	
	infile = (char *)NULL;
	outfile = (char *)NULL;
	

	/* Parse arg list:  */
	for (i = 1; i < argc; i++) {
		if (argv[i][0] == '-') {
			switch (argv[i][1]) {
				case 'D':
					depth = atof(&argv[i][2]);
					break;
				case 'F':
					do_flex = (argv[i][2] == 'h') ? 1 : -1;
					break;
				case 'N':
					nyfft = atoi(&argv[i][2]);
					break;
				case 'O':
					outfile = &argv[i][2];
					break;
				case 'P':
					error += ( (sscanf(&argv[i][2], "%d/%lf", &nparker, &parkscale)) != 2);
					break;
				case 'T':
					tracks = 1;
					break;
				case 'V':
					verbose = 1;
					break;
				case 'W':
					do_w2 = 1;
					break;
				default:
					fprintf (stderr, "img_parker:  Unsupported option:  %s\n", argv[i]);
					error ++;
					break;
			}
		}
		else {
			infile = argv[i];
		}
	}
	if (nyfft%2 != 0) {
		fprintf (stderr, "img_parker:  <nyfft> must be even.\n");
		error ++;
	}
	if (infile == (char *)NULL) {
		fprintf (stderr, "img_parker:  You must supply an input file.\n");
		error ++;
	}
	if (outfile == (char *)NULL) {
		fprintf (stderr, "img_parker:  You must supply -O<output file>.\n");
		error ++;
	}
	if (nparker <= 0 || parkscale <= 0.0) {
		fprintf (stderr, "img_parker:  You must supply -P<nterms>/<scale>.\n");
		error ++;
	}
	
	if (error) {
		fprintf (stderr, "usage:  img_parker <infile> -M<dxm> -N<nyfft> -O<outfile> -P<nterms>/<scale> [-D<d>] [-F<lorh>] [-T] [-V] [-W]\n\n");
		fprintf (stderr, "\t<infile> and <outfile> are img files with pixel width <dxm> minutes.\n");
		fprintf (stderr, "\t<nyfft> is the height of each FFT swath, in pixels.  Must be even.\n");
		fprintf (stderr, "\t-P<nterms>/<scale> sets # of terms in the Parker series, and scale factor out front.\n");
		fprintf (stderr, "\t\tUse a scale factor to convert img topo (input) to img gravity (output).  E.g.\n");
		fprintf (stderr, "\t\tfor topo of specific density 2.67 in seawater, input m, output mGal*10, use 0.6876\n");
		fprintf (stderr, "\tOption -D<d> downward continues to <d> km depth.\n");
		fprintf (stderr, "\tOption -Fh highpasses the flexure band (W1 by Eqn 9 of S&S 1994).\n");
		fprintf (stderr, "\tOption -Fl  lowpasses the flexure band (1-W1 by Eqn 9 of S&S 1994).\n");
		fprintf (stderr, "\tOption -T removes track coding from input file.\n");
		fprintf (stderr, "\tOption -V reports verbose operations to stderr.\n");
		fprintf (stderr, "\tOption -W applies W2 filter, including d if -D used (Eqn 11 of S&S 1994).\n");
		exit (-1);
	}


	if (check_imgfile_size (infile, &nximg, &nyimg) ) {
		fprintf (stderr, "img_parker:  FATAL ERROR.  Unrecognized file size for %s\n", infile);
		exit (EXIT_FAILURE);
	}
	
	if (nyfft == 0) {
		/* Use default value  */
		nyfft = (nximg == 21600) ? 1152 : 576;
	}

	/* Set up some stuff  */
	
	nyimghalf = nyimg/2;
	d2r = M_PI / 180.0;
	esq = f * (2.0 - f);
	radius = nximg/(2.0 * M_PI);
	if ( (2 * nyimg)%nyfft != 0) {
		fprintf (stderr, "img_parker:  nyfft = %d  does not fit nyimg = %d\n",
			nyfft, nyimg);
		exit (-1);
	}
	nswaths = (2 * nyimg) / nyfft;
	nyfftp1 = nyfft + 1;
	nyffthalf = nyfft / 2;


	/* Malloc stuff:  */
	kbyte = nximg*nyfft/2;
	if ( (blender = (float *)calloc(kbyte, (size_t)4) ) == NULL) {
		fprintf (stderr, "img_parker:  calloc failure for blender array.\n");
		exit (-1);
	}
	kbytesum = kbyte * 4;
	
	kbyte = nximg*nyfftp1*4;
	if ( (z = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for z array.\n");
		exit (-1);
	}
	kbytesum += kbyte;
	if ( (zsave = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for zsave array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = (2*nximg + 16) * 4;
	if ( (wrkx = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for wrkx array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = (3*nyfftp1 + 16) * 4;
	if ( (wrky = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for wrky array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = nyffthalf*4;
	if ( (tpr = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for tpr array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = nximg*4;
	if ( (kx = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for kx array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = nyfftp1*4;
	if ( (ky = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for ky array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = nyfftp1*4;
	if ( (transp = (float *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for transp array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	kbyte = nximg*nyfftp1*2;
	if ( (s = (short int *)malloc(kbyte) ) == NULL) {
		fprintf (stderr, "img_parker:  malloc failure for s array.\n");
		exit (-1);
	}
	kbytesum += kbyte;

	if (verbose) printf("img_parker.c  Files are %d by %d; nyfft is %d; %6.1lf Mbytes used.\n",
		nximg, nyimg, nyfft, kbytesum/1048576.0);



	/* Initialize blending taper:  */
	wf = M_PI/nyffthalf;
	for (i = 0; i < nyffthalf; i++) tpr[i] = (float) (0.5 * (1.0 + cos(wf*(i-nyffthalf+0.5))));
	
	
	/* Initialize FT routines  */
	costi_(&nyfftp1, wrky);
	rffti_(&nximg, wrkx);

	/* Open files:  */
	if ( (fpin = fopen(infile, "r")) == NULL) {
		fprintf (stderr, "img_parker:  Error.  Cannot open r %s\n", infile);
		exit (-1);
	}
	if ( (fpout = fopen(outfile, "w")) == NULL) {
		fprintf (stderr, "img_parker:  Error.  Cannot open w %s\n", outfile);
		exit (-1);
	}


	/* Do swaths  */
	
	nz = nximg * nyfftp1;
	
	for (iswath = 0; iswath <= nswaths; iswath++) {
		phi = 2.0 * atan(exp( (nyimghalf - iswath * (nyfft/2) ) / radius) ) - M_PI_2;
		sinphi = sin(phi);
		width = circum * cos(phi) / sqrt(1.0 - esq * sinphi * sinphi);
		height = (2.0 * nyfft) * (width / nximg);	/* Cos FT doubles Y length  */
		if (verbose) printf ("Strip %d of %d at latitude %8.4lf is %8.2lf by %7.2lf km.\n",
			iswath, nswaths, phi/d2r, width, height);
		
		load_data (iswath, nswaths, nximg, nyfftp1, tracks, s, z, fpin);
		load_wavenumbers (width, height, kx, ky, nximg, nyfftp1);
		
		iparkfact = 1;
		for (iparker = 1; iparker <= nparker; iparker++) {
		
			iparkfact *= iparker;
			
			if (iparker > 1) for (i = 0; i < nz; i++) z[i] = (float)pow((double)s[i], (double)iparker);

			rfft2dg1c2 (z, nximg, nyfftp1, wrkx, wrky, transp, -1);
			
			if (iparker == 1) {
				for (i = 0; i < nz; i++) zsave[i] = z[i];
			}
			else {
				for (j = 0, k = 0; j < nyfftp1; j++) {
					for (i = 0; i < nximg; i++, k++) {
						z[k] *= pow(1.0e-06*(ky[j] + kx[i]), 0.5*(iparker-1))/iparkfact;
						zsave[k] += z[k];
					}
				}
			}
		}
		for (i = 0; i < nz; i++) z[i] = zsave[i] * parkscale;

		filter (z, kx, ky, nximg, nyfftp1, do_flex, do_w2, depth, crust, flex_lambda, w2a);
		rfft2dg1c2 (z, nximg, nyfftp1, wrkx, wrky, transp, 1);
		
		blend_and_write (iswath, nswaths, nximg, nyfft, blender, tpr, z, fpout);
	}
	/* close files and free mallocs  */
	exit (0);
}

void	load_wavenumbers (width, height, kx, ky, nximg, nyfftp1)
double	width, height;
float	*kx, *ky;
int	nximg, nyfftp1;
{
	int	i, j;
	double	f;

	/* Define kx,ky as 1.0/wavelength.
	   Set up kx for rfft, ky for cost.
	   Put kx-squared in kx[]; ky-squared in ky[].  */
	 
	kx[0] = 0.0;
	f = 1.0 / width;
	for (i = 1; i < nximg; i+= 2) {
		j = (i + 1)/2;
		kx[i] = (float) (j * f);
		kx[i+1] = kx[i];
	}
	if (nximg%2 == 0) kx[nximg-1] = (float) ( (nximg/2) * f);
	
	for (i = 1; i < nximg; i++) kx[i] = (kx[i] * kx[i]);

	ky[0] = 0.0;
	f = 1.0 / height;
	for (i = 1; i < nyfftp1; i++) ky[i] = (float) (i * f);

	for (i = 1; i < nyfftp1; i++) ky[i] = (ky[i] * ky[i]);
	
	return;
}

void load_data (iswath, nswaths, nximg, nyfftp1, tracks, s, z, fp)
int	iswath, nswaths, nximg, nyfftp1, tracks;
short int *s;
float	*z;
FILE	*fp;
{
	/* Load short int array s[] with nyfftp1 rows of data.
		Shift and shuffle data as needed, reading from fpin
		to load new data into s after shuffling.
		If (tracks), subtract 1 from odd values after read.
		Then copy contents of s to z.
	*/

	int	k, kstart, kstop, kread, kmove, nyffthalf, nxy;

	nyffthalf = (nyfftp1 - 1)/2;
	nxy = nyfftp1 * nximg;
	
	if (iswath == 1 || iswath == nswaths) {
		/* Re-use previous data; no need to shuffle or read.  */
		for (k = 0; k < nxy; k++) z[k] = (float)s[k];
		return;
	}
	
	if (iswath == 0) {
		/* Need to load entire amount:  */
		kread = nximg * nyfftp1;
		if ( (fread( (void *)s, (size_t)2, kread, fp) ) != kread) {
			fprintf (stderr, "img_parker:  Read error at iswath = %d\n", iswath);
			exit (-1);
		}
		if (tracks) {
			/* Need to fix entire amount:  */
			for (k = 0; k < kread; k++) {
				if ( (abs((int)s[k]))%2 == 1) s[k] --;
			}
		}
		for (k = 0; k < nxy; k++) z[k] = (float)s[k];
		return;
	}

	if (iswath == nswaths - 1) {
		/* Discard, shuffle, and load one row less than usual.  */
		kstart = (nyffthalf - 1) * nximg;
		kmove = (nyffthalf + 2) * nximg * 2;
	}
	else {
		/* Need to discard first nyhalf rows, shuffle remainder, and load new on end.  */
		kstart = nyffthalf * nximg;
		kmove = (nyfftp1 - nyffthalf) * nximg * 2;
	}
	memmove ( (void *)s, (void *)&s[kstart], (size_t)kmove);
	kread = kstart;
	kstart = kmove / 2;
	if ( (fread( (void *)&s[kstart], (size_t)2, kread, fp) ) != kread) {
		fprintf (stderr, "img_parker:  Read error at iswath = %d\n", iswath);
		exit (-1);
	}
	if (tracks) {
		kstop = nyfftp1 * nximg;
		for (k = kstart; k < kstop; k++) {
			if ( (abs((int)s[k]))%2 == 1) s[k] --;
		}
	}
	for (k = 0; k < nxy; k++) z[k] = (float)s[k];
	return;
}

void blend_and_write (iswath, nswaths, nximg, nyfft, blender, tpr, z, fpout)
int	iswath, nswaths, nximg, nyfft;
float	blender[], tpr[], z[];
FILE	*fpout;
{
	int	nyffthalf, i, j, j1, j2, k, k1, k2;
	double	x;
	float	tpr1;
	short int h;

	nyffthalf = nyfft/2;
	
	if (iswath == 0) {
		/* Simply copy the first half of z to blender and return  */
		k2 = nyffthalf * nximg;
		for (k = 0; k < k2; k++) blender[k] = z[k];
		return;
	}
	
	
	j1 = (iswath >= nswaths - 1) ? 1 : 0;	/* Start of first half  */
	j2 = j1 + nyffthalf;			/* Start of 2nd half  */
	if (iswath == nswaths) {
		j1 = j2;	/* Need to blend 2nd half instead of 1st half  */
	}
	
	k1 = j1 * nximg;
	k2 = j2 * nximg;
	
	for (j = 0, k = 0; j < nyffthalf; j++) {
		tpr1 = 1.0 - tpr[j];
		for (i = 0; i < nximg; i++, k++) {
			x = rint(tpr[j] * z[k1+k] + tpr1 * blender[k]);
			if (fabs(x) > 32767.0) {
				fprintf (stderr, "img_parker:  Overflow problem.\n");
				x = copysign(32767.0, x);
			}
			h = (short int)x;
			if ( (fwrite( (void*)&h, (size_t)2, (size_t)1, fpout) ) != 1) {
				fprintf (stderr, "img_parker:  Write error.\n");
				exit (-1);
			}
			blender[k] = z[k2+k];
		}
	}
	return;
}


void	rfft2dg1c2 (a, n1, n2, wsave1, wsave2, transp, isgn)
float	*a, *wsave1, *wsave2, *transp;
int	n1, n2, isgn;
{
	/* A 2-d real array is stored contiguously in a[], with the n1
		dimension looped first (i.e., like fortran a(n1,n2) ).
		The n1 dimension is given a real fourier transform.
		The n2 dimension is given a real cosine  transform.
		If isgn = -1, the forward transforms are calculated;
		If isgn = +1, the inverse transforms are calculated,
		and the data are then normalized by the appropriate factor.

		The fftpack routines rfftf_, rfftb_, cost_ are used.
		It is assumed that wsave1 and wsave2 have been initialized
		by calling rffti_(&n1, wsave1) and costi_(&n2, wsave2), 
		respectively.
		
		transp[] is a work array of size n2 used to compute FTs in
		the n2 direction ("transpose").  */
		

	int	i, j, k, n;
	double	r;
	void	rfftf_(), rfftb_(), cost_();

	if (isgn < 0) {	

		/* Forward transform.  
			First do real ft of n1 dimension in place:  */
		n = n1;
		for (j = 0, k = 0; j < n2; j++, k+=n1) {
			rfftf_(&n, &a[k], wsave1);
		}

		/* Now do cosft of n2 dimension by column vectors in transp:  */
		n = n2;
		for (i = 0; i < n1; i++) {
			for (j = 0, k = i; j < n2; j++, k+=n1) {
				transp[j] = a[k];
			}
			cost_(&n, transp, wsave2);
			for (j = 0, k = i; j < n2; j++, k+=n1) {
				a[k] = transp[j];
			}
		}
	}
	else {
		/* Inverse transform.  Get normalization factor:  */
		n = 2*n1*(n2-1);
		r = 1.0/n;
		
		/* Do cosft of n2 dimension by column vectors in transp,
			and multiply by r to normalize:  */
		n = n2;
		for (i = 0; i < n1; i++) {
			for (j = 0, k = i; j < n2; j++, k+=n1) {
				transp[j] = (float) (r * a[k]);
			}
			cost_(&n, transp, wsave2);
			for (j = 0, k = i; j < n2; j++, k+=n1) {
				a[k] = transp[j];
			}
		}
		
		/* Now, do real ft of n1 dimension in place:  */
			
		n = n1;
		for (j = 0, k = 0; j < n2; j++, k+=n1) {
			rfftb_(&n, &a[k], wsave1);
		}

	}
	return;
}

void	filter (z, kx, ky, nx, ny, do_flex, do_w2, depth, crust, flex_lambda, w2a)
float	*z, *kx, *ky;
int	nx, ny, do_flex, do_w2;
double	depth, crust, flex_lambda, w2a;
{
	int	i, j, k;
	double	pow_flex_lam_4, two_pi_c, four_pi_d, two_pi_d;
	double	f, dky2, dkr2, dkr4, dkr, w2, w1, phi;


	two_pi_d = 2.0 * M_PI * depth;
	four_pi_d = 2.0 * two_pi_d;
	two_pi_c = 2.0 * M_PI * crust;
	pow_flex_lam_4 = pow(flex_lambda, 4.0);

	for (j = 0, k = 0; j < ny; j++) {
		dky2 = (double)ky[j];
		for (i = 0; i < nx; i++, k++) {
			dkr2 = dky2 + (double)kx[i];
			dkr4 = dkr2 * dkr2;
			dkr = sqrt(dkr2);
			f = (depth == 0.0) ? 1.0 : exp(two_pi_d * dkr);	/* Downward continue.  */
			if (do_w2) {
				w2 = 1.0 / (1.0 + w2a*dkr4*exp(dkr*four_pi_d) );
				f *= w2;
			}
			if (do_flex) {
				phi = 1.0 / (1.0 + pow_flex_lam_4 * dkr4);
				w1 = (do_flex < 0) ? phi * exp(-two_pi_c * dkr) : 1.0 - phi * exp(-two_pi_c * dkr);
				f *= w1;
			}
			z[k] = (float) (z[k] * f);
		}
	}
	return;
}