/*
 * refine-lmder.c
 *
 * Refinement by GSL's Levenberg-Marquardt LSQ
 *
 * (c) 2006-2007 Thomas White <taw27@cam.ac.uk>
 *
 *  synth2d - Two-Dimensional Crystallographic Fourier Synthesis
 *
 */

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

#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_multifit_nlin.h>
#include <gsl/gsl_blas.h>

#include "refine.h"
#include "model.h"
#include "reflist.h"
#include "statistics.h"

static int refine_lmder_f(const gsl_vector *coordinates, void *params, gsl_vector *f) {

	RefinementPair *pair;
	AtomicModel *model;
	ReflectionList *calc;
	unsigned int i, j, idx;
	double scale;
	ReflectionList *obs;
	
	pair = params;
	model = model_copy(pair->model);
	obs = pair->reflections;
	
	idx = 0;
	scale = gsl_vector_get(coordinates, idx++);
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) {
			if ( pair->spec & REFINE_SPEC_X ) model->atoms[j].x = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].x < 0 ) model->atoms[j].x = 1 + model->atoms[j].x;
			if ( pair->spec & REFINE_SPEC_Y ) model->atoms[j].y = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].y < 0 ) model->atoms[j].y = 1 + model->atoms[j].y;
			if ( pair->spec & REFINE_SPEC_Z ) model->atoms[j].z = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].z < 0 ) model->atoms[j].z = 1 + model->atoms[j].z;
			if ( pair->spec & REFINE_SPEC_B ) model->atoms[j].B = gsl_vector_get(coordinates, idx++);
			if ( model->atoms[j].B < 0 ) model->atoms[j].B = 0;
			if ( pair->spec & REFINE_SPEC_OCC ) model->atoms[j].occ = gsl_vector_get(coordinates, idx++);
			if ( model->atoms[j].occ < 0 ) model->atoms[j].occ = 0;
		}
	}
	
	/* Calculate reflections templated from the original list */
	calc = model_calculate_f(obs, model, 69);
	
	/* Calculate residuals */
	for ( i=1; i<calc->n_reflections; i++ ) {
		double am_calc, residual;
		am_calc = calc->refs[i].amplitude;
		if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
			/* ( k|Icalc| - |Iobs| ) / error(|Iobs|) */
			residual = (scale*am_calc*am_calc - obs->refs[i].amplitude*obs->refs[i].amplitude)/(sqrt(2)*obs->refs[i].am_error);
		} else {
			residual = (scale*am_calc - obs->refs[i].amplitude)/obs->refs[i].am_error; /* ( k|Fcalc| - |Fobs| ) / error(|Fobs|) */
		}
		gsl_vector_set(f, i-1, residual);
		// printf("Residual for %3i %3i %3i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, residual);
	}
	
	model_free(model);
	free(calc);
	
	return GSL_SUCCESS;
  
}

static int refine_lmder_df(const gsl_vector *coordinates, void *params, gsl_matrix *J) {

	RefinementPair *pair;
	AtomicModel *model_original;
	AtomicModel *model;
	ReflectionList *obs;
	ReflectionList *calc;
	unsigned int i, j, idx;
	double scale;
	
	pair = params;
	model = model_copy(pair->model);
	obs = pair->reflections;
	
	idx = 0;
	scale = gsl_vector_get(coordinates, idx++);
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) {
			if ( pair->spec & REFINE_SPEC_X ) model->atoms[j].x = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].x < 0 ) model->atoms[j].x = 1 + model->atoms[j].x;
			if ( pair->spec & REFINE_SPEC_Y ) model->atoms[j].y = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].y < 0 ) model->atoms[j].y = 1 + model->atoms[j].y;
			if ( pair->spec & REFINE_SPEC_Z ) model->atoms[j].z = fmod(gsl_vector_get(coordinates, idx++), 1);
			if ( model->atoms[j].z < 0 ) model->atoms[j].z = 1 + model->atoms[j].z;
			if ( pair->spec & REFINE_SPEC_B ) model->atoms[j].B = gsl_vector_get(coordinates, idx++);
			if ( model->atoms[j].B < 0 ) model->atoms[j].B = 0;
			if ( pair->spec & REFINE_SPEC_OCC ) model->atoms[j].occ = gsl_vector_get(coordinates, idx++);
			if ( model->atoms[j].occ < 0 ) model->atoms[j].occ = 0;
		}
	}
	
	/* Calculate reflections templated from the original list */
	calc = model_calculate_f(obs, model, 69);
	
	idx = 0;
	
	model_original = model_copy(model);

	/* Gradient of the scale factor */
	for ( i=1; i<calc->n_reflections; i++ ) {
		double am;
		am = calc->refs[i].amplitude;
		if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
			gsl_matrix_set(J, i-1, idx, am*am);
		} else {
			gsl_matrix_set(J, i-1, idx, am);
		}
	}
	idx++;

	/* Determine gradients of atomic coordinates and B-factors */
	for ( j=0; j<model->n_atoms; j++ ) {
		
		if ( !(model->atoms[j].refine) ) continue;
		double x, y, z;
		ReflectionList *reflections_new;
		
		x = model->atoms[j].x;
		y = model->atoms[j].y;
		z = model->atoms[j].z;
		
		if ( pair->spec & REFINE_SPEC_X ) {
			model->atoms[j].x += LSQ_MSLS_SHIFT;
			reflections_new = model_calculate_f(obs, model, 69);
			for ( i=1; i<obs->n_reflections; i++ ) {
			
				double derivative, am, residual_old, residual_new;
				
				am = calc->refs[i].amplitude;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_old = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_old = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				am = reflections_new->refs[i].amplitude;
				residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_new = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				
				derivative = (residual_new - residual_old)/LSQ_MSLS_SHIFT;
				gsl_matrix_set(J, i-1, idx, derivative);
				// printf("J for %3i %3i %3i reflection wrt x%i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, j, derivative);

			}
			idx++;
			reflist_free(reflections_new);
			model_move(model, model_original);
		}
		
		if ( pair->spec & REFINE_SPEC_Y ) {
			model->atoms[j].y += LSQ_MSLS_SHIFT;
			reflections_new = model_calculate_f(obs, model, 69);
			for ( i=1; i<obs->n_reflections; i++ ) {
			
				double derivative, am, residual_old, residual_new;
				
				am = calc->refs[i].amplitude;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_old = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_old = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				am = reflections_new->refs[i].amplitude;
				residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_new = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				
				derivative = (residual_new - residual_old)/LSQ_MSLS_SHIFT;
				gsl_matrix_set(J, i-1, idx, derivative);
				// printf("J for %3i %3i %3i reflection wrt y%i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, j, derivative);
				
			}
			idx++;
			reflist_free(reflections_new);
			model_move(model, model_original);
		}

		if ( pair->spec & REFINE_SPEC_Z ) {
			model->atoms[j].z += LSQ_MSLS_SHIFT;
			reflections_new = model_calculate_f(obs, model, 69);
			for ( i=1; i<obs->n_reflections; i++ ) {
			
				double derivative, am, residual_old, residual_new;
				
				am = calc->refs[i].amplitude;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_old = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_old = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				am = reflections_new->refs[i].amplitude;
				residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_new = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				
				derivative = (residual_new - residual_old)/LSQ_MSLS_SHIFT;
				gsl_matrix_set(J, i-1, idx, derivative);
				// printf("J for %3i %3i %3i reflection wrt z%i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, j, derivative);

			}
			idx++;
			reflist_free(reflections_new);
			model_move(model, model_original);
		}
		
		if ( pair->spec & REFINE_SPEC_B ) {
			model->atoms[j].B += LSQ_MSLS_SHIFT;
			reflections_new = model_calculate_f(obs, model, 69);
			for ( i=1; i<obs->n_reflections; i++ ) {
			
				double derivative, am, residual_old, residual_new;
				
				am = calc->refs[i].amplitude;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_old = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_old = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				am = reflections_new->refs[i].amplitude;
				residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_new = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				
				derivative = (residual_new - residual_old)/LSQ_MSLS_SHIFT;
				gsl_matrix_set(J, i-1, idx, derivative);
				// printf("J for %3i %3i %3i reflection wrt B%i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, j, derivative);

			}
			idx++;
			reflist_free(reflections_new);
			model_move(model, model_original);
		}
		
		if ( pair->spec & REFINE_SPEC_OCC ) {
			model->atoms[j].occ += LSQ_MSLS_SHIFT;
			reflections_new = model_calculate_f(obs, model, 69);
			for ( i=1; i<obs->n_reflections; i++ ) {
			
				double derivative, am, residual_old, residual_new;
				
				am = calc->refs[i].amplitude;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_old = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_old = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				am = reflections_new->refs[i].amplitude;
				residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
					residual_new = (scale*am*am - (obs->refs[i].amplitude*obs->refs[i].amplitude))/(sqrt(2)*obs->refs[i].am_error);
				} else {
					residual_new = (scale*am - obs->refs[i].amplitude)/obs->refs[i].am_error;
				}
				
				derivative = (residual_new - residual_old)/LSQ_MSLS_SHIFT;
				gsl_matrix_set(J, i-1, idx, derivative);
				// printf("J for %3i %3i %3i reflection wrt B%i = %f\n", obs->refs[i].h, obs->refs[i].k, obs->refs[i].l, j, derivative);

			}
			idx++;
			reflist_free(reflections_new);
			model_move(model, model_original);
		}
		
	}
	
	model_free(model_original);
	model_free(model);
	free(calc);
	
	return GSL_SUCCESS;
  
}

static int refine_lmder_fdf(const gsl_vector *coordinates, void *params, gsl_vector *f, gsl_matrix *J) {

	refine_lmder_f(coordinates, params, f);
	refine_lmder_df(coordinates, params, J);
	
	return GSL_SUCCESS;
	
}

void refine_lmder(AtomicModel *model, ReflectionList *reflections, RefinementSpec spec) {

	gsl_multifit_fdfsolver *s;
	gsl_multifit_function_fdf f;
	unsigned int iter = 0;
	gsl_vector *coordinates;
	unsigned int j;
	RefinementPair pair;
	int iter_status, conv_status;
	unsigned int n_params, n_atoms, idx;
	double scale;
	ReflectionList *calc;
	
	n_params = 0;
	if ( spec & REFINE_SPEC_X ) n_params++;
	if ( spec & REFINE_SPEC_Y ) n_params++;
	if ( spec & REFINE_SPEC_Z ) n_params++;
	if ( spec & REFINE_SPEC_B ) n_params++;
	if ( spec & REFINE_SPEC_OCC ) n_params++;
	
	/* Prevent origin from floating */
	if ( (spec & REFINE_SPEC_X) || (spec & REFINE_SPEC_Y) || (spec & REFINE_SPEC_Z) ) {
		model->atoms[0].refine = 0;
	}
	
	n_atoms = 0;
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) n_atoms++;
	}
	
	f.n = reflections->n_reflections-1;	f.p = 1 + n_params * n_atoms;
	s = gsl_multifit_fdfsolver_alloc(gsl_multifit_fdfsolver_lmsder, f.n, f.p);
	f.f = &refine_lmder_f;	f.df = &refine_lmder_df;	f.fdf = &refine_lmder_fdf;
	printf("Refining %i parameters against %i reflections\n", f.p, f.n);
	coordinates = gsl_vector_alloc(f.p);
	
	/* Set initial estimates */
	idx = 0;
	calc = model_calculate_f(reflections, NULL, 69);
	if ( spec & REFINE_SPEC_INTENSITIES ) {
		scale = stat_scale_intensity(reflections, calc);
	} else {
		scale = stat_scale(reflections, calc);
	}
	reflist_free(calc);
	gsl_vector_set(coordinates, idx++, scale);
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) {
			if ( spec & REFINE_SPEC_X ) gsl_vector_set(coordinates, idx++, model->atoms[j].x);
			if ( spec & REFINE_SPEC_Y ) gsl_vector_set(coordinates, idx++, model->atoms[j].y);
			if ( spec & REFINE_SPEC_Z ) gsl_vector_set(coordinates, idx++, model->atoms[j].z);
			if ( spec & REFINE_SPEC_B ) gsl_vector_set(coordinates, idx++, model->atoms[j].B);
			if ( spec & REFINE_SPEC_OCC ) gsl_vector_set(coordinates, idx++, model->atoms[j].occ);
		}
	}
	
	pair.model = model;  pair.reflections = reflections;  pair.spec = spec;
	f.params = &pair;
	
	gsl_multifit_fdfsolver_set(s, &f, coordinates);
	
	printf("initial: scale=%f, |f(x)|=%g\n", gsl_vector_get(s->x, 0), gsl_blas_dnrm2(s->f));
	
	do {
	
		ReflectionList *model_reflections;
	
		/* Iterate */
		printf("before iteration %i: scale=%f, |f(x)|=%g\n", iter, gsl_vector_get(s->x, 0), gsl_blas_dnrm2(s->f));
		iter_status = gsl_multifit_fdfsolver_iterate(s);
		printf("after iteration %i: scale=%f, |f(x)|=%g, status=%s, ", iter, gsl_vector_get(s->x, 0), gsl_blas_dnrm2(s->f), gsl_strerror(iter_status));
		iter++;

		/* Check for error */
		if ( iter_status ) {
			printf("\nError in LSQ refinement");
			break;
		}
		
		/* Test for convergence */
		conv_status = gsl_multifit_test_delta(s->dx, s->x, 1e-6, 1e-6);
		//printf("status after convergence test = %s\n", gsl_strerror(conv_status));
		
		idx = 0;
		scale = gsl_vector_get(s->x, idx++);
		for ( j=0; j<model->n_atoms; j++ ) {
			if ( model->atoms[j].refine ) {
				if ( spec & REFINE_SPEC_X ) model->atoms[j].x = fmod(gsl_vector_get(s->x, idx++), 1);
				if ( model->atoms[j].x < 0 ) model->atoms[j].x = 1 + model->atoms[j].x;
				if ( spec & REFINE_SPEC_Y ) model->atoms[j].y = fmod(gsl_vector_get(s->x, idx++), 1);
				if ( model->atoms[j].y < 0 ) model->atoms[j].y = 1 + model->atoms[j].y;
				if ( spec & REFINE_SPEC_Z ) model->atoms[j].z = fmod(gsl_vector_get(s->x, idx++), 1);
				if ( model->atoms[j].z < 0 ) model->atoms[j].z = 1 + model->atoms[j].z;
				if ( spec & REFINE_SPEC_B ) model->atoms[j].B = gsl_vector_get(s->x, idx++);
				if ( model->atoms[j].B < 0 ) model->atoms[j].B = 0;
				if ( spec & REFINE_SPEC_OCC ) model->atoms[j].occ = gsl_vector_get(s->x, idx++);
				if ( model->atoms[j].occ < 0 ) model->atoms[j].occ = 0;
			}
		}
		
		model_reflections = model_calculate_f(reflections, model, 69);
		if ( spec & REFINE_SPEC_INTENSITIES ) {
			printf("R2=%.2f%%\n\n", stat_r2(reflections, model_reflections)*100);  /* stat_r(Fobs, Fcalc) */
		} else {
			printf("R1=%.2f%%\n\n", stat_r(reflections, model_reflections)*100);  /* stat_r(Fobs, Fcalc) */
		}
		free(model_reflections);
		
		refine_schedule_update(model);
		
	} while ( (conv_status == GSL_CONTINUE) && (iter < MAX_REFINEMENT_ITERATIONS) && (model->refine_window->run_semaphore) );
	
	printf("%i iterations performed\n", iter);
	if ( iter == MAX_REFINEMENT_ITERATIONS ) printf("Reached maximum allowed number of iterations");

	gsl_multifit_fdfsolver_free(s);
	gsl_vector_free(coordinates);

}