/*
 * basis.c
 *
 * Find approximate lattices to feed various procedures
 *
 * (c) 2007 Thomas White <taw27@cam.ac.uk>
 *
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

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

#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "utils.h"
#include "reflections.h"
#include "basis.h"

static double basis_efom(ReflectionList *reflectionlist, Basis *basis) {

	int n_indexed, n_counted;
	Reflection *cur;
	
	cur = reflectionlist->reflections;
	n_indexed = 0;
	n_counted = 0;
	while ( cur ) {

		if ( cur->type == REFLECTION_NORMAL ) {
			
			/* Can this basis "approximately" account for this reflection? */
			double det;
			double a11, a12, a13, a21, a22, a23, a31, a32, a33;
			double h, k, l;
			
			/* Set up the coordinate transform from hkl to xyz */
			a11 = basis->a.x;  a12 = basis->a.y;  a13 = basis->a.z;
			a21 = basis->b.x;  a22 = basis->b.y;  a23 = basis->b.z;
			a31 = basis->c.x;  a32 = basis->c.y;  a33 = basis->c.z;
			
			/* Invert the matrix to get hkl from xyz */
			det = a11*(a22*a33 - a23*a32) - a12*(a21*a33 - a23*a31) + a13*(a21*a32 - a22*a31);
			h = ((a22*a33-a23*a32)*cur->x + (a23*a31-a21*a33)*cur->y + (a21*a32-a22*a31)*cur->z) / det;
			k = ((a13*a32-a12*a33)*cur->x + (a11*a33-a13*a31)*cur->y + (a12*a31-a11*a32)*cur->z) / det;
			l = ((a12*a23-a13*a22)*cur->x + (a13*a21-a11*a23)*cur->y + (a11*a22-a12*a21)*cur->z) / det;
			
			/* Calculate the deviations in terms of |a|, |b| and |c| */
			h = fabs(h);  k = fabs(k);  l = fabs(l);
			h -= floor(h);  k -= floor(k);  l -= floor(l);
			if ( h == 1.0 ) h = 0.0;
			if ( k == 1.0 ) k = 0.0;
			if ( l == 1.0 ) l = 0.0;
			
			/* Define "approximately" here.  Circle in basis space becomes an ellipsoid in reciprocal space */
			if ( h*h + k*k + l*l <= 0.1*0.1*0.1 ) n_indexed++;
			
			n_counted++;
		
		}
		
		cur = cur->next;

	}
	
	return (double)n_indexed / n_counted;

}

static int basis_lfom(ControlContext *ctx, double vx, double vy, double vz) {
	
	Reflection	*tcentre;
	int		lfom;
	double		tol;
	int		j;
	
	lfom = 0;
	tol = modulus(vx, vy, vz)/10.0;
	tcentre = ctx->reflectionlist->reflections;
	do {
	
		for ( j=-20; j<=20; j++ ) {
		
			Reflection *check;
			
			check = reflection_find_nearest(ctx->reflectionlist, tcentre->x+vx*j, tcentre->y+vy*j, tcentre->z+vz*j);
			if ( check && (distance3d(check->x, check->y, check->z, tcentre->x+vx*j, tcentre->y+vy*j, tcentre->z+vz*j) < tol) ) {
				lfom++;
			}
			
		}
		
		tcentre = tcentre->next;
		
	} while ( tcentre );

	return lfom;

}

/* Select a suitable number of sensible seed vectors */
static ReflectionList *basis_find_seeds(ControlContext *ctx) {

	double		tilt_min;
	double		tilt_max;
	double		tilt_mid;
	ImageRecord	*imagerecord;
	double		x_temp, y_temp, z_temp;
	double		scale;
	double		x, y, z;
	Reflection	*centre;
	int		i;
	ReflectionList *seeds;
	
	seeds = reflectionlist_new();
	
	/* Locate the 'plane' in the middle of the "wedge".
	 *	This whole procedure assumes there is just one tilt axis. */
	tilt_min = control_min_tilt(ctx);
	tilt_max = control_max_tilt(ctx);
	tilt_mid = tilt_min + (tilt_max-tilt_min)/2;
	imagerecord = control_image_nearest_tilt(ctx, tilt_mid);
	
	/* Apply the last two steps of the mapping transform to get the direction from the origin
	 *	towards the middle of the wedge */
	x_temp = 0.0;
	y_temp = cos(deg2rad(imagerecord->tilt));
	z_temp = -sin(deg2rad(imagerecord->tilt));
	x = x_temp*cos(-deg2rad(imagerecord->omega)) + y_temp*sin(-deg2rad(imagerecord->omega));
	y = -x_temp*sin(-deg2rad(imagerecord->omega)) + y_temp*cos(-deg2rad(imagerecord->omega));
	z = z_temp;
	
	/* Find the point in the middle of the "wedge" */
	scale = reflection_largest_g(ctx->reflectionlist)/6;
	x *= scale;
	y *= scale;
	z *= scale;
	reflection_add(ctx->reflectionlist, x, y, z, 1.0, REFLECTION_VECTOR_MARKER_2);
	
	/* Find an "origin" reflection */
	centre = reflection_find_nearest(ctx->reflectionlist, x, y, z);
	if ( !centre ) return NULL;
	centre->found = 1;
	reflection_add(ctx->reflectionlist, centre->x, centre->y, centre->z, 1.0, REFLECTION_GENERATED);
	
	for ( i=1; i<=10; i++ ) {
	
		Reflection *vector;
		Reflection *new_seed;
		int accept, lfom;
		double vx, vy, vz;
		
		do {
			
			Reflection	*check;
			
			accept = 1;
		
			/* Find a "candidate vector" reflection */
			vector = reflection_find_nearest_longer_unknown(ctx->reflectionlist, centre->x, centre->y, centre->z, 1e9); /* 0.5 A^-1 */
			if ( !vector ) {
				printf("BS: Only found %i seeds\n", i);
				return seeds;
			}
			vector->found = 1;
			
			/* Get vector components (not the coordinates the vector was calculated from!) */
			vx = vector->x - centre->x;
			vy = vector->y - centre->y;
			vz = vector->z - centre->z;
			
			/* Proximity test: don't duplicate seeds */
			check = reflection_find_nearest(seeds, vx, vy, vz);
			if ( check ) {
				if ( distance3d(vx, vy, vz, check->x, check->y, check->z) < 1e9 ) {
					/* Too close to another seed */
					accept = 0;
					continue;
				}
			}
			
			/* Record lFOM for later analysis */
			lfom = basis_lfom(ctx, vx, vy, vz);
			
		} while ( !accept );

		/* Add to the list of seeds */
		new_seed = reflection_add(seeds, vx, vy, vz, 1.0, REFLECTION_NORMAL);
		new_seed->lfom = lfom;
		
		/* Create a marker in the default list for visualisation */
		reflection_add(ctx->reflectionlist, vx, vy, vz, 1.0, REFLECTION_MARKER);

	}

	return seeds;

}

/* Assemble the most sensible basis from seeds */
static Basis *basis_assemble_seeds(ReflectionList *seeds) {

	Basis *basis;
	Reflection *ref;
	int i, j, b;
	ReflectionList *seeds_sorted;

	seeds_sorted = reflection_sort_lfom(seeds);
	
	basis = malloc(10*sizeof(Basis));
	
	for ( b=0; b<10; b++ ) {
	
		ref = seeds_sorted->reflections;
		j = 0;	/* Number of basis components already found */
		for ( i=1; i<=seeds_sorted->n_reflections; i++ ) {
		
			double vx, vy, vz;
		
			vx = ref->x;
			vy = ref->y;
			vz = ref->z;
			
			printf("Seed %2i: lFOM=%6i", i, ref->lfom);
			
			switch ( j ) {
				case 0 : {
					basis[b].a.x = vx;
					basis[b].a.y = vy;
					basis[b].a.z = vz;
					j++;
					printf(" *");
					break;
				}
				case 1 : {
					if ( (angle_between(vx, vy, vz, basis[b].a.x, basis[b].a.y, basis[b].a.z) > M_PI/6)
					  && (angle_between(vx, vy, vz, basis[b].a.x, basis[b].a.y, basis[b].a.z) < M_PI-M_PI/6) ) {
						basis[b].b.x = vx;
						basis[b].b.y = vy;
						basis[b].b.z = vz;
						j++;
						printf(" * (%4.1f deg)", rad2deg(angle_between(vx, vy, vz, basis[b].a.x, basis[b].a.y, basis[b].a.z)));
						break;
					}
				}
				case 2 : {
					double cx, cy, cz;
					cx = basis[b].a.y*basis[b].b.z - basis[b].a.z*basis[b].b.y;
					cy =  - basis[b].a.x*basis[b].b.z + basis[b].a.z*basis[b].b.x;
					cz = basis[b].a.x*basis[b].b.y - basis[b].a.y*basis[b].b.x;
					if ( (angle_between(vx, vy, vz, basis[b].a.x, basis[b].a.y, basis[b].a.z) > M_PI/6)
					  && (angle_between(vx, vy, vz, basis[b].b.x, basis[b].b.y, basis[b].b.z) > M_PI/6)
					  && (angle_between(vx, vy, vz, cx, cy, cz) < M_PI/2-M_PI/6) ) {
						basis[b].c.x = vx;
						basis[b].c.y = vy;
						basis[b].c.z = vz;
						j++;
						printf(" * (%4.1f deg)", rad2deg(angle_between(vx, vy, vz, cx, cy, cz)));
						break;
					}
				}
			}
			
			printf("\n");
			if ( j >= 3 ) break;
			ref = ref->next;
			
		}
		if ( j < 3 ) {
			break;
		}
	}
	
	return basis;

}

Basis *basis_find(ControlContext *ctx) {

	Basis *basis;
	ReflectionList *seeds;
	
	/* Get the shortlist of seeds */
	seeds = basis_find_seeds(ctx);
	if ( seeds->n_reflections < 3 ) {
		printf("BS: Not enough seeds to form a basis\n");
		return NULL;
	}
	printf("BS: Found %i seeds\n", seeds->n_reflections);

	basis = basis_assemble_seeds(seeds);
	
	printf("BS: eFOM = %7.3f %%\n", basis_efom(ctx->reflectionlist, basis)*100);
	
	return basis;

}