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

## ----packages-----------------------------------------------------------------
library(shrinkr)
library(distributional)
library(posterior)
library(dplyr)
library(ggplot2)
library(tidyr)

## ----network_setup------------------------------------------------------------
hospitals <- data.frame(
  site_id = 1:6,
  name = c("Metro General", "County Regional", "University Medical", 
           "Community Hospital", "Veterans Affairs", "Children's Specialty"),
  location = c("Urban", "Suburban", "Academic", "Rural", "Urban", "Urban"),
  n_patients = c(1500, 800, 2200, 350, 1100, 900),
  baseline_risk = c(0.15, 0.12, 0.18, 0.10, 0.20, 0.14)
)

print(hospitals)

## ----site_models--------------------------------------------------------------
set.seed(1104)

# True network-level parameters (unknown in practice)
true_mu_age <- 0.05    # log-OR per year
true_tau_age <- 0.015  # between-site heterogeneity
true_site_effects <- rnorm(6, true_mu_age, true_tau_age)

# Simulate Stage 1: Each site fits their model independently
# In reality: glm(), stan_glm(), or other Bayesian logistic regression

site_posteriors <- list()
site_sample_sizes <- hospitals$n_patients

for(i in 1:6) {
  # Posterior for age coefficient beta_1
  # SE inversely proportional to sqrt(sample size)
  se_i <- 0.02 * sqrt(800 / site_sample_sizes[i])
  
  site_posteriors[[hospitals$name[i]]] <- matrix(
    rnorm(4000, true_site_effects[i], se_i),
    ncol = 1
  )
}

# Each site computes summaries
site_summaries <- data.frame(
  site = hospitals$name,
  n_patients = site_sample_sizes,
  beta_age_mean = sapply(site_posteriors, mean),
  beta_age_se = sapply(site_posteriors, sd)
) %>%
  mutate(
    ci_lower = beta_age_mean - 1.96 * beta_age_se,
    ci_upper = beta_age_mean + 1.96 * beta_age_se
  )

print(site_summaries)

## ----share_posteriors---------------------------------------------------------
# Sites share posterior samples (4000 draws each)
# This is more informative but requires more data transfer

cat("Data shared per site:\n")
cat("  Posterior samples: 4000 draws\n")
cat("  Total data transfer:", 6 * 4000 * 8, "bytes (", 
    round(6 * 4000 * 8 / 1024, 1), "KB)\n")

## ----share_summaries----------------------------------------------------------
# Sites share only mean and SE
# Minimal data transfer, maximum privacy
# BUT: Only valid if posteriors are approximately normal!

summary_only <- site_summaries %>%
  select(site, beta_age_mean, beta_age_se)

cat("Data shared per site:\n")
cat("  Point estimate: 1 number\n")
cat("  Standard error: 1 number\n")
cat("  Total data transfer:", 6 * 2 * 8, "bytes (", 
    round(6 * 2 * 8 / 1024, 3), "KB)\n")

print(summary_only)

## ----path_a_mixture-----------------------------------------------------------
# Fit mixture approximation
mix <- fit_mixture(
  samples = site_posteriors,
  K_max = 2,  # Age effects should be fairly normal
  verbose = TRUE
)

# Check quality
plot(mix, draws = site_posteriors, type = "density")

## ----network_priors-----------------------------------------------------------
# Based on clinical knowledge:
# - Age effect should be positive but moderate
# - Some heterogeneity expected across hospital types

hierarchical_priors <- list(
  mu = dist_normal(0.05, 0.025),  # Centered on 5% increase per year
  tau = dist_truncated(dist_student_t(3, 0, 0.01), lower = 0)  # Modest heterogeneity
)

# Visualize prior implications
prior_pred <- sample_prior_predictive(
  hierarchical_priors = hierarchical_priors,
  n_groups = 6,
  n_draws = 1000
)

plot(prior_pred, type = "both")

## ----fit_path_a, results='hide'-----------------------------------------------
fit_full_post <- shrink(
  mixture = mix,
  hierarchical_priors = hierarchical_priors,
  chains = 4,
  iter = 2000,
  warmup = 1000,
  seed = 123
)

## ----path_a_results-----------------------------------------------------------
print(fit_full_post)

# Network-level estimates
mu_tau_full <- extract_mu_tau(fit_full_post)
cat("\nNetwork-level age effect (mu):\n")
cat("  Posterior mean:", round(mean(mu_tau_full$mu), 4), "\n")
cat("  95% CI: [", round(quantile(mu_tau_full$mu, 0.025), 4), ",",
    round(quantile(mu_tau_full$mu, 0.975), 4), "]\n")

cat("\nBetween-site heterogeneity (tau):\n")
cat("  Posterior mean:", round(mean(mu_tau_full$tau), 4), "\n")
cat("  95% CI: [", round(quantile(mu_tau_full$tau, 0.025), 4), ",",
    round(quantile(mu_tau_full$tau, 0.975), 4), "]\n")

## ----check_normality, fig.width=10, fig.height=8------------------------------
# Visual checks for approximate normality
oldpar <- par(no.readonly = TRUE)
par(mfrow = c(3, 2))
for(i in 1:6) {
  site_name <- names(site_posteriors)[i]
  samples_i <- as.vector(site_posteriors[[i]])
  
  # QQ plot against normal
  qqnorm(samples_i, main = paste("QQ Plot:", site_name))
  qqline(samples_i, col = "red", lwd = 2)
}
par(oldpar)

## ----normality_tests----------------------------------------------------------
# Quantitative checks
normality_checks <- data.frame(
  site = names(site_posteriors),
  skewness = sapply(site_posteriors, function(x) {
    m3 <- mean((x - mean(x))^3)
    s3 <- sd(x)^3
    m3 / s3
  }),
  kurtosis = sapply(site_posteriors, function(x) {
    m4 <- mean((x - mean(x))^4)
    s4 <- sd(x)^4
    m4 / s4 - 3  # Excess kurtosis
  })
)

print(normality_checks)

cat("\nNormality assessment:\n")
cat("  Skewness close to 0? (|skew| < 0.5 is good)\n")
cat("  Kurtosis close to 0? (|kurt| < 1.0 is good)\n")
cat("  All sites pass:", 
    all(abs(normality_checks$skewness) < 0.5 & abs(normality_checks$kurtosis) < 1.0),
    "\n")

## ----clt_failure_example------------------------------------------------------
# Simulated: posterior for a variance parameter (boundary at 0)
set.seed(999)
variance_posterior <- matrix(rchisq(4000, df = 5) / 5, ncol = 1)

# Compute summary statistics
var_mean <- mean(variance_posterior)
var_se <- sd(variance_posterior)

# Check normality
oldpar <- par(no.readonly = TRUE)
par(mfrow = c(1, 2))
hist(variance_posterior, breaks = 30, main = "Variance Posterior",
     xlab = "sigma^2", col = "lightblue")
qqnorm(variance_posterior, main = "QQ Plot")
qqline(variance_posterior, col = "red", lwd = 2)
par(oldpar)

# Skewness
skew <- mean((variance_posterior - var_mean)^3) / var_se^3
cat("Skewness:", round(skew, 2), "(should be about 0 for normal)\n")
cat("This posterior is right-skewed - CLT approximation would be poor!\n")

## ----path_b_mle---------------------------------------------------------------
# Extract means and variances
mle_estimates <- site_summaries$beta_age_mean
names(mle_estimates) <- site_summaries$site

mle_variances <- site_summaries$beta_age_se^2
names(mle_variances) <- site_summaries$site

# Fit using MLE path (CLT approximation)
fit_summaries <- shrink(
  mle = mle_estimates,
  var_matrix = mle_variances,
  hierarchical_priors = hierarchical_priors,
  chains = 4,
  iter = 2000,
  warmup = 1000,
  seed = 123,
  verbose = FALSE,
  refresh = 0
)

print(fit_summaries)

## ----compare_paths------------------------------------------------------------
mu_tau_summaries <- extract_mu_tau(fit_summaries)

comparison <- data.frame(
  parameter = c("mu", "tau"),
  full_posteriors = c(
    mean(mu_tau_full$mu),
    mean(mu_tau_full$tau)
  ),
  summaries_only = c(
    mean(mu_tau_summaries$mu),
    mean(mu_tau_summaries$tau)
  )
) %>%
  mutate(difference = abs(full_posteriors - summaries_only))

print(comparison)

cat("\nMaximum difference:", round(max(comparison$difference), 5), "\n")

## ----viz_shrinkage, fig.width=10, fig.height=6--------------------------------
# Using Path A results (nearly identical for Path B)
plot(fit_full_post)

## ----uncertainty_reduction----------------------------------------------------
# Get Stage 2 estimates
theta_post <- summarize_theta(fit_full_post)

# Compare Stage 1 vs Stage 2 uncertainty
uncertainty_comparison <- data.frame(
  site = site_summaries$site,
  n_patients = site_summaries$n_patients,
  stage1_se = site_summaries$beta_age_se,
  stage2_se = theta_post$sd
) %>%
  mutate(
    reduction_pct = 100 * (stage1_se - stage2_se) / stage1_se,
    stage1_ci_width = 2 * 1.96 * stage1_se,
    stage2_ci_width = 2 * 1.96 * stage2_se,
    ci_width_reduction = 100 * (stage1_ci_width - stage2_ci_width) / stage1_ci_width
  )

print(uncertainty_comparison)

## ----plot_uncertainty, fig.width=10, fig.height=6-----------------------------
# Prepare data for plotting
uncertainty_long <- uncertainty_comparison %>%
  select(site, n_patients, stage1_se, stage2_se) %>%
  pivot_longer(
    cols = c(stage1_se, stage2_se),
    names_to = "stage",
    values_to = "standard_error"
  ) %>%
  mutate(
    stage = factor(stage, 
                   levels = c("stage1_se", "stage2_se"),
                   labels = c("Stage 1 (Independent)", "Stage 2 (Shrunken)"))
  )

ggplot(uncertainty_long, aes(x = reorder(site, -n_patients), y = standard_error, 
                              fill = stage)) +
  geom_col(position = "dodge") +
  geom_text(aes(label = sprintf("%.4f", standard_error)),
            position = position_dodge(width = 0.9),
            vjust = -0.5, size = 3) +
  scale_fill_manual(values = c("Stage 1 (Independent)" = "steelblue",
                               "Stage 2 (Shrunken)" = "coral")) +
  labs(
    title = "Uncertainty Reduction Through Federated Learning",
    subtitle = "Sites ordered by sample size (largest to smallest)",
    x = "Hospital Site",
    y = "Standard Error",
    fill = NULL
  ) +
  theme_minimal() +
  theme(
    axis.text.x = element_text(angle = 45, hjust = 1),
    legend.position = "bottom"
  )

## ----stage1_predictions-------------------------------------------------------
# Example: 70-year-old patient
age <- 70
baseline_age <- 60  # Reference age

# Stage 1 predictions (independent)
stage1_log_or <- site_summaries$beta_age_mean * (age - baseline_age)
stage1_or <- exp(stage1_log_or)

stage1_preds <- data.frame(
  site = site_summaries$site,
  log_or = stage1_log_or,
  odds_ratio = stage1_or,
  stage = "Independent"
)

## ----stage2_predictions-------------------------------------------------------
# Stage 2 predictions (network-calibrated)
stage2_log_or <- theta_post$mean * (age - baseline_age)
stage2_or <- exp(stage2_log_or)

stage2_preds <- data.frame(
  site = theta_post$group,
  log_or = stage2_log_or,
  odds_ratio = stage2_or,
  stage = "Network-Calibrated"
)

# Combine
all_preds <- rbind(stage1_preds, stage2_preds)

## ----plot_predictions, fig.width=10, fig.height=6-----------------------------
ggplot(all_preds, aes(x = site, y = odds_ratio, fill = stage)) +
  geom_col(position = "dodge") +
  geom_hline(yintercept = 1, linetype = "dashed", color = "gray30") +
  geom_text(aes(label = sprintf("%.2f", odds_ratio)),
            position = position_dodge(width = 0.9),
            vjust = -0.5, size = 3) +
  scale_fill_manual(values = c("Independent" = "steelblue",
                               "Network-Calibrated" = "coral")) +
  labs(
    title = "Predicted Odds Ratio for 70 vs 60 Year-Old Patient",
    subtitle = "Network calibration stabilizes predictions across sites",
    x = "Hospital Site",
    y = "Odds Ratio",
    fill = NULL
  ) +
  theme_minimal() +
  theme(
    axis.text.x = element_text(angle = 45, hjust = 1),
    legend.position = "bottom"
  )

## ----privacy_summary----------------------------------------------------------
privacy_comparison <- data.frame(
  approach = c("Centralized Data", "Path A: Full Posteriors", "Path B: Summaries"),
  patient_data_shared = c("Yes - All records", "No", "No"),
  data_per_site = c("about 50-200 MB", "about 200 KB", "about 16 bytes"),
  privacy_risk = c("High", "Low", "Minimal"),
  validity = c("N/A", "Always valid", "Only if CLT holds"),
  when_to_use = c("Not for federated", "Default choice", "When posteriors normal")
)

knitr::kable(privacy_comparison, align = "lccccl")

## ----heterogeneous_models, eval=FALSE-----------------------------------------
# # Site 1: Linear model
# # Site 2: GLM with splines
# # Site 3: Bayesian hierarchical model
# 
# # As long as they all estimate the same parameter, shrinkr can combine them.
# 
# samples_heterogeneous <- list(
#   site1 = samples_from_lm,
#   site2 = samples_from_glm,
#   site3 = samples_from_bayes
# )
# 
# # Proceed with shrinkr as usual
# mix <- fit_mixture(samples_heterogeneous, K_max = 3)
# fit <- shrink(mix, hierarchical_priors = priors)

## ----meta_analysis, eval=FALSE------------------------------------------------
# # Extracted from publications
# published_estimates <- c(
#   "Smith et al. (2020)" = 0.45,
#   "Jones et al. (2021)" = 0.52,
#   "Garcia et al. (2022)" = 0.38,
#   "Williams et al. (2023)" = 0.48
# )
# 
# published_ses <- c(0.12, 0.15, 0.10, 0.13)
# 
# # Apply shrinkr for Bayesian meta-analysis
# # NOTE: This assumes published estimates are approximately normal.
# 
# fit_meta <- shrink(
#   mle = published_estimates,
#   var_matrix = published_ses^2,
#   hierarchical_priors = priors
# )

## ----iterative_updates, eval=FALSE--------------------------------------------
# # Initial network
# fit_initial <- shrink(samples_initial, priors)
# 
# # New site joins
# samples_updated <- c(samples_initial, list(new_site = new_samples))
# fit_updated <- shrink(samples_updated, priors)
# 
# # Compare network estimates before/after
# mu_before <- mean(extract_mu_tau(fit_initial)$mu)
# mu_after <- mean(extract_mu_tau(fit_updated)$mu)

## ----quality_control----------------------------------------------------------
# Example: Flag suspicious estimates
qc_results <- site_summaries %>%
  mutate(
    z_score = (beta_age_mean - median(beta_age_mean)) / mad(beta_age_mean),
    flag = ifelse(abs(z_score) > 3, "Review", "OK")
  )

cat("Quality control flags:\n")
print(qc_results %>% select(site, beta_age_mean, z_score, flag))

## ----sensitivity, results='hide'----------------------------------------------
# Alternative prior: More heterogeneity
priors_alt <- list(
  mu = dist_normal(0.05, 0.025),
  tau = dist_truncated(dist_student_t(3, 0, 0.02), lower = 0)
)

fit_alt <- shrink(
  mixture = mix,
  hierarchical_priors = priors_alt,
  chains = 2,
  iter = 1000,
  warmup = 500,
  seed = 456
)

## ----sensitivity_compare------------------------------------------------------
# Compare key results
mu_base <- mean(extract_mu_tau(fit_full_post)$mu)
mu_alt <- mean(extract_mu_tau(fit_alt)$mu)

tau_base <- mean(extract_mu_tau(fit_full_post)$tau)
tau_alt <- mean(extract_mu_tau(fit_alt)$tau)

cat("Sensitivity to prior on tau:\n")
cat("  mu: Base =", round(mu_base, 4), ", Alternative =", round(mu_alt, 4), "\n")
cat("  tau: Base =", round(tau_base, 4), ", Alternative =", round(tau_alt, 4), "\n")

## ----reporting----------------------------------------------------------------
# Create site-specific report
site_report <- data.frame(
  site = theta_post$group,
  original_estimate = site_summaries$beta_age_mean,
  original_se = site_summaries$beta_age_se,
  calibrated_estimate = theta_post$mean,
  calibrated_se = theta_post$sd,
  uncertainty_reduction = uncertainty_comparison$reduction_pct
) %>%
  mutate(across(where(is.numeric), ~round(.x, 4)))

cat("\nFederated Learning Results Report\n")
cat("==================================\n\n")
cat("Network-level estimate (mu):", round(mean(mu_tau_full$mu), 4), "\n")
cat("Between-site heterogeneity (tau):", round(mean(mu_tau_full$tau), 4), "\n\n")
cat("Site-specific calibrated estimates:\n")
print(site_report)

## ----resources, eval=FALSE----------------------------------------------------
# # See also:
# vignette("getting_started")
# vignette("tidy_bayesian_workflow")
# vignette("brms_integration")

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