/*
 * 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_equi_energy_simulation.h"
#include "dw_rand.h"

/*
   Frees all memory associated with TEquiEnergy structure.
*/
void FreeTEquiEnergy(TEquiEnergy *equi)
{
  if (equi)
    {
      FreeVector(equi->loglevels);
      FreeVector(equi->temperature);
      dw_free(equi->draw);
      dw_free(equi->new_draw);
      dw_free(equi);
    }
}

/*
   Creates and initializes TEquiEnergy structure.
*/
TEquiEnergy* CreateTEquiEnergy(TVector loglevels, TVector temperature, PRECISION p, int n, int n_theta, int level)
{
  TEquiEnergy *equi;
  if (!loglevels || !temperature || (DimV(loglevels) != DimV(temperature)) || (p < 0) || (p > 1.0) 
      || (level < 0) || (level >= DimV(loglevels)) || (n < 2))
    return (TEquiEnergy*)NULL;
  equi=(TEquiEnergy*)dw_malloc(sizeof(TEquiEnergy));
  equi->nlevels=DimV(loglevels);
  equi->loglevels=EquateVector((TVector)NULL,loglevels);
  equi->temperature=EquateVector((TVector)NULL,temperature);
  equi->p=p;
  equi->n=n;
  equi->n_theta=n_theta;
  equi->level=level;
  equi->draw=(PRECISION*)dw_malloc(n*sizeof(PRECISION));
  equi->new_draw=(PRECISION*)dw_malloc(n*sizeof(PRECISION));
  return equi;
}


/*
   Assumes
     loglevels - vector of length n in increasing order.
     logposteriorkernel - floating point.
     idx - starting point for search, 0 <= idx < n.

   Returns
     The index such that 
  
        loglevel(idx) >= logposteriorkernel > loglevel(idx-1)

     with the convention that loglevel(n) = infinity and logring(-1) = -infinity.
     The return value will be between 0 and n inclusive.
*/
int GetRingIndex(TVector loglevels, PRECISION logposteriorkernel, int idx)
{
  if (logposteriorkernel > loglevels(idx))
    {
      for (idx++; idx < n; idx++)
	if(loglevels(idx) >= logposteriorkernel) return idx;
    }
  else
    for ( ; idx > 0; idx--)
      if (logposteriorkernel > loglevels(idx-1)) return idx;
  return idx;
}

/*
   Gets from previous level and ring simulation.  Draw is stored in 
   equi->new_draw and equi->new_logposterior is set.  Returns 1 upon success and
   zero upon failure.  Call will fail if sdsm_get() fails.
*/
int GetDraw(TEquiEnergy *equi)
{
  sdsm_data_item *p;

  // Draw from correct level and ring. bin = (level - 1)*equi->nlevels + ring
  p=sdsm_get((equi->level - 1)*equi->nlevels + equi->ring);

  if (!p) return 0;

  memcpy(equi->new_draw,p->data,equi->n*sizeof(PRECISION));
  equi->new_logposteriorkernel=equi->new_draw[0];
  sdsm_release(p);
  return 1;
}

/*
   Puts draw in the appropriate level and ring.  Returns one upon success and
   zero upon failure.  Call will fail if an incorrect flag is passed or 
   sdsm_put() fails.  Sets equi->logposteriorkernel, equi->ring, and equi->draw,
   and parameters in model if necessary.
*/
int PutDraw(TEquiEnergy *equi, TStateModel *model, int flag)
{
  PRECISION *p;
  switch (flag)
    {
    case _EEFLAG_USEMODEL_:
      equi->draw[0]=equi->logposteriorkernel=LogPosterior_StatesIntegratedOut(model);
      equi->ring=GetRingIndex(equi->loglevels,equi->logposteriorkernel,equi->level);
      ConvertThetaToFreeParameters(model,equi->draw+1);
      ConvertQToFreeParameters(model,equi->draw+equi->n_theta+1);
      break;
    case _EEFLAG_USENEW_:
      equi->logposteriorkernel=equi->new_logposteriorkernel;
      p=equi->new_draw;
      equi->new_draw=equi->draw;
      equi->draw=p;
      ConvertFreeParametersToTheta(model,equi->draw+1);
      ConvertFreeParametersToQ(model,equi->draw+equi->n_theta+1);
      break;
    case _EEFLAG_USEOLD_:
      break;
    default:
      return 0;
    }
  // Write draw to correct level and ring. bin = level*equi->nlevels + ring
  sdsm_put(equi->level*equi->nlevels+equi->ring,equi->draw,equi->n,1.0);
  return 1;
}

/*
   Assumes
     equi  - valid pointer to TEquiEnergy structure.
     model - valid pointer to TStateModel structure.

   Results
     A new draw is made from the appropriate level.

   Returns
     one upon success and zero upon failure
*/
int draw_equi_energy(TEquiEnergy *equi, TStateModel *model)
{
  int flag=_EEFLAG_USEOLD_;
  if ((equi->level == 0) || (dw_uniform_rnd() > equi->p))
    {
      // Draw from posterior
      DrawAll(model);
      flag=_EEFLAG_USEMODEL_;
    }
  else
    {
      // Draw from correct level and ring = (level - 1)*equi->nlevels + ring
      GetDraw(equi);

      // Accept draw?
      if (equi->ring < equi->level)
	{
	  if (log(dw_uniform_rnd())/(1.0/ElementV(equi->temperatures,equi->level) - 1.0/ElementV(equi->temperatures,equi->level-1))
	      < equi->new_logposteriorkernel - equi->logposteriorkernel)
	    flag=_EEFLAG_USENEW_;
	}
      else
	if (equi->ring == equi->level)
	  {
	    if (ElementV(equi->temperatures,equi->level)*log(dw_uniform_rnd()) < logposteriorkernel - equi->logposteriorkernel)
	      flag=_EEFLAG_USENEW_;
	  }
	else
	  flag=_EEFLAG_USENEW_;
    }

  // Write draw to correct level and ring = level*equi->nlevels + ring
  return PutDraw(equi,model,flag);
}

int EquiEnergySimulationRun(int count, TStateMdoel *model, TVector loglevels, TVector temperatures, PRECISION p, int level)
{
  TEquiEnergy *equi=CreateTEquiEnergy(loglevels,temperatures,p,NumberFreeParametersTheta(model)+NumberFreeParametersQ(model)+1,
				      NumberFreeParametersTheta(model),level);

  if (!equi) return 0;

  for ( ; count > 0; ) count-=draw_equi_energy(equi,model);

  FreeTEquiEnergy(equi);

  return 1;
}

int dw_equi_energy_command_line(int nargs, char **args, TStateModel *model)
{
  if (!(model_id=ModelType(model)) || strcmp(model_id,"T_MSStateSpace"))
    {
      printf("dw_state_space_forecast_command_line(): Model type = %s\n",model_id ? model_id : "null");
      return 0;
    }



}
