/*
 * refine-cgrad.c
 *
 * Refinement by Conjugate Gradient Minimisation
 *
 * (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_blas.h>
#include <gsl/gsl_multimin.h>
#include <gsl/gsl_linalg.h>

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

static double refine_cgrad_f(const gsl_vector *coordinates, void *params) {

	double r;
	RefinementPair *pair = params;
	AtomicModel *model;
	unsigned int idx;
	unsigned int j;
	ReflectionList *obs;
	ReflectionList *calc;
	
	obs = pair->reflections;	
	model = model_copy(pair->model);
	
	idx = 0;
	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);
	
	if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
		r = stat_r2(obs, calc);	/* stat_r(Fobs, Fcalc) */
	} else {
		r = stat_r(obs, calc);	/* stat_r(Fobs, Fcalc) */
	}
	free(calc);
	model_free(model);
	
	return r;

}

static void refine_cgrad_df(const gsl_vector *coordinates, void *params, gsl_vector *df) {

	RefinementPair *pair;
	AtomicModel *model;
	ReflectionList *calc;
	unsigned int j;
	AtomicModel *model_original;
	unsigned int idx;
	ReflectionList *obs;
	double r_old;
	
	pair = params;
	model = model_copy(pair->model);
	obs = pair->reflections;
	
	idx = 0;
	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;
		}
	}
	
	/* Calculate reflections templated from the original list */
	calc = model_calculate_f(obs, model, 69);
	
	if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
		r_old = stat_r2(obs, calc);	/* stat_r(Fobs, Fcalc) */
	} else {
		r_old = stat_r(obs, calc);	/* stat_r(Fobs, Fcalc) */
	}
	reflist_free(calc);
	
	idx = 0;
	model_original = model_copy(model);

	/* Determine gradients */
	for ( j=0; j<model->n_atoms; j++ ) {
		
		ReflectionList *calc_new;
		
		if ( !(model->atoms[j].refine) ) continue;
		
		if ( pair->spec & REFINE_SPEC_X ) {
		
			double r_new, derivative;
			
			model->atoms[j].x += LSQ_MSLS_SHIFT;
			
			calc_new = model_calculate_f(obs, model, 69);
			
			if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
				r_new = stat_r2(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			} else {
				r_new = stat_r(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			}
			derivative = (r_new - r_old)/LSQ_MSLS_SHIFT;
			gsl_vector_set(df, idx, derivative);

			idx++;
			reflist_free(calc_new);
			model_move(model, model_original);
			
		}
		
		if ( pair->spec & REFINE_SPEC_Y ) {
		
			double r_new, derivative;
			
			model->atoms[j].y += LSQ_MSLS_SHIFT;
			
			calc_new = model_calculate_f(obs, model, 69);
			
			if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
				r_new = stat_r2(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			} else {
				r_new = stat_r(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			}
			derivative = (r_new - r_old)/LSQ_MSLS_SHIFT;
			gsl_vector_set(df, idx, derivative);

			idx++;
			reflist_free(calc_new);
			model_move(model, model_original);
			
		}
		
		if ( pair->spec & REFINE_SPEC_Z ) {
		
			double r_new, derivative;
			
			model->atoms[j].z += LSQ_MSLS_SHIFT;
			
			calc_new = model_calculate_f(obs, model, 69);
			
			if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
				r_new = stat_r2(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			} else {
				r_new = stat_r(obs, calc_new);	/* stat_r(Fobs, Fcalc) */
			}
			derivative = (r_new - r_old)/LSQ_MSLS_SHIFT;
			gsl_vector_set(df, idx, derivative);

			idx++;
			reflist_free(calc_new);
			model_move(model, model_original);
			
		}
		
		if ( pair->spec & REFINE_SPEC_B ) {
		
			double r_new, derivative;
			
			model->atoms[j].B += LSQ_MSLS_SHIFT;
			
			calc_new = model_calculate_f(obs, model, 69);
			
			if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
				r_new = stat_r2(obs, calc);	/* stat_r(Fobs, Fcalc) */
			} else {
				r_new = stat_r(obs, calc);	/* stat_r(Fobs, Fcalc) */
			}
			derivative = (r_new - r_old)/LSQ_MSLS_SHIFT;
			gsl_vector_set(df, idx, derivative);

			idx++;
			reflist_free(calc_new);
			model_move(model, model_original);
			
		}
		
		if ( pair->spec & REFINE_SPEC_OCC ) {
		
			double r_new, derivative;
			
			model->atoms[j].occ += LSQ_MSLS_SHIFT;
			
			calc_new = model_calculate_f(obs, model, 69);
			
			if ( pair->spec & REFINE_SPEC_INTENSITIES ) {
				r_new = stat_r2(obs, calc);	/* stat_r(Fobs, Fcalc) */
			} else {
				r_new = stat_r(obs, calc);	/* stat_r(Fobs, Fcalc) */
			}
			derivative = (r_new - r_old)/LSQ_MSLS_SHIFT;
			gsl_vector_set(df, idx, derivative);

			idx++;
			reflist_free(calc_new);
			model_move(model, model_original);
			
		}
		
	}
	
	model_free(model_original);
	model_free(model);
	
}

static void refine_cgrad_fdf(const gsl_vector *coordinates, void *params, double *f, gsl_vector *df) {

	*f = refine_cgrad_f(coordinates, params);
	refine_cgrad_df(coordinates, params, df);

}

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

	gsl_multimin_fdfminimizer *s;
	gsl_multimin_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;
	
	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 = n_params * n_atoms;
	s = gsl_multimin_fdfminimizer_alloc(gsl_multimin_fdfminimizer_conjugate_fr, f.n);
	f.f = &refine_cgrad_f;	f.df = &refine_cgrad_df;	f.fdf = &refine_cgrad_fdf;
	coordinates = gsl_vector_alloc(f.n);
	
	/* Set initial estimates */
	idx = 0;
	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_multimin_fdfminimizer_set(s, &f, coordinates, 1e-4, 1e-6);
	
	do {
	
		/* Iterate */
		iter_status = gsl_multimin_fdfminimizer_iterate(s);
		printf ("status = %s\n", gsl_strerror(iter_status));
		iter++;

		/* Check for error */
		if ( iter_status ) {
			printf("status after iteration = %s\n", gsl_strerror(iter_status));
			break;
		}
		
		/* Test for convergence */
		conv_status = gsl_multimin_test_gradient(s->gradient, 1e-6);
		printf("status after convergence test = %s\n", gsl_strerror(conv_status));
		
		idx = 0;
		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;
			}
		}
		
		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_multimin_fdfminimizer_free(s);
	gsl_vector_free(coordinates);

}