/*
 * Copyright (C) 1996-2011 Daniel Waggoner and Tao Zha
 *
 * This 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 3 of the License, or
 * (at your option) any later version.
 *
 * It 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.
 *
 * If you did not received a copy of the GNU General Public License
 * with this software, see <http://www.gnu.org/licenses/>.
 */

#include "dw_metropolis_theta.h"
#include "dw_rand.h"
#include "dw_matrix_rand.h"
#include "dw_matrix_array.h"
#include "dw_std.h"

#include <stdlib.h>
#include <memory.h>
#include <math.h>
#include <time.h>


/*
   Assumes:
     model: pointer to valid TStateModel structure with model->theta a TMixture
            structure.
     submodel : a submodel of model->theta
     logposterior_old : log posterior value of current parameters

   Results:
     A metropolis draw of theta is stored in submodel.

   Returns:
     The log posterior value of current draw.
*/
PRECISION draw_metropolis_theta_prior(TStateModel *model, PRECISION logposterior_old)
{
  // TStateModel specific calls
  TMetropolis_theta *metropolis=model->metropolis_theta;

  int i, k, n=metropolis->n_parameters;
  TVector theta, theta_old;
  PRECISION logposterior;

  theta_old=CreateVector(n);
  theta=CreateVector(n);

  // TStateModel specific calls
  ConvertThetaToFreeParameters(model,pElementV(theta_old));

  EquateVector(theta,theta_old);

  for (k=i=0; i < metropolis->n_blocks; i++)
    {
      for ( ; k < metropolis->block_offsets[i+1]; k++)
        if (metropolis->direction[k])
	  UpdateV(1.0,theta,dw_gaussian_rnd()*ElementV(metropolis->scale,k),metropolis->direction[k]);
        else
          ElementV(theta,k)+=dw_gaussian_rnd()*ElementV(metropolis->scale,k);
      
      // TStateModel specific calls
      ConvertFreeParametersToTheta(model,pElementV(theta));
      logposterior=LogPrior(model);

      if (log(dw_uniform_rnd()) <= logposterior - logposterior_old)
        {
          metropolis->jumps[i]++;
	  if (i < metropolis->n_blocks-1)
	    {
	      EquateVector(theta_old,theta);
	      logposterior_old=logposterior;
	    }
        }
      else
	if (i < metropolis->n_blocks-1)
	  EquateVector(theta,theta_old);
	else
	  {
	    // TStateModel specific calls
	    ConvertFreeParametersToTheta(model,pElementV(theta_old));
	  }

      metropolis->draws[i]++;
    }

  FreeVector(theta);
  FreeVector(theta_old);

  return logposterior_old;
}

/*
   Assumes:
     model: pointer to valid TStateModel structure with model->theta a TMixture
            structure.

   Results:
     A metropolis draw of all parameters are stored in model.  
*/
void draw_prior_all_metropolis(TStateModel *model)
{
  // Draw transition matrix parameters and compute log prior
  DrawTransitionMatrixParametersFromPrior(model);

  // Draw theta
  draw_metropolis_theta_prior(model,LogPrior(model));
}

/* /\* */
/*    Assumes: */
/*      model: pointer to valid TStateModel structure. */

/*    Results: */
/*      A metropolis draw of theta is stored in model */
/* *\/ */
/* void Draw_metropolis_theta_prior(TStateModel *model) */
/* { */
/*   // TStateModel specific calls */
/*   TMetropolis_theta *metropolis=model->metropolis_theta; */

/*   int i, k, n=metropolis->n_parameters; */
/*   TVector theta, theta_old; */
/*   PRECISION logposterior, logposterior_old; */

/*   theta_old=CreateVector(n); */
/*   theta=CreateVector(n); */

/*   // TStateModel specific calls */
/*   logposterior_old=LogPrior(model); */
/*   ConvertThetaToFreeParameters(model,pElementV(theta_old)); */

/*   EquateVector(theta,theta_old); */

/*   for (k=i=0; i < metropolis->n_blocks; i++) */
/*     { */
/*       for ( ; k < metropolis->block_offsets[i+1]; k++) */
/*         if (metropolis->direction[k]) */
/* 	  UpdateV(1.0,theta,dw_gaussian_rnd()*ElementV(metropolis->scale,k),metropolis->direction[k]); */
/*         else */
/*           ElementV(theta,k)+=dw_gaussian_rnd()*ElementV(metropolis->scale,k); */
      
/*       // TStateModel specific calls */
/*       ConvertFreeParametersToTheta(model,pElementV(theta)); */
/*       logposterior=LogPrior(model); */

/*       if (log(dw_uniform_rnd()) <= logposterior - logposterior_old) */
/*         { */
/*           metropolis->jumps[i]++; */
/* 	  if (i < metropolis->n_blocks-1) */
/* 	    { */
/* 	      EquateVector(theta_old,theta); */
/* 	      logposterior_old=logposterior; */
/* 	    } */
/*         } */
/*       else */
/* 	if (i < metropolis->n_blocks-1) */
/* 	  EquateVector(theta,theta_old); */
/* 	else */
/* 	  // TStateModel specific calls */
/* 	  ConvertFreeParametersToTheta(model,pElementV(theta_old)); */

/*       metropolis->draws[i]++; */
/*     } */

/*   FreeVector(theta); */
/*   FreeVector(theta_old); */
/* } */

/* /\* */
/*    Assumes: */
/*      model: pointer to valid TStateModel structure. */

/*    Results: */
/*      A metropolis draw of theta */
/* *\/ */
/* void Draw_metropolis_theta_prior_single_block(int block, TStateModel *model) */
/* { */
/*   // TStateModel specific calls */
/*   TMetropolis_theta *metropolis=model->metropolis_theta; */

/*   int k, n=metropolis->n_parameters; */
/*   TVector theta, theta_old; */
/*   PRECISION logposterior, logposterior_old; */

/*   theta_old=CreateVector(n); */
/*   theta=CreateVector(n); */

/*   // TStateModel specific calls */
/*   logposterior_old=LogPrior(model); */
/*   ConvertThetaToFreeParameters(model,pElementV(theta_old)); */

/*   EquateVector(theta,theta_old); */

/*   for (k=metropolis->block_offsets[block]; k < metropolis->block_offsets[block+1]; k++) */
/*     if (metropolis->direction[k]) */
/*       UpdateV(1.0,theta,dw_gaussian_rnd()*ElementV(metropolis->scale,k),metropolis->direction[k]); */
/*     else */
/*       ElementV(theta,k)+=dw_gaussian_rnd()*ElementV(metropolis->scale,k); */
      
/*   // TStateModel specific calls */
/*   ConvertFreeParametersToTheta(model,pElementV(theta)); */
/*   logposterior=LogPrior(model); */

/*   if (log(dw_uniform_rnd()) <= logposterior - logposterior_old) */
/*     metropolis->jumps[block]++; */
/*   else */
/*     // TStateModel specific calls */
/*     ConvertFreeParametersToTheta(model,pElementV(theta_old)); */

/*   metropolis->draws[block]++; */

/*   FreeVector(theta); */
/*   FreeVector(theta_old); */
/* } */
/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/

void AdaptiveScale_metropolis_theta_prior(TStateModel *model, PRECISION mid, int period, int max_period, int verbose)
{
  TMetropolis_theta *metropolis=model->metropolis_theta;
  PRECISION new_scale, old_scale, rescale, logposterior;
  TAdaptive **adaptive;
  int  done=0, i, k, begin_time=(int)time((time_t*)NULL), count=0, check=period;

  adaptive=(TAdaptive**)dw_CreateArray_pointer(metropolis->n_blocks,dw_free);
  for (k=metropolis->n_blocks-1; k >= 0; k--)
    InitializeAdaptive(period,k,adaptive[k]=CreateAdaptive(mid),metropolis);

  if (verbose)
    printf("Beginning adaptive burn in -- maximum period = %d.\n",max_period);

  logposterior=LogPrior(model);
  done=0;
  while (!done)
    {
      count++;

      draw_prior_all_metropolis(model);

      if (count == check)
	{
          done=1;

          if (verbose)
	    printf("\n%d iterations completed - elapsed time: %d seconds\n",count,(int)time((time_t*)NULL) - begin_time);

	  for (k=metropolis->n_blocks-1; k >= 0; k--)
            {
              old_scale=adaptive[k]->scale;
              new_scale=RecomputeScale(metropolis->jumps[k],metropolis->draws[k],adaptive[k]);
              if (new_scale != old_scale)
                for (rescale=new_scale/old_scale, i=metropolis->block_offsets[k]; i < metropolis->block_offsets[k+1]; i++)
                  ElementV(metropolis->scale,i)*=rescale;
              if (adaptive[k]->period <= max_period) done=0;
              if (verbose) PrintJumpingInfo(k,adaptive[k],metropolis);
            }
   
          check+=period;
	}
    }

  for (i=k=0; k < metropolis->n_blocks; k++)
    if (adaptive[k]->scale != adaptive[k]->best_scale)
      for (rescale=adaptive[k]->best_scale/adaptive[k]->scale; i < metropolis->block_offsets[k+1]; i++)
        ElementV(metropolis->scale,i)*=rescale;

  dw_FreeArray(adaptive);
}

/* void AdaptiveScale_metropolis_theta_prior_single_block(int block, TStateModel *model, PRECISION mid, int period, int max_period, int verbose) */
/* { */
/*   TMetropolis_theta *metropolis=model->metropolis_theta; */
/*   TAdaptive *adaptive; */
/*   int  done=0, i, begin_time=(int)time((time_t*)NULL), count=0, check=period; */
/*   PRECISION new_scale, old_scale, rescale; */

/*   if ((block < 0) || (block >= metropolis->n_blocks))  */
/*     { */
/*       if (verbose) printf("AdaptiveScale_metropolis_theta_singleblock(): Invalid value for block\n"); */
/*       return; */
/*     } */

/*   InitializeAdaptive(period,block,adaptive=CreateAdaptive(mid),metropolis); */

/*   if (verbose) */
/*     printf("Beginning adaptive burn in for block %d -- maximum period = %d.\n",block,max_period); */

/*   while (!done) */
/*     { */
/*       count++; */
      
/*       DrawTransitionMatrixParametersFromPrior(model); */
/*       Draw_metropolis_theta_prior_single_block(block,model); */

/*       if (count == check) */
/*         { */
/*           if (verbose) */
/*             printf("\n%d iterations completed - elapsed time: %d seconds\n",count,(int)time((time_t*)NULL) - begin_time); */

/*           old_scale=adaptive->scale; */
/*           new_scale=RecomputeScale(metropolis->jumps[block],metropolis->draws[block],adaptive); */
/*           if (new_scale != old_scale) */
/*             for (rescale=new_scale/old_scale, i=metropolis->block_offsets[block]; i < metropolis->block_offsets[block+1]; i++) */
/*               ElementV(metropolis->scale,i)*=rescale; */
/*           if (adaptive->period > max_period) done=1; */
/*           if (verbose) PrintJumpingInfo(block,adaptive,metropolis); */
                                   
/*           check+=period; */
/* 	} */
/*     } */

/*   if (adaptive->scale != adaptive->best_scale) */
/*     for (rescale=adaptive->best_scale/adaptive->scale, i=metropolis->block_offsets[block];  */
/* 	 i < metropolis->block_offsets[block+1]; i++) */
/*       ElementV(metropolis->scale,i)*=rescale; */

/*   dw_free(adaptive); */
/* } */

/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/

/*
   Assumes
    model      : pointer to valid TStateModel structure
    center     : target value for acceptance ration
    peroid     : initial length of sub-runs to calibrate scales
    max_period : maximum length of sub-runs to calibrate scales
    verbose    : non-zero for output to stdout

   Results
    The values of model->metropolis_scale are calibrated to obtain the
    desired acceptance ratio. 
*/
void Calibrate_metropolis_theta_prior(TStateModel *model, PRECISION center, int period, int max_period, int verbose)
{
  int n_theta=NumberFreeParametersTheta(model), n_q=NumberFreeParametersQ(model);
  TVector initial_theta=(TVector)NULL, initial_q=(TVector)NULL; 

  // Save initial parameters
  if (n_theta) ConvertThetaToFreeParameters(model,pElementV(initial_theta=CreateVector(n_theta)));
  if (n_q) ConvertQToFreeParameters(model,pElementV(initial_q=CreateVector(n_q)));

  // Calibrate
  AdaptiveScale_metropolis_theta_prior(model,center,period,max_period,verbose);

  // Reset initial parameters
  if (n_theta) 
    {
      ConvertFreeParametersToTheta(model,pElementV(initial_theta));
      FreeVector(initial_theta);
    }
  if (n_q) 
    {
      ConvertFreeParametersToQ(model,pElementV(initial_q));
      FreeVector(initial_q);
    }
}

/* /\* */
/*    Assumes */
/*     idx        : 0 <= idx < NumberFreeParametersTheta(model) */
/*     model      : pointer to valid TStateModel structure */
/*     center     : target acceptance rate */
/*     peroid     : initial length of sub-runs to calibrate scale for direction idx */
/*     max_period : maximum length of sub-runs to calibrate scale for direction idx */
/*     verbose    : non-zero for output to stdout */

/*    Returns */
/*     The total number of draws made.  A single draw consists of a draw of the  */
/*     states, a draw of the transition matrix parameters, and a draw of theta.  The */
/*     draw of theta is only in the direction idx. */

/*    Results */
/*     The value of ElementV(model->metropolis_scale_theta,idx) is calibrated to  */
/*     obtain the desired acceptance ratio. */
/* *\/ */
/* int Calibrate_metropolis_theta_prior_single_direction(int idx, TStateModel *model, PRECISION center, int period, int max_period, int verbose) */
/* { */
/*   TMetropolis_theta *metropolis=model->metropolis_theta; */
/*   int *dims, n_draws, n_blocks, n_theta=NumberFreeParametersTheta(model), n_q=NumberFreeParametersQ(model); */
/*   TVector initial_theta=(TVector)NULL, initial_q=(TVector)NULL;  */

/*   // save initial parameters */
/*   if (n_theta) ConvertThetaToFreeParameters(model,pElementV(initial_theta=CreateVector(n_theta))); */
/*   if (n_q) ConvertQToFreeParameters(model,pElementV(initial_q=CreateVector(n_q))); */

/*   // Save old blocks */
/*   n_blocks=metropolis->n_blocks; */
/*   dims=(int*)dw_malloc(n_blocks*sizeof(int)); */
/*   memcpy(dims,metropolis->block_dims,n_blocks*sizeof(int)); */

/*   // Setup blocks  */
/*   Setup_metropolis_theta_blocks_full(metropolis); */

/*   // Calibrate */
/*   AdaptiveScale_metropolis_theta_prior_single_block(idx,model,center,period,max_period,verbose); */
/*   n_draws=metropolis->draws[idx]; */

/*   // Reset starting value */
/*   if (n_theta)  */
/*     { */
/*       ConvertFreeParametersToTheta(model,pElementV(initial_theta)); */
/*       FreeVector(initial_theta); */
/*     } */
/*   if (n_q)  */
/*     { */
/*       ConvertFreeParametersToQ(model,pElementV(initial_q)); */
/*       FreeVector(initial_q); */
/*     } */

/*   // Reset blocks */
/*   Setup_metropolis_theta_blocks(metropolis,n_blocks,dims); */
/*   dw_free(dims); */

/*   return n_draws; */
/* } */

/* /\* */
/*    Assumes */
/*     model        : pointer to valid TStateModel structure */
/*     center_s     : target acceptance rate for each direction invidually */
/*     peroid_s     : initial length of sub-runs to calibrate scale for each direction individually */
/*     max_period_s : maximum length of sub-runs to calibrate scale for each direction individually */
/*     center_a     : target acceptance rate for each block */
/*     peroid_a     : initial length of sub-runs to calibrate scales for each block */
/*     max_period_a : maximum length of sub-runs to calibrate scales for each block */
/*     verbose      : non-zero for output to stdout */

/*    Returns */
/*     The total number of theta draws made. */

/*    Results */
/*     The values of model->metropolis_scale_theta are calibrated to obtain the  */
/*     desired acceptance ratio. */
/* *\/ */
/* int Calibrate_metropolis_theta_prior_two_pass(TStateModel *model, PRECISION center_s,	int period_s, int max_period_s,  */
/* 					      PRECISION center_a, int period_a, int max_period_a, int verbose) */
/* { */
/*   TMetropolis_theta *metropolis=model->metropolis_theta; */
/*   int *dims, n_draws=0, n_blocks, i, n_theta=NumberFreeParametersTheta(model), n_q=NumberFreeParametersQ(model); */
/*   TVector initial_theta=(TVector)NULL, initial_q=(TVector)NULL;  */

/*   // save initial value */
/*   if (n_theta) ConvertThetaToFreeParameters(model,pElementV(initial_theta=CreateVector(n_theta))); */
/*   if (n_q) ConvertQToFreeParameters(model,pElementV(initial_q=CreateVector(n_q))); */

/*   // Individual direction run */
/*   if (period_s > 0) */
/*     { */
/*       // Save old blocks */
/*       n_blocks=metropolis->n_blocks; */
/*       memcpy(dims=(int*)dw_malloc(n_blocks*sizeof(int)),metropolis->block_dims,n_blocks*sizeof(int)); */

/*       // Setup blocks  */
/*       Setup_metropolis_theta_blocks_full(metropolis); */

/*       // Calibrate */
/*       for (i=0; i < n_theta; i++) */
/*         { */
/*           AdaptiveScale_metropolis_theta_prior_single_block(i,model,center_s,period_s,max_period_s,verbose); */
/*           n_draws+=metropolis->draws[i]; */
/*           if (n_theta) ConvertFreeParametersToTheta(model,pElementV(initial_theta)); */
/*           if (n_q) ConvertFreeParametersToQ(model,pElementV(initial_q)); */
/*         } */
    
/*       // Reset blocks */
/*       Setup_metropolis_theta_blocks(metropolis,n_blocks,dims); */
/*       dw_free(dims); */
/*     } */
/*   else */
/*     ResetCounts_metropolis_theta(metropolis); */

/*   // Final run */
/*   if (period_a > 0) */
/*     { */
/*       AdaptiveScale_metropolis_theta(model,center_a,period_a,max_period_a,verbose); */
/*       for (i=0; i < metropolis->n_blocks; i++) n_draws+=metropolis->draws[i]; */
/*       if (n_theta) ConvertFreeParametersToTheta(model,pElementV(initial_theta)); */
/*       if (n_q) ConvertFreeParametersToQ(model,pElementV(initial_q)); */
/*     } */

/*   if (n_theta) FreeVector(initial_theta); */
/*   if (n_q) FreeVector(initial_q); */

/*   return n_draws; */
/* } */
