/*
 * 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 <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_linalg.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 {
	AXIS_X = 1,
	AXIS_Y = 2
} Axis;

/* 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;
	gsl_matrix *M;
	gsl_vector *q;
	gsl_vector *p;
	int j, ns;
	Axis axis;
	double dax = 0.0, dbx = 0.0, dcx = 0.0;
	double day = 0.0, dby = 0.0, dcy = 0.0;
	double dlx = 0.0, dly = 0.0, dlz = 0.0;
	
	flist = image->features;
	
	M = gsl_matrix_alloc(3, 3);
	p = gsl_vector_alloc(3);
	
	for ( axis = AXIS_X; axis <= AXIS_Y; axis++ ) {
	
		gsl_permutation *perm;
		int s;
		
		gsl_matrix_set_zero(M);
		gsl_vector_set_zero(p);
		
		ns = 0;
		for ( j=0; j<image->rflist->n_features; j++ ) {
		
			double val;
			ImageFeature *rf;
			signed int h, k, l;
			double xy;
			double dix, diy, dx, dy;
			double old_x, old_y;
			
			rf = &image->rflist->features[j];
			
			if ( !rf->partner ) continue;
			
			h = rf->reflection->h;
			k = rf->reflection->k;
			l = rf->reflection->l;
			
			/* Determine the difference vector */
			dix = rf->partner->x - rf->x;
			diy = rf->partner->y - rf->y;
			printf("RF: Feature %3i: %3i %3i %3i dev = %+9.5f %+9.5f px ", j, h, k, l, dix, diy);
			
			old_x = rf->partner->x;
			old_y = rf->partner->y;
			rf->partner->x = dix + rf->partner->parent->x_centre;
			rf->partner->y = diy + rf->partner->parent->y_centre;
			mapping_scale(rf->partner, &dx, &dy);
			rf->partner->x = old_x;
			rf->partner->y = old_y;
			printf("=> %+10.5f %+10.5f nm^-1\n", dx/1e9, dy/1e9);
			
			xy = 0;
			switch ( axis ) {
				case AXIS_X : xy = dx; break;
				case AXIS_Y : xy = dy; break;
			}
			
			/* Elements of "p" */
			val = gsl_vector_get(p, 0);	val += xy * h;	gsl_vector_set(p, 0, val);
			val = gsl_vector_get(p, 1);	val += xy * k;	gsl_vector_set(p, 1, val);
			val = gsl_vector_get(p, 2);	val += xy * l;	gsl_vector_set(p, 2, val);
			
			/* Elements of "M" */
			val = gsl_matrix_get(M, 0, 0);	val += h * h;	gsl_matrix_set(M, 0, 0, val);
			val = gsl_matrix_get(M, 0, 1);	val += k * h;	gsl_matrix_set(M, 0, 1, val);
			val = gsl_matrix_get(M, 0, 2);	val += l * h;	gsl_matrix_set(M, 0, 2, val);
			
			val = gsl_matrix_get(M, 1, 0);	val += h * k;	gsl_matrix_set(M, 1, 0, val);
			val = gsl_matrix_get(M, 1, 1);	val += k * k;	gsl_matrix_set(M, 1, 1, val);
			val = gsl_matrix_get(M, 1, 2);	val += l * k;	gsl_matrix_set(M, 1, 2, val);
			
			val = gsl_matrix_get(M, 2, 0);	val += h * l;	gsl_matrix_set(M, 2, 0, val);
			val = gsl_matrix_get(M, 2, 1);	val += k * l;	gsl_matrix_set(M, 2, 1, val);
			val = gsl_matrix_get(M, 2, 2);	val += l * l;	gsl_matrix_set(M, 2, 2, val);
			
			ns++;
		
		}
		
		if ( !ns ) {
			printf("RF: No partners found\n");
			gsl_matrix_free(M);
			gsl_vector_free(p);
			return NULL;
		}	
		
		/* Do the fitting */
		perm = gsl_permutation_alloc(M->size1);
		gsl_linalg_LU_decomp(M, perm, &s);
		
		q = gsl_vector_alloc(3);	/* This is the "answer" */
		gsl_vector_set_zero(q);
		gsl_linalg_LU_solve(M, perm, p, q);
		
		switch ( axis ) {
			case AXIS_X : {
				dax = gsl_vector_get(q, 0);
				dbx = gsl_vector_get(q, 1);	/* These are the deviations, in the direction "x" of the image coordinate  */
				dcx = gsl_vector_get(q, 2);	/* system, of a,b and c */
				break;
			}
			case AXIS_Y : {
				day = gsl_vector_get(q, 0);
				dby = gsl_vector_get(q, 1);	/* These are the deviations, in the direction "y" of the image coordinate  */
				dcy = gsl_vector_get(q, 2);	/* system, of a,b and c */
				break;
			}
		}
		
		gsl_permutation_free(perm);
		
	}
	
	gsl_matrix_free(M);
	gsl_vector_free(p);
	gsl_vector_free(q);
	
	printf("a should change by %+7.5f %+7.5f nm^-1 in the image plane\n", dax/1e9, day/1e9);
	printf("b should change by %+7.5f %+7.5f nm^-1 in the image plane\n", dbx/1e9, dby/1e9);
	printf("c should change by %+7.5f %+7.5f nm^-1 in the image plane\n", dcx/1e9, dcy/1e9);
	
	/* Update the cell */
	mapping_rotate(dax, day, 0.0, &dlx, &dly, &dlz, image->omega, image->tilt);
	printf("a changed by %+7.5f %+7.5f %+7.5f nm^-1 in reciprocal space (%+10.5f %+10.5f %+10.5f %%)\n", dlx/1e9, dly/1e9, dlz/1e9,
										  100*dlx/cell->a.x, 100*dly/cell->a.y, 100*dlz/cell->a.z);
	cell->a.x += dlx;  cell->a.y += dly;  cell->a.z += dlz;
	
	mapping_rotate(dbx, dby, 0.0, &dlx, &dly, &dlz, image->omega, image->tilt);
	printf("b changed by %+7.5f %+7.5f %+7.5f nm^-1 in reciprocal space (%+10.5f %+10.5f %+10.5f %%)\n", dlx/1e9, dly/1e9, dlz/1e9,
										  100*dlx/cell->b.x, 100*dly/cell->b.y, 100*dlz/cell->b.z);
	cell->b.x += dlx;  cell->b.y += dly;  cell->b.z += dlz;
	
	mapping_rotate(dcx, dcy, 0.0, &dlx, &dly, &dlz, image->omega, image->tilt);
	printf("c changed by %+7.5f %+7.5f %+7.5f nm^-1 in reciprocal space (%+10.5f %+10.5f %+10.5f %%)\n", dlx/1e9, dly/1e9, dlz/1e9,
										  100*dlx/cell->c.x, 100*dly/cell->c.y, 100*dlz/cell->c.z);
	cell->c.x += dlx;  cell->c.y += dly;  cell->c.z += dlz;
	
	return NULL;
	
}

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;
	
	/* Ensure that ctx->cell_lattice is set up */
	reproject_cell_to_lattice(ctx);
	
	for ( i=0; i<ctx->images->n_images; i++ ) {

		ImageRecord *image;
		int j, n;
		double image_dev_mean = 0;
		
		image = &ctx->images->images[i];
		
		/* Fit this image and update ctx->cell_lattice, index the selected pattern */
		if ( !image->rflist ) image->rflist = reproject_get_reflections(image, ctx->cell_lattice);
		refine_fit_image(ctx->cell, image, ctx->cell_lattice);
		reproject_cell_to_lattice(ctx);
		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++;
		
	}
	
	series_dev_mean /= series_dev_n;
	*fit = series_dev_mean;
	*warp = (series_dev_max - series_dev_min)/series_dev_min;
	
	return 0;

}

void refine_do_stack(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;
	
	fit_vals = malloc(401*sizeof(double));
	warp_vals = malloc(401*sizeof(double));
	idx = 0;

	if ( !ctx->cell_lattice ) {
		displaywindow_error("No reciprocal unit cell has been found.", ctx->dw);
		return;
	}
	
	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");
		}
		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));
		
	}
	
	displaywindow_update(ctx->dw);
	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);
	
}

void refine_do_image(ControlContext *ctx) {

	if ( !ctx->cell_lattice ) {
		displaywindow_error("No reciprocal unit cell has been found.", ctx->dw);
		return;
	}
	
	ImageFeature *fitted;
	
	fitted = refine_fit_image(ctx->cell, &ctx->images->images[ctx->dw->cur_image], ctx->cell_lattice);
	
	ctx->images->images[ctx->dw->cur_image].rflist = NULL;
	reproject_lattice_changed(ctx);
	displaywindow_update(ctx->dw);
		
}