/**************************************************************************
 *                                                                        *
 * Copyright (c) 1996 Michael Richmond and Richard Treffers               *
 *                                                                        *
 *                    This software may be copied and distributed for     *
 *                    educational, research and not for profit services   *
 *                    provided that this copyright and statement are      *
 *                    included in all such copies.                        *
 *                                                                        *
 **************************************************************************/


/****************************************************************************
 * Fit a low-order polynomial to the background of an image, and 
 *   subtract the background from the image.  Optionally, create
 *   a version of the image which is purely the model of the background
 *   fit.
 *
 *
 *   BACK filename [ngrid=] [order=] [outfile=] [map=] [verbose] [help]
 *
 * where
 *            ngrid        specifies the number of blocks in each direction
 *                              inside which we find a local background
 *                              value; we'll fit a polynomial to these values.
 *
 *            order        is the order of the polynomial fit.  In each
 *                              case, we pick the center of the image
 *                              as the fiducial spot, and normalize
 *                              by the size of the image in each 
 *                              direction.  So, for an image 100x200
 *                              pixels, the coords in polynomial are
 *
 *                                     x'  =  (x - 50)/50
 *                                     y'  =  (y - 100)/100
 *
 *                              0         a constant value = a
 *                              1         a + b*x' + c*y'
 *                              2         a + b*x' + c*y' + 
 *                                            + d*x'*x' + e*x'*y' + f*y'*y'
 *
 *            outfile      name of a new file into which we write the
 *                              background-subtracted image.  Default is
 *                              to overwrite the input image.
 *
 *            map          name of a new file into which to write an image
 *                              which is simply a map of the background 
 *                              model.  Default is not to create a map.
 *
 *            verbose      provide update on status of processing, and
 *                              print out polynomial coefficients when done
 *
 *            help         print a guide to use
 *
 * MWR 1/19/2003
 */

#include "pcvista.h"
#include "fits.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <math.h>

#undef DEBUG


#define USAGE \
 "usage: back filename [ngrid=] [order=] [outfile=] [map=] [verbose] [help]"


/* some default values */
#define NGRID   10        /* break image into this many pieces each way */
#define ORDER    1        /* means linear gradient in each direction */

#define MAXORDER   2      /* highest order we can handle */
#define MAXCOEFF   6      /* number of coeffs needed for highest order */

static int verbose = 0;
static char *progname = "back";
static char err_msg[BIG_BUF];
static char tempfile[] = "./back.tmp";

int 
main
	(
		int argc, 
		char *argv[]
	)
{
	char *gotit;
	char infile[BIG_BUF];
	char outfile[BIG_BUF];
	char mapfile[BIG_BUF];
	int i, j;
	int ngrid = NGRID;
	int order = ORDER;
	int nrow, ncol, nr, nc;
	int num_unknown;
	int num_box, count_box;
	int rsize, csize;
	int sr, sc;
	int nskip, binsize;
	int ibox, jbox;
	int16 sky;
	double central_row, central_col;
	double x, y;
	double *vector_x, *vector_b;
	double **matrix_a;
	float skysig;
	FITS_HANDLE infh;


	/*
	 * step 0: initialize some stuff 
	 */
	strcpy(outfile, "");
	strcpy(mapfile, "");

	/* 
	 * step 1: deal with command-line arguments 
	 */

	if (argc <= 1) {
		error(1, USAGE);
	}
	if (strcmp(argv[1], "help") == 0) {
		printf("%s\n", USAGE);
		exit(0);
	}
	if (strlen(argv[1]) >= BIG_BUF-1) {
		error(1, "back: input name is too long");
	}
	strncpy(infile, argv[1], BIG_BUF-1);

	for (i = 2; i < argc; i++) {
		if (keyword("help", argv[i])) {
			printf("%s\n", USAGE);
			exit(0);
		}
		if ((gotit = find("ngrid", argv[i])) != NULL) {
			ngrid = (int)(evaluate(gotit) + 0.5);	
			if (ngrid <= 0) {
				sprintf(err_msg, "%s: invalid ngrid value %d", progname, ngrid);
				error(1, err_msg);
			}
			continue;
		}
		if ((gotit = find("order", argv[i])) != NULL) {
			order = (int)(evaluate(gotit) + 0.5);
			if ((order < 0) || (order > MAXORDER)) {
				sprintf(err_msg, "%s: invalid order value %d", progname, order);
				error(1, err_msg);
			}
			continue;
		}
		if ((gotit = find("outfile", argv[i])) != NULL) {
			if (strlen(gotit) >= BIG_BUF-1) {
				error(1, "back: outfile name is too long");
			}
			strncpy(outfile, gotit, BIG_BUF-1);
			continue;
		}
		if ((gotit = find("map", argv[i])) != NULL) {
			if (strlen(gotit) >= BIG_BUF-1) {
				error(1, "back: map name is too long");
			}
			strncpy(mapfile, gotit, BIG_BUF-1);
			continue;
		}
		if (keyword("verbose", argv[i])) {
			verbose = 1;
			continue;
		}
		sprintf(err_msg, "%s: unknown argument '%s'\n", progname, argv[i]);
		error(1, err_msg);
	}

	/*
	 * step 2: allocate the arrays we need.
	 *         Also, set all elements to zero.
	 */
	num_box = ngrid*ngrid;
	switch (order) {
		case 0:   /* single constant value */
			num_unknown = 1;
			break;
		case 1:   /* const + linear gradient */ 
			num_unknown = 3;
			break;
		case 2:   /* const + linear gradient + quadratic terms */
			num_unknown = 6;
			break;
		default:
			sprintf(err_msg, "%s: bad value for order %d ??", progname, order);
			error(1, err_msg);
			break;
	}
	if ((matrix_a = (double **) malloc(num_unknown*sizeof(double *))) == NULL) {
		sprintf(err_msg, "%s: can't alloc for matrix_a", progname);
		error(1, err_msg);
	}
	for (i = 0; i < num_unknown; i++) {
		if ((matrix_a[i] = (double *) malloc(num_unknown*sizeof(double))) == NULL) {
			sprintf(err_msg, "%s: can't alloc for row %d of matrix_a",
							progname, i);
			error(1, err_msg);
		}
		for (j = 0; j < num_unknown; j++) {
			matrix_a[i][j] = 0.0;
		}
	}
	if ((vector_b = (double *) malloc(num_box*sizeof(double))) == NULL) {
		sprintf(err_msg, "%s: can't alloc for vector_b", progname);
		error(1, err_msg);
	}
	for (i = 0; i < num_unknown; i++) {
		vector_b[i] = 0.0;
	}
#ifdef DEBUG
	printf("%s: order %d  ngrid %d   num_unknown %d  num_box %d \n",
				progname, order, ngrid, num_unknown, num_box);
#endif


	/* 
	 * step 3: open input image
	 */
	infh = fits_open(argv[1], "r", &nrow, &ncol);


	/* 
	 * step 4: find "background" value in each sub-section of image.
	 *
	 *         Also, use the value in each image sub-section
	 *         together with various powers of the scaled coordinates:
	 *
	 *                x  =  (row - central_row) / (nrow/2) ;
	 *                y  =  (col - central_col) / (ncol/2) ;
	 *
	 *         to fill in the elements of the matrix "A" and vector "b",
	 *         so that we can solve for the unknowns.
	 */
	count_box = 0;
	nskip = 1;
	binsize = 1;
	rsize = floor(nrow/ngrid) + 1;
	csize = floor(ncol/ngrid) + 1;
	central_row = nrow/2;
	central_col = ncol/2;
	for (ibox = 0; ibox < ngrid; ibox++) {
		sr = rsize*ibox;
		nr = rsize;
		if ((sr + nr) >= nrow) {
			nr = nrow - sr;
		}
		x = (sr + nr/2);
		x = (x - central_row)/((double) central_row);
		for (jbox = 0; jbox < ngrid; jbox++) {
			sc = csize*jbox;
			nc = csize;
			if ((sc + nc) >= ncol) {
				nc = ncol - sc;
			}
			y = (sc + nc/2);
			y = (y - central_col)/((double) central_col);
#ifdef DEBUG
			printf("box %2d,%2d  has %4d %4d  %4d %4d ... ",
					ibox, jbox, sr, nr, sc, nc);
#endif
			if (find_sky(infh, sr, sc, nr, nc, nskip, binsize, &sky, &skysig) != 0) {
				sprintf(err_msg, "%s: warning: sky may be wrong in box %d,%d ",
								progname, ibox, jbox);
			}
#ifdef DEBUG
			printf("sky %6d \n", sky);
#endif

			/* 
			 * Now, we fill in elements of the matrix equation
			 *           A dot x  =  b
			 */

			/*
			 * this is the right-hand side of the matrix equation.
			 * The "vector_b" array has "num_unknown" equations.
			 */
			vector_b[0] += sky;
			if (order >= 1) {
				vector_b[1] += sky*x;
				vector_b[2] += sky*y;
			}
			if (order >= 2) {
				vector_b[3] += sky*x*x;
				vector_b[4] += sky*x*y;
				vector_b[5] += sky*y*y;
			}

			/* 
			 * and this is the matrix "A" on the left-hand side ...
			 * This matrix has "num_unknown"-by-"num_unknown" elements,
			 * so it is a bit unwieldy to fill it in ...
			 */
			if (order >= 0) {
				matrix_a[0][0] += 1;
			}
			if (order >= 1) {
				matrix_a[0][1] += x;
				matrix_a[0][2] += y;

				matrix_a[1][0] += x;
				matrix_a[1][1] += x*x;
				matrix_a[1][2] += x*y;

				matrix_a[2][0] += y;
				matrix_a[2][1] += x*y;
				matrix_a[2][2] += y*y;
			}
			if (order >= 2) {
				matrix_a[0][3] += x*x;
				matrix_a[0][4] += x*y;
				matrix_a[0][5] += y*y;

				matrix_a[1][3] += x*x*x;
				matrix_a[1][4] += x*x*y;
				matrix_a[1][5] += x*y*y;

				matrix_a[2][3] += x*x*y;
				matrix_a[2][4] += x*y*y;
				matrix_a[2][5] += y*y*y;

				matrix_a[3][0] += x*x;
				matrix_a[3][1] += x*x*x;
				matrix_a[3][2] += x*x*y;
				matrix_a[3][3] += x*x*x*x;
				matrix_a[3][4] += x*x*x*y;
				matrix_a[3][5] += x*x*y*y;

				matrix_a[4][0] += x*y;
				matrix_a[4][1] += x*x*y;
				matrix_a[4][2] += x*y*y;
				matrix_a[4][3] += x*x*x*y;
				matrix_a[4][4] += x*x*y*y;
				matrix_a[4][5] += x*y*y*y;

				matrix_a[5][0] += y*y;
				matrix_a[5][1] += x*y*y;
				matrix_a[5][2] += y*y*y;
				matrix_a[5][3] += x*x*y*y;
				matrix_a[5][4] += x*y*y*y;
				matrix_a[5][5] += y*y*y*y;
			}

			
			/* just a sanity check */
			count_box++;

			if (verbose > 0) {
				printf("\r  up to box %2d %2d  out of %2d %2d ",
						ibox, jbox, ngrid, ngrid);
			}

		}
	}
	if (verbose > 0) {
		printf("\n");
	}
	if (count_box != num_box) {
		sprintf(err_msg, "%s: actual box count %d != expected %d ?!",
					progname, count_box, num_box);
		error(1, err_msg);
	}


	/*
	 * step 5: fit polynomial to the background values. 
	 *         The values of the coefficients for which we've solved
	 *         will be placed into the "vector_b", overwriting
	 *         the values we placed there earlier.
	 */
	if (gauss_matrix(matrix_a, num_unknown, vector_b) != 0) {
		sprintf(err_msg, "%s: gauss_matrix fails to solve linear equations",
						progname);
		error(1, err_msg);
	}
	if (verbose != 0) {
		for (i = 0; i < num_unknown; i++) {
			printf("  vector_x[%d] = %12.5e\n", i, vector_b[i]);
		}
	}


	/* 
	 * step 6: (if desired) create background map image
	 */
	if (strlen(mapfile) != 0) {
		{
			int row, col;
			int16 data[NMAX];
			double x, y;
			double value;
			FITS_HANDLE mapfh;

			mapfh = fits_open(mapfile, "w", &nrow, &ncol);
			for (row = 0; row < nrow; row++) {
				x = (row - central_row)/((double) central_row);
	
				for (col = 0; col < ncol; col++) {
					y = (col - central_col)/((double) central_col);
	
					value = 0.0;
					if (order >= 0) {
						value += vector_b[0];
					}
					if (order >= 1) {
						value += vector_b[1]*x + vector_b[2]*y;
					}
					if (order >= 2) {
						value += vector_b[3]*x*x + vector_b[4]*x*y + vector_b[5]*y*y;
					}
	
					if (value >= MAX_DATA_VAL) {
						data[col] = MAX_DATA_VAL;
					} else if (value <= MIN_DATA_VAL) {
						data[col] = MIN_DATA_VAL;
					} else {
						data[col] = (int16) value;
					}
				}
	
				fits_put_data(mapfh, row, 0, data, ncol);
			}

			fits_close(mapfh);
		}
	}


	/* 
	 * step 7: subtract background model from image
	 *         We always place the difference into a temporary file;
	 *           then, after we're done, we either link this temp file
	 *           over the original, or into a new file if the user 
	 *           gave the "outfile=" option
	 */
	{
		int row, col;
		int16 indata[NMAX], outdata[NMAX];
		double x, y;
		double value;
		FITS_HANDLE outfh;

		outfh = fits_open(tempfile, "w", &nrow, &ncol);

		for (row = 0; row < nrow; row++) {
			fits_get_data(infh, row, 0, indata, ncol);

			x = (row - central_row)/((double) central_row);

			for (col = 0; col < ncol; col++) {
				y = (col - central_col)/((double) central_col);

				value = 0.0;
				if (order >= 0) {
					value += vector_b[0];
				}
				if (order >= 1) {
					value += vector_b[1]*x + vector_b[2]*y;
				}
				if (order >= 2) {
					value += vector_b[3]*x*x + vector_b[4]*x*y + vector_b[5]*y*y;
				}

				value = indata[col] - value;
				if (value >= MAX_DATA_VAL) {
					outdata[col] = MAX_DATA_VAL;
 				} else if (value <= MIN_DATA_VAL) {
					outdata[col] = MIN_DATA_VAL;
				} else {
					outdata[col] = (int16) value;
				}
			}
	
			fits_put_data(outfh, row, 0, outdata, ncol);
		}
		fits_close(outfh);

		/* 
		 * Now, either replace the input file with the temp file,
		 * or give the temp file the desired output file name
		 */
		if (strlen(outfile) == 0) {
			if (unlink(infile) != 0) {
				sprintf(err_msg, "%s: couldn't unlink input file %s \n", 
								progname, infile);
				error(1, err_msg);
			}
			if (link(tempfile, infile) != 0) {
				sprintf(err_msg, "%s: couldn't link temp file %s to %s \n",
								progname, tempfile, infile);
				error(1, err_msg);
			}
			if (unlink(tempfile) != 0) {
				sprintf(err_msg, "%s: couldn't unlink tempfile %s\n", 
								progname, tempfile);
				error(1, err_msg);
			}
		}
		else {
			/* we may need to delete an existing file here ... */
			if (access(outfile, F_OK) == 0) {
				/* yes, the file exists ... so we must unlink it */
				if (unlink(outfile) != 0) {
					sprintf(err_msg, "%s: couldn't unlink file %s \n",
									progname, outfile);
				 	error(1, err_msg);
				}
			}
			if (link(tempfile, outfile) != 0) {
				sprintf(err_msg, "%s: couldn't link temp file %s to %s \n",
								progname, tempfile, outfile);
				error(1, err_msg);
			}
			if (unlink(tempfile) != 0) {
				sprintf(err_msg, "%s: couldn't unlink tempfile %s\n", 
								progname, tempfile);
				error(1, err_msg);
			}
		}

	}


	/* 
	 * step 8: de-allocate, close open files, etc.
	 */
	fits_close(infh);
	for (i = 0; i < num_unknown; i++) {
		free(matrix_a[i]);
	}
	free(matrix_a);
	free(vector_x);
	free(vector_b);



	/* done */
	exit(0);
}