/*
 * ipr.c
 *
 * Iterative Prediction-Refinement Reconstruction
 *
 * (c) 2007 Thomas White <taw27@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 "structure.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();
	
}

/* 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);
		}

	}

}

/* Choose an initial set of three vectors for the expansion.
 * This would probably be just as easy for the user to do... */
static Basis *ipr_choose_initial_basis(ControlContext *ctx) {

	printf("ipr: c_i_b: starting\n");
	Basis *basis;
	Reflection *reflection;
	int cur_max;
	
	basis = malloc(sizeof(Basis));
	
	/* Determine a vaguely sensible starting 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;
	reflection = ctx->reflectionctx->reflections->next;	/* First reflection, skip initial placeholder */
	cur_max = 1;
	printf("ipr: c_i_b: starting loop\n");
	do {
		double mod;
		unsigned int angle_works;
		//printf("ipr: c_i_b: iteration - reflection (%f,%f,%f)\n",reflection->x,reflection->y,reflection->z);
		mod = modulus(reflection->x, reflection->y, reflection->z);
		
		angle_works = 1;
		if ( basis->a.modulus != 1e30 ) {
			double angle = angle_between(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(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(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) );
	
	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);
	
	reflection_clear(ctx->reflectionctx);
	reflection_add(ctx->reflectionctx, basis->a.x, basis->a.y, basis->a.z, 1, REFLECTION_GENERATED);
	reflection_add(ctx->reflectionctx, basis->b.x, basis->b.y, basis->b.z, 1, REFLECTION_GENERATED);
	reflection_add(ctx->reflectionctx, basis->c.x, basis->c.y, basis->c.z, 1, REFLECTION_GENERATED);
	
	ipr_try_compact(basis);
	
	return basis;

}

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

	ReflectionContext *ordered;
	double max_res;
	int max_order_a, max_order_b, max_order_c;
	signed int h, k, l;
	
	/* NB This assumes the diffraction pattern is "vaguely" centered... */
	if ( ctx->inputfiletype != INPUT_CACHE) {
		if ( ctx->fmode == FORMULATION_PIXELSIZE ) {
			max_res = sqrt(ctx->width*ctx->width + ctx->height*ctx->height) * ctx->pixel_size;
		} else { 
			max_res = sqrt(ctx->width*ctx->width + ctx->height*ctx->height) / (ctx->lambda * ctx->camera_length);
		}
		max_res = max_res / 4;
	} else {
		max_res = 2e10; //works until I put some more in the reflect.cache header
	}
	
	max_order_a = max_res/basis->a.modulus;
	max_order_b = max_res/basis->b.modulus;
	max_order_c = max_res/basis->c.modulus;
	
	ordered = reflection_init();
	
	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++ ) {
				/* IPR-generated reflections are stored as hkl because their
					g-vectors are going to be messed about with later. */
				reflection_add_index(ordered, h, k, l, 1, REFLECTION_GENERATED);
			}
		}
	}
	
	return ordered;

}

/* Sort the list of reflections into order of increasing g */
static void ipr_sort(ReflectionContext *rctx) {

	Reflection *prev;
	Reflection *this;
	Reflection *next;
	
	this = rctx->reflections;
	next = this->next;
	prev = NULL;
	
	while ( next ) {
	
		next = next->next;	/* ! */
		prev = prev->next;
		this = this->next;
		
	}

}

int ipr_reduce(ControlContext *ctx) {
	printf("ipr: starting\n");
	Basis *basis;
	
	printf("ipr: calling choose_initial_basis\n");
	basis = ipr_choose_initial_basis(ctx);
	ctx->reflectionctx->basis = basis;
	if ( !basis ) return -1;
	
	
	/* Get rid of the original list and replace it with the prediction list */
	printf("ipr: calling reflection_clear\n");
	reflection_clear(ctx->reflectionctx);
	free(ctx->reflectionctx);
	printf("ipr: calling ipr_generate\n");
	ctx->reflectionctx = ipr_generate(ctx, basis);
	ctx->reflectionctx->basis = basis;
	
	/* Sort the reflections into order of increasing g */
	printf("ipr: calling ipr_sort\n");
	ipr_sort(ctx->reflectionctx);
	
	return 0;

}