/*
 * clean.c
 *
 * CLEAN Algorithm
 *
 * (c) 2006-2007 Thomas White <taw27@cam.ac.uk>
 *
 *  Synth2D - two-dimensional Fourier synthesis
 *
 */

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

#define _GNU_SOURCE
#include <gtk/gtk.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>

#include "displaywindow.h"
#include "main.h"
#include "data.h"
#include "png-file.h"

static void clean_noise_open(fftw_complex *dirty, double max_am);

/* ........................ The Maths ........................ */


static double clean_tau = 0.1;

static void clean_reconvolve(fftw_complex *cleaned, double ax, double ay) {

	signed int h, k;
	unsigned int x, y;
	fftw_complex *in;
	signed int height = data_height();
	signed int width = data_width();
	fftw_plan plan, planr;
	fftw_complex *clean_beam_reciprocal;
	
	in = fftw_malloc(width*height*sizeof(fftw_complex));
	clean_beam_reciprocal = fftw_malloc(width*height*sizeof(fftw_complex));
	for ( x=0; x<width; x++ ) {
		for ( y=0; y<height; y++ ) {
		
			signed int xd, yd;
		
			xd = (((signed)x + width/2) % width)-width/2;
			yd = (((signed)y + height/2) % height)-height/2;
		
			in[y+height*x][0] = (ax*ay)/((ax*ax + xd*xd) + (ay*ay + yd*yd));
			in[y+height*x][1] = 0;
			
		}
	}
	plan = fftw_plan_dft_2d(width, height, in, clean_beam_reciprocal, FFTW_BACKWARD, FFTW_MEASURE);
	fftw_execute(plan);
	fftw_destroy_plan(plan);
	
	/* Transform to reciprocal space, and scale */
	plan = fftw_plan_dft_2d(width, height, cleaned, in, FFTW_BACKWARD, FFTW_MEASURE);
	fftw_execute(plan);
	fftw_destroy_plan(plan);
	for ( h=0; h<width; h++ ) {
		for ( k=0; k<height; k++ ) {
			in[k+height*h][0] = in[k+height*h][0] / (width*height);	/* Real */
			in[k+height*h][1] = in[k+height*h][1] / (width*height);	/* Imaginary */
		}
	}
	
	/* Perform the convolution */
	for ( h=-width/2; h<=width/2; h++ ) {
		for ( k=-height/2; k<=height/2; k++ ) {
			
			double re, im, am, ph;
			unsigned int hl, kl;
			
			if ( h < 0 ) hl = width+h; else hl = h;
			if ( k < 0 ) kl = height+k; else kl = k;
			re = in[kl+height*hl][0];
			im = in[kl+height*hl][1];
			am = sqrt(re*re + im*im);
			ph = atan2(im, re);
			
			am = am*clean_beam_reciprocal[kl+height*hl][0];	
			in[kl+height*hl][0] = am*cos(ph);
			in[kl+height*hl][1] = am*sin(ph);
			
		}
	}

	/* Transform back to real space */
	planr = fftw_plan_dft_2d(width, height, in, cleaned, FFTW_BACKWARD, FFTW_MEASURE);
	fftw_execute(planr);
	fftw_destroy_plan(planr);
	
	fftw_free(in);
	fftw_free(clean_beam_reciprocal);
	
}

/* Calculate and return the aperture function */
static fftw_complex *clean_aperture(ReflectionList *reflections, signed int width, signed int height) {

	fftw_complex *aperture;
	fftw_plan plan;
	double aperture_max;
	fftw_complex *delta_array;
	unsigned int i;
	unsigned int x, y;

	delta_array = fftw_malloc(width*height*sizeof(fftw_complex));
	for ( i=1; i<reflections->n_reflections; i++ ) {
		signed int h, k;
		h = reflections->refs[i].h;
		k = reflections->refs[i].k;
		if ( h < 0 ) { h = width+h; }
		if ( k < 0 ) { k = height+k; }
		delta_array[k + height*h][0] = 1;
		delta_array[k + height*h][1] = 0;	/* Phase = zero */
	}
	aperture = fftw_malloc(width*height*sizeof(fftw_complex));
	plan = fftw_plan_dft_2d(width, height, delta_array, aperture, FFTW_BACKWARD, FFTW_MEASURE);
	fftw_execute(plan);
	fftw_destroy_plan(plan);
	fftw_free(delta_array);
	
	/* Find the maximum *amplitude* in the function */
	aperture_max = 0;
	for ( x=0; x<width; x++ ) {
		for ( y=0; y<height; y++ ) {
			double re, im, am;
			re = aperture[y+height*x][0];  im = aperture[y+height*x][1];
			am = sqrt(re*re + im*im);
			if ( am > aperture_max ) aperture_max = am;
		}
	}
	
	/* Normalise the aperture function */
	for ( x=0; x<width; x++ ) {
		for ( y=0; y<height; y++ ) {
			aperture[y+height*x][0] = aperture[y+height*x][0] / aperture_max;
			aperture[y+height*x][1] = aperture[y+height*x][1] / aperture_max;
		}
	}
	
	return aperture;

}

static void clean_execute(fftw_complex *input, ReflectionList *reflections, double tau, double cutoff, int positive, int convolve, double ax, double ay) {

	fftw_complex *aperture;
	fftw_complex *dirty;
	signed int width = data_width();
	signed int height = data_height();
	
	unsigned int x, y;
	fftw_complex *cleaned;
	unsigned int npx = 0;
	
	double max_am;

	aperture = clean_aperture(reflections, width, height);
	
	/* The aperture function won't necessarily be real: if there was no measured reflection -h,-k,-l for a corresponding
		h,k,l, Friedel's Law will not apply to the delta array. */
	
	/* Initialise 'clean' array */
	printf("CL: Loop gain = %f, Cutoff = %f, %s\n", tau, cutoff, positive?"Constrained positive":"Negativity allowed");
	cleaned = fftw_malloc(width*height*sizeof(fftw_complex));
	for ( x=0; x<width; x++ ) {
		for ( y=0; y<height; y++ ) {
			cleaned[y+height*x][0] = 0;
			cleaned[y+height*x][1] = 0;
		}
	}
	
	/* Initialise 'dirty' array */
	dirty = malloc(width*height*sizeof(fftw_complex));
	memcpy(dirty, input, width*height*sizeof(fftw_complex));
	
	/* Iterate */
	do {
	
		double max_re = 0;
		double max_im = 0;
		unsigned int max_x, max_y;
	
		/* Find largest modulus, record its location */
		max_am = 0; max_x = 0; max_y = 0;
		for ( x=0; x<width; x++ ) {
			for ( y=0; y<height; y++ ) {
			
				double re, im, am;
				re = dirty[y+height*x][0];  im = dirty[y+height*x][1];
				am = sqrt(re*re + im*im);
				
				if ( fabs(am) > fabs(max_am) ) {
					if ( positive ) {
						/* See if this subtraction would result in a negative function */
						if ( (cleaned[y+height*x][0] + tau*re ) < 0 ) continue;
					}
					max_am = am;  max_re = re;  max_im = im;  max_x = x;  max_y = y;
				}
				
			}
		}
		
		/* Add tau*peak to clean map */
		cleaned[max_y+height*max_x][0] += tau * max_re;
		cleaned[max_y+height*max_x][1] += tau * max_im;
		
		/* Subtract tau*peak*aperture from the dirty map */
		for ( x=0; x<width; x++ ) {
			for ( y=0; y<height; y++ ) {
				unsigned int fx, fy;
				fx = (x + max_x) % width;
				fy = (y + max_y) % height;
				dirty[fy+height*fx][0] -= tau * max_am
						* ((aperture[y+height*x][0]*max_re - aperture[y+height*x][1]*max_im) / max_am);
				dirty[fy+height*fx][1] -= tau * max_am
						* ((aperture[y+height*x][0]*max_im - aperture[y+height*x][1]*max_re) / max_am);
			}
		}		

		/* Repeat */
		npx++;
			
	} while ( fabs(max_am) > fabs(cutoff) );
	
	fftw_free(aperture);

	printf("CL: %i iterations performed\n", npx);
	
	clean_noise_open(dirty, max_am);
	free(dirty);
	
	max_am = 0;
	for ( x=0; x<width; x++ ) {
		for ( y=0; y<height; y++ ) {
			double re, im, am;
			re = cleaned[y+height*x][0];  im = cleaned[y+height*x][1];
			am = sqrt(re*re + im*im);
			if ( fabs(am) > fabs(max_am) ) max_am = am;
		}
	}
	
	if ( convolve ) clean_reconvolve(cleaned, ax, ay);
	
	displaywindow_set_realspace(cleaned, DWR_CLEAN);
	displaywindow_forceview(DWV_REALSPACE);
	displaywindow_switchview();

}

/* ..................... The Other Stuff ..................... */

typedef struct {
	GtkWidget *tau;		/* Loop gain */
	GtkWidget *cutoff;	/* Cutoff */
	GtkWidget *conv;	/* Reconvolve? */
	GtkWidget *pos;		/* Constrain positive */
	GtkWidget *ax;		/* CLEAN beam ax */
	GtkWidget *ay;		/* CLEAN beam ay */
} CleanDialog;

CleanDialog *clean_dialog = NULL;

static void clean_gpwrite(FILE *gnuplot, const char *string) {
	fwrite(string, strlen(string), 1, gnuplot);
}

static void clean_plot_graph(char dim) {

	FILE *gnuplot;

	gnuplot = popen("gnuplot -persist -", "w");
	if ( !gnuplot ) {
		error_report("Couldn't invoke gnuplot.  Please check your PATH.");
		return;
	}
	
	clean_gpwrite(gnuplot, "set autoscale\n");
	clean_gpwrite(gnuplot, "unset log\n");
	clean_gpwrite(gnuplot, "unset label\n");
	clean_gpwrite(gnuplot, "set xtic auto\n");
	clean_gpwrite(gnuplot, "set ytic auto\n");
	clean_gpwrite(gnuplot, "set grid\n");
	clean_gpwrite(gnuplot, "set ylabel 'Amplitude' font \"Helvetica,14\"\n");
	if ( dim == 'x' ) clean_gpwrite(gnuplot, "set xlabel 'x / pixels' font \"Helvetica,14\"\n");
	if ( dim == 'y' ) clean_gpwrite(gnuplot, "set xlabel 'y / pixels' font \"Helvetica,14\"\n");
	if ( dim == 'x' ) clean_gpwrite(gnuplot, "set title 'x-component of aperture' font \"Helvetica,14\"\n");
	if ( dim == 'y' ) clean_gpwrite(gnuplot, "set title 'y-component of aperture' font \"Helvetica,14\"\n");
	clean_gpwrite(gnuplot, "plot 'synth2d-aperturefit.dat' with lines\n");
	clean_gpwrite(gnuplot, "replot 'synth2d-aperturefit-peaks.dat' with points\n");
	clean_gpwrite(gnuplot, "a=0.01\n");
	clean_gpwrite(gnuplot, "fit a**2/(a**2+x**2) 'synth2d-aperturefit-peaks.dat' via a\n");
	clean_gpwrite(gnuplot, "set label 'a = %3.5g px',a at graph 0.6, graph 0.7 font \"Helvetica,14\"\n");
	clean_gpwrite(gnuplot, "replot a**2/(a**2+x**2)\n");

	if ( pclose(gnuplot) == -1 ) {
		error_report("gnuplot returned an error code.");
		return;
	}

}

static gint clean_calculate_clean_beam(GtkWidget *cleanw_fit, gpointer data) {

	fftw_complex *aperture;
	FILE *fh;
	FILE *fh_fit;
	unsigned int i, width, height;
	double last_am;
	
	width = data_width();  height = data_height();
	aperture = clean_aperture(main_reflist(), width, height);
	
	fh = fopen("synth2d-aperturefit.dat", "w");
	fh_fit = fopen("synth2d-aperturefit-peaks.dat", "w");
	last_am = INFINITY;
	for ( i=0; i<width/2; i++ ) {

		double re, im, am, next_am;
		
		re = aperture[0+height*i][0];
		im = aperture[0+height*i][1];
		am = sqrt(re*re + im*im);
		if ( re < 0 ) am = -am;
		fprintf(fh, "%i %f\n", i, am);
		
		if ( i+1<width/2 ) {
			re = aperture[0+height*(i+1)][0];
			im = aperture[0+height*(i+1)][1];
			next_am = sqrt(re*re + im*im);
		} else {
			next_am = INFINITY;
		}
		
		if ( (am > last_am) && (am > next_am) ) fprintf(fh_fit, "%i %f\n", i, am);

		last_am = am;

	}
	fclose(fh);
	fclose(fh_fit);
	clean_plot_graph('x');	

	fh = fopen("synth2d-aperturefit.dat", "w");
	fh_fit = fopen("synth2d-aperturefit-peaks.dat", "w");
	last_am = INFINITY;
	for ( i=0; i<height/2; i++ ) {

		double re, im, am, next_am;
		
		re = aperture[i+height*0][0];
		im = aperture[i+height*0][1];
		am = sqrt(re*re + im*im);
		if ( re < 0 ) am = -am;
		fprintf(fh, "%i %f\n", i, am);
		
		if ( i+1<width/2 ) {
			re = aperture[(i+1)+height*0][0];
			im = aperture[(i+1)+height*0][1];
			next_am = sqrt(re*re + im*im);
		} else {
			next_am = INFINITY;
		}
		
		if ( (am > last_am) && (am > next_am) ) fprintf(fh_fit, "%i %f\n", i, am);

		last_am = am;

	}
	fclose(fh);
	fclose(fh_fit);
	clean_plot_graph('y');	

	fftw_free(aperture);
	return 0;

}

static gint clean_window_close(GtkWidget *widget, gint response, CleanDialog *this_clean_dialog) {

	if ( response == GTK_RESPONSE_OK ) {
		
		int n;
		double tau;
		float cutoff, ax, ay;
		const char *tmp;
		
		n = 1;
		tmp = gtk_entry_get_text(GTK_ENTRY(this_clean_dialog->cutoff));
		if ( sscanf(tmp, "%f", &cutoff) != 1 ) n = 0;
		tmp = gtk_entry_get_text(GTK_ENTRY(this_clean_dialog->ax));
		if ( sscanf(tmp, "%f", &ax) != 1 ) n = 0;
		tmp = gtk_entry_get_text(GTK_ENTRY(this_clean_dialog->ay));
		if ( sscanf(tmp, "%f", &ay) != 1 ) n = 0;
		
		if ( n == 1 ) {
			tau = gtk_range_get_value(GTK_RANGE(this_clean_dialog->tau));
			clean_execute(displaywindow_outarray(), main_reflist(),
					tau,
					cutoff,
					gtk_toggle_button_get_active(GTK_TOGGLE_BUTTON(this_clean_dialog->pos)),
					gtk_toggle_button_get_active(GTK_TOGGLE_BUTTON(this_clean_dialog->conv)),
					ax, ay);
		} else {
			return 1;
		}
		
	}
	
	gtk_widget_destroy(widget);
	free(clean_dialog);
	clean_dialog = NULL;
	
	return 0;

}

void clean_dialog_open(GtkWidget *widget, gpointer data) {

	GtkWidget *clean_window;
	GtkWidget *vbox;
	GtkWidget *hbox;
	GtkWidget *table;
	GtkWidget *cleanw_tau_label;
	GtkWidget *cleanw_cutoff_label;
	GtkWidget *cleanw_fit;
	GtkWidget *cleanw_xa_label;
	GtkWidget *cleanw_xa_unit_label;
	GtkWidget *cleanw_ya_label;
	GtkWidget *cleanw_ya_unit_label;
	double max_peak = 2000;
	char str[16];
	
	if ( clean_dialog ) return;
	clean_dialog = malloc(sizeof(CleanDialog));
	
	clean_window = gtk_dialog_new_with_buttons("CLEAN Image", GTK_WINDOW(displaywindow_gtkwindow()),
		GTK_DIALOG_DESTROY_WITH_PARENT,	GTK_STOCK_CANCEL, GTK_RESPONSE_CLOSE, GTK_STOCK_OK, GTK_RESPONSE_OK, NULL);
	
	vbox = gtk_vbox_new(FALSE, 0);
	hbox = gtk_hbox_new(TRUE, 0);
	gtk_box_pack_start(GTK_BOX(GTK_DIALOG(clean_window)->vbox), hbox, FALSE, FALSE, 7);
	gtk_box_pack_start(GTK_BOX(hbox), vbox, FALSE, FALSE, 5);
	table = gtk_table_new(7, 3, FALSE);
	gtk_table_set_row_spacings(GTK_TABLE(table), 5);
	gtk_table_set_col_spacings(GTK_TABLE(table), 5);
	gtk_box_pack_start(GTK_BOX(vbox), table, FALSE, FALSE, 5);

	cleanw_tau_label = gtk_label_new("Loop Gain:");
	gtk_misc_set_alignment(GTK_MISC(cleanw_tau_label), 1, 0.5);
	clean_dialog->tau = gtk_hscale_new_with_range(0, 1, 0.05);
	gtk_scale_set_value_pos(GTK_SCALE(clean_dialog->tau), GTK_POS_RIGHT);
	gtk_range_set_value(GTK_RANGE(clean_dialog->tau), clean_tau);
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_tau_label, 1, 2, 1, 2);
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->tau, 2, 4, 1, 2);
	
	cleanw_cutoff_label = gtk_label_new("Cutoff:");
	gtk_misc_set_alignment(GTK_MISC(cleanw_cutoff_label), 1, 0.5);
	clean_dialog->cutoff = gtk_entry_new();
	max_peak = displaywindow_maxpeak();
	snprintf(str, 15, "%f", max_peak/200);
	gtk_entry_set_text(GTK_ENTRY(clean_dialog->cutoff), str);
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_cutoff_label, 1, 2, 2, 3);
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->cutoff, 2, 4, 2, 3);

	clean_dialog->pos = gtk_check_button_new_with_label("Constrain to be Positive");
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->pos, 1, 4, 3, 4);
	
	clean_dialog->conv = gtk_check_button_new_with_label("Reconvolve with CLEAN beam");
	gtk_toggle_button_set_active(GTK_TOGGLE_BUTTON(clean_dialog->conv), TRUE);
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->conv, 1, 4, 4, 5);

	cleanw_xa_label = gtk_label_new("a_x = ");
	gtk_label_set_markup(GTK_LABEL(cleanw_xa_label), "a<sub>x</sub> = ");
	gtk_misc_set_alignment(GTK_MISC(cleanw_xa_label), 1, 0.5);
	clean_dialog->ax = gtk_entry_new();
	cleanw_xa_unit_label = gtk_label_new("px");
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_xa_label, 1, 2, 5, 6);
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->ax, 2, 3, 5, 6);
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_xa_unit_label, 3, 4, 5, 6);
	
	cleanw_ya_label = gtk_label_new("a_y = ");
	gtk_label_set_markup(GTK_LABEL(cleanw_ya_label), "a<sub>y</sub> = ");
	gtk_misc_set_alignment(GTK_MISC(cleanw_ya_label), 1, 0.5);
	clean_dialog->ay = gtk_entry_new();
	cleanw_ya_unit_label = gtk_label_new("px");
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_ya_label, 1, 2, 6, 7);
	gtk_table_attach_defaults(GTK_TABLE(table), clean_dialog->ay, 2, 3, 6, 7);
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_ya_unit_label, 3, 4, 6, 7);
	
	cleanw_fit = gtk_button_new_with_label("Determine CLEAN beam");
	gtk_table_attach_defaults(GTK_TABLE(table), cleanw_fit, 1, 4, 7, 8);
	g_signal_connect(G_OBJECT(cleanw_fit), "clicked", G_CALLBACK(clean_calculate_clean_beam), clean_dialog);	
	
	g_signal_connect(G_OBJECT(clean_window), "response", G_CALLBACK(clean_window_close), clean_dialog);	
	gtk_widget_show_all(clean_window);

}

void clean_aperture_open(ReflectionList *reflections) {

	GtkWidget *clean_aperture_window;
	GdkPixbuf *clean_aperture_pixbuf;
	GtkWidget *clean_aperture_pixmap_widget;
	fftw_complex *aperture;
	signed int width = data_width();
	signed int height = data_height();

	aperture = clean_aperture(reflections, width, height);
	
	clean_aperture_window = gtk_window_new(GTK_WINDOW_TOPLEVEL);
	gtk_window_set_title(GTK_WINDOW(clean_aperture_window), "Aperture Function");
	clean_aperture_pixbuf = displaywindow_render_pixbuf(aperture, 0.05, data_width(), data_height(), data_gamma(), 1, 1);
	fftw_free(aperture);
	clean_aperture_pixmap_widget = gtk_image_new_from_pixbuf(clean_aperture_pixbuf);
	gtk_container_add(GTK_CONTAINER(clean_aperture_window), clean_aperture_pixmap_widget);
	gtk_widget_show_all(clean_aperture_window);	
	
}

static void clean_noise_open(fftw_complex *dirty, double max_am) {

	GtkWidget *clean_aperture_window;
	GdkPixbuf *clean_aperture_pixbuf;
	GtkWidget *clean_aperture_pixmap_widget;

	clean_aperture_window = gtk_window_new(GTK_WINDOW_TOPLEVEL);
	gtk_window_set_title(GTK_WINDOW(clean_aperture_window), "CLEAN Residual");
	clean_aperture_pixbuf = displaywindow_render_pixbuf(dirty, max_am, data_width(), data_height(), data_gamma(), 1, 1);
	clean_aperture_pixmap_widget = gtk_image_new_from_pixbuf(clean_aperture_pixbuf);
	gtk_container_add(GTK_CONTAINER(clean_aperture_window), clean_aperture_pixmap_widget);
	gtk_widget_show_all(clean_aperture_window);	
	
}