/*
 * main.c
 *
 * The Top Level Source File
 *
 * (c) 2009 Thomas White <taw27@cam.ac.uk>
 *
 * Triclinator - solve nasty triclinic unit cells
 *
 */


#include <stdlib.h>
#include <stdio.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multifit_nlin.h>
#include <math.h>

#include "crystal.h"
#include "util.h"

#define DSTEP 0.0001
#define ASTEP 0.00001

struct data {
	size_t n;	/* Number of measurements */
	MVal *vals;	/* The measurements themselves */
};


int expb_f(const gsl_vector *x, void *data, gsl_vector *f)
{
	size_t n = ((struct data *)data)->n;
	MVal *vals = ((struct data *) data)->vals;
	Cell cell;
	size_t i;

	cell.a = gsl_vector_get(x, 0);
	cell.b = gsl_vector_get(x, 1);
	cell.c = gsl_vector_get(x, 2);
	cell.al = gsl_vector_get(x, 3);
	cell.be = gsl_vector_get(x, 4);
	cell.ga = gsl_vector_get(x, 5);

	for ( i=0; i<n; i++ ) {
		double calc = crystal_calc(vals[i], cell);
		gsl_vector_set(f, i, (calc - vals[i].meas)/vals[i].esd);
	}

	return GSL_SUCCESS;
}


int expb_df(const gsl_vector *x, void *data, gsl_matrix *J)
{
	size_t n = ((struct data *)data)->n;
	MVal *vals = ((struct data *) data)->vals;
	Cell cell;
	size_t i;

	cell.a = gsl_vector_get(x, 0);
	cell.b = gsl_vector_get(x, 1);
	cell.c = gsl_vector_get(x, 2);
	cell.al = gsl_vector_get(x, 3);
	cell.be = gsl_vector_get(x, 4);
	cell.ga = gsl_vector_get(x, 5);

	for ( i=0; i<n; i++ ) {

		double calc, calc2, s;
		Cell cell2;
		MVal val;

		val = vals[i];

		calc = crystal_calc(vals[i], cell);
		s = vals[i].esd;

		cell2 = cell; cell2.a += DSTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 0, (calc2 - calc)/(DSTEP*s)); /* df/da */

		cell2 = cell; cell2.b += DSTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 1, (calc2 - calc)/(DSTEP*s)); /* df/db */

		cell2 = cell; cell2.c += DSTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 2, (calc2 - calc)/(DSTEP*s)); /* df/dc */

		cell2 = cell; cell2.al += ASTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 3, (calc2 - calc)/(ASTEP*s)); /* df/dal */

		cell2 = cell; cell2.be += ASTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 4, (calc2 - calc)/(ASTEP*s)); /* df/dbe */

		cell2 = cell; cell2.ga += ASTEP; calc2 = crystal_calc(val, cell2);
		gsl_matrix_set(J, i, 5, (calc2 - calc)/(ASTEP*s)); /* df/dga */

	}

	return GSL_SUCCESS;
}


int expb_fdf(const gsl_vector *x, void *data, gsl_vector *f, gsl_matrix *J)
{
	expb_f(x, data, f);
	expb_df(x, data, J);
	return GSL_SUCCESS;
}


void print_state (size_t iter, gsl_multifit_fdfsolver *s)
{
	printf ("%3u %8.5f %8.5f %8.5f %8.5f %8.5f %8.5f", iter,
		gsl_vector_get(s->x, 0),
		gsl_vector_get(s->x, 1),
		gsl_vector_get(s->x, 2),
		gsl_vector_get(s->x, 3),
		gsl_vector_get(s->x, 4),
		gsl_vector_get(s->x, 5));

}


int main(void)
{
	const gsl_multifit_fdfsolver_type *T;
	gsl_multifit_fdfsolver *s;
	int status;
	unsigned int i, iter, done;
	size_t n;
	gsl_matrix *covar;
	MVal *vals;
	struct data d;
	gsl_multifit_function_fdf f;
	double x_init[6];
	gsl_vector_view x = gsl_vector_view_array(x_init, 6);
	Cell cell;

	/* Get the data to be fitted */
	done = 0; n = 0;
	vals = NULL;
	while ( !done ) {

		signed int h, k, l;
		char buf[64];
		int dspacing = 0;

//		printf("Enter the indicies of the first set of planes, "
//			" in the format 'h k l'.\n");
//		if ( n != 0 ) printf("If you have no further planes, "
//					"just hit enter.\n");
		printf("        First plane indicies: ");
		if ( fgets(buf, 63, stdin) != buf ) {
			fprintf(stderr, "Error reading from input\n");
		}
		if ( sscanf(buf, "%i %i %i", &h, &k, &l) != 3 ) {
			if ( buf[0] == '\n' ) {
				done = 1;
				continue;
			} else {
				printf("Invalid input, try again.\n");
			}
		}

		/* Allocate space for the new reading,
			now that we know it exists */
		vals = realloc(vals, (n+1)*sizeof(MVal));
		vals[n].h1 = h;
		vals[n].k1 = k;
		vals[n].l1 = l;

//		printf("Enter the indicies of the second set of planes in the "
//			"same format.\n");
//		if ( n != 0 ) printf("If you are trying to give a d-spacing "
//					"only, just hit enter here.\n");
		printf("       Second plane indicies: ");
		if ( fgets(buf, 63, stdin) != buf ) {
			fprintf(stderr, "Error reading from input\n");
		}
		if ( sscanf(buf, "%i %i %i", &h, &k, &l) != 3 ) {
			if ( buf[0] == '\n' ) {
				dspacing = 1;
			} else {
				printf("Invalid input, try again.\n");
			}
		}

		if ( dspacing ) {
			vals[n].h2 = 0;
			vals[n].k2 = 0;
			vals[n].l2 = 0;
			vals[n].meas = read_value("                   D-spacing: ");
		} else {
			vals[n].h2 = h;
			vals[n].k2 = k;
			vals[n].l2 = l;
			vals[n].meas = DEG2RAD(read_value("        Angle between planes: "));
		}

		vals[n].esd = read_value("Estimated standard deviation: ");

		printf("--------------------------------------\n");
		n++;

	}

	x_init[0] = read_value("    Initial guess for a: ");
	x_init[1] = read_value("    Initial guess for b: ");
	x_init[2] = read_value("    Initial guess for c: ");
	x_init[3] = read_value("Initial guess for alpha: ");
	x_init[4] = read_value(" Initial guess for beta: ");
	x_init[5] = read_value("Initial guess for gamma: ");

	printf("-----------------------------------------\n\n");

	printf("You gave me these measurements:\n");
	printf("--------------------------------------------------------------\n");
	printf(" h1  k1  l1   h2  k2  l2    Spacing   Angle          ESD\n");
	printf("--------------------------------------------------------------\n");
	for ( i=0; i<n; i++ ) {
		printf("%3i %3i %3i  ", vals[i].h1, vals[i].k1, vals[i].l1);
		if ( is_dspacing(vals[i]) ) {
			printf("  -   -   -   %8.5f    -     ", vals[i].meas);
		} else {
			printf("%3i %3i %3i        -    %8.5f", vals[i].h2,
				vals[i].k2, vals[i].l2, RAD2DEG(vals[i].meas));
		}
		printf(" +/- %8.5f", vals[i].esd);
		if ( !is_dspacing(vals[i]) ) {
			printf(" deg\n");
		} else {
			printf(" nm\n");
		}
	}
	printf("--------------------------------------------------------------\n");

	d.vals = vals;
	d.n = n;

	f.f = &expb_f;
	f.df = &expb_df;
	f.fdf = &expb_fdf;
	f.n = n;
	f.p = 6;
	f.params = &d;

	T = gsl_multifit_fdfsolver_lmsder;
	s = gsl_multifit_fdfsolver_alloc(T, n, 6);
	gsl_multifit_fdfsolver_set(s, &f, &x.vector);

	iter = 0;
	printf("\nIterating...\n");
	printf("---------------------------------------------------------\n");
	printf("Itn    a        b        c      alpha    beta    gamma\n");
	printf("---------------------------------------------------------\n");
	print_state(iter, s);

	do {

		iter++;
		status = gsl_multifit_fdfsolver_iterate(s);

		printf(" %s\n", gsl_strerror(status));
		print_state(iter, s);
		if ( status ) break;

		status = gsl_multifit_test_delta(s->dx, s->x, 1e-6, 1e-6);

	} while ( (status == GSL_CONTINUE) && (iter < 500) );
	printf("\n---------------------------------------------------------\n");
	printf ("Final status = %s\n", gsl_strerror (status));

	covar = gsl_matrix_alloc(6, 6);
	gsl_multifit_covar(s->J, 0.0, covar);

	#define FIT(i) gsl_vector_get(s->x, i)
	#define ERR(i) sqrt(gsl_matrix_get(covar,i,i))

	{
		double chi = gsl_blas_dnrm2(s->f);
		double dof = n - 6;
		double c = GSL_MAX_DBL(1, chi / sqrt(dof));

		printf("chisq/dof = %g / %g = %g\n",  pow(chi, 2.0), dof,
							pow(chi, 2.0) / dof);

		printf("\nI think the unit cell is:\n\n");
		printf("\t\t    a = %8.5f +/- %8.5f\n", FIT(0), c*ERR(0));
		printf("\t\t    b = %8.5f +/- %8.5f\n", FIT(1), c*ERR(1));
		printf("\t\t    c = %8.5f +/- %8.5f\n", FIT(2), c*ERR(2));
		printf("\t\talpha = %8.5f +/- %8.5f\n", FIT(3), c*ERR(3));
		printf("\t\t beta = %8.5f +/- %8.5f\n", FIT(4), c*ERR(4));
		printf("\t\tgamma = %8.5f +/- %8.5f\n\n", FIT(5), c*ERR(5));
	}

	gsl_multifit_fdfsolver_free(s);
	gsl_matrix_free(covar);

	cell.a = FIT(0);
	cell.b = FIT(1);
	cell.c = FIT(2);
	cell.al = FIT(3);
	cell.be = FIT(4);
	cell.ga = FIT(5);

	printf("\nYour measurements when calculated with my cell are:\n");
	printf("--------------------------------------------------------------------\n");
	printf(" h1  k1  l1   h2  k2  l2    Spacing   Angle   Deviation\n");
	printf("--------------------------------------------------------------------\n");
	for ( i=0; i<n; i++ ) {
		double dev, dr;
		printf("%3i %3i %3i  ", vals[i].h1, vals[i].k1, vals[i].l1);
		if ( is_dspacing(vals[i]) ) {
			double val;
			val = crystal_calc(vals[i], cell);
			printf("  -   -   -   %8.5f    -     ", val);
			dev = val - vals[i].meas;
			dr = dev / vals[i].meas;
		} else {
			double val;
			val = RAD2DEG(crystal_calc(vals[i], cell));
			printf("%3i %3i %3i        -    %8.5f", vals[i].h2,
				vals[i].k2, vals[i].l2, val);
			dev = val - RAD2DEG(vals[i].meas);
			dr = dev / RAD2DEG(vals[i].meas);
		}
		printf(" %+8.5f (%5.3f%%)", dev, fabs(dr*100.0));
		if ( !is_dspacing(vals[i]) ) {
			printf(" deg\n");
		} else {
			printf(" nm\n");
		}
	}
	printf("--------------------------------------------------------------------\n");


	return 0;
}