/* TwoAdvSelfSims.c

Simulation calculating fixation probability of second beneficial allele, given existing
ben allele at initial frequency p. For use in the study "Linkage and the limits to natural
selection in partially selfing populations".

Simulation uses routines found with the GNU Scientific Library (GSL)
(http://www.gnu.org/software/gsl/)
Since GSL is distributed under the GNU General Public License
(http://www.gnu.org/copyleft/gpl.html), you must download it
separately from this file.

This program can be compiled with e.g. GCC using a command like:
gcc TwoAdvSelfSims -lm -lgsl -lgslcblas -I/usr/local/include -L/usr/local/lib TwoAdvSelfSims.c

Then run by executing:
./TwoAdvSelfSims N self rec ha sa hb sb p reps
Where:
- N is the population size
- self is the rate of self-fertilisation
- rec is recombination rate
- ha, hb is dominance at the original, introduced beneficial allele
- sa, sb is selection coefficient of the original, introduced beneficial allele
- p is the initial frequency of the first beneficial allele (when the second is introduced)
- reps is how many times the second allele should FIX before simulation stops
(the number of actual runs is greater due to stochastic loss of second allele)

Note that haplotypes are defined as:
x1 = ab
x2 = Ab
x3 = aB
x4 = AB

Genotypes defined as:
g11 = g1 = ab/ab
g12 = g2 = Ab/ab
g13 = g3 = aB/ab
g14 = g4 = AB/ab
g22 = g5 = Ab/Ab
g23 = g6 = Ab/aB
g24 = g7 = Ab/AB
g33 = g8 = aB/aB
g34 = g9 = aB/AB
g44 = g10 = AB/AB

Output files are the parameters;
followed by number of times each haplotype fixed;
followed by average total generations elapsed in each case;
Then total number of simulations ran;
Then fixation prob of allele, both unscaled and scaled to unlinked case,
along with 95% CI intervals for the latter case;
then number of allele fixations.

Note that 'fixation' DIFFERS depending on the inputs of sa, sb.
If sa >= sb (interference case) then 'fixation' counts as fixation of second allele on any genetic background.
If sa < sb (replacement case) then 'fixation' only considers fixation of second allele with neutral haplotype
(I.e. where the 'less fit' neutral allele at locus A fixes, instead of selected allele).

*/

/* Preprocessor statements */
#include <stdio.h>
#include <time.h>
#include <math.h>
#include <stddef.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>

/* Function prototypes */
void geninit(double *geninit, double FIS, const gsl_rng *r);
void selection(double *geninit);
void reproduction(double *geninit);
unsigned int hcheck(double *geninit, double *haps, unsigned int *hf, unsigned int stype);

/* Global variable declaration */
unsigned int N = 0;		/* Pop size */
double rec = 0;			/* Recombination rate */
double self = 0;		/* Rate of self-fertilisation */
double ha = 0;			/* Dominance of site A */
double sa = 0;			/* Fitness of site A */
double hb = 0;			/* Dominance of site B */
double sb = 0;			/* Fitness of site B */
double pee = 0;			/* Freq of initial sweep */

/* Main program */
int main(int argc, char *argv[]){
	unsigned int i; 				/* A counter */
	unsigned int reps;				/* Length of simulation (no. of introductions of neutral site) */
	unsigned int stype = 0;			/* What type of sim (replacement or hitch-hiking)? */
	unsigned int nsfix = 0;			/* sims where target type fixed */
	unsigned int nstot = 0;			/* total sims ran */
	unsigned int gens = 0;			/* Number gens elapsed */
	unsigned int isfin = 0;			/* Is sim finished? */
	unsigned int hf = 0;			/* The hap that fixed */
	double pf = 0;					/* Overall fix prob */
	double FIS = 0;					/* Wright's FIS */
	double StdFix = 0;				/* Standard Fixation prob if unlinked */
	double citop = 0;
	double cibot = 0;
	double nsfix2 = 0;
	double nstot2 = 0;
	char selfchar[10];
	char recchar[15];
	char hchar[10];
	char pchar[10];
	char fname[64];
	FILE *ofp_tr;

	/* GSL random number definitions */
	const gsl_rng_type * T;
	gsl_rng * r;

	/* This reads in data from command line. */
	if(argc != 10){
		fprintf(stderr,"Invalid number of input values.\n");
		exit(1);
	}
	N = strtod(argv[1],NULL);
	self = strtod(argv[2],NULL);
	rec = strtod(argv[3],NULL);
	ha = strtod(argv[4],NULL);
	sa = strtod(argv[5],NULL);
	hb = strtod(argv[6],NULL);
	sb = strtod(argv[7],NULL);
	pee = strtod(argv[8],NULL);
	reps = strtod(argv[9],NULL);

	if(sa >= sb){
		stype = 0;
	}else if(sa < sb){
		stype = 1;
	}
	FIS = self/(2.0-self);

	/* Arrays definition and memory assignment */
	double *genotype = calloc(10,sizeof(double));				/* Genotype frequencies */
	unsigned int *gensamp = calloc(10,sizeof(unsigned int));	/* New population samples */
	double *haps = calloc(4,sizeof(double));					/* Haplotypes */
	unsigned int *pfix = calloc(4,sizeof(unsigned int));		/* Haplotypes that fix */
	unsigned int *tfix = calloc(4,sizeof(unsigned int));		/* time that haps fix */

	/* create a generator chosen by the
    environment variable GSL_RNG_TYPE */

	gsl_rng_env_setup();
	if (!getenv("GSL_RNG_SEED")) gsl_rng_default_seed = time(0);
	T = gsl_rng_default;
	r = gsl_rng_alloc(T);

	nsfix = 0;
	nstot = 0;

    while(nsfix < reps){
    	nstot++;

		/* Initialising genotypes */
		geninit(genotype,FIS,r);
		gens = 0;
		isfin = 0;
		hf = 0;

		while(isfin == 0){
    		/* Selection routine */
    		selection(genotype);

	    	/* Reproduction routine */
    		reproduction(genotype);

    		/* Sampling based on new frequencies */
	       	gsl_ran_multinomial(r,10,N,genotype,gensamp);
       		for(i = 0; i < 10; i++){
    	   		*(genotype + i) = (*(gensamp + i))/(1.0*N);
	       	}

	       	gens++;
	       	isfin = hcheck(genotype, haps, &hf, stype);
       	}

       	if(isfin == 1){
       		(*(pfix + hf))++;
    	  	(*(tfix + hf)) += gens;
	       	if(stype == 0){
       			nsfix += (*(haps + 2) + *(haps + 3));
       		}else if(stype == 1){
    	   		nsfix += (*(haps + 2));
	       	}
       	}

	}	/* End of simulation */

	pf = nsfix/(1.0*nstot);
	StdFix = 2*sb*((hb+FIS-hb*FIS)/(1+FIS));		/* Fix prob of new allele if unlinked */

	nstot2 = nstot + 3.84;
	nsfix2 = (1.0/(nstot2))*(nsfix + (3.84/2.0));
	citop = nsfix2 + 1.96*sqrt((1.0/(nstot2))*nsfix2*(1.0-nsfix2));
	cibot = nsfix2 - 1.96*sqrt((1.0/(nstot2))*nsfix2*(1.0-nsfix2));
	citop = citop/(1.0*StdFix);
	cibot = cibot/(1.0*StdFix);

	/* Printing solutions to file */
	/* First, converting values to strings */
    sprintf(selfchar, "%0.2lf",self);
    sprintf(recchar, "%0.7lf",rec);
    sprintf(hchar, "%0.3lf",ha);
    sprintf(pchar, "%0.5lf",pee);

	strcpy(fname,"sim_self");
	strcat(fname,selfchar);
	strcat(fname,"rec");
	strcat(fname,recchar);
	strcat(fname,"h");
	strcat(fname,hchar);
	strcat(fname,"p");
	strcat(fname,pchar);
	strcat(fname,".sim");

	ofp_tr = fopen(fname,"a+");
	fprintf(ofp_tr,"%d %lf %lf %lf %lf %lf %lf %lf ",N,self,rec,ha,sa,hb,sb,pee);
	for(i = 0; i < 4; i++){
		fprintf(ofp_tr,"%d ",*(pfix + i));
	}
	for(i = 0; i < 4; i++){
		fprintf(ofp_tr,"%lf ",((*(tfix + i)))/(1.0*(*(pfix + i))));
	}
	fprintf(ofp_tr,"%d %lf %lf %lf %lf %d\n",nstot,pf,pf/(1.0*StdFix),citop,cibot,reps);
	fclose(ofp_tr);

	/* Freeing memory and wrapping up */
 	gsl_rng_free(r);
 	free(tfix);
 	free(pfix);
 	free(haps);
 	free(gensamp);
	free(genotype);
	return 0;
}

/* Initialising genotypes */
void geninit(double *geninit, double FIS, const gsl_rng *r){

	unsigned int htype = 0;		/* Type of het assignment */

	double *ptype = calloc(3,sizeof(double));
	unsigned int *ctype = calloc(3,sizeof(double));

	/* Routine to determine initial background of second mutant, given first is at frequency p. */
	*(ptype + 0) = (1-pee)*(1-pee) + FIS*pee*(1-pee);
	*(ptype + 1) = 2*pee*(1-pee)*(1-FIS);
	*(ptype + 2) = pee*pee + FIS*pee*(1-pee);

	/* First initialise baseline freqs */
	*(geninit + 0) = *(ptype + 0);
	*(geninit + 1) = *(ptype + 1);
	*(geninit + 2) = 0;
	*(geninit + 3) = 0;
	*(geninit + 4) = *(ptype + 2);
	*(geninit + 5) = 0;
	*(geninit + 6) = 0;
	*(geninit + 7) = 0;
	*(geninit + 8) = 0;
	*(geninit + 9) = 0;

	/* Then decide where to add new mutant */
	gsl_ran_multinomial(r,3,1,ptype,ctype);
	if(*(ctype + 0) == 1){
		*(geninit + 0) -= 1/(1.0*N);
		*(geninit + 2) += 1/(1.0*N);
	}else if(*(ctype + 1) == 1){
		*(geninit + 1) -= 1/(1.0*N);
	  	htype = gsl_ran_bernoulli(r,0.5);
	  	if(htype == 0){
	  		*(geninit + 5) += 1/(1.0*N);
	  	}else if(htype == 1){
	  		*(geninit + 3) += 1/(1.0*N);
	  	}
	}else if(*(ctype + 2) == 1){
		*(geninit + 4) -= 1/(1.0*N);
		*(geninit + 6) += 1/(1.0*N);
	}

	free(ctype);
	free(ptype);

}	/* End of gen initiation routine */

/* Selection routine */
void selection(double *geninit){
	/* Fitness of each genotype */
	double W11, W12, W13, W14, W22, W23, W24, W33, W34, W44;
	double Wmean;				/* Mean fitness */

	W11 = 1;
	W12 = 1 + ha*sa;
	W13 = 1 + hb*sb;
	W14 = 1 + ha*sa + hb*sb;
	W22 = 1 + sa;
	W23 = 1 + ha*sa + hb*sb;
	W24 = 1 + sa + hb*sb;
	W33 = 1 + sb;
	W34 = 1 + ha*sa + sb;
	W44 = 1 + sa + sb;

	/* Mean fitness calculation */
	Wmean = ((*(geninit + 0))*W11) + ((*(geninit + 1))*W12) + ((*(geninit + 2))*W13) + ((*(geninit + 3))*W14) + ((*(geninit + 4))*W22) + ((*(geninit + 5))*W23) + ((*(geninit + 6))*W24) + ((*(geninit + 7))*W33) + ((*(geninit + 8))*W34) + ((*(geninit + 9))*W44);

	/* Changing frequencies by selection */
	*(geninit + 0) = ((*(geninit + 0))*W11)/Wmean;
	*(geninit + 1) = ((*(geninit + 1))*W12)/Wmean;
	*(geninit + 2) = ((*(geninit + 2))*W13)/Wmean;
	*(geninit + 3) = ((*(geninit + 3))*W14)/Wmean;
	*(geninit + 4) = ((*(geninit + 4))*W22)/Wmean;
	*(geninit + 5) = ((*(geninit + 5))*W23)/Wmean;
	*(geninit + 6) = ((*(geninit + 6))*W24)/Wmean;
	*(geninit + 7) = ((*(geninit + 7))*W33)/Wmean;
	*(geninit + 8) = ((*(geninit + 8))*W34)/Wmean;
	*(geninit + 9) = ((*(geninit + 9))*W44)/Wmean;

}	/* End of selection routine */

/* Reproduction routine */
void reproduction(double *geninit){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	/* Haplotypes */
	double x1, x2, x3, x4;

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Baseline change in haplotype frequencies */
	x1 = g11s + (g12s + g13s + g14s)/2.0 - ((g14s - g23s)*rec)/2.0;
	x2 = g22s + (g12s + g23s + g24s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x3 = g33s + (g13s + g23s + g34s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x4 = g44s + (g14s + g24s + g34s)/2.0 - ((g14s - g23s)*rec)/2.0;

	/* Change in SEXUAL frequencies (both outcrossing and selfing) */
	*(geninit + 0) = (g11s + (g12s + g13s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self + (1 - self)*pow(x1,2);
	*(geninit + 4) = (g22s + (g12s + g24s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self + (1 - self)*pow(x2,2);
	*(geninit + 7) = (g33s + (g13s + g34s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self + (1 - self)*pow(x3,2);
	*(geninit + 9) = (g44s + (g24s + g34s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self + (1 - self)*pow(x4,2);
	*(geninit + 1) = ((g12s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x1*x2;
	*(geninit + 2) = ((g13s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x1*x3;
	*(geninit + 3) = ((g14s*pow((1 - rec),2) + g23s*pow(rec,2))*self)/2.0 + 2.0*(1 - self)*x1*x4;
	*(geninit + 5) = ((g23s*pow((1 - rec),2) + g14s*pow(rec,2))*self)/2.0 + 2.0*(1 - self)*x2*x3;
	*(geninit + 6) = ((g24s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x2*x4;
	*(geninit + 8) = ((g34s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x3*x4;

}	/* End of reproduction routine */

/* Has any allele fixed or not? */
unsigned int hcheck(double *geninit, double *haps, unsigned int *hf, unsigned int stype){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	unsigned int retval = 0;

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Calculation of haplotypes */
	*(haps + 0) = g11s + (g12s + g13s + g14s)/2.0;
	*(haps + 1) = g22s + (g12s + g23s + g24s)/2.0;
	*(haps + 2) = g33s + (g13s + g23s + g34s)/2.0;
	*(haps + 3) = g44s + (g14s + g24s + g34s)/2.0;

/* 	printf("Haps are %lf %lf %lf %lf\n",*(haps + 0),*(haps + 1),*(haps + 2),*(haps + 3));*/

	if(*(haps + 0) == 1){
		retval = 1;
		*hf = 0;
	}
	else if(*(haps + 1) == 1){
		retval = 1;
		*hf = 1;
	}
	else if(*(haps + 2) == 1){
		retval = 1;
		*hf = 2;
	}
	else if(*(haps + 3) == 1){
		retval = 1;
		*hf = 3;
	}else if(stype == 1){
		if( (*(haps + 2) + *(haps + 3) ) == 0 ){
			retval = 2;
			*hf = 0;
		}
	}

	return retval;

}	/* End of hap check routine */

/* End of program */


/*
 * Copyright 2017 Daniel Eachern Huang
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef AUGUR_MCMCLIB_H
#define AUGUR_MCMCLIB_H

#include "augur_util.h"
#include "augur_math.h"
#include "augur_vecop.h"
#include "augur_blkop.h"
#include "augur_blkstk.h"
#include "augur_dist.h"
#include <augur_iface.h>
#include <math.h>
#include <gsl/gsl_rng.h>

typedef struct mcmc {
  AugurMod_t curr;
  AugurMod_t prop;
  AugurAux_t aux;
  int accept;
  double currLL;
  double propLL;
  double eps;
  double L;
  double leapLen;
  double deltaMax;
  int nutsalloc;
} mcmc_t;

mcmc_t MCMC;
gsl_rng* h_rng;

void dump_state(mcmc_t state) {
  printf("currLL: %f\n", state.currLL);
  printf("propLL: %f\n", state.propLL);
  printf("eps: %f\n", state.eps);
  printf("leapLen: %f\n", state.leapLen);
  printf("L: %f\n", state.L);

}

// double modObjFn(AugurMod_t curr, AugurAux_t aux);
EXTERNC void mcmcStep();

/*
void acceptAlways() {
  MCMC.currLL = modObjFn(MCMC.curr, MCMC.aux);
}
*/

/**
 * Returns TRUE if we should accept, FALSE if we should reject.
 */
Bool_t augur_mcmc_ar( double auxll
		    , double (*objFn)(AugurAux_t, AugurMod_t)
		    ) {
  double currll = objFn(MCMC.aux, MCMC.curr);
  double propll = objFn(MCMC.aux, MCMC.prop);
  double u = gsl_rng_uniform(h_rng);

  if (augur_log(u) < propll - currll + auxll) {
    return TRUE;
  }
  else {
    return FALSE;
  }
}

/**
 * Returns TRUE if we should accept, FALSE if we should reject.
 */
__HOSTORDEV__ Bool_t augur_mcmc_ar_idx
       ( double auxll
       , AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
       , double (*objFn)(AugurAux_t, AugurMod_t)
       ) {
  double currll = objFn(aux, curr);
  double propll = objFn(aux, prop);
  // double u = gsl_rng_uniform(h_rng);
  double u = augur_std_uniform_sample(aux.rng);

  if (augur_log(u) < propll - currll + auxll) {
    return TRUE;
  }
  else {
    return FALSE;
  }
}

int acceptReject(double auxLL, double (*objFn)(AugurAux_t, AugurMod_t)) {
  MCMC.propLL = objFn(MCMC.aux, MCMC.prop);
  double u = gsl_rng_uniform(h_rng);
  // printf("currLL: %f  propLL: %f  auxLL: %f   ar: %d\n", MCMC.currLL, MCMC.propLL, auxLL, log(u) < MCMC.propLL - MCMC.currLL + auxLL);
  if (log(u) < MCMC.propLL - MCMC.currLL + auxLL) {
    MCMC.accept += 1;
    AugurMod_t tmp = MCMC.curr;
    MCMC.curr = MCMC.prop;
    MCMC.prop = tmp;
    // double tmp2 = MCMC.currLL;
    MCMC.currLL = MCMC.propLL;
    // MCMC.propLL = tmp2;
    return 1;
  }
  else {
    return 0;
  }
}

void swapCurrProp(double (*objFn)(AugurAux_t, AugurMod_t)) {
  MCMC.propLL = objFn(MCMC.aux, MCMC.prop);
  AugurMod_t tmp = MCMC.curr;
  MCMC.curr = MCMC.prop;
  MCMC.prop = tmp;
  double tmp2 = MCMC.currLL;
  MCMC.currLL = MCMC.propLL;
  MCMC.propLL = tmp2;
}

void initMcmc() {
  const gsl_rng_type * T;
  gsl_rng_env_setup();
  T = gsl_rng_default;
  h_rng = gsl_rng_alloc (T);

  // MCMC.curr = curr;
  // MCMC.prop = prop;
  // MCMC.aux = aux;
  MCMC.accept = 0;
  // MCMC.currLL = objFn(MCMC.aux, MCMC.curr);
  MCMC.propLL = MCMC.currLL;
  MCMC.eps = 1.0;
  MCMC.L = 20;
  MCMC.leapLen = 10.0;
  MCMC.deltaMax = 1000.0;
  MCMC.nutsalloc = 0;
  return;
}

EXTERNC void setHmcParams(double eps, double leapLen) {
  MCMC.eps = eps;
  MCMC.leapLen = leapLen;
}

EXTERNC void setHmcParams2(double eps, int L) {
  MCMC.eps = eps;
  MCMC.L = L;
}

#define max(x, y) ( x < y ? y : x)

double h_sample_uniform(gsl_rng* rng, double left, double right) {
  double u = gsl_rng_uniform(rng);
  return left + u * (right - left);
}

/**
 * Metropololis-within-gibbs Kernel. (Note it is host/dev code)
 */
__HOSTORDEV__ void augur_mcmc_mwg
      ( AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
      , double* llCorrect
      , void (*proposal)(AugurAux_t, AugurMod_t, AugurMod_t)
      , void (*swap)(AugurAux_t, AugurMod_t, AugurMod_t, int)
      , double (*objFn)(AugurAux_t, AugurMod_t) ) {
  proposal(aux, curr, prop);
  if (augur_mcmc_ar_idx(*llCorrect, aux, curr, prop, objFn)) {
    // Accept, so copy proposal theta into current theta (maintain invariant)
    swap(aux, curr, prop, 0);
  }
  else {
    // Reject, so copy current theta into proposal theta (maintain invariant)
    swap(aux, curr, prop, 1);
  }
}

/**
 * Gradient checking. (Call from host code)
 */
void h_augur_chk_grad
      ( AugurMemLoc_t loc, AugurBlk_t* pt, AugurBlk_t* grad
      , AugurAux_t aux, AugurMod_t curr
      , double (*objFn)(AugurAux_t, AugurMod_t) ) {
  double eps = 10e-8;
  printf("Evaluating gradient at point:\n");
  h_augur_blk_dump(loc, pt);

  for (uint_t i = 0; i < pt->num_blks; i++) {
    switch (pt->typs[i]) {
    case AUGUR_DBL: {
      // TODO
      break;
    }
    case AUGUR_VEC: {
      AugurFlatVec_t* fvec = (AugurFlatVec_t*) pt->blks[i];
      for (int j = 0; j < fvec->base_elems; j++) {
	h_augur_flat_vec_basis_add(loc, fvec, j, eps);
	double plus = objFn(aux, curr);
	h_augur_flat_vec_basis_add(loc, fvec, j, -2.0 * eps);
	double minus = objFn(aux, curr);
	double fin_diff = (plus - minus) / (2.0 * eps);
	double ad;
	h_augur_flat_vec_getd_idx(loc, &ad, (AugurFlatVec_t*) (grad->blks[i]), j);
	printf("plus: %f, minus: %f\n", plus, minus);
	printf("auto-diff: %f, finite-diff: %f, error: %f\n", ad, fin_diff, ad - fin_diff);
	h_augur_flat_vec_basis_add(loc, fvec, j, eps);
      }
      break;
    }
    case AUGUR_MAT: {
      AugurMat_t* mat = (AugurMat_t*) pt->blks[i];
      for (int j = 0; j < mat->row * mat->col; j++) {
	h_augur_mat_basis_add(loc, mat, j, eps);
	double plus = objFn(aux, curr);
	h_augur_mat_basis_add(loc, mat, j, -2.0 * eps);
	double minus = objFn(aux, curr);
	double fin_diff = (plus - minus) / (2.0 * eps);
	double ad;
	h_augur_mat_getd_idx(loc, &ad, (AugurMat_t*) (grad->blks[i]), j);
	printf("auto-diff: %f, finite-diff: %f, error: %f\n", ad, fin_diff, ad - fin_diff);
	h_augur_mat_basis_add(loc, mat, j, eps);
      }
      break;
    }
    default: {
      // TODO
      break;
    }
    }
  }
}

/**
 * Calls a gradient function. (Call from host code)
 *
 * The struct curr contains "the point" we are evaluating gradfn at.
 * The result of the gradient computation is in grad, which is part of
 * the struct aux.
 */
void h_augur_mcmc_call_grad
      ( AugurMemLoc_t loc
      , AugurAux_t aux, AugurMod_t curr
      , AugurBlk_t* grad
      , void (*gradfn)(AugurAux_t, AugurMod_t)
      ) {
  h_augur_blk_zero(loc, grad);
  gradfn(aux, curr);
}

/**
 * The struct curr contains pt, "the point" we are evaluating gradfn at.
 */
void h_augur_mcmc_chkcall_grad
      ( AugurMemLoc_t loc
      , AugurAux_t aux, AugurMod_t curr
      , AugurBlk_t* grad
      , void (*gradfn)(AugurAux_t, AugurMod_t)
      , AugurBlk_t* pt
      , double (*objfn) (AugurAux_t, AugurMod_t)
      , Bool_t f_chkgrad
      ) {
  h_augur_mcmc_call_grad(loc, aux, curr, grad, gradfn);
  if (f_chkgrad) {
    h_augur_chk_grad(loc, pt, grad, aux, curr, objfn);
  }
}

/**
 * Leapfrog simulation. (Call from host code)
 *
 * Leapfrog with (POSITIVE) gradient of log-likelihood of objfn.
 * (Signs are important!!)
 *
 * theta0 / thetaStar: initial / proposal point
 * grad: result of gradient
 * p0 / pStar: initial / proposal momentum
 *
 * mod.thetaStar = thetaStar
 * aux.grad = grad
 *
 * gradfn(aux, mod) evaluates gradient at thetaStar, putting result in grad
 *
 */
void h_augur_mcmc_leapfrog1
      ( AugurMemLoc_t loc
      , AugurAux_t aux, AugurMod_t mod
      , AugurBlk_t* grad, AugurBlk_t* thetaStar, AugurBlk_t* pStar
      , void (*gradfn)(AugurAux_t, AugurMod_t)
      , double stepsize
      ) {
  // Simulate Hamiltonian dynamics first half step of momentum
  // pStar = p0 + stepsize / 2.0 * d_like(x, thetaStar)
  h_augur_mcmc_call_grad(loc, aux, mod, grad, gradfn);
  h_augur_blk_scale_plus(loc, pStar, pStar, stepsize / 2.0, grad);

  // Full step for position
  // Position: thetaStar = thetaStar + stepsize * pStar
  h_augur_blk_scale_plus(loc, thetaStar, thetaStar, stepsize, pStar);

  // Simulate last half step of momentum
  // pStar = pStar + stepsize / 2.0 * d_like(x, thetaStar)
  h_augur_mcmc_call_grad(loc, aux, mod, grad, gradfn);
  h_augur_blk_scale_plus(loc, pStar, pStar, stepsize / 2.0, grad);
}

/**
 * Leapfrog simulation. (Call from host code)
 *
 * Leapfrog with (POSITIVE) gradient of log-likelihood of objfn.
 * (Signs are important!!)
 *
 * theta0 / thetaStar: initial / proposal point
 * grad: result of gradient
 * p0 / pStar: initial / proposal momentum
 *
 * curr.theta0 = theta0, prop.theta0 = thetaStar, aux.grad = grad
 *
 * gradfn(aux, curr) evaluates gradient at theta0, putting result in grad
 * gradfn(aux, prop) evaluates gradient at thetaStar, putting result in grad
 *
 */
void h_augur_mcmc_leapfrog
      ( AugurMemLoc_t loc
      , AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
      , AugurBlk_t* theta0, AugurBlk_t* thetaStar
      , AugurBlk_t* grad, AugurBlk_t* p0, AugurBlk_t* pStar
      , void (*gradfn)(AugurAux_t, AugurMod_t)
      , double (*objfn) (AugurAux_t, AugurMod_t)
      , double stepsize, int simsteps
      ) {
  // Simulate Hamiltonian dynamics first half step of momentum
  // pStar = p0 + stepsize / 2.0 * d_like(x, theta0)
  h_augur_mcmc_call_grad(loc, aux, curr, grad, gradfn);
  h_augur_blk_scale_plus(loc, pStar, p0, stepsize / 2.0, grad);

  // Simulate full steps
  for (uint_t l = 0; l < simsteps; l++) {
    // Full step for position
    // Position: thetaStar = thetaStar + stepsize * pStar
    h_augur_blk_scale_plus(loc, thetaStar, thetaStar, stepsize, pStar);

    // Move momentum except at last step
    if (l != simsteps-1) {
      // Momentum: pStar = pStar + stepsize * d_like(x, thetaStar)
      h_augur_mcmc_call_grad(loc, aux, prop, grad, gradfn);
      h_augur_blk_scale_plus(loc, pStar, pStar, stepsize, grad);
    }
  }

  // Simulate last half step of momentum
  // pStar = pStar + stepsize / 2.0 * d_like(x, thetaStar)
  h_augur_mcmc_call_grad(loc, aux, prop, grad, gradfn);
  h_augur_blk_scale_plus(loc, pStar, pStar, stepsize / 2.0, grad);
}

/**
 * Hamiltonian Monte Carlo Kernel. (Call from host code)
 *
 * Suppose this kernel moves blk = [p_1, .., p_n].
 * Then we have the following equalities on the inputs:
 *
 * curr.blk = theta0
 * prop.blk = thetaStar
 *
 * For compositionality, we maintain the invariant that:
 * curr.blk = prop.blk
 */
void h_augur_mcmc_hmc
     ( AugurMemLoc_t loc
     , AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
     , AugurBlk_t* theta0, AugurBlk_t* thetaStar
     , AugurBlk_t* grad, AugurBlk_t* p0, AugurBlk_t* pStar
     , double simlen, double eps
     , void (*computeGrad)(AugurAux_t, AugurMod_t)
     , double (*objFn)(AugurAux_t, AugurMod_t)
     ) {
  // Initialize momentum
  h_augur_blk_randn(aux.rng, p0);

  // Randomize simsteps
  double scale_eps = h_augur_uniform_sample(h_rng, 0.8, 1.2);
  double stepsize = scale_eps * eps;
  double s = h_augur_uniform_sample(h_rng, 0.9, 1.1);
  uint_t simsteps = augur_max(1, round(s * simlen / stepsize));

  // Simulate Hamiltonian dynamics with Leapfrog method
  h_augur_mcmc_leapfrog(loc, aux, curr, prop, theta0, thetaStar, grad, p0, pStar, computeGrad, objFn, stepsize, simsteps);

  // Compute correction
  double initMomLL = h_augur_blk_dot(loc, p0, p0) / 2.0;
  double propMomLL = h_augur_blk_dot(loc, pStar, pStar) / 2.0;
  double auxll = initMomLL - propMomLL;

  // After trajectory, compute whether we should accept or reject
  if (augur_mcmc_ar(auxll, objFn)) {
    // Accept, so copy proposal theta into current theta (for invariant)
    h_augur_blk_cpy(loc, theta0, thetaStar);
  }
  else {
    // Reject, so copy current theta into proposal theta (for invariant)
    h_augur_blk_cpy(loc, thetaStar, theta0);
  }
}

/**
 * Elliptical Slice Sampling Kernel. (Note it is host/dev code)
 *
 * Suppose this kernel moves param = currTheta.
 * Then we have the following equalities on the inputs:
 *
 * curr.param = currTheta
 * prop.param = propTheta
 *
 * For compositionality, we maintain the invariant that:
 * curr.param = prop.param
 */
__HOSTORDEV__ void augur_mcmc_eslice
      ( AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
      , AugurVec_t* curr_theta, AugurVec_t* prop_theta
      , AugurVec_t* mean, AugurMat_t* cov
      , AugurMat_t* L, AugurVec_t* nu
      , double (*objfn)(AugurAux_t, AugurMod_t) ) {
  augur_mvnormal_sample(aux.rng, nu, mean, cov, L, nu);
  double u = augur_std_uniform_sample(aux.rng);
  double log_y = objfn(aux, curr) + augur_log(u);

  double theta = augur_uniform_sample(aux.rng, 0.0, 2.0 * AUGUR_PI);
  double theta_min = theta - 2.0 * AUGUR_PI;
  double theta_max = theta;

  while (TRUE) {
    double diff = theta_max - theta_min;
    if (diff < 1e-8) {
      // Maintain invariant that curr == prop
      augur_vec_cpy(prop_theta, curr_theta);
      break;
    }

    // propSt = (currSt - mean) * cos(theta) + (nu - mean) * sin(theta) + mean;
    for (uint_t i = 0; i < prop_theta->elems; i++) {
      double mean_i = AUGUR_VEC_GETD(mean, i);
      double a = (AUGUR_VEC_GETD(curr_theta, i) - mean_i) * augur_cos(theta);
      double b = (AUGUR_VEC_GETD(nu, i) - mean_i) * augur_sin(theta);
      AUGUR_VEC_SETD(prop_theta, i, a + b + mean_i);
    }

    double prop_ll = objfn(aux, prop);
    if (prop_ll > log_y) {
      // Maintain invariant that curr == prop
      augur_vec_cpy(curr_theta, prop_theta);
      break;
    }

    // Shrink bracket
    if (theta < 0) {
      theta_min = theta;
    }
    else {
      theta_max = theta;
    }
    theta = augur_uniform_sample(aux.rng, theta_min, theta_max);
  }
}

/**
 * Reflective Slice Kernel. (Call from host code)
 *
 * Suppose this kernel moves blk = [p_1, .., p_n].
 * Then we have the following equalities on the inputs:
 *
 * curr.blk = theta0
 * prop.blk = thetaStar
 *
 * For compositionality, we maintain the invariant that:
 * curr.blk = prop.blk
 */
void h_augur_mcmc_refl_slice
     ( AugurMemLoc_t loc
     , AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
     , AugurBlk_t* theta0, AugurBlk_t* thetaProp
     , AugurBlk_t* grad, AugurBlk_t* p0
     , double simlen, double eps
     , void (*computeGrad)(AugurAux_t, AugurMod_t)
     , double (*objFn)(AugurAux_t, AugurMod_t)
     ) {
  // Initialize momentum
  h_augur_blk_randn(aux.rng, p0);

  // Compute slice level
  double z = objFn(aux, curr) - gsl_ran_exponential(h_rng, 1.0);
  int cond = 0;
  double stepsize = eps;
  uint_t simsteps = augur_max(1, round(simlen / stepsize));

  for (int i = 0 ; i < simsteps; i++) {
    // Take scaled-step in direction of momentum
    h_augur_blk_scale_plus(loc, thetaProp, thetaProp, stepsize, p0);

    cond = objFn(aux, prop) < z;
    if (cond) {
      // Reflect (outside version) if we leave slice

      // Compute grad at outside point
      h_augur_blk_zero(loc, grad);
      computeGrad(aux, prop);

      // Compute reflection and update direction
      double angle1 = h_augur_blk_dot(loc, p0, grad);
      double angle2 = h_augur_blk_dot(loc, grad, grad);
      double scale = angle1 / angle2;
      h_augur_blk_scale_plus(loc, p0, p0, -2.0 * scale, grad);

    }
  }

  if (!cond) {
    // Copy if we are in slice
    h_augur_blk_cpy(loc, theta0, thetaProp);
  }
}

/* NUTS */

typedef enum NutsDir {
  NUTS_MINUS = 0,
  NUTS_PLUS = 1,
} NutsDir_t;

typedef struct NutsRet {
  int n;
  int s;
} NutsRet_t;

typedef struct NutsRT {
  AugurBlk_t* work;
  AugurBlkStk_t* theta_prop_stk;
  AugurBlkStk_t* theta_stk;
  AugurBlkStk_t* mom_stk;
} NutsRT_t;

typedef struct NutsInfo {
  AugurAux_t aux;
  AugurMod_t mod;
  AugurBlk_t* grad;
  double (*modObjFn)(AugurAux_t, AugurMod_t);
  void (*computeGrad)(AugurAux_t, AugurMod_t);
  double stepsize;
  double logu;
} NutsInfo_t;

int h_augur_nuts_update_s(AugurMemLoc_t loc, int s, AugurBlk_t* work, AugurBlk_t* theta_plus, AugurBlk_t* mom_plus, AugurBlk_t* theta_minus, AugurBlk_t* mom_minus) {
  if (s == 1) {
    h_augur_blk_minus(loc, work, theta_plus, theta_minus);
    int ut1 = augur_indicator(h_augur_blk_dot(loc, work, mom_plus) >= 0.0);
    int ut2 = augur_indicator(h_augur_blk_dot(loc, work, mom_minus) >= 0.0);
    if (ut1 && ut2) {
      return 1;
    }
  }
  return 0;
}

/**
 * info.mod.theta = theta (gradient is always evaluated at latest value of theta)
 *
 * Invariants:
 * Suppose dir = PLUS, then:
 *
 * theta: latest value of theta^+
 * mom: latest value of mom^+
 */
NutsRet_t h_augur_nuts_build_tree
      ( AugurMemLoc_t loc, NutsInfo_t* info , NutsRT_t* rt
      , AugurBlk_t* shape
      , AugurBlk_t* theta, AugurBlk_t* mom, int j, NutsDir_t dir
      , int tab) {
  if (j == 0) {
    // Simulate (updates theta, mom)
    h_augur_mcmc_leapfrog1(loc, info->aux, info->mod,
			   info->grad, theta, mom,
			   info->computeGrad,
			   info->stepsize);

    // Update return value
    augur_blk_stk_push(loc, rt->theta_prop_stk, shape, theta);

    // Check stopping condition
    NutsRet_t ret;
    double ll = info->modObjFn(info->aux, info->mod) - 0.5 * h_augur_blk_dot(loc, mom, mom);
    ret.n = augur_indicator(info->logu <= ll);
    ret.s = augur_indicator(info->logu < MCMC.deltaMax + ll);

    return ret;
  }
  else {
    // Build left subtree
    NutsRet_t left = h_augur_nuts_build_tree(loc, info, rt, shape, theta, mom, j - 1, dir, tab+2);

    // Build right subtree (if we should keep going)
    if (left.s == 1) {
      // Save "left"-most
      if (j == 1) {
	augur_blk_stk_push(loc, rt->theta_stk, shape, theta);
	augur_blk_stk_push(loc, rt->mom_stk, shape, mom);
      }

      NutsRet_t right = h_augur_nuts_build_tree(loc, info, rt, shape, theta, mom, j - 1, dir, tab+2);

      // Propogate proposal with probability proportional to size of subtrees
      AugurBlk_t* prop1 = augur_blk_stk_pop(rt->theta_prop_stk);
      AugurBlk_t* prop2 = augur_blk_stk_pop(rt->theta_prop_stk);
      double bias = ((double) right.n) / (left.n + right.n);
      if (h_augur_bernoulli_sample(h_rng, bias) == 0) {
	augur_blk_stk_push(loc, rt->theta_prop_stk, shape, prop2);
      }
      else {
	augur_blk_stk_push(loc, rt->theta_prop_stk, shape, prop1);
      }
      augur_blk_stk_dump(loc, rt->theta_prop_stk);

      // Propogate tree size
      left.n += right.n;

      // Check u-turn
      AugurBlk_t* theta_minus; AugurBlk_t* mom_minus;
      AugurBlk_t* theta_plus; AugurBlk_t* mom_plus;
      switch (dir) {
      case NUTS_MINUS: {
	theta_plus = augur_blk_stk_peek(rt->theta_stk);
	mom_plus = augur_blk_stk_peek(rt->mom_stk);
	theta_minus = theta;
	mom_minus = mom;
	break;
      }
      case NUTS_PLUS: {
	theta_minus = augur_blk_stk_peek(rt->theta_stk);
	mom_minus = augur_blk_stk_peek(rt->mom_stk);
	theta_plus = theta;
	mom_plus = mom;
	break;
      }
      }
      left.s = h_augur_nuts_update_s(loc, right.s, rt->work, theta_plus, mom_plus, theta_minus, mom_minus);

      // Pop "left"-most
      if (j > 1) {
	augur_blk_stk_pop(rt->theta_stk);
	augur_blk_stk_pop(rt->mom_stk);
      }
    }

    return left;
  }
}

/**
 * No-U-Turn Kernel. (Call from host code)
 *
 * Suppose this kernel moves blk = [p_1, .., p_n].
 * Then we have the following equalities on the inputs:
 *
 * curr.blk = theta0
 * prop.blk = thetaStar
 *
 * For compositionality, we maintain the invariant that:
 * curr.blk = prop.blk
 */
void h_augur_mcmc_nuts
      ( AugurMemLoc_t loc, AugurAux_t aux, AugurMod_t curr, AugurMod_t prop
      , AugurBlk_t* theta0, AugurBlk_t* theta_prop
      , AugurBlk_t* grad, AugurBlk_t* work
      , AugurBlk_t* theta_minus, AugurBlk_t* mom_minus
      , AugurBlk_t* theta_plus, AugurBlk_t* mom_plus
      , AugurBlk_t* shape
      , double stepsize
      , void (*computeGrad)(AugurAux_t, AugurMod_t)
      , double (*modObjFn)(AugurAux_t, AugurMod_t)
      ) {
  // Initialize simulation values
  h_augur_blk_cpy(loc, theta_minus, theta0);
  h_augur_blk_randn(aux.rng, mom_minus);
  h_augur_blk_cpy(loc, theta_plus, theta0);
  h_augur_blk_cpy(loc, mom_plus, mom_minus);

  // Initialize info
  NutsInfo_t info;
  info.aux = aux;
  info.mod = prop;
  info.grad = grad;
  info.modObjFn = modObjFn;
  info.computeGrad = computeGrad;
  info.logu = modObjFn(aux, curr) - 0.5 * h_augur_blk_dot(loc, mom_minus, mom_minus) - gsl_ran_exponential(h_rng, 1.0);

  // Initialize auxilliary nuts runtime values
  NutsRT_t rt;
  rt.work = work;
  AugurBlkStk_t theta_prop_stk = augur_blk_stk_stk_alloc();
  AugurBlkStk_t theta_stk = augur_blk_stk_stk_alloc();
  AugurBlkStk_t mom_stk = augur_blk_stk_stk_alloc();
  rt.theta_prop_stk = &theta_prop_stk;
  rt.theta_stk = &theta_stk;
  rt.mom_stk = &mom_stk;

  int n = 1;
  int j = 0;
  int s = 1;

  while (TRUE) {
    NutsRet_t nuts_ret;
    augur_blk_stk_reset(rt.theta_stk);
    augur_blk_stk_reset(rt.mom_stk);
    if (h_augur_bernoulli_sample(h_rng, 0.5)) {
      info.stepsize = -stepsize;
      h_augur_blk_cpy(loc, theta_prop, theta_minus);
      nuts_ret = h_augur_nuts_build_tree(loc, &info, &rt, shape, theta_prop, mom_minus, j, NUTS_MINUS, 0);
      h_augur_blk_cpy(loc, theta_minus, theta_prop);
    }
    else {
      info.stepsize = stepsize;
      h_augur_blk_cpy(loc, theta_prop, theta_plus);
      nuts_ret = h_augur_nuts_build_tree(loc, &info, &rt, shape, theta_prop, mom_plus, j, NUTS_PLUS, 0);
      h_augur_blk_cpy(loc, theta_plus, theta_prop);
    }

    AugurBlk_t* theta_prop_p = augur_blk_stk_pop(rt.theta_prop_stk);
    Bool_t f_cpy = FALSE;
    if (nuts_ret.s == 1) {
      double bias = augur_min(1.0, ((double) nuts_ret.n) / n);
      if (h_augur_bernoulli_sample(h_rng, bias) == 1) {
	// Update proposal
	h_augur_blk_cpy(loc, theta_prop, theta_prop_p);
	f_cpy = TRUE;
      }
    }

    n += nuts_ret.n;
    s = h_augur_nuts_update_s(loc, nuts_ret.s, rt.work, theta_plus, mom_plus, theta_minus, mom_minus);
    j += 1;

    if (s != 1) {
      // Maintain invariant
      if (f_cpy) {
	h_augur_blk_cpy(loc, theta0, theta_prop);
      }
      else {
	h_augur_blk_cpy(loc, theta_prop, theta0);
      }
      break;
    }
  }
}

#endif


/*
ODE: a program to get optime Runge-Kutta and multi-steps methods.

Copyright 2011-2019, Javier Burguete Tolosa.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

	1. Redistributions of source code must retain the above copyright notice,
		this list of conditions and the following disclaimer.

	2. Redistributions in binary form must reproduce the above copyright notice,
		this list of conditions and the following disclaimer in the
		documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY Javier Burguete Tolosa ``AS IS'' AND ANY EXPRESS OR
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
SHALL Javier Burguete Tolosa OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/**
 * \file rk_6_2.c
 * \brief Source file to optimize Runge-Kutta 6 steps 2nd order methods.
 * \author Javier Burguete Tolosa.
 * \copyright Copyright 2011-2019.
 */
#define _GNU_SOURCE
#include <string.h>
#include <math.h>
#include <libxml/parser.h>
#include <glib.h>
#include <libintl.h>
#include <gsl/gsl_rng.h>
#include "config.h"
#include "utils.h"
#include "optimize.h"
#include "rk.h"
#include "rk_6_2.h"

#define DEBUG_RK_6_2 0          ///< macro to debug.

/**
 * Function to obtain the coefficients of a 6 steps 2nd order Runge-Kutta
 * method.
 */
int
rk_tb_6_2 (Optimize * optimize) ///< Optimize struct.
{
  long double *tb, *r;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  b21 (tb) = r[2];
  t3 (tb) = r[3];
  b31 (tb) = r[4];
  b32 (tb) = r[5];
  t4 (tb) = r[6];
  b41 (tb) = r[7];
  b42 (tb) = r[8];
  b43 (tb) = r[9];
  t5 (tb) = r[10];
  b51 (tb) = r[11];
  b52 (tb) = r[12];
  b53 (tb) = r[13];
  b54 (tb) = r[14];
  b62 (tb) = r[15];
  b63 (tb) = r[16];
  b64 (tb) = r[17];
  b65 (tb) = r[18];
  b61 (tb) = (0.5L - b62 (tb) * t2 (tb) - b63 (tb) * t3 (tb)
              - b64 (tb) * t4 (tb) - b65 (tb) * t5 (tb)) / t1 (tb);
  if (isnan (b61 (tb)))
    return 0;
  rk_b_6 (tb);
#if DEBUG_RK_6_2
  rk_print_tb (optimize, "rk_tb_6_2", stderr);
  fprintf (stderr, "rk_tb_6_2: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 2nd order, 3rd order in
 * equations depending only in time, Runge-Kutta method.
 */
int
rk_tb_6_2t (Optimize * optimize)        ///< Optimize struct.
{
  long double *tb, *r;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2t: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  b21 (tb) = r[2];
  t3 (tb) = r[3];
  b31 (tb) = r[4];
  b32 (tb) = r[5];
  t4 (tb) = r[6];
  b41 (tb) = r[7];
  b42 (tb) = r[8];
  b43 (tb) = r[9];
  t5 (tb) = r[10];
  b51 (tb) = r[11];
  b52 (tb) = r[12];
  b53 (tb) = r[13];
  b54 (tb) = r[14];
  b61 (tb) = r[15];
  b62 (tb) = r[16];
  b63 (tb) = r[17];
  b64 (tb) = (1.L / 3.L - 0.5L * t5 (tb)
              - b61 (tb) * t1 (tb) * (t1 (tb) - t5 (tb))
              - b62 (tb) * t2 (tb) * (t2 (tb) - t5 (tb))
              - b63 (tb) * t3 (tb) * (t3 (tb) - t5 (tb)))
    / (t4 (tb) * (t4 (tb) - t5 (tb)));
  if (isnan (b64 (tb)))
    return 0;
  b65 (tb) = (0.5L - b61 (tb) * t1 (tb) - b62 (tb) * t2 (tb)
              - b63 (tb) * t3 (tb) - b64 (tb) * t4 (tb)) / t5 (tb);
  if (isnan (b65 (tb)))
    return 0;
  rk_b_6 (tb);
#if DEBUG_RK_6_2
  rk_print_tb (optimize, "rk_tb_6_2t", stderr);
  fprintf (stderr, "rk_tb_6_2t: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 1st-2nd order Runge-Kutta
 * pair.
 */
int
rk_tb_6_2p (Optimize * optimize)        ///< Optimize struct.
{
  long double *tb;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2p: start\n");
#endif
  if (!rk_tb_6_2 (optimize))
    return 0;
  tb = optimize->coefficient;
  e61 (tb) = e62 (tb) = e63 (tb) = e64 (tb) = 0.L;
  rk_e_6 (tb);
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2p: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 1st-2nd order, 1st-3rd order
 * in equations depending only in time, Runge-Kutta pair.
 */
int
rk_tb_6_2tp (Optimize * optimize)       ///< Optimize struct.
{
  long double *tb;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2tp: start\n");
#endif
  if (!rk_tb_6_2t (optimize))
    return 0;
  tb = optimize->coefficient;
  e61 (tb) = e62 (tb) = e63 (tb) = e64 (tb) = 0.L;
  rk_e_6 (tb);
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_tb_6_2tp: end\n");
#endif
  return 1;
}

/**
 * Function to calculate the objective function of a 6 steps 2nd order
 * Runge-Kutta method.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_2 (RK * rk)   ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_objective_tb_6_2: start\n");
#endif
  tb = rk->tb->coefficient;
  o = fminl (0.L, b20 (tb));
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b61 (tb) < 0.L)
    o += b61 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_objective_tb_6_2: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_2: end\n");
#endif
  return o;
}

/**
 * Function to calculate the objective function of a 6 steps 2nd order, third
 * order in equations depending only on time, Runge-Kutta method.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_2t (RK * rk)  ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_objective_tb_6_2t: start\n");
#endif
  tb = rk->tb->coefficient;
#if DEBUG_RK_6_2
  rk_print_tb (optimize, "rk_objective_tb_6_2t", stderr);
#endif
  o = fminl (0.L, b20 (tb));
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b64 (tb) < 0.L)
    o += b64 (tb);
  if (b65 (tb) < 0.L)
    o += b65 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_2
  fprintf (stderr, "rk_objective_tb_6_2t: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_2t: end\n");
#endif
  return o;
}


/* specfunc/elementary.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include "gsl_sf_elementary.h"

#include "error.h"
#include "check.h"

int
gsl_sf_multiply_e(const double x, const double y, gsl_sf_result * result)
{
  const double ax = fabs(x);
  const double ay = fabs(y);

  if(x == 0.0 || y == 0.0) {
    /* It is necessary to eliminate this immediately.
     */
    result->val = 0.0;
    result->err = 0.0;
    return GSL_SUCCESS;
  }
  else if((ax <= 1.0 && ay >= 1.0) || (ay <= 1.0 && ax >= 1.0)) {
    /* Straddling 1.0 is always safe.
     */
    result->val = x*y;
    result->err = 2.0 * GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
  else {
    const double f = 1.0 - 2.0 * GSL_DBL_EPSILON;
    const double min = GSL_MIN_DBL(fabs(x), fabs(y));
    const double max = GSL_MAX_DBL(fabs(x), fabs(y));
    if(max < 0.9 * GSL_SQRT_DBL_MAX || min < (f * DBL_MAX)/max) {
      result->val = GSL_COERCE_DBL(x*y);
      result->err = 2.0 * GSL_DBL_EPSILON * fabs(result->val);
      CHECK_UNDERFLOW(result);
      return GSL_SUCCESS;
    }
    else {
      OVERFLOW_ERROR(result);
    }
  }
}

int
gsl_sf_multiply_err_e(const double x, const double dx,
                         const double y, const double dy,
                         gsl_sf_result * result)
{
  int status = gsl_sf_multiply_e(x, y, result);
  result->err += fabs(dx*y) + fabs(dy*x);
  return status;
}

/*-*-*-*-*-*-*-*-*-* Functions w/ Natural Prototypes *-*-*-*-*-*-*-*-*-*-*/

#include "eval.h"

double gsl_sf_multiply(const double x, const double y)
{
  EVAL_RESULT(gsl_sf_multiply_e(x, y, &result));
}


#include <viaio/VImage.h>
#include <math.h>
#include <string.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_math.h>
#include "gsl_utils.h"

#define ABS(x) ((x) >= 0) ? (x) : -(x)
#define MIN(a,b) ((a) < (b)) ? (a) : (b)
#define MAX(a,b) ((a) > (b)) ? (a) : (b)
#define ROUND(x) (int)(((x) >= 0) ? ((x) + 0.5) : ((x) - 0.5))

void whitecov2(VImage effect_image,VImage rho_vol,
	       gsl_matrix_float* Y, gsl_matrix_float* X, gsl_matrix_float* con,
	       VFloat* Dfs, int numlags, int slice)
{
  int numcon = 1;
  int n = Y->size1;
  VFloat Df = Dfs[numcon];

  int i=0,j=0,k=0;

  /* some checks and warnings */
  if(X->size2 != n)
    VError("whitecov2: Warning! sizes of design and data matrix do not match, n=%d, %ld",n,X->size2);

  if(X->size1 != con->size2)
    VError("whitecov2: Warning! contrast and design do not match.");

  gsl_vector_float* irho = gsl_vector_float_alloc(VImageNBands(rho_vol));
  if(numlags == 1) {
    float* p = irho->data;
    const float drho = 0.01;
    for(i=0;i<irho->size;i++) {
      float tmp = ROUND((VPixel(rho_vol, i, slice, 0, VFloat) / drho)) * drho;
      tmp = MIN(tmp, (1-drho));
      tmp = MAX(tmp, (-1+drho));
      *p++ = tmp;
    }
  }
  else {
    float* p = irho->data;
    for(i=1;i<=irho->size;i++)
      *p++ = i;
  }

  /* : X=X';  */
  gsl_matrix_float* transX = gsl_matrix_float_alloc(X->size2, X->size1);
  gsl_matrix_float_transpose_memcpy(transX, X);

  /* : Xstar=X; */
  gsl_matrix_float* XStar = gsl_matrix_float_alloc(transX->size1, transX->size2);
  gsl_matrix_float_memcpy(XStar, transX);

  /* END allocate memory buffer */

  /****************************** order == 1 *****************************/
  if (numlags == 1) {

    /* allocate memory buffers */
    gsl_vector_float* pixBig = gsl_vector_float_alloc(irho->size);
    gsl_matrix_float* pinvXStar = gsl_matrix_float_alloc(XStar->size2, XStar->size1);
    gsl_matrix_float* V = gsl_matrix_float_alloc(pinvXStar->size1, pinvXStar->size1);
    gsl_matrix_float* buffer_01 = gsl_matrix_float_alloc(con->size1,V->size2);
    gsl_matrix_float* buffer_02 = gsl_matrix_float_alloc(con->size1,con->size1);
    gsl_matrix_float* cVcinv = gsl_matrix_float_alloc(con->size1, con->size1);

    gsl_vector_float* u = funique(irho);
    float* rho = u->data;
    for(i=0;i<u->size;i++) {

      /* :  pix=int16(find(irho==rho)); */
      float* p = irho->data;
      float* ppB = pixBig->data;
      int counter = 0;
      for(j=0;j<irho->size;j++) {
	if(*p++ == *rho) {
	  *ppB++ = (float)j;
	  counter++;
	}
      }
      gsl_vector_float* pix = gsl_vector_float_alloc(counter);
      ppB = pixBig->data;
      float* pp = pix->data;
      for(j=0;j<counter;j++) *pp++ = *ppB++;

      /* : Ystar=Y(:,pix); */
      gsl_matrix_float* YStar = fmat_subcols(Y,pix);

      /* : factor=1./sqrt(1-rho^2); */
      float factor = pow(1 - *rho * *rho, -0.5);

      /* : Ystar(k1,:)=(Y(k1,pix)-rho*Y(k1-1,pix))*factor; */
      /* point to last element */
      float* pStar = YStar->data+(YStar->size1*YStar->size2)-1;
      /* traverse matrix in reverse order */
      for (j = 0; j < (YStar->size1-1)*YStar->size2; ++j) {
	*pStar = (*pStar - *rho * *(pStar - YStar->size2)) * factor;
	pStar--;
      }

      /* : Xstar(k1,:)=(X(k1,:)-rho*X(k1-1,:))*factor; */
      pStar = XStar->data+XStar->size2;
      float* pX = transX->data+transX->size2;
      for (j = 0; j < (XStar->size1-1)*XStar->size2; ++j) {
	*pStar++ = (*pX - *rho * *(pX - transX->size2)) * factor;
	pX++;
      }

      /* : pinvXstar=pinv(Xstar); */
      fmat_PseudoInv(XStar,pinvXStar);

      /* : betahat=pinvXstar*Ystar; */
      gsl_matrix_float* betahat = fmat_x_mat(pinvXStar, YStar,NULL);

      /* : resid=Ystar-Xstar*betahat; */
      gsl_matrix_float* buffer = fmat_x_mat(XStar, betahat, NULL);
      gsl_matrix_float_sub(YStar, buffer);
      gsl_matrix_float* resid =  YStar;
      gsl_matrix_float_free(buffer);

      /* : SSE=sum(resid.^2,1); */
      buffer = gsl_matrix_float_alloc(resid->size1, resid->size2);
      gsl_matrix_float_memcpy (buffer, resid);
      gsl_matrix_float_mul_elements (buffer, buffer);
      gsl_vector_float* sse = fsum(buffer, 1,NULL);
      gsl_matrix_float_free(buffer);

      /* : sd=sqrt(SSE/Df); */
      gsl_vector_float* sd = gsl_vector_float_alloc(sse->size);
      float* pSSE = sse->data;
      float* pSD = sd->data;
      for (j = 0; j < sse->size; ++j) {
	*pSD++ = (float)sqrt(*pSSE++/Df);
      }

      /* :   V=pinvXstar*pinvXstar'; */
      fmat_x_matT(pinvXStar, pinvXStar, V);

      /* : mag_ef=contrast*betahat; */
      gsl_matrix_float* mag_ef = fmat_x_mat(con, betahat, NULL);

      /* : mag_sd=sqrt(diag(contrast*V*contrast'))*sd; */
      fmat_x_mat(con,V,buffer_01);
      fmat_x_matT(buffer_01, con, buffer_02);

      gsl_vector_float_view diag = gsl_matrix_float_diagonal(buffer_02);
      gsl_matrix_float* mag_sd = gsl_matrix_float_alloc(diag.vector.size, sd->size);
      float* pMag = mag_sd->data;
      for (j = 0; j < diag.vector.size; ++j) {
	float s = sqrt(gsl_vector_float_get(&diag.vector,j));
	pSD = sd->data;
	for (k = 0; k < sd->size; ++k) {
	  *pMag++ = s * (*pSD++);
	}
      }

      /* effect_slice(pix,1)= (mag_ef./(mag_sd+(mag_sd<=0)).*(mag_sd>0))'; */
      VFloat* peff = VPixelPtr(effect_image,slice,0,0);
      float* pmeff = mag_ef->data;
      float* pmsd = mag_sd->data;
      float* ppix = pix->data;
      for (j=0; j<pix->size; ++j) {
	*(peff+(int)*ppix++) = *pmeff / (*pmsd+(*pmsd<=0)) * (*pmsd>0);
	pmeff++; pmsd++;
      }

      /* free some memory */
      gsl_matrix_float_free(betahat);
      gsl_matrix_float_free(YStar);
      gsl_matrix_float_free(mag_ef);
      gsl_matrix_float_free(mag_sd);
      gsl_vector_float_free(pix);
      gsl_vector_float_free(sse);
      gsl_vector_float_free(sd);

      rho++;
    } /* END for(i=0;i<u->size;i++)  */

    /* free some memory */
    gsl_vector_float_free(pixBig);
    gsl_matrix_float_free(pinvXStar);
    gsl_matrix_float_free(V);
    gsl_matrix_float_free(buffer_01);
    gsl_matrix_float_free(buffer_02);
    gsl_matrix_float_free(cVcinv);

  } /******************* END order == 1 **********************************/

/*************************** order > 1 *********************************/
    else{

      /* allocate memory buffer */
      gsl_vector_float* Coradj_pix = gsl_vector_float_alloc(VImageNColumns(rho_vol)+1);
      gsl_matrix_float* Ainv = gsl_matrix_float_alloc(Coradj_pix->size, Coradj_pix->size);
      gsl_matrix* dbuff = gsl_matrix_alloc(Ainv->size1,Ainv->size2);
      gsl_matrix_float* A = gsl_matrix_float_alloc(Ainv->size1,Ainv->size2);
      int nl = Ainv->size2;
      gsl_matrix_float* buffer_01 = gsl_matrix_float_alloc(n-nl,1);
      gsl_matrix_float* buffer_02 = gsl_matrix_float_alloc(A->size1,1);
      gsl_matrix_float* buffer_03 = gsl_matrix_float_alloc(A->size1,transX->size2);
      gsl_matrix_float* B = gsl_matrix_float_alloc(buffer_01->size1, A->size2);
      gsl_matrix_float* Vmhalf = gsl_matrix_float_alloc(n-nl,n);
      gsl_matrix_float* buffer_04 = gsl_matrix_float_alloc(Vmhalf->size1, transX->size2);
      gsl_matrix_float* YStar = gsl_matrix_float_alloc(n,1);
      gsl_matrix_float* pinvXStar = gsl_matrix_float_alloc(XStar->size2, XStar->size1);
      gsl_matrix_float* betahat = gsl_matrix_float_alloc(pinvXStar->size1, YStar->size2);
      gsl_matrix_float* buffer_05 = gsl_matrix_float_alloc(XStar->size1, betahat->size2);
      gsl_vector_float* sse = gsl_vector_float_alloc(1);
      gsl_vector_float* sd = gsl_vector_float_alloc(YStar->size2);
      gsl_matrix_float* V = gsl_matrix_float_alloc(pinvXStar->size1, pinvXStar->size1);
      gsl_matrix_float* mag_ef = gsl_matrix_float_alloc(con->size1, betahat->size2);
      gsl_matrix_float* buffer_06 = gsl_matrix_float_alloc(con->size1, V->size2);
      gsl_matrix_float* buffer_07 = gsl_matrix_float_alloc(con->size1, con->size1);
      gsl_matrix_float* mag_sd = gsl_matrix_float_alloc(con->size1, sd->size);
      gsl_matrix_float* buffer_08 = gsl_matrix_float_alloc(mag_ef->size1, mag_ef->size2);
      gsl_vector_float* sst = gsl_vector_float_alloc(1);
      /* END allocate memory buffer */

     /* we know that irho only contains a list of indices for every voxel so we can skip
        the unique call and iterate over irho itself */
      float* rho = irho->data;
      for(i=0;i<irho->size;i++) {

        /* : Coradj_pix=squeeze(rho_vol(pix,slice,:)); */
        float* pCoradj = Coradj_pix->data+1;
        for(k=0;k<Coradj_pix->size-1;k++) {
          *(pCoradj++) = VPixel(rho_vol,
                                i,
                                slice,
                                k,
                                VFloat);
        }

        /* : [Ainvt posdef]=chol(toeplitz([1 Coradj_pix'])); */

        /* at first the toeplitz matrix */
        Coradj_pix->data[0] = 1;
        fmat_toeplitz(Coradj_pix, Ainv);
        /* the cholesky decomposition */
        /* double buffer */
        for(j=0;j<Ainv->size1;j++) {
          for(k=0;k<Ainv->size2;k++) {
            gsl_matrix_set(dbuff,j,k,(double)gsl_matrix_float_get(Ainv,j,k));
          }
        }

        if(gsl_linalg_cholesky_decomp(dbuff) == GSL_EDOM)
          VError("Calculation error, cholesky decomposition failed for pix=%d",i);

        /* note: the gsl cholesky factorization returns the inverse
         * of Ainvt in the lower triangular part of it's output matrix.
         * Since we will use the inverse matrix of Ainvt in the next steps we
         * will work with the lower triangular part in contrast to the matlab algorithm */

        /* convert back to float and remove upper triangular matrix */
        for(j=0;j<Ainv->size1;j++) {
          for(k=0;k<Ainv->size2;k++) {
            if(j<k) {
              gsl_matrix_float_set(Ainv,j,k,0);
            }
            else {
              gsl_matrix_float_set(Ainv,j,k,(float)gsl_matrix_get(dbuff,j,k));
            }
          }
        }

        /* :  A=inv(Ainvt'); */
        fInv(Ainv, A);

        /* : B=ones(n-nl,1)*A(nl,:); */
	gsl_matrix_float_set_all(buffer_01,1);
        gsl_matrix_float_view subm = gsl_matrix_float_submatrix(A,nl-1,0,1,A->size2);
        fmat_x_mat(buffer_01,&subm.matrix,B);

        /* : Vmhalf=spdiags(double(B),double(1:nl),double(n-nl),double(n)); */
        gsl_matrix_float_set_zero(Vmhalf);
        for(j=0;j<B->size1;j++) {
          for(k=0;k<B->size2;k++) {
            gsl_matrix_float_set(Vmhalf,j,k+1+j,gsl_matrix_float_get(B,j,k));
          }
        }

        /* : Ystar=single(zeros(n,1)); */
        gsl_matrix_float_set_zero(YStar);

        /* : Ystar(1:nl)=A*Y(1:nl,pix); */
        subm = gsl_matrix_float_submatrix(Y,0,i,nl,1);
        fmat_x_mat(A,&subm.matrix,buffer_02);
        for(j=0;j<nl;j++) {
          gsl_matrix_float_set(YStar,j,0,gsl_matrix_float_get(buffer_02,j,0));
        }

	/* : Ystar((nl+1):n)=single(Vmhalf*double(Y(:,pix))); */
	subm = gsl_matrix_float_submatrix(Y,0,i,Y->size1,1);
	fmat_x_mat(Vmhalf,&subm.matrix,buffer_01);
	for(j=nl;j<n;j++) {
	  gsl_matrix_float_set(YStar,j,0,gsl_matrix_float_get(buffer_01,j-nl,0));
	}

        /* : Xstar(1:nl,:)=A*X(1:nl,:); */
        subm = gsl_matrix_float_submatrix(transX,0,0,nl,transX->size2);
        fmat_x_mat(A, &subm.matrix, buffer_03);
        for(j=0;j<buffer_03->size1;j++) {

          for(k=0;k<buffer_03->size2;k++) {
            gsl_matrix_float_set(XStar,j,k,gsl_matrix_float_get(buffer_03,j,k));
          }
        }

        /* : Xstar((nl+1):n,:)=single(Vmhalf*double(X)); */
        fmat_x_mat(Vmhalf,transX,buffer_04);
        for(j=nl;j<n;j++) {
          for(k=0;k<transX->size2;k++) {
            gsl_matrix_float_set(XStar,j,k,gsl_matrix_float_get(buffer_04,j-nl,k));
          }
        }

        /* : pinvXstar=pinv(Xstar); */
        fmat_PseudoInv(XStar,pinvXStar);

        /* : betahat=pinvXstar*Ystar; */
        fmat_x_mat(pinvXStar, YStar,betahat);

        /* : resid=Ystar-Xstar*betahat; */
        fmat_x_mat(XStar, betahat, buffer_05);
        gsl_matrix_float_sub(YStar, buffer_05);
        gsl_matrix_float* resid =  YStar;

        /* : SSE=sum(resid.^2,1); */
        gsl_matrix_float_mul_elements (resid, resid);
        fsum(resid, 1,sse);

        /* : sd=sqrt(SSE/Df); */
        float* pSSE = sse->data;
        float* pSD = sd->data;
        for (j = 0; j < sse->size; ++j) {
          *pSD++ = (float)sqrt(*pSSE++/Df);
        }

        /* :   V=pinvXstar*pinvXstar'; */
        fmat_x_matT(pinvXStar, pinvXStar, V);

        /* : mag_ef=contrast*betahat; */
        fmat_x_mat(con, betahat, mag_ef);

        /* : mag_sd=sqrt(diag(contrast*V*contrast'))*sd; */
        fmat_x_mat(con,V,buffer_06);
        fmat_x_matT(buffer_06, con, buffer_07);

        gsl_vector_float_view diag = gsl_matrix_float_diagonal(buffer_07);
        float* pMag = mag_sd->data;
        for (j = 0; j < diag.vector.size; ++j) {
          float s = sqrt(gsl_vector_float_get(&diag.vector,j));
          pSD = sd->data;
          for (k = 0; k < sd->size; ++k) {
            *pMag++ = s * *pSD++;
          }
        }

	/* effect_slice(pix,1)= (mag_ef./(mag_sd+(mag_sd<=0)).*(mag_sd>0))'; */
	VFloat* peff = VPixelPtr(effect_image, slice, 0, 0);
	float* pmeff = mag_ef->data;
	float* pmsd = mag_sd->data;
	*(peff+i) = *pmeff / (*pmsd+(*pmsd<=0)) * (*pmsd>0);

        /* next value */
        rho++;
      } /* END for(i=0;i<irho->size;i++) */

      /* free some memory */
      gsl_vector_float_free(Coradj_pix);
      gsl_matrix_float_free(Ainv);
      gsl_matrix_free(dbuff);
      gsl_matrix_float_free(A);
      gsl_matrix_float_free(buffer_01);
      gsl_matrix_float_free(buffer_02);
      gsl_matrix_float_free(buffer_03);
      gsl_matrix_float_free(buffer_04);
      gsl_matrix_float_free(B);
      gsl_matrix_float_free(Vmhalf);
      gsl_matrix_float_free(YStar);
      gsl_matrix_float_free(pinvXStar);
      gsl_matrix_float_free(buffer_05);
      gsl_vector_float_free(sd);
      gsl_matrix_float_free(V);
      gsl_matrix_float_free(mag_ef);
      gsl_matrix_float_free(buffer_06);
      gsl_matrix_float_free(buffer_07);
      gsl_matrix_float_free(mag_sd);
      gsl_matrix_float_free(buffer_08);
      gsl_vector_float_free(sse);
      gsl_vector_float_free(sst);

    }/************************************** END order > 1 ***************************************/

  /* :
   * minus = ones(size(effect_slice),'single');
   * minus(find(effect_slice<0))=single(-1.0);
   * effect_slice=single(sqrt(Df .* log(1 + effect_slice.^2 ./ Df) .* (1 - 0.5 ./ Df)));
   * effect_slice=effect_slice .* minus;
   */
  double u=0;
  VFloat* peff = VPixelPtr(effect_image, slice,0,0);
  for (i = 0; i < VImageNColumns(effect_image) * VImageNRows(effect_image); ++i) {
    VFloat oldVal = *peff;
    Df = Dfs[0];
    u = (double)(*peff);
    u = sqrt(Df * log(1.0 + pow(u,2.0) / Df) * (1 - 0.5 / Df));
    u = (oldVal < 0) ? u * (-1.0) : u;
    if (gsl_isnan(u) || gsl_isinf(u)) u = 0;
    (*peff) = (float)u;
    peff++;
  }

  /* free allocated memory */
  gsl_matrix_float_free(XStar);
  gsl_matrix_float_free(transX);
  /* END free allocated memory */
}


/* Copyright (c) 2011-2012, Jérémy Fix. All rights reserved. */

/* Redistribution and use in source and binary forms, with or without */
/* modification, are permitted provided that the following conditions are met: */

/* * Redistributions of source code must retain the above copyright notice, */
/* this list of conditions and the following disclaimer. */
/* * Redistributions in binary form must reproduce the above copyright notice, */
/* this list of conditions and the following disclaimer in the documentation */
/* and/or other materials provided with the distribution. */
/* * None of the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. */

/* THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND */
/* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED */
/* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE */
/* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE */
/* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL */
/* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR */
/* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER */
/* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, */
/* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE */
/* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */

#ifndef UKF_PARAMETER_NDIM_H
#define UKF_PARAMETER_NDIM_H

#include <gsl/gsl_linalg.h> // For the Cholesky decomposition
#include <gsl/gsl_math.h>
#include <gsl/gsl_blas.h>

#include "ukf_types.h"

namespace ukf
{
    namespace parameter
    {

        /**
          * @short Allocation of the vectors/matrices and initialization
          *
          */
        void ukf_init(ukf_param &p, ukf_state &s)
        {
            // Init the lambda
            p.lambda = p.alpha * p.alpha * (p.n + p.kpa) - p.n;
            p.gamma = sqrt(p.n + p.lambda);
            p.nbSamples = 2 * p.n + 1;

            // Init the matrices used to iterate
            s.Kk = gsl_matrix_alloc(p.n,p.no); // Kalman gain
            gsl_matrix_set_zero(s.Kk);

            s.Kk_T = gsl_matrix_alloc(p.no,p.n);
            gsl_matrix_set_zero(s.Kk_T);

            s.Pwdk = gsl_matrix_alloc(p.n,p.no);
            gsl_matrix_set_zero(s.Pwdk);

            // Whatever the type of evolution noise, its covariance is set to evolution_noise
            s.Prrk = gsl_matrix_alloc(p.n,p.n);
            p.evolution_noise->init(p,s);

            // Whatever the type of observation noise, its covariance is set to observation_noise
            s.Peek = gsl_matrix_alloc(p.no,p.no);
            gsl_matrix_set_identity(s.Peek);
            gsl_matrix_scale(s.Peek, p.observation_noise);

            s.Pddk = gsl_matrix_alloc(p.no, p.no); // Covariance of the output
            gsl_matrix_set_zero(s.Pddk);

            s.w = gsl_vector_alloc(p.n); // Parameter vector
            gsl_vector_set_zero(s.w);

            s.wk = gsl_vector_alloc(p.n); // Vector holding one sigma point
            gsl_vector_set_zero(s.wk);

            s.Pk = gsl_matrix_alloc(p.n,p.n); // Covariance matrix
            gsl_matrix_set_identity(s.Pk);
            gsl_matrix_scale(s.Pk,p.prior_pi);

            s.Sk = gsl_matrix_alloc(p.n,p.n); // Matrix holding the cholesky decomposition of Pk
            // Initialize Sk to the cholesky decomposition of Pk
            gsl_matrix_memcpy(s.Sk, s.Pk);
            gsl_linalg_cholesky_decomp(s.Sk);
            // Set all the elements of Lpi strictly above the diagonal to zero
            for(int k = 0 ; k < p.n ; k++)
                for(int j = 0 ; j < k ; j++)
                    gsl_matrix_set(s.Sk,j,k,0.0);

            s.cSk = gsl_vector_alloc(p.n); // Vector holding one column of Lpi
            gsl_vector_set_zero(s.cSk);

            s.wm  = gsl_vector_alloc(p.nbSamples); // Weights used to compute the mean of the sigma points images
            s.wc = gsl_vector_alloc(p.nbSamples); // Weights used to update the covariance matrices

            // Set the weights
            gsl_vector_set(s.wm, 0, p.lambda / (p.n + p.lambda));
            gsl_vector_set(s.wc, 0, p.lambda / (p.n + p.lambda) + (1.0 - p.alpha*p.alpha + p.beta));
            for(int j = 1 ; j < p.nbSamples; j ++)
            {
                gsl_vector_set(s.wm, j, 1.0 / (2.0 * (p.n + p.lambda)));
                gsl_vector_set(s.wc, j, 1.0 / (2.0 * (p.n + p.lambda)));
            }

            s.dk = gsl_matrix_alloc(p.no, p.nbSamples); // Holds the image of the sigma points
            gsl_matrix_set_zero(s.dk);

            s.ino_dk = gsl_vector_alloc(p.no); // Holds the inovation
            gsl_vector_set_zero(s.ino_dk);

            s.d_mean = gsl_vector_alloc(p.no); // Holds the mean of the sigma points images
            gsl_vector_set_zero(s.d_mean);

            s.sigmaPoints = gsl_matrix_alloc(p.n,p.nbSamples); // Holds the sigma points in the columns
            gsl_matrix_set_zero(s.sigmaPoints);

            // Temporary vectors/matrices
            s.vec_temp_n = gsl_vector_alloc(p.n);
            s.vec_temp_output = gsl_vector_alloc(p.no);

            s.mat_temp_n_1 = gsl_matrix_alloc(p.n,1);
            s.mat_temp_n_output = gsl_matrix_alloc(p.n, p.no);
            s.mat_temp_output_n = gsl_matrix_alloc(p.no, p.n);
            s.mat_temp_1_output = gsl_matrix_alloc(1,p.no);
            s.mat_temp_output_1 = gsl_matrix_alloc(p.no, 1);
            s.mat_temp_output_output = gsl_matrix_alloc(p.no, p.no);
            s.mat_temp_n_n = gsl_matrix_alloc(p.n, p.n);
        }

        /**
          * @short Free of memory allocation
          *
          */
        void ukf_free(ukf_param &p, ukf_state &s)
        {
            gsl_matrix_free(s.Kk);
            gsl_matrix_free(s.Kk_T);
            gsl_matrix_free(s.Pwdk);
            gsl_matrix_free(s.Pddk);
            gsl_matrix_free(s.Peek);
            gsl_matrix_free(s.Prrk);

            gsl_vector_free(s.w);
            gsl_vector_free(s.wk);

            gsl_matrix_free(s.Pk);
            gsl_matrix_free(s.Sk);
            gsl_vector_free(s.cSk);

            gsl_vector_free(s.wm);
            gsl_vector_free(s.wc);

            gsl_matrix_free(s.dk);
            gsl_vector_free(s.ino_dk);
            gsl_vector_free(s.d_mean);

            gsl_matrix_free(s.sigmaPoints);

            gsl_vector_free(s.vec_temp_n);
            gsl_vector_free(s.vec_temp_output);

            gsl_matrix_free(s.mat_temp_n_1);
            gsl_matrix_free(s.mat_temp_n_output);
            gsl_matrix_free(s.mat_temp_output_n);
            gsl_matrix_free(s.mat_temp_1_output);
            gsl_matrix_free(s.mat_temp_output_1);

            gsl_matrix_free(s.mat_temp_output_output);
            gsl_matrix_free(s.mat_temp_n_n);
        }

        /**
          * @short Iteration for the statistical linearization
          *
          */
	template<typename GFUNC>
	void ukf_iterate(ukf_param &p, ukf_state &s,
			 GFUNC g,
			 gsl_vector * xk, gsl_vector* dk)
        {

            // Here, we implement the UKF for parameter estimation in the vectorial case
            // The notations follow p93 of the PhD thesis of Van Der Merwe, "Sigma-Point Kalman Filters for Probabilistic Inference in Dynamic State-Space Models"

            // ************************************************** //
            // ************ Time update equations    ************ //
            // ************************************************** //
            // Add the evolution noise to the parameter covariance Eq 3.137
            gsl_matrix_add(s.Pk, s.Prrk);

            // ************************************************** //
            // ************ Compute the sigma points ************ //
            // ************************************************** //
            // Equations 3.138
            // w_k^j = w_(k-1)  <-- this is here denoted s.w
            // w_k^j = w_(k-1) + gamma Sk_j for 1 <= j <= n
            // w_k^j = w_(k-1) - gamma Sk_j for n+1 <= j <= 2n

            // Perform a cholesky decomposition of Pk
            gsl_matrix_memcpy(s.Sk, s.Pk);
            gsl_linalg_cholesky_decomp(s.Sk);
            // Set all the elements of Lpi strictly above the diagonal to zero
            for(int k = 0 ; k < p.n ; ++k)
                for(int j = 0 ; j < k ; ++j)
                    gsl_matrix_set(s.Sk,j,k,0.0);

            gsl_matrix_set_col(s.sigmaPoints,0, s.w);
            for(int j = 1 ; j < p.n+1 ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(s.sigmaPoints, i, j, gsl_vector_get(s.w, i) + p.gamma * gsl_matrix_get(s.Sk,i,j-1));

            for(int j = p.n+1 ; j < p.nbSamples ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(s.sigmaPoints, i, j, gsl_vector_get(s.w, i) - p.gamma * gsl_matrix_get(s.Sk,i,j-(p.n+1)));

            /**************************************************/
            /***** Compute the images of the sigma points *****/
            /**************************************************/

            // Compute the images of the sigma points
            // and the mean of the dk
            gsl_vector_set_zero(s.d_mean);
            for(int j = 0 ; j < p.nbSamples ; j++)
            {
                // Equation 3.129
                gsl_matrix_get_col(s.wk, s.sigmaPoints,j);
                g(s.wk,xk, s.vec_temp_output);
                gsl_matrix_set_col(s.dk, j, s.vec_temp_output);

                // Equation 3.140
                // Update the mean : y_mean = sum_[j=0..2n] w_j y_j
                gsl_vector_scale(s.vec_temp_output, gsl_vector_get(s.wm,j));
                gsl_vector_add(s.d_mean, s.vec_temp_output);
            }

            /**************************************************/
            /************** Update the statistics *************/
            /**************************************************/

            gsl_matrix_set_zero(s.Pwdk);
            gsl_matrix_memcpy(s.Pddk, s.Peek); // Add R^e_k to Pddk, Eq 3.142
            for(int j = 0 ; j < p.nbSamples ; ++j)
            {
                // Update of Pwdk
                // (wk - w)
                gsl_matrix_get_col(s.wk, s.sigmaPoints,j);
                gsl_vector_sub(s.wk, s.w);
                gsl_matrix_set_col(s.mat_temp_n_1, 0, s.wk);

                // (dk - d_mean)
                gsl_matrix_get_col(s.vec_temp_output, s.dk, j);
                gsl_vector_sub(s.vec_temp_output, s.d_mean);
                gsl_matrix_set_col(s.mat_temp_output_1, 0, s.vec_temp_output);

                // compute wc_j . (wk - w_mean) * (dk - d_mean)^T
                gsl_blas_dgemm(CblasNoTrans, CblasTrans, gsl_vector_get(s.wc,j) , s.mat_temp_n_1, s.mat_temp_output_1, 0.0, s.mat_temp_n_output);

                // Equation 3.142
                // And add it to Pwdk
                gsl_matrix_add(s.Pwdk, s.mat_temp_n_output);

                // Equation 3.143
                // Update of Pddk
                gsl_matrix_get_col(s.vec_temp_output, s.dk, j);
                gsl_vector_sub(s.vec_temp_output, s.d_mean);
                gsl_matrix_set_col(s.mat_temp_output_1, 0, s.vec_temp_output);
                gsl_blas_dgemm(CblasNoTrans, CblasTrans, gsl_vector_get(s.wc,j) , s.mat_temp_output_1, s.mat_temp_output_1, 0.0, s.mat_temp_output_output);

                gsl_matrix_add(s.Pddk, s.mat_temp_output_output);
            }

            // ************************************************** //
            // ******* Kalman gain and parameters update ******** //
            // ************************************************** //

            //*** Ki = Pwdk Pddk^-1
            // Compute the inverse of Pddk
            gsl_matrix_memcpy(s.mat_temp_output_output, s.Pddk);
            gsl_linalg_cholesky_decomp(s.mat_temp_output_output);
            gsl_linalg_cholesky_invert(s.mat_temp_output_output);

            // Compute the product : Pwdk . Pddk^-1
            // Equation 3.144
            gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0 , s.Pwdk, s.mat_temp_output_output, 0.0, s.Kk);

            // Update of the parameters
            // wk = w_(k-1) + Kk * (dk - d_mean)
            // Equation 3.145

            // Set the inovations
            /*for(int i = 0 ; i < p.no; ++i)
                s.ino_dk->data[i] = dk->data[i] - s.d_mean->data[i];
            gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0 , s.Kk, &gsl_matrix_view_array(s.ino_dk->data,p.no,1).matrix, 0.0, s.mat_temp_n_1);
            gsl_matrix_get_col(s.vec_temp_n, s.mat_temp_n_1, 0);
            gsl_vector_add(s.w, s.vec_temp_n);*/

            for(int i = 0 ; i < p.no; ++i)
                s.ino_dk->data[i] = dk->data[i] - s.d_mean->data[i];
            gsl_blas_dgemv(CblasNoTrans, 1.0, s.Kk, s.ino_dk, 1.0, s.w);

            // Update of the parameter covariance
            // Pk = P_(k-1) - Kk Pddk Kk^T
            // Equation 3.146
            gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0 , s.Kk, s.Pddk, 0.0, s.mat_temp_n_output);
            gsl_blas_dgemm(CblasNoTrans, CblasTrans, 1.0, s.mat_temp_n_output , s.Kk, 0.0, s.mat_temp_n_n);
            gsl_matrix_sub(s.Pk, s.mat_temp_n_n);

            // Update of the evolution noise
            p.evolution_noise->updateEvolutionNoise(p, s);
        }

        /**
          * @short Evaluation of the output from the sigma points
          *
          */
      void ukf_evaluate(ukf_param &p, ukf_state &s,
			void(*g)(gsl_vector*, gsl_vector*, gsl_vector*),
			gsl_vector * xk, gsl_vector * dk)
        {
            // ************************************************** //
            // ************ Compute the sigma points ************ //
            // ************************************************** //
            // Equations 3.138
            // w_k^j = w_(k-1)  <-- this is here denoted s.w
            // w_k^j = w_(k-1) + gamma Sk_j for 1 <= j <= n
            // w_k^j = w_(k-1) - gamma Sk_j for n+1 <= j <= 2n

            // Perform a cholesky decomposition of Pk
            gsl_matrix_memcpy(s.mat_temp_n_n, s.Pk);
            gsl_linalg_cholesky_decomp(s.mat_temp_n_n);
            // Set all the elements of Lpi strictly above the diagonal to zero
            for(int k = 0 ; k < p.n ; ++k)
                for(int j = 0 ; j < k ; ++j)
                    gsl_matrix_set(s.mat_temp_n_n,j,k,0.0);

            gsl_matrix_set_col(s.sigmaPoints,0, s.w);
            for(int j = 1 ; j < p.n+1 ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(s.sigmaPoints, i, j, gsl_vector_get(s.w, i) + p.gamma * gsl_matrix_get(s.mat_temp_n_n,i,j-1));

            for(int j = p.n+1 ; j < p.nbSamples ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(s.sigmaPoints, i, j, gsl_vector_get(s.w, i) - p.gamma * gsl_matrix_get(s.mat_temp_n_n,i,j-(p.n+1)));

            /**************************************************/
            /***** Compute the images of the sigma points *****/
            /**************************************************/

            // Compute the images of the sigma points
            // and the mean of the dk
            gsl_vector_set_zero(dk);
            for(int j = 0 ; j < p.nbSamples ; j++)
            {
                // Equation 3.129
                gsl_matrix_get_col(s.wk, s.sigmaPoints,j);
                g(s.wk,xk, s.vec_temp_output);
                gsl_matrix_set_col(s.dk, j, s.vec_temp_output);

                // Equation 3.140
                // Update the mean : y_mean = sum_[j=0..2n] w_j y_j
                gsl_vector_scale(s.vec_temp_output, gsl_vector_get(s.wm,j));
                gsl_vector_add(dk, s.vec_temp_output);
            }
        }

        /**
          * @short Returns a set of sigma points
          */
        void getSigmaPoints(ukf_param &p, ukf_state &s, gsl_matrix * sigmaPoints)
        {
            //gsl_matrix * sigmaPoints = gsl_matrix_alloc(p.n, p.nbSamples);

            // ************************************************** //
            // ************ Compute the sigma points ************ //
            // ************************************************** //
            // Equations 3.138
            // w_k^j = w_(k-1)  <-- this is here denoted s.w
            // w_k^j = w_(k-1) + gamma Sk_j for 1 <= j <= n
            // w_k^j = w_(k-1) - gamma Sk_j for n+1 <= j <= 2n

            // Perform a cholesky decomposition of Pk
            gsl_matrix_memcpy(s.mat_temp_n_n, s.Pk);
            gsl_linalg_cholesky_decomp(s.mat_temp_n_n);
            // Set all the elements of Lpi strictly above the diagonal to zero
            for(int k = 0 ; k < p.n ; ++k)
                for(int j = 0 ; j < k ; ++j)
                    gsl_matrix_set(s.mat_temp_n_n,j,k,0.0);

            gsl_matrix_set_col(sigmaPoints,0, s.w);
            for(int j = 1 ; j < p.n+1 ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(sigmaPoints, i, j, gsl_vector_get(s.w, i) + p.gamma * gsl_matrix_get(s.mat_temp_n_n,i,j-1));

            for(int j = p.n+1 ; j < p.nbSamples ; ++j)
                for(int i = 0 ; i < p.n ; ++i)
                    gsl_matrix_set(sigmaPoints, i, j, gsl_vector_get(s.w, i) - p.gamma * gsl_matrix_get(s.mat_temp_n_n,i,j-(p.n+1)));

        }

    } // parameter
} // ukf

#endif // SL_PARAMETER_NDIM_H


/* StdVarTrajs_Sampling.c

Simulation for calculating allele trajectories, adaptation from standing variation
To be fed into MSMS

Simulation uses routines found with the GNU Scientific Library (GSL)
(http://www.gnu.org/software/gsl/)
Since GSL is distributed under the GNU General Public License
(http://www.gnu.org/copyleft/gpl.html), you must download it
separately from this file.

This program can be compiled with e.g. GCC using a command like:
gcc StdVarTrajs_Sampling -lm -lgsl -lgslcblas -I/usr/local/include -L/usr/local/lib StdVarTrajs_Sampling.c

Then run by executing:
./StdVarTrajs_Sampling N s h x0 reps cutoff
Where:
- N is the (haploid) size
- s is the selection coefficient of the beneficial allele
- h is dominance of the beneficial allele
- x0 is frequency of the allele when it started to be selected for
- reps is how many times the second allele should FIX before simulation stops
(the number of actual runs is greater due to stochastic loss)
- cutoff denotes how high the derived allele can be in the neutral phase ('0' for fixation, '1' for x0 + 1/2N)
*/

/* Preprocessor statements */
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include <stddef.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <sys/stat.h>
#include <sys/types.h>

#define INITTS 310000

/* Function prototypes */
void geninit(double *geninit, double pee);
void selection(double *geninit, double h, double s);
void reproduction(double *geninit);
void bentraj(double *XFreq, unsigned int *btime, double *genotype, double x0, unsigned int rep, unsigned int N, double s, double h, const gsl_rng *r);
void neutraj(double *NFreq, unsigned int *ntime, double *genotype, double x0, unsigned int rep, unsigned int N, unsigned int cutoff, const gsl_rng *r);
void trajprint(double **AFreq, unsigned int *ttime, unsigned int rep);
void Wait();

/* Initialising genotypes */
void geninit(double *geninit, double pee){

	*(geninit + 0) = (1-pee)*(1-pee);
	*(geninit + 1) = 2*pee*(1-pee);
	*(geninit + 2) = pee*pee;

}	/* End of gen initiation routine */

/* Selection routine (single locus) */
void selection(double *geninit, double h, double s){
	/* Fitness of each genotype */
	double WAA, WAa, Waa;
	double Wmean;				/* Mean fitness */

	Waa = 1;
	WAa = 1 + h*s;
	WAA = 1 + s;

	/* Mean fitness calculation */
	Wmean = ((*(geninit + 0))*Waa) + ((*(geninit + 1))*WAa) + ((*(geninit + 2))*WAA);

	/* Changing frequencies by selection */
	*(geninit + 0) = ((*(geninit + 0))*Waa)/Wmean;
	*(geninit + 1) = ((*(geninit + 1))*WAa)/Wmean;
	*(geninit + 2) = ((*(geninit + 2))*WAA)/Wmean;

}	/* End of selection routine */

/* Reproduction routine */
void reproduction(double *geninit){

	/* Fed-in genotype frequencies (for ease of programming) */
	double gaas, gAas, gAAs;
	/* Haplotypes */
	double xA, xa;

	/* Initial definition of genotypes */
	gaas = *(geninit + 0);
	gAas = *(geninit + 1);
	gAAs = *(geninit + 2);

	/* Baseline change in haplotype frequencies with obligate outcrossing */
	xA = gAAs + (gAas)/2.0;
	xa = gaas + (gAas)/2.0;

	/* Change in frequencies (HWE) */
	*(geninit + 0) = pow(xa,2);
	*(geninit + 1) = 2.0*xA*xa;
	*(geninit + 2) = pow(xA,2);

}	/* End of reproduction routine */

/* Routine to simulate beneficial allele trajectory */
void bentraj(double *XFreq, unsigned int *btime, double *genotype, double x0, unsigned int rep, unsigned int N, double s, double h, const gsl_rng *r){

	unsigned int gt = 1;		/* Generation time */
	unsigned int done = 0;		/* Counter if routine done or not */
	unsigned int exitl = 0;		/* Decide if exit loop or not */
	unsigned int i;				/* Genotype counter */
	double x = 0;				/* Beneficial allele frequency */

	unsigned int *gensamp = calloc(3,sizeof(unsigned int));		/* New population samples */

	while(done == 0){
		exitl = 0;
		gt = 1;
		x = x0;
		*(XFreq + 0) = x0;
		/* Initiating genotypes */
		geninit(genotype,x0);
		while(exitl == 0){

			/* Selection routine */
    		selection(genotype,h,s);

			/* Reproduction routine */
    		reproduction(genotype);

    		/* Sampling based on new frequencies */
	       	gsl_ran_multinomial(r,3,N,genotype,gensamp);
       		for(i = 0; i < 3; i++){
    	   		*(genotype + i) = (*(gensamp + i))/(1.0*N);
	       	}

	       	x = 0.5*(*(genotype + 1)) + (*(genotype + 2));
			*(XFreq + gt) = x;
			gt += 1;
			if(gt > INITTS){
				fprintf(stderr,"Number of generations in selection trajectory exceed vector length (INITTS).\n");
				exit(1);
			}

			if( (x == 1) || (x == 0) ){
				exitl = 1;
			}
		}

		if(x == 1){
			done = 1;
			*(btime + rep) = (gt - 1);
		}
	}

	free(gensamp);

}

/* Routine to simulate neutral allele trajectory */
void neutraj(double *NFreq, unsigned int *ntime, double *genotype, double x0, unsigned int rep, unsigned int N, unsigned int cutoff, const gsl_rng *r){

	unsigned int gt = 1;		/* Generation time */
	unsigned int done = 0;		/* Counter if routine done or not */
	unsigned int exitl = 0;		/* Decide if exit loop or not */
	unsigned int i;				/* Genotype counter */
	double x = 0;				/* Neutral allele frequency */
	double xmax = 1;			/* Value of x to use as maximum cutoff */

	unsigned int *gensamp = calloc(3,sizeof(unsigned int));		/* New population samples */

	if(cutoff == 0){
		xmax = 1;
	}else if(cutoff == 1){
		xmax = x0 + (1.0/(2.0*N));
		if(xmax > 1){
			xmax = 1;
		}
	}
/*	printf("xmax is %lf\n",xmax);	*/

	while(done == 0){
		exitl = 0;
		gt = 0;		/* Note start from zero here since x0 case defined in selection trajectory */
		x = x0;

		/* Initiating genotypes */
		geninit(genotype,x0);

		while(exitl==0){

			/* Reproduction routine */
    		reproduction(genotype);

			/* Random sampling of allele frequencies (drift) */
	       	gsl_ran_multinomial(r,3,N,genotype,gensamp);
       		for(i = 0; i < 3; i++){
    	   		*(genotype + i) = (*(gensamp + i))/(1.0*N);
	       	}
	       	x = 0.5*(*(genotype + 1)) + (*(genotype + 2));
			*(NFreq + gt) = x;

			gt++;
			if(gt > INITTS){
				printf("Number of generations in neutral trajectory exceed vector length (INITTS).\n");
				exit(1);
			}

			if( (x == 0) || (x >= xmax) ){
				exitl = 1;
			}
		}

		if(x == 0){
			done = 1;
			*(ntime + rep) = (gt - 1);
		}
	}

	free(gensamp);

}

/* Print out trajectories to file */
void trajprint(double **AFreq, unsigned int *ttime, unsigned int rep){
	unsigned int j;
	char filename[32];
	FILE *ofp_tr;				/* Pointer for file output */

	sprintf(filename,"Traj/ATraj%d.dat",(rep + 1));
	ofp_tr = fopen(filename,"w");

	for(j = 0; j < *(ttime + rep); j++){
		fprintf(ofp_tr,"%lf %lf %lf\n",*((*(AFreq + 0)) + j),*((*(AFreq + 1)) + j),*((*(AFreq + 2)) + j));
	}

	fclose(ofp_tr);

}

void Wait(){
	printf("Press Enter to Continue");
	while( getchar() != '\n' );
	printf("\n");
}

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

	/* Declare variables here */
	unsigned int i, a;			/* Rep counter, memory counter */
	unsigned int N = 0;			/* Population Size */
	unsigned int Nreps = 0;		/* Number of repetitions */
	unsigned int acc = 0;		/* Accumulator (when merging two trajectories) */
	unsigned int es = 10;		/* How many extra steps to put in the trajectory files */
	unsigned int cutoff = 0;	/* What kind of back-in-time cutoff to use */
	int j;						/* Timestep counter */
	double s = 0;				/* Strength of selection */
	double h = 0;				/* Dominance level */
	double x0 = 0;				/* Initial allele frequency */
	double lim = 0;				/* Limit of acceptable allele frequencies */
	double dx = 0;				/* Rescaled timestep so can be used by MSMS */
	FILE *ofp_sd;				/* Pointer for seed output */

	/* GSL random number definitions */
	const gsl_rng_type * T;
	gsl_rng * r;

	/* Reading in data from command line */
	/* <Program> N s h x0 NReps cutoff */
	if(argc != 7){
		fprintf(stderr,"Six inputs are needed (N s h x0 Reps cutoff).\n");
		exit(1);
	}

	N = atoi(argv[1]);
	if(N <= 0){
		fprintf(stderr,"Total Population size N is zero or negative, not allowed.\n");
		exit(1);
	}

	/* Defining min frequency limit ** based on unscaled N ** */
	lim = (1.0/(1.0*N));

	/* Defining timestep, dx = 1/4N since that's what MSMS uses */
	dx = (1.0/(4.0*N));

	s = strtod(argv[2],NULL);
	if(s < 0){
		fprintf(stderr,"Allele strength s is negative, not allowed.\n");
		exit(1);
	}

	h = strtod(argv[3],NULL);
	if(h < 0 || h > 1){
		fprintf(stderr,"Dominance value must lie between 0 and 1.\n");
		exit(1);
	}

	x0 = strtod(argv[4],NULL);
	if(x0 < lim || x0 > (1-lim)){
		fprintf(stderr,"Initial mutant frequency must lie between %0.5lf and %0.5lf.\n",lim,1-lim);
		exit(1);
	}

	/* Number of samples/reps to take */
	Nreps = atoi(argv[5]);
	if(Nreps <= 0){
		fprintf(stderr,"Must set positive number of repetitions.\n");
		exit(1);
	}

	/* Determining cutoff */
	cutoff = atoi(argv[6]);
	if(!(cutoff == 0 || cutoff == 1)){
		fprintf(stderr,"'Cutoff' switch must equal 0 or 1.\n");
		exit(1);
	}

	unsigned int *btime = calloc(Nreps,sizeof(unsigned int));			/* Number of generations needed for ben allele fixation */
	unsigned int *ntime = calloc(Nreps,sizeof(unsigned int));			/* Number of generations needed for neut allele loss */
	unsigned int *ttime = calloc(Nreps,sizeof(unsigned int));			/* Number of generations needed for whole processes */
	double *genotype = calloc(3,sizeof(double));						/* Genotype frequencies */

	/* create a generator chosen by the
    environment variable GSL_RNG_TYPE */

	gsl_rng_env_setup();
	if (!getenv("GSL_RNG_SEED")) gsl_rng_default_seed = time(0);
	T = gsl_rng_default;
	r = gsl_rng_alloc(T);
	ofp_sd = fopen("Seed.dat","w");
	fprintf(ofp_sd,"%lu\n",gsl_rng_default_seed);
	fclose(ofp_sd);

	mkdir("Traj/", 0777);

	/* Executing simulation */
	for(i = 0; i < Nreps; i++){

		if(Nreps > 100){
			if(i%100 == 0){
				printf("Starting run %d\n",i);
			}
		}

		double *XFreq = calloc(INITTS,sizeof(double));			/* Beneficial allele frequency */
		double *NFreq = calloc(INITTS,sizeof(double));			/* Neutral allele frequency */

		/* Simulating beneficial trajectories */
		bentraj(XFreq, btime, genotype, x0, i, N, s, h, r);

		/* Simulation neutral trajectory */
		neutraj(NFreq, ntime, genotype, x0, i, N, cutoff, r);

		*(ttime + i) = (*(btime+i)) + (*(ntime+i) + 2 + es);
/*		printf("%d %d %d\n",*(btime+i),*(ntime+i),*(ttime+i));	*/

		/* Defining memory for AFreq */
		double **AFreq = calloc(3,sizeof(double *));					/* Allele frequency over time */
		for(a = 0; a < 3; a++){											/* Assigning space for each run */
			AFreq[a] = calloc(*(ttime + i),sizeof(double));
		}

		/* Merging trajectories
		Timing is set up so tfix = 0
		And discrete generations is scaled by timestep, dx
		*/

		/* First add ten timesteps for x = 0
		So trajectory fully covered in MSMS */
		acc = 0;
		for(j = 0; j < es; j++){
			*((*(AFreq + 0)) + j) = ((*(ttime + i) - es - 1)*dx) + (es - j);
			*((*(AFreq + 1)) + j) = 1.0;
			*((*(AFreq + 2)) + j) = 0.0;
		}
		for(j = *(ntime + i); j >= 0; j--){
			*((*(AFreq + 0)) + acc + es) = (*(ttime + i) - acc - es - 1)*dx;
			*((*(AFreq + 1)) + acc + es) = 1.0-(*(NFreq + j));
			*((*(AFreq + 2)) + acc + es) = *(NFreq + j);
			acc++;
		}
		for(j = 0; j <= *(btime + i); j++){
			*((*(AFreq + 0)) + acc + es) = (*(ttime + i) - acc - es - 1)*dx;
			*((*(AFreq + 1)) + acc + es) = 1.0-(*(XFreq + j));
			*((*(AFreq + 2)) + acc + es) = *(XFreq + j);
			acc++;
		}
		/*
		*(ttime + i) = (acc-1);
		printf("%d\n",*(ttime+i));
		*/

		free(NFreq);
		free(XFreq);

		/* Printout of trajectories */
		trajprint(AFreq, ttime, i);

		for(a = 0; a < 3; a++){
	 		free(AFreq[a]);
		}
		free(AFreq);

	}

	gsl_rng_free(r);
	free(genotype);
	free(ttime);
	free(ntime);
	free(btime);

	return 0;

}	/* End of main program */

/* End of File */


#include <stdlib.h>
#include <cblas.h>

#include "tasks.h"

void syrk_task_par_reconfigure(int nth)
{
    // empty
}

void syrk_task_par_finalize(void)
{
    // empty
}

void syrk_task_par(void *ptr, int nth, int me)
{
    struct syrk_task_arg *arg = (struct syrk_task_arg*) ptr;

    int n       = arg->n;
    int k       = arg->k;
    double *A21 = arg->A21;
    double *A22 = arg->A22;
    int ldA     = arg->ldA;

    // Balance the load by flops.
    int part[nth + 1];
    const int total_work = n * (n + 1) / 2;
    const int ideal_part_work = total_work / nth;
    part[0] = 0;
    part[nth] = n;
    for (int k = 1; k < nth; ++k) {
        part[k] = part[k - 1];
        int work = 0;
        while (work < ideal_part_work && part[k] < n) {
            work += n - part[k];
            part[k] += 1;
        }
    }

    const int my_first_col = part[me];
    const int my_num_cols  = part[me + 1] - part[me];
    const int i1 = my_first_col;
    const int i2 = i1 + my_num_cols;
    const int m2 = n - i2;

    // TODO Insert picture (see krnl_syrk in task_chol_par).

    if (my_num_cols > 0) {
        cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans,
                    my_num_cols, k,
                    -1.0, A21 + i1,                           ldA,
                     1.0, A22 + i1 + my_first_col * ldA,      ldA);
    }

    if (m2 > 0) {
        cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans,
                    m2, my_num_cols, k,
                    -1.0, A21 + i2,                      ldA,
                          A21 + my_first_col,            ldA,
                     1.0, A22 + i2 + my_first_col * ldA, ldA);
    }
}