## ----setup, include = FALSE---------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 8,
  fig.height = 6,
  warning = FALSE,
  message = FALSE
)

# Load pre-computed Stage 1 results (Stan samples, mixture, prior predictive)
# The expensive Stan stage 1 fit is pre-computed; shrink() runs live below
r <- get("getting_started_results", envir = asNamespace("shrinkr"))

## ----packages, message=FALSE--------------------------------------------------
library(shrinkr)
library(distributional)
library(posterior)
library(tidyverse)

## ----generate_data, eval=FALSE------------------------------------------------
# library(rstan)
# set.seed(1104)
# 
# true_mu <- 0.5
# true_tau <- 0.3
# true_effects <- c(0.45, 0.72, 0.38, 0.55, 0.61)
# 
# regions <- c("North", "South", "East", "West", "Central")
# n_per_region <- c(100, 80, 120, 90, 70)
# 
# trial_data <- lapply(seq_along(regions), function(i) {
#   n <- n_per_region[i]
#   data.frame(
#     region = regions[i],
#     treatment = rep(c(0, 1), each = n/2),
#     outcome = c(
#       rnorm(n/2, mean = 0, sd = 1),
#       rnorm(n/2, mean = true_effects[i], sd = 1)
#     )
#   )
# }) %>% bind_rows()

## ----show_data_hidden, include=FALSE------------------------------------------
trial_data <- r$trial_data

## ----show_data, eval=FALSE----------------------------------------------------
# head(trial_data)
# table(trial_data$region, trial_data$treatment)

## ----show_data_run, echo=FALSE------------------------------------------------
head(r$trial_data)
table(r$trial_data$region, r$trial_data$treatment)

## ----stan_model, eval=FALSE---------------------------------------------------
# stan_code <- "
# data {
#   int<lower=0> N;
#   int<lower=1> G;
#   vector[N] y;
#   vector[N] treatment;
#   array[N] int<lower=1,upper=G> region;
# }
# parameters {
#   vector[G] beta_region;
#   real<lower=0> sigma;
# }
# model {
#   // IMPORTANT: Flat prior on beta_region - critical for the two-stage approach!
#   sigma ~ normal(0, 2);
#   for (n in 1:N) {
#     y[n] ~ normal(treatment[n] * beta_region[region[n]], sigma);
#   }
# }
# "

## ----fit_stan, eval=FALSE-----------------------------------------------------
# regions <- c("North", "South", "East", "West", "Central")
# region_indices <- as.integer(factor(trial_data$region, levels = regions))
# 
# fit_stan <- stan(
#   model_code = stan_code,
#   data = list(
#     N = nrow(trial_data), G = length(regions),
#     y = trial_data$outcome, treatment = trial_data$treatment,
#     region = region_indices
#   ),
#   chains = 4, iter = 2000, warmup = 1000, refresh = 0, seed = 123
# )
# 
# beta_draws <- rstan::extract(fit_stan, pars = "beta_region")$beta_region
# samples_list <- lapply(seq_along(regions), function(i) beta_draws[, i])
# names(samples_list) <- regions
# samples <- lapply(samples_list, function(x) matrix(x, ncol = 1))

## ----examine_stage1_hidden, include=FALSE-------------------------------------
samples <- r$samples
regions <- names(samples)

## ----examine_stage1, eval=FALSE-----------------------------------------------
# stage1_summary <- data.frame(
#   region = regions,
#   mean   = sapply(samples, mean),
#   sd     = sapply(samples, sd),
#   lower  = sapply(samples, function(x) quantile(x, 0.025)),
#   upper  = sapply(samples, function(x) quantile(x, 0.975))
# )
# 
# print(stage1_summary)

## ----examine_stage1_run, echo=FALSE-------------------------------------------
samples <- r$samples
regions <- names(samples)

stage1_summary <- data.frame(
  region = regions,
  mean   = sapply(samples, mean),
  sd     = sapply(samples, sd),
  lower  = sapply(samples, function(x) quantile(x, 0.025)),
  upper  = sapply(samples, function(x) quantile(x, 0.975))
)

print(stage1_summary)

## ----plot_stage1, fig.width=10, fig.height=5----------------------------------
ggplot(stage1_summary, aes(x = region, y = mean)) +
  geom_hline(yintercept = 0, linetype = "dashed", color = "gray50") +
  geom_pointrange(aes(ymin = lower, ymax = upper),
                  size = 0.8, color = "steelblue") +
  labs(
    title    = "Stage 1: Independent Regional Estimates",
    subtitle = "Each region analyzed separately with flat priors",
    x = "Region", y = "Treatment Effect",
    caption = "Points show posterior means; bars show 95% credible intervals"
  ) +
  theme_minimal(base_size = 12)

## ----fit_mixture, eval=FALSE--------------------------------------------------
# mix <- fit_mixture(samples = samples, K_max = 3, verbose = TRUE)

## ----show_mixture_hidden, include=FALSE---------------------------------------
mix <- r$mix

## ----show_mixture, eval=FALSE-------------------------------------------------
# print(mix)

## ----show_mixture_run, echo=FALSE---------------------------------------------
print(r$mix)

## ----check_mixture, fig.width=12, fig.height=8--------------------------------
# Blue line should overlay the histogram well
plot(mix, draws = samples, type = "density")

# Points should fall near the diagonal
plot(mix, draws = samples, type = "qq")

## ----specify_priors-----------------------------------------------------------
hierarchical_priors <- list(
  mu  = dist_normal(0, 1),
  tau = dist_truncated(dist_student_t(3, 0, 0.5), lower = 0)
)

## ----prior_predictive, eval=FALSE---------------------------------------------
# prior_pred <- sample_prior_predictive(
#   hierarchical_priors = hierarchical_priors,
#   n_groups = 5,
#   n_draws  = 1000
# )

## ----show_prior_pred_hidden, include=FALSE------------------------------------
prior_pred <- r$prior_pred

## ----show_prior_pred, eval=FALSE----------------------------------------------
# cat("Prior on tau (between-region SD):\n")
# cat("  Median:", round(median(prior_pred$tau), 2), "\n")
# cat("  95% interval:", round(quantile(prior_pred$tau, c(0.025, 0.975)), 2), "\n\n")
# 
# cat("Implied variation in regional effects:\n")
# cat("  Typical range:", round(median(prior_pred$implied_range), 2), "\n")
# cat("  95% interval:", round(quantile(prior_pred$implied_range, c(0.025, 0.975)), 2), "\n")

## ----show_prior_pred_run, echo=FALSE------------------------------------------
cat("Prior on tau (between-region SD):\n")
cat("  Median:", round(median(r$prior_pred$tau), 2), "\n")
cat("  95% interval:", round(quantile(r$prior_pred$tau, c(0.025, 0.975)), 2), "\n\n")

cat("Implied variation in regional effects:\n")
cat("  Typical range:", round(median(r$prior_pred$implied_range), 2), "\n")
cat("  95% interval:", round(quantile(r$prior_pred$implied_range, c(0.025, 0.975)), 2), "\n")

## ----plot_prior, fig.width=10, fig.height=8-----------------------------------
plot(prior_pred)

## ----pairwise_prior, fig.width=8, fig.height=5--------------------------------
pw <- prior_pairwise_differences(prior_pred)
print(pw)

## ----plot_pairwise, fig.width=8, fig.height=5---------------------------------
# Pooled histogram of |theta_i - theta_j| across all pairs
plot(pw)

## ----fit_shrink---------------------------------------------------------------
fit <- shrink(
  mixture             = mix,
  hierarchical_priors = hierarchical_priors,
  chains  = 4,
  iter    = 2000,
  warmup  = 1000,
  seed    = 456,
  refresh = 0
)

print(fit)

## ----extract_hyperparams------------------------------------------------------
mu_tau <- extract_mu_tau(fit)

cat("Overall treatment effect (mu):\n")
cat("  Mean:", round(mean(mu_tau$mu), 3), "\n")
cat("  95% CI:", round(quantile(mu_tau$mu, c(0.025, 0.975)), 3), "\n\n")

cat("Between-region heterogeneity (tau):\n")
cat("  Mean:", round(mean(mu_tau$tau), 3), "\n")
cat("  95% CI:", round(quantile(mu_tau$tau, c(0.025, 0.975)), 3), "\n")

## ----summarize_theta----------------------------------------------------------
theta_summary <- summarize_theta(fit)
print(theta_summary)

## ----plot_shrink, fig.width=10, fig.height=6----------------------------------
plot(fit)

## ----plot_diagnostics, fig.width=12, fig.height=8-----------------------------
plot(fit, type = "diagnostics")

## ----mle_approach-------------------------------------------------------------
mle_estimates <- sapply(samples, mean)
mle_variances <- sapply(samples, var)

fit_mle <- shrink(
  mle                 = mle_estimates,
  var_matrix          = mle_variances,
  hierarchical_priors = hierarchical_priors,
  chains  = 4,
  iter    = 2000,
  warmup  = 1000,
  seed    = 456,
  refresh = 0
)

## ----show_mle-----------------------------------------------------------------
mu_tau_mle <- extract_mu_tau(fit_mle)

cat("Mixture approach:\n")
cat("  mu =",  round(mean(mu_tau$mu), 3), "\n")
cat("  tau =", round(mean(mu_tau$tau), 3), "\n\n")

cat("MLE approach:\n")
cat("  mu =",  round(mean(mu_tau_mle$mu), 3), "\n")
cat("  tau =", round(mean(mu_tau_mle$tau), 3), "\n")

## ----workflow_summary, eval=FALSE---------------------------------------------
# # 1. Write Stan model with FLAT PRIORS on parameters of interest
# stan_code <- "
# data {
#   int<lower=0> N;
#   int<lower=1> G;
#   vector[N] y;
#   vector[N] treatment;
#   array[N] int<lower=1,upper=G> group;
# }
# parameters {
#   vector[G] theta;  // Group-specific effects - NO PRIOR SPECIFIED
#   real<lower=0> sigma;
# }
# model {
#   sigma ~ normal(0, 2);
#   for (n in 1:N) {
#     y[n] ~ normal(treatment[n] * theta[group[n]], sigma);
#   }
# }
# "
# 
# # 2. Fit model once to get all group effects, extract posteriors
# group_indices <- as.integer(factor(data$group, levels = groups))
# fit_stan <- stan(
#   model_code = stan_code,
#   data = list(N = nrow(data), G = length(groups),
#               y = data$y, treatment = data$treatment,
#               group = group_indices),
#   chains = 4, iter = 2000, warmup = 1000, refresh = 0
# )
# 
# theta_draws <- extract(fit_stan)$theta
# samples <- lapply(seq_along(groups), function(i) matrix(theta_draws[, i], ncol = 1))
# names(samples) <- groups
# 
# # 3. Fit mixture approximation and check quality
# mix <- fit_mixture(samples, K_max = 3)
# plot(mix, draws = samples)
# 
# # 4. Specify and check hierarchical priors
# priors <- list(
#   mu  = dist_normal(0, 5),
#   tau = dist_truncated(dist_student_t(3, 0, 1), lower = 0)
# )
# plot(sample_prior_predictive(priors, n_groups = length(groups)))
# 
# # 5. Fit, extract, and visualize
# fit <- shrink(mixture = mix, hierarchical_priors = priors)
# plot(fit)
# summarize_theta(fit)
# extract_mu_tau(fit)

## ----sessioninfo--------------------------------------------------------------
sessionInfo()