#include<RcppArmadillo.h>
#include<Rmath.h>
#include<stdio.h>
#include"BVCUtilities.h"
// [[Rcpp::depends(RcppArmadillo)]]

using namespace Rcpp;
using namespace arma;
//using namespace R;


// [[Rcpp::export()]]
Rcpp::List BBL(arma::mat xx, arma::vec y, arma::vec ystar,arma::mat W, int maxSteps, arma::vec hatBeta, arma::vec hatAlpha, arma::vec hatInvTauSq, arma::mat invSigAlpha0, double hatLambdaSqStar, double hatSigmaSq, double aStar, double bStar, double alpha, double gamma, int progress)
{
  unsigned int n = xx.n_rows, s = xx.n_cols, clc = W.n_cols;
  arma::mat gsAlpha(maxSteps, clc),
  gsBeta(maxSteps, s),
  gsystar(maxSteps, n),
  gsInvTauSq(maxSteps, s);
  
  arma::vec gsLambdaStar(maxSteps),
  gsSigmaSq(maxSteps),
  gsMSE(maxSteps);
  
  arma::mat tWW = W.t()*W, varAlpha, matBeta;
  arma::vec res, meanAlpha, tRsRs, muInvTauSq,mustar,sdstar;
  double tempS, varRs, meanRs, lInvTauSq;
  
  arma::vec tBrBrDiag = sum(square(xx), 0).t();
  
  for (int k = 0; k < maxSteps; k++) {
    
    // Rcpp::Rcout << "ystar" << std::endl;
    mustar = xx * hatBeta;
    sdstar = sqrt(hatSigmaSq)*arma::ones(n);
    
    for (int i = 0; i < n; i++) {
      double mean = mustar(i);
      double sd = sdstar(i);
      
      if (y(i) == 0) {
        ystar(i) = rtnorm2(0, false, mean, sd);
      } else {
        ystar(i) = rtnorm2(0, true, mean, sd);
      }
    }
    
    gsystar.row(k) = ystar.t();
    
    // alpha|
    varAlpha = arma::inv(tWW/hatSigmaSq + invSigAlpha0);
    res = ystar - xx * hatBeta;
    meanAlpha = varAlpha * (W.t() * res/hatSigmaSq);
    hatAlpha = mvrnormCpp(meanAlpha, varAlpha);
    res -= W * hatAlpha;
    gsAlpha.row(k) = hatAlpha.t();
    
    for(unsigned int j=0; j<s; j++){		
      tempS = 1/(tBrBrDiag(j) + hatInvTauSq(j));
      varRs = hatSigmaSq * tempS;
      res += xx.col(j) * hatBeta(j);
      meanRs = arma::as_scalar(tempS * xx.col(j).t() * res);
      hatBeta(j) = R::rnorm(meanRs, std::sqrt(varRs));
      res -= xx.col(j) * hatBeta(j);
    }
    gsBeta.row(k) = hatBeta.t();
    
    // sigma.sq|
    double shapeSig = alpha + (n+s)/2;
    double rateSig = gamma + 0.5*(arma::accu(arma::square(res)) + 
                                  arma::accu(square(hatBeta) % hatInvTauSq));
    hatSigmaSq = 1/R::rgamma(shapeSig, 1/rateSig);
    gsSigmaSq(k) = hatSigmaSq;
    
    
    // invTAUsq.star|
    lInvTauSq = hatLambdaSqStar;
    tRsRs = arma::square(hatBeta);
    muInvTauSq = arma::sqrt(hatLambdaSqStar * hatSigmaSq / tRsRs);
    for(unsigned int j = 0; j<s; j++){
      hatInvTauSq(j) = rinvgaussian(muInvTauSq(j), lInvTauSq);
    }
    gsInvTauSq.row(k) = hatInvTauSq.t();
    
    
    // lambda.star|
    double shapeS = aStar + s;
    double rateS = bStar + arma::accu(1/hatInvTauSq)/2;
    hatLambdaSqStar = R::rgamma(shapeS, 1/rateS);
    gsLambdaStar(k) = hatLambdaSqStar;
    
    gsMSE(k) = arma::mean(arma::square(res));
    if(k % 100 == 0){
      Rcpp::checkUserInterrupt();
    }
    if(progress != 0 && k % progress == 0){
      Rcpp::Rcout << "\nIter." << k << "  mse: " << gsMSE(k) << std::endl;
    }
  }
  
  return Rcpp::List::create(Rcpp::Named("GS.alpha") = gsAlpha,
                            Rcpp::Named("GS.beta") = gsBeta,
                            Rcpp::Named("GS.invTAUsq") = gsInvTauSq,
                            Rcpp::Named("GS.ystar") = gsystar,
                            Rcpp::Named("GS.lambda.sq") = gsLambdaStar,
                            Rcpp::Named("GS.sigma.sq") = gsSigmaSq,
                            Rcpp::Named("GS.mse") = gsMSE);
}
