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

run_expensive <- identical(Sys.getenv("SHRINKR_RUN_VIGNETTES"), "true")

## ----packages-----------------------------------------------------------------
library(shrinkr)
library(brms)
library(tidybayes)
library(distributional)
library(tidyverse)
library(survival)
library(posterior)
library(patchwork)

theme_set(theme_minimal(base_size = 12))

cell_types <- c("squamous", "smallcell", "adeno", "large")

prior_specs <- list(
  very_strong = list(name = "Very Strong", scale = 0.1),
  strong = list(name = "Strong", scale = 0.25),
  moderate = list(name = "Moderate", scale = 0.5),
  weak = list(name = "Weak", scale = 1.0),
  very_weak = list(name = "Very Weak", scale = 2.0)
)

## ----load_cached_results, include=FALSE---------------------------------------
if (!run_expensive) {
  veteran_analysis <- get("veteran_analysis", envir = asNamespace("shrinkr"))
}

## ----explore_data-------------------------------------------------------------
data(veteran, package = "survival")

head(veteran)
table(veteran$celltype, veteran$trt)

veteran %>%
  group_by(celltype, trt) %>%
  summarise(
    n = n(),
    deaths = sum(status),
    median_time = median(time),
    .groups = "drop"
  )

## ----fit_brms_uninformative, eval=run_expensive-------------------------------
# brms_uninformative <- brm(
#   time | cens(1 - status) ~ trt:celltype + karno + age,
#   data = veteran,
#   family = cox(),
#   chains = 4,
#   iter = 4000,
#   warmup = 1000,
#   seed = 123
# )
# 
# brms_uninformative_summary <- capture.output(print(summary(brms_uninformative)))

## ----fit_brms_uninformative_fallback, include=FALSE---------------------------
if (!run_expensive) {
  brms_uninformative_summary <- veteran_analysis$brms_uninformative_summary
}

## ----show_brms_uninformative--------------------------------------------------
cat(brms_uninformative_summary, sep = "\n")

## ----extract_posteriors, eval=run_expensive-----------------------------------
# brms_posteriors <- brms_uninformative %>%
#   spread_draws(`b_trt:celltypesquamous`, `b_trt:celltypesmallcell`,
#                `b_trt:celltypeadeno`, `b_trt:celltypelarge`) %>%
#   select(-c(.chain, .iteration, .draw)) %>%
#   pivot_longer(everything(), names_to = "celltype", values_to = "value") %>%
#   mutate(celltype = gsub("b_trt:celltype", "", celltype)) %>%
#   group_by(celltype) %>%
#   summarise(draws = list(matrix(value, ncol = 1)), .groups = "drop") %>%
#   deframe()

## ----extract_posteriors_fallback, include=FALSE-------------------------------
if (!run_expensive) {
  brms_posteriors <- veteran_analysis$brms_posteriors
}

## ----fit_mixture_explain, eval=run_expensive----------------------------------
# mix_brms <- fit_mixture(brms_posteriors, K_max = 3, verbose = TRUE)

## ----fit_mixture_explain_fallback, include=FALSE------------------------------
if (!run_expensive) {
  mix_brms <- veteran_analysis$mix_brms
}

## ----show_mixture-------------------------------------------------------------
print(mix_brms)
plot(mix_brms, draws = brms_posteriors)

## ----define_moderate_prior----------------------------------------------------
priors_moderate <- list(
  mu = dist_normal(0, 1),
  tau = dist_truncated(dist_normal(0, 0.5), lower = 0)
)

## ----shrink_explain, eval=run_expensive---------------------------------------
# fit_twostage_brms <- shrink(
#   mixture = mix_brms,
#   hierarchical_priors = priors_moderate,
#   chains = 4,
#   iter = 4000,
#   warmup = 1000,
#   seed = 456
# )
# 
# moderate_brms_output <- capture.output(print(fit_twostage_brms))

## ----shrink_explain_fallback, include=FALSE-----------------------------------
if (!run_expensive) {
  moderate_brms_output <- veteran_analysis$sensitivity_summaries$moderate_brms$print_output
}

## ----show_twostage_brms-------------------------------------------------------
cat(moderate_brms_output, sep = "\n")

## ----fit_brms_hierarchical, eval=run_expensive--------------------------------
# brms_hierarchical <- brm(
#   time | cens(1 - status) ~ trt + (0 + trt | celltype) + karno + age,
#   data = veteran,
#   family = cox(),
#   prior = c(
#     prior(normal(0, 1), class = b, coef = "trt"),
#     prior(normal(0, 0.5), class = sd, group = "celltype", lb = 0)
#   ),
#   chains = 4,
#   iter = 4000,
#   warmup = 1000,
#   seed = 123
# )
# 
# brms_hierarchical_summary <- capture.output(print(summary(brms_hierarchical)))
# 
# brms_hier_effects <- brms_hierarchical %>%
#   spread_draws(r_celltype[celltype, term], b_trt) %>%
#   filter(term == "trt") %>%
#   mutate(theta = b_trt + r_celltype) %>%
#   group_by(celltype) %>%
#   summarise(
#     hr_mean = exp(mean(theta)),
#     hr_lower = exp(quantile(theta, 0.025)),
#     hr_upper = exp(quantile(theta, 0.975)),
#     .groups = "drop"
#   )

## ----fit_brms_hierarchical_fallback, include=FALSE----------------------------
if (!run_expensive) {
  brms_hierarchical_summary <- veteran_analysis$brms_hierarchical_summary
  brms_hier_effects <- veteran_analysis$brms_hier_effects
}

## ----show_brms_hierarchical---------------------------------------------------
cat(brms_hierarchical_summary, sep = "\n")

## ----fit_cox, eval=run_expensive----------------------------------------------
# cox_model <- coxph(
#   Surv(time, status) ~ trt:celltype + karno + age,
#   data = veteran
# )
# 
# cox_summary <- summary(cox_model)
# 
# trt_idx <- grep("^trt:celltype", names(coef(cox_model)))
# 
# trt_effects <- coef(cox_model)[trt_idx]
# trt_vcov <- vcov(cox_model)[trt_idx, trt_idx, drop = FALSE]
# 
# names(trt_effects) <- gsub("^trt:celltype", "", names(trt_effects))
# rownames(trt_vcov) <- colnames(trt_vcov) <- names(trt_effects)

## ----fit_cox_fallback, include=FALSE------------------------------------------
if (!run_expensive) {
  cox_summary <- veteran_analysis$cox_summary
  trt_effects <- veteran_analysis$trt_effects
  trt_vcov <- veteran_analysis$trt_vcov
}

## ----show_cox-----------------------------------------------------------------
print(cox_summary)

## ----show_cox_effects---------------------------------------------------------
print("Treatment effects (log HR):")
print(trt_effects)

print("\nStandard errors:")
print(sqrt(diag(trt_vcov)))

## ----shrink_freq, eval=run_expensive------------------------------------------
# fit_twostage_freq <- shrink(
#   mle = trt_effects,
#   var_matrix = trt_vcov,
#   hierarchical_priors = priors_moderate,
#   chains = 4,
#   iter = 4000,
#   warmup = 1000,
#   seed = 456
# )
# 
# moderate_freq_output <- capture.output(print(fit_twostage_freq))

## ----shrink_freq_fallback, include=FALSE--------------------------------------
if (!run_expensive) {
  moderate_freq_output <- veteran_analysis$sensitivity_summaries$moderate_freq$print_output
}

## ----show_twostage_freq-------------------------------------------------------
cat(moderate_freq_output, sep = "\n")

## ----comparison_table, eval=run_expensive-------------------------------------
# theta_brms <- summary(fit_twostage_brms)$theta %>%
#   transmute(
#     celltype = group,
#     twostage_brms = mean
#   )
# 
# theta_freq <- summary(fit_twostage_freq)$theta %>%
#   transmute(
#     celltype = group,
#     twostage_freq = mean
#   )
# 
# comparison <- brms_hier_effects %>%
#   transmute(
#     celltype,
#     full_hier_brms = log(hr_mean)
#   ) %>%
#   left_join(theta_brms, by = "celltype") %>%
#   left_join(theta_freq, by = "celltype") %>%
#   mutate(
#     diff_two_stage_vs_full = twostage_brms - full_hier_brms
#   )

## ----comparison_table_fallback, include=FALSE---------------------------------
if (!run_expensive) {
  comparison <- veteran_analysis$comparison
}

## ----comparison_table_show----------------------------------------------------
knitr::kable(
  comparison[, 1:4],
  digits = 3,
  caption = "Comparison of treatment effects (log HR scale)"
)

## ----compare_approaches, fig.width=12, fig.height=8, eval=run_expensive-------
# theta_brms_plot <- summary(fit_twostage_brms)$theta %>%
#   mutate(
#     approach = "Two-Stage (brms + shrinkr)",
#     hr_mean = exp(mean),
#     hr_lower = exp(q2.5),
#     hr_upper = exp(q97.5),
#     celltype = group
#   ) %>%
#   select(celltype, approach, hr_mean, hr_lower, hr_upper)
# 
# theta_freq_plot <- summary(fit_twostage_freq)$theta %>%
#   mutate(
#     approach = "Two-Stage (Frequentist + shrinkr)",
#     hr_mean = exp(mean),
#     hr_lower = exp(q2.5),
#     hr_upper = exp(q97.5),
#     celltype = group
#   ) %>%
#   select(celltype, approach, hr_mean, hr_lower, hr_upper)
# 
# all_approaches <- bind_rows(
#   theta_brms_plot,
#   brms_hier_effects %>% mutate(approach = "Full Hierarchical (brms)"),
#   theta_freq_plot
# ) %>%
#   mutate(
#     approach = factor(approach, levels = c(
#       "Two-Stage (brms + shrinkr)",
#       "Full Hierarchical (brms)",
#       "Two-Stage (Frequentist + shrinkr)"
#     ))
#   )

## ----compare_approaches_fallback, include=FALSE-------------------------------
if (!run_expensive) {
  theta_brms_plot <- veteran_analysis$sensitivity_summaries$moderate_brms$theta_summary %>%
    mutate(
      approach = "Two-Stage (brms + shrinkr)",
      hr_mean = exp(mean),
      hr_lower = exp(q2.5),
      hr_upper = exp(q97.5),
      celltype = group
    ) %>%
    select(celltype, approach, hr_mean, hr_lower, hr_upper)

  theta_freq_plot <- veteran_analysis$sensitivity_summaries$moderate_freq$theta_summary %>%
    mutate(
      approach = "Two-Stage (Frequentist + shrinkr)",
      hr_mean = exp(mean),
      hr_lower = exp(q2.5),
      hr_upper = exp(q97.5),
      celltype = group
    ) %>%
    select(celltype, approach, hr_mean, hr_lower, hr_upper)

  all_approaches <- bind_rows(
    theta_brms_plot,
    veteran_analysis$brms_hier_effects %>% mutate(approach = "Full Hierarchical (brms)"),
    theta_freq_plot
  ) %>%
    mutate(
      approach = factor(approach, levels = c(
        "Two-Stage (brms + shrinkr)",
        "Full Hierarchical (brms)",
        "Two-Stage (Frequentist + shrinkr)"
      ))
    )
}

## ----compare_approaches_show, fig.width=12, fig.height=8----------------------
ggplot(all_approaches, aes(x = celltype, y = hr_mean, color = approach)) +
  geom_hline(yintercept = 1, linetype = "dashed", alpha = 0.5) +
  geom_pointrange(
    aes(ymin = hr_lower, ymax = hr_upper),
    position = position_dodge(width = 0.5),
    size = 0.8
  ) +
  scale_y_log10() +
  scale_color_brewer(palette = "Set1") +
  labs(
    title = "Comparison of Three Modeling Approaches",
    subtitle = "Treatment effects by cell type (hazard ratios)",
    x = "Cell Type",
    y = "Hazard Ratio (log scale)",
    color = "Approach"
  ) +
  theme(
    legend.position = "bottom",
    panel.grid.minor = element_blank()
  )

## ----show_priors--------------------------------------------------------------
prior_summary <- tibble(
  Strength = c("Very Strong", "Strong", "Moderate", "Weak", "Very Weak"),
  Prior = c(
    "Half-Normal(0, 0.1)",
    "Half-Normal(0, 0.25)",
    "Half-Normal(0, 0.5)",
    "Half-Normal(0, 1.0)",
    "Half-Normal(0, 2.0)"
  ),
  Scale = c(0.1, 0.25, 0.5, 1.0, 2.0),
  Interpretation = c(
    "Very similar effects expected",
    "Similar effects expected",
    "Moderate heterogeneity allowed",
    "Substantial differences allowed",
    "Large differences allowed"
  )
)

knitr::kable(prior_summary)

## ----sensitivity_fits, eval=run_expensive-------------------------------------
# all_priors <- list(
#   very_strong = list(
#     mu = dist_normal(0, 1),
#     tau = dist_truncated(dist_normal(0, 0.1), lower = 0)
#   ),
#   strong = list(
#     mu = dist_normal(0, 1),
#     tau = dist_truncated(dist_normal(0, 0.25), lower = 0)
#   ),
#   moderate = list(
#     mu = dist_normal(0, 1),
#     tau = dist_truncated(dist_normal(0, 0.5), lower = 0)
#   ),
#   weak = list(
#     mu = dist_normal(0, 1),
#     tau = dist_truncated(dist_normal(0, 1.0), lower = 0)
#   ),
#   very_weak = list(
#     mu = dist_normal(0, 1),
#     tau = dist_truncated(dist_normal(0, 2.0), lower = 0)
#   )
# )
# 
# # --- brms fits ---
# sensitivity_fits_brms <- lapply(all_priors, function(prior) {
#   shrink(
#     mixture = mix_brms,
#     hierarchical_priors = prior,
#     chains = 4,
#     iter = 4000,
#     warmup = 1000
#   )
# })
# 
# # --- frequentist fits ---
# sensitivity_fits_freq <- lapply(all_priors, function(prior) {
#   shrink(
#     mle = trt_effects,
#     var_matrix = trt_vcov,
#     hierarchical_priors = prior,
#     chains = 4,
#     iter = 4000,
#     warmup = 1000
#   )
# })
# 
# # --- summaries ---
# sensitivity_summaries <- c(
#   purrr::imap(sensitivity_fits_brms, function(fit, nm) {
#     summ <- summary(fit)
#     list(
#       theta_summary = summ$theta,
#       mu_tau_summary = summ$mu_tau,
#       print_output = capture.output(print(fit))
#     )
#   }),
#   purrr::imap(sensitivity_fits_freq, function(fit, nm) {
#     summ <- summary(fit)
#     list(
#       theta_summary = summ$theta,
#       mu_tau_summary = summ$mu_tau,
#       print_output = capture.output(print(fit))
#     )
#   })
# )
# 
# # --- name them clearly ---
# names(sensitivity_summaries) <- c(
#   paste0(names(all_priors), "_brms"),
#   paste0(names(all_priors), "_freq")
# )

## ----sensitivity_fits_fallback, include=FALSE---------------------------------
if (!run_expensive) {
  sensitivity_summaries <- veteran_analysis$sensitivity_summaries
  prior_specs <- veteran_analysis$prior_specs
}

## ----prior_densities, fig.width=10, fig.height=5------------------------------
tau_seq <- seq(0, 3, length.out = 200)

prior_densities <- lapply(names(prior_specs), function(spec_name) {
  spec <- prior_specs[[spec_name]]
  tibble(
    tau = tau_seq,
    density = dnorm(tau_seq, 0, spec$scale) * 2,
    prior_strength = spec$name,
    scale = spec$scale
  )
}) %>%
  bind_rows() %>%
  mutate(
    prior_strength = factor(prior_strength, levels = c(
      "Very Strong", "Strong", "Moderate", "Weak", "Very Weak"
    ))
  )

ggplot(prior_densities, aes(x = tau, y = density, color = prior_strength)) +
  geom_line(linewidth = 1.2) +
  scale_color_brewer(palette = "RdYlBu", direction = -1) +
  labs(
    title = "Prior Densities for the Heterogeneity Parameter (tau)",
    subtitle = "Half-Normal(0, sigma) priors with increasing scale",
    x = "tau",
    y = "Density",
    color = "Prior Strength"
  ) +
  theme(legend.position = "right")

## ----tau_sensitivity----------------------------------------------------------
tau_results <- lapply(names(sensitivity_summaries), function(fit_name) {
  summary_obj <- sensitivity_summaries[[fit_name]]
  prior_name <- sub("_(brms|freq)$", "", fit_name)
  approach <- if (grepl("_brms$", fit_name)) "brms + shrinkr" else "Frequentist + shrinkr"

  summary_obj$mu_tau_summary %>%
    filter(parameter == "tau") %>%
    mutate(
      prior_strength = prior_specs[[prior_name]]$name,
      prior_scale = prior_specs[[prior_name]]$scale,
      approach = approach
    )
}) %>%
  bind_rows() %>%
  mutate(
    prior_strength = factor(
      prior_strength,
      levels = c("Very Strong", "Strong", "Moderate", "Weak", "Very Weak")
    )
  )

if (all(c("q2.5", "q97.5") %in% names(tau_results))) {
  tau_results <- tau_results %>%
    mutate(lower = `q2.5`, upper = `q97.5`)
} else if (all(c("q5", "q95") %in% names(tau_results))) {
  tau_results <- tau_results %>%
    mutate(lower = q5, upper = q95)
} else {
  stop(
    "Could not find interval columns in sensitivity_summaries$mu_tau_summary. ",
    "Available columns are: ",
    paste(names(tau_results), collapse = ", ")
  )
}

ggplot(tau_results, aes(x = prior_scale, y = mean, color = approach)) +
  geom_point(size = 3, position = position_dodge(width = 0.1)) +
  geom_errorbar(
    aes(ymin = lower, ymax = upper),
    width = 0.1,
    linewidth = 1,
    position = position_dodge(width = 0.1)
  ) +
  geom_line(aes(group = approach), position = position_dodge(width = 0.1)) +
  scale_x_log10(breaks = c(0.1, 0.25, 0.5, 1.0, 2.0)) +
  scale_color_brewer(palette = "Set2") +
  labs(
    title = "Sensitivity of the Heterogeneity Parameter (tau)",
    subtitle = "How prior scale affects the estimated between-cell-type variation",
    x = "Prior Scale (log scale)",
    y = "Posterior tau",
    color = "Stage 1 Approach"
  ) +
  theme(legend.position = "bottom")

## ----theta_sensitivity_prep---------------------------------------------------
theta_sensitivity <- lapply(names(sensitivity_summaries), function(fit_name) {
  summary_obj <- sensitivity_summaries[[fit_name]]
  prior_name <- sub("_(brms|freq)$", "", fit_name)
  approach <- if (grepl("_brms$", fit_name)) "brms + shrinkr" else "Frequentist + shrinkr"

  summary_obj$theta_summary %>%
    mutate(
      prior_strength = prior_specs[[prior_name]]$name,
      prior_scale = prior_specs[[prior_name]]$scale,
      approach = approach,
      hr_mean = exp(mean),
      hr_lower = exp(q2.5),
      hr_upper = exp(q97.5)
    )
}) %>%
  bind_rows() %>%
  mutate(
    prior_strength = factor(prior_strength, levels = c(
      "Very Strong", "Strong", "Moderate", "Weak", "Very Weak"
    ))
  )

## ----theta_sensitivity_plot, fig.width=12, fig.height=10----------------------
ggplot(theta_sensitivity, aes(x = prior_scale, y = hr_mean, color = approach)) +
  geom_hline(yintercept = 1, linetype = "dashed", alpha = 0.5) +
  geom_point(size = 2, position = position_dodge(width = 0.1)) +
  geom_errorbar(
    aes(ymin = hr_lower, ymax = hr_upper),
    width = 0.1,
    position = position_dodge(width = 0.1)
  ) +
  geom_line(aes(group = approach), position = position_dodge(width = 0.1)) +
  facet_wrap(~group, ncol = 2, scales = "free_y") +
  scale_x_log10(breaks = c(0.1, 0.25, 0.5, 1.0, 2.0)) +
  scale_y_log10() +
  scale_color_brewer(palette = "Set2") +
  labs(
    title = "Sensitivity Analysis: Cell Type-Specific Treatment Effects",
    subtitle = "How the prior scale affects hazard ratio estimates",
    x = "Prior Scale (log scale)",
    y = "Hazard Ratio (log scale)",
    color = "Stage 1 Approach"
  ) +
  theme(
    legend.position = "bottom",
    panel.grid.minor = element_blank()
  )

## ----session------------------------------------------------------------------
sessionInfo()