/*
 * refine-lsq.c
 *
 * Refinement by Simple 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_blas.h>
#include <gsl/gsl_linalg.h>
#include <stdio.h>
#include <stdlib.h>

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

static AtomicModel *refine_model_from_parameters(const gsl_vector *parameters, RefinementSpec spec, AtomicModel *template) {

	unsigned int j, idx;
	AtomicModel *model;
	
	model = model_copy(template);

	idx = 1; /* Ignore scale factor */
	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(parameters, 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(parameters, 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(parameters, 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(parameters, idx++);
			if ( model->atoms[j].B < 0 ) model->atoms[j].B = 0;
			if ( spec & REFINE_SPEC_OCC ) model->atoms[j].occ = gsl_vector_get(parameters, idx++);
			if ( model->atoms[j].occ < 0 ) model->atoms[j].occ = 0;
		}
	}
	
	return model;

}

static double refine_lsq_gradient(const gsl_vector *parameters, const ReflectionList *obs, RefinementSpec spec,
					size_t param_idx, size_t obs_idx, ReflectionList *f_calc, AtomicModel *model) {

	unsigned int j, idx;
	double f_calc_new, gradient;
	AtomicModel *model_new;
	double scale;
	
	if ( spec & REFINE_SPEC_INTENSITIES ) {
		if ( param_idx == 0 ) return f_calc->refs[obs_idx].amplitude*f_calc->refs[obs_idx].amplitude;	/* Scale factor */
	} else {
		if ( param_idx == 0 ) return f_calc->refs[obs_idx].amplitude;	/* Scale factor */
	}
	
	scale = gsl_vector_get(parameters, 0);
	idx = 0;
	model_new = model_copy(model);
	for ( j=0; j<model->n_atoms; j++ ) {
		
		if ( !(model->atoms[j].refine) ) continue;
		
		if ( spec & REFINE_SPEC_X ) idx++;
		if ( idx == param_idx ) {
			model->atoms[j].x += LSQ_MSLS_SHIFT;
			break;
		}
		
		if ( spec & REFINE_SPEC_Y ) idx++;
		if ( idx == param_idx ) {
			model->atoms[j].y += LSQ_MSLS_SHIFT;
			break;
		}
		
		if ( spec & REFINE_SPEC_Z ) idx++;
		if ( idx == param_idx ) {
			model->atoms[j].z += LSQ_MSLS_SHIFT;
			break;
		}
		
		if ( spec & REFINE_SPEC_B ) idx++;
		if ( idx == param_idx ) {
			model->atoms[j].B += LSQ_MSLS_SHIFT;
			break;
		}
	
		if ( spec & REFINE_SPEC_OCC ) idx++;
		if ( idx == param_idx ) {
			model->atoms[j].occ += LSQ_MSLS_SHIFT;
			break;
		}
	
	}
	
	f_calc_new = model_mod_f(model, f_calc->refs[obs_idx].h, f_calc->refs[obs_idx].k, f_calc->refs[obs_idx].l);
	
	if ( spec & REFINE_SPEC_INTENSITIES ) {
		gradient = ((f_calc->refs[obs_idx].amplitude*f_calc->refs[obs_idx].amplitude) - (f_calc_new*f_calc_new))/LSQ_MSLS_SHIFT;
	} else {
		gradient = (f_calc->refs[obs_idx].amplitude - f_calc_new)/LSQ_MSLS_SHIFT;
	}
	model_free(model_new);
	
	return scale*gradient;

}

static gsl_matrix *refine_lsq_B(const gsl_vector *parameters, const ReflectionList *obs, RefinementSpec spec, AtomicModel *model, ReflectionList *f_calc) {

	gsl_matrix *B;
	size_t j;
	int n_params, n_obs;
	
	n_params = parameters->size;
	n_obs = obs->n_reflections - 1;

	B = gsl_matrix_alloc(n_params, n_params);
	for ( j=0; j<n_params; j++ ) {
		size_t k;
		for ( k=0; k<n_params; k++ ) {
			int i;
			double total_grad = 0;
			for ( i=1; i<=n_obs; i++ ) {
				double grad1, grad2;
				grad1 = refine_lsq_gradient(parameters, obs, spec, j, i, f_calc, model);
				grad2 = refine_lsq_gradient(parameters, obs, spec, k, i, f_calc, model);
				total_grad += grad1 * grad2;
			}
			gsl_matrix_set(B, j, k, total_grad);
		}
	}
	
	return B;
	
}

static gsl_matrix *refine_lsq_D(const gsl_vector *parameters, ReflectionList *obs, RefinementSpec spec, AtomicModel *model, ReflectionList *f_calc) {

	gsl_matrix *D;
	size_t j;
	int n_params, n_obs;
	double scale;
	
	n_params = parameters->size;
	n_obs = obs->n_reflections - 1;
	scale = gsl_vector_get(parameters, 0);
	
	D = gsl_matrix_alloc(n_params, 1);
	for ( j=0; j<n_params; j++ ) {
		int i;
		double total = 0;
		for ( i=1; i<=n_obs; i++ ) {
		
			double grad;
			double res;
			
			grad = refine_lsq_gradient(parameters, obs, spec, j, i, f_calc, model);
			
			if ( spec & REFINE_SPEC_INTENSITIES ) {
				res = (obs->refs[i].amplitude*obs->refs[i].amplitude) - scale*(f_calc->refs[i].amplitude*f_calc->refs[i].amplitude);
			} else {
				res = obs->refs[i].amplitude - scale*f_calc->refs[i].amplitude;
			}
			
			total += grad*res;
			
		}
		gsl_matrix_set(D, j, 0, total);
	
	}
	
	return D;
	
}

/* Determine and apply shift */
static void refine_lsq_iterate(gsl_vector *parameters, ReflectionList *obs, RefinementSpec spec, AtomicModel *model_template) {

	gsl_matrix *B;
	gsl_matrix *D;
	gsl_matrix *Binv;
	gsl_matrix *shift;
	gsl_permutation *perm;
	int s, n_params;
	size_t i;
	AtomicModel *model;
	ReflectionList *f_calc;
	
	n_params = parameters->size;
	model = refine_model_from_parameters(parameters, spec, model_template);
	f_calc = model_calculate_f(obs, model, 69);
	
	/* Form 'B' */
	//printf("LSQ: Forming B\n");
	B = refine_lsq_B(parameters, obs, spec, model, f_calc);
	
	/* Form 'D' */
	//printf("LSQ: Forming D\n");
	D = refine_lsq_D(parameters, obs, spec, model, f_calc);
	
	/* Invert 'B' */
	//printf("LSQ: Inverting B\n");
	perm = gsl_permutation_alloc(B->size1);
	Binv = gsl_matrix_alloc(B->size1, B->size2);
	gsl_linalg_LU_decomp(B, perm, &s);
	gsl_linalg_LU_invert(B, perm, Binv);
	gsl_permutation_free(perm);
	gsl_matrix_free(B);
	
	/* shift = B^-1 D */
	//printf("LSQ: Calculating shift\n");
	shift = gsl_matrix_alloc(n_params, 1);
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, Binv, D, 0.0, shift);
	gsl_matrix_free(Binv);
	gsl_matrix_free(D);
	
	/* Apply shifts */
	//printf("LSQ: Applying shifts\n");
	for ( i=0; i<n_params; i++ ) {
		double old, new;
		old = gsl_vector_get(parameters, i);
		new = old + gsl_matrix_get(shift, i, 0);
		gsl_vector_set(parameters, i, new);
		//printf("Parameter %i: %f -> %f\n", i, old, new);
	}
	
	gsl_matrix_free(shift);
	model_free(model);
	//printf("LSQ: Done\n");


}

void refine_lsq(AtomicModel *model_orig, ReflectionList *reflections, RefinementSpec spec) {

	gsl_vector *parameters;
	unsigned int n_obs, n_params, n_atoms, idx, j, iter;
	double scale;
	ReflectionList *calc;
	AtomicModel *model;
	
	model = model_copy(model_orig);
	
	/* Count the number of parameters to refine */
	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++;
	
	/* Count the number of atoms to refine (having first fixed the origin) */
	n_atoms = 0;
	if ( (spec & REFINE_SPEC_X) || (spec & REFINE_SPEC_Y) || (spec & REFINE_SPEC_Z) ) model->atoms[0].refine = 0;
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) n_atoms++;
	}	
	
	n_obs = reflections->n_reflections-1;	/* Number of observations */
	n_params = 1 + n_params * n_atoms;		/* Number of parameters */
	printf("LSQ: Refining %i parameters against %i observations\n", n_params, n_obs);
	
	/* Set initial estimates of the parameters */
	parameters = gsl_vector_alloc(n_params);
	idx = 0;
	
	/* Initial scale factor */
	calc = model_calculate_f(reflections, NULL, 69);
	scale = stat_scale(reflections, calc);
	reflist_free(calc);
	gsl_vector_set(parameters, idx++, scale);
	
	/* Initial atomic coordinates and (iso-)B-factors */
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].refine ) {
			if ( spec & REFINE_SPEC_X ) gsl_vector_set(parameters, idx++, model->atoms[j].x);
			if ( spec & REFINE_SPEC_Y ) gsl_vector_set(parameters, idx++, model->atoms[j].y);
			if ( spec & REFINE_SPEC_Z ) gsl_vector_set(parameters, idx++, model->atoms[j].z);
			if ( spec & REFINE_SPEC_B ) gsl_vector_set(parameters, idx++, model->atoms[j].B);
			if ( spec & REFINE_SPEC_OCC ) gsl_vector_set(parameters, idx++, model->atoms[j].occ);
		}
	}
	
	iter = 0;
	do {
		
		/* Iterate */
		refine_lsq_iterate(parameters, reflections, spec, model);
		printf("LSQ: after iteration %i:\tscale=%f\n", iter, gsl_vector_get(parameters, 0));
		iter++;
		
		/* Put the parameter vector back into the atomic model */
		model_free(model);
		model = refine_model_from_parameters(parameters, spec, model);
		model_move(model_orig, model);
		
		refine_schedule_update(model);
		
	} while ( (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_vector_free(parameters);

}