#include "i3_image_fits.h"
#include "i3_multixray.h"
#include "i3_math.h"
#include "i3_model_tools.h"
#include "i3_multixray_definition.h"
#include "i3_options.h"

#define BEERMAT_MAX_E   0.95  // THIS VALUE SHOULD BE TESTED/OPTIMIZED
#define BEERMAT_R_PARAM 0.2   // THIS VALUE SHOULD BE TESTED/OPTIMIZED

i3_flt i3_multixray_model_exposure_tiles_like(i3_multixray_parameter_set * p, i3_data_set * dataset, i3_image * tiled_image)
{

    int x_start = 0;
    int i;
    i3_image_zero(tiled_image);
    i3_flt total_like = 0.0;
    for (i = 0; i< dataset->n_exposure; i++){
        i3_image * single_tile = i3_image_like(dataset->exposure_images[i]);
        i3_multixray_model_image(p, dataset, single_tile, i);
        i3_image_copy_into_part(single_tile, x_start, 0, tiled_image);
        x_start += single_tile->nx;
        i3_flt like;
		if (dataset->options->poisson_errors){
			like = i3_poisson_like_exposure_background(
				single_tile, dataset->exposure_images[i], 
				dataset->exposure_weights[i], p->background[i]);
		}
		else{
			i3_flt chi2 = i3_chi2_weight_map(single_tile, 
				single_tile, dataset->exposure_weights[i] );
			like = -chi2/2.0;
		}
		total_like += like;
        i3_image_destroy(single_tile);
    }
    return total_like;
}

void i3_multixray_map(i3_multixray_parameter_set * input, i3_multixray_parameter_set * output, i3_options * options)
{
	memcpy(output, input, sizeof(i3_multixray_parameter_set));
	i3_e12_to_e12beermat(input->e1, input->e2, BEERMAT_MAX_E, &(output->e1), &(output->e2));
	i3_r_to_rbeermat(input->rc, BEERMAT_R_PARAM, &(output->rc));
}

i3_flt i3_multixray_likelihood(i3_image * model_image, i3_multixray_parameter_set * p, i3_data_set * dataset)
{
	i3_multixray_parameter_set * p_beermat = malloc(sizeof(i3_multixray_parameter_set));
	*p_beermat = *p;
	i3_e12_to_e12beermat(p->e1, p->e2, 0.95, &(p_beermat->e1), &(p_beermat->e2));

	//Generate the model image	
	i3_flt like = i3_multixray_model_exposure_tiles_like(p_beermat, dataset, model_image);
	// Return log-like
	if (isnan(like)) like = -1e10;
	free(p_beermat);
	return like;

}

void i3_multixray_model_image(i3_multixray_parameter_set * p, 
	i3_data_set * dataset, i3_image * model_image, int exposure_index)
{
	// Get the image size parameters
    int n_sub = dataset->upsampling;
    int n_pix = dataset->stamp_size;
    int n_pad = dataset->padding;
    int n_all = n_pix+n_pad;

    // Make an empty high-res image
    i3_image * image_hires = i3_image_create(n_all*n_sub,n_all*n_sub);
    i3_image_zero(image_hires);

    // ab is in pixels right now.
    // we want it in degrees.
    // you know what, actually - let's just specify that it's in arcmin to start with, or something


    // Transform parameters in the sky frame to the tangent plane
    i3_transform transform = dataset->exposure_transforms[exposure_index];
    double alpha_deg = p->alpha_arcsec*ARCSEC2DEG/transform.cosdec0;
    double delta_deg = p->delta_arcsec*ARCSEC2DEG;
    i3_flt rc2_w = (p->rc*p->rc*ARCSEC2DEG*ARCSEC2DEG);
    i3_flt x0, y0, e1, e2, rc2;
    i3_wcs_to_image(alpha_deg, delta_deg, p->e1, p->e2, rc2_w, transform, &x0, &y0, &e1, &e2, &rc2);
    
    i3_flt ab_scaling  = rc2 / rc2_w;

    // Pad and upsample the centroid coordinates
    x0 = (x0+(n_pad+1)/2) * n_sub;
    y0 = (y0+(n_pad+1)/2) * n_sub;
    x0 += 0.5;
    y0 += 0.5;
	// Get the upsampled radius parameter
    i3_flt rc = i3_sqrt(rc2) * n_sub;

	//Add the profile
	i3_add_real_space_beta_profile(rc, e1, e2, p->A/ab_scaling, p->beta, x0, y0, image_hires);

	// Convolve with the PSF and down-sample into the model image

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

	int cut_start = (n_pad/2)*n_sub + n_sub/2;
	int cut_step = n_sub;

	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, 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, 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, model_image, cut_start, cut_start, cut_step);
	}else {
		printf("model_image %d image_hires_pad %d",(int)model_image->nx,(int)image_hires->nx);
		I3_FATAL("model_image size is not correct\n",1);
	};

	//Free hires image
	i3_image_destroy(image_hires);

}