/*
 * model.c
 *
 * Atomic Model Structures
 *
 * (c) 2006-2007 Thomas White <taw27@cam.ac.uk>
 *
 *  synth2d - Two-Dimensional Crystallographic Fourier Synthesis
 *
 */

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

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

#include "data.h"
#include "elements.h"
#include "main.h"
#include "statistics.h"
#include "symmetry.h"
#include "model.h"
#include "displaywindow.h"
#include "model-editor.h"
#include "multislice.h"
#include "dpsynth.h"
#include "normalise.h"

/* This is the SPOT for the "current atomic model",  but don't assume it's the only model around... */
AtomicModel *model_current = NULL;

/* --------------------------------------------------- Loading and saving of models --------------------------------------------------- */

/* NB This doesn't exactly replicate Perl's chomp()... */
static void chomp(char *line) {

	size_t i;

	for ( i=0; i<strlen(line); i++ ) {
		if ( line[i] == '\n' ) {
			line[i] = '\0';
			return;
		}
		if ( line[i] == '\r' ) {
			line[i] = '\0';
			return;
		}
	}

}

AtomicModel *model_load(const char *filename) {

	FILE *fh;
	AtomicModel *model;
	unsigned int i;
	char line[256];

	model = model_new();

	fh = fopen(filename, "r");
	i = 0;
	while ( fgets(line, 255, fh) != NULL ) {

		float x, y, z, B, occ, thickness;
		char active, refine;
		char tmp[256];

		chomp(line);
		if ( strlen(line) == 0 ) continue;

		/* The '%s' is guaranteed to be shorter than the 'line' which contains it */
		if ( sscanf(line, "%f %f %f %s %f %f %c %c", &x, &y, &z, tmp, &B, &occ, &active, &refine) == 8 ) {
			model->atoms[model->n_atoms].x = x;
			model->atoms[model->n_atoms].y = y;
			model->atoms[model->n_atoms].z = z;
			model->atoms[model->n_atoms].B = B;
			model->atoms[model->n_atoms].occ = occ;
			model->atoms[model->n_atoms].ref = elements_lookup(tmp);
			model->atoms[model->n_atoms].active = (active == 'A')?1:0;  /* Anything other than "A" means "inactive" */
			model->atoms[model->n_atoms].refine = (refine == 'R')?1:0;  /* Anything other than "R" means "don't refine" */
			model->n_atoms++;
		} else if ( strncmp(line, "symmetry ", 9) == 0 ) {
			model->sym = symmetry_encode(line+9);
		} else if ( sscanf(line, "thickness %f nm", &thickness) == 1 ) {
			model->thickness = thickness;
		} else {
			fprintf(stderr, "Unrecognised line in model file (line %i)\n", i);
		}

		i++;

	}
	fclose(fh);

	return model;

}

void model_load_as_current(const char *filename) {

	AtomicModel *new;
	double new_thickness;
	Symmetry new_sym;

	new = model_load(filename);
	new_thickness = new->thickness;
	new_sym = new->sym;

	model_move(model_current, new);	/* Preserves lock_count, editor, refinementwindow, thickness and sym */
	model_free(new);

	/* Restore thickness and symmetry */
	model_current->thickness = new_thickness;
	model_current->sym = new_sym;

	model_notify_update(model_current);

}

void model_save(const char *filename, AtomicModel *model) {

	FILE *fh;
	size_t i;

	fh = fopen(filename, "w");
	fprintf(fh, "symmetry %s\n", symmetry_decode(model->sym));
	fprintf(fh, "thickness %.2f nm\n", model->thickness);
	for ( i=0; i<model->n_atoms; i++ ) {

		fprintf(fh, "%f %f %f %s %f %f %c %c\n", model->atoms[i].x, model->atoms[i].y, model->atoms[i].z,
							elements[model->atoms[i].ref].element_name, model->atoms[i].B,
							model->atoms[i].occ,
							model->atoms[i].active?'A':'a', model->atoms[i].refine?'R':'r');

	}

	fclose(fh);

}

/* --------------------------------------------------- Basic model handling functions ------------------------------------------------- */

AtomicModel *model_new() {

	AtomicModel *model;

	model = malloc(sizeof(AtomicModel));

	model->n_atoms = 0;
	model->sym = PLANEGROUP_P1;
	model->thickness = 0.0;
	model->point_atoms = 0;

	model->editor = NULL;
	model->refine_window = NULL;
	model->lock_count = 0;

	return model;

}

AtomicModel *model_copy(const AtomicModel *a) {
	AtomicModel *model;
	model = malloc(sizeof(AtomicModel));
	memcpy(model, a, sizeof(AtomicModel));
	return model;
}

void model_move(AtomicModel *to, AtomicModel *from) {

	ModelEditor *old_editor;
	RefinementWindow *old_refine;
	int old_lock_count;
	Symmetry old_sym;
	double old_thickness;

	old_editor = to->editor;
	old_refine = to->refine_window;
	old_lock_count = to->lock_count;
	old_sym = to->sym;
	old_thickness = to->thickness;

	memcpy(to, from, sizeof(AtomicModel));

	to->editor = old_editor;
	to->refine_window = old_refine;
	to->lock_count = old_lock_count;
	to->sym = old_sym;
	to->thickness = old_thickness;

}

void model_free(AtomicModel *model) {
	free(model);
}

/* --------------------------------------------------- Structure-Factor Calculations -------------------------------------------------- */

/* Generate a template for writing structure factors when there's nothing better to use */
extern void model_generate_template(ReflectionList *reflections, signed int layer) {

	signed int h;
	signed int k;
	for ( h=-40; h<=40; h++ ) {
		for ( k=-40; k<=40; k++ ) {
			reflist_addref(reflections, h, k, layer);
		}
	}

}

/* Return f(j, hkl) */
static double model_sfac(AtomicModel *model, unsigned int j, signed int h, signed int k, signed int l) {

	unsigned int ref;
	double x, y, z;
	double sfac;
	double s;
	double a = data_a() * 10;
	double b = data_b() * 10;
	double c = data_c() * 10; /* Change to Angstroms */
	double gamma = data_gamma();
	double B;

	if ( model->point_atoms ) return 1;

	x = model->atoms[j].x;	y = model->atoms[j].y;	z = model->atoms[j].z;
	ref = model->atoms[j].ref;	B = model->atoms[j].B;

	s = resolution(h, k, l, a, b, c, gamma);

	/* Calculate f(j,hkl) */
	sfac = elements[ref].sfac_c;
	sfac += elements[ref].sfac_a1 * exp(-elements[ref].sfac_b1 * s * s);
	sfac += elements[ref].sfac_a2 * exp(-elements[ref].sfac_b2 * s * s);
	sfac += elements[ref].sfac_a3 * exp(-elements[ref].sfac_b3 * s * s);
	sfac += elements[ref].sfac_a4 * exp(-elements[ref].sfac_b4 * s * s);

	/* Thermal factor */
	sfac = sfac * exp(- B * s * s);

	return sfac;

}

/* Real part of a particular atom's (and it's equivalents') contribution to F */
static double model_f_contribution_re(AtomicModel *model, unsigned int j, signed int h, signed int k, signed int l) {

	double sfac, cont;
	AtomicModel *equivalents;
	size_t i;

	sfac = model_sfac(model, j, h, k, l);
	equivalents = symmetry_generate_equivalent_atoms(model, j, MODEL_TARGET_CALCULATION);

	cont = 0;
	for ( i=0; i<equivalents->n_atoms; i++ ) {
		double dot, x, y, z;
		x = equivalents->atoms[i].x;	y = equivalents->atoms[i].y;	z = equivalents->atoms[i].z;
		dot = h * (1-x) + k * (1-y) + l * z;
		cont += model->atoms[j].occ * sfac * cos(-2 * M_PI * dot);
	}

	model_free(equivalents);

	return cont;

}

/* Imaginary part of a particular atom's (and it's equivalents') contribution to F */
static double model_f_contribution_im(AtomicModel *model, unsigned int j, signed int h, signed int k, signed int l) {

	double sfac, cont;
	AtomicModel *equivalents;
	size_t i;

	sfac = model_sfac(model, j, h, k, l);
	equivalents = symmetry_generate_equivalent_atoms(model, j, MODEL_TARGET_CALCULATION);

	cont = 0;
	for ( i=0; i<equivalents->n_atoms; i++ ) {
		double dot, x, y, z;
		x = equivalents->atoms[i].x;	y = equivalents->atoms[i].y;	z = equivalents->atoms[i].z;
		dot = h * (1-x) + k * (1-y) + l * z;
		cont += model->atoms[j].occ * sfac * sin(-2 * M_PI * dot);
	}

	model_free(equivalents);

	return cont;

}

double model_mod_f(AtomicModel *model, signed int h, signed int k, signed int l) {

	double sf_re, sf_im;
	unsigned int j;

	sf_re = 0;	sf_im = 0;
	/* Work out the total contribution */
	for ( j=0; j<model->n_atoms; j++ ) {
		if ( model->atoms[j].active ) {
			sf_re += model_f_contribution_re(model, j, h, k, l);
			sf_im += model_f_contribution_im(model, j, h, k, l);
		}
	}

	return sqrt((sf_re*sf_re) + (sf_im*sf_im));

}

/* Calculate kinematical structure factor amplitudes.
 * "layer" is meaningless and is ignored if "given_template" is not NULL */
ReflectionList *model_calculate_f(ReflectionList *given_template, AtomicModel *given_model, signed int layer) {

	ReflectionList *model_reflections;
	unsigned int i;
	AtomicModel *model;
	ReflectionList *template;

	model_reflections = reflist_new();

	if ( !given_template ) {
		template = reflist_new();
		model_generate_template(template, layer);
	} else {
		template = given_template;
	}

	if ( !given_model ) {
		model = model_get_current();
	} else {
		model = given_model;
	}

	if ( model->thickness > 0.0001 ) {

		/* Dynamical diffraction amplitudes */
		model_reflections = multislice_calculate_f_dyn(model, template);

	} else {

		/* Kinematical structure factors */
		for ( i=1; i<template->n_reflections; i++ ) {	/* 'hkl' loop */

			signed int h = template->refs[i].h;
			signed int k = template->refs[i].k;
			signed int l = template->refs[i].l;
			double sf_re = 0;
			double sf_im = 0;
			unsigned int j;

			/* Work out the total contribution */
			for ( j=0; j<model->n_atoms; j++ ) {
				if ( model->atoms[j].active ) {
					sf_re += model_f_contribution_re(model, j, h, k, l);
					sf_im += model_f_contribution_im(model, j, h, k, l);
				}
			}

			//printf("%3i %3i %3i am=%f ph=%f re=%f im=%f\n", h, k, l, sqrt((sf_re*sf_re) + (sf_im*sf_im)), atan2(sf_im, sf_re), sf_re, sf_im);
			reflist_addref_am_ph(model_reflections, h, k, l, sqrt((sf_re*sf_re) + (sf_im*sf_im)), atan2(sf_im, sf_re));
			reflist_set_components(model_reflections, h, k, l, sf_re, sf_im);

		}

	}

	if ( given_template == NULL ) free(template);

	return model_reflections;

}

#if 0
/* Apply a smoothing filter to the amplitudes */
static void model_smooth(ReflectionList *reflections) {

	unsigned int i;
	float a = data_a();
	float b = data_b();
	float c = data_c();
	double sigma = 5.19;

	for ( i=1; i<reflections->n_reflections; i++ ) {

		double s;
		signed int h, k, l;

		h = reflections->refs[i].h;
		k = reflections->refs[i].k;
		l = reflections->refs[i].l;

		s = sqrt(((h*h)/(a*a)) + ((k*k)/(b*b)) + ((l*l)/(c*c)));

		reflections->refs[i].amplitude *= exp(-(s*s)/(2*sigma*sigma));

	}


}
#endif

/* For the difference synthesis */
void model_calculate_difference_coefficients(ReflectionList *reflections) {

	unsigned int i;
	double scale;
	ReflectionList *model_reflections;

	if ( reflections->n_reflections == 0 ) return;	/* No reflections */
	model_reflections = model_calculate_f(reflections, NULL, 69);
	if ( !model_reflections ) return;

	scale = stat_scale(reflections, model_reflections);
	//printf("Scale for difference coefficients = %f\n", scale);

	for ( i=0; i<reflections->n_reflections; i++ ) {	/* 'hkl' loop */

		if ( model_reflections->refs[i].amplitude > 0.0000001 ) {
			reflections->refs[i].amplitude = reflections->refs[i].amplitude - (scale * model_reflections->refs[i].amplitude);
			reflections->refs[i].phase_known = model_reflections->refs[i].phase_known;
			reflections->refs[i].phase_known_set = 1;
		} else {
			reflections->refs[i].phase_known_set = 0;
		}

	}

	//model_smooth(reflections);

	free(model_reflections);

}

/* For the Fourier refinement synthesis */
void model_calculate_refinement_coefficients(ReflectionList *reflections) {

	unsigned int i;
	ReflectionList *model_reflections;

	if ( reflections->n_reflections == 0 ) return;	/* No reflections */
	model_reflections = model_calculate_f(reflections, NULL, 69);
	if ( !model_reflections ) return;

	for ( i=0; i<reflections->n_reflections; i++ ) {	/* 'hkl' loop */

		reflections->refs[i].amplitude = reflections->refs[i].amplitude;
		reflections->refs[i].phase_known = model_reflections->refs[i].phase_known;

	}

	//model_smooth(reflections);

	free(model_reflections);

}

/* ---------------------------------------------------------------- Misc -------------------------------------------------------------- */

void model_open_editor() {
	model_current->editor = model_editor_open(model_current);
}

AtomicModel *model_get_current() {
	return model_current;
}

/* Notify that a model has been changed.
 *  Special version for when a ModelEditor initiates the update
 *  (don't try and update the editor in this case!) */
void model_notify_update_editor(AtomicModel *model) {

	if ( model == model_current ) {
		displaywindow_switchview();
	}

}

/* Notify that a model has been changed  */
void model_notify_update(AtomicModel *model) {

	model_notify_update_editor(model);

	if ( model->editor ) {
		model_editor_get_model(model->editor);
	}

}

/* Initial setting of the current model.  Many functions assume that
 *  a model structure exists. */
void model_default() {
	model_current = model_new();
}

void model_lock(AtomicModel *model) {
	model->lock_count++;
	if ( model->editor ) {
		model_editor_lock(model->editor);
	}
}

void model_unlock(AtomicModel *model) {
	model->lock_count--;
	if ( model->lock_count == 0 ) {
		if ( model->editor ) {
			model_editor_unlock(model->editor);
		}
	}
}

int model_current_is_blank() {
	if ( model_get_current()->n_atoms) {
		return 0;
	}
	return 1;
}