# This script reproduces all BSTS model fitting, diagnostic statistics (Table S1),
# diagnostic plots (Figure S1), and projection results reported in the main text.
# Input:  forecast_ASR.csv
# Output: BSTS_Diagnostics.pdf  → Supplementary Figure S1
#         BSTS_Vietnam_Results.csv → Source data for main text figures
#
# =============================================================================
# Bayesian Structural Time Series (BSTS) Pipeline
# Vietnam GBD 2023 – Mental Disorders
# =============================================================================

rm(list = ls())

# -----------------------------------------------------------------------------
# 0. LIBRARIES
# -----------------------------------------------------------------------------
required_packages <- c("bsts", "tidyverse", "coda")

installed <- rownames(installed.packages())

for (pkg in required_packages) {
  if (!pkg %in% installed) install.packages(pkg)
  library(pkg, character.only = TRUE)
}

# -----------------------------------------------------------------------------
# 1. GLOBAL CONFIGURATION
# -----------------------------------------------------------------------------
CONFIG <- list(
  country = "Viet Nam",
  cause = "Mental disorders",
  measures = c("Incidence", "Prevalence", "DALYs"),
  
  covid_years = 2019:2023,
  
  covid_models = list(
    Incidence = "pulse",
    Prevalence = "step",
    DALYs = "step"
  ),
  
  forecast_years = 2024:2040,
  n_iter = 20000,
  burn_in = 2000,
  seed = 123
)

# -----------------------------------------------------------------------------
# 2. LOAD DATA
# -----------------------------------------------------------------------------
load_data <- function(file_path) {
  read_csv(file_path, show_col_types = FALSE)
}

# -----------------------------------------------------------------------------
# 3. FILTER TIME SERIES
# -----------------------------------------------------------------------------
filter_time_series <- function(df, measure, config) {
  df %>%
    filter(
      location_name == config$country,
      str_detect(measure_name, measure),
      cause_name == config$cause
    ) %>%
    arrange(year)
}

# -----------------------------------------------------------------------------
# 4. BUILD COVID REGRESSOR
# -----------------------------------------------------------------------------
build_covid_regressor <- function(years, measure, config) {
  model_type <- config$covid_models[[measure]]
  
  if (model_type == "pulse") {
    if_else(years %in% config$covid_years, 1, 0)
  } else {
    if_else(years >= min(config$covid_years), 1, 0)
  }
}

# -----------------------------------------------------------------------------
# 5. FIT BSTS MODEL
# -----------------------------------------------------------------------------
fit_bsts_model <- function(y, xreg, config) {
  ss <- AddSemilocalLinearTrend(list(), y)
  
  bsts(
    y ~ xreg,
    state.specification = ss,
    niter = config$n_iter,
    ping = 0,
    seed = config$seed
  )
}

# -----------------------------------------------------------------------------
# 6. COMPUTE DIAGNOSTICS
# -----------------------------------------------------------------------------
compute_diagnostics <- function(model, y, config) {
  burn <- config$burn_in
  resids_matrix <- residuals(model)
  
  post_resids <- resids_matrix[(burn + 1):nrow(resids_matrix), , drop = FALSE]
  avg_resids <- colMeans(post_resids)
  
  fitted_vals <- y - avg_resids
  
  coef_cols <- colnames(model$coefficients)
  target_coef <- coef_cols[coef_cols != "(Intercept)"][1]
  
  coeff_samples <- model$coefficients[
    (burn + 1):nrow(model$coefficients),
    target_coef
  ]
  
  ess_val <- effectiveSize(as.mcmc(coeff_samples))
  
  list(
    MAPE = mean(abs(avg_resids) / y) * 100,
    Bayesian_R2 = var(fitted_vals) / (var(fitted_vals) + var(avg_resids)),
    ESS = ess_val,
    residuals = avg_resids,
    coeff_samples = coeff_samples
  )
}

# -----------------------------------------------------------------------------
# 7. DIAGNOSTIC PANEL
# -----------------------------------------------------------------------------
plot_diagnostics <- function(diagnostics, model, measure) {
  
  old_par <- par(no.readonly = TRUE)
  on.exit(par(old_par))
  
  par(
    mfrow = c(1, 3),
    mar = c(5, 5, 4, 1),
    oma = c(0, 0, 3, 0),
    cex.axis = 1.5,   
    cex.lab = 1.5,    
    cex.main = 1.5    
  )
  
  # ---------------------------------------------------
  # 1. Trace Plot
  # ---------------------------------------------------
  plot(
    diagnostics$coeff_samples,
    type = "l",
    col = "darkblue",
    xlab = "Iteration",
    ylab = "Coefficient",
    main = "Trace Plot"
  )
  
  abline(
    h = mean(diagnostics$coeff_samples),
    col = "red",
    lwd = 2,
    lty = 2
  )
  
  # ---------------------------------------------------
  # 2. BSTS Residual Plot
  # ---------------------------------------------------
  plot(model, "residuals", main = "")
  title(
    main = "Residual Distribution",
    cex.main = 1.4
  )
  
  # ---------------------------------------------------
  # 3. Normal Q-Q Plot
  # ---------------------------------------------------
  qqnorm(
    diagnostics$residuals,
    main = "Normal Q-Q Plot",
    cex = 1.8,
    pch = 16
  )
  
  qqline(
    diagnostics$residuals,
    col = "red",
    lwd = 2.5
  )
  
  mtext(
    paste("BSTS Diagnostic Panel -", measure),
    outer = TRUE,
    cex = 1.5,
    font = 2
  )
}
# -----------------------------------------------------------------------------
# 8. FORECASTING
# -----------------------------------------------------------------------------
forecast_bsts <- function(model, measure, config) {
  future_years <- config$forecast_years
  n_ahead <- length(future_years)
  
  model_type <- config$covid_models[[measure]]
  
  xreg_future <- if (model_type == "pulse") rep(0, n_ahead) else rep(1, n_ahead)
  
  pred <- predict(
    model,
    horizon = n_ahead,
    newdata = data.frame(xreg = xreg_future),
    burn = config$burn_in
  )
  
  list(
    years = future_years,
    mean = apply(pred$distribution, 2, mean),
    lower = apply(pred$distribution, 2, quantile, probs = 0.025),
    upper = apply(pred$distribution, 2, quantile, probs = 0.975)
  )
}

# -----------------------------------------------------------------------------
# 9. POLICY SCENARIOS
# -----------------------------------------------------------------------------
apply_policy_scenarios <- function(forecast) {
  n <- length(forecast$years)
  
  multipliers <- list(
    baseline = rep(1, n),
    policy_intervention = 1 - 0.17 * ((0:(n - 1)) / (n - 1))^1.3
  )
  
  bind_rows(lapply(names(multipliers), function(sc_name) {
    mult <- multipliers[[sc_name]]
    
    data.frame(
      Year = forecast$years,
      Value = forecast$mean * mult,
      Lower_95 = forecast$lower * mult,
      Upper_95 = forecast$upper * mult,
      Scenario = sc_name,
      Type = "Forecast"
    )
  }))
}

# -----------------------------------------------------------------------------
# 10. MAIN PIPELINE
# -----------------------------------------------------------------------------
run_bsts_pipeline <- function(df, config, diagnostics_pdf = TRUE) {
  
  if (diagnostics_pdf) {
    pdf("BSTS_Diagnostics.pdf", width = 14, height = 5)
  }
  
  results_list <- list()
  
  for (measure in config$measures) {
    
    cat("\n========================================\n")
    cat("Processing:", measure, "\n")
    
    ts_data <- filter_time_series(df, measure, config)
    
    if (nrow(ts_data) < 10) {
      warning(paste("Skipping", measure, "- insufficient data"))
      next
    }
    
    y <- ts_data$val
    xreg <- build_covid_regressor(ts_data$year, measure, config)
    
    model <- fit_bsts_model(y, xreg, config)
    diagnostics <- compute_diagnostics(model, y, config)
    
    cat(sprintf(
      "MAPE: %.2f%% | Bayesian R2: %.4f | ESS: %.1f\n",
      diagnostics$MAPE,
      diagnostics$Bayesian_R2,
      diagnostics$ESS
    ))
    
    plot_diagnostics(diagnostics, model, measure)
    
    forecast <- forecast_bsts(model, measure, config)
    scenario_df <- apply_policy_scenarios(forecast)
    
    hist_df <- data.frame(
      Year = ts_data$year,
      Value = y,
      Lower_95 = y,
      Upper_95 = y,
      Scenario = "historical",
      Type = "Historical"
    )
    
    final_df <- bind_rows(hist_df, scenario_df) %>%
      mutate(Measure = measure)
    
    results_list[[measure]] <- final_df
  }
  
  if (diagnostics_pdf) dev.off()
  
  bind_rows(results_list)
}

# -----------------------------------------------------------------------------
# 11. RUN PIPELINE
# -----------------------------------------------------------------------------
df <- load_data("forecast_ASR.csv")

final_results <- run_bsts_pipeline(df, CONFIG)

# -----------------------------------------------------------------------------
# 12. EXPORT RESULTS
# -----------------------------------------------------------------------------
write_csv(final_results, "BSTS_Vietnam_Results.csv")