/*------------------------------------------------------
Maximum likelihood estimation
of migration rate  and effectice population size
using a Metropolis-Hastings Monte Carlo algorithm
-------------------------------------------------------
    slice sampler routines   R O U T I N E S

    Peter Beerli 2006, Tallahassee
    beerli@scs.fsu.edu

    Copyright 2006 Peter Beerli, Tallahassee

    This software is distributed free of charge for non-commercial use
    and is copyrighted. Of course, we do not guarantee that the software
    works and are not responsible for any damage you may cause or have.


 $Id$
    */
/*! \file slice.c 
this file implements the slice sampler that can replace the proposal distribution
for parameters as an alternative to the more standard proposal. The code is based on this
mathematica function (created after talking to Paul Lewis):

slice[start_, stick_] := Module[{x = start, r, y, z, z2},
    r = f[x];
y = Random[Real, {0, r}];
r2 = Random[];
z = {x - (1 - r2)*stick, x + r2 * stick};
z2 = Map[f, z];
While[y < z2 [[1]] ,
      z2[[1]] = f[z[[1]] -= stick]];
While[z2 [[2]] >  y, z2[[2]] = f[z[[2]] += stick]];
    x = Random[Real, z];
While[(f[x]) < y && count++ < 100,
      If[(x - z[[1]]) < z[[2]] - x, z[[1]] = x, z[[2]] = x];
      (*Print[{x, f[x]}];*)
      x = Random[Real, z];
      ];
      {x, y}]

stick is arbitrary in size, too small to much work, too large to much work
*/
#include <stdlib.h>
#include "migration.h"
#include "bayes.h"
#include "slice.h"
#include "random.h"
#include "tools.h"
#include "sighandler.h"
#include "world.h"
#ifdef PRETTY
#include "pretty.h"
#endif
#ifdef MPI
#include "migrate_mpi.h"
#else
extern int myID;
#endif
extern MYREAL norm (MYREAL *d, long size);

void set_slice_param(MYREAL newparam, long which, world_fmt * world)
{
  long numpop2 = world->numpop2;  
  // set new parameters and allow for custom migration matrix
  if(which<numpop2)
    {
      if(newparam<0)
	newparam = world->bayes->minparam[which];
      bayes_set_param(world->param0,newparam,which,world->options->custm2, world->numpop);
      reprecalc_world(world, which);
      }
    else
      {
	world->options->mu_rates[world->locus] = newparam;
      }
}


void set_multi_slice_param(MYREAL *newparam, long numparam, world_fmt * world)
{
  long numpop2 = world->numpop2;  
  // set new parameters and allow for custom migration matrix
  memcpy(world->param0,newparam,sizeof(MYREAL)*numpop2);
  precalc_world(world);
  if(world->bayes->mu)
    {
      world->options->mu_rates[world->locus] = newparam[numpop2];
    }
}


///
/// slice() calculates a new parameter value that comes from the proposal distribution
/// returns log(prob) and replaces startval
MYREAL slice (MYREAL *startval, long which, world_fmt * world, MYREAL  (*func) (MYREAL, MYREAL, bayes_fmt *, long))
{
  bayes_fmt * bayes = world-> bayes;
  MYREAL *minparam = bayes->minparam;
  MYREAL *maxparam = bayes->maxparam;
  MYREAL newstartval = *startval;
  MYREAL log_newstartval;
  MYREAL ra;
  MYREAL x;
  MYREAL fx;
  MYREAL y;
  MYREAL r;
  MYREAL zl;
  MYREAL zr;
  MYREAL ezl;
  MYREAL ezr;
  MYREAL ex;
  MYREAL fzl;
  MYREAL fzr;
  MYREAL priorratio;
  MYREAL stick = (which < world->numpop) ? 1.0 : 2.0 ;
  long count;
  // safeguard against slicing towards zero
  if(minparam[which] < SMALL_VALUE)
    minparam[which] = SMALL_VALUE;
  // replace the parameter value at which with startval
  if(maxparam[which] -  minparam[which] < SMALLEPSILON)
    {
       return world->bayes->oldval;
    }
  // normal execution pattern
  if((newstartval > maxparam[which]) || (newstartval < minparam[which]))
    {
      fprintf(stdout,"Adjusting startvalue in slice sampler oldval=%f",newstartval);
      newstartval = minparam[which] + RANDUM() * (maxparam[which] - minparam[which]);
      fprintf(stdout," newval=%f\n", newstartval);
    }
  log_newstartval = log(newstartval);
  x = log_newstartval;
  priorratio = (MYREAL) func(newstartval, -1. , bayes, which);
  set_slice_param(newstartval,which, world);
  // calculate the function
  fx = probg_treetimes(world)+ priorratio;
  // find a random point between zero and fx
  ra = log(RANDUM());
  y  = ra + fx ;
  // position horizontal stick on coordinates (x,y)
  ra = RANDUM();
  r  = ra * stick;
  zl = x - r;
  zr = zl + stick;
  // extend the stick until it crosses the function
  ezl = EXP(zl);
  ezr = EXP(zr);
  priorratio = func(ezl, -1. , bayes, which);
  set_slice_param(ezl,which, world);
  fzl= probg_treetimes(world) + priorratio;
  priorratio = func(ezr, -1. , bayes, which);
  set_slice_param(ezr,which, world);
  fzr= probg_treetimes(world) + priorratio;
  while(y < fzl)
    {
      zl -= stick;
      ezl = EXP(zl);
      priorratio = func(ezl, -1. , bayes, which);
      set_slice_param(ezl,which, world);
      fzl= probg_treetimes(world) + priorratio;
    }
  while(y < fzr)
    {
      zr += stick;
      ezr = EXP(zr);
      priorratio = func(ezr, -1. ,bayes, which);
      set_slice_param(ezr,which, world);
      fzr= probg_treetimes(world) + priorratio;
    }
  // pick new value at random
  ra = RANDUM();
  x  = ra * (zr - zl) + zl;
  ex = EXP(x);
  priorratio = func(ex, -1. , bayes, which);
  set_slice_param(ex,which, world);
  fx = probg_treetimes(world) + priorratio;
  count = 0;
  while(fx < y && ((zr-zl) > EPSILON))
    {
      count++;
      //      if((x - zl) < (zr - x))
      if(x < log_newstartval)
	{
	  zl = x;
	}
      else
	{
	  zr = x;
	}
      ra = RANDUM();
      x  = ra * (zr-zl)+ zl;
      priorratio = func(EXP(x), -1. , bayes, which);
      set_slice_param(EXP(x),which, world);
      fx = probg_treetimes(world) + priorratio;
    }
  //fprintf(stdout,"{{%f,%f},{%f,%f},{%f,%f}} (* y=%f, c=%li, %li, lv=%f *)\n",zl,fzl,x,fx,zr,fzr,y, count, which,startval);
  *startval = EXP(x);
  return fx - priorratio;
}

///
/// slice() calculates a new parameter value  that comes from the proposal distribution
/// returns the log(probability) and changes startval]
MYREAL expslice (MYREAL *startval, long which, world_fmt * world, MYREAL  (*func) (MYREAL, MYREAL, bayes_fmt *, long))
{
  bayes_fmt * bayes = world-> bayes;
  MYREAL *minparam = bayes->minparam;
  MYREAL *maxparam = bayes->maxparam;
  MYREAL newstartval = *startval;
  //  MYREAL log_newstartval;
  MYREAL ra;
  MYREAL x;
  MYREAL fx;
  MYREAL y;
  MYREAL r;
  MYREAL zl;
  MYREAL zr;
  //  MYREAL ezl;
  //  MYREAL ezr;
  //  MYREAL ex;
  MYREAL fzl;
  MYREAL fzr;
  MYREAL priorratio;
  //  MYREAL stick = (maxparam[which] - minparam[which]) / 20. ;
  MYREAL stick = *startval / 10.;
  long count = 3L;
#ifdef SLICEREPORTER
  fprintf(stdout,"# start slice --------------------\n");
#endif
  // safeguard against slicing towards zero
  //if(minparam[which] < SMALL_VALUE)
  //  minparam[which] = SMALL_VALUE;
  // replace the parameter value at which with startval
  if(maxparam[which] -  minparam[which] < SMALLEPSILON)
    {
       return world->bayes->oldval;
    }
  // normal execution pattern
  // replace the parameter value at which with startval
  if((newstartval > maxparam[which]) || (newstartval < minparam[which]))
  {
    fprintf(stdout,"Adjusting startvalue in slice sampler oldval=%f",newstartval);
    newstartval = minparam[which] + RANDUM() * (maxparam[which] - minparam[which]);
    fprintf(stdout," newval=%f\n", newstartval);
  }
  x = newstartval;
  priorratio = (MYREAL) func(newstartval, -1. , bayes, which);
  set_slice_param(newstartval,which, world);
  // calculate the function
  fx = probg_treetimes(world)+ priorratio;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s\n", x, fx,"startpoint");
#endif

  // find a random point between zero and fx
  ra = log(RANDUM());
  y  = ra + fx ;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s\n",y , fx,"y");
#endif
  // position horizontal stick on coordinates (x,y)
  ra = RANDUM();
  r  = ra * stick;
  zl = x - r;
  zr = zl + stick;
  // extend the stick until it crosses the function
  //  ezl = EXP(zl);
  //ezr = EXP(zr);
  priorratio = func(zl, -1. , bayes, which);
  set_slice_param(zl,which, world);
  fzl= probg_treetimes(world) + priorratio;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s\n", zl, fzl,"lower");
#endif
  priorratio = func(zr, -1. , bayes, which);
  set_slice_param(zr,which, world);
  fzr= probg_treetimes(world) + priorratio;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s\n", zr, fzr,"upper");
#endif
  while(y < fzl)
    {
      zl -= stick;
      //      ezl = EXP(zl);
      priorratio = func(zl, -1. , bayes, which);
      set_slice_param(zl,which, world);
      fzl= probg_treetimes(world) + priorratio;
      count++;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", zl, fzl,"lower", count);
#endif
    }
  while(y < fzr)
    {
      zr += stick;
      //      ezr = EXP(zr);
      priorratio = func(zr, -1. ,bayes, which);
      set_slice_param(zr,which, world);
      fzr= probg_treetimes(world) + priorratio;
      count++;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", zr, fzr,"upper", count);
#endif
    }
  // pick new value at random
  ra = RANDUM();
  x  = ra * (zr - zl) + zl;
  //  ex = EXP(x);
  priorratio = func(x, -1. , bayes, which);
  set_slice_param(x,which, world);
  fx = probg_treetimes(world) + priorratio;
  count++;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", x, fx,"new_x", count);
#endif
  while(fx < y && ((zr-zl) > EPSILON))
    {
      count++;
      //      if((x - zl) < (zr - x))
      if(x < newstartval)
	{
	  zl = x;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", x, fx,"new_lower", count);
#endif
	}
      else
	{
	  zr = x;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", x, fx,"new_upper", count);
#endif
	}
      ra = RANDUM();
      x  = ra * (zr-zl)+ zl;
      priorratio = func(x, -1. , bayes, which);
      set_slice_param(x,which, world);
      fx = probg_treetimes(world) + priorratio;
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n", x, fx,"new_x", count);
#endif
    }
  //fprintf(stdout,"{{%f,%f},{%f,%f},{%f,%f}} (* y=%f, c=%li, %li, lv=%f *)\n",zl,fzl,x,fx,zr,fzr,y, count, which,startval);
  //fprintf(stdout,"%3li: %f %f %li\n",which, x, fx, count);
#ifdef SLICEREPORTER
  fprintf(stdout,"#@ %f %f %s %li\n#----------------------\n", x, fx,"return_x", count);
#endif
  *startval = x;
  return fx - priorratio;;
}

//#ifdef TESTING2
// log_prior_ratioall needs a body in bayes.c
void random_direction(MYREAL *coord,long dim)
{
  long i;
  MYREAL ci;
  MYREAL norm=0.0;
  for(i=0;i<dim;i++)
    {
      ci = rannor(0.0,1.0);
      coord[i] = ci;
      norm += ci * ci;
    }
  norm = sqrt(norm);
  for(i=0;i<dim;i++)
    {
      coord[i] /= norm;
    }
}

void sphere_random(MYREAL *lower, MYREAL *values, MYREAL *upper, 
		   MYREAL *r, MYREAL stick, long numparam, MYREAL *minparam, MYREAL *maxparam)
{
  long i;
  random_direction(r,numparam);
  for(i=0;i<numparam;i++)
    {
      lower[i] = values[i] - (r[i] * stick);
      if(lower[i] < minparam[i])
	lower[i]=minparam[i];
      upper[i] = lower[i] + stick;
      if(upper[i]>maxparam[i])
	upper[i] = maxparam[i];
    }
}

void sphere_x_random(MYREAL *lower, MYREAL *values, MYREAL *upper, 
		   MYREAL *r, long numparam, MYREAL * minparam, MYREAL * maxparam)
{
  long i;
  random_direction(r,numparam);
  for(i=0;i<numparam;i++)
    {
      values[i] = (RANDUM() * (upper[i]-lower[i]) + lower[i]);
      if(values[i] < minparam[i])
	values[i]=minparam[i];
      else
	{
	  if(values[i] > maxparam[i])
	    values[i] = maxparam[i];
	}
    }
}

void add_stick(MYREAL *values, MYREAL stick, MYREAL *r, long numparam, MYREAL * minparam, MYREAL * maxparam)
{
  long i;
  for(i=0;i<numparam;i++)
    {
      values[i] = values[i] + (r[i] * stick);
    }
  if(values[i] < minparam[i])
    values[i]=minparam[i];
  else
    {
      if(values[i] > maxparam[i])
	values[i] = maxparam[i];
    }
}

///
/// slice() calculates a new point that comes from the proposal distribution changing all parameters
/// at once using Michal's sphere approach, the parameter which is not really needed in this context
/// is kept to make the function look the same as the other single-dimensional ones 
void expallslice (MYREAL * values, MYREAL *startval, long which, world_fmt * world)
{
  long i;
  long numpop2 = world->numpop2;
  long numparam = numpop2 + ((which>=numpop2) ? 1 : 0);
  MYREAL norm_newstartval;
  MYREAL *minparam = world->bayes->minparam;
  MYREAL *maxparam = world->bayes->maxparam;
  MYREAL *newstartval;
  //  MYREAL log_newstartval;
  MYREAL ra;
  MYREAL *x;
  MYREAL fx;
  MYREAL y;
  MYREAL *r;
  MYREAL *zl;
  MYREAL *zr;
  //  MYREAL ezl;
  //  MYREAL ezr;
  //  MYREAL ex;
  MYREAL *v;
  MYREAL fzl;
  MYREAL fzr;
  MYREAL priorratio;
  MYREAL stick = 1.0; //this needs a remedy! (maxparam[which] - minparam[which]) / 20. ;
  long count = 3L;

  v = (MYREAL *) mycalloc(numparam * 5, sizeof(MYREAL));
  r = v;
  x = r + numparam;
  zl = x + numparam;
  zr = zl + numparam;
  newstartval = zr + numparam;
  memcpy(newstartval,startval,sizeof(MYREAL)*numparam);
  for(i=0; i < numpop2; i++)
    {
      // safeguard against slicing towards zero, eplxore whether zero is harmful, should not be!
      if(minparam[i] < SMALL_VALUE)
	minparam[i] = SMALL_VALUE;
      // replace the parameter value at which with startval
      if((newstartval[i] > maxparam[i]) || (newstartval[i] < minparam[i]))
	{

	  newstartval[i] = minparam[i] + RANDUM() * (maxparam[i] - minparam[i]);
	  fprintf(stdout,"Adjusting startvalue in slice sampler oldval[%li]=%f",i,startval[i]);
	  fprintf(stdout," newval[%li]=%f\n", i, newstartval[i]);
	}
      else
	{
	  newstartval[i] = startval[i];
	}
      x[i] = newstartval[i];
    }
  //  norm_newstartval = norm(newstartval,numparam, minparam, maxparam);
  norm_newstartval = norm(newstartval,numparam);
  priorratio = (MYREAL) log_prior_ratio_all(world,x);
  set_multi_slice_param(newstartval,numparam, world);
  // calculate the function
  fx = probg_treetimes(world)+ priorratio;
  // find a random point between zero and fx
  ra = log(RANDUM());
  y  = ra + fx ;
  // position horizontal stick on coordinates (x1,x2,...,xn,y)
  sphere_random(zl,x,zr,r,stick, numparam, minparam, maxparam);//zl = x - (r=sphererand()) * stick, zr = zl + stick
  //  zl = x - r;
  //zr = zl + stick;
  // extend the stick until it crosses the function
  //  ezl = EXP(zl);
  //ezr = EXP(zr);
  priorratio = log_prior_ratio_all(world,zl);
  set_multi_slice_param(zl, numparam, world);
  fzl= probg_treetimes(world) + priorratio;
  priorratio = log_prior_ratio_all(world,zr);
  set_multi_slice_param(zr, numparam, world);
  fzr= probg_treetimes(world) + priorratio;
  while(y < fzl)
    {
      add_stick(zl, -stick, r, numparam, minparam, maxparam);//zl -= stick;
      //      ezl = EXP(zl);
      priorratio = log_prior_ratio_all(world,zl);
      set_multi_slice_param(zl,numparam, world);
      fzl= probg_treetimes(world) + priorratio;
      count++;
    }
  while(y < fzr)
    {
      add_stick(zr, stick, r, numparam, minparam, maxparam); //zr += stick;
      //      ezr = EXP(zr);
      priorratio = log_prior_ratio_all(world,zr);
      set_multi_slice_param(zr,which, world);
      fzr= probg_treetimes(world) + priorratio;
      count++;
    }
  // pick new value at random
  sphere_x_random(zl,x,zr,r, numparam, minparam, maxparam);
  //ra = RANDUM();
  //x  = ra * (zr - zl) + zl;
  //  ex = EXP(x);
  priorratio = log_prior_ratio_all(world,x);
  set_multi_slice_param(x,which, world);
  fx = probg_treetimes(world) + priorratio;
  count++;
  while(fx < y /*&& ((zr-zl) > EPSILON)*/)
    {
      count++;
      //      if((x - zl) < (zr - x))
      if(norm(x,numparam)<norm_newstartval)//vector comparison?
	{
	  memcpy(zl,x,sizeof(MYREAL)*numpop2);
	}
      else
	{
	  memcpy(zr,x,sizeof(MYREAL)*numpop2);
	}
      sphere_x_random(zl,x, zr,r, numparam, minparam, maxparam);
      //      ra = RANDUM();
      //x  = ra * (zr-zl)+ zl;
      priorratio = log_prior_ratio_all(world,x);
      set_multi_slice_param(x,which, world);
      fx = probg_treetimes(world) + priorratio;
    }
  //fprintf(stdout,"{{%f,%f},{%f,%f},{%f,%f}} (* y=%f, c=%li, %li, lv=%f *)\n",zl,fzl,x,fx,zr,fzr,y, count, which,startval);
  //fprintf(stdout,"%3li: %f %f %li\n",which, x, fx, count);
  memcpy(values,x,sizeof(MYREAL)*numparam);
  myfree(v);
}
//#endif