/*
 * reproject.c
 *
 * Synthesize diffraction patterns
 *
 * (c) 2007-2008 Thomas White <taw27@cam.ac.uk>
 *
 *  dtr - Diffraction Tomography Reconstruction
 *
 */

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

#include "control.h"
#include "reflections.h"
#include "utils.h"
#include "imagedisplay.h"
#include "displaywindow.h"
#include "image.h"

/* Attempt to find partners from the feature list of "image" for each feature in "flist". */
void reproject_partner_features(ImageFeatureList *rflist, ImageRecord *image) {

	int i;

	for ( i=0; i<rflist->n_features; i++ ) {

		//if ( !reflection_is_easy(rflist->features[i].reflection) ) continue;

		double d = 0.0;
		ImageFeature *partner;
		int idx;

		partner = image_feature_closest(image->features, rflist->features[i].x, rflist->features[i].y, &d, &idx);

		if ( (d <= 20.0) && partner ) {
			rflist->features[i].partner = partner;
			rflist->features[i].partner_d = d;
		} else {
			rflist->features[i].partner = NULL;
		}

	}

}

ImageFeatureList *reproject_get_reflections(ImageRecord *image, ReflectionList *reflectionlist) {

	ImageFeatureList *flist;
	Reflection *reflection;
	double smax = 0.1e9;
	double tilt, omega;
	double nx, ny, nz, nxt, nyt, nzt;	/* "normal" vector (and calculation intermediates) */
	double kx, ky, kz;			/* Electron wavevector ("normal" times 1/lambda */
	double ux, uy, uz, uxt, uyt, uzt;	/* "up" vector (and calculation intermediates) */
	//int done_debug = 0;

	flist = image_feature_list_new();

	tilt = image->tilt;
	omega = image->omega;

	/* Calculate the (normalised) incident electron wavevector
	 *	n is rotated "into" the reconstruction, so only one omega step. */
	nxt = 0.0;  nyt = 0.0;  nzt = 1.0;
	nx = nxt;  ny = cos(tilt)*nyt + sin(tilt)*nzt;  nz = -sin(tilt)*nyt + cos(tilt)*nzt;
	nxt = nx;  nyt = ny;  nzt = nz;
	nx = nxt*cos(-omega) + nyt*sin(-omega);  ny = -nxt*sin(-omega) + nyt*cos(-omega);  nz = nzt;
	kx = nx / image->lambda;
	ky = ny / image->lambda;
	kz = nz / image->lambda;	/* This is the centre of the Ewald sphere */
	//reflection_add(ctx->reflectionlist, kx, ky, kz, 1, REFLECTION_VECTOR_MARKER_1);

	/* Determine where "up" is
	 *	See above. */
	uxt = 0.0;  uyt = 1.0;  uzt = 0.0;
	ux = uxt;  uy = cos(tilt)*uyt + sin(tilt)*uzt;  uz = -sin(tilt)*uyt + cos(tilt)*uzt;
	uxt = ux;  uyt = uy;  uzt = uz;
	ux = uxt*cos(-omega) + uyt*-sin(omega);  uy = -uxt*sin(omega) + uyt*cos(omega);  uz = uzt;
	//reflection_add(ctx->reflectionlist, ux*50e9, uy*50e9, uz*50e9, 1, REFLECTION_VECTOR_MARKER_2);

	reflection = reflectionlist->reflections;
	do {

		double xl, yl, zl;
		double a, b, c;
		double A1, A2, s1, s2, s;

		/* Get the coordinates of the reciprocal lattice point */
		xl = reflection->x;
		yl = reflection->y;
		zl = reflection->z;

		/* Next, solve the relrod equation to calculate the excitation error */
		a = 1.0;
		b = -2.0*(xl*nx + yl*ny + zl*nz);
		c = xl*xl + yl*yl + zl*zl - 1.0/(image->lambda*image->lambda);
		A1 = (-b + sqrt(b*b-4.0*a*c))/(2.0*a);
		A2 = (-b - sqrt(b*b-4.0*a*c))/(2.0*a);
		s1 = 1.0/image->lambda - A1;
		s2 = 1.0/image->lambda - A2;
		if ( fabs(s1) < fabs(s2) ) s = s1; else s = s2;

		/* Skip this reflection if s is large */
		if ( fabs(s) <= smax ) {

			double xi, yi, zi;
			double gx, gy, gz;
			double cx, cy, cz;
			double theta;
			double x, y;
			double rx, ry, rz;

			/* Determine the intersection point */
			xi = xl + s*nx;  yi = yl + s*ny;  zi = zl + s*nz;

			/* Calculate Bragg angle */
			gx = xi - kx;
			gy = yi - ky;
			gz = zi - kz;	/* This is the vector from the centre of the sphere to the intersection */
			theta = angle_between(-kx, -ky, -kz, gx, gy, gz);

			if ( theta > 0.0 ) {

				double psi, disc;

				//reflection_add(ctx->reflectionlist, xl, yl, zl, 1, REFLECTION_GENERATED);
				//reflection_add(ctx->reflectionlist, xi, yi, zi, 1, REFLECTION_MARKER);

				/* Calculate azimuth of point in image (clockwise from "up", will be changed later) */
				cx = yi*nz-zi*ny;  cy = nx*zi-nz*xi;  cz = ny*xi-nx*yi;  /* c = i x n */
				psi = angle_between(cx, cy, cz, ux, uy, uz);

				/* Finally, resolve the +/- Pi ambiguity from the previous step */
				rx = cy*nz-cz*ny;  ry = nx*cz-nz*cx;  rz = ny*cx-nx*cy;  /* r = [i x n] x n */
				disc = angle_between(rx, ry, rz, ux, uy, uz);
			//	if ( (i==20) && !done_debug ) {
			//		reflection_add(ctx->reflectionlist, xi, yi, zi, 1, REFLECTION_VECTOR_MARKER_3);
			//		reflection_add(ctx->reflectionlist, cx, cy, cz, 1, REFLECTION_VECTOR_MARKER_4);
			//		reflection_add(ctx->reflectionlist, rx, ry, rz, 1, REFLECTION_VECTOR_MARKER_4);
			//		printf("psi=%f deg, disc=%f deg\n", rad2deg(psi), rad2deg(disc));
			//	}
				if ( (psi >= M_PI_2) && (disc >= M_PI_2) ) {
					psi -= M_PI_2;		/* Case 1 */
			//		if ( (i==20) && !done_debug ) printf("case 1\n");
				} else if ( (psi >= M_PI_2) && (disc < M_PI_2) ) {
					psi = 3*M_PI_2 - psi;	/* Case 2 */
			//		if ( (i==20) && !done_debug ) printf("case 2\n");
				} else if ( (psi < M_PI_2) && (disc < M_PI_2) ) {
					psi = 3*M_PI_2 - psi;	/* Case 3 */
			//		if ( (i==20) && !done_debug ) printf("case 3\n");
				} else if ( (psi < M_PI_2) && (disc >= M_PI_2) ) {
					psi = 3*M_PI_2 + psi;	/* Case 4 */
			//		if ( (i==20) && !done_debug ) printf("case 4\n");
				}

			//	if ( (i==20) && !done_debug ) printf("final psi=%f clockwise from 'up'\n", rad2deg(psi));
				psi = M_PI_2 - psi;	/* Anticlockwise from "+x" instead of clockwise from "up" */
			//	if ( (i==20) && !done_debug ) printf("final psi=%f anticlockwise from +x\n", rad2deg(psi));

				psi += omega;
			//	if ( (i==20) && !done_debug ) printf("final psi=%f anticlockwise from +tilt axis\n", rad2deg(psi));
			//	if ( (i==20) && !done_debug ) done_debug = 1;

				/* Calculate image coordinates from polar representation */
				if ( image->fmode == FORMULATION_CLEN ) {
					x = image->camera_len*tan(theta)*cos(psi);
					y = image->camera_len*tan(theta)*sin(psi);
					x *= image->resolution;
					y *= image->resolution;
				} else if ( image->fmode == FORMULATION_PIXELSIZE ) {
					x = tan(theta)*cos(psi) / image->lambda;
					y = tan(theta)*sin(psi) / image->lambda;
					x /= image->pixel_size;
					y /= image->pixel_size;
				} else {
					fprintf(stderr, "Unrecognised formulation mode in reproject_get_reflections()\n");
					return NULL;
				}

				x += image->x_centre;
				y += image->y_centre;

				/* Sanity check */
				if ( (x>=0) && (x<image->width) && (y>=0) && (y<image->height) ) {

					/* Record the reflection */
					image_add_feature_reflection(flist, x, y, image, reflection->intensity, reflection);
					/* Intensity should be multiplied by relrod spike function except
					 * reprojected reflections aren't used quantitatively for anything */

					//printf("Reflection at %f, %f\n", x, y);

				} /* else it's outside the picture somewhere */

			} /* else this is the central beam so don't worry about it */

		}

		reflection = reflection->next;

	} while ( reflection );

	/* Partner features only if the image has a feature list.  This allows the test
	 *	program to use this function to generate simulated data. */
	if ( image->features != NULL ) {
		reproject_partner_features(flist, image);
	}

	return flist;

}

/* Ensure ctx->cell_lattice matches ctx->cell */
void reproject_cell_to_lattice(ControlContext *ctx) {

	int i;

	if ( ctx->cell_lattice ) {
		reflection_list_from_new_cell(ctx->cell_lattice, ctx->cell);
	} else {
		ctx->cell_lattice = reflection_list_from_cell(ctx->cell);
	}

	displaywindow_enable_cell_functions(ctx->dw, TRUE);

	/* Invalidate all the reprojected feature lists */
	for ( i=0; i<ctx->images->n_images; i++ ) {
		image_feature_list_free(ctx->images->images[i].rflist);
		ctx->images->images[i].rflist = NULL;
	}

}

/* Notify that ctx->cell has changed (also need to call displaywindow_update) */
void reproject_lattice_changed(ControlContext *ctx) {

	reproject_cell_to_lattice(ctx);
	displaywindow_update_imagestack(ctx->dw);

}