#include "i3_model.h"
#include "i3_model_tools.h"
#include "i3_psf.h"
#include <time.h>
#ifdef MPI
#include "mpi.h"
#endif

int main(){

	i3_fftw_load_wisdom();
	atexit(i3_fftw_save_wisdom);

// set up 

	i3_init_rng();
//	clock_t time_start;

// set up the postage stamp

	int N_add = 1;
	int N_obs = 47;
	int N_mod =  N_obs +  N_add*2; 
	int N_sub = 7;
	int N = N_sub*N_mod;
	int N_win = (2*N_add + 1);  
	i3_flt snr = 20;

	int N_sub2 = N_sub*N_sub;	
//	int N_obs2 = N_obs*N_obs;
	int N_mod2 = N_mod*N_mod;
	int N_win2 = N_win*N_win;

// set up the temp object

	i3_image * image_tmp = i3_image_create(N,N);
	i3_image_zero(image_tmp);

// warm up the fft

	// i3_fourier * fft_warm_up = i3_image_fourier( image_tmp );
	// i3_fourier_destroy(fft_warm_up);

// create B+D galaxy

	i3_flt e1 = 0.4;
	i3_flt e2 = 0.;
	i3_flt x0 = N/2;
	i3_flt y0 = N/2;
	i3_flt bulge_amplitude = 0.4;
	i3_flt disc_amplitude = 0.6;
	i3_flt bulge_radius = 2.715*N_sub;
	i3_flt disc_radius = 1.845*N_sub;
	
	i3_flt bulge_sersic_index = 4.0;
	i3_flt disc_sersic_index = 1.0;
		
// set up the gaussian psf 

	i3_flt psf_sersic_n = 0.5f; 
	i3_flt psf_A=1.0f;
	i3_flt psf_x0=N/2;
	i3_flt psf_y0=N/2;
	i3_flt psf_fwhm = 2.85f*N_sub;
	i3_flt psf_ab = i3_fwhm_to_ab(psf_fwhm,psf_sersic_n);
	i3_flt psf_e1 = 0.0f;
	i3_flt psf_e2 = 0.0f;

// create the galaxy
		
	i3_image * image_mod = i3_image_create(N,N);
	i3_image * image_mod_dvc = i3_image_create(N,N);
	i3_image * image_mod_exp = i3_image_create(N,N);

	i3_image_zero(image_mod);
	i3_add_real_space_sersic(image_mod, bulge_radius*bulge_radius, e1,e2, bulge_amplitude,x0,y0, bulge_sersic_index);
	i3_image_save_fits(image_mod,"image_dvc.fits");
	
	i3_image_moments mom;
	i3_image_compute_moments(image_mod,&mom);
	printf("moments: e1 = %f \t e2 = %f\n",mom.e1,mom.e2);

	i3_image_zero(image_mod);
	i3_add_real_space_sersic(image_mod, disc_radius*disc_radius, e1,e2, disc_amplitude,x0,y0, disc_sersic_index);
	i3_image_save_fits(image_mod,"image_exp.fits");
	
	i3_image_zero(image_mod);
	i3_add_real_space_sersic(image_mod_dvc,bulge_radius*bulge_radius, e1,e2, bulge_amplitude,x0,y0, bulge_sersic_index);
	i3_flt sum_dvc = i3_image_sum( image_mod_dvc );	i3_image_scale( image_mod_dvc, 1./sum_dvc * bulge_amplitude );
	
	i3_add_real_space_sersic(image_mod_exp,disc_radius*disc_radius, e1,e2, disc_amplitude,x0,y0, disc_sersic_index);
	i3_flt sum_exp = i3_image_sum( image_mod_exp );	i3_image_scale( image_mod_exp, 1./sum_exp * disc_amplitude );

	i3_image_add_images_into( image_mod_dvc, image_mod_exp, image_mod );

	i3_image_save_fits(image_mod,"image_bd.fits");

// create psf
	
	i3_image * image_psf = i3_image_create(N,N);	i3_image_zero(image_psf);
	i3_add_real_space_sersic( image_psf,psf_ab,psf_e1,psf_e2, psf_A, psf_x0, psf_y0, psf_sersic_n );
	
	i3_image_centered_into( image_psf, image_tmp );

	i3_fourier * fourier_psf = i3_image_fourier( image_tmp );

	i3_image_save_fits( image_tmp,"image_tmp.fits" );
	i3_image_save_fits( image_psf,"image_psf.fits" );
	

	i3_image * abs_psf = i3_fourier_absolute ( fourier_psf );
	i3_image_save_fits( abs_psf, "image_abs_psf.fits" );
	printf("psf fourier size n_x = %d n_y = %d \n", (int)fourier_psf->nx_f, (int)fourier_psf->ny_f );
	
// set up the pixel integration kernel
	 
	i3_image * top_hat = i3_image_get_top_hat( N_mod,N_mod,N_sub );
	i3_image_centered_into( top_hat, image_tmp );
	i3_fourier * fourier_pix = i3_image_fourier( image_tmp );
	
	i3_image * abs_pix = i3_fourier_absolute ( fourier_pix );
	i3_image_save_fits( abs_pix, "image_abs_pix.fits" );
	printf("pix fourier size n_x = %d n_y = %d \n", (int)fourier_pix->nx_f, (int)fourier_pix->ny_f );

// combine the kernels

	printf("combine psf and pixel kernel\n");
	i3_fourier * fourier_ker = i3_fourier_create(N,N);
	i3_multiply_fourier_fourier_into(fourier_pix,fourier_psf,fourier_ker);
	i3_image * abs_ker  = i3_fourier_absolute ( fourier_ker );
	i3_image_save_fits( abs_ker, "image_abs_ker.fits" );

// create observed image

	int conv_level = 2;

	if(conv_level == 1) i3_image_convolve_fourier( image_mod, fourier_pix );
	if(conv_level == 2) i3_image_convolve_fourier( image_mod, fourier_ker );
			
	i3_image_save_fits( image_mod, "image_gal.fits" );
	
	i3_image * image_obs1 = i3_image_create(N_mod,N_mod);
	i3_image * image_obs2 = i3_image_create(N_obs,N_obs);
	i3_image * image_obs3 = i3_image_create(N_obs,N_obs);

	i3_image_dsample_into( image_mod, image_obs1);
	i3_image_copy_part_into(image_obs1, N_add, N_add, image_obs2);
		

	i3_image_save_fits( image_obs1, "image_obs1.fits" );
	i3_image_save_fits( image_obs2, "image_obs2.fits" );

	i3_image_copy_into( image_obs2, image_obs3);
	i3_flt noise_std = i3_snr_to_noise_std( snr, image_obs3 );
	printf("noise_std = %2.4e \n",noise_std);
	i3_image_add_white_noise( image_obs3, noise_std );

	i3_image_save_fits( image_obs3, "image_obs3.fits" );

// test the log-pdf in xy0 space and save the pdfs 

	//i3_flt noise_std = 1.;
	printf("getting the logL \n");
	i3_flt * logL_xy0 = i3_logL_in_xy0(image_mod,image_obs3,noise_std,N_add,N_sub);
	i3_flt * pdf_xy0 = i3_pdf_from_logL(logL_xy0, N_sub2*N_mod2);
 
	// for test only
	FILE * file;

	file = fopen("logL_xy0.dat","w"); 
	fwrite(logL_xy0,sizeof(i3_flt),N_win2*N_sub2,file);
	fclose(file);
	
	file = fopen("pdf_xy0.dat","w"); 
	fwrite(pdf_xy0,sizeof(i3_flt),N_win2*N_sub2,file);
	fclose(file);	
	

// use the actual function to get the log-p of image given model with marginalised amplitude and xy0 

	i3_flt logL_marg_xy0 = i3_logL_white_noise_marg_A_xy0(image_mod,image_obs3,noise_std,N_add,N_sub);

	printf("i3_logL_white_noise_marg_ampl_xy0 function output: xy0 = %2.6f \n",logL_marg_xy0);

	printf("Now you can run plots_i3_xy0_marg_test.m. Remember to check if input parameters are the same. \n");
}