// cSSVSquantreg.cc is a function that uses stochastic search variable selection
// to select promising models at a pre-specified quantile.
//
// The initial version of this file was generated by the
// auto.Scythe.call() function in the MCMCpack R package
// written by:
//
// Andrew D. Martin
// Dept. of Political Science
// Washington University in St. Louis
// admartin@wustl.edu
//
// Kevin M. Quinn
// Dept. of Government
// Harvard University
// kevin_quinn@harvard.edu
//
// This software is distributed under the terms of the GNU GENERAL
// PUBLIC LICENSE Version 2, June 1991.  See the package LICENSE
// file for more information.
//
// Copyright (C) 2009 Andrew D. Martin and Kevin M. Quinn
//
// This file was initially generated on June 1 2009
//
// The function was rewritten by:
//
// Craig Reed
// Department of Mathematical Sciences
// Brunel University
// craig.reed@brunel.ac.uk

#ifndef CSSVSQUANTREG_CC
#define CSSVSQUANTREG_CC

#include "MCMCrng.h"
#include "MCMCfcds.h"
#include "matrix.h"
#include "distributions.h"
#include "stat.h"
#include "la.h"
#include "ide.h"
#include "smath.h"
#include "rng.h"

#include <R.h>           // needed to use Rprintf()
#include <R_ext/Utils.h> // needed to allow user interrupts

using namespace std;
using namespace scythe;

struct COV_TRIAL_PREP{
Matrix<> C;
Matrix<> U;
double logdetminhalf;
};

// Helper function to prepare for the sequence of draws from the covariate indicators

static inline COV_TRIAL_PREP
QR_SSVS_covariate_trials_prep(const Matrix<>& X_gamma, const Matrix<>& Y, const Matrix<bool>& gamma, const Matrix<>& weights, const Matrix<>& lambda, double tau, unsigned int n_cov, unsigned int q){

const unsigned int n_obs=Y.rows();
Matrix<> U(Y);

// obtain cholesky decomposition of (Xtilde, utilde)
	if (tau!=0.5){
	U -= (1.0-2.0*tau)*weights;
	}
	Matrix<> XU(U.rows(),n_cov+1,false);
        if (n_cov == 0){
	XU = U;
	}
	else{
	XU = cbind(X_gamma,U);
	}

        Matrix<> XUtwXU(n_cov+1,n_cov+1,false);

	double temp_xu = 0.0;

//Calculate XUtwXU
  for (unsigned int i=0; i<n_cov+1; ++i){
    for (unsigned int j=0; j<(i+1); ++j){
      for (unsigned int m=0; m<n_obs; ++m){
	temp_xu += XU(m,i) * XU(m,j) / weights(m);
      }
      XUtwXU(i,j) = temp_xu;
      XUtwXU(j,i) = temp_xu;
      temp_xu = 0.0;
    }
  }

XUtwXU *= 0.5;
const Matrix<> lambda_gamma = selif(lambda,gamma(q,0,gamma.rows()-1,0));

for (unsigned int j=q; j<n_cov; ++j){

XUtwXU(j,j) += lambda_gamma(j-q);
}

Matrix<> C = cholesky(XUtwXU);

double logdetminhalf = 0.0;
// Work out -1/2 log(det(C'C))
for (unsigned int r=0; r<n_cov; ++r){
    logdetminhalf -= std::log(C(r,r));
}
COV_TRIAL_PREP result;
result.logdetminhalf = logdetminhalf;
result.U = U;
result.C = C;
return result;
}

/* SSVSquantreg implementation.  Takes Matrix<> reference which it
 * fills with the posterior of the indicator variables.
 */
template <typename RNGTYPE>
void SSVSquantreg_impl (rng<RNGTYPE>& stream, double tau, Matrix<>& Y,
    const Matrix<>& X, unsigned int q,
    double pi0a0, double pi0b0,
    unsigned int burnin, unsigned int mcmc,
    unsigned int thin, unsigned int verbose,
    Matrix<>& result)

{
   // define constants
   unsigned int n_obs = X.rows();
   unsigned int k = X.cols();
   double tau_onemintau = tau*(1.0-tau);
   const unsigned int tot_iter = burnin + mcmc; //total iterations
   const unsigned int nstore = mcmc / thin; // number of draws to store

   // storage matrices
   Matrix<bool> gamma(k,1,true,true);
   Matrix<> gamma_matrix (k, nstore);
   Matrix<> beta(k,1,true,0.0);
   Matrix<> beta_matrix(k, nstore);

   COV_TRIAL trial;
   COV_TRIAL_PREP trial_prep;

   unsigned int n_cov = k;

   // Matrices that will change dimension
   Matrix<> *X_gamma = new Matrix<>(X);
   Matrix<> *C = new Matrix<>(n_cov+1,n_cov+1,false);
   Matrix<> *beta_gamma = new Matrix<>(n_cov,1,false);

   // Matrix with same dimensions
   Matrix<> U(n_obs,1,false);

   // Gibbs sampler
   unsigned int count = 0;
   unsigned int col_index = q;

   // Initial value of pi0
   double pi0 = stream.rbeta(pi0a0,pi0b0);
   // Initial values for lambda
   Matrix<> lambda = stream.rgamma(n_cov-q,1,0.5,0.5);
   // Initial values for weights
   Matrix<> weights = stream.rexp(n_obs,1,tau_onemintau);
   for (unsigned int iter = 0; iter < tot_iter; ++iter) {

trial_prep = QR_SSVS_covariate_trials_prep(*X_gamma, Y, gamma, weights, lambda, tau, n_cov, q);
C = new Matrix<>(trial_prep.C);
double logdetminhalf = trial_prep.logdetminhalf;
U = trial_prep.U;

// updating the indicator variables corresponding to each covariate

    for (unsigned int j=q; j<k; ++j){

	//obtain column index
	for (unsigned int m=q; m<j; ++m){
	if (gamma(m) == true){
	++col_index;
	}
	}

	if (gamma(j) == false){
	trial = QR_SSVS_covariate_trials_draw_absent(*C,*X_gamma,trial_prep.U,X(_,j),col_index,weights,pi0,lambda(j-q),logdetminhalf,stream);
        gamma(j) = trial.newtrial;
        logdetminhalf = trial.logdetminhalf;
	if (trial.newtrial == true){
	delete C;
        C = new Matrix<> (trial.Cnew);
	delete X_gamma;
	X_gamma = new Matrix<> (t(selif(t(X),gamma)));
	++n_cov;
	}
        }

	else{
	trial = QR_SSVS_covariate_trials_draw_present(*C,col_index,n_obs,pi0,lambda(j-q),logdetminhalf,stream);
	gamma(j) = trial.newtrial;
	logdetminhalf = trial.logdetminhalf;
        if (trial.newtrial == false){
	delete C;
	C = new Matrix<> (trial.Cnew);
	delete X_gamma;
	X_gamma = new Matrix<> (t(selif(t(X),gamma)));
        --n_cov;
	}
	}

     col_index = q;

     } // end covariate trials loop

if (n_cov == 0){ // Support for null model
weights = ALaplaceIGaussregress_weights_draw (fabs(Y), stream);
lambda = stream.rexp(k-q,1,0.5);
}

else{
beta_gamma = new Matrix<> (QR_SSVS_beta_draw (*C, stream));

if (q != 0){
beta(0,0,q-1,0) = (*beta_gamma)(0,0,q-1,0);
}

for (unsigned int j=q; j<k; ++j){

if (gamma(j) == true){
beta(j) = (*beta_gamma)(col_index);
++col_index;
}
}

col_index = q;

Matrix<> e = gaxpy(*X_gamma, (-1*(*beta_gamma)), Y);
Matrix<> abse = fabs(e);

weights = ALaplaceIGaussregress_weights_draw (abse, stream);

lambda = QR_SSVS_lambda_draw(*beta_gamma, gamma, k, q, stream);

delete beta_gamma;

}

delete C;

pi0 = QR_SSVS_pi0_draw(n_cov-q, k-q, pi0a0, pi0b0, stream);

    // store draws in storage matrices
       if (iter >= burnin && (iter % thin == 0)) {
       gamma_matrix(_, count) = gamma;
       beta_matrix(_, count) = beta;
       ++count;
       }

     // print output to stdout
     if(verbose > 0 && iter % verbose == 0) {
       Rprintf("\n\nSSVSquantreg iteration %i of %i \n",
         (iter+1), tot_iter);
       Rprintf("gamma = \n");
       for (unsigned int r=0; r<k; ++r)
         // Rprintf("%10.5f\n", gamma(r));
        Rprintf("beta = \n");
       for (unsigned int r=0; r<k; ++r)
        Rprintf("%10.5f\n", beta(r));
     }

     beta = Matrix<>(k,1,true,0.0);

     R_CheckUserInterrupt(); // allow user interrupts

   } // end MCMC loop

   delete X_gamma;

   result = cbind(t(gamma_matrix),t(beta_matrix));
 } // end SSVSquantreg_impl

extern "C" {
   void cSSVSquantreg(double *sampledata, const int *samplerow,
		    const int *samplecol, const double *tau, const double *Ydata, const int *Yrow,
		    const int *Ycol, const double *Xdata, const int *Xrow,
		    const int *Xcol, const int *q, const int *burnin, const int *mcmc, const int *thin,
		    const int *uselecuyer, const int *seedarray,
		    const int *lecuyerstream, const int *verbose,
		    const double *pi0a0, const double *pi0b0)
   {
     // pull together Matrix objects
     Matrix<> Y(*Yrow, *Ycol, Ydata);
     Matrix<> X(*Xrow, *Xcol, Xdata);
     Matrix<> storagematrix;
     MCMCPACK_PASSRNG2MODEL(SSVSquantreg_impl, *tau, Y, X, *q,
                            *pi0a0, *pi0b0, *burnin, *mcmc, *thin, *verbose,
                            storagematrix);

     const unsigned int size = *samplerow * *samplecol;
     for (unsigned int h = 0; h < size; ++h){
       sampledata[h] = storagematrix(h);
   }
}
}

#endif

