############################################################################################################################
# Fitting a model in Stan: from the Vignettes in RStan: https://cran.r-project.org/web/packages/rstan/vignettes/rstan.html #
############################################################################################################################
library(rstudioapi)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path)) 

library(rstan) 

# Eight Schools example:  
# A hierarchical meta-analysis model  to model the effect of coaching programs on college admissions tests

# Data declaration 
schools_data <- list(
  J = 8,                                     # Number of schools  
  y = c(28,  8, -3,  7, -1,  1, 18, 12),     # Vector of estimates
  sigma = c(15, 10, 16, 11,  9, 11, 10, 18)  # Vector of standard errors of the estimates
)


# Sample from the Posterior Distribution
fit1 <- stan(
  file = "schools.stan",  # Stan program
  data = schools_data,    # named list of data
  chains = 4,             # number of Markov chains (default)
  warmup = 1000,          # number of warmup iterations per chain 
  iter = 2000,            # total number of iterations per chain (default)
  cores = 2,              # number of cores (could use one per chain) 
  refresh = 0             # no progress shown
)


help(stan)

# extracting the estimates
sims <- extract(fit1)
str(sims)

# extracting the estimates in a matrix format 
sims2 <- as.matrix(fit1) 
str(sims2)
head(sims2)

# Print the summaries 
print(fit1)
print(fit1, pars = c("theta", "mu", "tau", "lp__"), probs = c(.1, .5, .9))
print(fit1, pars = c("mu", "tau"))

# Plot of osterior uncertainty intervals (by default 80% (inner) and 95% (outer)) and the posterior median for all the parameters
plot(fit1)

# Trace plot 
traceplot(fit1, pars = c("mu", "tau"), inc_warmup = TRUE, nrow = 2)

# Using the bayes plot package : see the vignettes https://cran.r-project.org/web/packages/bayesplot/vignettes/plotting-mcmc-draws.html
library(bayesplot)
library(ggplot2)
??bayesplot
theme_set(bayesplot::theme_default())

#Extract the posterior draws 
posterior <- as.array(fit1)

# mcmc_<something> is able to produce specific plots
# However, it is better to use the arguments pars or regex_pars for selecting the desired parameters

# posterior intervals
mcmc_intervals(posterior, regex_pars = c("theta", "tau", "mu" )) +
  ggtitle("Posterior intervals") 

# posterior areas
mcmc_areas(posterior, pars=c(  "theta[1]", "theta[2]", 
                               "theta[3]", "theta[4]", "theta[5]", "theta[6]",
                               "theta[7]", "theta[8]",
                               "tau", "mu" )) +
  ggtitle("Posterior areas")

# marginal posteriors
mcmc_dens(posterior)

# marginal posteriors separated for each chain
mcmc_dens_overlay(posterior)

# bivariate plots
mcmc_pairs(posterior, pars = c("mu", "tau"))

# trace plots
mcmc_trace(posterior)