/*
 * 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 "imagedisplay.h"
#include "reproject.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) {
	
	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;
	
	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) );
	
	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);
	
	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;
	signed int h, k, l;
	
	/* 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();
//	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;
	for ( h=-30; h<=30; h++ ) {
		for ( k=-30; k<=30; k++ ) {
			for ( l=-30; l<=30; 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 int ipr_random_image(ControlContext *ctx) {
	double n;
	n = (double)random()/RAND_MAX;
	n *= ctx->n_images;
	return (int)n;
}

int ipr_refine(ControlContext *ctx) {
	
	Basis *basis;
	ReflectionContext *lat;
	int finished;
	basis = ipr_choose_initial_basis(ctx);
	ctx->reflectionctx->basis = basis;
	if ( !basis ) return -1;
	
	lat = ipr_generate(ctx, basis);
	ctx->reflectionctx->basis = basis;
	
	finished = 0;
	do {
	
		int i;
		size_t n, j;
		ImageDisplay *cur;
		ImageReflection *refl;
		
		/* Select an image */
		i = ipr_random_image(ctx);
		cur = imagedisplay_open(ctx->images[i], "Current Image", IMAGEDISPLAY_SHOW_CENTRE | IMAGEDISPLAY_SHOW_TILT_AXIS);
		
		refl = reproject_get_reflections(ctx->images[i], &n, lat, ctx);
		for ( j=0; j<n; j++ ) {
			imagedisplay_mark_circle(cur, refl[j].x, refl[j].y);
		}
		
		finished = 1;
		
	} while ( !finished );
	reflection_free(lat);
	
	return 0;

}