#include <RcppArmadillo.h>
//[[Rcpp::depends(RcppArmadillo)]]
using namespace Rcpp;
using namespace arma;  

// Mathematical constants
#define MATH_PI        3.141592653589793238462643383279502884197169399375105820974
#define MATH_PI_2      1.570796326794896619231321691639751442098584699687552910487
#define MATH_2_PI      0.636619772367581343075535053490057448137838582961825794990
#define MATH_PI2       9.869604401089358618834490999876151135313699407240790626413
#define MATH_PI2_2     4.934802200544679309417245499938075567656849703620395313206
#define MATH_SQRT1_2   0.707106781186547524400844362104849039284835937688474036588
#define MATH_SQRT_PI_2 1.253314137315500251207882642405522626503493370304969158314
#define MATH_LOG_PI    1.144729885849400174143427351353058711647294812915311571513
#define MATH_LOG_2_PI  -0.45158270528945486472619522989488214357179467855505631739
#define MATH_LOG_PI_2  0.451582705289454864726195229894882143571794678555056317392


// Define helper functions
namespace help{   


// Call sample function from R
double sample_cpp( IntegerVector x ){
  Function f("sample");
  IntegerVector sampled = f(x, Named("size") = 1);
  return sampled[0];
}

// Make a sequence
IntegerVector myseq( int first, int last) {
  IntegerVector y(abs(last - first) + 1);
  if (first < last) 
    std::iota(y.begin(), y.end(), first);
  else {
    std::iota(y.begin(), y.end(), last);
    std::reverse(y.begin(), y.end());
  }
  return y;
}

// Simulate MVT normal data
arma::mat mvrnormArma( int n, arma::vec mu, arma::mat sigma ) {
  int ncols = sigma.n_cols;
  arma::mat Y = arma::randn( n, ncols );
  return arma::repmat( mu, 1, n ).t()  + Y * arma::chol( sigma );
}

// Generate exponential distribution random variates
double exprnd(double mu)
{
  return -mu * (double)std::log(1.0 - (double)R::runif(0.0,1.0));
}

// Function a_n(x) defined in equations (12) and (13) of
// Bayesian inference for logistic models using Polya-Gamma latent variables
// Nicholas G. Polson, James G. Scott, Jesse Windle
// arXiv:1205.0310
//
// Also found in the PhD thesis of Windle (2013) in equations
// (2.14) and (2.15), page 24
double aterm(int n, double x, double t)
{
  double f = 0;
  if(x <= t) {
    f = MATH_LOG_PI + (double)std::log(n + 0.5) + 1.5*(MATH_LOG_2_PI- (double)std::log(x)) - 2*(n + 0.5)*(n + 0.5)/x;
  }
  else {
    f = MATH_LOG_PI + (double)std::log(n + 0.5) - x * MATH_PI2_2 * (n + 0.5)*(n + 0.5);
  }    
  return (double)exp(f);
}

// Generate inverse gaussian random variates
double randinvg(double mu)
{
  // sampling
  double u = R::rnorm(0.0,1.0);
  double V = u*u;
  double out = mu + 0.5*mu * ( mu*V - (double)std::sqrt(4.0*mu*V + mu*mu * V*V) );
  
  if(R::runif(0.0,1.0) > mu /(mu+out)) {    
    out = mu*mu / out; 
  }    
  return out;
}

// Sample truncated gamma random variates
// Ref: Chung, Y.: Simulation of truncated gamma variables 
// Korean Journal of Computational & Applied Mathematics, 1998, 5, 601-610
double truncgamma()
{
  double c = MATH_PI_2;
  double X, gX;
  
  bool done = false;
  while(!done)
  {
    X = help::exprnd(1.0) * 2.0 + c;
    gX = MATH_SQRT_PI_2 / (double)std::sqrt(X);
    
    if(R::runif(0.0,1.0) <= gX) {
      done = true;
    }
  }
  
  return X;  
}

// Sample truncated inverse Gaussian random variates
// Algorithm 4 in the Windle (2013) PhD thesis, page 129
double tinvgauss(double z, double t)
{
  double X, u;
  double mu = 1.0/z;
  
  // Pick sampler
  if(mu > t) {
    // Sampler based on truncated gamma 
    // Algorithm 3 in the Windle (2013) PhD thesis, page 128
    while(1) {
      u = R::runif(0.0, 1.0);
      X = 1.0 / help::truncgamma();
      
      if ((double)std::log(u) < (-z*z*0.5*X)) {
        break;
      }
    }
  }  
  else {
    // Rejection sampler
    X = t + 1.0;
    while(X >= t) {
      X = help::randinvg(mu);
    }
  }    
  return X;
}

// Sample PG(1,z)
// Based on Algorithm 6 in PhD thesis of Jesse Bennett Windle, 2013
// URL: https://repositories.lib.utexas.edu/bitstream/handle/2152/21842/WINDLE-DISSERTATION-2013.pdf?sequence=1
double samplepg(double z)
{
  //  PG(b, z) = 0.25 * J*(b, z/2)
  z = (double)std::fabs((double)z) * 0.5;
  
  // Point on the intersection IL = [0, 4/ log 3] and IR = [(log 3)/pi^2, \infty)
  double t = MATH_2_PI;
  
  // Compute p, q and the ratio q / (q + p)
  // (derived from scratch; derivation is not in the original paper)
  double K = z*z/2.0 + MATH_PI2/8.0;
  double logA = (double)std::log(4.0) - MATH_LOG_PI - z;
  double logK = (double)std::log(K);
  double Kt = K * t;
  double w = (double)std::sqrt(MATH_PI_2);
  
  double logf1 = logA + R::pnorm(w*(t*z - 1),0.0,1.0,1,1) + logK + Kt;
  double logf2 = logA + 2*z + R::pnorm(-w*(t*z+1),0.0,1.0,1,1) + logK + Kt;
  double p_over_q = (double)std::exp(logf1) + (double)std::exp(logf2);
  double ratio = 1.0 / (1.0 + p_over_q); 
  
  double u, X;
  
  // Main sampling loop; page 130 of the Windle PhD thesis
  while(1) 
  {
    // Step 1: Sample X ? g(x|z)
    u = R::runif(0.0,1.0);
    if(u < ratio) {
      // truncated exponential
      X = t + help::exprnd(1.0)/K;
    }
    else {
      // truncated Inverse Gaussian
      X = help::tinvgauss(z, t);
    }
    
    // Step 2: Iteratively calculate Sn(X|z), starting at S1(X|z), until U ? Sn(X|z) for an odd n or U > Sn(X|z) for an even n
    int i = 1;
    double Sn = help::aterm(0, X, t);
    double U = R::runif(0.0,1.0) * Sn;
    int asgn = -1;
    bool even = false;
    
    while(1) 
    {
      Sn = Sn + asgn * help::aterm(i, X, t);
      
      // Accept if n is odd
      if(!even && (U <= Sn)) {
        X = X * 0.25;
        return X;
      }
      
      // Return to step 1 if n is even
      if(even && (U > Sn)) {
        break;
      }
      
      even = !even;
      asgn = -asgn;
      i++;
    }
  }
  return X;
}

// Function :: Propose a new beta
double beta_prop_cpp( double beta ){
  // Simulate proposal with jitter around current beta_gamma
  // Make sure that beta_gamma is specified with iteration iterations and branch b
  double beta_prop = beta + sqrt(0.5)*rnorm( 1 )[ 0 ];
  
  // Return output
  return beta_prop ;
}

// Function :: Calculate log beta contribution for individual p
double log_beta_p_cpp( double beta, 
                       double sigma2_beta){
  
  double log_beta = -0.50*log( 2*atan(1)*4*sigma2_beta ) - 1/( 2*sigma2_beta )*pow( beta, 2 );
  
  // Return output
  return log_beta;
}

// Function :: Calculate individual zeta
double log_zeta_eta_p_cpp( double t_p, 
                       double a_zeta_eta, 
                       double b_zeta_eta
){
  double post_a_zeta_eta = t_p + a_zeta_eta;
  double post_b_zeta_eta = 1 - t_p + b_zeta_eta;
  double log_zeta_eta_p = lgamma( post_a_zeta_eta ) + lgamma( post_b_zeta_eta ) - lgamma( post_a_zeta_eta + post_b_zeta_eta ) - ( lgamma( a_zeta_eta ) + lgamma( b_zeta_eta ) - lgamma( a_zeta_eta + b_zeta_eta ) );
  
  // Return output
  return log_zeta_eta_p ;
}



// Joint update for alpha_t and zeta
List between_alpha_t_zeta_update_cov_cpp( arma::mat Z, 
                                          arma::vec Z_sum, 
                                          arma::mat X_eta, 
                                          arma::mat beta_eta, 
                                          arma::mat alpha_t, 
                                          arma::mat zeta,  
                                          arma::mat loggamma,
                                          int N 
){
  // Get dimensions
  int TT = zeta.n_cols;   
  
  arma::mat zeta_accept( N, TT );
  zeta_accept.zeros();  
  
  // Add/Delete for zeta and alpha_t
  for( int t = 0; t < TT; ++t ){   // For updating all of the indicators 
    
    // Get current zeta value and mu_t
    arma::uvec zeta_index_t = find( Z.col( t ) == 0 );
    int num_zero_t = zeta_index_t.size();   
    for( int n = 0; n < num_zero_t; ++n ){ 
      int i_index = zeta_index_t[ n ]; 
      int zeta_c = zeta( i_index, t ); 
      double alpha_t_c = alpha_t( i_index, t ); 
      double z_dot_n = Z_sum[ i_index ];  
      double T_n = sum( alpha_t.row( i_index ) ); 
      arma::mat theta_i = exp( X_eta.row( i_index )*beta_eta.row( t ).t() )/( 1 + exp( X_eta.row( i_index )*beta_eta.row( t ).t() ) ); 
      double a_c_prop_scaled = exp( loggamma( i_index, t ) ) ; 
      
      // Delete Step:
      if ( zeta_c == 1 ){
        if( sum( zeta.row( i_index )) > 1 ){ 
          // Propose 
          double alpha_prop = 0;
          double zeta_prop = 0;
          double r = -z_dot_n*log( 1 - alpha_t_c/T_n ) + lgamma( exp( loggamma(  i_index, t ) ) ) - lgamma( a_c_prop_scaled ) + ( a_c_prop_scaled - exp( loggamma(  i_index, t ) ) )*log( alpha_t_c ) + log( 1 - theta_i[ 0 ] ) - log( theta_i[ 0 ] );
          
          //Calculate acceptance probability
          double a  = log( runif( 1 )[ 0 ] );
          
          // Determine acceptance
          if(a < r){
            alpha_t( i_index, t ) = alpha_prop;
            zeta( i_index, t ) = zeta_prop;
            zeta_accept( i_index, t ) = 1;
          }
        }else{
          zeta_accept( i_index, t ) = 0;
        }
        
      }else{ // Add step
        
        // Propose
        double alpha_prop = rgamma( 1, a_c_prop_scaled, 1 )[ 0 ];
        double zeta_prop = 1;
        
        double r = -z_dot_n*log( 1 + alpha_prop/T_n ) + lgamma( a_c_prop_scaled ) - lgamma( exp( loggamma(  i_index, t ) ) ) + ( exp( loggamma( i_index, t ) ) - a_c_prop_scaled )*log( alpha_prop ) - log( 1 - theta_i[ 0 ] ) + log( theta_i[ 0 ] );
        
        // Calculate acceptance probability
        double a  = log( runif( 1 )[ 0 ] );
        
        // Determine acceptance
        if(a < r){
          alpha_t( i_index, t ) = alpha_prop;
          zeta( i_index, t ) = zeta_prop;
          zeta_accept( i_index, t ) = 1;
        }
      }
    } 
  }
  
  // Return output
  List alpha_t_zeta_out( 2 );
  alpha_t_zeta_out[ 0 ] = alpha_t;
  alpha_t_zeta_out[ 1 ] = zeta;  
  return alpha_t_zeta_out; 
}

arma::mat alpha_t_update_cov_cpp( arma::mat Z, arma::mat zeta, arma::mat loggamma, arma::vec mu, int N ){
  int TT = zeta.n_cols;    
  
  arma::mat zeta_accept( N, TT );
  zeta_accept.zeros(); 
  
  arma::mat alpha_t( N, TT );
  alpha_t.zeros();
  
  for( int n = 0; n < N; ++n ){
    for( int t = 0; t < TT; ++t ){
      if( zeta( n, t ) == 1 ){ 
        alpha_t( n, t ) = rgamma( 1, Z( n, t ) +  exp(loggamma( n, t ) ), 1/( mu[ n ] + 1 )  )[ 0 ];  
        
        if( alpha_t( n, t ) < pow( 10, -100 ) ){
          alpha_t( n, t ) = pow( 10, -100 );
        }
      }
    } 
  } 
  
  return alpha_t;
}
 
arma::mat alpha_t_update_DM_cov_cpp( arma::mat Z, arma::mat loggamma, arma::vec mu, int N ){
   int TT = Z.n_cols;     
   
   arma::mat alpha_t( N, TT );
   alpha_t.zeros();
   
   for( int n = 0; n < N; ++n ){
     for( int t = 0; t < TT; ++t ){ 
         alpha_t( n, t ) = rgamma( 1, Z( n, t ) +  exp(loggamma( n, t ) ), 1/( mu[ n ] + 1 )  )[ 0 ];  
         
         if( alpha_t( n, t ) < pow( 10, -100 ) ){
           alpha_t( n, t ) = pow( 10, -100 );
         } 
     } 
   } 
   
   return alpha_t;
 }


 
 arma::vec mu_update_cpp( arma::mat Z, arma::vec Z_sum, arma::mat alpha_t ){  
   int N = alpha_t.n_rows;   
   
   arma::vec mu( N );
   mu.zeros();
   
   for( int n = 0; n < N; ++n ){
     double z_dot_n = Z_sum[ n ];
     double T_n = sum( alpha_t.row( n ) );
     
     mu[ n ] = rgamma( 1, z_dot_n, 1/T_n )[ 0 ]; 
   }
   
   return mu; 
 }

// Update auxiliary parameters omega 
arma::mat update_omega_eta( arma::mat X_eta, arma::mat beta_eta, int T ){ 
  int N  = X_eta.n_rows;  
  
  // Make a home for omega_eta updates
  arma::mat updated_omega_eta( N, T ); 
  updated_omega_eta.zeros();
  
  // Update each ometa_eta_ij individually
  for( int t = 0; t < T; ++t ){  // For each true classification
    for( int l = 0; l < N; ++l ){ // For each site 
      arma::mat X_eta_sub = X_eta.row( l );   
      arma::vec beta_site = beta_eta.row( t ).t();   
      
      arma::mat phi_ij( 1, 1 ); 
      
      // Get species-specific regression coefficients 
      phi_ij = X_eta_sub*beta_site;  
      updated_omega_eta( l, t ) = help::samplepg( phi_ij[ 0 ] ); 
    }
  }
  return updated_omega_eta;
}

// Function :: Update beta_eta  
arma::mat update_beta_eta_population( arma::mat X_eta, arma::mat omega_eta, arma::mat zeta, arma::mat mu_beta_eta_prior, double sigma2_beta_eta, int T, int N ){
  int P_eta = X_eta.n_cols;  
  
  // Make place for update
  arma::mat beta_eta_update( T, P_eta );
  beta_eta_update.zeros();
  
  for( int t = 0; t < T; ++t ){
    // Get all visits at that site  
    arma::vec zeta_t = zeta.col( t );
    arma::vec omega_eta_t = omega_eta.col( t ); 
    
    arma::vec H_t = ( zeta_t  - 0.5 )/omega_eta_t ;  
    
    // Get matrix of sigma2_beta_eta
    arma::mat sigma2_beta_eta_MAT( P_eta , P_eta  );
    sigma2_beta_eta_MAT.zeros();
    
    for( int p = 0; p < P_eta; ++p ){
      sigma2_beta_eta_MAT( p, p ) = 1/sigma2_beta_eta;
    }
    
    // Make V_beta_eta and mu_beta_eta (and inside)
    arma::mat V_beta_eta( P_eta, P_eta );
    arma::mat mu_beta_eta( P_eta, 1 );
    arma::mat mu_beta_eta_inside( P_eta, 1 );
    V_beta_eta.zeros();
    mu_beta_eta.zeros();
    mu_beta_eta_inside.zeros(); 
    
    for( int n = 0; n < N; ++n ){    
      V_beta_eta += omega_eta_t[ n ]*X_eta.row( n ).t()*X_eta.row( n );  
      mu_beta_eta_inside += omega_eta_t[ n ]*X_eta.row( n ).t()*H_t[ n ]; 
      
    }
    
    // Subset prior    
    arma::vec prior_means = mu_beta_eta_prior.row( t ).t(); // Just the first row, since we are treating them all the same
    V_beta_eta += sigma2_beta_eta_MAT;
    V_beta_eta = inv( V_beta_eta ); 
    mu_beta_eta = mu_beta_eta_inside + sigma2_beta_eta_MAT*prior_means;  
    
    mu_beta_eta = V_beta_eta*mu_beta_eta; 
    
    // Update bth element
    arma::mat beta_eta_t_update( 1, P_eta );
    beta_eta_t_update.zeros();
    
    if( P_eta != 1 ){
      beta_eta_t_update =  help::mvrnormArma( 1, mu_beta_eta, V_beta_eta );
    }else{
      beta_eta_t_update = R::rnorm( mu_beta_eta[ 0 ], V_beta_eta[ 0 ] );
    }
    
    beta_eta_update.row( t ) = beta_eta_t_update;   
    
  }
  return beta_eta_update;
}



// Function :: Between Step (jointly update beta_theta and zeta) with add/delete using BB prior
List between_beta_eta_zeta_update_BB_cpp(
    arma::mat X_eta,
    arma::mat zeta_eta, 
    arma::mat beta_eta, 
    arma::mat alpha_t,  
    double sigma2_beta_eta,
    double a_zeta_eta,
    double b_zeta_eta
){
  
  // Between Model Step
  // Initiate memory space
  int P_eta = zeta_eta.n_cols;  
  int B = zeta_eta.n_rows;
  int obs = X_eta.n_rows;
  
  // Choose a compositional element
  int b = help::sample_cpp( help::myseq( 0, B - 1) ); 
  //for( int b = 1; b < B; ++b ){
  
  
  // For each covariate, update
  for( int p = 1; p < P_eta; ++p ){   // p = 1 skips intercept term 
    //int p = help::sample_cpp( help::myseq( 1, P_theta - 1) );
    
    // Calculate current log likelihood of eta
    double log_like = 0;  
    
    for( int n = 0; n < obs; ++n ){ 
      arma::mat theta_i = exp( X_eta.row( n )*beta_eta.row( b ).t() )/( 1 + exp( X_eta.row( n )*beta_eta.row( b ).t() ) );
      log_like = log_like  + alpha_t( n, b )*log( theta_i[ 0 ] ) + ( 1 - alpha_t( n, b ) )*log( 1 - theta_i[ 0 ] ) ;
    }
    
    // Calculate proposed likelihood 
    arma::mat beta_eta_prop = beta_eta.row( b );
    arma::mat zeta_eta_prop = zeta_eta.row( b );
    
    if( zeta_eta( b, p ) == 1 ){
      beta_eta_prop[ p ] = 0.0;
      zeta_eta_prop[ p ] = 0;
    }else{
      beta_eta_prop[ p ] = help::beta_prop_cpp( beta_eta( b, p ) ); 
      zeta_eta_prop[ p ] = 1;
    } 
    
    double log_like_prop = 0; 
    
    for( int n = 0; n < obs; ++n ){ 
      arma::mat theta_i_prop = exp( X_eta.row( n )*beta_eta_prop.t() )/( 1 + exp( X_eta.row( n )*beta_eta_prop.t() ) );
      log_like_prop = log_like_prop  + alpha_t( n, b )*log( theta_i_prop[ 0 ] ) + ( 1 - alpha_t( n, b ) )*log( 1 - theta_i_prop[ 0 ] ) ;
    }
    
    // Calculate ratio for single varphi
    double r = 0; 
    
    // If BB
    if( zeta_eta( b, p ) == 1 ){
      r = log_like_prop + help::log_zeta_eta_p_cpp( 0, a_zeta_eta, b_zeta_eta) - ( log_like + help::log_beta_p_cpp( beta_eta( b, p ), sigma2_beta_eta ) + help::log_zeta_eta_p_cpp( 1, a_zeta_eta, b_zeta_eta ) );
    }else{
      r = log_like_prop + help::log_beta_p_cpp( beta_eta_prop[ p ], sigma2_beta_eta ) + help::log_zeta_eta_p_cpp( 1, a_zeta_eta, b_zeta_eta ) - ( log_like + help::log_zeta_eta_p_cpp( 0, a_zeta_eta, b_zeta_eta ) );
    }
    
    // Calculate acceptance probability
    double aa  = log( runif( 1 )[ 0 ] ); 
    
    // Determine acceptance
    if( aa < r ){
      zeta_eta( b, p ) = zeta_eta_prop[ p ];
      beta_eta( b, p ) = beta_eta_prop[ p ];
    }
  }
  //}
  
  // Return output
  List between( 2 );
  between[ 0 ] = zeta_eta;
  between[ 1 ] = beta_eta; 
  return between;
}





// Function :: Calculate individual varphi
double log_zeta_gamma_pj_cpp( double t_pj, 
                          double a_zeta_gamma, 
                          double b_zeta_gamma
){
  double post_a_zeta_gamma = t_pj + a_zeta_gamma;
  double post_b_zeta_gamma = 1 - t_pj + b_zeta_gamma;
  double log_zeta_gamma_pj = lgamma( post_a_zeta_gamma ) + lgamma( post_b_zeta_gamma ) - lgamma( post_a_zeta_gamma + post_b_zeta_gamma ) - ( lgamma( a_zeta_gamma ) + lgamma( b_zeta_gamma ) - lgamma( a_zeta_gamma + b_zeta_gamma ) );
  
  // Return output
  return log_zeta_gamma_pj ;
}


// Function :: Calculate log beta contribution for individual p
double log_beta_p_cpp( double beta, 
                       double prior_beta,
                       double sigma2_beta){
  
  double log_beta = -0.50*log( 2*atan(1)*4*sigma2_beta ) - 1/( 2*sigma2_beta )*pow( beta - prior_beta, 2 );
  
  // Return output
  return log_beta;
}
 

// Function ::  Update beta_gamma 
List within_beta_gamma_augment_cpp(
    arma::mat X_gamma, 
    arma::mat beta_gamma, 
    arma::mat loggamma,
    arma::mat alpha_t, 
    arma::mat zeta,   
    arma::mat mu_beta_gamma_prior,
    double sigma2_beta_gamma,
    double N
){
  
  // Initiate memory space
  int T = beta_gamma.n_rows;
  int P = beta_gamma.n_cols; 
  
  
  // Update all coefficients 
  for( int t = 0; t < T; ++t ){
    for( int p = 0; p < P; ++p ){   
      
      // Calculate current log likelihood 
      double log_like = 0; 
      
      for( int n = 0; n < N; ++n ){
        if( zeta( n, t ) == 1 ){
          log_like = log_like - lgamma( exp( loggamma( n , t ) ) ) + exp( loggamma( n , t ) )*log( alpha_t( n, t ) );
        }  
      }
      
      // Calculate proposed likelihood 
      double beta_gamma_prop = 0; 
      
      beta_gamma_prop = help::beta_prop_cpp( beta_gamma( t, p ) ); 
      
      arma::vec loggamma_prop( N );
      loggamma_prop.zeros();
      
      for( int n = 0 ; n < N ; ++n ){
        loggamma_prop[ n ] = loggamma( n, t ) - beta_gamma( t, p ) * X_gamma( n, p ) + beta_gamma_prop * X_gamma( n, p );
      }
      
      double log_like_prop = 0; 
      for( int n = 0; n < N; ++n ){
        if( zeta( n, t ) == 1 ){
          log_like_prop = log_like_prop - lgamma( exp( loggamma_prop[ n ] ) ) + exp( loggamma_prop[ n ] )*log( alpha_t( n, t ) );
        }
      }
      
      // Calculate ratio
      double r = 0;  
      r = log_like_prop + help::log_beta_p_cpp( beta_gamma_prop, mu_beta_gamma_prior( t, p), sigma2_beta_gamma) - ( log_like + help::log_beta_p_cpp( beta_gamma( t, p ), mu_beta_gamma_prior( t, p), sigma2_beta_gamma )  );
      
      // Calculate acceptance probability
      double a  = log( runif( 1 )[ 0 ] );
      
      // Determine acceptance
      if(a < r){ 
        beta_gamma( t, p ) = beta_gamma_prop;
        loggamma.col( t ) = loggamma_prop;
      }
 
    }
  } 
  
  // Return output
  List beta_gamma_out( 2 );
  beta_gamma_out[ 0 ] = beta_gamma;
  beta_gamma_out[ 1 ] = loggamma;
  return beta_gamma_out;
}

// Function ::  Update beta_gamma 
List within_beta_gamma_BVS_augment_cpp(
    arma::mat X_gamma, 
    arma::mat zeta_gamma,
    arma::mat beta_gamma, 
    arma::mat loggamma,
    arma::mat alpha_t, 
    arma::mat zeta,   
    arma::mat mu_beta_gamma_prior,
    double sigma2_beta_gamma,
    double N
){
  
  // Initiate memory space
  int T = beta_gamma.n_rows;
  int P = beta_gamma.n_cols; 
  
  
  // Update all coefficients 
  for( int t = 0; t < T; ++t ){
    for( int p = 0; p < P; ++p ){   
      
      if( zeta_gamma( t, p ) == 1 ){
        // Calculate current log likelihood 
        double log_like = 0; 
      
        for( int n = 0; n < N; ++n ){
          if( zeta( n, t ) == 1 ){
            log_like = log_like - lgamma( exp( loggamma( n , t ) ) ) + exp( loggamma( n , t ) )*log( alpha_t( n, t ) );
          }  
        }
      
        // Calculate proposed likelihood 
        double beta_gamma_prop = 0; 
      
        beta_gamma_prop = help::beta_prop_cpp( beta_gamma( t, p ) ); 
      
        arma::vec loggamma_prop( N );
        loggamma_prop.zeros();
      
        for( int n = 0 ; n < N ; ++n ){
          loggamma_prop[ n ] = loggamma( n, t ) - beta_gamma( t, p ) * X_gamma( n, p ) + beta_gamma_prop * X_gamma( n, p );
        }
      
        double log_like_prop = 0; 
        for( int n = 0; n < N; ++n ){
          if( zeta( n, t ) == 1 ){
            log_like_prop = log_like_prop - lgamma( exp( loggamma_prop[ n ] ) ) + exp( loggamma_prop[ n ] )*log( alpha_t( n, t ) );
          }
        }
      
        // Calculate ratio
        double r = 0;  
        r = log_like_prop + help::log_beta_p_cpp( beta_gamma_prop, mu_beta_gamma_prior( t, p), sigma2_beta_gamma) - ( log_like + help::log_beta_p_cpp( beta_gamma( t, p ), mu_beta_gamma_prior( t, p), sigma2_beta_gamma )  );
      
        // Calculate acceptance probability
        double a  = log( runif( 1 )[ 0 ] );
      
        // Determine acceptance
        if(a < r){ 
          beta_gamma( t, p ) = beta_gamma_prop;
          loggamma.col( t ) = loggamma_prop;
      }
      }
    }
  } 
  
  // Return output
  List beta_gamma_out( 2 );
  beta_gamma_out[ 0 ] = beta_gamma;
  beta_gamma_out[ 1 ] = loggamma;
  return beta_gamma_out;
}
 
// Function :: Between Step (jointly update beta_gamma and zeta_gamma) with add/delete using BB prior
List between_beta_gamma_zeta_augment_update_BB_cpp(
    arma::mat X_gamma,
    arma::mat zeta_gamma, 
    arma::mat beta_gamma, 
    arma::mat loggamma, 
    arma::mat alpha_t,
    arma::mat zeta, 
    double sigma2_beta_gamma,
    double a_zeta_gamma,
    double b_zeta_gamma,
    int N
){
  
  // Between Model Step
  // Initiate memory space
  int J = beta_gamma.n_rows;
  int P = beta_gamma.n_cols;  
 
  // Get a covariate p to update in branch j 
  
  int j = help::sample_cpp( help::myseq( 0, J - 1) ); 
  for( int p = 1; p < P; ++p ){   // p = 1 skips intercept term
  
    // Calculate current log likelihood 
    double log_like = 0; 
 
    for( int n = 0; n < N; ++n ){
      if( zeta( n, j ) == 1 ){
        log_like = log_like - lgamma( exp( loggamma( n , j ) ) ) + exp( loggamma( n , j ) )*log( alpha_t( n, j ) );
      }
    }
 
    // Calculate proposed likelihood 
    double beta_gamma_prop;
    int zeta_gamma_prop; 
    
    if( zeta_gamma( j, p ) == 1 ){
      beta_gamma_prop = 0.0;
      zeta_gamma_prop = 0;
    }else{
      beta_gamma_prop = help::beta_prop_cpp( beta_gamma( j, p ) ); 
      zeta_gamma_prop = 1;
    }
    
    arma::vec loggamma_prop( N );
    loggamma_prop.zeros();
    
    for( int n = 0 ; n < N ; ++n ){
      loggamma_prop[ n ] = loggamma( n, j ) - beta_gamma( j, p ) * X_gamma( n, p ) + beta_gamma_prop * X_gamma( n, p );
    }
    
    double log_like_prop = 0; 
    for( int n = 0; n < N; ++n ){
      if( zeta( n, j ) == 1 ){
        log_like_prop = log_like_prop - lgamma( exp( loggamma_prop[ n ] ) ) + exp( loggamma_prop[ n ] )*log( alpha_t( n, j ) );
      }
    }
    
    // Calculate ratio for single varphi
    double r;
 
    // If BB
    if( zeta_gamma( j, p ) == 1 ){
      r = log_like_prop + help::log_zeta_gamma_pj_cpp( 0, a_zeta_gamma, b_zeta_gamma) - ( log_like + help::log_beta_p_cpp( beta_gamma( j, p ),0, sigma2_beta_gamma ) + help::log_zeta_gamma_pj_cpp( 1, a_zeta_gamma, b_zeta_gamma ) );
    }else{
      r = log_like_prop + help::log_beta_p_cpp( beta_gamma_prop,0, sigma2_beta_gamma ) + help::log_zeta_gamma_pj_cpp( 1, a_zeta_gamma, b_zeta_gamma) - ( log_like + help::log_zeta_gamma_pj_cpp( 0, a_zeta_gamma, b_zeta_gamma ) );
    }
 
    
    // Calculate acceptance probability
    double aa  = log( runif( 1 )[ 0 ] );
    
    // Determine acceptance
    if( aa < r ){
      zeta_gamma( j, p ) = zeta_gamma_prop;
      beta_gamma( j, p ) = beta_gamma_prop;
      loggamma.col( j ) = loggamma_prop;
    }
    
  }
  
  // Return output
  List between( 3 );
  between[ 0 ] = zeta_gamma;
  between[ 1 ] = beta_gamma;
  between[ 2 ] = loggamma;
  return between;
}

// Function ::  Update beta_gamma 
List within_beta_gamma_DM_augment_cpp(
    arma::mat X_gamma, 
    arma::mat beta_gamma, 
    arma::mat loggamma,
    arma::mat alpha_t,  
    arma::mat mu_beta_gamma_prior,
    double sigma2_beta_gamma,
    double N
){
  
  // Initiate memory space
  int T = beta_gamma.n_rows;
  int P = beta_gamma.n_cols; 
  
  
  // Update all coefficients 
  for( int t = 0; t < T; ++t ){
    for( int p = 0; p < P; ++p ){   
      
      // Calculate current log likelihood 
      double log_like = 0; 
      
      for( int n = 0; n < N; ++n ){ 
          log_like = log_like - lgamma( exp( loggamma( n , t ) ) ) + exp( loggamma( n , t ) )*log( alpha_t( n, t ) );
      }
      
      // Calculate proposed likelihood 
      double beta_gamma_prop = 0; 
      
      beta_gamma_prop = help::beta_prop_cpp( beta_gamma( t, p ) ); 
      
      arma::vec loggamma_prop( N );
      loggamma_prop.zeros();
      
      for( int n = 0 ; n < N ; ++n ){
        loggamma_prop[ n ] = loggamma( n, t ) - beta_gamma( t, p ) * X_gamma( n, p ) + beta_gamma_prop * X_gamma( n, p );
      }
      
      double log_like_prop = 0; 
      for( int n = 0; n < N; ++n ){
          log_like_prop = log_like_prop - lgamma( exp( loggamma_prop[ n ] ) ) + exp( loggamma_prop[ n ] )*log( alpha_t( n, t ) );
      }
      
      // Calculate ratio
      double r = 0;  
      r = log_like_prop + help::log_beta_p_cpp( beta_gamma_prop, mu_beta_gamma_prior( t, p), sigma2_beta_gamma) - ( log_like + help::log_beta_p_cpp( beta_gamma( t, p ), mu_beta_gamma_prior( t, p), sigma2_beta_gamma )  );
      
      // Calculate acceptance probability
      double a  = log( runif( 1 )[ 0 ] );
      
      // Determine acceptance
      if(a < r){ 
        beta_gamma( t, p ) = beta_gamma_prop;
        loggamma.col( t ) = loggamma_prop;
      }
      
    }
  } 
  
  // Return output
  List beta_gamma_out( 2 );
  beta_gamma_out[ 0 ] = beta_gamma;
  beta_gamma_out[ 1 ] = loggamma;
  return beta_gamma_out;
}



// Multinomial update
IntegerVector oneMultinomCalt(int count, NumericVector probs) {
  int k = probs.size();
  IntegerVector ans(k);
  rmultinom(count, probs.begin(), k, ans.begin());
  return(ans);
}


// Function :: Update Z
List update_Z( arma::mat Y, arma::mat alpha_t, arma::mat a_tc, int N, int T, int C ){
  
  arma::mat Z_update( N, T );
  Z_update.zeros();
  arma::mat tc_update( T, C );
  tc_update.zeros();
  
  // Make theta for each t
  arma::mat theta( T, C );
  for( int t = 0; t < T; ++t ){
    double sum_theta_t = sum( a_tc.row( t ) );
    theta.row( t ) = a_tc.row( t )/sum_theta_t;
  }  
   
  // Set Theta for each site
  arma::mat Theta( N, T );
  for( int n = 0; n < N; ++n ){
    double sum_alpha = sum( alpha_t.row( n ) );
    Theta.row( n ) = alpha_t.row( n )/sum_alpha; 
  } 
  
  for( int n = 0; n < N; ++n ){
    for( int c = 0; c < C; ++c ){
      int Ync = Y( n, c );
      arma::vec Theta_n = Theta.row( n ).t() % theta.col( c );

      NumericVector probs = as<NumericVector>( wrap( Theta_n ) );

      probs = probs/sum( probs ); 
      IntegerVector ans( T );
 
      rmultinom( Ync, probs.begin(), T, ans.begin());
      
      arma::vec count = as<arma::vec>( wrap( ans ) );

      Z_update.row( n ) = Z_update.row( n ) + count.t();
      tc_update.col( c ) = tc_update.col( c ) + count;
    }
    
  } 
  
  // Return output
  List Z_out( 2 );
  Z_out[ 0 ] = Z_update;
  Z_out[ 1 ] = tc_update;  
  return Z_out; 
}

arma::mat a_tc_update_cpp( arma::mat tc, arma::mat nu, arma::vec u_t, int T, int C ){
  arma::mat update_a_tc( T, C );
  update_a_tc.zeros();
  
  // Get Matrix of classification counts over T and  C
  // Then update a_tc
  
  for( int t = 0; t < T; ++t ){
    for( int c = 0; c < C; ++c ){ 
      update_a_tc( t, c ) = rgamma( 1, tc( t, c ) + nu( t, c ), 1/( u_t[ t ] + 1 )  )[ 0 ];  
      
      if( update_a_tc( t, c ) < pow( 10, -100 ) ){
        update_a_tc( t, c ) = pow( 10, -100 );
      }
    }
  } 
  
  return update_a_tc;
}

arma::vec u_t_update_cpp( arma::mat Z, arma::mat a_tc, int T ){ 
  
  arma::vec update_u_t( T );
  update_u_t.zeros();
  
  for( int t = 0; t < T; ++t ){
    double a_bar_t = sum( a_tc.row( t ) );
    
    update_u_t[ t ] =  rgamma( 1, sum( Z.col( t ) ), 1/a_bar_t )[ 0 ]; 
    if( update_u_t[ t ] < pow( 10, -100 ) ){
      update_u_t[ t ] = pow( 10, -100 );
    }
  }
  return update_u_t; 
}



}  // END OF NAMESPACE

 
 
// [[Rcpp::export]]
List MicroMiss( 
    int iterations,
    int thin, 
    int N, // Number of hosts
    int T, // Number of true classifications
    int C, // Number of observed classifications 
    int P_eta, // Number of columns in design matrix for probability of occurence 
    int P_gamma, // Number of columns in design matrix for probability of relative abundance
    arma::mat Y, // Observed classifications N X C - matrix
    arma::cube Z, // True classification  N X T X M - array 
    arma::vec Z_sum, // Sum of observations N - vector
    arma::cube alpha_t, // Aux var for occupancy probability N X T X M - cube
    arma::cube zeta, // Occupancy indicator N X T X M - cube
    arma::cube beta_eta, // Probability of occupancy T X P_eta X M - cube
    arma::mat X_eta, // Design matrix for occupancy N X P_eta - matrix
    arma::cube omega_eta, // Aux var for occupancy indicator N X T X M - cube
    arma::cube beta_gamma, // Probability of relative abundance T X P_gamma X M - cube
    arma::mat X_gamma, // Design matrix for relative abundance N X P_gamma - matrix
    arma::mat mu, // Aux var for DM trick for true classification prob N - vector (not saved over M) 
    arma::cube a_tc, //  Aux var for missclassification probability T X C X M - cube
    arma::vec u_t, // Aux var for DM trick for missclassification prob T - vector (not saved)
    arma::mat nu, // Concentration parameter for a_tc T X C - matrix (fixed)
    arma::mat mu_beta_eta_prior, // Prior mean of beta_eta
    arma::mat mu_beta_gamma_prior, // Prior mean of beta_gamma 
    double sigma2_beta_eta, // Prior variance of beta_eta 
    double sigma2_beta_gamma // Prior variance of beta_gamma
){
  
  // Initiate memory
  arma::mat temp_Z( N, T ); 
  arma::mat temp_alpha_t( N, T ); 
  arma::mat temp_zeta( N, T ); 
  arma::mat temp_beta_eta( T, P_eta );
  arma::mat temp_omega_eta( N, T );
  arma::vec temp_mu( N );
  arma::mat temp_a_tc( T, C );
  arma::vec temp_u_t( T );
  arma::mat temp_beta_gamma( T, P_gamma ); 
  arma::mat temp_tc( T, C );
  
  temp_Z = Z.slice( 0 );  
  temp_alpha_t = alpha_t.slice( 0 );
  temp_zeta = zeta.slice( 0 );
  temp_beta_eta = beta_eta.slice( 0 );
  temp_omega_eta = omega_eta.slice( 0 );
  temp_mu = mu.col( 0 );
  temp_a_tc = a_tc.slice( 0 );
  temp_u_t = u_t;
  temp_beta_gamma = beta_gamma.slice( 0 );
  temp_tc.zeros();
  
  // Set initial loggamma
  arma::mat temp_loggamma( N, T );
  temp_loggamma.zeros();
  for( int n = 0; n < N; ++n ){
    for( int t = 0; t < T; ++t ){
      temp_loggamma( n, t ) = ( X_gamma.row( n )*temp_beta_gamma.row( t ).t() ).eval()( 0, 0 );
    }
  }
  
  // Looping over the number of iterations specified by user
  for( int iter = 0; iter < iterations; ++iter ){
    
    // Joint update for alpha_t and zeta
    // Between Step: Expand or Contract   
    List between_return = help::between_alpha_t_zeta_update_cov_cpp( temp_Z, Z_sum, X_eta, temp_beta_eta, temp_alpha_t, temp_zeta, temp_loggamma, N );
    temp_alpha_t = as<arma::mat>( between_return[ 0 ] );
    temp_zeta = as<arma::mat>( between_return[ 1 ] );
    
    // Within Step
    temp_alpha_t = help::alpha_t_update_cov_cpp( temp_Z, temp_zeta, temp_loggamma, temp_mu, N );

    // Update mu
    temp_mu = help::mu_update_cpp( temp_Z, Z_sum, temp_alpha_t );

    // Update omega_eta
    temp_omega_eta = help::update_omega_eta( X_eta, temp_beta_eta, T );

    // Update beta_eta
    temp_beta_eta = help::update_beta_eta_population( X_eta, temp_omega_eta, temp_zeta, mu_beta_eta_prior, sigma2_beta_eta, T, N );

    // Update beta_gamma
    List beta_gamma_out = help::within_beta_gamma_augment_cpp( X_gamma, temp_beta_gamma, temp_loggamma, temp_alpha_t, temp_zeta, mu_beta_gamma_prior, sigma2_beta_gamma, N );
    temp_beta_gamma =  as<arma::mat>( beta_gamma_out[ 0 ] );
    temp_loggamma = as<arma::mat>( beta_gamma_out[ 1 ] );
 
    // Update Z
    List Z_out = help::update_Z( Y, temp_alpha_t, temp_a_tc, N, T, C );
    temp_Z =  as<arma::mat>( Z_out[ 0 ] );
    temp_tc = as<arma::mat>( Z_out[ 1 ] ); 

    // Update a_tc
    temp_a_tc = help::a_tc_update_cpp( temp_tc, nu, temp_u_t, T, C );
  
    // Update u_t
    temp_u_t = help::u_t_update_cpp( temp_Z, temp_a_tc, T );
    
    // Set the starting values for the next iteration
    if( (iter + 1) % thin == 0 ){
      Z.slice( (iter + 1)/thin - 1 ) = temp_Z; 
      alpha_t.slice(  (iter + 1)/thin - 1 ) = temp_alpha_t;
      mu.col( (iter + 1)/thin - 1 ) = temp_mu;
      zeta.slice(  (iter + 1)/thin - 1 ) = temp_zeta;
      beta_eta.slice(  (iter + 1)/thin - 1 ) = temp_beta_eta;
      omega_eta.slice(  (iter + 1)/thin - 1 ) = temp_omega_eta;
      beta_gamma.slice( (iter +1)/thin - 1 ) = temp_beta_gamma;
      a_tc.slice(  (iter + 1)/thin - 1 ) = temp_a_tc;
    }
    
    // Print out progress
    double printer = iter % 250;
    
    if( printer == 0 ){
      Rcpp::Rcout << "Iteration = " << iter << std::endl;
    }
  }
  
  // Return output
  List output( 8 ); 
  output[ 0 ] = Z;  
  output[ 1 ] = alpha_t;
  output[ 2 ] = zeta;
  output[ 3 ] = beta_eta;
  output[ 4 ] = omega_eta;
  output[ 5 ] = a_tc;
  output[ 6 ] = beta_gamma;
  output[ 7 ] = mu;
  return output ; 
}

// [[Rcpp::export]]
List MicroMissBVS( 
    int iterations,
    int thin, 
    int N, // Number of hosts
    int T, // Number of true classifications
    int C, // Number of observed classifications 
    int P_eta, // Number of columns in design matrix for probability of occurence 
    int P_gamma, // Number of columns in design matrix for probability of relative abundance
    arma::mat Y, // Observed classifications N X C - matrix
    arma::cube Z, // True classification  N X T X M - array 
    arma::vec Z_sum, // Sum of observations N - vector
    arma::cube alpha_t, // Aux var for occupancy probability N X T X M - cube
    arma::cube zeta, // Occupancy indicator N X T X M - cube
    arma::cube beta_eta, // Probability of occupancy T X P_eta X M - cube
    arma::cube zeta_eta, // Inclusion indicator for beta_eta T X P_eta X M - cube
    arma::mat X_eta, // Design matrix for occupancy N X P_eta - matrix
    arma::cube omega_eta, // Aux var for occupancy indicator N X T X M - cube
    arma::cube beta_gamma, // Probability of relative abundance T X P_gamma X M - cube
    arma::cube zeta_gamma, // Inclusion indicator for beta_gamma T X P_gamma X M - cube
    arma::mat X_gamma, // Design matrix for relative abundance N X P_gamma - matrix
    arma::mat mu, // Aux var for DM trick for true classification prob N - vector (not saved over M) 
    arma::cube a_tc, //  Aux var for missclassification probability T X C X M - cube
    arma::vec u_t, // Aux var for DM trick for missclassification prob T - vector (not saved)
    arma::mat nu, // Concentration parameter for a_tc T X C - matrix (fixed)
    arma::mat mu_beta_eta_prior, // Prior mean of beta_eta
    arma::mat mu_beta_gamma_prior, // Prior mean of beta_gamma 
    double sigma2_beta_eta, // Prior variance of beta_eta 
    double sigma2_beta_gamma, // Prior variance of beta_gamma
    double a_zeta_gamma,
    double b_zeta_gamma,
    double a_zeta_eta,
    double b_zeta_eta 
){
  
  // Initiate memory
  List between_beta_gamma_zeta( 3 ); 
  List between_beta_eta_zeta( 2 );
  arma::mat temp_Z( N, T ); 
  arma::mat temp_alpha_t( N, T ); 
  arma::mat temp_zeta( N, T ); 
  arma::mat temp_beta_eta( T, P_eta );
  arma::mat temp_zeta_eta( T, P_eta );
  arma::mat temp_omega_eta( N, T );
  arma::vec temp_mu( N );
  arma::mat temp_a_tc( T, C );
  arma::vec temp_u_t( T );
  arma::mat temp_beta_gamma( T, P_gamma ); 
  arma::mat temp_zeta_gamma( T, P_gamma ); 
  arma::mat temp_tc( T, C );
  
  temp_Z = Z.slice( 0 );  
  temp_alpha_t = alpha_t.slice( 0 );
  temp_zeta = zeta.slice( 0 );
  temp_beta_eta = beta_eta.slice( 0 );
  temp_zeta_eta = zeta_eta.slice( 0 );
  temp_omega_eta = omega_eta.slice( 0 );
  temp_mu = mu.col( 0 );
  temp_a_tc = a_tc.slice( 0 );
  temp_u_t = u_t;
  temp_beta_gamma = beta_gamma.slice( 0 );
  temp_zeta_gamma = zeta_gamma.slice( 0 );
  temp_tc.zeros();
  
  // Set initial loggamma
  arma::mat temp_loggamma( N, T );
  temp_loggamma.zeros();
  for( int n = 0; n < N; ++n ){
    for( int t = 0; t < T; ++t ){
      temp_loggamma( n, t ) = ( X_gamma.row( n )*temp_beta_gamma.row( t ).t() ).eval()( 0, 0 );
    }
  }
  
  // Looping over the number of iterations specified by user
  for( int iter = 0; iter < iterations; ++iter ){
  
    // Joint update for alpha_t and zeta
    // Between Step: Expand or Contract   
    List between_return = help::between_alpha_t_zeta_update_cov_cpp( temp_Z, Z_sum, X_eta, temp_beta_eta, temp_alpha_t, temp_zeta, temp_loggamma, N );
    temp_alpha_t = as<arma::mat>( between_return[ 0 ] );
    temp_zeta = as<arma::mat>( between_return[ 1 ] );
 
    // Within Step
    temp_alpha_t = help::alpha_t_update_cov_cpp( temp_Z, temp_zeta, temp_loggamma, temp_mu, N );
 
    // Update mu
    temp_mu = help::mu_update_cpp( temp_Z, Z_sum, temp_alpha_t );
    
    // Update omega_eta
    temp_omega_eta = help::update_omega_eta( X_eta, temp_beta_eta, T );
    
    // Update beta_eta within
    temp_beta_eta = help::update_beta_eta_population( X_eta, temp_omega_eta, temp_zeta, mu_beta_eta_prior, sigma2_beta_eta, T, N );
 
    // Add/Delete beta_eta 
    between_beta_eta_zeta = help::between_beta_eta_zeta_update_BB_cpp( X_eta, temp_zeta_eta, temp_beta_eta, temp_alpha_t, sigma2_beta_eta, a_zeta_eta, b_zeta_eta );
    temp_zeta_eta = as<arma::mat>( between_beta_eta_zeta[ 0 ] );
    temp_beta_eta = as<arma::mat>( between_beta_eta_zeta[ 1 ] );
  
    // Update beta_gamma within
    List beta_gamma_out = help::within_beta_gamma_BVS_augment_cpp( X_gamma, temp_zeta_gamma, temp_beta_gamma, temp_loggamma, temp_alpha_t, temp_zeta, mu_beta_gamma_prior, sigma2_beta_gamma, N );
    temp_beta_gamma =  as<arma::mat>( beta_gamma_out[ 0 ] );
    temp_loggamma = as<arma::mat>( beta_gamma_out[ 1 ] );
 
    // Add/Delete beta_gamma
    between_beta_gamma_zeta = help::between_beta_gamma_zeta_augment_update_BB_cpp( X_gamma, temp_zeta_gamma, temp_beta_gamma, temp_loggamma, temp_alpha_t, temp_zeta, sigma2_beta_gamma, a_zeta_gamma, b_zeta_gamma, N );
    temp_zeta_gamma = as<arma::mat>( between_beta_gamma_zeta[ 0 ] );
    temp_beta_gamma = as<arma::mat>( between_beta_gamma_zeta[ 1 ] );
    temp_loggamma = as<arma::mat>( between_beta_gamma_zeta[ 2 ] ); 
  
    // Update Z
    List Z_out = help::update_Z( Y, temp_alpha_t, temp_a_tc, N, T, C );
    temp_Z =  as<arma::mat>( Z_out[ 0 ] );
    temp_tc = as<arma::mat>( Z_out[ 1 ] ); 
    
    // Update a_tc
    temp_a_tc = help::a_tc_update_cpp( temp_tc, nu, temp_u_t, T, C );
    
    // Update u_t
    temp_u_t = help::u_t_update_cpp( temp_Z, temp_a_tc, T );
    
    // Set the starting values for the next iteration
    if( (iter + 1) % thin == 0 ){
      Z.slice( (iter + 1)/thin - 1 ) = temp_Z; 
      alpha_t.slice(  (iter + 1)/thin - 1 ) = temp_alpha_t;
      mu.col( (iter + 1)/thin - 1 ) = temp_mu;
      zeta.slice(  (iter + 1)/thin - 1 ) = temp_zeta;
      beta_eta.slice(  (iter + 1)/thin - 1 ) = temp_beta_eta;
      omega_eta.slice(  (iter + 1)/thin - 1 ) = temp_omega_eta;
      beta_gamma.slice( (iter +1)/thin - 1 ) = temp_beta_gamma;
      zeta_gamma.slice( (iter +1)/thin - 1 ) = temp_zeta_gamma;
      zeta_eta.slice( (iter +1)/thin - 1 ) = temp_zeta_eta;
      a_tc.slice(  (iter + 1)/thin - 1 ) = temp_a_tc;
    }
    
    // Print out progress
    double printer = iter % 250;
    
    if( printer == 0 ){
      Rcpp::Rcout << "Iteration = " << iter << std::endl;
    }
  }
  
  // Return output
  List output( 10 ); 
  output[ 0 ] = Z;  
  output[ 1 ] = alpha_t;
  output[ 2 ] = zeta;
  output[ 3 ] = beta_eta;
  output[ 4 ] = omega_eta;
  output[ 5 ] = a_tc;
  output[ 6 ] = beta_gamma;
  output[ 7 ] = mu;
  output[ 8 ] = zeta_eta;
  output[ 9 ] = zeta_gamma;
  return output ; 
}

// 
// [[Rcpp::export]]
List ZIDM( 
    int iterations,
    int thin, 
    int N, // Number of hosts
    int T, // Number of true classifications
    int P_eta, // Number of columns in design matrix for probability of occurence 
    int P_gamma, // Number of columns in design matrix for probability of relative abundance
    arma::mat Y, // Observed classifications N X C - matrix
    arma::vec Z_sum, // Sum of observations N - vector
    arma::cube alpha_t, // Aux var for occupancy probability N X T X M - cube
    arma::cube zeta, // Occupancy indicator N X T X M - cube
    arma::cube beta_eta, // Probability of occupancy T X P_eta X M - cube
    arma::mat X_eta, // Design matrix for occupancy N X P_eta - matrix
    arma::cube omega_eta, // Aux var for occupancy indicator N X T X M - cube
    arma::cube beta_gamma, // Probability of relative abundance T X P_gamma X M - cube
    arma::mat X_gamma, // Design matrix for relative abundance N X P_gamma - matrix
    arma::vec mu, // Aux var for DM trick for true classification prob N - vector (not saved over M) 
    arma::mat mu_beta_eta_prior, // Prior mean of beta_eta
    arma::mat mu_beta_gamma_prior, // Prior mean of beta_gamma 
    double sigma2_beta_eta, // Prior variance of beta_eta 
    double sigma2_beta_gamma // Prior variance of beta_gamma
){
  
  // Initiate memory
  // arma::mat temp_Z( N, T ); 
  arma::mat temp_alpha_t( N, T ); 
  arma::mat temp_zeta( N, T ); 
  arma::mat temp_beta_eta( T, P_eta );
  arma::mat temp_omega_eta( N, T );
  arma::vec temp_mu( N ); 
  arma::mat temp_beta_gamma( T, P_gamma );  
  
  // temp_Z = Z.slice( 0 );  
  temp_alpha_t = alpha_t.slice( 0 );
  temp_zeta = zeta.slice( 0 );
  temp_beta_eta = beta_eta.slice( 0 );
  temp_omega_eta = omega_eta.slice( 0 );
  temp_mu = mu; 
  temp_beta_gamma = beta_gamma.slice( 0 ); 
  
  // Set initial loggamma
  arma::mat temp_loggamma( N, T );
  temp_loggamma.zeros();
  for( int n = 0; n < N; ++n ){
    for( int t = 0; t < T; ++t ){
      temp_loggamma( n, t ) = ( X_gamma.row( n )*temp_beta_gamma.row( t ).t() ).eval()( 0, 0 );
    }
  }
  
  // Looping over the number of iterations specified by user
  for( int iter = 0; iter < iterations; ++iter ){
    
    // Joint update for alpha_t and zeta
    // Between Step: Expand or Contract   
    List between_return = help::between_alpha_t_zeta_update_cov_cpp( Y, Z_sum, X_eta, temp_beta_eta, temp_alpha_t, temp_zeta, temp_loggamma, N );
    temp_alpha_t = as<arma::mat>( between_return[ 0 ] );
    temp_zeta = as<arma::mat>( between_return[ 1 ] );
    
    // Within Step
    temp_alpha_t = help::alpha_t_update_cov_cpp( Y, temp_zeta, temp_loggamma, temp_mu, N );
    
    // Update mu
    temp_mu = help::mu_update_cpp( Y, Z_sum, temp_alpha_t );
    
    // Update omega_eta
    temp_omega_eta = help::update_omega_eta( X_eta, temp_beta_eta, T );
    
    // Update beta_eta
    temp_beta_eta = help::update_beta_eta_population( X_eta, temp_omega_eta, temp_zeta, mu_beta_eta_prior, sigma2_beta_eta, T, N );
    
    // Update beta_gamma
    List beta_gamma_out = help::within_beta_gamma_augment_cpp( X_gamma, temp_beta_gamma, temp_loggamma, temp_alpha_t, temp_zeta, mu_beta_gamma_prior, sigma2_beta_gamma, N );
    temp_beta_gamma =  as<arma::mat>( beta_gamma_out[ 0 ] );
    temp_loggamma = as<arma::mat>( beta_gamma_out[ 1 ] ); 
    
    // Set the starting values for the next iteration
    if( (iter + 1) % thin == 0 ){
      // Z.slice( (iter + 1)/thin - 1 ) = temp_Z; 
      alpha_t.slice(  (iter + 1)/thin - 1 ) = temp_alpha_t;
      zeta.slice(  (iter + 1)/thin - 1 ) = temp_zeta;
      beta_eta.slice(  (iter + 1)/thin - 1 ) = temp_beta_eta;
      omega_eta.slice(  (iter + 1)/thin - 1 ) = temp_omega_eta;
      beta_gamma.slice( (iter +1)/thin - 1 ) = temp_beta_gamma; 
    }
    
    // Print out progress
    double printer = iter % 250;
    
    if( printer == 0 ){
      Rcpp::Rcout << "Iteration = " << iter << std::endl;
    }
  }
  
  // Return output
  List output( 5 );  
  output[ 0 ] = alpha_t;
  output[ 1 ] = zeta;
  output[ 2 ] = beta_eta;
  output[ 3 ] = omega_eta; 
  output[ 4 ] = beta_gamma;
  return output ; 
}


// 
// [[Rcpp::export]]
List DM( 
    int iterations,
    int thin, 
    int N, // Number of hosts
    int T, // Number of true classifications 
    int P_gamma, // Number of columns in design matrix for probability of relative abundance
    arma::mat Y, // Observed classifications N X C - matrix
    arma::vec Z_sum, // Sum of observations N - vector
    arma::cube alpha_t, // Aux var for occupancy probability N X T X M - cube 
    arma::cube beta_gamma, // Probability of relative abundance T X P_gamma X M - cube
    arma::mat X_gamma, // Design matrix for relative abundance N X P_gamma - matrix
    arma::vec mu, // Aux var for DM trick for true classification prob N - vector (not saved over M)  
    arma::mat mu_beta_gamma_prior, // Prior mean of beta_gamma  
    double sigma2_beta_gamma // Prior variance of beta_gamma
){
  
  // Initiate memory
  // arma::mat temp_Z( N, T ); 
  arma::mat temp_alpha_t( N, T );  
  arma::mat temp_omega_eta( N, T );
  arma::vec temp_mu( N ); 
  arma::mat temp_beta_gamma( T, P_gamma );  
  
  // temp_Z = Z.slice( 0 );  
  temp_alpha_t = alpha_t.slice( 0 ); 
  temp_mu = mu; 
  temp_beta_gamma = beta_gamma.slice( 0 ); 
  
  // Set initial loggamma
  arma::mat temp_loggamma( N, T );
  temp_loggamma.zeros();
  for( int n = 0; n < N; ++n ){
    for( int t = 0; t < T; ++t ){
      temp_loggamma( n, t ) = ( X_gamma.row( n )*temp_beta_gamma.row( t ).t() ).eval()( 0, 0 );
    }
  }
  
  // Looping over the number of iterations specified by user
  for( int iter = 0; iter < iterations; ++iter ){
    
    // Within Step
    temp_alpha_t = help::alpha_t_update_DM_cov_cpp( Y, temp_loggamma, temp_mu, N );
    
    // Update mu
    temp_mu = help::mu_update_cpp( Y, Z_sum, temp_alpha_t ); 
    
    // Update beta_gamma
    List beta_gamma_out = help::within_beta_gamma_DM_augment_cpp( X_gamma, temp_beta_gamma, temp_loggamma, temp_alpha_t, mu_beta_gamma_prior, sigma2_beta_gamma, N );
    temp_beta_gamma =  as<arma::mat>( beta_gamma_out[ 0 ] );
    temp_loggamma = as<arma::mat>( beta_gamma_out[ 1 ] ); 
    
    // Set the starting values for the next iteration
    if( (iter + 1) % thin == 0 ){
      // Z.slice( (iter + 1)/thin - 1 ) = temp_Z; 
      alpha_t.slice(  (iter + 1)/thin - 1 ) = temp_alpha_t; 
      beta_gamma.slice( (iter +1)/thin - 1 ) = temp_beta_gamma; 
    }
    
    // Print out progress
    double printer = iter % 250;
    
    if( printer == 0 ){
      Rcpp::Rcout << "Iteration = " << iter << std::endl;
    }
  }
  
  // Return output
  List output( 2 );  
  output[ 0 ] = alpha_t; 
  output[ 1 ] = beta_gamma;
  return output ; 
}

// Function ::  Calculate true multiplicative change in psi_j with 1 unit increase in covariate 1 over MCMC 
// [[Rcpp::export]] 
arma::cube effects_all( arma::cube beta, arma::mat cov, int T, int M, int P ) {
  int N = cov.n_rows;
  
  arma::cube estimate( N, T, M );
  estimate.zeros();
  for( int m = 0; m < M; ++m ){
    for( int n = 0; n < N; ++n){
      arma::vec ind_cov = cov.row( n ).t();
      arma::vec ind_cov_plus = ind_cov;
      ind_cov_plus[ P ] = ind_cov_plus[ P ] + 1;  
      for( int t = 0; t < T; ++t){ 
        estimate( n, t, m ) = exp( beta( t, P, m ) )*sum( exp( beta.slice( m )*ind_cov ) ) / sum( exp( beta.slice( m )*ind_cov_plus ) );
      }
    }
  } 
  return estimate;
}
 