#include "i3_image_fits.h"
#include "i3_multisersics.h"
#include "i3_mcmc.h"
#include "math.h"
#include "i3_math.h"
#include "float.h"
#include "i3_model_tools.h"
#include "i3_multisersics_definition.h"
#include "i3_options.h"
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include "i3_catalog.h"
#include "i3_analyze.h"


#define BEERMAT_MAX_E   0.95  // THIS VALUE SHOULD BE TESTED/OPTIMIZED
#define BEERMAT_R_PARAM 0.05   // THIS VALUE SHOULD BE TESTED/OPTIMIZED
/*
 * For a discussion of the meaning of the beermat R param see... TODO PUT SOME DESCRIPTION AND REF
 */

#define ARCSEC2DEG 0.000277777777777778

double i3_multisersics_model_exposure_tile_chi2(i3_multisersics_parameter_set * p, i3_data_set * dataset, i3_image * tiled_image)
{
    double chi2 = 0.0;
    int x_start = 0;
    int i;
    i3_image_zero(tiled_image);
    for (i = 0; i< dataset->n_exposure; i++){
        i3_image * single_tile = i3_image_like(dataset->exposure_images[i]);
        i3_multisersics_model_image(p, dataset, single_tile, i);
        i3_image_copy_into_part(single_tile, x_start, 0, tiled_image);
        x_start += single_tile->nx;
#ifdef SERSICS_TRUE_LIKE        
        i3_flt chi2_exposure = i3_chi2_weight_map(single_tile, dataset->exposure_images[i], dataset->exposure_weights[i]);
        chi2 += chi2_exposure;
        dataset->exposure_chi2[i] = chi2_exposure;
#endif    
        i3_image_destroy(single_tile);
    }
    return chi2;
}


void i3_multisersics_model_image(i3_multisersics_parameter_set * p, i3_data_set * dataset, i3_image * model_image, int exposure_index)
{
  if ( dataset->options->use_hankel_transform )
    i3_multisersics_model_image_using_hankel_transform(p, dataset, model_image, false, exposure_index);
  else
    i3_multisersics_model_image_save_components(p, dataset, model_image, false, exposure_index);
}


void i3_multisersics_model_image_using_hankel_transform(i3_multisersics_parameter_set * p, i3_data_set * dataset, i3_image * model_image, bool save_components, int exposure_index){

   I3_FATAL("Have not written code for Hankel multiple exposure - call back later.",1);

}

void i3_multisersics_model_image_save_components(i3_multisersics_parameter_set * p, i3_data_set * dataset, i3_image * model_image, bool save_components, int exposure_index){

// get the image parameters 
    i3_transform transform = dataset->exposure_transforms[exposure_index];
    int n_sub = dataset->upsampling;
    int n_pix = dataset->stamp_size;
    int n_pad = dataset->padding;
    int n_all = n_pix+n_pad;

// allocate memory for high res image

    i3_image * image_hires_pad = i3_image_create(n_all*n_sub,n_all*n_sub);
    i3_image_zero(image_hires_pad);

    // Parameter units are in arcseconds in tangent plane coordinates.
    // so that delta ra ~ delta dec in pixels wherever you are on the plane.
    // Our WCS info is in arcsec so no need to scale.
    double ra_deg = p->ra_as;
    double dec_deg = p->dec_as;
    i3_flt x0, y0, e1_bulge, e2_bulge, ab_dvc;
    i3_wcs_to_image(ra_deg, dec_deg, p->e1, p->e2, p->radius*p->radius, transform, &x0, &y0, &e1_bulge, &e2_bulge, &ab_dvc);

    i3_flt ab_scaling  = ab_dvc / (p->radius*p->radius);
    

    //Save to the data set.
    dataset->exposure_x[exposure_index] = x0;
    dataset->exposure_y[exposure_index] = y0;
    dataset->exposure_e1[exposure_index] = e1_bulge;
    dataset->exposure_e2[exposure_index] = e2_bulge;
    // x0 += (n_pix/2.0);
    // y0 += (n_pix/2.0);

    i3_flt sersic_dvc = p->bulge_index;
    i3_flt sersic_exp = p->disc_index;

    ab_dvc = ab_dvc*n_sub*n_sub;
    i3_flt bulge_disc_ratio = p->radius_ratio;

    i3_flt ab_exp = ab_dvc/bulge_disc_ratio/bulge_disc_ratio;

    i3_flt A_dvc = p->bulge_A/ab_scaling;
    i3_flt A_exp = p->disc_A/ab_scaling;

    // These are pixel center coordinates
    // in the high-res image
    x0 = (x0+(n_pad+1)/2) * n_sub;
    y0 = (y0+(n_pad+1)/2) * n_sub;
    x0 += 0.5;
    y0 += 0.5;
    
    i3_flt e1_disc = e1_bulge;
    i3_flt e2_disc = e2_bulge;
    
    // The default value of these options is zero, meaning no truncation
    i3_flt bulge_truncation = dataset->options->sersics_bulge_truncation * i3_sqrt(ab_dvc);
    i3_flt disc_truncation = dataset->options->sersics_disc_truncation * i3_sqrt(ab_exp);


// make the components
        if(dataset->central_pixel_upsampling){

            int n_central_upsampling = dataset->n_central_pixel_upsampling;
            int n_central_pixels_to_upsample = dataset->n_central_pixels_to_upsample;
	
    		if(save_components){
    			// If you are saving the components then you cannot be too bothered about speed.
                // In that case we make the images twice, saving them once
    			if (A_dvc!=0.0) i3_add_real_space_sersic_truncated_radius_upsample_central(image_hires_pad, ab_dvc, e1_bulge, e2_bulge, A_dvc, x0, y0, sersic_dvc, bulge_truncation, n_central_pixels_to_upsample, n_central_upsampling);
                char filename[256];
                snprintf(filename, 256, "%s/bulge_%ld_%d.fits", dataset->options->output_directory, dataset->identifier,exposure_index);
    			i3_image_save_fits(image_hires_pad,filename);
    			i3_image_zero(image_hires_pad);
    			if (A_exp!=0.0) i3_add_real_space_sersic_truncated_radius_upsample_central(image_hires_pad, ab_exp,  e1_disc,  e2_disc, A_exp, x0, y0, sersic_exp, disc_truncation, n_central_pixels_to_upsample, n_central_upsampling);

                snprintf(filename, 256, "%s/disc_%ld_%d.fits", dataset->options->output_directory, dataset->identifier,exposure_index);
                i3_image_save_fits(image_hires_pad,filename);
    			i3_image_zero(image_hires_pad);
    	
    		}

			if (A_dvc!=0.0) i3_add_real_space_sersic_truncated_radius_upsample_central(image_hires_pad, ab_dvc, e1_bulge, e2_bulge, A_dvc, x0, y0, sersic_dvc, bulge_truncation, n_central_pixels_to_upsample, n_central_upsampling);
			if (A_exp!=0.0) i3_add_real_space_sersic_truncated_radius_upsample_central(image_hires_pad, ab_exp,  e1_disc,  e2_disc, A_exp, x0, y0, sersic_exp, disc_truncation, n_central_pixels_to_upsample, n_central_upsampling);


        }else{
            if(save_components){
                    
                if (A_dvc!=0.0) i3_add_real_space_sersic_truncated_radius(image_hires_pad, ab_dvc, e1_bulge, e2_bulge, A_dvc, x0, y0, sersic_dvc, bulge_truncation);
                char filename[256];
                snprintf(filename, 256, "bulge_%ld_%d.fits", dataset->identifier,exposure_index);
                i3_image_save_fits(image_hires_pad,filename);
                i3_image_zero(image_hires_pad);
                if (A_exp!=0.0) i3_add_real_space_sersic_truncated_radius(image_hires_pad, ab_exp,  e1_disc,  e2_disc, A_exp, x0, y0, sersic_exp, disc_truncation);
                snprintf(filename, 256, "disc_%ld_%d.fits", dataset->identifier,exposure_index);
                i3_image_save_fits(image_hires_pad,filename);
                i3_image_zero(image_hires_pad);
        
            }

                if (A_dvc!=0.0) i3_add_real_space_sersic_truncated_radius(image_hires_pad, ab_dvc, e1_bulge, e2_bulge, A_dvc, x0, y0, sersic_dvc, bulge_truncation);
                if (A_exp!=0.0) i3_add_real_space_sersic_truncated_radius(image_hires_pad, ab_exp,  e1_disc,  e2_disc, A_exp, x0, y0, sersic_exp, disc_truncation);
        }

// Convolve with the pre-prepared point spread function and Sinc kernel

	if( dataset->exposure_kernels[exposure_index] ){
        	i3_image_convolve_fourier( image_hires_pad, dataset->exposure_kernels[exposure_index] );
	}

// downsample

	int cut_start = 0;
	int cut_step = 0;
	//printf("tk3 %d %d %d %d \n",model_image->nx,n_all*n_sub,n_pix*n_sub,n_pix);
	
	if(model_image->nx == n_pix){
		cut_start = (n_pad/2)*n_sub + n_sub/2;
		cut_step = n_sub;
	       	i3_image_dsample_cut_into( image_hires_pad, model_image, cut_start, cut_start, cut_step);
	}else if(model_image->nx == n_pix*n_sub){
		cut_start = n_sub/2;
		cut_step = 1;
	       	i3_image_dsample_cut_into( image_hires_pad, model_image, cut_start, cut_start, cut_step);
	}else if(model_image->nx == n_all*n_sub){
		cut_start = 0;
		cut_step = 1;
		i3_image_dsample_cut_into( image_hires_pad, model_image, cut_start, cut_start, cut_step);
		//i3_image_copy_into( image_hires_pad, model_image );
	}else {
		printf("model_image %d image_hires_pad %d",(int)model_image->nx,(int)image_hires_pad->nx);
		I3_FATAL("model_image size is not correct\n",1);
	};

// clean up

	i3_image_destroy( image_hires_pad );

}


void i3_multisersics_start(i3_multisersics_parameter_set * start, i3_data_set * data_set, i3_options * options){
        
}

i3_flt i3_multisersics_likelihood(i3_image * model_image, i3_multisersics_parameter_set * paramset, i3_data_set * dataset){

    // Some initial sanity checks
    if (!(dataset->image) || !(dataset->weight)) I3_FATAL("Image or weight not set for run of sersics model",1);
    if (model_image->nx != dataset->image->nx  || model_image->ny != dataset->image->ny ) I3_FATAL("Model image and data image different sizes",1);
    if (model_image->nx != dataset->weight->nx || model_image->ny != dataset->weight->ny) I3_FATAL("Model image and weight image different sizes. Possibly you have not set the weight image at all?",1);

// calculate how many likelihood evaluations are taken

        dataset->n_logL_evals++;

// prepare the model image 

    	i3_image_zero( model_image );

// prepare the beermatted param set

	   i3_multisersics_parameter_set * paramset_beermat = i3_multisersics_beermat_params(paramset,dataset->options);
        
// get the model image
        
    i3_flt chi2 = i3_multisersics_model_exposure_tile_chi2( paramset_beermat, dataset, model_image);

// Get the chi^2 value

#ifdef SERSICS_TRUE_LIKE
#else
	chi2=666.;
#endif
        // clean up
       

		if (dataset->options->verbosity>11 && dataset->options->save_images) {
			char filename[256]; 
			snprintf( filename, 256, "%s/trial_model_image_%04d.fits", dataset->options->output_directory, dataset->n_logL_evals);	
        	i3_image_save_fits(model_image,filename);
		}
		
	free(paramset_beermat);

        return -chi2/2.0;
}

void i3_multisersics_beermat_mapping(i3_multisersics_parameter_set * input, i3_multisersics_parameter_set * output, i3_options * options)
{
        // Most parametersare the same in the two parameterizations
        memcpy(output, input, sizeof(i3_multisersics_parameter_set));  // BARNEY: DO WE NEED TO REMEMBER TO FREE HP LATER ON USING FREE()?

        // NOW include the beermat changes, mapping only these internal variables        
        i3_e12_to_e12beermat(input->e1, input->e2, BEERMAT_MAX_E, &(output->e1), &(output->e2));
        i3_r_to_rbeermat(input->radius, BEERMAT_R_PARAM, &(output->radius));

        if(options->sersics_beermat_amplitudes_positive){ 
                // output->bulge_A = i3_fabs(input->bulge_A);
                // output->disc_A = i3_fabs(output->disc_A);
                i3_flt eps = 1e-1;
                output->bulge_A = sqrt(input->bulge_A* input->bulge_A  + eps);
                output->disc_A  = sqrt(input->disc_A * input->disc_A   + eps);
        }

}

i3_multisersics_parameter_set * i3_multisersics_beermat_params(i3_multisersics_parameter_set * p , i3_options * options){

	i3_multisersics_parameter_set * hp = malloc(sizeof(i3_multisersics_parameter_set));
	memcpy(hp, p, sizeof(i3_multisersics_parameter_set));  // BARNEY: DO WE NEED TO REMEMBER TO FREE HP LATER ON USING FREE()?

    i3_flt e1 = p->e1;
    i3_flt e2 = p->e2;
    i3_flt r = p->radius;

    // NOW include the beermat changes, mapping only these internal variables
    i3_flt e1_beermat;
    i3_flt e2_beermat;
    i3_flt r_beermat;

    i3_e12_to_e12beermat(e1, e2, BEERMAT_MAX_E, &e1_beermat, &e2_beermat);
    i3_r_to_rbeermat(r, BEERMAT_R_PARAM, &r_beermat);

	hp->e1 = e1_beermat;
	hp->e2 = e2_beermat;
	hp->radius = r_beermat;
	if(options->sersics_beermat_amplitudes_positive){ 
		// hp->bulge_A = i3_fabs(p->bulge_A);
		// hp->disc_A = i3_fabs(p->disc_A);
        // hp->bulge_A = p->bulge_A * p->bulge_A;
        // hp->disc_A  = p->disc_A * p->disc_A;
        // fprintf(stdout,"aaa\n");

        i3_flt eps = 1e-4;
        hp->bulge_A = sqrt(p->bulge_A* p->bulge_A  + eps);
        hp->disc_A  = sqrt(p->disc_A * p->disc_A   + eps);

        
	}
	
	return hp;
}

// void i3_multisersics_get_image_info(i3_multisersics_parameter_set * params, i3_data_set * dataset,  i3_flt * image_args, i3_flt * image_info){

// // get the image parameters
 
// 	int original_dataset_upsampling=dataset->upsampling ;
// 	dataset->upsampling = 9;
//         int n_sub = dataset->upsampling;
// 	int n_pix = dataset->image->nx; 
//         int n_pad = dataset->padding;
//         int n_all = n_pix+n_pad;

// // prepare the circular psf
	
//         i3_moffat_psf psf_form_circular;
//         psf_form_circular.beta = dataset->psf_form->beta;
//         psf_form_circular.fwhm = dataset->psf_form->fwhm;
// 	psf_form_circular.e1 = 0.;
//        	psf_form_circular.e2 = 0.;
        
// // prepare the parameter set

// 	i3_multisersics_parameter_set * params_circular = malloc(sizeof(i3_multisersics_parameter_set));
// 	memcpy(params_circular,params,sizeof(i3_multisersics_parameter_set));
// 	params_circular->x0 = n_pix/2.;
// 	params_circular->y0 = n_pix/2.;
// 	params_circular->e1 = 0;
// 	params_circular->e2 = 0;
	
// // create image space

// 	//i3_image * image_npsf = i3_image_create( n_all*n_sub, n_all*n_sub ); i3_image_zero( image_npsf );
// 	i3_image * image_wpsf = i3_image_create( n_all*n_sub, n_all*n_sub ); i3_image_zero( image_wpsf );
// 	i3_fourier * fourier_ker = dataset->psf_downsampler_kernel;
// 	i3_fourier * kernel_circular = i3_fourier_conv_kernel_moffat(n_pix, n_sub, n_pad, psf_form_circular.beta, psf_form_circular.fwhm, psf_form_circular.e1, psf_form_circular.e2, dataset->options->psf_truncation_pixels, GREAT10_PSF_TYPE_MOFFAT);


// // get the unconvolved image

// // get the convolved image

// 	dataset->psf_downsampler_kernel = kernel_circular;
//         i3_multisersics_model_image(params_circular,dataset,image_wpsf);

// // measure what needs to be measured
	
// 	image_info[0] = i3_image_get_fwhm(image_wpsf, 0.5, params_circular->x0/2.)/(i3_flt)n_sub ;
// // clean up

// 	dataset->psf_downsampler_kernel = fourier_ker;
// 	dataset->upsampling = original_dataset_upsampling;

// 	i3_image_destroy( image_wpsf );
// 	//i3_image_destroy( image_npsf );
// 	i3_fourier_destroy( kernel_circular );
// 	//i3_fourier_destroy( kernel_elliptic );
// 	free(params_circular);
	
// }