########## R script: stacklossLOOanaWithInteracs ##########

# For Bayesian model comparison via the leave-one-out
# information criterion as implemented in the package
# 'loo'.

# Last changed: 15 JUN 2020

# Set tasks:

doMCMC <- TRUE
obtainLOOvalues <- TRUE

# Load required packages:

library(rstan)
library(loo)

# Set MCMC parameters:

nWarm <- 10000        # Length of burn-in.
nIter <- 10000        # Size of the kept sample.
nThin <- 1            # Thinning factor. 

# Load in stackloss data:

stackLoss <- stackloss$stack.loss
airFlow <- stackloss$Air.Flow
waterTemp <- stackloss$Water.Temp
acidConc <- stackloss$Acid.Conc

# Obtain standardised versions of response variable (taste)
# and the predictor variables (airFlow, waterTemp and acidConc) for
# Bayesian analysis:

y <- (stackLoss - mean(stackLoss))/sd(stackLoss)
x1 <- (airFlow - mean(airFlow))/sd(airFlow)
x2 <- (waterTemp - mean(waterTemp))/sd(waterTemp)
x3 <- (acidConc - mean(acidConc))/sd(acidConc)

# Set number of models:

numModels <- 18

if (doMCMC)
{
   # Define model for general linear regression:

   linRegModel <- 
     'data
      { 
         int<lower=1> n;             vector[n] y;    
         int<lower=1> ncX;           matrix[n,ncX] X;       
         real<lower=0> sigmabeta;    real<lower=0> ssigma;     
      }
      parameters 
      {
         vector[ncX] beta;      real<lower=0> sigma; 
      }
      model 
      {
         y ~ normal(X*beta,sigma);
         beta ~ normal(0,sigmabeta); sigma ~ cauchy(0,ssigma); 
      }
      generated quantities
      {
         vector[n] log_lik;
         for (i in 1:n)
         {
            log_lik[i] = normal_lpdf(y[i]|X[i]*beta,sigma);
         }
      }'

   stanObjList <- vector("list",numModels)

   for (iModel in 1:numModels)
   {
      if (iModel==1) X <- as.matrix(rep(1,length(y)))
      if (iModel==2) X <- cbind(1,x1)
      if (iModel==3) X <- cbind(1,x2)
      if (iModel==4) X <- cbind(1,x3)
      if (iModel==5) X <- cbind(1,x1,x2)
      if (iModel==6) X <- cbind(1,x1,x3)
      if (iModel==7) X <- cbind(1,x2,x3)
      if (iModel==8) X <- cbind(1,x1,x2,x3)
      if (iModel==9) X <- cbind(1,x1,x2,x1*x2)
      if (iModel==10) X <- cbind(1,x1,x3,x1*x3)
      if (iModel==11) X <- cbind(1,x2,x3,x2*x3)
      if (iModel==12) X <- cbind(1,x1,x2,x3,x1*x2)
      if (iModel==13) X <- cbind(1,x1,x2,x3,x1*x3)
      if (iModel==14) X <- cbind(1,x1,x2,x3,x2*x3)
      if (iModel==15) X <- cbind(1,x1,x2,x3,x1*x2,x1*x3)
      if (iModel==16) X <- cbind(1,x1,x2,x3,x1*x2,x2*x3)
      if (iModel==17) X <- cbind(1,x1,x2,x3,x1*x3,x2*x3)
      if (iModel==18) X <- cbind(1,x1,x2,x3,x1*x2,x1*x3,x2*x3)

      allData <- list(n=length(y),ncX=ncol(X),y=y,X=X,sigmabeta=1e5,ssigma=1e5)
      
      stanCompilObj <- stan(model_code=linRegModel,data=allData,
                               iter=1,chains=1)
      stanObjList[[iModel]] <-  stan(model_code=linRegModel,data=allData,
                                  warmup=nWarm,iter=(nWarm+nIter),
                                  chains=1,thin=nThin,refresh=1000,fit=stanCompilObj)
   }
}

# Obtain LOO values:

if (obtainLOOvalues)
{
   LOOvec <- rep(NA,numModels)
  
   for (iModel in 1:numModels)
   {
      currLOOobj <- loo(extract_log_lik(stanObjList[[iModel]])) 

      currEst <- currLOOobj$estimates

      LOOvec[iModel] <- currEst[3,1]
  
   }

   # Convert LOO values to a matrix form and print out:

   LOOmat <- data.frame(1:numModels,signif(LOOvec,4))
   names(LOOmat) <- c("model number","LOO information criterion")

   cat("\n\n\n\n\n")
   print(LOOmat)
   cat("\n\n\n\n\n")
}

########## End of stacklossLOOanaWithInteracs ##########