/*
 * itrans.c
 *
 * Parameterise features in an image for reconstruction
 *
 * (c) 2007 Thomas White <taw27@cam.ac.uk>
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

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

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

#define MAX_LSQ_ITERATIONS 100
#define PEAK_WINDOW_SIZE 20

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

static unsigned int itrans_peaksearch_zaefferer(int16_t *image, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {

	int x, y;
	unsigned int n_reflections;
	
	n_reflections = 0;
	for ( x=1; x<ctx->width-1; x++ ) {
		for ( y=1; y<ctx->height-1; y++ ) {

			double dx1, dx2, dy1, dy2;
			double dxs, dys;
			double grad;

			if ( ctx->max_d ) {
				if ( (x-ctx->x_centre)*(x-ctx->x_centre) + (y-ctx->y_centre)*(y-ctx->y_centre) > ctx->max_d*ctx->max_d ) {
					continue;
				}
			}
			
			/* Get gradients */
			dx1 = image[x+ctx->width*y] - image[(x+1)+ctx->width*y];
			dx2 = image[(x-1)+ctx->width*y] - image[x+ctx->width*y];
			dy1 = image[x+ctx->width*y] - image[(x+1)+ctx->width*(y+1)];
			dy2 = image[x+ctx->width*(y-1)] - image[x+ctx->width*y];
			
			/* Average gradient measurements from both sides */
			dxs = ((dx1*dx1) + (dx2*dx2)) / 2;
			dys = ((dy1*dy1) + (dy2*dy2)) / 2;
			
			/* Calculate overall gradient */
			grad = dxs + dys;

			if ( grad > 400 ) {
			
				unsigned int mask_x, mask_y;
				unsigned int sx, sy;
				double max;
				unsigned int did_something = 1;
				
				mask_x = x;
				mask_y = y;
				
				while ( (did_something) && (distance(mask_x, mask_y, x, y)<50) ) {
					max = image[mask_x+ctx->width*mask_y];
					did_something = 0;
					for ( sy=biggest(mask_y-PEAK_WINDOW_SIZE/2, 0); sy<smallest(mask_y+PEAK_WINDOW_SIZE/2, ctx->height); sy++ ) {
						for ( sx=biggest(mask_x-PEAK_WINDOW_SIZE/2, 0); sx<smallest(mask_x+PEAK_WINDOW_SIZE/2, ctx->width); sx++ ) {
							if ( image[sx+ctx->width*sy] > max ) {
								max = image[sx+ctx->width*sy];
								mask_x = sx;
								mask_y = sy;
								did_something = 1;
							}
						}
					}
				}
				
				if ( !did_something ) {
					if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
						reflection_add_from_reciprocal(ctx, (signed)(mask_x-ctx->width/2)*ctx->pixel_size,
										(signed)(mask_y-ctx->height/2)*ctx->pixel_size,
										tilt_degrees, image[mask_x + ctx->width*mask_y]);
					} else {
						reflection_add_from_dp(ctx, (signed)(mask_x-ctx->width/2), (signed)(mask_y-ctx->height/2),
										tilt_degrees, image[mask_x + ctx->width*mask_y]);
					}
					if ( ctx->first_image ) {
						imagedisplay_mark_point(imagedisplay, mask_x, mask_y);
					}
					n_reflections++;
				}
				
			}
		}
	}

	return n_reflections;

}

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(int16_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;

}

static unsigned int itrans_peaksearch_threshold(int16_t *image, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {

	int width, height;
	int x, y;
	unsigned int n_reflections = 0;
	
	width = ctx->width;	height = ctx->height;

	/* Simple Thresholding */
	for ( y=0; y<height; y++ ) {
		for ( x=0; x<width; x++ ) {

			if ( image[x + width*y] > 10 ) {
				if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
					reflection_add_from_reciprocal(ctx, (signed)(x-ctx->x_centre)*ctx->pixel_size, (signed)(y-ctx->y_centre)*ctx->pixel_size,
									tilt_degrees, image[x + width*y]);
				} else {
					reflection_add_from_dp(ctx, (signed)(x-ctx->x_centre), (signed)(y-ctx->y_centre), tilt_degrees, image[x + width*y]);
				}
				if ( ctx->first_image ) {
					imagedisplay_mark_point(imagedisplay, x, y);
				}
				n_reflections++;
			}
		
		}
	}
	
	return n_reflections;
	
}

static unsigned int itrans_peaksearch_adaptive_threshold(int16_t *image, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {

	int16_t max_val = 0;
	int width, height;
	unsigned int n_reflections = 0;
	
	width = ctx->width;	height = ctx->height;

	/* Adaptive Thresholding */
	do {
	
		int max_x = 0;
		int max_y = 0;;
		int x, y;
	
		/* Locate the highest point */
		max_val = 0;
		for ( y=0; y<height; y++ ) {
			for ( x=0; x<width; x++ ) {

				if ( image[x + width*y] > max_val ) {
					max_val = image[x + width*y];
					max_x = x;  max_y = y;
				}
			
			}
		}
		
		if ( max_val > 50 ) {
			if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
				reflection_add_from_reciprocal(ctx, (signed)(max_x-ctx->x_centre)*ctx->pixel_size, (signed)(max_y-ctx->y_centre)*ctx->pixel_size,
								tilt_degrees, max_val);
			} else {
				reflection_add_from_dp(ctx, (signed)(max_x-ctx->x_centre), (signed)(max_y-ctx->y_centre), tilt_degrees, max_val);
			}
			if ( ctx->first_image ) {
				imagedisplay_mark_point(imagedisplay, max_x, max_y);
			}
			n_reflections++;
			
			/* Remove it and its surroundings */
			for ( y=((max_y-10>0)?max_y-10:0); y<((max_y+10)<height?max_y+10:height); y++ ) {
				for ( x=((max_x-10>0)?max_x-10:0); x<((max_x+10)<width?max_x+10:width); x++ ) {
					image[x + width*y] = 0;
				}
			}
		}
		
	} while ( max_val > 50 );
	
	return n_reflections;

}

static unsigned int itrans_peaksearch_lsq(int16_t *image, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {

	int16_t max_val = 0;
	int width, height;
	unsigned int n_reflections = 0;
	
	width = ctx->width;	height = ctx->height;

	/* 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;
		int16_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;
			
				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);
				if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
					reflection_add_from_reciprocal(ctx, (x-ctx->x_centre)*ctx->pixel_size, (y-ctx->y_centre)*ctx->pixel_size,
										tilt_degrees, image[x + width*y]);
				} else {
					reflection_add_from_dp(ctx, (x-ctx->x_centre), (y-ctx->y_centre), tilt_degrees, image[x + width*y]);
				}
				if ( ctx->first_image ) {
					imagedisplay_mark_point(imagedisplay, x, y);
				}
				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;
}

static unsigned int itrans_peaksearch_iterative(int16_t *image, ControlContext *ctx, double tilt_degrees, ImageDisplay *imagedisplay) {
	printf("Starting itrans_peaksearch_iterative\n");
	int n_reflections;	
	gsl_matrix *m;
	gsl_matrix *p;
	printf("Calling createImageMatrix\n");
	m = createImageMatrix(ctx,image);
	printf("Calling iterate\n");
	p = iterate(m, &n_reflections);
	int i;
	double px,py;
	printf("Found %d reflections\n",n_reflections);
	for (i=0;i<n_reflections;i++) {
		px = gsl_matrix_get(p,0,i);
		py = gsl_matrix_get(p,1,i);
		printf("Reflection %d (%lf,%lf)\n",i,px,py);
		if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
			reflection_add_from_reciprocal(ctx, (px-ctx->x_centre)*ctx->pixel_size, (py-ctx->y_centre)*ctx->pixel_size,
								tilt_degrees, 1.0);
		} else {
			reflection_add_from_dp(ctx, (px-ctx->x_centre), (py-ctx->y_centre), tilt_degrees, 1.0);
		}
		if (ctx->first_image) imagedisplay_mark_point(imagedisplay, (unsigned int)px, (unsigned int)py);
	}
	gsl_matrix_free(m);
	gsl_matrix_free(p);
	return (unsigned)n_reflections;
}

void itrans_process_image(int16_t *image, ControlContext *ctx, double tilt_degrees) {
	
	unsigned int n_reflections;
	ImageDisplay *imagedisplay = NULL;
	
	if ( ctx->first_image ) {
		imagedisplay = imagedisplay_open(image, ctx->width, ctx->height, "Image Display");
		imagedisplay_add_tilt_axis(imagedisplay, ctx, ctx->omega);
	}
	
	switch( ctx->psmode ) {
		case PEAKSEARCH_THRESHOLD : n_reflections = itrans_peaksearch_threshold(image, ctx, tilt_degrees, imagedisplay); break;
		case PEAKSEARCH_ADAPTIVE_THRESHOLD : n_reflections = itrans_peaksearch_adaptive_threshold(image, ctx, tilt_degrees, imagedisplay); break;
		case PEAKSEARCH_LSQ : n_reflections = itrans_peaksearch_lsq(image, ctx, tilt_degrees, imagedisplay); break;
		case PEAKSEARCH_ZAEFFERER : n_reflections = itrans_peaksearch_zaefferer(image, ctx, tilt_degrees, imagedisplay); break;
		case PEAKSEARCH_ITERATIVE : n_reflections = itrans_peaksearch_iterative(image, ctx, tilt_degrees, imagedisplay); break;
		default: n_reflections = 0;
	}
	
	if ( ctx->rmode == RECONSTRUCTION_PREDICTION ) {
		
	}
	
	ctx->first_image = 0;
	printf(" ..... %i peaks found\n", n_reflections);

}