/*
 * refine.c
 *
 * Refine the reconstruction
 *
 * (c) 2007 Thomas White <taw27@cam.ac.uk>
 *
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

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

#include <gtk/gtk.h>
#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>

#include "displaywindow.h"
#include "gtk-valuegraph.h"
#include "basis.h"
#include "reflections.h"
#include "image.h"
#include "reproject.h"
#include "control.h"
#include "mapping.h"
#include "imagedisplay.h"
#include "utils.h"

typedef enum {
	INDEX_A = 1<<0,
	INDEX_B = 1<<1,
	INDEX_C = 1<<2
} RefinementIndex;

#if 0
static const char *refine_decode(RefinementIndex i) {

	switch ( i ) {
	
		case INDEX_A			 : return "a--";
		case INDEX_B			 : return "-b-";
		case INDEX_C			 : return "--c";
		case INDEX_A | INDEX_B		 : return "ab-";
		case INDEX_A | INDEX_C		 : return "a-c";
		case INDEX_B | INDEX_C		 : return "-bc";
		case INDEX_A | INDEX_B | INDEX_C : return "abc";
	
		default				 : return "---";
	
	}

}
#endif

/* Use the IPR algorithm to make "cell" fit the given image */
static ImageFeature *refine_fit_image(Basis *cell, ImageRecord *image, ReflectionList *cell_lattice) {
	
	ImageFeatureList *flist;
	int i;
	Basis cd; /* Cell delta */
	int n_a = 0;
	int n_b = 0;
	int n_c = 0;
	int done = FALSE;
	
	cd.a.x = 0.0;	cd.a.y = 0.0;	cd.a.z = 0.0;
	cd.b.x = 0.0;	cd.b.y = 0.0;	cd.b.z = 0.0;
	cd.c.x = 0.0;	cd.c.y = 0.0;	cd.c.z = 0.0;
	
	if ( !image->rflist ) {
		image->rflist = reproject_get_reflections(image, cell_lattice);
	}
	flist = image->features;
	
	for ( i=0; i<image->rflist->n_features; i++ ) {
		
		double dix, diy;
		double dx, dy, dz, twotheta;
		double old_x, old_y;
		RefinementIndex index, delta, shared;
		signed int h, k, l;
		double a11, a12, a13, a21, a22, a23, a31, a32, a33, det;
		double dh, dk, dl;
		
		/* Skip if no partner */
		if ( !image->rflist->features[i].partner ) continue;
		
		h = image->rflist->features[i].reflection->h;
		k = image->rflist->features[i].reflection->k;
		l = image->rflist->features[i].reflection->l;
		
		/* Determine the difference vector */
		dix = image->rflist->features[i].partner->x - image->rflist->features[i].x;
		diy = image->rflist->features[i].partner->y - image->rflist->features[i].y;
	//	printf("RF: Feature %3i: %3i %3i %3i dev = %f %f px\n", i, h, k, l, dix, diy);
		
		/* Map the difference vector to the relevant tilted plane */
		old_x = image->rflist->features[i].partner->x;
		old_y = image->rflist->features[i].partner->y;
		image->rflist->features[i].partner->x = dix + image->rflist->features[i].partner->parent->x_centre;
		image->rflist->features[i].partner->y = diy + image->rflist->features[i].partner->parent->y_centre;
		mapping_map_to_space(image->rflist->features[i].partner, &dx, &dy, &dz, &twotheta);
		image->rflist->features[i].partner->x = old_x;
		image->rflist->features[i].partner->y = old_y;
	//	printf("RF: dev=%8e %8e %8e (%5f mrad)\n", dx, dy, dz, twotheta*1e3);
		
		/* Select the basis vectors which are allowed to be altered */
		index = 0;
		if ( h ) index |= INDEX_A;
		if ( k ) index |= INDEX_B;
		if ( l ) index |= INDEX_C;
		assert(index != 0);	/* Can't refine using the central beam! */
		
		/* Set up the coordinate transform from hkl to xyz */
		a11 = cell->a.x;  a12 = cell->a.y;  a13 = cell->a.z;
		a21 = cell->b.x;  a22 = cell->b.y;  a23 = cell->b.z;
		a31 = cell->c.x;  a32 = cell->c.y;  a33 = cell->c.z;
		
		/* Invert the matrix to get dh,dk,dl from dx,dy,dz */
		det = a11*(a22*a33 - a23*a32) - a12*(a21*a33 - a23*a31) + a13*(a21*a32 - a22*a31);
		dh = ((a22*a33-a23*a32)*dx + (a23*a31-a21*a33)*dy + (a21*a32-a22*a31)*dz) / det;
		dk = ((a13*a32-a12*a33)*dx + (a11*a33-a13*a31)*dy + (a12*a31-a11*a32)*dz) / det;
		dl = ((a12*a23-a13*a22)*dx + (a13*a21-a11*a23)*dy + (a11*a22-a12*a21)*dz) / det;
	//	printf("RF: dev(hkl) = %f %f %f\n", dh, dk, dl);
		
		delta = 0;
		if ( fabs(dh)/(fabs(dh)+fabs(dk)+fabs(dl)) > 0.001 ) delta |= INDEX_A;
		if ( fabs(dk)/(fabs(dh)+fabs(dk)+fabs(dl)) > 0.001 ) delta |= INDEX_B;
		if ( fabs(dl)/(fabs(dh)+fabs(dk)+fabs(dl)) > 0.001 ) delta |= INDEX_C;
		
		shared = index & delta;
		
	//	printf("RF: index: %s\nRF: delta: %s\nRF: shared: %s\n", refine_decode(index), refine_decode(delta), refine_decode(shared));
		
		if ( shared == 0 ) {
			/* No indices left - 'pure shear' (delta is perpendicular (in the abc basis) to index) */
			shared = index;
		//	printf("RF: Pure shear.\n");
		}
		
		if ( shared & INDEX_A ) {
			double w = (double)abs(h) / (abs(h)+abs(k)+abs(l));
		//	printf("RF: w(a) = %f\n", w);
			cd.a.x += w*dx / (double)h;
			cd.a.y += w*dy / (double)h;
			cd.a.z += w*dz / (double)h;
			n_a++;
		}
		if ( shared & INDEX_B ) {
			double w = (double)abs(k) / (abs(h)+abs(k)+abs(l));
		//	printf("RF: w(b) = %f\n", w);
			cd.b.x += w*dx / (double)k;
			cd.b.y += w*dy / (double)k;
			cd.b.z += w*dz / (double)k;
			n_b++;
		}
		if ( shared & INDEX_C ) {
			double w = (double)abs(l) / (abs(h)+abs(k)+abs(l));
		//	printf("RF: w(c) = %f\n", w);
			cd.c.x += w*dx / (double)l;
			cd.c.y += w*dy / (double)l;
			cd.c.z += w*dz / (double)l;
			n_c++;
		}
		
	//	printf("RF: Distortion along x: %+8e = %+8e\n", h*cd.a.x + k*cd.b.x + l*cd.c.x, dx);
	//	printf("RF: Distortion along y: %+8e = %+8e\n", h*cd.a.y + k*cd.b.y + l*cd.c.y, dy);
	//	printf("RF: Distortion along z: %+8e = %+8e\n", h*cd.a.z + k*cd.b.z + l*cd.c.z, dz);
		
		done = TRUE;
		
	//	break;
		
	}
	
	if ( !done ) {
		printf("RF: No partners found.\n");
		return NULL;
	}
	
	//printf("RF: n_a=%i, n_b=%i, n_c=%i\n", n_a, n_b, n_c);
	if ( n_a ) {
		cd.a.x /= (double)n_a;	cd.a.y /= (double)n_a;	cd.a.z /= (double)n_a;
	}
	if ( n_b ) {
		cd.b.x /= (double)n_b;	cd.b.y /= (double)n_b;	cd.b.z /= (double)n_b;
	}
	if ( n_c ) {
		cd.c.x /= (double)n_c;	cd.c.y /= (double)n_c;	cd.c.z /= (double)n_c;
	}
	
//	printf("RF: Total distortion(a) = %+8e %+8e %+8e\n", cd.a.x, cd.a.y, cd.a.z);
//	printf("RF: Total distortion(b) = %+8e %+8e %+8e\n", cd.b.x, cd.b.y, cd.b.z);
//	printf("RF: Total distortion(c) = %+8e %+8e %+8e\n", cd.c.x, cd.c.y, cd.c.z);
	
	cell->a.x += cd.a.x;	cell->a.y += cd.a.y;	cell->a.z += cd.a.z;
	cell->b.x += cd.b.x;	cell->b.y += cd.b.y;	cell->b.z += cd.b.z;
	cell->c.x += cd.c.x;	cell->c.y += cd.c.y;	cell->c.z += cd.c.z;
	
	//return image->rflist->features[i].partner;
	
	return NULL;
	
}

static gint refine_step(GtkWidget *step_button, ControlContext *ctx) {

	if ( (ctx->reproject_id) && (ctx->cell_lattice) ) {
		
		ImageFeature *fitted;
		
		fitted = refine_fit_image(ctx->cell, &ctx->images->images[ctx->reproject_cur_image], ctx->cell_lattice);
		
		ctx->images->images[ctx->reproject_cur_image].rflist = NULL;
		reproject_lattice_changed(ctx);
		displaywindow_update(ctx->dw);
		
		if ( fitted ) {
			imagedisplay_add_mark(ctx->reproject_id, fitted->x,fitted->y, IMAGEDISPLAY_MARK_CIRCLE_3);
		}
		
	} else {
		displaywindow_error("Please first open the reprojection window and select the image to fit", ctx->dw);
	}

	return 0;

}

static int refine_sequence_sweep(ControlContext *ctx, double *fit, double *warp) {

	int i;
	double series_dev_max = 0;
	double series_dev_min = +HUGE_VAL;
	double series_dev_mean = 0;
	int series_dev_n = 0;
	
	for ( i=0; i<ctx->images->n_images; i++ ) {

		/* Ensure lattice is up to date */
		reproject_lattice_changed(ctx);
		ctx->images->images[i].rflist = NULL;	/* Invalidate reprojection for this image - it's wrong */
		
		if ( is_odd(i) ) {
			
			/* Odd-numbered images: measure */
			ImageRecord *image;
			int j, n;
			double image_dev_mean = 0;
			
			image = &ctx->images->images[i];
			if ( !image->rflist ) {
				image->rflist = reproject_get_reflections(image, ctx->cell_lattice);
			}
			
			n = 0;
			for ( j=0; j<image->rflist->n_features; j++ ) {
				
				double dix, diy;
				
				/* Skip if no partner */
				if ( !image->rflist->features[j].partner ) continue;
				
				/* Determine the difference vector */
				dix = image->rflist->features[j].partner->x - image->rflist->features[j].x;
				diy = image->rflist->features[j].partner->y - image->rflist->features[j].y;
				
				image_dev_mean += sqrt(dix*dix + diy*diy);
				n++;
				
			}
			image_dev_mean /= n;
			
			if ( image_dev_mean > series_dev_max ) series_dev_max = image_dev_mean;
			if ( image_dev_mean < series_dev_min ) series_dev_min = image_dev_mean;
			series_dev_mean += image_dev_mean;
			series_dev_n++;
			
		} else {
			
			/* Even-numbered images: fit */
			refine_fit_image(ctx->cell, &ctx->images->images[i], ctx->cell_lattice);
			
		}

	}
	
	series_dev_mean /= series_dev_n;
	*fit = series_dev_mean;
	*warp = (series_dev_max - series_dev_min)/series_dev_min;
	
	return 0;

}

static gint refine_sequence(GtkWidget *step_button, ControlContext *ctx) {

	double omega_offs;
	GtkWidget *window_fit;
	GtkWidget *graph_fit;
	double *fit_vals;
	GtkWidget *window_warp;
	GtkWidget *graph_warp;
	double *warp_vals;
	size_t idx;
	ImageDisplay *id;
	
	fit_vals = malloc(401*sizeof(double));
	warp_vals = malloc(401*sizeof(double));
	idx = 0;

	if ( !ctx->cell ) {
		displaywindow_error("No reciprocal unit cell has been found.", ctx->dw);
		return 0;
	}
	
	/* Temporarily disable ImageDisplay stuff */
	id = ctx->reproject_id;
	ctx->reproject_id = NULL;
	
	for ( omega_offs=-2.0; omega_offs<=2.0; omega_offs+=0.01 ) {
	
		double fit, warp;
		int i;
		Basis cell_copy;
		
		memcpy(&cell_copy, ctx->cell, sizeof(Basis));
		for ( i=0; i<ctx->images->n_images; i++ ) {
			ctx->images->images[i].omega += omega_offs;
		}
		
		if ( refine_sequence_sweep(ctx, &fit, &warp) ) {
			printf("RF: Sequencer sweep failed\n");
			return 0;
		}
		printf("RF: omega_offs=%f, fit=%f, warp=%f\n", omega_offs, fit, warp);
		fit_vals[idx] = fit;
		warp_vals[idx++] = warp;

		for ( i=0; i<ctx->images->n_images; i++ ) {
			ctx->images->images[i].omega -= omega_offs;
		}
		memcpy(ctx->cell, &cell_copy, sizeof(Basis));
		
	}
	
	ctx->reproject_id = id;
	reproject_lattice_changed(ctx);
	
	window_fit = gtk_window_new(GTK_WINDOW_TOPLEVEL);
	gtk_window_set_default_size(GTK_WINDOW(window_fit), 640, 256);
	gtk_window_set_title(GTK_WINDOW(window_fit), "Omega-Search Graph: Fit");
	graph_fit = gtk_value_graph_new();
	gtk_value_graph_set_data(GTK_VALUE_GRAPH(graph_fit), fit_vals, idx);
	gtk_container_add(GTK_CONTAINER(window_fit), graph_fit);
	gtk_widget_show_all(window_fit);
	
	window_warp = gtk_window_new(GTK_WINDOW_TOPLEVEL);
	gtk_window_set_default_size(GTK_WINDOW(window_warp), 640, 256);
	gtk_window_set_title(GTK_WINDOW(window_warp), "Omega-Search Graph: Warp");
	graph_warp = gtk_value_graph_new();
	gtk_value_graph_set_data(GTK_VALUE_GRAPH(graph_warp), warp_vals, idx);
	gtk_container_add(GTK_CONTAINER(window_warp), graph_warp);
	gtk_widget_show_all(window_warp);
	
	return 0;

}

static gint refine_response(GtkWidget *refine_window, gint response, ControlContext *ctx) {

	ctx->refine_window = NULL;
	gtk_widget_destroy(refine_window);
	
	return 0;

}

void refine_open(ControlContext *ctx) {

	GtkWidget *vbox;
	GtkWidget *hbox;
	GtkWidget *table;
	GtkWidget *label;
	GtkWidget *step_button;
	GtkWidget *sequence_button;
	
	if ( ctx->refine_window ) return;
	
	ctx->refine_window = gtk_dialog_new_with_buttons("Refine Reconstruction", GTK_WINDOW(ctx->dw->window),
		GTK_DIALOG_DESTROY_WITH_PARENT,	GTK_STOCK_CLOSE, GTK_RESPONSE_CLOSE, NULL);
	gtk_window_set_default_size(GTK_WINDOW(ctx->refine_window), 256, -1);
	
	vbox = gtk_vbox_new(FALSE, 0);
	hbox = gtk_hbox_new(TRUE, 0);
	gtk_box_pack_start(GTK_BOX(GTK_DIALOG(ctx->refine_window)->vbox), GTK_WIDGET(hbox), FALSE, FALSE, 7);
	gtk_box_pack_start(GTK_BOX(hbox), GTK_WIDGET(vbox), FALSE, FALSE, 5);
	
	table = gtk_table_new(5, 1, FALSE);
	gtk_table_set_row_spacings(GTK_TABLE(table), 5);
	gtk_table_set_col_spacings(GTK_TABLE(table), 5);
	gtk_box_pack_start(GTK_BOX(vbox), GTK_WIDGET(table), FALSE, FALSE, 0);
	
	label = gtk_label_new("Steps");
	gtk_label_set_markup(GTK_LABEL(label), "<span weight=\"bold\">Steps</span>");
	gtk_misc_set_alignment(GTK_MISC(label), 0.0, 0.5);
	gtk_table_attach_defaults(GTK_TABLE(table), label, 1, 2, 1, 2);
	step_button = gtk_button_new_with_label("Refine Lattice to Fit Current Pattern");
	gtk_table_attach_defaults(GTK_TABLE(table), step_button, 1, 2, 2, 3);
	g_signal_connect(G_OBJECT(step_button), "clicked", G_CALLBACK(refine_step), ctx);
	
	label = gtk_label_new("Sequencing");
	gtk_misc_set_alignment(GTK_MISC(label), 0.0, 0.5);
	gtk_label_set_markup(GTK_LABEL(label), "<span weight=\"bold\">Sequencing</span>");
	gtk_table_attach_defaults(GTK_TABLE(table), label, 1, 2, 4, 5);
	sequence_button = gtk_button_new_with_label("Run Sequencer");
	gtk_table_attach_defaults(GTK_TABLE(table), sequence_button, 1, 2, 5, 6);
	g_signal_connect(G_OBJECT(sequence_button), "clicked", G_CALLBACK(refine_sequence), ctx);
		
	g_signal_connect(G_OBJECT(ctx->refine_window), "response", G_CALLBACK(refine_response), ctx);
	gtk_widget_show_all(ctx->refine_window);
	

}