#include "i3_global.h"
#include "i3_mcmc.h"
#include "i3_math.h"
#include "i3_model.h"
#include "float.h"
#include "strings.h"


/** 
* The Metropolis-Hastings MCMC acceptance criterion. 
* Called internally by the MCMC engine.
**/
bool i3_mcmc_should_accept_jump(i3_flt current_likelihood,i3_flt proposed_likelihood){
	if (proposed_likelihood>current_likelihood) return true;
	i3_flt alpha = i3_log(i3_random_uniform());
	if (proposed_likelihood-current_likelihood>alpha) return true;
	return false;
	
}

/**
* Find the index of the best sample in an MCMC that has already been run.
**/
int i3_mcmc_best_index(i3_mcmc * mcmc){
	float best_like=BAD_LIKELIHOOD;
	int best_index=-1;
	for(int i=0;i<mcmc->nSamples;i++){
		if (mcmc->likelihoods[i]>best_like){
			best_index=i;
			best_like=mcmc->likelihoods[i];
		}
	}
	return best_index;

}

/*
* Return a pointer to the best sample with the selected number in the MCMC chain.
*/
i3_parameter_set * i3_mcmc_best_sample(i3_mcmc * mcmc){
	int best_index = i3_mcmc_best_index(mcmc);
	if (best_index==-1) return NULL;
	return i3_mcmc_sample_number(mcmc,best_index);
}


/*
* Return a pointer to the sample with the selected number in the MCMC chain.
* Return NULL if there is no such sample (if n<0 or > number of samples).
*/
i3_parameter_set * i3_mcmc_sample_number(i3_mcmc * mcmc, int n){
	char * samples = (char *) mcmc->samples;
	if (n<0||n>mcmc->nSamples) return NULL;
	return (i3_parameter_set*) (samples+n*mcmc->model->nbytes);	
}

i3_flt i3_mcmc_like_number(i3_mcmc * mcmc, int n){
	if (n<0||n>mcmc->nSamples) return BAD_LIKELIHOOD;
	return mcmc->likelihoods[n];
}

/*
* Return a pointer to the last sample in the MCMC chain.
*/
i3_parameter_set * i3_mcmc_last_sample(i3_mcmc * mcmc){
	return i3_mcmc_sample_number(mcmc,mcmc->nSamples-1);
}

i3_flt i3_mcmc_last_like(i3_mcmc * mcmc){
	return mcmc->likelihoods[mcmc->nSamples-1];
	
}

/**
* Create an i3_mcmc instance, based on the selected model with the selected options.
*/
i3_mcmc * i3_mcmc_create(i3_model * model, i3_options * options){
	i3_mcmc * mcmc = malloc(sizeof(i3_mcmc));
	mcmc->model = model;
	int max_samples = options->number_samples>options->tuning_phase ?
	                  options->number_samples : 
	                  options->tuning_phase;
	mcmc->width = i3_model_option_widths(model,options);
	mcmc->nParam = i3_model_number_varied_params_nonzero_width(model, mcmc->width);
	mcmc->samples = malloc(max_samples * model->nbytes );
	int ntotal_proposals = options->number_samples+options->burn_length+options->tuning_phase;
	mcmc->acceptances = malloc(ntotal_proposals * sizeof(char) );
	mcmc->likelihoods = malloc(max_samples * sizeof(i3_flt));
	mcmc->random_rotation = malloc(mcmc->nParam * mcmc->nParam * sizeof(i3_flt));
	mcmc->scaling_parameter = options->scaling;
	mcmc->model_image = NULL;
	mcmc->nSamples = 0;
	mcmc->nAccept = 0;
	mcmc->cov=NULL;
	mcmc->cov_sqrt=NULL;
	mcmc->propose_rotation=options->do_rotation;
	mcmc->temperature=1.0;
	mcmc->filename_base = options->mcmc_filename_base;
	if (mcmc->propose_rotation) 
		i3_random_rotation(mcmc->random_rotation,mcmc->nParam);
	else
		i3_identity_matrix(mcmc->random_rotation,mcmc->nParam);
	return mcmc;
}


/**
* Free the memory associated with the selected MCMC instance.
* BEWARE: This will also invalidate any pointers previously returned by i3_mcmc_sample_number or i3_mcmc_best_sample
* I will rethink this behaviour.
*/
void i3_mcmc_destroy(i3_mcmc * mcmc){
	if(mcmc->samples) free(mcmc->samples);  mcmc->samples=NULL;
	if (mcmc->cov) free(mcmc->cov);  mcmc->cov=NULL;
	if (mcmc->cov_sqrt) free(mcmc->cov_sqrt);  mcmc->cov_sqrt=NULL;
	if (mcmc->likelihoods) free(mcmc->likelihoods);  mcmc->likelihoods=NULL;
	if (mcmc->random_rotation) free(mcmc->random_rotation);  mcmc->random_rotation=NULL;
	if (mcmc->model_image) i3_image_destroy(mcmc->model_image);  mcmc->model_image=NULL;
	if (mcmc->width) free(mcmc->width); mcmc->width=NULL;
	if (mcmc->acceptances) free(mcmc->acceptances); mcmc->acceptances=NULL;
	mcmc->nSamples=0;
	free(mcmc);
}


void i3_mcmc_record_parameters(i3_mcmc * mcmc, i3_parameter_set * current_parameters, i3_flt current_likelihood, bool accepted){	
	i3_parameter_set * next_sample = i3_mcmc_sample_number(mcmc,mcmc->nSamples);
	i3_model_copy_parameters(mcmc->model,next_sample, current_parameters);
	mcmc->likelihoods[mcmc->nSamples]=current_likelihood;
	mcmc->nSamples++;
	if(accepted) mcmc->nAccept++;
	
}



void i3_mcmc_save_results(i3_mcmc * mcmc, char * filename){
	/* Save the samples from an MCMC to a flat text file.  This is model-specific*/
	FILE * output = fopen(filename,"w");
	char parameter_string[i3_max_string_length];
	for(int i=0;i<mcmc->nSamples;i++){
		i3_parameter_set * sample = i3_mcmc_sample_number(mcmc,i);
		i3_model_parameter_string(mcmc->model,sample,parameter_string,i3_max_string_length);
		fprintf(output,"%e  %s\n",mcmc->likelihoods[i], parameter_string);
	}
	fclose(output);
}

void i3_mcmc_copy_results(i3_mcmc * mcmc, i3_parameter_set * buffer, int nsample_max, int * nsample)
{
	int ncopy = (mcmc->nSamples<=nsample_max) ? mcmc->nSamples : nsample_max;
	int nbytes = ncopy * mcmc->model->nbytes;
	*nsample = ncopy;
	memcpy(buffer, mcmc->samples, nbytes);
}


void i3_mcmc_resize(i3_mcmc * mcmc, int N){
	mcmc->samples = realloc(mcmc->samples,N*mcmc->model->nbytes);
	mcmc->width = realloc(mcmc->width,mcmc->model->nbytes);
	mcmc->likelihoods = realloc(mcmc->likelihoods, N*sizeof(i3_flt));
}


void i3_mcmc_compute_statistics(i3_mcmc * mcmc, i3_parameter_set * means, 
	    i3_parameter_set * stdevs){
	int N=mcmc->nSamples;
	i3_flt row[mcmc->nParam];
	i3_flt sums[mcmc->nParam];
	i3_flt sums2[mcmc->nParam];

	for(int p=0;p<mcmc->nParam;p++){
		sums[p] = 0.0;
		sums2[p] = 0.0;
	}

	for(int s=0;s<N;s++)	{
		i3_model_extract_varied_nonzero_width_parameters(mcmc->model, 
			i3_mcmc_sample_number(mcmc,s), 
			mcmc->width, row);
		for(int p=0;p<mcmc->nParam;p++){
			sums[p] += row[p];
			sums2[p] += i3_pow(row[p],2);
		}
	}

	for(int p=0;p<mcmc->nParam;p++){
		sums[p]/=N;
		sums2[p]=i3_sqrt(sums2[p]/N - sums[p]*sums[p]);
		if (isnan(sums2[p])) sums2[p]=0.0;
	}
	if (stdevs != NULL){
		bzero(stdevs, mcmc->model->nbytes);
		i3_model_input_varied_nonzero_width_parameters(mcmc->model, 
			sums2, stdevs, mcmc->width, mcmc->samples);
	}
	if (means != NULL){
		i3_model_copy_parameters(mcmc->model, means, mcmc->samples);
		i3_model_input_varied_nonzero_width_parameters(mcmc->model, 
			sums, means, mcmc->width, mcmc->samples);		
	}

}

i3_flt i3_mcmc_acceptance_rate(i3_mcmc * mcmc)
{
	return (1.0 * mcmc->nAccept)/(mcmc->nSamples);
	
}

int i3_mcmc_run(i3_mcmc * mcmc, i3_parameter_set * initial_parameters, i3_data_set * data, i3_options * options)
{
	return i3_mcmc_run_callback(mcmc, initial_parameters, data, options, NULL);
}


i3_flt i3_mcmc_annealing_schedule_linear(i3_mcmc * mcmc, i3_options * options){
	int n = mcmc->nSamples;
	int nmax = options->number_samples;

	/* 
	In this linear model for the annealing temperature we have:
		T(r) = Tmax(1-r)
	where r is the fraction of the way through the chain we have moved.
	*/

	i3_flt Tmax=options->max_annealing_temperature;
	i3_flt r = (1.0*n)/nmax;
	
	return Tmax*(1-r);
}


i3_flt i3_mcmc_annealing_schedule_quadratic(i3_mcmc * mcmc, i3_options * options){
	int n = mcmc->nSamples;
	int nmax = options->number_samples;

	/* 
	In this quadratic model for the annealing temperature we have:
		T(r) = a + alpha*r + beta*r**2
	where r is the fraction of the way through the chain we have moved.
	
	We therefore need three constraints to specify the model.
	The first is that at at the end of the chain r=1, T=0.
	For the other two we use the initial temperature r=0, T=T_max
	and the point where r=p T=1.
	*/

	i3_flt Tmax=options->max_annealing_temperature;
	i3_flt p = options->annealing_one_point;


	
	i3_flt alpha = (1-Tmax*(1-p*p))/(p-p*p);
	i3_flt beta = -(alpha+Tmax);

	i3_flt r = (1.0*n)/nmax;
	
	return Tmax + alpha*r + beta*r*r;
}

void i3_mcmc_anneal_callback(i3_mcmc * mcmc, i3_data_set * data_set, i3_options * options){
	mcmc->temperature=i3_mcmc_annealing_schedule_linear(mcmc,options);
}

void i3_mcmc_anneal(i3_mcmc * mcmc, i3_parameter_set * initial_parameters, i3_data_set * data, i3_options * options)
{
	mcmc->temperature=options->max_annealing_temperature;
	i3_mcmc_run_callback(mcmc, initial_parameters, data, options, &i3_mcmc_anneal_callback);
}

void i3_mcmc_adaptive_callback(i3_mcmc * mcmc, i3_data_set * data_set, i3_options * options){
	i3_parameter_set * means = alloca(mcmc->model->nbytes);
	if (mcmc->nSamples>0 && mcmc->nSamples%options->adaptive_metropolis_interval==0){
		i3_mcmc_compute_statistics(mcmc, means, mcmc->width);
	}
}

/*
Simulated annealing algorithm, from wikipedia:
s ← s0; e ← E(s)                                // Initial state, energy.
sbest ← s; ebest ← e                            // Initial "best" solution
k ← 0                                           // Energy evaluation count.
while k < kmax and e < emax                     // While time left & not good enough:
  snew ← neighbour(s)                           // Pick some neighbour.
  enew ← E(snew)                                // Compute its energy.
  if P(e, enew, temp(k/kmax)) > random() then   // Should we move to it?
    s ← snew; e ← enew                          // Yes, change state.
  if enew < ebest then                          // Is this a new best?
    sbest ← snew; ebest ← enew                  // Save 'new neighbour' to 'best found'.
  k ← k + 1                                     // One more evaluation done
return sbest                                    // Return the best solution found.
*/


i3_flt i3_mcmc_recent_acceptance(i3_mcmc * mcmc, int start, int end){
	if (start<0) start=0;
	int nAccept=0;
	for (int i=start; i<end; i++) nAccept+= (int) mcmc->acceptances[i];
	return ((i3_flt) nAccept)/(end-start);
}
#ifdef USE_LAPACK
int i3_mcmc_update_covariance_matrix(i3_mcmc * mcmc){
	int np = mcmc->nParam;
	i3_flt row[np];
	i3_flt means[np];
	i3_flt cov[np*np];
	i3_flt cov_sqrt[np*np];
	int nb = np*np*sizeof(i3_flt);

	// Initialize empty covmat and decomposition
	bzero(cov, nb);
	bzero(cov_sqrt, nb);
	i3_parameter_set * mean_p = malloc(mcmc->model->nbytes);


	// Get the means of the parameters so that we can compute the covariance part
	i3_mcmc_compute_statistics(mcmc, mean_p, NULL);
	i3_model_extract_varied_nonzero_width_parameters(mcmc->model, 
		mean_p, mcmc->width, means);



	//Loop through the samples adding the covariance estimate from that
	//row to the estimate
	for (int i=0; i<mcmc->nSamples; i++){
		i3_parameter_set * sample = i3_mcmc_sample_number(mcmc, i);
		i3_model_extract_varied_nonzero_width_parameters(mcmc->model, sample, mcmc->width, row);
		// i3_model_extract_varied_nonzero_width_parameters(mcmc->model, sample, mcmc->width, row);
		for (int i=0; i<np; i++){
			for (int j=0; j<np; j++){
				cov[i*np+j] += (row[i]-means[i])*(row[j]-means[j]);
			}
		}
	}
	// Divide by the number of samples to get the mean.
	if (i3_should_print(i3_verb_standard)){
		for (int i=0; i<np; i++) printf("means[%d] = %le\n",i,means[i]);
	
		printf("mcmc->nSamples = %d\n",mcmc->nSamples);
		printf("\n");
		for (int i=0; i<np; i++){
			for (int j=0; j<np; j++){
				cov[i*np+j]/=mcmc->nSamples;
				printf("%le    ",cov[i*np+j]);
			}
			printf("\n");
		}
		printf("\n");
	}

	// Try the Cholesky decomp.
	// This is not guaranteed to work because the estimate is noisy.
	int status = i3_matrix_cholesky_sqrt(cov, cov_sqrt, np);

	if (status){
		// If the covmat is bad from noise this may not work, so status!=0
		// We will not update the cov mat at all
		i3_print(i3_verb_standard, "Warning: Failure in Cholesky decomposition of covmat (bad covmat): %d", status);
	}
	else{
		//Status==0 and the covmat is good.
		//Copy covmat and covmat_sqrt to the MCMC
		memcpy(mcmc->cov, cov, nb);
		memcpy(mcmc->cov_sqrt, cov_sqrt, nb);
	}
	return status;
}
#endif

void i3_rescale_proposal(i3_mcmc * mcmc, int i)
{
	int tuning_frequency=50;
	i3_flt rescaling=1.0;
	i3_flt acc = i3_mcmc_recent_acceptance(mcmc, i-tuning_frequency, i);
	// These are empirical.
	// Or rather just completely made up.
	if (acc<0.1) rescaling/=2.0;
	if (acc<0.01) rescaling/=2.0;
	if (acc>0.9) rescaling*=2.0;
	if (acc>0.99) rescaling*=2.0;
	mcmc->scaling_parameter*=rescaling;
	i3_print(i3_verb_noisy, "Rescaling MCMC by factor %lf to %lf (acceptance factor %.2lf)", rescaling, mcmc->scaling_parameter, acc);

}

void i3_tune_proposal(i3_mcmc * mcmc, int i, bool tune_covmat)
{
	int tuning_frequency=50;
	// We only tune the covariance matrix at certain
	// intervals.  If no tuning is to be done just return straight away
	if (i==0 || i%tuning_frequency!=0) return;

	// There are two ways we can do this.
	// either we tune the entire covariance matrix
	// using the current samples...
	if (tune_covmat){
		// Report recent acceptance - handy
		i3_flt acc = i3_mcmc_recent_acceptance(mcmc, i-tuning_frequency, i);
		i3_print(i3_verb_noisy, " Recent (last %d) acceptance factor %.2lf)", tuning_frequency, acc);

#ifdef USE_LAPACK		
		int status = i3_mcmc_update_covariance_matrix(mcmc);
		if (status && acc==0.0){
			i3_rescale_proposal(mcmc, i);
		}
#else
		I3_FATAL("Sorry - you cannot have tune_covmat=True without compiling with -DUSE_LAPACK",1);
#endif

	}
	// Or we just scale the jump sizes collectively
	else{
		i3_rescale_proposal(mcmc, i);
	}


}

int i3_mcmc_setup_covmat(i3_mcmc * mcmc)
{
	// Get varied parameters from width into w.
	i3_flt w[mcmc->nParam];
	i3_model_extract_varied_nonzero_width_parameters(mcmc->model, mcmc->width,mcmc->width,  w);
	
	//Set up space for covmat, if not already present
	int fresh_covmat=0;
	if (mcmc->cov==NULL){
		fresh_covmat=1;
		mcmc->cov=calloc(mcmc->nParam*mcmc->nParam,sizeof(i3_flt));
		// Set the covariance matrix to diagonal with specified sigma**2
		for (int i=0; i<mcmc->nParam; i++) mcmc->cov[i+mcmc->nParam*i]=w[i]*w[i];
	}
	//If needed, allocate space for the cholesky decomposition
	if (mcmc->cov_sqrt==NULL){
		mcmc->cov_sqrt=calloc(mcmc->nParam*mcmc->nParam,sizeof(i3_flt));
	}
#ifdef USE_LAPACK	
	// do the decomposition.
	int status = i3_matrix_cholesky_sqrt(mcmc->cov, mcmc->cov_sqrt, mcmc->nParam);

	if (status){
		// If the covmat is bad from noise this may not work, so status!=0
		i3_print(i3_verb_quiet, "Warning: Failure in Cholesky decomposition of covmat (bad covmat): %d", status);
		// If this is a brand new covmat then the problem must be with the
		// widths.  We can do nothing about that here to give up
		if (fresh_covmat) return status;

		//Otherwise it is from an input covmat 
		//and we can fall back to the widths.
		for (int i=0; i<mcmc->nParam; i++) mcmc->cov[i+mcmc->nParam*i]=w[i]*w[i];
		// And try again.
		status = i3_matrix_cholesky_sqrt(mcmc->cov, mcmc->cov_sqrt, mcmc->nParam);
		if (status){
			i3_print(i3_verb_quiet, "Double covmat failure! %d", status);
		}
	}
#else 
	int status=0;
	for (int i=0; i<mcmc->nParam; i++) {
		mcmc->cov_sqrt[i+mcmc->nParam*i]=i3_sqrt(mcmc->cov[i+mcmc->nParam*i]);
	}
#endif	
	return status;

	
}

int i3_mcmc_run_callback(i3_mcmc * mcmc, i3_parameter_set * initial_parameters, i3_data_set * data, i3_options * options, i3_mcmc_callback_function callback)
{

	i3_mcmc_setup_image(mcmc, data->image);
	int status = i3_mcmc_setup_covmat(mcmc);
	if (status) return status;  // Failure from bad covmat input
	i3_model_prior_violations(mcmc->model, initial_parameters, stdout);

	int ntotal_proposals = options->number_samples+options->burn_length+options->tuning_phase;
	i3_print(i3_verb_noisy, "Running MCMC on %d parameters burn+tune+samples = %d + %d + %d = %d", mcmc->nParam, options->burn_length, options->tuning_phase, options->number_samples, ntotal_proposals);
	
	/* Set initial parameter values */
	i3_parameter_set * current_parameters = malloc(mcmc->model->nbytes);
	i3_parameter_set * proposed_parameters = malloc(mcmc->model->nbytes);	
	i3_model_copy_parameters(mcmc->model,current_parameters,initial_parameters);

	mcmc->current_likelihood=i3_model_posterior(mcmc->model,mcmc->model_image,current_parameters,data);

	/*
	*/
	FILE * samples_file = NULL;
	if (strcmp(mcmc->filename_base, "")){
        char mcmc_filename[i3_max_string_length];
        snprintf(mcmc_filename, i3_max_string_length, mcmc->filename_base, data->identifier);
        samples_file = fopen(mcmc_filename, "w");
	}


	for(int i=0; i<ntotal_proposals;i++){

		/* Propose, and get the likelihood of the new proposal. */
		mcmc->model->propose(mcmc,current_parameters,proposed_parameters);
		mcmc->proposed_likelihood = i3_model_posterior(mcmc->model,mcmc->model_image,proposed_parameters,data);
		/* If we accept the proposal, update the current parameters. */
		bool accept;
		if (options->climbing_only) {
			accept = mcmc->proposed_likelihood > mcmc->current_likelihood;
		}
		else{
			accept = i3_mcmc_should_accept_jump(mcmc->current_likelihood/mcmc->temperature,mcmc->proposed_likelihood/mcmc->temperature);
		}
		mcmc->acceptances[i]=(char)accept;

		if (accept){
			i3_model_copy_parameters(mcmc->model,current_parameters,proposed_parameters);
			mcmc->current_likelihood=mcmc->proposed_likelihood;
		}
		if (options->verbosity>=i3_verb_debug){
			printf("MCMC %d  %le ",i, mcmc->current_likelihood);
			i3_model_printf_parameters(mcmc->model, current_parameters);
			printf("\n");
		}
		/* If we are in the right phase then record the jump and perhaps tune the parameters. */
		if (i>options->burn_length) i3_mcmc_record_parameters(mcmc,current_parameters,mcmc->current_likelihood,accept);
		if ((i>options->burn_length) && (i<options->tuning_phase+options->burn_length)) i3_tune_proposal(mcmc,i,options->mcmc_tune_covmat);


		if (samples_file && (i>options->burn_length+options->tuning_phase)){
			char parameter_string[i3_max_string_length];
			i3_model_parameter_string(mcmc->model,current_parameters,parameter_string,i3_max_string_length);
			fprintf(samples_file,"%le  %s\n",mcmc->current_likelihood, parameter_string);
			fflush(samples_file);
		}

		if (callback) callback(mcmc,data,options);
		// Reset if we are at the end of the tuning phase.
		// This should be all we need to do.
		if (i==options->burn_length+options->tuning_phase){
			mcmc->nSamples=0;
			mcmc->nAccept=0;
		}
	}
	if (samples_file){
		fclose(samples_file);
	}
	free(current_parameters);
	free(proposed_parameters);
	return 0;
}

void i3_mcmc_setup_image(i3_mcmc * mcmc, i3_image * data_image)
{

	if (mcmc->model_image) i3_image_destroy(mcmc->model_image);
	if (data_image) mcmc->model_image = i3_image_like(data_image);
}