/*
 * 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/>.
 */

#ifndef __DIRICHLET_RESTRICTIONS__
#define __DIRICHLET_RESTRICTIONS__

#ifdef __cplusplus
extern "C" {
#endif

#include "dw_matrix.h"
#include "dw_switch.h"


typedef struct
{
  //=== Number of states and quasi-free variables ===
  int nstates;                     // number states
  int total_Dirichlet_parameters;  // total_Dirichlet_parameters = sv->nfree + dw_DimA(DirichletDim)

  //=== Workspace for Dirichlet parameters ===
  PRECISION *B;               // Workspace for the dirichlet parameters. Dimension of B is total_dirichlet_parameters.
  PRECISION **b;              // Dirichlet variables.  Points to the buffer B.  

  //=== Prior information ===
  TMatrix  Prior;             // Restricted prior on the columns of Q.  Must be nstates x nstates with positive elements.
  PRECISION *Prior_B;         // Workspace for prior on the Dirichlet parameters
  PRECISION **Prior_b;        // Dirichlet prior parameters.  Points to the buffer Prior_B.

  //=== Restrictions ===
  int *DirichletDim;          // DirichletDim[k] = number Dirichlet parameters for the kth Dirichlet random variable
  int **NonZeroIndex;         // nstates x nstates : Q(i,j) = (NonZeroIndex[i][j] >= 0) ? B[NonZeroIndex[i][j]]*MQ(i,j) : 0.0; 
  TMatrix MQ;                 // nstates x nstates : Q(i,j) = (NonZeroIndex[i][j] >= 0) ? B[NonZeroIndex[i][j]]*MQ(i,j) : 0.0; 

  //=== Control variables ===
  int UseErgodic;

  //=== Workspace ===
  PRECISION LogPriorConstant;
  PRECISION scale;

} TDirichletRestrictions;

void FreeDirichletRestrictions(TDirichletRestrictions *restrictions);
TDirichletRestrictions* CreateDirichletRestrictions(int nstates, TMatrix Prior, int* DirichletDim, 
						    int** NonZeroIndex, TMatrix MQ);
TDirichletRestrictions* DuplicateDirichletRestrictions(TDirichletRestrictions *restrictions);

TMarkovStateVariable* CreateMarkovStateVariable_Mixture(int nstates, int nlags_encoded, TMatrix Prior);
TMarkovStateVariable* CreateMarkovStateVariable_NoRestrictions(int nstates, int nlags_encoded, TMatrix Prior);
TMarkovStateVariable* CreateMarkovStateVariable_Exclusion(int nstates, int nlags_encoded, TMatrix Prior, TMatrix Exclusion);
TMarkovStateVariable* CreateMarkovStateVariable_SimpleRestrictions(int nstates, int nlags_encoded, TMatrix Prior, TMatrix* Restrictions);
TMarkovStateVariable* CreateMarkovStateVariable_ConstantState(int nlags_encoded);


//====== QRoutines ======//
QRoutines* dirichlet_restrictions_CreateQRoutines(void);

int dirichlet_restrictions_ComputeTransitionMatrix(int t, TMarkovStateVariable *sv, TStateModel *model);
PRECISION dirichlet_restrictions_LogPrior_q( TMarkovStateVariable *sv);
int dirichlet_restrictions_Probability_s0(TVector p, TMarkovStateVariable *sv);

PRECISION dirichlet_restrictions_DrawProposalTransitionMatrixParameters(int *S, TMarkovStateVariable *sv, TStateModel *model);
int dirichlet_restrictions_DrawTransitionMatrixParameters(int *S, TMarkovStateVariable *sv, TStateModel *model);
int dirichlet_restrictions_DrawTransitionMatrixParametersFromPrior(TMarkovStateVariable *sv);

int dirichlet_restrictions_DefaultTransitionMatrixParameters(TMarkovStateVariable *sv, TStateModel *model);
int dirichlet_restrictions_Update_q_from_BaseTransitionMatrix(int t, TMarkovStateVariable *sv, TStateModel *model);

//====== Utility Routines ======//
TMatrix dirichlet_restrictions_CreatePriorFromDuration(int nbasestates, PRECISION duration);
void dirichlet_restrictions_SetLogPriorConstant(TDirichletRestrictions *restrictions);
void dirichlet_restrictions_SetTransitionMatrixParametersToPriorMean(PRECISION *q, TDirichletRestrictions *restrictions);
int dirichlet_restrictions_Populate_B(PRECISION *q, TDirichletRestrictions *restrictions);
void dirichlet_restrictions_Populate_q(PRECISION *q, TDirichletRestrictions *restrictions);
int dirichlet_restrictions_CheckRestrictions(int* DirichletDim, int** NonZeroIndex, TMatrix MQ, TMatrix Prior, int nstates);

//====== Routines to simulate and evalute Dirichlet density ======//
void DrawDirichlet(PRECISION *b, PRECISION *alpha, int n);
void DrawDirichlet_free(PRECISION *q, PRECISION *alpha, int n);
PRECISION DrawDirichlet_kernel(PRECISION *b, PRECISION *alpha, int n);
PRECISION LogDirichlet_pdf(PRECISION *q, PRECISION *alpha, int n);
PRECISION LogDirichlet_kernel(PRECISION *q, PRECISION *alpha, int n);
void DrawIndependentDirichlet_free(PRECISION *q, PRECISION *alpha, int *dirichlet_dims, int m);
PRECISION LogIndependentDirichlet_pdf(PRECISION *q, PRECISION *alpha, int *dirichlet_dims, int m);

  
#ifdef __cplusplus
}
#endif  
#endif
