/*
 * itrans-lsq.c
 *
 * Peak detection by least-squares fitting
 *
 * (c) 2007 Thomas White <taw27@cam.ac.uk>
 *
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <gsl/gsl_multifit_nlin.h>
#include <stdint.h>

#include "control.h"
#include "imagedisplay.h"
#include "reflections.h"

#define MAX_LSQ_ITERATIONS 100

typedef struct struct_sampledpoints {
	gsl_matrix *coordinates;
	gsl_vector *values;
	unsigned int n_samples;
} SampledPoints;

static int itrans_peaksearch_lsq_f(const gsl_vector *gaussian, void *params, gsl_vector *residuals) {

	double cal;
	double obs;
	SampledPoints *samples = params;
	double a, b, c, d, e, f;
	unsigned int sample;
	
	a = gsl_vector_get(gaussian, 0);  b = gsl_vector_get(gaussian, 1);  c = gsl_vector_get(gaussian, 2);
	d = gsl_vector_get(gaussian, 3);  e = gsl_vector_get(gaussian, 4);  f = gsl_vector_get(gaussian, 5);
	
	//printf("Trial parameters: a=%f b=%f c=%f d=%f e=%f f=%f\n", a, b, c, d, e, f);
	
	for ( sample=0; sample<samples->n_samples; sample++ ) {
		int dx = gsl_matrix_get(samples->coordinates, sample, 0);
		int dy = gsl_matrix_get(samples->coordinates, sample, 1);
		cal = exp(a + b*dx + c*dy + d*dx*dx + e*dy*dy + f*dx*dy);
		obs = gsl_vector_get(samples->values, sample);
		gsl_vector_set(residuals, sample, (cal-obs));
		//printf("At %3i,%3i: residual=%f - %f = %f\n", dx, dy, cal, obs, cal-obs);
	}
	
	return GSL_SUCCESS;

}

static int itrans_peaksearch_lsq_df(const gsl_vector *gaussian, void *params, gsl_matrix *J) {

	double a, b, c, d, e, f;
	unsigned int sample;
	SampledPoints *samples = params;
	
	a = gsl_vector_get(gaussian, 0);  b = gsl_vector_get(gaussian, 1);  c = gsl_vector_get(gaussian, 2);
	d = gsl_vector_get(gaussian, 3);  e = gsl_vector_get(gaussian, 4);  f = gsl_vector_get(gaussian, 5);

	//printf("------a------------b------------c------------d------------e------------f-------\n");
	for ( sample=0; sample<samples->n_samples; sample++ ) {
	
		double ex;
		int dx, dy;
		double derivative;
		
		dx = gsl_matrix_get(samples->coordinates, sample, 0);
		dy = gsl_matrix_get(samples->coordinates, sample, 1);
		ex = exp(a + b*dx + c*dy + d*dx*dx + e*dy*dy + f*dx*dy);
		
		derivative = ex;		/* wrt a */
		gsl_matrix_set(J, sample, 0, derivative);
		derivative = dx * ex;		/* wrt b */
		gsl_matrix_set(J, sample, 1, derivative);
		derivative = dy * ex;		/* wrt c */
		gsl_matrix_set(J, sample, 2, derivative);
		derivative = dx * dx * ex;	/* wrt d */
		gsl_matrix_set(J, sample, 3, derivative);
		derivative = dy * dy * ex;	/* wrt e */
		gsl_matrix_set(J, sample, 4, derivative);
		derivative = dx * dy * ex;	/* wrt f */
		gsl_matrix_set(J, sample, 5, derivative);

		/*if ( (dx == 0) && (dy == 0) ) {
			printf("> ");
		} else {
			printf("| ");
		}
		printf("%10.2e | %10.2e | %10.2e | %10.2e | %10.2e | %10.2e", gsl_matrix_get(J, sample, 0), gsl_matrix_get(J, sample, 1),
				gsl_matrix_get(J, sample, 2), gsl_matrix_get(J, sample, 3), gsl_matrix_get(J, sample, 4), gsl_matrix_get(J, sample, 5));
		if ( (dx == 0) && (dy == 0) ) {
			printf(" <\n");
		} else {
			printf(" |\n");
		}*/

	}
	//printf("-------------------------------------------------------------------------------\n");
	
	return GSL_SUCCESS;

}

static int itrans_peaksearch_lsq_fdf(const gsl_vector *gaussian, void *params, gsl_vector *f, gsl_matrix *J) {

	itrans_peaksearch_lsq_f(gaussian, params, f);
	itrans_peaksearch_lsq_df(gaussian, params, J);
	
	return GSL_SUCCESS;

}

static void itrans_interpolate(uint16_t *image, int width, int x, int y) {

	int a, b, c, d;
	double av_horiz, av_vert, av;
	printf("Interpolating...\n");
	a = image[(x-1) + width*y];  b = image[(x+1) + width*y];
	c = image[x + width*(y-1)];  d = image[x + width*(y+1)];
	av_horiz = (a+b)/2;	av_vert = (c+d)/2;
	av = (av_horiz+av_vert) / 2;
	image[x + width*y] = av;

}

unsigned int itrans_peaksearch_lsq(ImageRecord imagerecord, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {

	uint16_t max_val = 0;
	int width, height;
	uint16_t *image;
	unsigned int n_reflections = 0;
	
	width = imagerecord.width;
	height = imagerecord.height;
	image = imagerecord.image;

	/* Least-Squares Craziness. NB Doesn't quite work... */
	do {
	
		int max_x = 0;
		int max_y = 0;;
		int x, y;
		gsl_multifit_fdfsolver *s;
		gsl_multifit_function_fdf f;
		unsigned int iter;
		gsl_vector *gaussian;
		int iter_status, conv_status;
		gsl_matrix *covar;
		SampledPoints samples;
		int sample;
		double ga, gb, gc, gd, ge, gf;
		gboolean sanity;
		int dx, dy;
		int bb_lh, bb_rh, bb_top, bb_bot;
		uint16_t ival;
		double sx, sy;
	
		/* Locate the highest point */
		max_val = 0;
		for ( y=10; y<height-10; y++ ) {
			for ( x=10; x<width-10; x++ ) {

				if ( image[x + width*y] > max_val ) {
					max_val = image[x + width*y];
					max_x = x;  max_y = y;
				}
			
			}
		}
		x = max_x; y = max_y;
		
		/* Try fitting a Gaussian to this region and see what happens... */
		f.p = 6;
		
		/* Set initial estimate of the fit parameters.  I = exp(a + bx + cy + dx^2 + ey^2 + fxy) */
		gaussian = gsl_vector_alloc(f.p);
		gsl_vector_set(gaussian, 0, log(max_val));	/* a */
		gsl_vector_set(gaussian, 1, -0.01);		/* b */
		gsl_vector_set(gaussian, 2, -0.01);		/* c */
		gsl_vector_set(gaussian, 3, -0.01);		/* d */
		gsl_vector_set(gaussian, 4, -0.01);		/* e */
		gsl_vector_set(gaussian, 5, -0.01);		/* f */
		
		/* Determine the bounding box which contains most of the peak */
		/* Left */	dx = 0;	do { ival = image[(x+dx) + width*(y+0)]; dx--; } while ( (ival>max_val/300) && (x+dx>0) );	bb_lh = dx;
		/* Right */	dx = 0;	do { ival = image[(x+dx) + width*(y+0)]; dx++; } while ( (ival>max_val/300) && (x+dx<width) );	bb_rh = dx;
		/* Top */	dy = 0;	do { ival = image[(x+0) + width*(y+dy)]; dy++; } while ( (ival>max_val/300) && (y+dy<height) );	bb_top = dy;
		/* Bottom */	dy = 0;	do { ival = image[(x+0) + width*(y+dy)]; dy--; } while ( (ival>max_val/300) && (y+dy>0) );	bb_bot = dy;
		printf("Peak %i,%i: bounding box %i<dx<%i, %i<dy<%i\n", x, y, bb_lh, bb_rh, bb_bot, bb_top);
		sanity = TRUE;
		if ( (bb_rh-bb_lh) > 300 ) sanity = FALSE;
		if ( (bb_top-bb_bot) > 300 ) sanity = FALSE;
		if ( sanity ) {
		
			/* Choose which points inside this bounding box to use for curve fitting */
			samples.n_samples = ((bb_top-bb_bot)+1)*((bb_rh-bb_lh)+1);
			samples.coordinates = gsl_matrix_alloc(samples.n_samples, 2);
			sample = 0;
			for ( sx=bb_lh; sx<=bb_rh; sx++ ) {
				for ( sy=bb_bot; sy<=bb_top; sy++ ) {
					gsl_matrix_set(samples.coordinates, sample, 0, sx);	gsl_matrix_set(samples.coordinates, sample, 1, sy);
					sample++;
				}
			}
			
			samples.values = gsl_vector_alloc(samples.n_samples);
			f.n = samples.n_samples;
			for ( sample=0; sample<samples.n_samples; sample++ ) {
				int dx = gsl_matrix_get(samples.coordinates, sample, 0);
				int dy = gsl_matrix_get(samples.coordinates, sample, 1);
				gsl_vector_set(samples.values, sample, image[(x+dx) + width*(y+dy)]);
				//printf("At %3i,%3i: value=%i\n", dx, dy, image[(x+dx) + width*(y+dy)]);
			}
			f.params = &samples;
			
			/* Execute the LSQ fitting procedure */
			s = gsl_multifit_fdfsolver_alloc(gsl_multifit_fdfsolver_lmsder, f.n, f.p);
			f.f = &itrans_peaksearch_lsq_f;	f.df = &itrans_peaksearch_lsq_df;	f.fdf = &itrans_peaksearch_lsq_fdf;
			gsl_multifit_fdfsolver_set(s, &f, gaussian);
			iter = 0; conv_status = GSL_CONTINUE;
			do {
			
				/* Iterate */
				iter_status = gsl_multifit_fdfsolver_iterate(s);
				iter++;		
				
				/* Check for error */
				if ( iter_status ) {
					break;
				}
				
				/* Test for convergence */
				conv_status = gsl_multifit_test_delta(s->dx, s->x, 1e-6, 1e-6);
										
			} while ( (conv_status == GSL_CONTINUE) && (iter < MAX_LSQ_ITERATIONS) );
			
			/* See how good the fit is */
			covar = gsl_matrix_alloc(f.p, f.p);
			gsl_multifit_covar(s->J, 0.0, covar);
			ga = gsl_vector_get(s->x, 0);  gb = gsl_vector_get(s->x, 1);  gc = gsl_vector_get(s->x, 2);
			gd = gsl_vector_get(s->x, 3);  ge = gsl_vector_get(s->x, 4);  gf = gsl_vector_get(s->x, 5);
			if ( fabs(exp(ga + gb*100 + gc*100 + gd*100*100 + ge*100*100 + gf*100*100)) > 10 ) {
				printf("Failed sanity check: %3i,%3i, a=%f b=%f c=%f d=%f e=%f f=%f\n", x, y, ga, gb, gc, gd, ge, gf);
				sanity = FALSE;
			} else {
				sanity = TRUE;
			}
			if ( (conv_status == GSL_SUCCESS) && (!iter_status) && (sanity) ) {
				
				/* Good fit! */
				int dx, dy;
				int bb_top, bb_bot, bb_lh, bb_rh;
				double frac;
				double brightness;
			
				printf("Fit converged after %i iterations: Centre %3i,%3i, a=%f b=%f c=%f d=%f e=%f f=%f\n",
						iter, x, y, ga, gb, gc, gd, ge, gf);
				brightness = image[x + width*y];
				reflection_add_from_dp(ctx, (x-imagerecord.x_centre), (y-imagerecord.y_centre), imagerecord, brightness);
				n_reflections++;
				
				/* Remove this peak from the image */
				/* Find right-hand edge of bounding box */
				dx = 0;	do {
					double pval, ival;
					pval = exp(ga + gb*dx + gc*0 + gd*dx*dx + ge*0*0 + gf*dx*-0);
					ival = image[(x+dx) + width*(y+0)];
					frac = pval / ival;
					dx++;
				} while ( (frac > 0.1) && (x+dx<width) );
				bb_rh = dx;
				/* Find left-hand edge of bounding box */
				dx = 0;	do {
					double pval, ival;
					pval = exp(ga + gb*dx + gc*0 + gd*dx*dx + ge*0*0 + gf*dx*-0);
					ival = image[(x+dx) + width*(y+0)];
					frac = pval / ival;
					dx--;
				} while ( (frac > 0.1) && (x+dx>0) );
				bb_lh = dx;
				/* Find top edge of bounding box */
				dy = 0;	do {
					double pval, ival;
					pval = exp(ga + gb*0 + gc*dy + gd*0*0 + ge*dy*dy + gf*0*dy);
					ival = image[(x+0) + width*(y+dy)];
					frac = pval / ival;
					dy++;
				} while ( (frac > 0.1) && (y+dy<height) );
				bb_top = dy;
				/* Find bottom edge of bounding box */
				dy = 0;	do {
					double pval, ival;
					pval = exp(ga + gb*0 + gc*dy + gd*0*0 + ge*dy*dy + gf*0*dy);
					ival = image[(x+0) + width*(y+dy)];
					frac = pval / ival;
					dy--;
				} while ( (frac > 0.1) && (y+dy>0) );
				bb_bot = dy;
				printf("Fitted peak bounding box %i<dx<%i, %i<dy<%i\n", bb_lh, bb_rh, bb_bot, bb_top);
				for ( dx=bb_lh; dx<bb_rh; dx++ ) {
					for ( dy=bb_bot; dy<bb_top; dy++ ) {
						double pval;
						pval = exp(ga + gb*dx + gc*dy + gd*dx*dx + ge*dy*dy + gf*dx*dy);
						image[(x+dx) + width*(y+dy)] -= pval;
					}
				}
				
			} else {
				itrans_interpolate(image, width, x, y);
			}
			
			gsl_matrix_free(covar);
			gsl_multifit_fdfsolver_free(s);
			gsl_vector_free(gaussian);
			gsl_matrix_free(samples.coordinates);
			gsl_vector_free(samples.values);
		
		} else {
			printf("Failed bounding box sanity check\n");
			itrans_interpolate(image, width, x, y);
		}
		
	
	} while ( max_val > 50 );
	
	return n_reflections;
}