/*
 * 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_dirichlet_restrictions.h"
#include "dw_switch.h"
#include "dw_switchio.h"
#include "dw_matrix_array.h"
#include "dw_error.h"
#include "dw_array.h"
#include "dw_rand.h"
#include "dw_std.h"

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


void FreeDirichletRestrictions(TDirichletRestrictions *restrictions)
{
  if (restrictions)
    {
      if (restrictions->B) dw_free(restrictions->B);
      if (restrictions->b) dw_free(restrictions->b);
      if (restrictions->Prior_B) dw_free(restrictions->Prior_B);
      if (restrictions->Prior_b) dw_free(restrictions->Prior_b);
      FreeMatrix(restrictions->Prior);
      dw_FreeArray(restrictions->DirichletDim);
      dw_FreeArray(restrictions->NonZeroIndex);
      FreeMatrix(restrictions->MQ);
      dw_free(restrictions);
    }
}

/*
   Assumes
    nstates      : positive integer
    Prior        : nstates x nstates matrix
    DirichletDim      : integer array
    NonZeroIndex : nstates x nstates integer matrix
    MQ           : nstates x nstates matrix

   Returns
    A valid pointer to a TDirichletRestrictions structure.

   Notes
    This is the basic constructor for the TDirichletRestrictions structure.  Upon
    error, this procedure terminates the program.
*/
TDirichletRestrictions* CreateDirichletRestrictions(int nstates, TMatrix Prior, int* DirichletDim, 
						    int** NonZeroIndex, TMatrix MQ)
{
  TDirichletRestrictions *restrictions;
  int i, j, k, q, total_Dirichlet_parameters=0, terminal_errors;

  if ((nstates <= 0))
    {
      printf("CreateMarkovStateVariable(): number of states must be positive\n");
      dw_exit(0);
    }

  if (!dirichlet_restrictions_CheckRestrictions(DirichletDim,NonZeroIndex,MQ,Prior,nstates))
    {
      printf("CreateMarkovStateVariable_Single(): Error in restrictions\n");
      dw_exit(0);
    }

  if (!(restrictions=(TDirichletRestrictions*)dw_malloc(sizeof(TDirichletRestrictions))))
    {
      printf("CreateMarkovStateVariable(): out of memory\n");
      dw_exit(0);
    }

  //=== Set to terminate on memory error ===
  terminal_errors=dw_SetTerminalErrors(MEM_ERR);

  //=== Compute total number of quasi-free transition matrix parameters
  for (k=dw_DimA(DirichletDim)-1; k >= 0; k--) total_Dirichlet_parameters+=DirichletDim[k];

  // Workspace for quasi-free variables
  restrictions->B=(PRECISION*)dw_malloc(total_Dirichlet_parameters*sizeof(PRECISION));
  restrictions->b=(PRECISION**)dw_malloc(dw_DimA(DirichletDim)*sizeof(PRECISION*));
  for (q=k=0; k < dw_DimA(DirichletDim); k++)
    {
      restrictions->b[k]=restrictions->B + q;
      q+=DirichletDim[k];
    }

  //=== Sizes ===
  restrictions->nstates=nstates;
  restrictions->total_Dirichlet_parameters=total_Dirichlet_parameters;

  //=== Prior information ===
  restrictions->Prior=EquateMatrix((TMatrix)NULL,Prior);

  restrictions->Prior_B=(PRECISION*)dw_malloc(total_Dirichlet_parameters*sizeof(PRECISION));
  for (i=total_Dirichlet_parameters-1; i >= 0; i--) restrictions->Prior_B[i]=1.0;
  for (j=nstates-1; j >= 0; j--)
    for (i=nstates-1; i >= 0; i--)
      if ((k=NonZeroIndex[i][j]) >= 0)
	restrictions->Prior_B[k]+=ElementM(Prior,i,j)-1.0;

  restrictions->Prior_b=(PRECISION**)dw_malloc(dw_DimA(DirichletDim)*sizeof(PRECISION*));
  for (q=k=0; k < dw_DimA(DirichletDim); k++)
    {
      restrictions->Prior_b[k]=restrictions->Prior_B + q;
      q+=DirichletDim[k];
    }

  //=== Restrictions ===
  restrictions->DirichletDim=(int*)dw_CopyArray(NULL,DirichletDim);
  restrictions->NonZeroIndex=(int**)dw_CopyArray(NULL,NonZeroIndex);
  restrictions->MQ=EquateMatrix((TMatrix)NULL,MQ);

  //=== Control variables ===
  restrictions->UseErgodic=0;

  //=== Dirichlet scale
  restrictions->scale=1;

  //=== Set Constants ===
  dirichlet_restrictions_SetLogPriorConstant(restrictions);

  //=== Reset terminal errors ===
  dw_SetTerminalErrors(terminal_errors);

  return restrictions;
}

TDirichletRestrictions* DuplicateDirichletRestrictions(TDirichletRestrictions *restrictions)
{
  int q, k;
  TDirichletRestrictions *new_restrictions;

  if (!(new_restrictions=(TDirichletRestrictions*)dw_malloc(sizeof(TDirichletRestrictions)))) return (TDirichletRestrictions*)NULL;
  memcpy(new_restrictions,restrictions,sizeof(TDirichletRestrictions));

  //=== Quasi-free workspace ===
  new_restrictions->B=(PRECISION*)dw_malloc(restrictions->total_Dirichlet_parameters*sizeof(PRECISION));
  new_restrictions->b=(PRECISION**)dw_malloc(dw_DimA(restrictions->DirichletDim)*sizeof(PRECISION*));
  for (q=k=0; k < dw_DimA(restrictions->DirichletDim); k++)
    {
      new_restrictions->b[k]=new_restrictions->B + q;
      q+=restrictions->DirichletDim[k];
    }
  memcpy(new_restrictions->B,restrictions->B,restrictions->total_Dirichlet_parameters*sizeof(PRECISION));

  //=== Prior information ===
  new_restrictions->Prior=EquateMatrix((TMatrix)NULL,restrictions->Prior);
  new_restrictions->Prior_B=(PRECISION*)dw_malloc(restrictions->total_Dirichlet_parameters*sizeof(PRECISION));
  new_restrictions->Prior_b=(PRECISION**)dw_malloc(dw_DimA(restrictions->DirichletDim)*sizeof(PRECISION*));
  for (q=k=0; k < dw_DimA(restrictions->DirichletDim); k++)
    {
      new_restrictions->Prior_b[k]=new_restrictions->Prior_B + q;
      q+=restrictions->DirichletDim[k];
    }
  memcpy(new_restrictions->Prior_B,restrictions->Prior_B,restrictions->total_Dirichlet_parameters*sizeof(PRECISION));

  //=== Restrictions ===
  new_restrictions->DirichletDim=(int*)dw_CopyArray(NULL,restrictions->DirichletDim);
  new_restrictions->NonZeroIndex=(int**)dw_CopyArray(NULL,restrictions->NonZeroIndex);
  new_restrictions->MQ=EquateMatrix((TMatrix)NULL,restrictions->MQ);

  return new_restrictions;
}

TMarkovStateVariable* CreateMarkovStateVariable_Mixture(int nstates, int nlags_encoded, TMatrix Prior)
{
  int i, j;
  TDirichletRestrictions *restrictions;
  int* DirichletDim;
  int** NonZeroIndex;
  TMatrix MQ;
  NonZeroIndex=dw_CreateRectangularArray_int(nstates,nstates);
  for (i=nstates-1; i >= 0; i--)
    for (j=nstates-1; j >= 0; j--)
      NonZeroIndex[i][j]=i;
  MQ=InitializeMatrix(CreateMatrix(nstates,nstates),1.0);
  DirichletDim=dw_CreateArray_int(1);
  DirichletDim[0]=nstates;
  restrictions=CreateDirichletRestrictions(nstates,Prior,DirichletDim,NonZeroIndex,MQ);
  dw_FreeArray(NonZeroIndex);
  FreeMatrix(MQ);
  dw_FreeArray(DirichletDim);
  return CreateMarkovStateVariable_Single(nstates,restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim),
					  nlags_encoded,dirichlet_restrictions_CreateQRoutines(),(void*)restrictions);
}

TMarkovStateVariable* CreateMarkovStateVariable_NoRestrictions(int nstates, int nlags_encoded, TMatrix Prior)
{
  int i, j;
  TDirichletRestrictions *restrictions;
  int* DirichletDim;
  int** NonZeroIndex;
  TMatrix MQ;
  NonZeroIndex=dw_CreateRectangularArray_int(nstates,nstates);
  for (i=nstates-1; i >= 0; i--)
    for (j=nstates-1; j >= 0; j--)
      NonZeroIndex[i][j]=i+nstates*j;
  InitializeMatrix(MQ=CreateMatrix(nstates,nstates),1.0);
  DirichletDim=dw_CreateArray_int(nstates);
  for (i=nstates-1; i >= 0; i--) DirichletDim[i]=nstates;
  restrictions=CreateDirichletRestrictions(nstates,Prior,DirichletDim,NonZeroIndex,MQ);
  dw_FreeArray(NonZeroIndex);
  FreeMatrix(MQ);
  dw_FreeArray(DirichletDim);
  return CreateMarkovStateVariable_Single(nstates,restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim),
					  nlags_encoded,dirichlet_restrictions_CreateQRoutines(),(void*)restrictions);
}

TMarkovStateVariable* CreateMarkovStateVariable_Exclusion(int nstates, int nlags_encoded, TMatrix Prior, TMatrix Exclusion)
{
  int i, j, k, q;
  TDirichletRestrictions *restrictions;
  int* DirichletDim;
  int** NonZeroIndex;
  TMatrix MQ;
  dw_InitializeArray_int(NonZeroIndex=dw_CreateRectangularArray_int(nstates,nstates),-1);
  MQ=InitializeMatrix(CreateMatrix(nstates,nstates),0.0);
  DirichletDim=dw_CreateArray_int(nstates);
  for (k=j=0; j < nstates; j++)
    {
      for (q=i=0; i < nstates; i++)
	if (ElementM(Exclusion,i,j) > 0)
	  {
	    NonZeroIndex[i][j]=k++;
	    ElementM(MQ,i,j)=1.0;
	    q++;
	  }
      DirichletDim[j]=q;
    }
  restrictions=CreateDirichletRestrictions(nstates,Prior,DirichletDim,NonZeroIndex,MQ);
  dw_FreeArray(NonZeroIndex);
  FreeMatrix(MQ);
  dw_FreeArray(DirichletDim);
  return CreateMarkovStateVariable_Single(nstates,restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim),
					  nlags_encoded,dirichlet_restrictions_CreateQRoutines(),(void*)restrictions);
}

TMarkovStateVariable* CreateMarkovStateVariable_SimpleRestrictions(int nstates, int nlags_encoded, TMatrix Prior, TMatrix* Restriction)
{
  int i, j, k, q;
  TDirichletRestrictions *restrictions;
  int* DirichletDim;
  int** NonZeroIndex;
  TMatrix MQ;
  dw_InitializeArray_int(NonZeroIndex=dw_CreateRectangularArray_int(nstates,nstates),-1);
  InitializeMatrix(MQ=CreateMatrix(nstates,nstates),0.0);
  DirichletDim=dw_CreateArray_int(nstates);
  for (q=k=0; k < nstates; k++)
    {
      DirichletDim[k]=ColM(Restriction[k]);
      for (i=0; i < nstates; i++)
        {
          for (j=0; j < ColM(Restriction[k]); j++)
            if (ElementM(Restriction[k],i,j) > 0)
              {
                NonZeroIndex[i][k]=q+j;
                ElementM(MQ,i,k)=ElementM(Restriction[k],i,j);
                break;
              }
        }
      q+=DirichletDim[k];
    }
  restrictions=CreateDirichletRestrictions(nstates,Prior,DirichletDim,NonZeroIndex,MQ);
  dw_FreeArray(NonZeroIndex);
  FreeMatrix(MQ);
  dw_FreeArray(DirichletDim);
  return CreateMarkovStateVariable_Single(nstates,restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim),
					  nlags_encoded,dirichlet_restrictions_CreateQRoutines(),(void*)restrictions);
}

TMarkovStateVariable* CreateMarkovStateVariable_ConstantState(int nlags_encoded)
{
  TDirichletRestrictions *restrictions;
  int* DirichletDim;
  int** NonZeroIndex;
  TMatrix MQ, Prior;
  dw_InitializeArray_int(NonZeroIndex=dw_CreateRectangularArray_int(1,1),0);
  InitializeMatrix(MQ=CreateMatrix(1,1),1.0);
  dw_InitializeArray_int(DirichletDim=dw_CreateArray_int(1),1);
  InitializeMatrix(Prior=CreateMatrix(1,1),1.0);
  restrictions=CreateDirichletRestrictions(1,Prior,DirichletDim,NonZeroIndex,MQ);
  dw_FreeArray(NonZeroIndex);
  FreeMatrix(MQ);
  dw_FreeArray(DirichletDim);
  FreeMatrix(Prior);
  return CreateMarkovStateVariable_Single(1,restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim),
					  nlags_encoded,dirichlet_restrictions_CreateQRoutines(),(void*)restrictions);
}

/*******************************************************************************/
/*********************** TDirichletRestrictions Routines ***********************/
/*******************************************************************************/
QRoutines* dirichlet_restrictions_CreateQRoutines(void)
{
  QRoutines* pRoutines=CreateQRoutines_empty();
  
  if (pRoutines)
    {
      pRoutines->ComputeTransitionMatrix=dirichlet_restrictions_ComputeTransitionMatrix;
      pRoutines->pLogPrior=dirichlet_restrictions_LogPrior_q;
      pRoutines->Probability_s0=dirichlet_restrictions_Probability_s0;

      pRoutines->DrawTransitionMatrixParameters=dirichlet_restrictions_DrawTransitionMatrixParameters;
      pRoutines->DrawProposalTransitionMatrixParameters=dirichlet_restrictions_DrawProposalTransitionMatrixParameters;
      pRoutines->pDrawTransitionMatrixParametersFromPrior=dirichlet_restrictions_DrawTransitionMatrixParametersFromPrior;

      pRoutines->FreeInfo=(void (*)(void*))FreeDirichletRestrictions;
      pRoutines->DuplicateInfo=(void* (*)(void*))DuplicateDirichletRestrictions;

      pRoutines->WriteSpecification=WriteSpecification_dirichlet_restrictions;

      pRoutines->DefaultTransitionMatrixParameters=dirichlet_restrictions_DefaultTransitionMatrixParameters;
      pRoutines->Update_q_from_BaseTransitionMatrix=dirichlet_restrictions_Update_q_from_BaseTransitionMatrix;
    }
  
  return pRoutines;
}

//------------------------------ Required routines ------------------------------

/*
   Assumes:
     t     : ignored
     sv    : pointer to valid TMarkovStateVariable structure.
     model : ignored

   Results:
     Computes the all transition matrices from the free parameters sv->q.  Sets 
     sv->t0 and sv->t1.

   Returns:
     Pointer to sv->baseQ upon success and null otherwise.
*/
int dirichlet_restrictions_ComputeTransitionMatrix(__attribute__ ((unused)) int t, TMarkovStateVariable *sv, __attribute__ ((unused)) TStateModel *model)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  int i, j, k;

  if (!dirichlet_restrictions_Populate_B(sv->q,restrictions)) 
    {
      sv->t1=-1;
      return 0;
    }

  for (j=sv->nbasestates-1; j >= 0; j--)
    for (i=sv->nbasestates-1; i >= 0; i--)
      ElementM(sv->baseQ,i,j)=((k=restrictions->NonZeroIndex[i][j]) >= 0) ? ElementM(restrictions->MQ,i,j)*restrictions->B[k] : 0.0;

  if (sv->nlags_encoded > 0) ConvertBaseTransitionMatrix(sv->Q,sv->baseQ,sv->nlags_encoded);

  sv->t0=0;
  sv->t1=INT_MAX;

  return 1;
}

PRECISION dirichlet_restrictions_LogPrior_q(TMarkovStateVariable *sv)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  PRECISION *q, log_prior=restrictions->LogPriorConstant;
  int k, m=dw_DimA(restrictions->DirichletDim);

  // Accumulate Dirichlet kernel
  for (q=sv->q, k=0; k < m; q+=restrictions->DirichletDim[k]-1, k++)
    log_prior+=LogDirichlet_kernel(q,restrictions->Prior_b[k],restrictions->DirichletDim[k]);

  return log_prior;
}

/*
   Assumes:
     p  : vector of length sv->nstates.
     sv : pointer to valid TMarkovStateVariable structure

   Returns:
     One upon success and zero upon failure.

   Results:
     If p is null, it is created.  The p[i] is the probability that s0 is equal 
     to i.
*/
int dirichlet_restrictions_Probability_s0(TVector p, TMarkovStateVariable *sv)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  if (restrictions->UseErgodic)
    {
      return 0;
    }
  else
    {
      InitializeVector(p,1.0/(PRECISION)(sv->nstates));
      return 1;
    }
}

//----------------------------- Simulation routines -----------------------------
/*
   Assumes:
     S     : integer array of length model->nobs + 1 satisfying 0 <= S[t] < sv->nstates.
     sv    : pointer to valid TMarkovStateVariable structure
     model : pointer to valid TStateModel structure

   Returns:
     ln(p(q_old | S[T])) - ln(p(q_new | S[T])) where p( | ) is the conditional
     proposal density.

   Results:
    Draws sv->q from the proposal density conditional on the values of S.  Sets 
    sv->t1 to -1 which invalidates the transition matrices.
*/
PRECISION dirichlet_restrictions_DrawProposalTransitionMatrixParameters(int *S, TMarkovStateVariable *sv, TStateModel *model)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  PRECISION *q, log_proposal_ratio=0.0;
  int k, t, m=dw_DimA(restrictions->DirichletDim);

  // Set Dirichlet parameters
  memcpy(restrictions->B,restrictions->Prior_B,restrictions->total_Dirichlet_parameters*sizeof(PRECISION));
  for (t=model->nobs; t > 0; t--)
    if ((k=restrictions->NonZeroIndex[S[t]][S[t-1]]) >= 0)
      restrictions->B[k]+=1.0;

  // Rescale if necessary
  if (restrictions->scale != 1.0)
    for (k=restrictions->total_Dirichlet_parameters-1; k >= 0; k--) restrictions->B[k]*=restrictions->scale;

  // Compute p(q_old | S[T])
  for (q=sv->q, k=0; k < m; q+=restrictions->DirichletDim[k]-1, k++)
    log_proposal_ratio-=LogDirichlet_kernel(q,restrictions->b[k],restrictions->DirichletDim[k]);

  // Draw b[j] and compute p(q_new | S[T])
  for (k=dw_DimA(restrictions->DirichletDim)-1; k >= 0; k--)  
    log_proposal_ratio+=DrawDirichlet_kernel(restrictions->b[k],restrictions->b[k],restrictions->DirichletDim[k]);

  // Populate q
  dirichlet_restrictions_Populate_q(sv->q,restrictions);

  // Set sv->t1
  sv->t1=-1;

  return log_proposal_ratio;
}

/*
   Assumes:
     S     : integer array of length model->nobs + 1 satisfying 0 <= S[t] < sv->nstates.
     sv    : pointer to valid TMarkovStateVariable structure
     model : pointer to valid TStateModel structure

   Returns:
     Always returns one.

   Results:
     Draws sv->q from the conditional posterior density 
     p(q | Y[T], Z[T], S[T], theta) = p(q | S[T]).   
*/
int dirichlet_restrictions_DrawTransitionMatrixParameters(int *S, TMarkovStateVariable *sv, TStateModel *model)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  int k, t;

  // Set Dirichlet parameters
  memcpy(restrictions->B,restrictions->Prior_B,restrictions->total_Dirichlet_parameters*sizeof(PRECISION));
  for (t=model->nobs; t > 0; t--)
    if ((k=restrictions->NonZeroIndex[S[t]][S[t-1]]) >= 0)
      restrictions->B[k]+=1.0;

  // Draw b[k]
  for (k=dw_DimA(restrictions->DirichletDim)-1; k >= 0; k--)
    DrawDirichlet(restrictions->b[k],restrictions->b[k],restrictions->DirichletDim[k]);

  // Populate q
  dirichlet_restrictions_Populate_q(sv->q,restrictions);

  // Set sv->t1
  sv->t1=-1;

  return 1;
}

/*
   Assumes:
     sv    : pointer to valid TMarkovStateVariable structure

   Returns:
     Always returns one.

   Results:
     Draws sv->q from the prior density p(q), which is assumed to be Dirichlet.   
*/
int dirichlet_restrictions_DrawTransitionMatrixParametersFromPrior(TMarkovStateVariable *sv)
{
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);
  int k;

  // Draw b[k] from prior
  for (k=dw_DimA(restrictions->DirichletDim)-1; k >= 0; k--)
    DrawDirichlet(restrictions->b[k],restrictions->Prior_b[k],restrictions->DirichletDim[k]);

  // Populate q
  dirichlet_restrictions_Populate_q(sv->q,restrictions);

  // Set sv->t1
  sv->t1=-1;

  return 1;
}

//--------------------------- Miscelleanous routines ----------------------------

/*
   Assumes:
     sv    : pointer to valid TMarkovStateVariable structure
     model : ignored

  Returns: 
    One, failure cannot occur,

  Results:
    Sets sv->q to its prior mean.
*/
int dirichlet_restrictions_DefaultTransitionMatrixParameters(TMarkovStateVariable *sv, __attribute__ ((unused)) TStateModel *model)
{
  dirichlet_restrictions_SetTransitionMatrixParametersToPriorMean(sv->q,(TDirichletRestrictions*)(sv->info));
  return 1;
}

/*
   Assumes
     sv->baseQ has been updated and that sv->n_state_variables = 1.

   Returns
     One upon success and zero upon failure

   Results
     Updates sv->q from sv->baseQ.  If necessary, sv->Q is computed.  Both sv->t0 
     and sv->t1 are also updated.

   Notes
     If any of the elements of sv->q[i] is negative or their sum differs from one
     by more than DimV(sv->b[i])*MACHINE_EPSILON, then the routine returns zero.

     If any element of the updated sv->Q differs from the original sv->Q by more 
     than SQRT_MACHINE_EPSILON, then the routine returns zero.

     Because of the checking done, care should be taken in calling this routine 
     in instances where efficiency is important.
*/
int dirichlet_restrictions_Update_q_from_BaseTransitionMatrix(__attribute__ ((unused)) int t, TMarkovStateVariable *sv, __attribute__ ((unused)) TStateModel *model)
{
  int i, j, k;
  PRECISION sum, *scale;
  TDirichletRestrictions *restrictions=(TDirichletRestrictions*)(sv->info);

  // Compute B
  scale=(PRECISION*)dw_malloc(restrictions->total_Dirichlet_parameters*sizeof(PRECISION));
  for (k=restrictions->total_Dirichlet_parameters-1; k >= 0; k--) scale[k]=restrictions->B[k]=0.0;
  for (i=restrictions->nstates-1; i >= 0; i--)
    for (j=restrictions->nstates-1; j >= 0; j--)
      if ((k=restrictions->NonZeroIndex[i][j]) >= 0)
	{
	  restrictions->B[k]+=ElementM(sv->baseQ,i,j);
	  scale[k]+=ElementM(restrictions->MQ,i,j);
	}
  for (k=restrictions->total_Dirichlet_parameters-1; k >= 0; k--)
    if (scale[k] > 0) 
      restrictions->B[k]/=scale[k];
  dw_free(scale);

  // Check computed B
  for (i=dw_DimA(restrictions->DirichletDim)-1; i >= 0; i--)
    {
      for (sum=0.0, j=restrictions->DirichletDim[i]-1; j >= 0; j--)
	if (restrictions->b[i][j] < 0)
	  {
	    sv->t1=-1;
	    return 0;
	  }
	else
	  sum+=restrictions->b[i][j];
      if (fabs(sum-1.0) > SQRT_MACHINE_EPSILON)
	{
	  sv->t1=-1;
	  return 0;
	}
      for (sum=1/sum, j=restrictions->DirichletDim[i]-1; j >= 0; j--) restrictions->b[i][j]/=sum;
    }

  // Check computed Q
  for (i=restrictions->nstates-1; i >= 0; i--)
    for (j=restrictions->nstates-1; j >= 0; j--)
      {
	if ((k=restrictions->NonZeroIndex[i][j]) >= 0) ElementM(sv->baseQ,i,j)-=ElementM(restrictions->MQ,i,j)*restrictions->B[k];
	if (fabs(ElementM(sv->baseQ,i,j)) > SQRT_MACHINE_EPSILON)
	  {
	    sv->t1=-1;
	    return 0;
	  }
	else
	  ElementM(sv->baseQ,i,j)=(k >= 0) ? ElementM(restrictions->MQ,i,j)*restrictions->B[k] : 0.0;
      }

  // Success
  dirichlet_restrictions_Populate_q(sv->q,restrictions);
  if (sv->nlags_encoded > 0) ConvertBaseTransitionMatrix(sv->Q,sv->baseQ,sv->nlags_encoded);
  sv->t0=0;
  sv->t1=INT_MAX;

  return 1;
}
/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/

/*******************************************************************************/
/****************************** Utility Routines *******************************/
/*******************************************************************************/
TMatrix dirichlet_restrictions_CreatePriorFromDuration(int nbasestates, PRECISION duration)
{
  int i;
  TMatrix Prior=InitializeMatrix(Prior=CreateMatrix(nbasestates,nbasestates),1.0);
  PRECISION d=(nbasestates - 1)*duration;
  for (i=nbasestates-1; i >= 0; i--)
    ElementM(Prior,i,i)=d;
  return Prior;
}

void dirichlet_restrictions_SetLogPriorConstant(TDirichletRestrictions *restrictions)
{
  PRECISION sum, x;
  int i, j;
  restrictions->LogPriorConstant=0.0;
  for (i=dw_DimA(restrictions->DirichletDim)-1; i >= 0; i--)
    {
      for (sum=0.0, j=restrictions->DirichletDim[i]-1; j >= 0; j--)
	{
	  sum+=(x=restrictions->Prior_b[i][j]);
	  restrictions->LogPriorConstant-=dw_log_gamma(x);
	}
      restrictions->LogPriorConstant+=dw_log_gamma(sum);
    }
}

/*
   Assumes:
     q : array of length restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim)
     restrictions : pointer to valid TDirichletRestrictions structure

   Results:
     The elements of q, which are the free parameters, are set to the prior mean 
*/
void dirichlet_restrictions_SetTransitionMatrixParametersToPriorMean(PRECISION *q, TDirichletRestrictions *restrictions)
{
  int i, j;
  PRECISION sum;
  for (j=dw_DimA(restrictions->DirichletDim)-1; j >= 0; j--)
    {
      for (sum=restrictions->Prior_b[j][0], i=restrictions->DirichletDim[j]-1; i > 0; i--) sum+=restrictions->Prior_b[j][i];
      for (sum=1.0/sum, i=restrictions->DirichletDim[j]-1; i >= 0; i--)
	restrictions->b[j][i]=sum*restrictions->Prior_b[j][i];
    }
  dirichlet_restrictions_Populate_q(q,restrictions);
}

/*
   Assumes:
     q  : array of length restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim)
     restrictions : pointer to valid TDirichletRestrictions structure

   Returns:
     One upon success and zero upon failure.

   Results:
     Populates the Dirichlet parameters in restrictions->B and restrictions->b from 
     the array of free parameters q. 
 
   Notes:
     The elements of restrictions->b[k] are restrictions to be non-negative and their 
     sum must be one.  Checking is performed to ensure that this holds, up to 
     some tolerance.  
*/
int dirichlet_restrictions_Populate_B(PRECISION *q, TDirichletRestrictions *restrictions)
{
  int i, j, k, m=dw_DimA(restrictions->DirichletDim);
  PRECISION scale;
  for (i=0; i < m; i++)
    if ((k=restrictions->DirichletDim[i]-1) > 0)
      {
	memcpy(restrictions->b[i],q,k*sizeof(PRECISION));

	for (scale=1.0, j=k-1; j >= 0; j--) 
	  if (q[j] < 0)
	    return 0;
	  else
	    scale-=q[j];

	q+=k;

	if (scale < 0)
	  if (scale < -restrictions->DirichletDim[i]*MACHINE_EPSILON)
	    return 0;
	  else
	     restrictions->b[i][k]=0.0;
	else
	  restrictions->b[i][k]=scale;	
      }
    else
      restrictions->b[i][k]=1.0;
  return 1;
}

/*
   Assumes:
     q  : (restrictions->total_Dirichlet_parameters - dw_DimA(restrictions->DirichletDim)) array
     restrictions : pointer to valid TDirichletRestrictions structure

   Results:
     Populates the free parameters in the array q from the quasi-free parameters
     in restrictions->B. 
 
   Notes:
     The elements of restrictions->b[k] are restrictioned to be non-negative and their sum 
     must be one.  No checking is performed to ensure this.
*/
void dirichlet_restrictions_Populate_q(PRECISION *q, TDirichletRestrictions *restrictions)
{
  int i, k;
  for (i=0; i < dw_DimA(restrictions->DirichletDim); i++)
    if ((k=restrictions->DirichletDim[i]-1) > 0)
      {
	memcpy(q,restrictions->b[i],k*sizeof(PRECISION));
	q+=k;
      }
}

/*
   Checks that DirichletDim, NonZeroIndex, and MQ satisfies the appropriate
   conditions.  Returns 1 if the conditions are satisfied and 0 otherwise.
*/
int dirichlet_restrictions_CheckRestrictions(int* DirichletDim, int** NonZeroIndex, TMatrix MQ, TMatrix Prior, int nstates)
{
  int i, j, k, q, r, total_Dirichlet_parameters;
  PRECISION alpha, sum, common_sum=0.0, total_sum;

  //=== Check for null pointer and sizes 
  if (!Prior || (RowM(Prior) != nstates) || (ColM(Prior) != nstates)) return 0;
  if (!MQ || (RowM(MQ) != nstates) || (ColM(MQ) != nstates)) return 0;
  if (!DirichletDim) 
    return 0;
  if (!NonZeroIndex || (dw_DimA(NonZeroIndex) != nstates))
    return 0;
  else
    for (i=nstates-1; i >= 0; i--)
      if (dw_DimA(NonZeroIndex[i]) != nstates) return 0;

  //=== Checks DirichletDim[i] > 0 
  //=== Computes total_Dirichlet_parameters = DirichletDim[0] + ... + DirichletDim[dw_DimA(DirichletDim)-1] 
  for (total_Dirichlet_parameters=0, i=dw_DimA(DirichletDim)-1; i >= 0; i--)
    if (DirichletDim[i] <= 0)
      return 0;
    else
      total_Dirichlet_parameters+=DirichletDim[i];

  //=== Checks -1 <= NonZeroIndex[i][j] < total_Dirichlet_parameters.
  //=== Checks NonZeroIndex[i][j] >= 0, ==> MQ[i][j] > 0.
  //=== Checks NonZeroIndex[i][j] = -1, ==> MQ[i][j] = 0.
  for (j=0; j < nstates; j++)
    for (i=0; i < nstates; i++)
      if ((NonZeroIndex[i][j] < -1) || (total_Dirichlet_parameters <= NonZeroIndex[i][j]))
        return 0;
      else
        if (NonZeroIndex[i][j] >= 0)
          {
            if (ElementM(MQ,i,j) <= 0.0) return 0;
          }
        else
          {
            if (ElementM(MQ,i,j) != 0.0) return 0;
          }

  //=== Check that column sums are correct 
  for (j=0; j < nstates; j++)
    {
      total_sum=0.0;
      for (q=r=k=0; k < total_Dirichlet_parameters; k++)
        if (k == q)
          {
            common_sum=0.0;
            for (i=0; i < nstates; i++)
              if (NonZeroIndex[i][j] == k) 
		common_sum+=ElementM(MQ,i,j);
            q+=DirichletDim[r++];
            total_sum+=common_sum;
          }
        else
          {
            sum=0.0;
            for (i=0; i < nstates; i++)
	      if (NonZeroIndex[i][j] == k) 
		sum+=ElementM(MQ,i,j);
            if (fabs(sum - common_sum) > SQRT_MACHINE_EPSILON) return 0;
          }
      if (fabs(total_sum - 1.0) > SQRT_MACHINE_EPSILON) return 0;
    }

  //=== Check that prior elements are positive
  for (i=0; i < nstates; i++)
    for (j=0; j < nstates; j++)
      if (ElementM(Prior,i,j) <= 0) return 0;

  //=== Check that prior elements are large enough given the restrictions
  for (q=total_Dirichlet_parameters-1; q >= 0; q--)
    {
      alpha=1.0;
      for (j=0; j < nstates; j++)
        for (i=0; i < nstates; i++)
          if (NonZeroIndex[i][j] == q)
            alpha+=ElementM(Prior,i,j)-1.0;
      if (alpha <= 0) return 0;
    }

  return 1;
}

/*
   Assumes:
     b     : n dimensional array
     alpha : n dimensional array with positive elements
     n     : positive integer

   Results:
     Fills the vector b with a draw from the Dirichlet distribution with 
     parameters given by alpha.  The density of the Dirichlet distribution is


       Gamma(Sum(alpha[i], 0 <= i < n))
    -------------------------------------- Product(b[i]^(alpha[i]-1), 0 <= i < n)
     Product(Gamma(alpha[i]), 0 <= i < n)  

   Notes:
     The arguments b and alpha do not have to be distinct.

     A gamma deviate is positive with probability one.  This code calls
     dw_gamma_rnd(), but does not check if the returned value is positive.
     If all calls of dw_gamma_rnd() returned zero, then the return vector would
     identically zero, which is not a permissible value for a Dirichlet random
     vector.
*/
void DrawDirichlet(PRECISION *b, PRECISION *alpha, int n)
{
  int i;
  PRECISION scale=0.0;
  for (i=n-1; i >= 0; i--) scale+=(b[i]=dw_gamma_rnd(alpha[i]));
  if (scale > 0)
    for (scale=1.0/scale, i=n-1; i >= 0; i--) b[i]*=scale;
}

/*
   Assumes:
     b     : n dimensional array
     alpha : n dimensional array with positive elements
     n     : positive integer

   Results:
     Fills the vector b with a draw from the Dirichlet distribution with
     parameters given by alpha.  The density of the Dirichlet distribution is


       Gamma(Sum(alpha[i], 0 <= i < n))
    -------------------------------------- Product(b[i]^(alpha[i]-1), 0 <= i < n)
     Product(Gamma(alpha[i]), 0 <= i < n)

   Returns:
     Returns the value of

          ln(Product(b[i]^(alpha[i]-1), 0 <= i < n))

     evaluated at the draw.

   Notes:
     The arguments b and alpha do not have to be distinct.

     A gamma deviate is positive with probability one.  This code calls
     dw_gamma_rnd(), but does not check if the returned value is positive.
     A return of zero would cause this routine to crash.
*/
PRECISION DrawDirichlet_kernel(PRECISION *b, PRECISION *alpha, int n)
{
  int i;
  PRECISION log_kernel, sum_alpha, scale, a;
  for (log_kernel=sum_alpha=scale=0.0, i=n-1; i >= 0; i--)
    {
      a=alpha[i]-1;
      scale+=(b[i]=dw_gamma_rnd(alpha[i]));
      sum_alpha+=a;
      log_kernel+=a*log(b[i]);
    }
  for (scale=1.0/scale, i=n-1; i >= 0; i--) b[i]*=scale;
  return log_kernel + log(scale)*sum_alpha;
}

/*
   Assumes:
     q     : n-1 dimensional array
     alpha : n dimensional array with positive elements

   Results:
     Fills q with a draw from the Dirichlet distribution with parameters given by 
     alpha.  The density of the Dirichlet distribution is

      Gamma(Sum(alpha[i],0 <= i < n))  
  -------------------------------------- Product(q[i]^(alpha[i] - 1), 0 <= i < n)
   Product(Gamma(alpha[i]), 0 <= i < n)

   Notes:
     Only the first n-1 elements are returned.  The nth element will be equal to
     one minus the sum of the first n-1 elements.
*/
void DrawDirichlet_free(PRECISION *q, PRECISION *alpha, int n)
{
  int i;
  PRECISION scale=dw_gamma_rnd(alpha[n-1]);

  for (i=n-2; i >= 0; i--)
    scale+=(q[i]=dw_gamma_rnd(alpha[i]));

  if (scale > 0)
    for (scale=1.0/scale, i=n-2; i >= 0; i--) q[i]*=scale;
}

/*
   Assumes:
     q     : n-1 dimensional array with non-negative elements whose sum is less 
             than or equal to one.
     alpha : n dimensional array with positive elements

   Returns:
     The log value of the Dirichlet density with parameters alpha evaluated at q.
     Return MINUS_INFINITY if the restrictions on q are violated.

   Notes:
     The Dirichlet density is given by

         Gamma(Sum(alpha[i],0 <= i < n))
      -------------------------------------- Product(q[i]^(alpha[i] - 1), 0 <= i < n)
       Product(Gamma(alpha[i]), 0 <= i < n)
*/
PRECISION LogDirichlet_pdf(PRECISION *q, PRECISION *alpha, int n)
{
  PRECISION sum_q=1.0, sum_alpha=alpha[n-1], log_pdf=-dw_log_gamma(alpha[n-1]);
  int i;
  for (i=n-2; i >= 0; i--)
    if (q[i] > 0.0)
      {
	sum_q-=q[i];
	sum_alpha+=alpha[i];
	log_pdf+=(alpha[i]-1.0)*log(q[i]) - dw_log_gamma(alpha[i]);
      }
    else
      if (alpha[i] != 1.0)
	{
	  log_pdf=(alpha[i]-1.0);
	  for ( ; i >= 0; i--)
	    if (q[i] > 0.0) 
	      sum_q-=q[i];
	    else
	      log_pdf+=(alpha[i]-1.0);
	  if (sum_q <= 0.0)
	    log_pdf+=alpha[n-1]-1.0;
	  return (log_pdf > 0) ? MINUS_INFINITY : PLUS_INFINITY;
	}
  if (sum_q > 0.0)
    return log_pdf + (alpha[n-1]-1.0)*log(sum_q) + dw_log_gamma(sum_alpha);
  else
    if (alpha[n-1] == 1.0)
      return log_pdf + dw_log_gamma(sum_alpha);
    else
      return (alpha[n-1] > 1.0) ? MINUS_INFINITY : PLUS_INFINITY;//}
}

/*
   Assumes:
     q     : n-1 dimensional vector with non-negative elements whose sum is less
             than or equal to one.
     alpha : n dimensional array with positive elements

   Returns:
     The log value of a Dirichlet kernal with parameters alpha evaluated at q.  
     Returns MINUS_INFINITY if the restrictions on q are violated.

   Notes:
     The Dirichlet density is given by

        Gamma(Sum(alpha[i],0 <= i < n))
    -------------------------------------- Product(q[i]^(alpha[i]-1), 0 <= i < n)
     Product(Gamma(alpha[i]), 0 <= i < n)

     The kernal returned is given by

                  ln(Product(q[i]^(alpha[i]-1), 0 <= i < n))

     q[n-1] in the above formulas is defined to be 1 - Sum(q[i], 0 <= i < n-1).

     It is more efficient to evaluate the kernel than the pdf since the gamma
     functions do not have to be evaluated.
*/
PRECISION LogDirichlet_kernel(PRECISION *q, PRECISION *alpha, int n)
{
  PRECISION sum=1.0, log_kernel=0.0;
  int i;
  for (i=n-2; i >= 0; i--)
    if (q[i] > 0.0)
      {
	sum-=q[i];
	log_kernel+=(alpha[i]-1.0)*log(q[i]);
      }
    else
      if (alpha[i] != 1.0)
	{
	  log_kernel=(alpha[i]-1.0);
	  for ( ; i >= 0; i--)
	    if (q[i] > 0.0) 
	      sum-=q[i];
	    else
	      log_kernel+=(alpha[i]-1.0);
	  if (sum <= 0.0)
	    log_kernel+=alpha[n-1]-1.0;
	  return (log_kernel > 0) ? MINUS_INFINITY : PLUS_INFINITY;
	}
  if (sum > 0.0)
    return log_kernel + (alpha[n-1]-1.0)*log(sum);
  else
    if (alpha[n-1] == 1.0)
      return log_kernel;
    else
      return (alpha[n-1] > 1.0) ? MINUS_INFINITY : PLUS_INFINITY; 
}

/*
   Assumes:
     q     : (m-n)-dimensional array
     alpha : m-dimensional array with positive elements greater

   Results:
     Fills the q with independent draws from the Dirichlet distribution with
     parameters given by alpha. 

   Notes:
     m is equal to the sum of the elements of dirichlet_dims. 
*/
void DrawIndependentDirichlet_free(PRECISION *q, PRECISION *alpha, int *dirichlet_dims, int n)
{
  int i, j, k;
  for (i=j=k=0; i < n; j+=dirichlet_dims[i], k+=dirichlet_dims[i]-1, i++)
    DrawDirichlet_free(q+k,alpha+j,dirichlet_dims[i]);
}

/*
   Assumes:
     q     : (m-n)-dimensional array
     alpha : m-dimensional array with positive elements

   Returns:
     The sum of the log values of the Dirichlet densities with parameters alpha
     evaluated at q.

   Notes:
     m is equal to the sum of the elements of dirichlet_dims. 
*/
PRECISION LogIndependentDirichlet_pdf(PRECISION *q, PRECISION *alpha, int *dirichlet_dims, int n)
{
  int i, j, k;
  PRECISION log_pdf=0.0;
  for (i=j=k=0; i < n; j+=dirichlet_dims[i], k+=dirichlet_dims[i]-1, i++)
    log_pdf+=LogDirichlet_pdf(q+k,alpha+j,dirichlet_dims[i]);
  return log_pdf;
}
/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/
