# To propely install cmdstan, follow the instructions at:
# https://mc-stan.org/cmdstanr/articles/cmdstanr.html

library(cmdstanr)
library(posterior)
library(bayesplot)
library(dplyr)
library(ggplot2)
library(tidyverse)

color_scheme_set("brightblue")

source("prep_data.R")
source("prep_data.R")
source("utils_plot.R")

##########################
## TOY EXAMPLE: Bernoulli
##########################

file <- file.path(cmdstan_path(), "examples", "bernoulli", "bernoulli.stan")
mod <- cmdstan_model(file)
# names correspond to the data block in the Stan program
# The $sample() method for CmdStanModel objects runs Stan's default MCMC algorithm
data_list <- list(N = 10, y = c(0,1,0,0,0,0,0,0,0,1))

fit <- mod$sample(
  data = data_list,
  seed = 123,
  chains = 4,
  parallel_chains = 4,
  refresh = 500 # print update every 500 iters
)

fit$summary()
fit$summary(variables = c("theta", "lp__"), "mean", "sd")

# The $draws() method can be used to extract the posterior draws in formats provided by the posterior package.
draws_arr <- fit$draws() # or format="array"
str(draws_arr)
# plotting draws
mcmc_hist(fit$draws("theta"))
# The $sampler_diagnostics() method extracts the values of the sampler parameters (treedepth__, divergent__, etc.) in formats supported by the posterior package. 
str(fit$sampler_diagnostics(format = "df"))

# Optimization and VI
fit_mle <- mod$optimize(data = data_list, seed = 123)
fit_mle$print() # includes lp__ (log prob calculated by Stan program)
fit_map <- mod$optimize(
  data = data_list,
  jacobian = TRUE,
  seed = 123
)

# Laplace approximation
# The $laplace() method produces a sample from a normal approximation centered at the mode of a distribution in the unconstrained space.
fit_laplace <- mod$laplace(
  mode = fit_map,
  draws = 4000,
  data = data_list,
  seed = 123,
  refresh = 1000
)
fit_laplace$print("theta")
mcmc_hist(fit_laplace$draws("theta"), binwidth = 0.025)

# We can run Stan's experimental Automatic Differentiation Variational Inference (ADVI) using the $variational() method. For details on the ADVI algorithm
fit_vb <- mod$variational(
  data = data_list,
  seed = 123,
  draws = 4000
)
fit_vb$print("theta")

# Variational 'pathfinder'
fit_pf <- mod$pathfinder(
  data = data_list,
  seed = 123,
  draws = 4000
)
mcmc_hist(fit_pf$draws("theta"), binwidth = 0.025) +
  ggplot2::labs(subtitle = "Approximate posterior from pathfinder") +
  ggplot2::xlim(0, 1)

mcmc_hist(fit_vb$draws("theta"), binwidth = 0.025) +
  ggplot2::labs(subtitle = "Approximate posterior from variational") +
  ggplot2::xlim(0, 1)

mcmc_hist(fit_laplace$draws("theta"), binwidth = 0.025) +
  ggplot2::labs(subtitle = "Approximate posterior from Laplace") +
  ggplot2::xlim(0, 1)

mcmc_hist(fit$draws("theta"), binwidth = 0.025) +
  ggplot2::labs(subtitle = "Posterior from MCMC") +
  ggplot2::xlim(0, 1)


############################
## Bayesian mixtures
###########################

# We first generate the data
n_gaussian <- 2
sigma <- .4
gaussian_data <- random_gaussian_mixture(
  n_gaussian=n_gaussian,
  n_per_gaussian=rep(25, n_gaussian),
  means=c(-1,1),
  sigmas=rep(sigma, n_gaussian)
)
gaussian_data$sigma <- sigma
hist(gaussian_data$y, breaks = 20)


file_mix <- file.path("gaussian_mixture.stan")
mod_mix <- cmdstan_model(file_mix)

## fit 
# mcmc
fit_mcmc <- mod_mix$sample(
  data = gaussian_data,
  seed = 123,
  chains = 4,
  parallel_chains = 4,
  refresh = 500 # print update every 500 iters
)

# pathfinder
fit_pf <- mod_mix$pathfinder(
  data = gaussian_data,
  seed = 123,
  draws = 4000
)

plot_comparison_mixture_means(fit_mcmc, fit_pf)
ggsave(file="mix_means.pdf", width =8, height = 7)
plot_comparison_mixture_proportions(fit_mcmc, fit_pf)
ggsave(file="mix_prop.pdf", width =8, height = 7)

# print results
fit_pf$print()
