## ----setup, include = FALSE---------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 5,
  dev.args = list(bg = "white")
)
if (requireNamespace("thematic", quietly = TRUE)) {
  thematic::thematic_off()
}

# Define a knitr hook to fold source code into a <details> block below the plot
local({
  hook_source <- knitr::knit_hooks$get("source")
  knitr::knit_hooks$set(source = function(x, options) {
    res <- hook_source(x, options)
    if (isTRUE(options$fold) || !is.null(options$code_summary)) {
      label <- options$code_summary
      if (is.null(label)) label <- "Show code"
      sprintf("\n<details><summary><b>%s</b></summary>\n\n%s\n</details>\n", label, res)
    } else {
      res
    }
  })

  hook_chunk <- knitr::knit_hooks$get("chunk")
  knitr::knit_hooks$set(chunk = function(x, options) {
    res <- hook_chunk(x, options)
    if (isTRUE(options$fold)) {
      sprintf("\n<div class=\"chunk-plot-first\">\n%s\n</div>\n", res)
    } else {
      res
    }
  })
})

library(imuGAP)
library(data.table)
library(ggplot2)
library(knitr)

# Always enforce white background and black text on plots
ggplot2::theme_set(
  ggplot2::theme_bw() +
    ggplot2::theme(
      plot.background = ggplot2::element_rect(fill = "white", colour = NA),
      panel.background = ggplot2::element_rect(fill = "white", colour = NA),
      legend.background = ggplot2::element_rect(fill = "white", colour = NA),
      legend.box.background = ggplot2::element_blank(),
      text = ggplot2::element_text(colour = "black"),
      axis.text = ggplot2::element_text(colour = "black"),
      axis.title = ggplot2::element_text(colour = "black"),
      plot.title = ggplot2::element_text(colour = "black"),
      plot.subtitle = ggplot2::element_text(colour = "black"),
      strip.text = ggplot2::element_text(colour = "black"),
      strip.background = ggplot2::element_rect(fill = "grey90", colour = "grey20")
    )
)

## ----hierarchy-overview, echo = FALSE-----------------------------------------
state_name <- locations_sim[is.na(parent_id), loc_id]
counties <- locations_sim[parent_id == state_name, loc_id]
schools_by_county <- split(
  locations_sim[parent_id %in% counties, loc_id],
  locations_sim[parent_id %in% counties, parent_id]
)

max_len <- max(vapply(schools_by_county, length, integer(1)))
pad <- function(x, n) c(x, rep("", n - length(x)))

horiz_df <- as.data.frame(
  lapply(schools_by_county, pad, n = max_len),
  col.names = counties,
  check.names = FALSE
)

kable(horiz_df, col.names = paste(counties, "County"))

## ----locs-nlayer--------------------------------------------------------------
locs_3layer <- canonicalize_locations(locations_sim) # all the data
locs_2layer <- locs_3layer[layer <= 2] # state and county only
locs_1layer <- locs_3layer[layer <= 1] # state only

## ----obs-nlayer---------------------------------------------------------------
data("observations_sim", package = "imuGAP")
data("populations_sim", package = "imuGAP")
data("latent_params_sim", package = "imuGAP")

# 3-Layer views (State, County, School)
# 1. Filter populations to locations present in locs_3layer
pops_3layer <- populations_sim[loc_id %in% locs_3layer$loc_id]
# 2. Extract relevant observations corresponding to filtered populations
obs_3layer <- observations_sim[obs_id %in% pops_3layer$obs_id]
# 3. Confirm extracted observations don't have unpresent locations in original population data
stopifnot(!populations_sim[obs_id %in% obs_3layer$obs_id, any(!loc_id %in% locs_3layer$loc_id)])

# 2-Layer views (State and County)
# 1. Filter populations to locations present in locs_2layer
pops_2layer <- populations_sim[loc_id %in% locs_2layer$loc_id]
# 2. Extract relevant observations corresponding to filtered populations
obs_2layer <- observations_sim[obs_id %in% pops_2layer$obs_id]
# 3. Confirm extracted observations don't have unpresent locations in original population data
stopifnot(!populations_sim[obs_id %in% obs_2layer$obs_id, any(!loc_id %in% locs_2layer$loc_id)])

# 1-Layer views (State only)
# 1. Aggregate and filter populations to State level (locs_1layer)
pops_1layer <- copy(populations_sim)[, loc_id := "State"]
pops_1layer <- pops_1layer[, .(weight = sum(weight)), by = .(obs_id, loc_id, cohort, age, dose)]
pops_1layer <- pops_1layer[loc_id %in% locs_1layer$loc_id]
# 2. Extract relevant observations corresponding to filtered populations
obs_1layer <- observations_sim[obs_id %in% pops_1layer$obs_id]
# 3. Confirm extracted observations don't have unpresent locations in population data
stopifnot(!pops_1layer[obs_id %in% obs_1layer$obs_id, any(!loc_id %in% locs_1layer$loc_id)])

## ----fit-code, eval = FALSE---------------------------------------------------
# st_opts <- stan_options(iter = 1000, chains = 4, seed = 1L)
# 
# # various layer fits
# fit_3layer <- sampling(obs_3layer, pops_3layer, locs_3layer, stan_opts = st_opts)
# fit_2layer <- sampling(obs_2layer, pops_2layer, locs_2layer, stan_opts = st_opts)
# fit_1layer <- sampling(obs_1layer, pops_1layer, locs_1layer, stan_opts = st_opts)

## ----load-fits----------------------------------------------------------------
data("fit_sim", package = "imuGAP")
data("fit_sim_2layer", package = "imuGAP")
data("fit_sim_1layer", package = "imuGAP")

## ----predict-code, eval = FALSE-----------------------------------------------
# data("target_sim", package = "imuGAP")
# 
# # 3-Layer Prediction (State, County, School)
# target_3layer <- target_sim[loc_id %in% locs_3layer$loc_id]
# predict_3layer <- predict(object = fit_3layer, target = target_3layer, posterior_size = 100)
# 
# # 2-Layer Prediction (State and County)
# target_2layer <- target_sim[loc_id %in% locs_2layer$loc_id]
# predict_2layer <- predict(object = fit_2layer, target = target_2layer, posterior_size = 100)
# 
# # 1-Layer Prediction (State only)
# target_1layer <- target_sim[loc_id %in% locs_1layer$loc_id]
# predict_1layer <- predict(object = fit_1layer, target = target_1layer, posterior_size = 100)

## ----load-predictions---------------------------------------------------------
data("predict_sim", package = "imuGAP")
data("predict_sim_2layer", package = "imuGAP")
data("predict_sim_1layer", package = "imuGAP")

## ----summarize----------------------------------------------------------------
summary_3layer <- summary(predict_sim)
summary_2layer <- summary(predict_sim_2layer)
summary_1layer <- summary(predict_sim_1layer)

## ----state-plot, fold = TRUE, code_summary = "Show plot code", fig.cap = "State-level coverage comparison across 1-layer, 2-layer, and 3-layer model fits against true values.", fig.width = 8, fig.height = 4----
state_name <- locations_sim[is.na(parent_id), loc_id]

state_1layer_pred <- summary_1layer[
  loc_id == state_name & dose == 2 & age > 4
][, hierarchy := "1-Layer (State Only)"]
state_2layer_pred <- summary_2layer[
  loc_id == state_name & dose == 2 & age > 4
][, hierarchy := "2-Layer (State/County)"]
state_3layer_pred <- summary_3layer[
  loc_id == state_name & dose == 2 & age > 4
][, hierarchy := "3-Layer (State/County/School)"]

state_cov_pred <- rbindlist(list(
  state_1layer_pred,
  state_2layer_pred,
  state_3layer_pred
))
h_levels_3 <- c(
  "1-Layer (State Only)",
  "2-Layer (State/County)",
  "3-Layer (State/County/School)"
)
state_cov_pred[, hierarchy := factor(hierarchy, levels = h_levels_3)]

state_idx <- predict_sim$target[
  loc_id == state_name & dose == 2 & age > 4,
  which = TRUE
]
true_state_df <- data.table(
  age = predict_sim$target[state_idx, age],
  coverage = latent_params_sim$coverage[state_idx]
)

true_state_faceted <- rbindlist(lapply(h_levels_3, function(h) {
  df <- copy(true_state_df)
  df[, hierarchy := factor(h, levels = h_levels_3)]
  df
}))

ggplot() +
  geom_ribbon(
    data = state_cov_pred,
    aes(x = age, ymin = q2_5, ymax = q97_5, fill = "Estimated (95% CI)"),
    alpha = 0.2
  ) +
  geom_line(
    data = state_cov_pred,
    aes(x = age, y = q50, color = "Estimated (Median)"),
    linewidth = 1
  ) +
  geom_point(
    data = true_state_faceted,
    aes(x = age, y = coverage, shape = "True Coverage"),
    color = "black",
    size = 2.2
  ) +
  facet_grid(. ~ hierarchy) +
  scale_shape_manual(name = "", values = c("True Coverage" = 17)) +
  scale_color_manual(name = "", values = c("Estimated (Median)" = "#1b9e77")) +
  scale_fill_manual(name = "", values = c("Estimated (95% CI)" = "#1b9e77")) +
  theme_bw() +
  scale_x_continuous(breaks = seq(5, 18, by = 3), minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.02, 0.05),
    legend.justification.inside = c(0, 0),
    legend.direction = "horizontal",
    legend.box = "horizontal"
  ) +
  labs(
    x = "Age",
    y = "State-Level Two-Dose Coverage"
  )

## ----county-plot, fold = TRUE, code_summary = "Show plot code", fig.cap = "County-level coverage comparison: 2-layer vs. 3-layer model estimates against true values.", fig.width = 8, fig.height = 4----
counties <- locations_sim[parent_id == state_name, loc_id]

county_2layer_pred <- summary_2layer[
  loc_id %in% counties & dose == 2 & age > 4
][, hierarchy := "2-Layer (State/County)"]
county_3layer_pred <- summary_3layer[
  loc_id %in% counties & dose == 2 & age > 4
][, hierarchy := "3-Layer (State/County/School)"]

county_cov_pred <- rbindlist(list(county_2layer_pred, county_3layer_pred))
h_levels_county <- c("2-Layer (State/County)", "3-Layer (State/County/School)")
county_cov_pred[, hierarchy := factor(hierarchy, levels = h_levels_county)]
county_cov_pred[, loc_id := factor(
  loc_id,
  levels = c("Simone", "Watson", "Scruggs")
)]

county_idx <- predict_sim$target[
  loc_id %in% counties & dose == 2 & age > 4,
  which = TRUE
]
true_county_df <- data.table(
  loc_id = factor(
    predict_sim$target[county_idx, loc_id],
    levels = c("Simone", "Watson", "Scruggs")
  ),
  age = predict_sim$target[county_idx, age],
  coverage = latent_params_sim$coverage[county_idx]
)

true_county_faceted <- rbindlist(lapply(h_levels_county, function(h) {
  df <- copy(true_county_df)
  df[, hierarchy := factor(h, levels = h_levels_county)]
  df
}))

ggplot() +
  geom_ribbon(
    data = county_cov_pred,
    aes(x = age, ymin = q2_5, ymax = q97_5, fill = loc_id),
    alpha = 0.15
  ) +
  geom_line(
    data = county_cov_pred,
    aes(x = age, y = q50, color = loc_id),
    linewidth = 0.9
  ) +
  geom_point(
    data = true_county_faceted,
    aes(x = age, y = coverage, color = loc_id, shape = "True Coverage"),
    size = 2
  ) +
  facet_grid(. ~ hierarchy) +
  scale_shape_manual(name = "", values = c("True Coverage" = 18)) +
  scale_color_discrete(NULL, aesthetics = c("color", "fill")) +
  scale_x_continuous(breaks = seq(5, 18, by = 3), minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  theme_bw() +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.02, 0.05),
    legend.justification.inside = c(0, 0),
    legend.direction = "horizontal",
    legend.box = "horizontal"
  ) +
  labs(
    x = "Age",
    y = "County-Level Two-Dose Coverage"
  )

## ----lambda-plot, fold = TRUE, code_summary = "Show plot code", fig.cap = "Force of vaccination (lambda) estimates across data resolutions compared to true values.", fig.width = 7, fig.height = 4----
extract_lambda_summary <- function(fit, label) {
  draws <- rstan::extract(fit$stanfit, pars = "lambda_raw")$lambda_raw
  doses_factor <- factor(c(
    rep("Dose 1", nrow(draws)),
    rep("Dose 2", nrow(draws))
  ))
  data.table(
    hierarchy = label,
    dose = doses_factor,
    lambda_raw = c(draws[, 1], draws[, 2])
  )[, .(
    q50 = stats::median(lambda_raw),
    q2_5 = stats::quantile(lambda_raw, 0.025),
    q97_5 = stats::quantile(lambda_raw, 0.975)
  ), by = .(hierarchy, dose)]
}

lambda_est <- rbindlist(list(
  extract_lambda_summary(fit_sim_1layer, "1-Layer"),
  extract_lambda_summary(fit_sim_2layer, "2-Layer"),
  extract_lambda_summary(fit_sim, "3-Layer")
))
h_levels_lambda <- c("1-Layer", "2-Layer", "3-Layer")
lambda_est[, hierarchy := factor(hierarchy, levels = h_levels_lambda)]

true_lambda <- data.table(
  dose = factor(c("Dose 1", "Dose 2")),
  hierarchy = factor("1-Layer", levels = h_levels_lambda),
  true_val = log(latent_params_sim$lambda),
  label = sprintf("True~lambda == %.1f", latent_params_sim$lambda)
)

ggplot(lambda_est, aes(x = hierarchy, y = q50)) +
  geom_hline(
    data = true_lambda,
    aes(yintercept = true_val),
    color = "firebrick",
    linetype = "dashed",
    linewidth = 0.8
  ) +
  geom_label(
    data = true_lambda,
    aes(x = hierarchy, y = true_val, label = label),
    parse = TRUE,
    color = "firebrick",
    fill = ggplot2::alpha("white", 0.75),
    linewidth = NA,
    vjust = -0.3,
    hjust = 0.1,
    size = 3.2
  ) +
  geom_pointrange(
    aes(ymin = q2_5, ymax = q97_5),
    size = 0.7
  ) +
  facet_wrap(~dose) +
  coord_cartesian(ylim = c(0.5, 1.5)) +
  scale_y_continuous(
    transform = "exp",
    labels = function(x) sprintf("%.2f", exp(x))
  ) +
  theme_bw() +
  labs(
    x = "Data Resolution (Model Hierarchy)",
    y = "Uptake Rate (exponential scale)"
  )