## Package and options
options(prompt = "R> ", continue = "+  ", width = 77, useFancyQuotes = FALSE)
RhpcBLASctl::blas_set_num_threads(1)

library("shrinkTVP")

# Baseline model
set.seed(123)
sim <- simTVP(theta = c(0.2, 0, 0), beta_mean = c(1.5, -0.3, 0))
data <- sim$data
res <- shrinkTVP(y ~ x1 + x2, data = data)

# Bayesian hierarchical Lasso
res_hierlasso <- shrinkTVP(y ~ x1 + x2, data = data, learn_a_xi = FALSE,
  learn_a_tau = FALSE, a_xi = 1, a_tau = 1, display_progress = FALSE)

# Bayesian non hierarchical Lasso
res_lasso <- shrinkTVP(y ~ x1 + x2, data = data, learn_a_xi = FALSE,
  learn_a_tau = FALSE, a_xi = 1, a_tau = 1, learn_kappa2_B = FALSE,
  learn_lambda2_B = FALSE, kappa2_B = 100, lambda2_B = 100,
  display_progress = FALSE)

# Triple gamma
res_tg <- shrinkTVP(y ~ x1 + x2, data = data, mod_type = "triple",
  display_progress = FALSE)

# Non-hierarchical horseshoe
res_hs <- shrinkTVP(y ~ x1 + x2, data = data, mod_type = "triple",
  learn_a_xi = FALSE, learn_a_tau = FALSE, a_xi = 0.5, a_tau = 0.5,
  learn_c_xi = FALSE, learn_c_tau = FALSE, c_xi = 0.5, c_tau = 0.5,
  learn_kappa2_B = FALSE, learn_lambda2_B = FALSE, display_progress = FALSE)

# Adding stochastic volatility
res_sv <- shrinkTVP(y ~ x1 + x2, data = data, sv = TRUE,
  display_progress = FALSE)

# Changing some of the hyperparameters
res_hyp <- shrinkTVP(y ~ x1 + x2, data = data,
  hyperprior_param = list(beta_a_xi = 5, alpha_a_xi = 10),
  display_progress = FALSE)

# Tuning the MH steps
res_MH <- shrinkTVP(y ~ x1 + x2, data = data,
  MH_tuning = list(a_xi_adaptive = FALSE, a_tau_max_adapt = 0.001,
  a_tau_batch_size = 20), display_progress = FALSE)

# Summary method
summary(res, showprior = FALSE)

# Making all the plots
plot(res)


library("RColorBrewer")
color <- brewer.pal(5, "RdBu")
plot(res, pars = "beta", xlim = c(100, 200), probs = seq(0.1, 0.9, by = 0.1),
  quantlines = TRUE, quantcol = color[5:2], quantlty = 1, quantlwd = 3,
  col = color[1], lwd = 3, shadecol = "gold1")

plot(res, pars = "theta_sr")

# Show how to calculate LPDS
res_LPDS <- shrinkTVP(y ~ x1 + x2, data = data[1:(nrow(data) - 1), ],
  display_progress = FALSE)
LPDS(res_LPDS, data[nrow(data), ])

# Evaluate predictive density
eval_pred_dens(1:3, res_LPDS, data[nrow(data), ])

# Plot predictive density
curve(eval_pred_dens(x, res_LPDS, data[nrow(data), ]), to = 12,
  ylab = bquote("p(" *
    y[t[0] + 1] * "｜" * y[1] * "," * ldots * "," ~ y[t[0]] *
    "," ~ x[t[0] + 1] *
    ")"), xlab = expression(y[t[0] + 1]), lwd = 2.5, col = "skyblue",
  axes = FALSE)
abline(v = data$y[nrow(data)])
axis(1)
axis(2)


# Example with usmacro.update
library("bvarsv")
data("usmacro.update", package = "bvarsv")

# Create matrix of lags and create final data set
lags <- usmacro.update[1:(nrow(usmacro.update) - 1), ]
colnames(lags) <- paste0(colnames(lags), "_lag")
us_data <- data.frame(inf = usmacro.update[2:nrow(usmacro.update), "inf"], lags)

# Baseline specification
us_res <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, us_data, niter = 60000,
  nburn = 10000, nthin = 10, display_progress = FALSE)

# Summary
summary(us_res, showprior = FALSE)

# Plot beta
plot(us_res, nplot = 4, xlab = "t")

# Plot theta
plot(us_res, "theta_sr", show.obs = FALSE)

# Calculate LPDS in multicore Load libraries for multicore computations
library("doParallel")
library("foreach")

# For manipulating dates
library("zoo")

# Load library for controlling number of BLAS threads
library("RhpcBLASctl")

# Define how many periods to calculate LPDS for
Tmax <- nrow(us_data) - 1
T0 <- Tmax - 49

# Determine number of cores to be used and register parallel backend
ncores <- 4
cl <- makeCluster(ncores)
registerDoParallel(cl)

lpds <- foreach(t = T0:Tmax, .combine = "cbind",
  .packages = c("RhpcBLASctl", "shrinkTVP"),
  .errorhandling = "pass") %dopar% {

  set.seed(t)

  niter <- 30000
  nburn <- 15000
  nthin <- 5

  # Set number of BLAS threads, so they dont interfere with each other
  blas_set_num_threads(1)

  # Create data_t from all data up to time t and y_test and x_test from data at
  # time t+1
  data_test <- us_data[t + 1, ]
  data_t <- us_data[1:t, ]

  # Run MCMC to calculate all LPDS Fully hierarchical triple gamma
  res_FH_TG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "triple", niter = niter, nburn = nburn, nthin = nthin)

  # Hierarchical triple gamma
  res_H_TG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "triple", niter = niter, nburn = nburn, nthin = nthin,
    learn_a_xi = FALSE, learn_a_tau = FALSE, learn_c_xi = FALSE,
    learn_c_tau = FALSE)

  # Non-hierarchical triple gamma
  res_TG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "triple", niter = niter, nburn = nburn, nthin = nthin,
    learn_kappa2_B = FALSE, learn_lambda2_B = FALSE, learn_a_xi = FALSE,
    learn_a_tau = FALSE, learn_c_xi = FALSE, learn_c_tau = FALSE)

  # Hierarchical horseshoe
  res_H_HS <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "triple",
    niter = niter, nburn = nburn, nthin = nthin, learn_a_xi = FALSE,
    learn_a_tau = FALSE, learn_c_xi = FALSE, learn_c_tau = FALSE, a_xi = 0.5,
    a_tau = 0.5, c_xi = 0.5, c_tau = 0.5)

  # Non-hierarchical horseshoe
  res_HS <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "triple", niter = niter, nburn = nburn, nthin = nthin,
    learn_kappa2_B = FALSE, learn_lambda2_B = FALSE,
    learn_a_xi = FALSE, learn_a_tau = FALSE, learn_c_xi = FALSE,
    learn_c_tau = FALSE, a_xi = 0.5, a_tau = 0.5, c_xi = 0.5, c_tau = 0.5)

  # Fully hierarchical double gamma
  res_FH_DG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    niter = niter, nburn = nburn, nthin = nthin,
    hyperprior_param = list(alpha_a_tau = 1, alpha_a_xi = 1))

  # Hierarchical double gamma
  res_H_DG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    niter = niter,
    nburn = nburn, nthin = nthin, learn_a_xi = FALSE, learn_a_tau = FALSE,
    a_xi = 0.1, a_tau = 0.1)

  # Non-hierarchical double gamma
  res_DG <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    niter = niter, nburn = nburn, nthin = nthin, learn_a_xi = FALSE,
    learn_a_tau = FALSE, a_xi = 0.1, a_tau = 0.1, learn_kappa2_B = FALSE,
    learn_lambda2_B = FALSE)

  # Hierarchical Lasso
  res_H_LS <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    niter = niter, nburn = nburn, nthin = nthin, learn_a_xi = FALSE,
    learn_a_tau = FALSE, a_xi = 1, a_tau = 1)

  # Non-hierarchical Lasso
  res_LS <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    niter = niter, nburn = nburn, nthin = nthin, learn_a_xi = FALSE,
    learn_a_tau = FALSE, a_xi = 1, a_tau = 1, learn_kappa2_B = FALSE,
    learn_lambda2_B = FALSE)

  # Fixed variance prior
  res_FV <- shrinkTVP(inf ~ inf_lag + une_lag + tbi_lag, data = data_t,
    mod_type = "ridge", niter = niter, nburn = nburn, nthin = nthin)


  lpds_res <- c(LPDS(res_FH_TG, data_test), LPDS(res_H_TG, data_test),
    LPDS(res_TG, data_test), LPDS(res_H_HS, data_test), 
    LPDS(res_HS, data_test), LPDS(res_FH_DG, data_test),
    LPDS(res_H_DG, data_test), LPDS(res_DG, data_test),
    LPDS(res_H_LS, data_test), LPDS(res_LS, data_test), LPDS(res_FV, data_test))

  rm(list = ls()[!ls() %in% c("lpds_res", "us_data")])
  return(lpds_res)
}
stopCluster(cl)


cumu_lpds <- apply(lpds, 1, cumsum)
color <- c(rep("cyan3", 3), rep("firebrick3", 2), rep("forestgreen", 3),
  rep("yellow2", 2), "black")
lty <- c(1:2, 5, 1:2, 1:2, 5, 1:2, 1)

# Plot results
par(mar = c(6, 4, 1, 1))
colnames(cumu_lpds) <- c("Fully hierarchical NGG", "Hierarchical NGG", "NGG",
   "Hierarchical Horseshoe", "Horseshoe", "Fully hierarchical NG",
   "Hierarchical NG", "NG", "Hierarchical Lasso", "Lasso", "Ridge")

matplot(cumu_lpds, type = "l", ylab = "Cumulative LPDS", xaxt = "n", xlab = "",
  col = color, lty = lty, lwd = 1.5)

# Extract labels from time series
labs <- as.yearmon(time(usmacro.update))[T0:Tmax + 1][c(FALSE, TRUE)]

# Create custom axis labels
axis(1, at = (1:length(T0:Tmax))[c(FALSE, TRUE)], labels = FALSE)
text(x = (1:length(T0:Tmax))[c(FALSE, TRUE)], y = par()$usr[3] - 
  0.05 * (par()$usr[4] - 
  par()$usr[3]), labels = labs, srt = 45, adj = 1, xpd = TRUE)

# Add legend
legend("topright", colnames(cumu_lpds), col = color, lty = lty, lwd = 1.5,
  bty = "n", cex = 0.8)