## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 5,
  out.width = "100%",
  dev.args = list(bg = "white")
)

## ----setup, echo = FALSE, message = FALSE, warning = FALSE--------------------
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
    }
  })
})

# List all suggested packages needed for this vignette
required_pkgs <- c(
  "imuGAP", "data.table", "bayesplot", "ggplot2",
  "tidyr", "dplyr"
)

# Check if they are available
pkgs_available <- all(sapply(required_pkgs, requireNamespace, quietly = TRUE))

# If not available, hide the rest of the document or knit dynamically
if (!pkgs_available) {
  knitr::opts_chunk$set(eval = FALSE)
  message("Some suggested packages are missing. Vignette code will not be executed.")
} else {
  # Assign the result so the chunk does not auto-print the logical vector
  # sapply() returns (the setup block should render no output).
  chk <- suppressPackageStartupMessages(
    sapply(required_pkgs, require, character.only = TRUE)
  )

  # 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")
      )
  )
}

## ----locations----------------------------------------------------------------
data("locations_sim", package = "imuGAP")
head(locations_sim)

# Canonicalize and validate
canonical_locations <- canonicalize_locations(locations_sim)
head(canonical_locations)

## ----observations-------------------------------------------------------------
data("observations_sim", package = "imuGAP")
head(observations_sim[, .(obs_id, loc_id, positive, sample_n, censored)])

# Canonicalize and validate
canonical_observations <- canonicalize_observations(observations_sim)
head(canonical_observations)

## ----populations--------------------------------------------------------------
data("populations_sim", package = "imuGAP")
head(populations_sim)

# Canonicalize and validate
canonical_populations <- canonicalize_populations(
  populations_sim, observations_sim, locations_sim
)
head(canonical_populations)

## ----validation-failure-obs---------------------------------------------------
# Create a copy with an invalid observation (positive > sample_n)
invalid_obs <- copy(observations_sim[, .(obs_id, loc_id, positive, sample_n, censored)])
invalid_obs[1, positive := sample_n + 10]

# This will fail validation and throw an error:
tryCatch(
  canonicalize_observations(invalid_obs),
  error = function(e) message("Caught expected error: ", e$message)
)

## ----validation-failure-locs--------------------------------------------------
# Create a copy with a duplicate location ID
invalid_locs <- rbind(
  locations_sim,
  data.frame(loc_id = "Scruggs", parent_id = "State"),
  fill = TRUE
)

# This will fail validation:
tryCatch(
  canonicalize_locations(invalid_locs),
  error = function(e) message("Caught expected error: ", e$message)
)

## ----synthetic-state-viz, fold = TRUE, code_summary = "Show plot code"--------
data("latent_params_sim", package = "imuGAP")

# Categorize state-level observation sources
state_obs <- copy(observations_sim[loc_id == "State"])
state_obs[, source := factor(
  fcase(
    dose == 1 & age_min == 2, "ChildVaxView (Dose 1, Age 2)",
    dose == 1 & age_min == 3, "ChildVaxView (Dose 1, Age 3)",
    age_min == 5, "SchoolVaxView (Dose 2, Age 5)",
    default = "TeenVaxView (Dose 2, Ages 14-18)"
  ),
  levels = c(
    "ChildVaxView (Dose 1, Age 2)",
    "ChildVaxView (Dose 1, Age 3)",
    "SchoolVaxView (Dose 2, Age 5)",
    "TeenVaxView (Dose 2, Ages 14-18)"
  )
)]
state_obs[, obs_prop := positive / sample_n]

# Split single-cohort point observations vs multi-cohort cross-sectional survey snapshots
single_cohort_obs <- state_obs[is.na(age_max) | age_max == age_min + 1L]
multi_cohort_obs <- copy(state_obs[!is.na(age_max) & age_max > age_min + 1L])
multi_cohort_obs[, cohort_max := cohort_min + (age_max - 1L) - age_min]

# True state lifetime propensity across cohorts
latent_state <- data.table(
  cohort_min = seq_along(latent_params_sim$phi_state),
  phi = latent_params_sim$phi_state
)

# Latent milestone coverage curves corresponding to each observation source
n_c <- length(latent_params_sim$phi_state)
latent_curves <- rbindlist(list(
  data.table(
    cohort_min = seq_len(n_c),
    latent_cov = latent_params_sim$phi_state *
      latent_params_sim$uptake[2, 1] *
      latent_params_sim$censor_reduction,
    source = "ChildVaxView (Dose 1, Age 2)"
  ),
  data.table(
    cohort_min = seq_len(n_c),
    latent_cov = latent_params_sim$phi_state *
      latent_params_sim$uptake[3, 1] *
      latent_params_sim$censor_reduction,
    source = "ChildVaxView (Dose 1, Age 3)"
  ),
  data.table(
    cohort_min = seq_len(28),
    latent_cov = latent_params_sim$phi_state[1:28] *
      latent_params_sim$uptake[5, 2],
    source = "SchoolVaxView (Dose 2, Age 5)"
  ),
  data.table(
    cohort_min = seq_len(15),
    latent_cov = latent_params_sim$phi_state[1:15] *
      mean(latent_params_sim$uptake[14:18, 2]),
    source = "TeenVaxView (Dose 2, Ages 14-18)"
  )
))
latent_curves[, source := factor(source, levels = levels(state_obs$source))]

ggplot() +
  geom_line(
    data = latent_state,
    aes(x = cohort_min, y = phi, linetype = "True Lifetime Propensity (phi)"),
    color = "gray40",
    linewidth = 0.8,
    alpha = 0.5
  ) +
  geom_line(
    data = latent_curves,
    aes(x = cohort_min, y = latent_cov, color = source),
    linetype = "dashed",
    linewidth = 0.7,
    alpha = 0.4
  ) +
  geom_segment(
    data = multi_cohort_obs,
    aes(
      x = cohort_min,
      xend = cohort_max,
      y = obs_prop,
      yend = obs_prop,
      color = source
    ),
    linewidth = 1.1,
    alpha = 0.95
  ) +
  geom_point(
    data = single_cohort_obs,
    aes(x = cohort_min, y = obs_prop, color = source, shape = source),
    size = 2.4,
    alpha = 0.95
  ) +
  theme_bw() +
  scale_x_continuous(
    limits = c(0, 30),
    breaks = seq(0, 30, by = 5),
    minor_breaks = seq(1, 30, by = 1)
  ) +
  scale_y_continuous(limits = c(0.4, 1.0)) +
  scale_linetype_manual(
    name = NULL,
    values = c("True Lifetime Propensity (phi)" = "dashed")
  ) +
  scale_color_brewer(name = "Data Source", palette = "Dark2") +
  scale_shape_manual(
    name = "Data Source",
    values = c(
      "ChildVaxView (Dose 1, Age 2)" = 16,
      "ChildVaxView (Dose 1, Age 3)" = 17,
      "SchoolVaxView (Dose 2, Age 5)" = 15,
      "TeenVaxView (Dose 2, Ages 14-18)" = 18
    )
  ) +
  guides(
    color = guide_legend(
      override.aes = list(
        shape = c(16, 17, 15, NA),
        linetype = c("blank", "blank", "blank", "solid"),
        linewidth = c(0, 0, 0, 1.1),
        alpha = 1
      )
    ),
    shape = "none"
  ) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.98, 0.02),
    legend.justification.inside = c(1, 0)
  ) +
  labs(
    x = "Birth Cohort Index",
    y = "Vaccination Proportion"
  )

## ----synthetic-county-viz, fold = TRUE, code_summary = "Show plot code"-------
county_obs <- copy(observations_sim[loc_id %in% c("Scruggs", "Simone", "Watson")])
county_obs[, obs_prop := positive / sample_n]

# Analytical county-level latent curves for 6th grade survey (age 11, dose 2, censored)
county_latent <- rbindlist(lapply(names(latent_params_sim$off_cnty), function(cnty) {
  cohorts <- seq_len(19)
  c_idx <- match(cnty, names(latent_params_sim$off_cnty))
  offset <- latent_params_sim$off_cnty[c_idx]
  phi_shifted <- plogis(qlogis(latent_params_sim$phi_state[cohorts]) + offset)
  cov_true <- phi_shifted * latent_params_sim$uptake[11, 2] * latent_params_sim$censor_reduction
  data.table(loc_id = cnty, cohort_min = cohorts, latent_cov = cov_true)
}))

ggplot() +
  geom_point(
    data = county_obs,
    aes(x = cohort_min, y = obs_prop),
    color = "steelblue", size = 2, alpha = 0.85
  ) +
  geom_line(
    data = county_latent,
    aes(x = cohort_min, y = latent_cov, color = "True Latent Coverage"),
    linetype = "dashed", linewidth = 0.9
  ) +
  facet_wrap(~loc_id) +
  theme_bw() +
  scale_x_continuous(
    limits = c(0, 30),
    breaks = seq(0, 30, by = 5),
    minor_breaks = seq(1, 30, by = 1)
  ) +
  scale_y_continuous(limits = c(0.4, 1.0)) +
  scale_color_manual(name = NULL, values = c("True Latent Coverage" = "firebrick")) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.85, 0.15),
    legend.justification.inside = c(1, 0)
  ) +
  labs(
    x = "Birth Cohort Index",
    y = "6th Grade Survey Coverage (Dose 2, Age 11)"
  )

## ----synthetic-school-viz, fold = TRUE, code_summary = "Show plot code"-------
# Select representative schools at the 0, 0.25, 0.5, 0.75, and 1 quantiles of school offsets
sch_info <- locations_sim[!loc_id %in% c("State", "Scruggs", "Simone", "Watson")]
sch_info[, off := latent_params_sim$off_sch[loc_id]]

probs <- c(0, 0.25, 0.5, 0.75, 1)
labels <- c("0% (Min)", "25% (Q1)", "50% (Median)", "75% (Q3)", "100% (Max)")

sel_schools <- sch_info[, {
  q_vals <- quantile(off, probs = probs, type = 7)
  chosen_idx <- sapply(q_vals, function(qv) which.min(abs(off - qv)))
  .(
    quantile_label = factor(labels, levels = labels),
    loc_id = loc_id[chosen_idx],
    off = off[chosen_idx]
  )
}, by = parent_id]

# Filter school observations to the selected quantile schools
school_obs <- merge(
  observations_sim,
  sel_schools[, .(parent_id, loc_id, quantile_label)],
  by = c("parent_id", "loc_id")
)
school_obs[, obs_prop := positive / sample_n]

sch_cohorts <- 1:28

# 1. State-level lifetime propensity reference
state_sch_propensity <- rbindlist(lapply(
  c("Scruggs", "Simone", "Watson"),
  function(cnty) {
    data.table(
      parent_id = cnty,
      cohort_min = sch_cohorts,
      phi = latent_params_sim$phi_state[sch_cohorts]
    )
  }
))

# 2. County-level latent milestone trajectory (age 5, dose 2)
county_sch_latent <- rbindlist(lapply(
  names(latent_params_sim$off_cnty),
  function(cnty) {
    c_idx <- match(cnty, names(latent_params_sim$off_cnty))
    offset <- latent_params_sim$off_cnty[c_idx]
    phi_shifted <- plogis(
      qlogis(latent_params_sim$phi_state[sch_cohorts]) + offset
    )
    cov_true <- phi_shifted * latent_params_sim$uptake[5, 2]
    data.table(parent_id = cnty, cohort_min = sch_cohorts, latent_cov = cov_true)
  }
))

# 3. School-level latent milestone trajectories for selected quantile schools
school_sch_latent <- rbindlist(lapply(
  seq_len(nrow(sel_schools)),
  function(i) {
    row <- sel_schools[i]
    cnty <- row$parent_id
    s_name <- row$loc_id
    q_lab <- row$quantile_label
    c_offset <- latent_params_sim$off_cnty[cnty]
    s_offset <- latent_params_sim$off_sch[s_name]
    phi_sch <- plogis(
      qlogis(latent_params_sim$phi_state[sch_cohorts]) + c_offset + s_offset
    )
    cov_sch <- phi_sch * latent_params_sim$uptake[5, 2]
    data.table(
      parent_id = cnty,
      loc_id = s_name,
      quantile_label = q_lab,
      cohort_min = sch_cohorts,
      latent_cov = cov_sch
    )
  }
))

ggplot() +
  # State lifetime propensity reference
  geom_line(
    data = state_sch_propensity,
    aes(x = cohort_min, y = phi, linetype = "True State Lifetime Propensity (phi)"),
    color = "gray40",
    linewidth = 0.8,
    alpha = 0.5
  ) +
  # County latent curve
  geom_line(
    data = county_sch_latent,
    aes(x = cohort_min, y = latent_cov, linetype = "True County Latent Coverage"),
    color = "firebrick",
    linewidth = 0.9
  ) +
  # School latent curves
  geom_line(
    data = school_sch_latent,
    aes(x = cohort_min, y = latent_cov, color = quantile_label, group = loc_id),
    linetype = "dashed",
    linewidth = 0.7,
    alpha = 0.8
  ) +
  # School observation points (faded)
  geom_point(
    data = school_obs,
    aes(x = cohort_min, y = obs_prop, color = quantile_label),
    size = 1.8,
    alpha = 0.6
  ) +
  facet_wrap(~parent_id) +
  theme_bw() +
  scale_x_continuous(
    limits = c(0, 30),
    breaks = seq(0, 30, by = 5),
    minor_breaks = seq(1, 30, by = 1)
  ) +
  scale_y_continuous(limits = c(0.4, 1.0)) +
  scale_color_viridis_d(name = "School Quantile", option = "plasma", end = 0.9) +
  scale_linetype_manual(
    name = "Reference Curves",
    values = c(
      "True State Lifetime Propensity (phi)" = "dotted",
      "True County Latent Coverage" = "solid"
    )
  ) +
  guides(
    color = guide_legend(reverse = TRUE, order = 1),
    linetype = guide_legend(order = 2)
  ) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.02, 0.05),
    legend.justification.inside = c(0, 0),
    legend.background = element_rect(fill = alpha("white", 0.8), color = NA),
    legend.box = "horizontal",
    legend.spacing.x = unit(0.3, "cm")
  ) +
  labs(
    x = "Birth Cohort Index",
    y = "Kindergarten Entry Coverage (Dose 2, Age 5)"
  )

## ----sampling-code, eval = FALSE----------------------------------------------
# fit_sim <- sampling(
#   observations_sim, populations_sim, locations_sim,
#   stan_opts = stan_options(
#     iter = 2000, chains = 4, refresh = 0, seed = 1L
#   )
# )

## ----load-fit-----------------------------------------------------------------
data("fit_sim", package = "imuGAP")

## ----extract-params-----------------------------------------------------------
beta_draws <- extract_imugap(fit_sim, pars = "beta_bs")
str(beta_draws)

## ----trace-plot-beta, fold = TRUE, code_summary = "Show plot code"------------
bayesplot::mcmc_trace(
  fit_sim$stanfit,
  pars = c(
    "beta_bs[1]", "beta_bs[2]", "beta_bs[3]",
    "beta_bs[4]", "beta_bs[5]"
  )
) +
  theme_bw() +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.9, 0.1),
    legend.justification.inside = c(1, 0)
  )

## ----trace-plot-sigmas, fold = TRUE, code_summary = "Show plot code"----------
sigma_ref <- data.frame(
  parameter = c("sigma_layer[1]", "sigma_layer[2]"),
  true_val = c(latent_params_sim$sigma_cnty, latent_params_sim$sigma_sch),
  label = sprintf(
    "True~sigma == %.2f",
    c(latent_params_sim$sigma_cnty, latent_params_sim$sigma_sch)
  )
)

bayesplot::mcmc_trace(
  fit_sim$stanfit,
  pars = c("sigma_layer[1]", "sigma_layer[2]"),
  facet_args = list(labeller = ggplot2::as_labeller(c(
    "sigma_layer[1]" = "sigma[County]",
    "sigma_layer[2]" = "sigma[School]"
  ), default = ggplot2::label_parsed))
) +
  geom_hline(
    data = sigma_ref,
    aes(yintercept = true_val),
    color = "firebrick",
    linetype = "dashed",
    linewidth = 0.8
  ) +
  geom_label(
    data = sigma_ref,
    aes(x = 100, y = true_val, label = label),
    parse = TRUE,
    color = "firebrick",
    fill = ggplot2::alpha("white", 0.75),
    linewidth = NA,
    vjust = -0.3,
    hjust = 0,
    size = 3.2
  ) +
  coord_cartesian(ylim = c(0, 2.5)) +
  theme_bw() +
  theme(legend.position = "bottom")

## ----trace-plot-lambdas, fold = TRUE, code_summary = "Show plot code"---------
lambda_ref <- data.frame(
  parameter = c("lambda_raw[1]", "lambda_raw[2]"),
  true_val = log(latent_params_sim$lambda),
  label = sprintf("True~lambda == %.1f", latent_params_sim$lambda)
)

bayesplot::mcmc_trace(
  fit_sim$stanfit,
  pars = c("lambda_raw[1]", "lambda_raw[2]"),
  facet_args = list(labeller = ggplot2::as_labeller(c(
    "lambda_raw[1]" = "lambda[1]~(Dose~1)",
    "lambda_raw[2]" = "lambda[2]~(Dose~2)"
  ), default = ggplot2::label_parsed))
) +
  geom_hline(
    data = lambda_ref,
    aes(yintercept = true_val),
    color = "firebrick",
    linetype = "dashed",
    linewidth = 0.8
  ) +
  geom_label(
    data = lambda_ref,
    aes(x = 100, y = true_val, label = label),
    parse = TRUE,
    color = "firebrick",
    fill = ggplot2::alpha("white", 0.75),
    linewidth = NA,
    vjust = -0.3,
    hjust = 0,
    size = 3.2
  ) +
  coord_cartesian(ylim = c(0.5, 1.5)) +
  scale_y_continuous(
    transform = "exp",
    labels = function(x) sprintf("%.2f", exp(x))
  ) +
  labs(y = "Uptake Rate (exponential scale)") +
  theme_bw() +
  theme(legend.position = "bottom")

## ----target-------------------------------------------------------------------
target_sim <- create_target(
  location = unique(locations_sim$loc_id), age = 1:18,
  cohort = max(populations_sim$cohort) - 18, dose = c(1, 2), mode = "snapshot"
)
head(target_sim)

## ----predict-code, eval = FALSE-----------------------------------------------
# predict_sim <- predict(object = fit_sim, target = target_sim, posterior_size = 100)

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

## ----summarize-predictions----------------------------------------------------
# Calculate the posterior mean coverage probability for each location and dose at age 5
summary_predict <- summary(predict_sim)
head(summary_predict)

## ----state-viz, fold = TRUE, code_summary = "Show plot code"------------------
data("latent_params_sim", package = "imuGAP")

# Filter predictions for the State level, dose 2, and ages > 4
state_predict <- summary_predict[loc_id == "State" & dose == 2 & age > 4]

# Create the lookup index for the matching target populations to attach true latent values
state_idx <- predict_sim$target[loc_id == "State" & dose == 2 & age > 4, which = TRUE]
state_predict[, latent := latent_params_sim$coverage[state_idx]]

ggplot(state_predict) +
  aes(x = age) +
  geom_ribbon(aes(ymin = q2_5, ymax = q97_5, fill = "95% Credible Interval"), alpha = 0.25) +
  geom_line(aes(y = q50, color = "Posterior Median"), linewidth = 0.8) +
  geom_line(aes(y = latent, color = "True Latent"), linetype = "dashed", linewidth = 0.8) +
  theme_bw() +
  scale_x_continuous(breaks = 5:18, minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  scale_color_manual(
    name = NULL,
    values = c("Posterior Median" = "black", "True Latent" = "firebrick")
  ) +
  scale_fill_manual(name = NULL, values = c("95% Credible Interval" = "grey50")) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.05, 0.05),
    legend.justification.inside = c(0, 0)
  ) +
  labs(x = "Age", y = "State-Level Two-Dose Coverage")

## ----county-viz, fold = TRUE, code_summary = "Show plot code"-----------------
summary_predict |>
  subset(loc_id %in% c("Scruggs", "Simone", "Watson") & dose == 2 & age > 4) |>
  transform(loc_id = factor(loc_id, levels = c("Simone", "Watson", "Scruggs"))) |>
  ggplot() +
  aes(x = age) +
  geom_line(aes(y = q50, color = loc_id)) +
  geom_ribbon(aes(ymin = q2_5, ymax = q97_5, fill = loc_id), alpha = 0.2) +
  theme_bw() +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.12, 0.05),
    legend.justification.inside = c(0, 0)
  ) +
  scale_x_continuous(breaks = 5:18, minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  scale_color_discrete(NULL, aesthetics = c("color", "fill")) +
  labs(
    x = "Age", y = "County-Level Two-Dose Coverage"
  )

## ----grade-viz, fold = TRUE, code_summary = "Show plot code"------------------
scruggs_schools <- locations_sim[parent_id == "Scruggs", loc_id]
off_scruggs <- latent_params_sim$off_sch[scruggs_schools]
med_sch <- names(off_scruggs)[which.min(
  abs(off_scruggs - stats::quantile(off_scruggs, 0.5))
)]

predict_sch <- subset(predict_sim, loc_id == med_sch & dose == 2 & age > 4)
draws_sch <- as.data.frame(predict_sch)

sch_target <- predict_sch$target
sch_target$latent <- latent_params_sim$coverage[sch_target$obs_id]

summary_sch <- summary(predict_sch)
summary_sch$latent <- sch_target$latent

ggplot() +
  geom_line(
    data = draws_sch,
    aes(
      x = age,
      y = coverage,
      group = interaction(chain, iteration),
      color = "Posterior Draws"
    ),
    alpha = 0.12,
    linewidth = 0.4
  ) +
  geom_line(
    data = summary_sch,
    aes(x = age, y = q50, color = "Posterior Median"),
    linewidth = 0.9
  ) +
  geom_line(
    data = summary_sch,
    aes(x = age, y = latent, color = "True Latent"),
    linetype = "dashed",
    linewidth = 0.9
  ) +
  theme_bw() +
  scale_x_continuous(breaks = 5:18, minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  scale_color_manual(
    name = NULL,
    values = c(
      "Posterior Median" = "black",
      "True Latent" = "firebrick",
      "Posterior Draws" = "steelblue"
    ),
    guide = guide_legend(override.aes = list(
      linewidth = c(0.9, 0.9, 0.8),
      linetype = c("solid", "dashed", "solid"),
      alpha = c(1, 1, 0.6)
    ))
  ) +
  annotate(
    "text",
    x = 18, y = 0.99,
    label = sprintf("%s (50%% Quantile School)", med_sch),
    hjust = 1, vjust = 1,
    size = 3.5, fontface = "italic"
  ) +
  theme(
    legend.position = "inside",
    legend.position.inside = c(0.05, 0.05),
    legend.justification.inside = c(0, 0)
  ) +
  labs(
    x = "Age",
    y = "Two-Dose Coverage"
  )

## ----school-viz, fold = TRUE, code_summary = "Show plot code"-----------------
schools <- c(
  "Towhee Children's Academy", # ~380 per grade
  "Flycatcher Elementary", # ~110 per grade
  "Sparrow School" # ~60 per grade
)

# Subset to targets of interest (all retained posterior draws)
predict_sub <- predict_sim |>
  subset(loc_id %in% schools & dose == 2 & age > 4)

# Get the pre-computed background coverage matching the subsetted target
target_idx <- predict_sim$target[loc_id %in% schools & dose == 2 & age > 4, which = TRUE]
latent_ref <- copy(predict_sub$target)
latent_ref$coverage <- latent_params_sim$coverage[target_idx]

# Convert predictions to a long-format data.frame
draws_df <- as.data.frame(predict_sub)

# Now plot it all
ggplot() +
  aes(age, coverage, color = loc_id) +
  geom_point(
    data = draws_df,
    alpha = 0.15, shape = 16, size = 1.2,
    position = position_jitterdodge(
      dodge.width = 0.5,
      jitter.width = 0.15
    )
  ) +
  geom_point(
    data = latent_ref,
    mapping = aes(shape = "True value"),
    size = 2.5,
    stroke = 1.1,
    position = position_dodge(width = 0.5)
  ) +
  theme_bw() +
  scale_shape_manual(
    name = "",
    values = c("True value" = 24)
  ) +
  scale_color_discrete(NULL, aesthetics = c("color", "fill")) +
  scale_x_continuous(breaks = 5:18, minor_breaks = NULL) +
  scale_y_continuous(limits = c(0.8, 1.0)) +
  theme(legend.position = "bottom") +
  labs(color = "School", x = "Age", y = "Two-Dose Coverage")