#include "i3_model.h"
#include "i3_model_tools.h"
#include "i3_load_data.h"
#include "i3_data_set.h"
#include "i3_great.h"
#include "i3_mcmc.h"
#include "i3_math.h"
#include "i3_minimizer.h"
#include "i3_fisher.h"
#include "i3_psf.h"
#include "i3_image_fits.h"
#include "i3_model_tools.h"
#include "i3_utils.h"
#include "i3_options.h"
#include "i3_analyze.h"
#include <time.h>

void set_start_params(i3_model * model, i3_data_set * dataset, i3_options * options, i3_sersics_parameter_set * start, i3_sersics_parameter_set * truth){

	i3_model_copy_parameters(model, start, truth);
	i3_sersics_start(start,dataset,options);

        start->e1 = 0.;
        start->e2 = 0.;
	start->bulge_A = truth->bulge_A;
        start->disc_A  = truth->disc_A;
	start->x0 = truth->x0;
	start->y0 = truth->y0;
	start->radius = truth->radius;
	
	i3_flt change_r  = 0.8;
	i3_flt change_xy0 = 0.1;
	i3_flt change_A = 0.8;
	
        start->radius *=	(1.+	change_r	*(i3_random_uniform()-0.5));
        start->x0 *= 		(1.+	change_xy0	*(i3_random_uniform()-0.5));
        start->y0 *= 		(1.+	change_xy0	*(i3_random_uniform()-0.5));
        start->bulge_A *=	(1.+	change_A	*(i3_random_uniform()-0.5));
        start->disc_A  *= 	(1.+	change_A	*(i3_random_uniform()-0.5));
   
        
        }
        
   
        
int count_same_parameters(i3_sersics_parameter_set * p1,  i3_sersics_parameter_set * p2, i3_flt delta){
	
	// in here we assume same delta for each parameter, what is not generally correct
		
	//this should be coded in a better way, using a loop
	int n_same = 0;
	if( i3_fabs(p1->x0 - p2->x0) < delta ) n_same++;
	if( i3_fabs(p1->y0 - p2->y0) < delta ) n_same++;
	if( i3_fabs(p1->e1 - p2->e1) < delta ) n_same++;
	if( i3_fabs(p1->e2 - p2->e2) < delta ) n_same++;
	if( i3_fabs(p1->radius  - p2->radius)  < delta ) n_same++;
	if( i3_fabs(p1->bulge_A - p2->bulge_A) < delta ) n_same++;
	if( i3_fabs(p1->disc_A  - p2->disc_A)  < delta ) n_same++;
	
	return n_same;

}

static void save_result(i3_model * model, i3_parameter_set * params, i3_flt likelihood, FILE * output_file, int identifier, i3_data_set * dataset){

        i3_sersics_parameter_set * hp = i3_sersics_beermat_params(params);
                
        i3_flt flux_ratio = hp->bulge_A/(hp->bulge_A + hp->disc_A); 

        fprintf(output_file, "%d\t%e",identifier,likelihood);
        //i3_model_fprintf_parameters( model, output_file, hp); 
	fprintf(output_file, "\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e",hp->x0,hp->y0,hp->e1,hp->e2,hp->radius,hp->bulge_A,hp->disc_A);
	fprintf(output_file, "\t% 2.4f",flux_ratio);
        fprintf(output_file, "\t% 2.4f",dataset->signal_to_noise);
        fprintf(output_file, "\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e\t% 2.4e\t% d\t% d\t% d\t% d\t% d",dataset->levmar_info[0],dataset->levmar_info[1],dataset->levmar_info[2],dataset->levmar_info[3],dataset->levmar_info[4],(int)dataset->levmar_info[5],(int)dataset->levmar_info[6],(int)dataset->levmar_info[7],(int)dataset->levmar_info[8],(int)dataset->levmar_info[9]);
        fprintf(output_file, "\n");
        fflush(output_file);
}

int main(int argc, char *argv[]){

// init 

        //Initialize FFTW
        i3_fftw_load_wisdom();
        atexit(i3_fftw_save_wisdom);

// check the arguments
        
        int n_args = 16+1;
        if(argc !=n_args) {
        
                printf("You supplied %d of %d args.\n",argc,n_args);
                printf("1  filename_ini - name of the ini file\n");
                printf("2  dir_out - output directory \n");
                printf("3  proc_id - task id [int]\n");
                printf("4  params_set_id - id of the parameter set id [int]\n");
                printf("5  n_reals - number of noise realisation per core [int] (if 0, then just save the model images and quit)\n");
                printf("6  psf_type - 0 for Moffat, 1 for Airy\n");
                printf("7  psf_beta\n");
                printf("8  psf_fwhm\n");
                printf("9  psf_e1\n");
                printf("10 psf_e2\n");
                printf("11 gal_r - r_bulge, r_disc will be calculated\n");
                printf("12 gal_r_ratio - r_bulge/r_disc (overwrites the vals in .ini file)\n");
                printf("13 gal_f_ratio - flux ratio defined as flux_bulge/(flux_bulge+flux_disc)\n");
                printf("14 gal_e1\n");
                printf("15 gal_e2\n");
                printf("16 SNR - signal to noise ratio\n");

                return 0;
        }
        
                
// read the arguments

        char * filename_ini = argv[1];
        char * dir_out = argv[2];
        int proc_id = atoi(argv[3]);
        int params_set_id = atoi(argv[4]);
        int n_reals = atoi(argv[5]);
        int psf_type = atoi(argv[6]);
        i3_flt psf_beta = atof(argv[7]);
        i3_flt psf_fwhm = atof(argv[8]);
        i3_flt psf_e1 = atof(argv[9]);
        i3_flt psf_e2 = atof(argv[10]);
        i3_flt gal_r = atof(argv[11]);
        i3_flt gal_r_ratio = atof(argv[12]);
        i3_flt gal_f_ratio = atof(argv[13]);            
        i3_flt gal_e1 = atof(argv[14]);
        i3_flt gal_e2 = atof(argv[15]);
        i3_flt snr = atof(argv[16]);
        

// init files

        char filename_nbc_params[128]; snprintf(filename_nbc_params,128,"%s/nbc_params_%04d.dat",dir_out,params_set_id);
        FILE * file_nbc_params = fopen(filename_nbc_params,"w");

        char * str_fmt; str_fmt = "%d\t%d\t%f\t%f\t%2.3f\t%2.3f\t%2.6f\t%2.6f\t%2.6e\t%2.6e\t%2.6e\t%f\n";
        fprintf(file_nbc_params,str_fmt,params_set_id,psf_type,psf_beta,psf_fwhm,psf_e1,psf_e2,gal_e1,gal_e2,gal_r,gal_r_ratio,gal_f_ratio,snr);
        fclose(file_nbc_params);

// load the options
 
        i3_options * options = i3_options_default();
        i3_options_read(options, filename_ini);

        int n_pix = options->stamp_size;
        int verb  = options->verbosity;

        i3_model * model = i3_model_create(options->model_name,options);
                        
	i3_init_rng_multiprocess(proc_id);
	//i3_gsl_init_rng_environment();
	
// io stuff

        char filename_nbc_result1[128]; snprintf(filename_nbc_result1,128,"%s/nbc_result1_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_result1 = fopen(filename_nbc_result1,"a");
        char filename_nbc_starts1[128]; snprintf(filename_nbc_starts1,128,"%s/nbc_starts1_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_starts1 = fopen(filename_nbc_starts1,"a");

	char filename_nbc_result2[128]; snprintf(filename_nbc_result2,128,"%s/nbc_result2_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_result2 = fopen(filename_nbc_result2,"a");
        char filename_nbc_starts2[128]; snprintf(filename_nbc_starts2,128,"%s/nbc_starts2_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_starts2 = fopen(filename_nbc_starts2,"a");

	char filename_nbc_result3[128]; snprintf(filename_nbc_result3,128,"%s/nbc_result3_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_result3 = fopen(filename_nbc_result3,"a");
        char filename_nbc_starts3[128]; snprintf(filename_nbc_starts3,128,"%s/nbc_starts3_%06d_%06d.dat",dir_out,params_set_id,proc_id);
        FILE * file_nbc_starts3 = fopen(filename_nbc_starts3,"a");

// verb 

        if(verb > 0) printf("[%d] running noise bias calibration with n_reals=%d minimizer_loops=%d verb=%d\n",proc_id,n_reals,options->minimizer_loops,verb);

// loop timer
     
        clock_t time_start, time_end;
        double cpu_time_used;
        time_start = clock();

// make psf
        
        i3_moffat_psf psf_form;
        psf_form.beta = psf_beta;
        psf_form.e1 = psf_e1;
        psf_form.e2 = psf_e2;
        psf_form.fwhm = psf_fwhm;

// prep image used as noisy

        i3_image * image_noisy = i3_image_create( n_pix, n_pix ); i3_image_zero( image_noisy );
        i3_data_set * dataset = i3_build_dataset(options, 0, image_noisy, NULL, &psf_form, psf_type);
        
// prepare the paramters

        i3_sersics_parameter_set * params_gal = i3_model_option_starts(model, options);
        
        params_gal->e1 = gal_e1;
        params_gal->e2 = gal_e2;
        params_gal->radius= gal_r;
        params_gal->radius_ratio = gal_r_ratio;
        i3_flt bulge_A = gal_f_ratio/(1. - gal_f_ratio);
        i3_flt disc_A  = 1.;
        params_gal->bulge_A = bulge_A/(bulge_A + disc_A);
        params_gal->disc_A  = disc_A /(bulge_A + disc_A);
        
        
        
#ifndef _H_I3_SERSICS
        I3_FATAL("Noise bias calibration is currently set up to work on the milestone model, which was not compiled.  Alterations pending.",1);
#else
        if (strcmp(model->name,"sersics")) I3_FATAL("nbc_main is currently set up to work on the sersics model.  Alterations pending.",2);
//      options->sersics_radius_ratio_start = gal_r_ratio;
#endif  

// set up space fot the best found

        i3_flt bestL = 0;
        i3_image * bestI = i3_image_copy( image_noisy );

// run the noise realisations loop

        for(int nr=0;nr<n_reals;nr++){

        // create the true image
                
                params_gal->x0= n_pix/2 + i3_random_normal()*0.5; 
                params_gal->y0= n_pix/2 + i3_random_normal()*0.5;
                //params_gal->x0= n_pix/2; 
                //params_gal->y0= n_pix/2; 
                        
                i3_image_zero( image_noisy );
                i3_sersics_model_image(params_gal,dataset,image_noisy);
		
	// calculate the noise and create the weight map
                
                i3_flt std_noise = 1e-8;
                if(snr>0) std_noise = i3_snr_to_noise_std(snr, image_noisy);
                options->noise_sigma = std_noise;  
                i3_image * weight= i3_image_create(n_pix,n_pix);
                i3_image_fill(weight,pow(std_noise,-2));
                //i3_image_fill(weight,1.);             
                dataset->weight = weight;
                dataset->noise = std_noise;
                
                // add noise to image           
                
                if(snr>0) i3_image_add_white_noise( image_noisy,std_noise );
                
                if(verb>0){ char filename[128]; snprintf( filename, 128, "image_noisy_%04d.fits", proc_id );   i3_image_save_fits( image_noisy, filename );}   
                
	
		int n_loops=5;
		for(int nl=0;nl<n_loops;nl++){

			i3_sersics_parameter_set * params_start = i3_model_option_starts(model, options);
			set_start_params(model, dataset, options, params_start, params_gal);
			       
			// set up the start positions
			options->minimizer_max_iterations=10;
			{
				
				i3_sersics_parameter_set * ps = i3_model_option_starts(model, options);
				i3_model_copy_parameters(model,ps,params_start);
				dataset->start_params = ps;

				i3_parameter_set * params_ml = i3_analyze_dataset_method(dataset, model, options, &bestL, bestI, i3_minimizer_method_levmar);
				
				if(verb>0){
					
					
					save_result(model,params_gal	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_start	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_ml	,bestL,stdout,params_set_id,dataset);
					printf("Levmar finished: tolerance %2.1e snr=%2.4f\tnoise_std= % 2.4e\n",options->minimizer_tolerance,snr,std_noise);
				}
				
				save_result(model,params_start,bestL,file_nbc_starts1,params_set_id,dataset);
				save_result(model,params_ml   ,bestL,file_nbc_result1,params_set_id,dataset);
							
				free(params_ml);
				free(ps);
			}

			options->minimizer_max_iterations=20;
			{

				i3_sersics_parameter_set * ps = i3_model_option_starts(model, options);
				i3_model_copy_parameters(model,ps,params_start);
				dataset->start_params = ps;
								
				i3_parameter_set * params_ml = i3_analyze_dataset_method(dataset, model, options, &bestL, bestI, i3_minimizer_method_levmar);
				
				if(verb>0){
					
					
					save_result(model,params_gal	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_start	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_ml	,bestL,stdout,params_set_id,dataset);
					printf("Levmar finished: tolerance %2.1e snr=%2.4f\tnoise_std= % 2.4e\n",options->minimizer_tolerance,snr,std_noise);
				}
				
				save_result(model,params_start,bestL,file_nbc_starts2,params_set_id,dataset);
				save_result(model,params_ml   ,bestL,file_nbc_result2,params_set_id,dataset);
							
				free(params_ml);
				free(ps);
			}

			options->minimizer_max_iterations=100;
			{

				i3_sersics_parameter_set * ps = i3_model_option_starts(model, options);
				i3_model_copy_parameters(model,ps,params_start);
				dataset->start_params = ps;
								
				i3_parameter_set * params_ml = i3_analyze_dataset_method(dataset, model, options, &bestL, bestI, i3_minimizer_method_levmar);
				
				if(verb>0){
					
					
					save_result(model,params_gal	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_start	,bestL,stdout,params_set_id,dataset);
					save_result(model,params_ml	,bestL,stdout,params_set_id,dataset);
					printf("Levmar finished: tolerance %2.1e snr=%2.4f\tnoise_std= % 2.4e\n",options->minimizer_tolerance,snr,std_noise);
				}
				
				save_result(model,params_start,bestL,file_nbc_starts3,params_set_id,dataset);
				save_result(model,params_ml   ,bestL,file_nbc_result3,params_set_id,dataset);
							
				free(params_ml);
				free(ps);
			}

			free(params_start);
		

		}

                if(verb>-1 && nr%100 == 0) printf("[%d:%d] %d / %d\n",params_set_id,proc_id,nr,n_reals); 
                        
        }

        i3_image_destroy( image_noisy );

        // show the time
        fclose(file_nbc_result1);
	fclose(file_nbc_starts1);
	fclose(file_nbc_result2);
	fclose(file_nbc_starts2);
        time_end = clock();
        cpu_time_used = ((double) (time_end - time_start)) / CLOCKS_PER_SEC;
        if(verb>0) printf("time total %4.2fs\ttime per galaxy %2.2fs\n",cpu_time_used,cpu_time_used/(i3_flt)n_reals);
        
}