/*
 * 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_switch.h"
#include "dw_switchio.h"
#include "dw_MSStateSpace.h"
#include "dw_histogram.h"
#include "dw_parse_cmd.h"
#include "dw_ascii.h"
#include "dw_error.h"
#include "dw_math.h"
#include "dw_std.h"

#include <stdlib.h>
#include <string.h>
#include <math.h>

/*
   Normalizes variance decomposition.  Converts cummulative contributions to the 
   total variance into cummulative percentage contributions.
*/
void NormalizeVarianceDecomposition(TMatrix VD, int nvars)
{
  PRECISION tmp;
  int nshocks=ColM(VD)/nvars, i, j, t;

  // normalizes to percentages
  for (t=RowM(VD)-1; t >= 0; t--)
    for (i=nvars-1; i >= 0; i--)
      {
	for (tmp=0.0, j=nshocks-1; j >= 0; j--) tmp+=ElementM(VD,t,j+i*nshocks);
	if (tmp > 0)
	  for (tmp=1.0/tmp, j=nshocks-1; j >= 0; j--) ElementM(VD,t,j+i*nshocks)*=tmp;
      }
}

/*
   Assumes
    draws : number of draws of the regime path to make for each posterior draw
    posterior_file : FILE pointer to file containing posterior draws.  If null, 
                     current parameters are used.
    thin : thinning factor to use if posterior_file is not null.
    horizon : non-negative integer
    ergodic : if non-zero, uses the ergodic distribution to draw the initial 
              regime, otherwise uses the filtered probabilities at time nobs.
    model : point to valid TStateModel structure
    flag : equal to IR_STATES or IR_OBSERVABLES

   Results:
    Creates matrix containing the mean variance decompositions.

   Returns:
    Pointer to matrix upon success and null pointer on failure.

  Notes:
     The element in position (t,j+i*nshocks) of returned matrix is the 
     cummulative contribution of the jth shock to the total variance of ith 
     variable to the  at horizon t. Horizon zero is the contemporaneous 
     contribution.

     If flag is equal to IR_STATES then the variables are the state variables and
     the shocks are the fundamental shocks.  If flag is equal to IR_OBSERVABLES
     then the variables are the observables and the shocks are both the 
     fundamental shocks and the measurement error.  The fundamental shocks are 
     ordered first.
*/
TMatrix state_space_variance_decomposition_mean(int draws, FILE *posterior_file, int thin, int ergodic, int horizon, TStateModel *model, int flag)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int done=0, *S, i, j, k, m, n=1000, T=model->nobs, count=0, nvars, nshocks;
  TMatrix vd, tvd;
  TVector init_prob, prob;

  // set sizes
  if (flag & IR_STATES)
    {
      nvars=statespace->nz;
      nshocks=statespace->nepsilon;
    }
  else
    {
      nvars=statespace->ny;
      nshocks=statespace->nepsilon + statespace->nu;
    }

  // allocate memory
  S=(int*)dw_malloc(horizon*sizeof(int));
  vd=CreateMatrix(horizon,nvars*nshocks);
  InitializeMatrix(tvd=CreateMatrix(horizon,nvars*nshocks),0.0);
  init_prob=CreateVector(model->sv->nbasestates);
  prob=CreateVector(model->sv->nbasestates);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!dw_ReadPosteriorDraw(posterior_file,model))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	  // Use ergodic initial probability or filtered probability at time T
	  if (ergodic)
	    Ergodic(init_prob,model->sv->baseQ);
	  else
	    for (j=model->sv->nstates-1; j >= 0; j--)
	      ElementV(init_prob,j)=ProbabilityBaseStateConditionalCurrent(j,T,model);

	  for (k=draws; k > 0; k--)
	    {
	      // Draw time T regime
	      m=DrawDiscrete(init_prob);
              
	      // Draw regimes from time T+1 through T+h inclusive
	      for (j=0; j < horizon; j++)
		{
		  ColumnVector(prob,model->sv->baseQ,m);
		  S[j]=m=DrawDiscrete(prob);
		}

	      // Compute variance decomposition contributions
	      dw_state_space_variance_decomposition(vd,S,model,flag);

	      // Accumulate variance decomposition contributions
	      AddMM(tvd,tvd,vd);

	      count++;
	    }
	}
    }

  FreeMatrix(vd);
  dw_free(S);
  FreeVector(prob);
  FreeVector(init_prob);

  ProductMS(tvd,tvd,1.0/(PRECISION)count);
  return tvd;
}

/*
   Assumes
    posterior_file : FILE pointer to file containing posterior draws.  If null, 
                     current parameters are used.
    thin : thinning factor to use if posterior_file is not null.
    s : fixed regime
    horizon : non-negative integer
    model : point to valid TStateModel structure
    flag : determines the type of posterior draws file passed.

   Results:
    Creates matrix containing the mean variance decompositions.

   Returns:
    Pointer to matrix upon success and null pointer on failure.

  Notes:
     The element in position (k,i+j*nvars) of tvd is the cummulative contributions
     to the total variance of ith variable to the jth shock at horizon k. Horizon 
     zero is the contemporaneous contribution.
*/
TMatrix state_space_variance_decomposition_mean_regime(FILE *posterior_file, int thin, int s, int horizon, TStateModel *model, int flag)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int done=0, *S, i, n=1000, count=0, nvars, nshocks;
  TMatrix tvd, vd;

  // set sizes
  if (flag & IR_STATES)
    {
      nvars=statespace->nz;
      nshocks=statespace->nepsilon;
    }
  else
    {
      nvars=statespace->ny;
      nshocks=statespace->nepsilon + statespace->nu;
    }

  // allocate memory
  S=(int*)dw_malloc(horizon*sizeof(int));
  for (i=0; i < horizon; i++) S[i]=s;
  vd=CreateMatrix(horizon,nvars*nshocks);
  InitializeMatrix(tvd=CreateMatrix(horizon,nvars*nshocks),0.0);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!dw_ReadPosteriorDraw(posterior_file,model))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	  // Compute variance decomposition contributions
	  dw_state_space_variance_decomposition(vd,S,model,flag);

	  // Accumulate variance decomposition contributions
	  AddMM(tvd,tvd,vd);

	  count++;
	}
    }

  FreeMatrix(vd);
  dw_free(S);

  ProductMS(tvd,tvd,1.0/(PRECISION)count);
  return tvd;
}

/*
   Assumes
    draws : number of draws of the regime path to make for each posterior draw
    posterior_file : FILE pointer to file containing posterior draws.  If null, 
                     current parameters are used.
    thin : thinning factor to use if posterior_file is not null.
    horizon : non-negative integer
    ergodic : if non-zero, uses the ergodic distribution to draw the initial 
              regime, otherwise uses the filtered probabilities at time nobs.
    model : point to valid TStateModel structure
    flag : determines the type of posterior draws file passed.

   Results:
    Creates matrix histogram containing the cummulative percentage contributions 
    to the total variance.

   Returns:
    Pointer to histogram upon success and null pointer on failure.

   Notes:
     The element in position (k,i+j*nvars) is the cummulative contribution to the
     total variance of ith variable to the jth shock at horizon k. Horizon zero
     is the contemporaneous contribution.
*/
TMatrixHistogram *state_space_variance_decomposition_percentiles(int draws, FILE *posterior_file, int thin, int ergodic, int horizon, TStateModel *model, int flag)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int done=0, *S, i, j, k, m, n=1000, T=model->nobs, nvars, nshocks;
  TMatrix vd;
  TVector init_prob, prob;
  TMatrixHistogram *histogram;

  // set sizes
  if (flag & IR_STATES)
    {
      nvars=statespace->nz;
      nshocks=statespace->nepsilon;
    }
  else
    {
      nvars=statespace->ny;
      nshocks=statespace->nepsilon + statespace->nu;
    }

  // allocate memory
  S=(int*)dw_malloc(horizon*sizeof(int));
  vd=CreateMatrix(horizon,nvars*nshocks);
  histogram=CreateMatrixHistogram(horizon,nvars*nshocks,100,HISTOGRAM_VARIABLE);
  init_prob=CreateVector(model->sv->nbasestates);
  prob=CreateVector(model->sv->nbasestates);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!dw_ReadPosteriorDraw(posterior_file,model))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	  // Use ergodic initial probability or filtered probability at time T
	  if (ergodic)
	    Ergodic(init_prob,model->sv->baseQ);
	  else
	    for (j=model->sv->nstates-1; j >= 0; j--)
	      ElementV(init_prob,j)=ProbabilityBaseStateConditionalCurrent(j,T,model);

	  for (k=draws; k > 0; k--)
	    {
	      // Draw time T regime
	      m=DrawDiscrete(init_prob);
              
	      // Draw regimes from time T+1 through T+h inclusive
	      for (j=0; j < horizon; j++)
		{
		  ColumnVector(prob,model->sv->baseQ,m);
		  S[j]=m=DrawDiscrete(prob);
		}

	      // Compute variance decomposition contributions
	      dw_state_space_variance_decomposition(vd,S,model,flag);

	      // Convert to percentages
	      NormalizeVarianceDecomposition(vd,nvars);

	      // Accumulate variance decomposition percentages
	      AddMatrixObservation(vd,histogram);
	    }
	}
    }

  FreeMatrix(vd);
  dw_free(S);
  FreeVector(prob);
  FreeVector(init_prob);

  return histogram;
}

/*
   Assumes
    posterior_file : FILE pointer to file containing posterior draws.  If null, 
                     current parameters are used.
    thin : thinning factor to use if posterior_file is not null.
    s : fixed regime
    horizon : non-negative integer
    model : point to valid TStateModel structure
    flag : determines the type of posterior draws file passed.

   Results:
    Creates matrix histogram containing the cummulative percentage contributions 
    to the total variance.

   Returns:
    Pointer to histogram upon success and null pointer on failure.

   Notes:
     The element in position (k,i+j*nvars) is the cummulative contribution to the
     total variance of ith variable to the jth shock at horizon k. Horizon zero
     is the contemporaneous contribution.
*/
TMatrixHistogram *state_space_variance_decomposition_percentiles_regime(FILE *posterior_file, int thin, int s, int horizon, TStateModel *model, int flag)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int nvars, nshocks, done=0, *S, i, n=1000;
  TMatrix vd;
  TMatrixHistogram *histogram;

 // set sizes
  if (flag & IR_STATES)
    {
      nvars=statespace->nz;
      nshocks=statespace->nepsilon;
    }
  else
    {
      nvars=statespace->ny;
      nshocks=statespace->nepsilon + statespace->nu;
    }

  // allocate memory
  S=(int*)dw_malloc(horizon*sizeof(int));
  for (i=0; i < horizon; i++) S[i]=s;
  vd=CreateMatrix(horizon,nvars*nshocks);
  histogram=CreateMatrixHistogram(horizon,nvars*nshocks,100,HISTOGRAM_VARIABLE);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!dw_ReadPosteriorDraw(posterior_file,model))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	  // Compute variance decomposition contributions
	  dw_state_space_variance_decomposition(vd,S,model,flag);

	  // Convert to percentages
	  NormalizeVarianceDecomposition(vd,nvars);

	  // Accumulate variance decomposition percentages
	  AddMatrixObservation(vd,histogram);
	}
    }

  FreeMatrix(vd);
  dw_free(S);

  return histogram;
}
