#include "i3_mcmc.h"
#include "i3_load_data.h"
#include "i3_math.h"
#include "i3_model.h"
#include "i3_minimizer.h"
#include "i3_fisher.h"
#include "i3_mcmc.h"
#include "i3_utils.h"
#include "i3_great.h"
#include "i3_image_fits.h"
#include "i3_analyze.h"
#include "i3_catalog.h"
#include "models/i3_sersics.h"
#include "i3_version.h"
#include "i3_pywrapper.h"
#include "i3_meds.h"
#include <time.h>



void i3_analyze_wrapper(i3_image * galaxy, i3_image * psf_image, i3_image * weight_image, 
	i3_image * mask_image, i3_image * bestI, i3_options * options, 
	gal_id identifier, i3_flt x, i3_flt y, i3_result * result, int starting_params_only){


	//We run a few functions to initialize FFTW and the RNG.
	//i3_fftw_load_wisdom();
	//atexit(i3_fftw_save_wisdom);
	i3_init_rng();
	// SLB: should we be doing this instead:? i3_init_rng_multiprocess(proc_id);

	//Set verbosity level
	i3_set_verbosity(options->verbosity);		
	
	//Create basic information and space for the chosen model (e.g. milestone)
	i3_model * model = i3_model_create(options->model_name,options);
	i3_parameter_set * params_ml;

	i3_image * weight_stamp;
	if (weight_image){
		weight_stamp = i3_image_cut_out_stamp(weight_image,options->stamp_size);
	}
	else {
		//Make a dummy weight image. Assume noise level is constant across image
		weight_stamp = i3_image_create(options->stamp_size,options->stamp_size);
		if(options->noise_sigma>0) i3_image_fill(weight_stamp,pow(options->noise_sigma,-2));
		else i3_image_fill(weight_stamp,1.);
	}


	//Set the minimizer method as levmar
	i3_minimizer_method method;
	method = (i3_minimizer_method) options->minimizer_method;
	// is this an OK place to put this?? I guess we don't want to set in ini file?
	
	//Set up space for the best likelihood, best parameters and best image
	i3_flt bestL = 0;
	
	FILE * output;
	if (options->save_output){
	  char output_filename[256];
	  snprintf(output_filename,256,"%s/%s",options->output_directory,options->output_filename);

	  //Prepare the output file 
	  output = fopen(output_filename,"a");
	  if(!output) I3_FATAL("Cannot open output file for writing",1);
	}

	clock_t start_time = clock();
    
	//The concept of a different image and stamp size is now obselete (given that we have the catalogue)
	// int image_size = options->stamp_size;
					
	// Arghh, we need to fit a background level for ClusterSTEP, DES etc
	// Bodge this in really messily for now.... Sorry... SLB.
	// This is not a bad fix for ClusterSTEP but will bias GREAT runs which have no background
	// Ideally the options file would include a switch for whether or not we estimate the background 
	// This "background_estimate" step should go into a function really...
	if (options->background_subtract) {

	  i3_flt background_estimate = i3_image_subtract_background(galaxy);
	  i3_print(i3_verb_noisy,"background_estimate = %20.10f",background_estimate);
	}

	// Not convinced the below line is currently doing this (seems to run with ix<0)
	i3_image * stamp = i3_image_cut_out_stamp(galaxy,options->stamp_size);
	if (!stamp) I3_FATAL("Failed to cut stamp out from galaxy",3);

	if (mask_image) {
		i3_image * mask_stamp = i3_image_cut_out_stamp(mask_image,options->stamp_size);
		i3_image_multiply_mask(weight_stamp, mask_stamp);
		i3_image_destroy(mask_stamp);
	}
		
	//Scale the image and noise to unit maximum
	if (options->rescale_stamp) {
	  i3_flt scaling = fabs(1.0/i3_image_sum_masked(stamp, weight_stamp));
	  i3_print(i3_verb_noisy, "Image scaling = %le", scaling);
	  i3_image_scale(stamp,scaling);
	  i3_flt weight_scaling = pow(scaling,-2);
	  i3_image_scale(weight_stamp,weight_scaling);
	  //if(options->noise_sigma>0) i3_image_fill(weight_stamp,pow(options->noise_sigma*scaling,-2));
	}

	if (options->circularize_mask){
		i3_image_circular_mask(weight_stamp);
	}


					
	i3_data_set * data_set; 
	data_set = i3_build_dataset_with_psf_image(options, identifier, stamp, weight_stamp, psf_image);

	//Modify the PSF FWHM definition to match the GREAT08 (and therefore also GREAT10?) simulation code
	//this_psf->fwhm = this_psf->fwhm * pow(1 - this_psf->e1*this_psf->e1 - this_psf->e2*this_psf->e2 , 0.5);
	//printf("%2.2f %2.2f %2.2f %2.2f \n",this_psf->beta,this_psf->fwhm,this_psf->e1,this_psf->e2);
			
	//Now run the main image analysis on this postage stamp - see i3_analyze.c
	//i3_flt mincheck = i3_image_min(stamp);
	//printf("mincheck = %20.10f\n",mincheck);
	//i3_data_set * data_set = i3_build_dataset(options,stamp,weight_stamp,this_psf,psf_type);	      

	/* for debugging (Michael)
	char optname[256];
	snprintf(optname,256,"%s/py_options.txt",options->output_directory);
	i3_options_save(optname,data_set->options);
	i3_image_print_first_element_of_each_row(stamp);
	*/
     
	// argh haven't thought about the value of weight_stamp for non-GREAT10 images...


	if (starting_params_only) {
		// skip the analysis and just run on the starting parameters,
		// for example to get images for a previously found best fit.
        params_ml = i3_model_option_starts(model, options);
        // Get likelihood and model image
        bestL = i3_model_likelihood(model, bestI, params_ml, data_set);
	}
	else
	{
		//Do the proper analysis
		params_ml = i3_analyze_dataset_method(data_set,model,options,&bestL,bestI,method);
	}
	//Calculate minimum and maximum of residuals map
	i3_image * residuals = i3_image_weighted_add_image(stamp,1.0,bestI,-1.0);
	i3_flt min_residuals = i3_image_min_masked(residuals, data_set->weight);	
	i3_flt max_residuals = i3_image_max_masked(residuals, data_set->weight);	

	// get additional image information

	// i3_flt image_args = 0.5;
	i3_flt image_info = 0.0;

	//  if(0==strcmp(model->name,"sersics") &&  0==strcmp(options->psf_input,"moffat_catalog")){
	//  i3_sersics_get_image_info(params_ml, data_set, &image_args, &image_info);
	//  }

	i3_flt model_min = i3_image_min(bestI);
	i3_flt model_max = i3_image_max(bestI);
	

	clock_t end_time = clock();
	double time_taken = (double)(end_time-start_time) / CLOCKS_PER_SEC;			

	i3_catalog_row * row = malloc(sizeof(i3_catalog_row));       
	row->x = x;
	row->y = y;
	row->id = identifier;
	row->catalog_row = identifier;

	if (options->save_output){      			
	  //Save the results
	  save_result_im3shape(options,model,params_ml,bestL,output,row,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	  if(options->verbosity>=i3_verb_noisy){
	    printf("im3shape ML parameters: \n");
	  }
	}
				
	if(options->verbosity>=i3_verb_standard){
	  save_result_im3shape(options,model,params_ml,bestL,stdout,row,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	}			
	
	copy_result_im3shape_to_struct(options,model,params_ml,bestL,result,identifier,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	
	if (options->save_images){
			
	  char filename[256];

	  snprintf(filename,256,"%s/model.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(bestI,filename);

	  snprintf(filename,256,"%s/image.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(stamp,filename);
	  
	  snprintf(filename,256,"%s/psf.%ld.fits",options->output_directory,identifier);
	  if (psf_image) i3_image_save_fits(psf_image,filename);
	  
	  snprintf(filename,256,"%s/weight.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(weight_stamp,filename);

	}
		
	//Clean up
	i3_dataset_destroy(data_set);
	i3_image_destroy(stamp);
	i3_image_destroy(residuals);			


	if (options->save_output){
	  fflush(output);
	  fclose(output);
	}

	free(params_ml);
	i3_model_destroy(model);
	i3_image_destroy(weight_stamp);
	free(row);
	
	end_time = clock();
	time_taken = (double)(end_time-start_time) / CLOCKS_PER_SEC;			
	if(options->timeit){
	  fprintf(stdout,"Total time for processing galaxy image in C: %2.2fs\n",time_taken);
	};

	// return (i3_result*) result; //(i3_parameter_set*) params_ml;
}



void i3_analyze_exposures_wrapper(int n_exposures, i3_image ** images, 
	i3_image ** psf_images, i3_image ** weights, i3_image * bestI, 
	i3_transform * transforms, i3_options * options, gal_id identifier, 
	i3_flt x, i3_flt y, i3_result * result, int starting_params_only){

#ifdef USE_WCS

    i3_parameter_set * params_ml;

	//We run a function to initialize the RNG.
	i3_init_rng();

	//Set verbosity level
	i3_set_verbosity(options->verbosity);		
	
	//Create basic information and space for the chosen model (e.g. milestone)
	i3_model * model = i3_model_create(options->model_name,options);
	
	//Set the minimizer method as levmar
	i3_minimizer_method method = (i3_minimizer_method) options->minimizer_method;
	
	//Set up space for the best likelihood, best parameters and best image
	i3_flt bestL = 0;
	
	FILE * output;
	if (options->save_output){
	  char output_filename[256];
	  snprintf(output_filename,256,"%s/%s",options->output_directory,options->output_filename);

	  //Prepare the output file 
	  output = fopen(output_filename,"a");
	  if(!output) I3_FATAL("Cannot open output file for writing",1);
	}

	if (options->background_subtract) {
		for (int i=0; i<n_exposures; i++){
			i3_flt bg = i3_image_subtract_weighted_background(images[i], weights[i]);
			result->exposure_background[i] = bg;
		}
	}


	//Scale the image and noise to unit maximum
	//We take the mean flux across the images.
	i3_flt mean_flux=0.0;
	if (options->rescale_stamp) {
		i3_flt total_flux=0.0;
		for (int i=0; i<n_exposures; i++){
			total_flux += i3_image_sum_masked(images[i], weights[i]);
		}
		mean_flux = fabs(total_flux / n_exposures);
		i3_flt scaling = 1.0/mean_flux;
		i3_flt weight_scaling = pow(scaling,-2);
		for (int i=0; i<n_exposures; i++){
			i3_image_scale(images[i],scaling);
			i3_image_scale(weights[i],weight_scaling);
		}
	}

	if (options->circularize_mask){
		for (int i=0; i<n_exposures; i++){
			i3_image_circular_mask(weights[i]);
		}
	}

	i3_data_set * data_set =
	i3_build_multiple_exposure_dataset(options, identifier, n_exposures, 
		images, weights, psf_images, transforms);

	if (starting_params_only) {
		// skip the analysis and just run on the starting parameters,
		// for example to get images for a previously found best fit.
        params_ml = i3_model_option_starts(model, options);
        // Get likelihood and model image
        bestL = i3_model_likelihood(model, bestI, params_ml, data_set);
	}
	else
	{
		//Otherwise run the standard analysis
		params_ml = i3_analyze_dataset_method(data_set,model,options,&bestL,bestI,method);

	}
					

	//Calculate minimum and maximum of residuals map
	i3_image * residuals = i3_image_weighted_add_image(data_set->image,1.0,bestI,-1.0);		
	i3_flt min_residuals = i3_image_min_masked(residuals, data_set->weight);	
	i3_flt max_residuals = i3_image_max_masked(residuals, data_set->weight);	

	// get additional image information
	i3_flt image_info = 0.0;

	i3_flt model_min = i3_image_min_masked(bestI, data_set->weight);
	i3_flt model_max = i3_image_max_masked(bestI, data_set->weight);
	
	i3_catalog_row * row = malloc(sizeof(i3_catalog_row));       
	row->x = x;
	row->y = y;
	row->id = identifier;
	row->catalog_row = identifier;
	double time_taken = 0.0; //JAZ Doing this in C is problematically inconsistent.
	// We should do it in python

	if (options->save_output){      			
	  //Save the results
	  save_result_im3shape(options,model,params_ml,bestL,output,row,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	  if(options->verbosity>=i3_verb_noisy){
	    printf("im3shape ML parameters: \n");
	  }
	}
				
	if(options->verbosity>=i3_verb_standard){
	  save_result_im3shape(options,model,params_ml,bestL,stdout,row,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	}			
	
	copy_result_im3shape_to_struct(options,model,params_ml,bestL,result,identifier,data_set,min_residuals,max_residuals, model_min, model_max, time_taken, &image_info);
	
	for (int i=0; i<data_set->n_exposure; i++){
	  result->exposure_x[i] = data_set->exposure_x[i];
	  result->exposure_y[i] = data_set->exposure_y[i];
	  result->exposure_e1[i] = data_set->exposure_e1[i];
	  result->exposure_e2[i] = data_set->exposure_e2[i];
	  result->exposure_chi2[i] = data_set->exposure_chi2[i];
	  result->exposure_residual_stdev[i] = data_set->exposure_residual_stdev[i];
	}
	result->mean_flux = mean_flux;

	if (options->save_images){
			
	  char filename[256];

	  snprintf(filename,256,"%s/model.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(bestI,filename);

	  snprintf(filename,256,"%s/image.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(data_set->image,filename);
	  

	  i3_image * psf_stack = i3_create_tiled_image(n_exposures, psf_images);
	  snprintf(filename,256,"%s/psf.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(psf_stack,filename);
	  i3_image_destroy(psf_stack);
	  
	  snprintf(filename,256,"%s/weight.%ld.fits",options->output_directory,identifier);
	  i3_image_save_fits(data_set->weight,filename);

	}
		
	//Clean up
	i3_multiple_exposure_dataset_destroy(data_set);
	i3_image_destroy(residuals);			
	free(params_ml);

	if (options->save_output){
	  fflush(output);
	  fclose(output);
	}

	i3_model_destroy(model);
	free(row);
#else
	I3_FATAL("Tried to use multiple exposures without support for WCS built.", 2);
#endif
}