//////////////////////////////////////////////////////////////////////////
// cMCMCregress.cc is a program that simualates draws from the posterior
// density of a linear regression model with Gaussian errors.
//
// 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.
//
// This file was initially generated on Fri Jul 23 15:07:21 2004
//
// ADM and KQ 10/10/2002 [ported to Scythe0.3]
// ADM 6/2/04 [re-written using template]
// KQ 6/18/04 [modified to meet new developer specification]
// ADM 7/22/04 [modified to work with new Scythe and rngs]
// DBP 7/1/07 [ported to scythe 1.0.x]
//
// Copyright (C) 2003-2007 Andrew D. Martin and Kevin M. Quinn
// Copyright (C) 2007-present Andrew D. Martin, Kevin M. Quinn,
//    and Jong Hee Park
//////////////////////////////////////////////////////////////////////////

#ifndef CMCMCREGRESS_CC
#define CMCMCREGRESS_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;

static double digamma(double theta, double a, double b) {
  double logf =  a * log(b) - lngammafn(a) + -(a+1) * log(theta) +
                 -b/theta;
  return exp(logf);
  //pow(b, a) / gammafn(a) * pow(theta, -(a+1)) * exp(-b/theta);
}


/* cMCMCregress implementation.  Takes Matrix<> reference which it
 * fills with the posterior.  The logmarklike double reference is
 * filled with the log marginal likelihood if asked for.
 */
template <typename RNGTYPE>
void MCMCregress_impl (rng<RNGTYPE>& stream, const Matrix<>& Y,
		       const Matrix<>& X, Matrix<>& beta, const Matrix<>& b0,
		       const Matrix<>& B0, double c0, double d0,
		       unsigned int burnin, unsigned int mcmc, unsigned int thin,
		       unsigned int verbose, bool chib,
		       Matrix<>& result, double& logmarglike)
{
   // define constants and form cross-product matrices
   const unsigned int tot_iter = burnin + mcmc; //total iterations
   const unsigned int nstore = mcmc / thin; // number of draws to store
   const unsigned int k = X.cols ();
   const Matrix<> XpX = crossprod(X);
   const Matrix<> XpY = t(X) * Y;

   // storage matrices
   Matrix<> betamatrix (k, nstore);
   Matrix<> sigmamatrix (1, nstore);

   // Gibbs sampler
   unsigned int count = 0;
   for (unsigned int iter = 0; iter < tot_iter; ++iter) {
     double sigma2 = NormIGregress_sigma2_draw (X, Y, beta, c0, d0,
                                                stream);
     beta = NormNormregress_beta_draw (XpX, XpY, b0, B0, sigma2, stream);

     // store draws in storage matrix (or matrices)
     if (iter >= burnin && (iter % thin == 0)) {
       sigmamatrix (0, count) = sigma2;
       betamatrix(_, count) = beta;
       ++count;
     }

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

     R_CheckUserInterrupt(); // allow user interrupts

   } // end MCMC loop

   if (chib == 1) {
     // marginal likelihood calculation stuff starts here
     const double sigma2star = meanc(t(sigmamatrix))(0);
     const double sigmastar = sqrt(sigma2star);
     const Matrix<> betastar = t(meanc(t(betamatrix)));

     // step 1
     Matrix<> sigma2_density(nstore, 1);

     // second set of Gibbs scans
     for (unsigned int iter = 0; iter < nstore; ++iter) {
       // double sigma2 = sigmamatrix(iter);
       // beta = NormNormregress_beta_draw (XpX, XpY, b0, B0, sigma2,  stream);
       beta = betamatrix (_, iter);
       const Matrix<> e = gaxpy(X, (-1*beta), Y);
       const Matrix<> SSE = crossprod (e);
       const double c_post = (c0 + X.rows ()) * 0.5;
       const double d_post = (d0 + SSE[0]) * 0.5;
       sigma2_density(iter) = digamma(sigma2star, c_post, d_post);
       R_CheckUserInterrupt();
     } // end MCMC loop
     double pdf_sigma2 = log(mean(sigma2_density));

     // step 2
     const Matrix<> Bn = invpd (B0 + XpX /sigma2star);
     const Matrix<> bn = Bn * gaxpy(B0, b0, XpY/sigma2star);
     double pdf_beta = 0;
     if (k == 1){
       pdf_beta = log(dnorm(betastar(0), bn(0), sqrt(Bn(0))));
     }
     else {
       pdf_beta = lndmvn(betastar, bn, Bn);
     }

     // calculate loglikelihood at (betastar, sigma2star)
     Matrix<> eta = X * betastar;
     double loglike_sum = 0.0;
     for (unsigned int i = 0; i < X.rows(); ++i) {
       loglike_sum += lndnorm(Y(i), eta(i), sigmastar);
     }
     double loglike = loglike_sum;

     // calculate log prior ordinate
     double logprior = 0;
     if (k == 1){
       logprior = log(digamma(sigma2star, c0/2.0, d0/2.0)) +
	 log(dnorm(betastar(0), b0(0), 1/sqrt(B0(0))));
     }
     else{
       logprior = log(digamma(sigma2star, c0/2.0, d0/2.0)) +
                           lndmvn(betastar, b0, invpd(B0));
     }

     // put pieces together and print the marginal likelihood
     logmarglike = loglike + logprior - pdf_beta - pdf_sigma2;
     if (verbose >0 ){
       Rprintf("\nlogmarglike = %10.5f\n", logmarglike);
       Rprintf("loglike = %10.5f\n", loglike);
       Rprintf("log_prior = %10.5f\n", logprior);
       Rprintf("log_beta = %10.5f\n", pdf_beta);
       Rprintf("log_sigma2 = %10.5f\n", pdf_sigma2);
     }
   }

   result = cbind(t(betamatrix), t(sigmamatrix));
 } // end cMCMCregress

extern "C" {
   void cMCMCregress(double *sampledata, const int *samplerow,
		    const int *samplecol, const double *Ydata, const int *Yrow,
		    const int *Ycol, const double *Xdata, const int *Xrow,
		    const int *Xcol, const int *burnin, const int *mcmc,
		    const int *thin, const int *uselecuyer, const int *seedarray,
		    const int *lecuyerstream, const int *verbose,
		    const double *betastartdata, const int *betastartrow,
		    const int *betastartcol, const double *b0data,
		    const int *b0row, const int *b0col,
		    const double *B0data, const int *B0row,
		    const int *B0col, const double *c0, const double *d0,
		    double* logmarglikeholder, const int* chib)
   {
     // pull together Matrix objects
     Matrix<> Y(*Yrow, *Ycol, Ydata);
     Matrix<> X(*Xrow, *Xcol, Xdata);
     Matrix<> betastart(*betastartrow, *betastartcol, betastartdata);
     Matrix<> b0(*b0row, *b0col, b0data);
     Matrix<> B0(*B0row, *B0col, B0data);

     double logmarglike;
     Matrix<> storagematrix;
     MCMCPACK_PASSRNG2MODEL(MCMCregress_impl, Y, X, betastart, b0, B0,
                            *c0, *d0, *burnin, *mcmc, *thin, *verbose,
                            *chib, storagematrix, logmarglike);
     logmarglikeholder[0] = logmarglike;
     const unsigned int size = *samplerow * *samplecol;
     for (unsigned int i = 0; i < size; ++i)
       sampledata[i] = storagematrix(i);
   }
}

#endif
