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

#if 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 "control.h"
#include "displaywindow.h"
#include "image.h"
#include "reproject.h"
#include "mapping.h"
#include "refine.h"
#include "gtk-valuegraph.h"
#include "utils.h"

/* Divide numbers by this for display */
#define DISPFACTOR 1.0e9

/* Number of parameters */
#define NUM_PARAMS 9

/* Refine debug */
#define REFINE_DEBUG_MORE 0

#define DELTA (0.0001e9)

typedef struct {
	/* Indicies of reflection */
	signed int h;	signed int k;	signed int l;
	/* Deviation vector for this reflection in nm^-1 */
	double dx;	double dy;
	/* x,y,z components of "image x" */
	double xx;	double xy;	double xz;
	/* x,y,z components of "image y" */
	double yx;	double yy;	double yz;

	#if REFINE_DEBUG_MORE
	/* Image number this dev came from, for debugging purposes */
	int image;
	double rdx;	double rdy;
	#endif /* REFINE_DEBUG_MORE */

} Deviation;

static void refine_printf(DisplayWindow *dw, char *fmt, ...)
{
	va_list ap;
	char tmp[1024];

	va_start(ap, fmt);
	vsnprintf(tmp, 1023, fmt, ap);
	va_end(ap);

	printf("%s", tmp);
	displaywindow_message(dw, tmp);
}

void refine_do_sequence(ControlContext *ctx)
{
	double omega_offs;
	int idx;
	double *fit_vals;
	GtkWidget *window_fit;
	GtkWidget *graph_fit;
	double fit_best, omega_offs_best;
	int j;

	fit_vals = malloc(401*sizeof(double));
	idx = 0;

	fit_best = 1000.0e9;
	omega_offs_best = 0.0;
	for ( omega_offs=-deg2rad(2.0); omega_offs<=deg2rad(2.0);
						omega_offs+=deg2rad(0.01) ) {

		double fit;
		int i;
		Basis cell_copy;

		cell_copy = *ctx->cell;

		for ( i=0; i<ctx->images->n_images; i++ ) {
			ctx->images->images[i].omega += omega_offs;
		}
		reproject_lattice_changed(ctx);

		fit = refine_do_cell(ctx);
		refine_printf(ctx->dw, "                           "
					"omega_offs=%f deg, fit=%f nm^-1\n",
			      rad2deg(omega_offs), fit/DISPFACTOR);
		fit_vals[idx++] = fit;
		if ( fit < fit_best ) {
			fit_best = fit;
			omega_offs_best = omega_offs;
		}

		for ( i=0; i<ctx->images->n_images; i++ ) {
			ctx->images->images[i].omega -= omega_offs;
		}
		*ctx->cell = cell_copy;

	}

	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);

	/* Perform final refinement */
	refine_printf(ctx->dw, "Best omega offset = %f deg (%f nm^-1)\n",
		      rad2deg(omega_offs_best), fit_best/DISPFACTOR);
	for ( j=0; j<ctx->images->n_images; j++ ) {
		ctx->images->images[j].omega += omega_offs_best;
	}
	refine_do_cell(ctx);
	reproject_lattice_changed(ctx);
	mapping_adjust_axis(ctx, omega_offs_best);
}

static double refine_mean_dev(Deviation *d, int nf, Basis *devcell, int disp)
{
	double fom = 0.0;
	int f;

	for ( f=0; f<nf; f++ ) {

		double xdf, ydf, zdf;
		double dx, dy;

		/* Calculate 3D deviation vector */
		xdf = d[f].h*devcell->a.x + d[f].k*devcell->b.x
							+ d[f].l*devcell->c.x;
		ydf = d[f].h*devcell->a.y + d[f].k*devcell->b.y
							+ d[f].l*devcell->c.y;
		zdf = d[f].h*devcell->a.z + d[f].k*devcell->b.z
							+ d[f].l*devcell->c.z;

		/* Project into original image */
		dx = xdf*d[f].xx + ydf*d[f].xy + zdf*d[f].xz;
		dy = xdf*d[f].yx + ydf*d[f].yy + zdf*d[f].yz;

		dx -= d[f].dx;
		dy -= d[f].dy;

		#if REFINE_DEBUG_MORE
		if ( disp ) {
			d[f].rdx = dx;
			d[f].rdy = dy;
			printf("Image %2i ref %3i %3i %3i dev %8.5f %8.5f\n",
				d[f].image, d[f].h, d[f].k, d[f].l,
				dx/DISPFACTOR, dy/DISPFACTOR);
		}
		#endif /* REFINE_DEBUG_MORE */

		fom += sqrt(dx*dx + dy*dy);

	}

	return fom/nf;

}

static void refine_cell_delta(Basis *devcell_try, int comp) {

	switch ( comp ) {
		case  0 : break;
		case  1 : devcell_try->a.x += DELTA; break;
		case  2 : devcell_try->a.y += DELTA; break;
		case  3 : devcell_try->a.z += DELTA; break;
		case  4 : devcell_try->b.x += DELTA; break;
		case  5 : devcell_try->b.y += DELTA; break;
		case  6 : devcell_try->b.z += DELTA; break;
		case  7 : devcell_try->c.x += DELTA; break;
		case  8 : devcell_try->c.y += DELTA; break;
		case  9 : devcell_try->c.z += DELTA; break;
		case 10 : devcell_try->a.x -= DELTA; break;
		case 11 : devcell_try->a.y -= DELTA; break;
		case 12 : devcell_try->a.z -= DELTA; break;
		case 13 : devcell_try->b.x -= DELTA; break;
		case 14 : devcell_try->b.y -= DELTA; break;
		case 15 : devcell_try->b.z -= DELTA; break;
		case 16 : devcell_try->c.x -= DELTA; break;
		case 17 : devcell_try->c.y -= DELTA; break;
		case 18 : devcell_try->c.z -= DELTA; break;
		default : fprintf(stderr,
				"refine_cell_delta: argument out of range\n");
	}

}

static void refine_show_cell(DisplayWindow *dw, Basis cell) {

	refine_printf(dw, "a: %+10.8f %+10.8f %+10.8f\n",
			cell.a.x/DISPFACTOR, cell.a.y/DISPFACTOR,
							cell.a.z/DISPFACTOR);
	refine_printf(dw, "b: %+10.8f %+10.8f %+10.8f\n",
			cell.b.x/DISPFACTOR, cell.b.y/DISPFACTOR,
							cell.b.z/DISPFACTOR);
	refine_printf(dw, "c: %+10.8f %+10.8f %+10.8f\n",
			cell.c.x/DISPFACTOR, cell.c.y/DISPFACTOR,
							cell.c.z/DISPFACTOR);

}

static Basis refine_test_confidence(Deviation *d, int nf) {

	Basis *devcell;
	double val, new_val;
	Basis confidence;

	devcell = malloc(sizeof(Basis));
	devcell->a.x = 0.0e9;  devcell->b.x = 0.0e9;  devcell->c.x = 0.0e9;
	devcell->a.y = 0.0e9;  devcell->b.y = 0.0e9;  devcell->c.y = 0.0e9;
	devcell->a.z = 0.0e9;  devcell->b.z = 0.0e9;  devcell->c.z = 0.0e9;
	val = refine_mean_dev(d, nf, devcell, 0);

	devcell->a.x = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.a.x = new_val-val;
	devcell->a.x = 0.0;

	devcell->a.y = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.a.y = new_val-val;
	devcell->a.y = 0.0;

	devcell->a.z = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.a.z = new_val-val;
	devcell->a.z = 0.0;

	devcell->b.x = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.b.x = new_val-val;
	devcell->b.x = 0.0;

	devcell->b.y = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.b.y = new_val-val;
	devcell->b.y = 0.0;

	devcell->b.z = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.b.z = new_val-val;
	devcell->b.z = 0.0;

	devcell->c.x = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.c.x = new_val-val;
	devcell->c.x = 0.0;

	devcell->c.y = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.c.y = new_val-val;
	devcell->c.y = 0.0;

	devcell->c.z = 1.0e9;
	new_val = refine_mean_dev(d, nf, devcell, 0);
	confidence.c.z = new_val-val;
	devcell->c.z = 0.0;

	return confidence;

}

double refine_do_cell(ControlContext *ctx) {

	Deviation *d;
	int i, nf, f;
	Basis *devcell;
	Basis *devcell_try;
	Basis *devcell_try_best;
	double mean_dev;
	double mean_dev_try;
	int it;
	Basis conf;
	double conf_threshold;

	if ( !ctx->cell_lattice ) {
		displaywindow_error("No reciprocal unit cell has been found.",
								ctx->dw);
		return -1;
	}

	if ( ctx->images->n_images == 0 ) {
		displaywindow_error("There are no images to refine against.",
								ctx->dw);
		return -1;
	}

	/* Determine the size of the 'deviation table' */
	nf = 0;
	for ( i=0; i<ctx->images->n_images; i++ ) {

		int j;

		if ( !ctx->images->images[i].rflist ) {
			ctx->images->images[i].rflist =
			 reproject_get_reflections(&ctx->images->images[i],
							ctx->cell_lattice, ctx);
		}

		for ( j=0; j<ctx->images->images[i].rflist->n_features; j++ ) {
			if ( ctx->images->images[i].rflist->features[j].partner
								!= NULL ) nf++;
		}

		refine_printf(ctx->dw, "%i features from image %i\n", nf, i);

	}
	refine_printf(ctx->dw, "There are %i partnered features in total\n",
									nf);

	/* Initialise the 'deviation table' */
	d = malloc(nf*sizeof(Deviation));
	f = 0;
	for ( i=0; i<ctx->images->n_images; i++ ) {

		ImageRecord *image;
		int j;

		image = &ctx->images->images[i];

		for ( j=0; j<ctx->images->images[i].rflist->n_features; j++ ) {

			ImageFeature *rf;
			double dix, diy, dx, dy;
			double old_x, old_y;

			rf = &image->rflist->features[j];
			if ( !rf->partner ) continue;

			d[f].h = rf->reflection->h;
			d[f].k = rf->reflection->k;
			d[f].l = rf->reflection->l;

			/* Determine the difference vector */
			dix = rf->partner->x - rf->x;
			diy = rf->partner->y - rf->y;
			refine_printf(ctx->dw,
				"%3i: %3i: %3i %3i %3i dev %+9.5f %+9.5f px ",
				i, j, d[f].h, d[f].k, d[f].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;
			double mod = sqrt(dx*dx + dy*dy)/DISPFACTOR;
			refine_printf(ctx->dw,
				"= %+10.5f %+10.5f nm^-1 (length %7.5f nm^1)\n",
				dx/DISPFACTOR, dy/DISPFACTOR, mod);

			d[f].dx = dx;
			d[f].dy = dy;

			/* Store plane normal for use in projection later */
			mapping_rotate(1.0, 0.0, 0.0,
					&d[f].xx, &d[f].xy, &d[f].xz,
					image->omega, image->tilt);
			mapping_rotate(0.0, 1.0, 0.0,
					&d[f].yx, &d[f].yy, &d[f].yz,
					image->omega, image->tilt);

			#if REFINE_DEBUG_MORE
			d[f].image = i;
			#endif /* REFINE_DEBUG_MORE */

			f++;

		}

	}
	assert( f == nf );

	/* Initial situation */
	devcell = malloc(sizeof(Basis));
	devcell->a.x = 0.0e9;  devcell->b.x = 0.0e9;  devcell->c.x = 0.0e9;
	devcell->a.y = 0.0e9;  devcell->b.y = 0.0e9;  devcell->c.y = 0.0e9;
	devcell->a.z = 0.0e9;  devcell->b.z = 0.0e9;  devcell->c.z = 0.0e9;
	mean_dev = refine_mean_dev(d, nf, devcell, 0);
	refine_printf(ctx->dw, "Initial mean deviation: %13.8f nm^1\n",
							mean_dev/DISPFACTOR);

	/* Test confidence */
	conf = refine_test_confidence(d, nf);

	/* Determine direction of steepest gradient */
	devcell_try = malloc(sizeof(Basis));
	devcell_try_best = malloc(sizeof(Basis));
	for ( it=1; it<=10000; it++ ) {

		int found = 0;
		int comp1, comp2, comp3;

		#if REFINE_DEBUG_MORE
		refine_mean_dev(d, nf, devcell, 1);
		printf(ctx->dw, "Iteration %i starts at dev %14.12f\n",
						it, mean_dev/DISPFACTOR);
		printf(ctx->dw, "Current dev cell:\n");
		printf(ctx->dw, "a = %8.5f %8.5f %8.5f\n",
				devcell->a.x/DISPFACTOR,
				devcell->a.y/DISPFACTOR,
				devcell->a.z/DISPFACTOR);
		printf(ctx->dw, "b = %8.5f %8.5f %8.5f\n",
				devcell->b.x/DISPFACTOR,
				devcell->b.y/DISPFACTOR,
				devcell->b.z/DISPFACTOR);
		printf(ctx->dw, "c = %8.5f %8.5f %8.5f\n",
				devcell->c.x/DISPFACTOR,
				devcell->c.y/DISPFACTOR,
				devcell->c.z/DISPFACTOR);
		#endif /* REFINE_DEBUG_MORE */

		for ( comp1=1; comp1<19; comp1++ ) {
		for ( comp2=0; comp2<19; comp2++ ) {
		for ( comp3=0; comp3<19; comp3++ ) {

			memcpy(devcell_try, devcell, sizeof(Basis));
			refine_cell_delta(devcell_try, comp1);
			refine_cell_delta(devcell_try, comp2);
			refine_cell_delta(devcell_try, comp3);
			mean_dev_try = refine_mean_dev(d, nf, devcell_try, 0);

			/* Improvement greater than the tolerance? */
			if ( mean_dev_try < mean_dev-0.001  ) {
				mean_dev = mean_dev_try;
				memcpy(devcell_try_best, devcell_try,
							sizeof(Basis));
				found = 1;
			}

		}
		}
		}
//		printf("mean_dev = %f\n", mean_dev);
		if ( found ) {
			memcpy(devcell, devcell_try_best, sizeof(Basis));
		} else {
			printf("No further change after %i iterations\n", it);
			break;
		}
		if ( !(it % 1000) ) {
			refine_printf(ctx->dw,
			     "After %5i iterations: mean dev = %13.8f nm^1\n",
			     it, mean_dev/DISPFACTOR);
		}

	}
	free(devcell_try);
	free(devcell_try_best);
	refine_printf(ctx->dw,
			"Final mean dev (%5i iterations) = %13.8f nm^1\n",
			it, mean_dev/DISPFACTOR);
	refine_printf(ctx->dw, "Final cell deviation:\n");
	refine_show_cell(ctx->dw, *devcell);
	refine_printf(ctx->dw, "Confidence:\n");
	refine_show_cell(ctx->dw, conf);

	/* Apply the final values to the cell */
	conf_threshold = 2.5;
	if ( conf.a.x > conf_threshold ) ctx->cell->a.x += devcell->a.x;
	if ( conf.a.y > conf_threshold ) ctx->cell->a.y += devcell->a.y;
	if ( conf.a.z > conf_threshold ) ctx->cell->a.z += devcell->a.z;
	if ( conf.b.x > conf_threshold ) ctx->cell->b.x += devcell->b.x;
	if ( conf.b.y > conf_threshold ) ctx->cell->b.y += devcell->b.y;
	if ( conf.b.z > conf_threshold ) ctx->cell->b.z += devcell->b.z;
	if ( conf.c.x > conf_threshold ) ctx->cell->c.x += devcell->c.x;
	if ( conf.c.y > conf_threshold ) ctx->cell->c.y += devcell->c.y;
	if ( conf.c.z > conf_threshold ) ctx->cell->c.z += devcell->c.z;

	ctx->images->images[ctx->dw->cur_image].rflist = NULL;
	reproject_lattice_changed(ctx);
	displaywindow_update(ctx->dw);

	return mean_dev;

}