/*
 * ipr.c
 *
 * Iterative Prediction-Refinement Reconstruction
 *
 * (c) 2007 Thomas White <taw27@cam.ac.uk>
 * 	    Gordon Ball <gfb21@cam.ac.uk>
 *
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

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

#include <stdlib.h>
#include <stdio.h>
#include <math.h>

#include "control.h"
#include "reflections.h"
#include "itrans.h"
#include "utils.h"
#include "imagedisplay.h"
#include "reproject.h"
#include "ipr.h"
#include "displaywindow.h"

static int ipr_choose_max(Basis *basis) {

	if ( (basis->a.modulus >= basis->b.modulus) && (basis->a.modulus >= basis->c.modulus) ) return 1;
	if ( (basis->b.modulus >= basis->a.modulus) && (basis->b.modulus >= basis->c.modulus) ) return 2;
	if ( (basis->c.modulus >= basis->a.modulus) && (basis->c.modulus >= basis->b.modulus) ) return 3;
	
	/* Never reaches here... */
	fprintf(stderr, "IP: Couldn't determine maximum basis vector modulus\n");
	abort();
	
}

static int ipr_random_image(ControlContext *ctx) {
	double n;
	n = (double)random()/RAND_MAX;
	n *= ctx->n_images;
	return (int)n;
}

static Reflection *ipr_random_reflection(ReflectionContext *rctx) {
	
	double n;
	int i;
	Reflection *cur;
	
	n = (double)random()/RAND_MAX;
	n *= rctx->n_reflections/2;
	cur = rctx->reflections;	/* Jump a maximum of half way in */
	for ( i=1; i<n; i++ ) {		/* Skip initial placeholder */
		cur = cur->next;
	}
	
	return cur;
	
}

/* Try to turn the basis into a more sensible one */
static void ipr_try_compact(Basis *basis) {

	int i;
	
	for ( i=1; i<5; i++ ) {
	
		if ( modulus(basis->a.x - basis->b.x, basis->a.y - basis->b.y, basis->a.z - basis->b.z) < basis->a.modulus ) {
			basis->a.x = basis->a.x - basis->b.x; basis->a.y = basis->a.y - basis->b.y; basis->a.y = basis->a.y - basis->b.y;
			basis->a.modulus = modulus(basis->a.x, basis->a.y, basis->a.z);
		}
		if ( modulus(basis->a.x - basis->c.x, basis->a.z - basis->c.z, basis->a.z - basis->c.z) < basis->a.modulus ) {
			basis->a.x = basis->a.x - basis->c.x; basis->a.y = basis->a.y - basis->c.y; basis->a.y = basis->a.y - basis->c.y;
			basis->a.modulus = modulus(basis->a.x, basis->a.y, basis->a.z);
		}
		if ( modulus(basis->a.x + basis->b.x, basis->a.y + basis->b.y, basis->a.z + basis->b.z) < basis->a.modulus ) {
			basis->a.x = basis->a.x + basis->b.x; basis->a.y = basis->a.y + basis->b.y; basis->a.y = basis->a.y + basis->b.y;
			basis->a.modulus = modulus(basis->a.x, basis->a.y, basis->a.z);
		}
		if ( modulus(basis->a.x + basis->c.x, basis->a.z + basis->c.z, basis->a.z + basis->c.z) < basis->a.modulus ) {
			basis->a.x = basis->a.x + basis->c.x; basis->a.y = basis->a.y + basis->c.y; basis->a.y = basis->a.y + basis->c.y;
			basis->a.modulus = modulus(basis->a.x, basis->a.y, basis->a.z);
		}

		if ( modulus(basis->b.x - basis->a.x, basis->b.y - basis->a.y, basis->b.z - basis->a.z) < basis->b.modulus ) {
			basis->b.x = basis->b.x - basis->a.x; basis->b.y = basis->b.y - basis->a.y; basis->b.y = basis->b.y - basis->a.y;
			basis->b.modulus = modulus(basis->b.x, basis->b.y, basis->b.z);
		}
		if ( modulus(basis->b.x - basis->c.x, basis->b.y - basis->c.y, basis->b.z - basis->c.z) < basis->b.modulus ) {
			basis->b.x = basis->b.x - basis->c.x; basis->b.y = basis->b.y - basis->c.y; basis->b.y = basis->b.y - basis->c.y;
			basis->b.modulus = modulus(basis->b.x, basis->b.y, basis->b.z);
		}
		if ( modulus(basis->b.x + basis->a.x, basis->b.y + basis->a.y, basis->b.z + basis->a.z) < basis->b.modulus ) {
			basis->b.x = basis->b.x + basis->a.x; basis->b.y = basis->b.y + basis->a.y; basis->b.y = basis->b.y + basis->a.y;
			basis->b.modulus = modulus(basis->b.x, basis->b.y, basis->b.z);
		}
		if ( modulus(basis->b.x + basis->c.x, basis->b.y + basis->c.y, basis->b.z + basis->c.z) < basis->b.modulus ) {
			basis->b.x = basis->b.x + basis->c.x; basis->b.y = basis->b.y + basis->c.y; basis->b.y = basis->b.y + basis->c.y;
			basis->b.modulus = modulus(basis->b.x, basis->b.y, basis->b.z);
		}

		if ( modulus(basis->c.x - basis->a.x, basis->c.y - basis->a.z, basis->c.z - basis->a.z) < basis->c.modulus ) {
			basis->c.x = basis->c.x - basis->a.x; basis->c.y = basis->c.y - basis->a.y; basis->c.y = basis->c.y - basis->a.y;
			basis->c.modulus = modulus(basis->c.x, basis->c.y, basis->c.z);
		}
		if ( modulus(basis->c.x - basis->b.x, basis->c.y - basis->b.y, basis->c.z - basis->b.z) < basis->c.modulus ) {
			basis->c.x = basis->c.x - basis->b.x; basis->c.y = basis->c.y - basis->b.y; basis->c.y = basis->c.y - basis->b.y;
			basis->c.modulus = modulus(basis->c.x, basis->c.y, basis->c.z);
		}
		if ( modulus(basis->c.x + basis->a.x, basis->c.y + basis->a.z, basis->c.z + basis->a.z) < basis->c.modulus ) {
			basis->c.x = basis->c.x + basis->a.x; basis->c.y = basis->c.y + basis->a.y; basis->c.y = basis->c.y + basis->a.y;
			basis->c.modulus = modulus(basis->c.x, basis->c.y, basis->c.z);
		}
		if ( modulus(basis->c.x + basis->b.x, basis->c.y + basis->b.y, basis->c.z + basis->b.z) < basis->c.modulus ) {
			basis->c.x = basis->c.x + basis->b.x; basis->c.y = basis->c.y + basis->b.y; basis->c.y = basis->c.y + basis->b.y;
			basis->c.modulus = modulus(basis->c.x, basis->c.y, basis->c.z);
		}

	}

}

static Basis *ipr_choose_random_basis(ControlContext *ctx) {
	
	Basis *basis;
	Reflection *reflection;
	int cur_max;
	
	basis = malloc(sizeof(Basis));
	
	basis->a.x = 0;	basis->a.y = 0;	basis->a.z = 0;	basis->a.modulus = 1e30;
	basis->b.x = 0;	basis->b.y = 0;	basis->b.z = 0;	basis->b.modulus = 1e30;
	basis->c.x = 0;	basis->c.y = 0;	basis->c.z = 0;	basis->c.modulus = 1e30;
	cur_max = 1;
	
	reflection = ipr_random_reflection(ctx->reflectionctx);
	
	do {
		double mod;
		unsigned int angle_works;
		
		mod = modulus(reflection->x, reflection->y, reflection->z);
		
		angle_works = 1;
		if ( basis->a.modulus != 1e30 ) {
			double angle = angle_between_d(reflection->x, reflection->y, reflection->z, basis->a.x, basis->a.y, basis->a.z);
			if ( angle < 20 ) angle_works = 0;
			if ( angle > 160 ) angle_works = 0;
		}
		if ( basis->b.modulus != 1e30 ) {
			double angle = angle_between_d(reflection->x, reflection->y, reflection->z, basis->b.x, basis->b.y, basis->b.z);
			if ( angle < 20 ) angle_works = 0;
			if ( angle > 160 ) angle_works = 0;
		}
		if ( basis->c.modulus != 1e30 ) {
			double angle = angle_between_d(reflection->x, reflection->y, reflection->z, basis->c.x, basis->c.y, basis->c.z);
			if ( angle < 20 ) angle_works = 0;
			if ( angle > 160 ) angle_works = 0;
		}
		if ( (mod > 1e9) && angle_works ) {
			cur_max = ipr_choose_max(basis);
			if ( (mod < basis->a.modulus) && (mod < basis->b.modulus) && (mod < basis->c.modulus) && (cur_max == 1) ) {
				basis->a.x = reflection->x;
				basis->a.y = reflection->y;
				basis->a.z = reflection->z;
				basis->a.modulus = mod;
			}
			if ( (mod < basis->a.modulus) && (mod < basis->b.modulus) && (mod < basis->c.modulus) && (cur_max == 2) ) {
				basis->b.x = reflection->x;
				basis->b.y = reflection->y;
				basis->b.z = reflection->z;
				basis->b.modulus = mod;
			}
			if ( (mod < basis->a.modulus) && (mod < basis->b.modulus) && (mod < basis->c.modulus) && (cur_max == 3) ) {
				basis->c.x = reflection->x;
				basis->c.y = reflection->y;
				basis->c.z = reflection->z;
				basis->c.modulus = mod;
			}
		}
		reflection = reflection->next;
		
		/* End of list? */
		if ( !reflection ) {
			if ( (basis->a.modulus == 1e30) || (basis->b.modulus == 1e30) || (basis->c.modulus == 1e30) ) {
				//printf("IP: Failed to find appropriate starting basis!\n");
				free(basis);
				return NULL;
			}
		}
		
	} while ( reflection || (basis->a.modulus == 1e30) || (basis->b.modulus == 1e30) || (basis->c.modulus == 1e30) );
	
	ipr_try_compact(basis);
	
//	basis->a.x = 1.842e9;  basis->a.y = 0.0;  basis->a.z = 0.0;
//	basis->b.x = 0.0;  basis->b.y = 1.842e9;  basis->b.z = 0.0;
//	basis->c.x = 0.0;  basis->c.y = 0.0;  basis->c.z = 1.842e9;
	
	return basis;

}

static double ipr_efom(ReflectionContext *rctx, Basis *basis) {

	int n_indexed, n_counted;
	Reflection *cur;
	
	cur = rctx->reflections;
	n_indexed = 0;
	n_counted = 0;
	while ( cur ) {

		if ( cur->type == REFLECTION_NORMAL ) {
			
			/* Can this basis approximately account for this reflection? */
			double det;
			double a11, a12, a13, a21, a22, a23, a31, a32, a33;
			double h, k, l;
			
			/* Set up the coordinate transform from hkl to xyz */
			a11 = basis->a.x;  a12 = basis->a.y;  a13 = basis->a.z;
			a21 = basis->b.x;  a22 = basis->b.y;  a23 = basis->b.z;
			a31 = basis->c.x;  a32 = basis->c.y;  a33 = basis->c.z;
			
			/* Invert the matrix to get hkl from xyz */
			det = a11*(a22*a33 - a23*a32) - a12*(a21*a33 - a23*a31) + a13*(a21*a32 - a22*a31);
			h = ((a22*a33-a23*a32)*cur->x + (a23*a31-a21*a33)*cur->y + (a21*a32-a22*a31)*cur->z)/ det;
			k = ((a13*a32-a12*a33)*cur->x + (a11*a33-a13*a31)*cur->y + (a12*a31-a11*a32)*cur->z) / det;
			l = ((a12*a23-a13*a22)*cur->x + (a13*a21-a11*a23)*cur->y + (a11*a22-a12*a21)*cur->z) / det;
			
			//printf("%12.0f %12.0f %12.0f ----> %3.2f %3.2f %3.2f", cur->x, cur->y, cur->z, h, k, l);
			//printf("%f %f     %f %f     %f %f\n", cur->x, h*basis->a.x + k*basis->b.x + l*basis->c.x,
			//					cur->y, h*basis->a.y + k*basis->b.y + l*basis->c.y,
			//					cur->z, h*basis->a.z + k*basis->b.z + l*basis->c.z);
			
			/* Calculate the deviations in terms of |a|, |b| and |c| */
			h = fabs(h);  k = fabs(k);  l = fabs(l);
			h -= floor(h);  k -= floor(k);  l -= floor(l);
			if ( h == 1.0 ) h = 0.0;
			if ( k == 1.0 ) k = 0.0;
			if ( l == 1.0 ) l = 0.0;
			//printf(" ----> %3.2f %3.2f %3.2f\n", h, k, l);
			if ( h*h + k*k + l*l <= 0.1*0.1*0.1 ) n_indexed++;
			
			n_counted++;
		
		}
		
		cur = cur->next;

	}
	
	return (double)n_indexed / n_counted;

}

static Basis *ipr_choose_initial_basis(ControlContext *ctx) {

	Basis *basis;
	Basis *best_basis;
	double best_efom;
	int i;
	
	/* Track the best basis throughout the whole procedure */
	best_basis = NULL;  best_efom = 0;
	
	printf("IP: Finding initial basis using eFOM...\n");
	for ( i=1; i<=1000; i++ ) {
		
		double efom;
		
		do { 
			basis = ipr_choose_random_basis(ctx);
		} while ( !basis );
		efom = ipr_efom(ctx->reflectionctx, basis);
		
		if ( (efom > best_efom) || !best_basis ) {
			if ( best_basis ) free(best_basis);
			best_efom = efom;
			best_basis = basis;
			printf("IP:    %3i: eFOM = %f %%\n", i, efom*100);
		
		}
		
	}
	
	printf("IP: Initial basis:\n");
	printf("IP: a = [%10.2e %10.2e %10.2e] mod=%10.2e\n", basis->a.x, basis->a.y, basis->a.z, basis->a.modulus);
	printf("IP: b = [%10.2e %10.2e %10.2e] mod=%10.2e\n", basis->b.x, basis->b.y, basis->b.z, basis->b.modulus);
	printf("IP: c = [%10.2e %10.2e %10.2e] mod=%10.2e\n", basis->c.x, basis->c.y, basis->c.z, basis->c.modulus);
	
	return basis;
}

/* Generate list of reflections to analyse */
static ReflectionContext *ipr_generate(ControlContext *ctx, Basis *basis) {

	ReflectionContext *ordered;
	double max_res;
	signed int h, k, l;
	int max_order_a, max_order_b, max_order_c;
	
	/* NB This assumes the diffraction pattern is "vaguely" centered... */
	if ( ctx->inputfiletype != INPUT_CACHE) {
		int max_width, max_height;
		max_width = ctx->images[0].width;
		max_height = ctx->images[0].height;
		if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
			max_res = sqrt(max_width*max_width + max_height*max_height) * ctx->images[0].pixel_size;
		} else { 
			max_res = sqrt(max_width*max_width + max_height*max_height) / (ctx->images[0].lambda * ctx->images[0].camera_len);
		}
		max_res = max_res / 2;
	} else {
		max_res = 2e10; //works until I put some more in the reflect.cache header
	}
	
	ordered = reflection_init();
	
	max_order_a = max_res/basis->a.modulus;
	max_order_b = max_res/basis->b.modulus;
	max_order_c = max_res/basis->c.modulus;
	for ( h=-max_order_a; h<=max_order_a; h++ ) {
		for ( k=-max_order_b; k<=max_order_b; k++ ) {
			for ( l=-max_order_c; l<=max_order_c; l++ ) {

				double x, y, z;
				
				x = h*basis->a.x + k*basis->b.x + l*basis->c.x;
				y = h*basis->a.y + k*basis->b.y + l*basis->c.y;
				z = h*basis->a.z + k*basis->b.z + l*basis->c.z;
				
				if ( ( x*x + y*y + z*z ) <= max_res*max_res ) {
					Reflection *ref;
					ref = reflection_add(ordered, x, y, z, 1, REFLECTION_GENERATED);
					if ( ref ) {	/* Check it's not NULL */
						ref->h = h;	ref->k = k;	ref->l = l;
					}
				}
				
			}
		}
	}
	
	return ordered;

}

static gint ipr_clicked(GtkWidget *widget, GdkEventButton *event, ControlContext *ctx) {

	ImageReflection *refl;
	size_t n, j;
	
	ctx->ipr_cur_image++;
	if ( ctx->ipr_cur_image == ctx->n_images ) ctx->ipr_cur_image = 0;
	
	imagedisplay_clear_circles(ctx->ipr_id);
	reflection_clear_markers(ctx->reflectionctx);
	
	refl = reproject_get_reflections(ctx->images[ctx->ipr_cur_image], &n, ctx->ipr_lat, ctx);
	for ( j=0; j<n; j++ ) {
		imagedisplay_mark_circle(ctx->ipr_id, refl[j].x, refl[j].y);
	}
	
	imagedisplay_put_data(ctx->ipr_id, ctx->images[ctx->ipr_cur_image]);
	displaywindow_update(ctx->dw);
	
	return 0;

}

int ipr_refine(ControlContext *ctx) {
	
	Basis *basis;
	int finished;
	
	srand(time(NULL));
	
	basis = ipr_choose_initial_basis(ctx);
	if ( !basis ) return -1;
	
	ctx->ipr_basis = basis;
	ctx->ipr_lat = ipr_generate(ctx, basis);
	
	finished = 0;
	do {
	
		size_t n, j;
		ImageReflection *refl;
		
		/* Select an image */
		ctx->ipr_cur_image = 0;//ipr_random_image(ctx);	Temporary!
		
		ctx->ipr_id = imagedisplay_open_with_message(ctx->images[ctx->ipr_cur_image], "Current Image", "Click to change image",
					IMAGEDISPLAY_SHOW_CENTRE | IMAGEDISPLAY_SHOW_TILT_AXIS, G_CALLBACK(ipr_clicked), ctx);
		
		refl = reproject_get_reflections(ctx->images[ctx->ipr_cur_image], &n, ctx->ipr_lat, ctx);
		for ( j=0; j<n; j++ ) {
			imagedisplay_mark_circle(ctx->ipr_id, refl[j].x, refl[j].y);
		}
		
		finished = 1;
		
	} while ( !finished );
	
	return 0;

}