# ═══════════════════════════════════════════════════════════════
# h3n2 non-responders analysis
# Author: Jose Victor Zambrana & Boshu Chen (University of Michigan - Ann Arbor)
# Created: 2025-03-28 | Modified: 06/12/2026
# ═══════════════════════════════════════════════════════════════

# ── Libraries ──────────────────────────────────────────────────────

library(tidyverse)   # v2.0.0 – Data wrangling, reshaping (pivot_longer), and ggplot2 plotting
library(gtsummary)   # v2.4.0 – Builds formatted summary/comparison tables (tbl_summary, tbl_merge)
library(broom)       # v1.0.12 – Tidies model output into data frames (tidy() for glm/multinom)
library(nnet)        # v7.3.20 – Fits multinomial logistic regression (multinom)
library(splines)     # v4.5.0 – Piecewise-linear B-splines in regression models (bs())
library(patchwork)   # v1.3.2 – Combines multiple ggplot panels into composite figures (/ and plot_layout)
library(ggpubr)      # v0.6.0 – Adds statistical test annotations to plots (stat_compare_means)
library(ggvenn)      # v0.1.19 – Venn diagram of overlapping antibody response sets
library(ggupset)     # v0.4.1 – UpSet plot showing combination frequencies of responses
library(DiagrammeR)  # 1.0.12 - Renders Graphviz flowcharts (grViz) for study enrollment diagram



# ── Load data ─────────────────────────────────────────────────

load("Data/simulated_data.RData")


# ═══════════════════════════════════════════════════════════════
# Data prep for tables ---------------------
# ═══════════════════════════════════════════════════════════════

c05 <- c05 %>%
  mutate(
    group1 = if_else(hk14_4fold, "HAI Responder", "HAI Non-Responder"),
    group2 = relevel(factor(group2), ref = "No Response"),
    index  = probable_index1 == 1,
    age_group = factor(if_else(age_15, "15+", "0-14"), levels = c("0-14", "15+")),
    male   = sexo == "M",
    log2_pre_hai   = log2(pmax(pre_hai_hk14, 5)), # tweak on simulated data only to have data floored at 5
    log2_flha_pre  = log2(pmax(flha_auc_pre, 5)),
    log2_stalk_pre = log2(pmax(stalk_auc_pre, 5)),
    log2_na_pre    = log2(pmax(na_auc_pre, 5)),
    pre_hai_hk14_cat  = factor(case_when(
      pre_hai_hk14 == 5 ~ "5", pre_hai_hk14 <= 80 ~ "6-80", .default = ">80"
    ), levels = c("5", "6-80", ">80")),
    flha_auc_pre_cat  = factor(case_when(
      flha_auc_pre == 5 ~ "5", flha_auc_pre <= 80 ~ "6-80", .default = ">80"
    ), levels = c("5", "6-80", ">80")),
    stalk_auc_pre_cat = factor(case_when(
      stalk_auc_pre == 5 ~ "5", stalk_auc_pre <= 80 ~ "6-80", .default = ">80"
    ), levels = c("5", "6-80", ">80")),
    na_auc_pre_cat    = factor(case_when(
      na_auc_pre == 5 ~ "5", na_auc_pre <= 80 ~ "6-80", .default = ">80"
    ), levels = c("5", "6-80", ">80"))
  ) 


# ═══════════════════════════════════════════════════════════════
# Builder for Tables 1 and 2 ------
# ═══════════════════════════════════════════════════════════════

c05 <- c05 %>%
  mutate(
    any_alt_rise = flha_4fold | stalk_4fold | na_4fold,  # adjust variable names
    group3 = case_when(
      group2 == "HAI Responder" & any_alt_rise  ~ "Full Responder",
      group2 == "HAI Responder" & !any_alt_rise ~ "HAI Responder Only",
      group2 == "Alternate Response"            ~ "Alternate Response Only",
      group2 == "No Response"                   ~ "No Response"
    )
  )


# Creates a comparison table of characteristics across response groups
make_comparison_table <- function(data, 
                                  response_var = "group1",
                                  vars = NULL,
                                  overall = TRUE) {
  
  # Default set of variables to include in the table
  all_vars <- c("impage","age_group", "ever_vax", "male", "fever", "ili", "ari", "index",
                "cough_duration", "CT_low",
                "log2_pre_hai", "log2_flha_pre", "log2_stalk_pre", "log2_na_pre"
  )
  # Override defaults if I specified which variables to use
  if (!is.null(vars)) all_vars <- vars
  all_vars <- setdiff(all_vars, response_var)
  
  # Determine the response group levels and how many there are
  resp_levels <- if (is.factor(data[[response_var]])) levels(data[[response_var]]) else sort(unique(data[[response_var]]))
  n_resp <- length(resp_levels)
  
  # Use Wilcoxon for 2-group comparisons, Kruskal-Wallis for 3+ groups
  cont_test <- if (n_resp == 2) "wilcox.test" else "kruskal.test"
  
  # Classify variables by type so each gets appropriate summary statistics and tests
  dicho_vars <- intersect(c("male", "ever_vax", "index", "fever", "ili", "ari"), all_vars)
  cat_vars   <- intersect(c("age_group"), all_vars)
  cont_vars  <- intersect(c("cough_duration", "CT_low", "log2_pre_hai",
                            "log2_flha_pre", "log2_stalk_pre", "log2_na_pre"), all_vars)
  
  # Map variable names to human-readable labels for the table output
  label_map <- c(
    impage = "Age (years)",
    log2_pre_hai = "log\u2082 HAI titer",
    log2_flha_pre = "log\u2082 FLHA titer",
    log2_stalk_pre = "log\u2082 HA Stalk titer",
    log2_na_pre = "log\u2082 NA titer",
    age_group = "Age group",
    male = "Male",
    index = "Index case",
    fever = "Fever",
    ili = "ILI",
    ari = "ARI",
    cough_duration = "Cough duration (days)",
    CT_low = "Min Ct value",
    ever_vax = "Ever vaccinated"
  )
  
  # Prepare the data: set response variable as the grouping factor
  df <- data %>%
    mutate(grp = factor(.data[[response_var]], levels = resp_levels)) %>%
    select(grp, all_of(all_vars))
  
  for (v in intersect(names(df), names(label_map))) {
    attr(df[[v]], "label") <- label_map[[v]]
  }
  
  # Tell gtsummary how to treat each variable type
  type_list <- list()
  if (length(dicho_vars) > 0) type_list <- c(type_list, list(all_of(dicho_vars) ~ "dichotomous"))
  if (length(cat_vars) > 0)   type_list <- c(type_list, list(all_of(cat_vars) ~ "categorical"))
  if (length(cont_vars) > 0)  type_list <- c(type_list, list(all_of(cont_vars) ~ "continuous"))
  
  # Apply the appropriate non-parametric test for continuous variables
  test_list <- list()
  if (length(cont_vars) > 0) test_list <- c(test_list, list(all_of(cont_vars) ~ cont_test))
  
  # Build the summary table: counts/percentages for categorical, medians/IQR for continuous, plus p-values
  df %>%
    tbl_summary(by = grp, type = type_list,
                statistic = list(all_dichotomous() ~ "{n} ({p}%)",
                                 all_categorical() ~ "{n} ({p}%)",
                                 all_continuous()  ~ "{median} ({p25} - {p75})"),
                digits = list(all_continuous() ~ 1, all_dichotomous() ~ c(0, 1), all_categorical() ~ c(0,1)),
                missing = "no") %>%
    # Optionally add an "Overall" column combining all groups
    {if (overall) add_overall(.) else .} %>%
    add_stat_label(location = "row") %>%
    add_p(test = test_list) %>%
    modify_fmt_fun(p.value ~ function(x) formatC(signif(x, 3), format = "fg")) %>%
    separate_p_footnotes()
}


# ═══════════════════════════════════════════════════════════════
# Table 1 - HAI Responder versus Non-Responder characteristics -----
# ═══════════════════════════════════════════════════════════════

table1 = make_comparison_table(c05, response_var = "group1")

# ═══════════════════════════════════════════════════════════════
# Table 2 - HAI Responders versus Alternate responders versus No response ----
table2 = make_comparison_table(c05, response_var = "group2", overall = FALSE)
# ═══════════════════════════════════════════════════════════════

# ═══════════════════════════════════════════════════════════════
# Regression function -------------
# ═══════════════════════════════════════════════════════════════

# Fits a logistic (2-level) or multinomial (3+ level) regression and returns odds ratios
fit_response_model <- function(data,
                               outcome    = "hk14_4fold",
                               assay_cat  = "pre_hai_hk14_cat",
                               symptom    = NULL,
                               covariates = c("age_group", "male", "index"),
                               ref_level  = NULL) {
  
  # Keep only the columns needed for the model
  model_vars <- c(outcome, covariates, assay_cat, symptom)
  df <- data %>%
    select(all_of(model_vars)) %>%
    drop_na()
  
  # Set the reference level for categorical assay titer groups (it defaults to "6-80")
  if (is.character(df[[assay_cat]]) || is.factor(df[[assay_cat]])) {
    df[[assay_cat]] <- relevel(factor(df[[assay_cat]]), ref = "6-80")
  }
  
  # Determine whether the outcome is binary or multi-level
  outcome_vals <- unique(df[[outcome]])
  n_levels <- length(outcome_vals)
  
  # Build the regression formula from covariates, assay term, and symptom variable
  rhs <- paste(c(covariates, assay_cat, symptom), collapse = " + ")
  fml <- as.formula(paste(outcome, "~", rhs))
  
  # Binary outcome: fit standard logistic regression
  if (n_levels == 2) {
    # Convert logical outcome to a labeled factor
    if (is.logical(df[[outcome]])) {
      df[[outcome]] <- factor(df[[outcome]], levels = c(FALSE, TRUE),
                              labels = c("Non-Responder", "Responder"))
    }
    # Flexibility to specify which level is the reference category
    if (!is.null(ref_level)) {
      lvls <- c(ref_level, setdiff(unique(df[[outcome]]), ref_level))
      df[[outcome]] <- factor(df[[outcome]], levels = lvls)
    }
    fit <- glm(fml, data = df, family = binomial)
    
    # Extract exponentiated coefficients (odds ratios) with confidence intervals
    results <- tidy(fit, conf.int = TRUE, exponentiate = TRUE) %>%
      filter(term != "(Intercept)") %>%
      transmute(
        outcome_model = outcome,
        comparison    = paste(levels(factor(df[[outcome]])), collapse = " vs "),
        assay         = assay_cat,
        symptom_var   = ifelse(is.null(symptom), "none", symptom),
        term,
        OR    = estimate,
        lower = conf.low,
        upper = conf.high,
        p     = round(p.value, 4),
        n     = nobs(fit)
      )
  } else {
    # Multi-level outcome: fit multinomial logistic regression
    if (!is.null(ref_level)) {
      lvls <- c(ref_level, setdiff(unique(df[[outcome]]), ref_level))
      df[[outcome]] <- factor(df[[outcome]], levels = lvls)
    } else {
      df[[outcome]] <- factor(df[[outcome]])
    }
    fit <- multinom(fml, data = df, trace = FALSE)
    
    # Extract odds ratios for each non-reference outcome level
    results <- tidy(fit, conf.int = TRUE, exponentiate = TRUE) %>%
      filter(term != "(Intercept)") %>%
      transmute(
        outcome_model = outcome,
        comparison    = y.level,
        assay         = assay_cat,
        symptom_var   = ifelse(is.null(symptom), "none", symptom),
        term,
        OR    = estimate,
        lower = conf.low,
        upper = conf.high,
        p     = round(p.value, 4),
        n     = nrow(df)
      )
  }
  results
}


# Runs fit_response_model across all single assay models 
run_all_models <- function(data,
                           outcome    = "hk14_4fold",
                           assays     = c("pre_hai_hk14_cat"),
                           symptoms   = list(NULL),
                           covariates = c("age_group", "male", "index"),
                           ref_level  = NULL) {
  
  # Create every assay × symptom combination to iterate over
  combos <- expand_grid(assay_cat = assays, symptom = symptoms)
  # Fit a model for each combination and stack results into one data frame
  results <- map2_dfr(combos$assay_cat, combos$symptom, function(a, s) {
    fit_response_model(
      data = data, outcome = outcome, assay_cat = a,
      symptom = s, covariates = covariates, ref_level = ref_level

    )
  })
  results
}

# ═══════════════════════════════════════════════════════════════
# Run models -------------------------------
# ═══════════════════════════════════════════════════════════════

# Set the reference level for the 3-level response outcome
c05 <- c05 %>%
  mutate(group2 = factor(group2, levels = c("No Response", "HAI Responder", "Alternate Response")))

cat_assays  <- c("pre_hai_hk14_cat", "flha_auc_pre_cat",
                 "stalk_auc_pre_cat", "na_auc_pre_cat")
cont_assays <- c("log2_pre_hai", "log2_flha_pre", "log2_stalk_pre", "log2_na_pre")

# Primary analysis: multinomial regression with all 3 response groups, adjusting for ILI
res_2.1 <- run_all_models(
  data = c05, outcome = "group2",
  assays = cont_assays, symptoms = list("ili")
)

# binomial version
# Sensitivity: drop HAI Responders to compare Alternate Response vs No Response only
res_2.2 <- run_all_models(
  data = c05 %>% filter(group2 != "HAI Responder"),
  outcome = "group2", assays = cont_assays, symptoms = list("ili")
)

# Sensitivity: drop Alternate Response to compare HAI Responder vs No Response only
res_2.3 <- run_all_models(
  data = c05 %>% filter(group2 != "Alternate Response"),
  outcome = "group2", assays = cont_assays, symptoms = list("ili")
)

# Formats raw regression output into a publication-ready table with odds ratios, p-values, and FDR-adjusted p-values
format_reg_table <- function(data) {
  data %>%
    # Flag titer terms so FDR correction is applied separately from covariate terms
    mutate(is_titer = str_detect(term, "log2")) %>%
    group_by(is_titer) %>%
    mutate(adj.p = p.adjust(p, method = "fdr")) %>%
    ungroup() %>%
    mutate(
      # Format OR and 95% CI as a single string
      or_label = paste0(
        sprintf("%.2f", OR), " (",
        sprintf("%.2f", lower), "-",
        sprintf("%.2f", upper), ")"),
      # Format raw and adjusted p-values with 3 significant figures
      p_label = ifelse(signif(p, 3) < 0.001, "p<0.001", paste0("p=", formatC(signif(p, 3), format = "fg", digits = 3))),
      adj_p_label = case_when(
        is.na(adj.p) ~ NA_character_,
        signif(adj.p, 3) < 0.001 ~ "p<0.001",
        TRUE ~ paste0("p=", formatC(signif(adj.p, 3), format = "fg", digits = 3))
      ),
      # Clean up comparison and term names for display
      comparison = str_remove(comparison, "\nvs.*"),
      term = recode(term, "age_group15+" = "Age ≥15", "maleTRUE" = "Male",
                    "indexTRUE" = "Index case", "iliTRUE" = "ILI", "ariTRUE" = "ARI"),
      assay_label = recode(assay, "log2_pre_hai" = "HAI", "log2_flha_pre" = "FLHA",
                           "log2_stalk_pre" = "HA Stalk", "log2_na_pre" = "NA"),
      # Replace raw log2 variable names with short assay labels for titer rows
      term = ifelse(is_titer, assay_label, term)
    ) %>%
    select(comparison, assay_label, term, or_label, p_label, adj_p_label)
}

# Format each model's results into presentation-ready tables

tablesx2 = format_reg_table(res_2.2) 
tablesx3 = format_reg_table(res_2.3)






# ═══════════════════════════════════════════════════════════════
# Figure 1 -----------------------------------------
# ═══════════════════════════════════════════════════════════════

# ═══════════════════════════════════════════════════════════════
## Figure 1A - Distribution and overlap of ≥4-fold antibody responses by assay among A/H3N2-infected participants -------------------------------
# ═══════════════════════════════════════════════════════════════
# Simulate h3n2_titer_long2 from c07 (one row per antibody per episode)
h3n2_titer_long2 <- c05 %>%
  mutate(codigo2 = paste0(codigo, "_1")) %>%
  select(codigo2, hk14_4fold, flha_4fold, stalk_4fold, na_4fold) %>%
  pivot_longer(
    cols = ends_with("_4fold"),
    names_to = "antibody",
    values_to = "response"
  ) %>%
  mutate(response = as.logical(response))

# Reshape to wide: one row per episode, one logical column per antibody indicating ≥4-fold response
wide <- h3n2_titer_long2 %>%
  mutate(response = as.logical(response)) %>%
  group_by(codigo2, antibody) %>%
  summarise(response = any(response), .groups = "drop") %>%
  pivot_wider(names_from = antibody, values_from = response, values_fill = list(response = FALSE))



# Prepare data for Venn diagram, ensuring all 4 antibody columns exist
venn_df <- wide %>% rename(value = codigo2)
target_antibodies <- c("flha_4fold", "hk14_4fold", "na_4fold", "stalk_4fold")
missing <- setdiff(target_antibodies, names(venn_df))
for (m in missing) venn_df[[m]] <- FALSE
antibodies <- target_antibodies
fill_cols <- c("#377EB8", "#E41A1C", "#4DAF4A", "#984EA3")

# Figure 1A: Venn diagram showing overlap of ≥4-fold responses across the 4 antibody assays
ggvenn(
  data              = venn_df,
  columns           = antibodies,
  show_elements     = FALSE,
  show_stats        = "cp",
  show_set_totals   = "none",
  show_outside      = "always",
  show_percentage   = TRUE,
  digits            = 1,
  fill_color        = fill_cols,
  fill_alpha        = 0.4,
  stroke_color      = "grey30",
  stroke_size       = 0.6,
  set_name_color    = fill_cols,
  set_name_size     = 4,
  text_color        = "black",
  text_size         = 3.5
) +
  theme_void() +
  ggtitle(sprintf(
    "Antibody response overlap (n, %% of all episodes, N = %d)",
    nrow(venn_df)
  )) +
  theme(
    plot.title      = element_text(hjust = 0.5, size = 16),
    legend.position = "none"
  )


# ═══════════════════════════════════════════════════════════════
##  Figure 1B: UpSet plot showing combination frequencies of ≥4-fold responses ------------------
# ═══════════════════════════════════════════════════════════════

# Convert 4-fold indicators into a list-column of response labels per episode
us_h3 <- c05 %>%
  select(codigo,  hk14_4fold, flha_4fold, stalk_4fold, na_4fold) %>%
  mutate(
    hk14_4fold  = if_else(hk14_4fold, "HAI", NA),
    flha_4fold  = if_else(flha_4fold, "Full-length HA", NA),
    stalk_4fold = if_else(stalk_4fold, "HA Stalk", NA),
    na_4fold    = if_else(na_4fold, "NA", NA)
  ) %>%
  pivot_longer(cols = contains("4fold"), names_to = "antibody", values_to = "response") %>%
  group_by(codigo) %>%
  reframe(responses = list(response)) %>%
  mutate(virus = "H3N2")


# Figure 1B
us_h3_plot <- us_h3 %>%
  ggplot(aes(x = responses, y = after_stat(count / sum(count)) * 100)) +
  geom_bar(fill = "#1F78B4", alpha = 0.8) +
  geom_text(
    aes(label = paste0(round(after_stat(count / sum(count)) * 100, 2), "%")),
    stat = "count", vjust = -0.5, size = 3
  ) +
  scale_x_upset(order_by = "freq", name = "Antibody Type") +
  theme_minimal() +
  theme(
    axis.text.x = element_text(angle = 45, hjust = 1),
    plot.title = element_text(size = 14, face = "bold")
  ) +
  labs(x = "Antibody Type", y = "Percentage (%)",
       title = "H3N2: ≥4-fold Antibody Responses")
us_h3_plot

# ═══════════════════════════════════════════════════════════════
# Figure 2 Pre-exposure antibody titers by assay among HAI responders, alternate responders, and no response individuals ------------------------------------------
# ═══════════════════════════════════════════════════════════════

# ═══════════════════════════════════════════════════════════════
## Fig 2A ----------------------------------
# ═══════════════════════════════════════════════════════════════

# Reshape pre-exposure titer data from wide to long for violin plots across response groups
q4_h3 <- c05 %>%
  select(
    log2_pre_hai, log2_flha_pre, log2_stalk_pre, log2_na_pre,
    group2, season_order, age_15
  ) %>% 
  pivot_longer(
    cols = contains("pre"),
    names_to = "antibody",
    values_to = "titer"
  ) %>%
  mutate(
    # Rename assay variables to publication-friendly labels
    antibody = factor(
      antibody,
      levels = c("log2_pre_hai", "log2_flha_pre", "log2_stalk_pre", "log2_na_pre")
    ) %>%
      fct_recode(
        "HAI" = "log2_pre_hai",
        "Full-length HA" = "log2_flha_pre",
        "Stalk" = "log2_stalk_pre",
        "NA" = "log2_na_pre"
      ),
    group2 = factor(
      group2,
      levels = c("HAI Responder", "Alternate Response", "No Response")
    ),
    virus = "H3N2"
  )



# Fig 2A: Violin + jittered points showing titer distributions by response group, one panel per assay
Fig2A <- q4_h3 %>%
  mutate(group2 = factor(group2, levels = c("HAI Responder", "Alternate Response", "No Response"))) %>%
  ggplot(aes(x = group2, y = titer, fill = group2)) +
  geom_violin(alpha = 0.4, width = 0.7, linewidth = 0, position = position_dodge(width = 0.8), scale = "width") +
  geom_point(aes(colour = after_scale(fill)), alpha = 0.6, size = 1, position = position_jitterdodge(jitter.width = 0.5)) +
  # Black diamond marks the median of each group
  stat_summary(fun = median, geom = "point", shape = 18, size = 2, 
               color = "black", position = position_dodge(width = 0.8)) +
  # Kruskal-Wallis p-value comparing the three groups
  stat_compare_means(method = "kruskal.test", label = "p.format", size = 3) +
  scale_fill_brewer(palette = "Set2") +
  facet_wrap(~antibody, scales = "free") +
  theme_minimal(base_size = 14) +
  theme(
    axis.text.x        = element_blank(),
    plot.title          = element_text(face = "bold", size = 15),
    legend.position     = "bottom",
    strip.text          = element_text(size = 13, face = "bold"),
    panel.grid.minor    = element_blank(),
    panel.grid.major.x  = element_blank()
  ) +
  labs(
    x = NULL,
    y = expression("Pre-exposure Ab levels (log"[2]*" scale)"),
    fill = "Group"
  )

# ═══════════════════════════════════════════════════════════════
## Fig 2B ----------------------------------
# ═══════════════════════════════════════════════════════════════

# Prepare the titer-only regression results for forest plot display
plot_data <- res_2.1 %>%
  # Keep only the assay titer terms (drop covariates like age, sex, index)
  filter(grepl("^log2_", term)) %>%
  # Apply Benjamini-Hochberg correction for multiple comparisons. This answers
  # across all assay data/combination is there something truly statistically significant?
  mutate(adj.p = p.adjust(p, method = "BH")) %>%
  mutate(
    # Replace variable names with short assay labels
    assay_label = case_when(
      term == "log2_pre_hai"   ~ "HAI",
      term == "log2_flha_pre"  ~ "FLHA",
      term == "log2_stalk_pre" ~ "HA Stalk",
      term == "log2_na_pre"    ~ "NA"
    ),
    # Order assays for consistent y-axis positioning
    assay_label = factor(assay_label, levels = c("NA", "HA Stalk", "FLHA", "HAI")),
    # Label each panel by the multinomial comparison it represents
    comparison = factor(comparison,
                        levels = c("HAI Responder", "Alternate Response"),
                        labels = c("HAI Responder\nvs No Response", "Alternate Response\nvs No Response")
    ),
    # Flag significance for point styling
    sig = ifelse(adj.p < 0.05, "p < 0.05", "NS"),
    # Format OR (95% CI) and adjusted p-value as display strings
    or_label = paste0(
      sprintf("%.2f", OR), " (",
      sprintf("%.2f", lower), "-",
      sprintf("%.2f", upper), ")"),
    p_label = case_when(
      is.na(adj.p) ~ NA_character_,
      signif(adj.p, 3) < 0.001 ~ "p<0.001",
      TRUE ~ paste0("p=", formatC(signif(adj.p, 3), format = "fg", digits = 3))
    )
  )

# Fig 2B: Forest plot of adjusted ORs from the multinomial model, annotated with OR (95% CI) and adjusted p
Fig2B <- ggplot(plot_data, aes(x = OR, y = assay_label, color = comparison)) +
  geom_vline(xintercept = 1, linetype = "dashed", color = "grey50") +
  geom_pointrange(
    aes(xmin = lower, xmax = upper),
    position = position_dodge(width = 0.5),
    size = 0.8, linewidth = 0.8
  ) +
  geom_text(
    aes(label = paste0(round(OR, 2), " (", round(lower, 2), "-", round(upper, 2), ") ", p_label)),
    position = position_dodge(width = 0.5),
    size = 3.5, vjust = -1.1, show.legend = FALSE
  ) +
  scale_color_manual(
    values = c("HAI Responder\nvs No Response" = "#2166AC",
               "Alternate Response\nvs No Response" = "#B2182B"),
    name = "Comparison"
  ) +
  scale_x_continuous(
    trans = "log2",
    breaks = c(0.5, 0.625, 0.75, 1),
    labels = c("0.50", "0.625", "0.75", "1.00"),
    limits = c(0.35, 1.15)
  ) +
  labs(
    x = expression("OR per "*log[2]*" unit increase (95% CI)"),
    y = NULL
  ) +
  theme_minimal(base_size = 14) +
  theme(
    panel.grid.major.y  = element_blank(),
    panel.grid.minor    = element_blank(),
    legend.position     = "bottom",
    legend.box          = "vertical",
    plot.title          = element_text(face = "bold", size = 15),
    axis.text.y         = element_text(size = 13, face = "bold")
  )

# Combine panels A and B vertically into a single composite figure
fig2 <- Fig2A / Fig2B + plot_annotation(tag_levels = "A") & 
  theme(plot.tag = element_text(size = 16, face = "bold"))
ggsave("New Figure2.pdf", fig2, width = 10, height = 12, units = "in")







# Supplemental material ----------------------------------------

# ═══════════════════════════════════════════════════════════════
# Table S1 - Proportion of 4-fold rises among HAI Responders vs Alternate Responders by alternate assays -----------------------------------
# ═══════════════════════════════════════════════════════════════

c09 <- c05 %>% filter(group2 != "No Response")

response_vars <- c("flha_4fold", "stalk_4fold", "na_4fold", "alternate_4fold")

ts1 <- c09 %>%
  group_by(group2) %>%
  summarise(
    n = n(),
    across(all_of(response_vars),
           list(n = sum, pct = ~ round(sum(.x) / n() * 100, 1)),
           .names = "{.col}___{.fn}")
  ) %>%
  pivot_longer(-group2, names_to = "var_fn", values_to = "value") %>%
  separate_wider_delim(var_fn, delim = "___", names = c("variable", "fn"), too_few = "align_end") %>%
  pivot_wider(id_cols = c(group2, variable), names_from = fn, values_from = value) %>%
  mutate(
    variable = if_else(is.na(variable), "n", variable),
    pct = if_else(variable == "n", NA_character_, paste0("(", round(pct, 2), "%)"))
  ) %>%
  pivot_wider(names_from = group2, values_from = c(n, pct)) %>%
  select(variable, contains("HAI Responder"), everything()) %>%
  # Chi-squared test for each seroconversion outcome comparing the two response groups
  mutate(
    p = map_dbl(variable, ~ if (.x %in% setdiff(response_vars, "alternate_4fold"))
      chisq.test(table(c09[[.x]], c09$group2))$p.value else NA_real_)
  )

# ═══════════════════════════════════════════════════════════════
# Table S2 - Distribution of ≥4-fold antibody responses by number of overlapping assays ------
# ═══════════════════════════════════════════════════════════════

response_vars_all <- c("hk14_4fold", "flha_4fold", "stalk_4fold", "na_4fold")

ts2 <- c05 %>%
  summarise(across(all_of(response_vars_all),
                   list(n = ~ sum(.x, na.rm = TRUE),
                        pct = ~ round(sum(.x, na.rm = TRUE) / n() * 100, 2)),
                   .names = "{.col}___{.fn}")) %>%
  pivot_longer(everything(), names_to = c("variable", ".value"), names_sep = "___") %>%
  mutate(overall = paste0(n, " (", pct, "%)"))


# Table builder for tables S3 and S4 -----------------

c05 <- c05 %>%
  mutate(group2 = factor(group2, levels = c("HAI Responder", "Alternate Response", "No Response")))

c05 <- c05 %>%
  mutate(group1 = factor(group1, levels = c("HAI Responder", "HAI Non-Responder")))


# Creates a summary table comparing characteristics across groups, split by a stratifying variable (e.g., ILI status)
make_stratified_table <- function(data, 
                                  group_var  = "group1",
                                  stratify_by = "ili",
                                  group_levels = NULL,
                                  vars = NULL) {
  
  # Default set of variables to include in the table
  all_vars <- c("age_group", "male", "fever", "ili", "ari", "index",
                "cough_duration", "CT_low",
                "log2_pre_hai", "log2_flha_pre", "log2_stalk_pre", "log2_na_pre",
                "pre_hai_hk14", # just to exctract natural scale medians
                'flha_auc_pre',
                "flha_auc_pre",
                "stalk_auc_pre",
                "na_auc_pre"
  )
  
  # Override defaults if I want other specified which variables to use
  if (!is.null(vars)) all_vars <- vars
  all_vars <- setdiff(all_vars, stratify_by)
  
  # Determine the unique values of the grouping variable (e.g., seroconversion groups)
  strat_vals <- sort(unique(data[[group_var]]))
  
  if (!is.null(group_levels)) {
    strat_vals <- group_levels
  } else {
    strat_vals <- sort(unique(data[[group_var]]))
  }
  
  # Classify variables by type so each gets appropriate summary statistics and tests
  dicho_vars <- intersect(c("male", "index", "fever", "ili", "ari"), all_vars)
  cat_vars   <- intersect(c("age_group"), all_vars)
  cont_vars  <- intersect(c("cough_duration", "CT_low", "log2_pre_hai", 
                            "log2_flha_pre", "log2_stalk_pre", "log2_na_pre"), all_vars)
  
  # Inner function that builds one summary sub-table for a single group level
  build_tbl <- function(df) {
    
    # Map variable names to human-readable labels for the table output
    label_map <- c(
      impage = "Age (years)",
      log2_pre_hai = "log\u2082 HAI titer",
      log2_flha_pre = "log\u2082 FLHA titer",
      log2_stalk_pre = "log\u2082 HA Stalk titer",
      log2_na_pre = "log\u2082 NA titer",
      age_group = "Age group",
      male = "Male",
      index = "Index case",
      fever = "Fever",
      ili = "ILI",
      ari = "ARI",
      cough_duration = "Cough duration (days)",
      CT_low = "Min Ct value",
      ever_vax = "Ever vaccinated"
    )
    for (v in intersect(names(df), names(label_map))) {
      attr(df[[v]], "label") <- label_map[[v]]
    }
    
    # Tell gtsummary how to treat each variable type
    type_list <- list()
    if (length(dicho_vars) > 0) type_list <- c(type_list, list(all_of(dicho_vars) ~ "dichotomous"))
    if (length(cat_vars) > 0)   type_list <- c(type_list, list(all_of(cat_vars) ~ "categorical"))
    if (length(cont_vars) > 0)  type_list <- c(type_list, list(all_of(cont_vars) ~ "continuous"))
    
    # Use Wilcoxon rank-sum test for continuous variables (non-parametric)
    test_list <- list()
    if (length(cont_vars) > 0) test_list <- c(test_list, list(all_of(cont_vars) ~ "wilcox.test"))
    
    # Build the summary table: counts/percentages for categorical, medians/IQR for continuous, plus p-values
    suppressMessages((
      df %>%
        tbl_summary(by = grp, type = type_list,
                    statistic = list(all_dichotomous() ~ "{n} ({p}%)",
                                     all_categorical() ~ "{n} ({p}%)",
                                     all_continuous()  ~ "{median} ({p25} - {p75})"),
                    digits = list(all_continuous() ~ 1, all_dichotomous() ~ c(0, 2)),
                    missing = "no") %>%
        add_p(test = test_list) %>%
        add_stat_label(location = "row") %>%
        modify_fmt_fun(p.value ~ function(x) formatC(signif(x, 3), format = "fg")) %>%
        separate_p_footnotes()
    ))
  }
  
  # Loop over each group level, subset the data, label the stratifying variable, and build a sub-table
  tbls <- lapply(strat_vals, function(val) {
    strat_levels <- sort(unique(data[[stratify_by]]))
    # Convert TRUE/FALSE into meaningful labels (e.g., "ILI" / "Non-ILI")
    labels <- if (is.logical(data[[stratify_by]])) {
      switch(stratify_by,
             index = c("Index", "Non-Index"),
             ili   = c("ILI", "Non-ILI"),
             fever = c("Febrile", "Non-Febrile"),
             ari   = c("ARI", "Non-ARI"),
             male  = c("Male", "Female"),
             c("TRUE", "FALSE")
      )
    } else as.character(strat_levels)
    
    # Filter to this group level and recode the stratifying variable as a labeled factor
    df <- data %>%
      filter(.data[[group_var]] == val) %>%
      mutate(grp = factor(
        as.character(.data[[stratify_by]]),
        levels = c("TRUE", "FALSE"),
        labels = if (is.logical(data[[stratify_by]])) {
          switch(stratify_by,
                 index = c("Index", "Non-Index"),
                 ili   = c("ILI", "Non-ILI"),
                 fever = c("Febrile", "Non-Febrile"),
                 ari   = c("ARI", "Non-ARI"),
                 male  = c("Male", "Female"),
                 c("TRUE", "FALSE")
          )
        } else strat_levels
      )) %>%
      select(grp, all_of(all_vars))
    build_tbl(df)
  })
  
  # Merge all sub-tables side by side, with bold group names as column spanners
  spanners <- paste0("**", strat_vals, "**")
  tbl_merge(tbls, tab_spanner = spanners)
}



# ═══════════════════════════════════════════════════════════════
# table S3 - Demographic and clinical characteristics by response group, stratified by index case status -------
# ═══════════════════════════════════════════════════════════════

tableS3 = (make_stratified_table(c05, group_var =  "group2", stratify_by = "index"))

# ═══════════════════════════════════════════════════════════════
# table S4 - Demographic and clinical characteristics by response group, stratified by ILI status ---------------
# ═══════════════════════════════════════════════════════════════

tableS4 = (make_stratified_table(c05,  group_var =  "group2", stratify_by = "ili"))

# ═══════════════════════════════════════════════════════════════
# Table S5 - Multinomial Logistic Regression of pre-exposure Antibody Titers by Assay Among HAI Responders, Alternate Responders, and No Response Individuals ------------------------------
# ═══════════════════════════════════════════════════════════════

tables5 = format_reg_table(res_2.1) 

# ═══════════════════════════════════════════════════════════════
# Table S6 - Serological and Symptom Characteristics associated with a ≥4-Fold NA Rise --------------------------------
# ═══════════════════════════════════════════════════════════════

# NA 4-fold analysis - by pre existing HAI
res_NA <- run_all_models(
  data = c05,
  outcome = "na_4fold", assays = cat_assays, symptoms = list("ari")
)

tables6 = format_reg_table(res_NA) 

# ═══════════════════════════════════════════════════════════════
# Figure S1 - Participant selection and classification of HAI and alternate antibody responses. --------------------------------------
# ═══════════════════════════════════════════════════════════════


grViz("
digraph flowchart {
  graph [rankdir=TB, fontname=Helvetica, bgcolor=white]
  node  [fontname=Helvetica, fontsize=11, style=filled, fillcolor='#f0f0f0',
         shape=box, color='#333333', penwidth=1.2]
  edge  [fontname=Helvetica, fontsize=10]

  all      [label='All Participants\\n(N = 899)']
  pcr_pos  [label='PCR Positive\\n(N = 329)']
  pcr_neg  [label='PCR Negative\\n(N = 570)', fillcolor='#d9d9d9']
  hai_inv  [label='Lack of HAI Titer\\n(N = 23)', fillcolor='#d9d9d9']
  hai_v    [label='Valid HAI Titer\\n(N = 306)']

  hai_4    [label='HAI ≥4-fold?', shape=diamond, fillcolor='#fff2cc']
  hai_r    [label='HAI Responder\\n(N = 235)', fillcolor='#c6e0b4']
  hai_nr   [label='HAI Non-Responder\\n(N = 71)']

  alt_4    [label='Alternate Ab\\n≥4-fold?', shape=diamond, fillcolor='#fff2cc']
  alt_r    [label='Alternate Response\\n(N = 44)', fillcolor='#c6e0b4']
  no_r     [label='No Response\\n(N = 27)', fillcolor='#f4cccc']

  all     -> pcr_pos
  all     -> pcr_neg  [style=dashed]
  pcr_pos -> hai_inv  [style=dashed]
  pcr_pos -> hai_v
  hai_v   -> hai_4
  hai_4   -> hai_r    [label='Yes']
  hai_4   -> hai_nr   [label='No']
  hai_nr  -> alt_4
  alt_4   -> alt_r    [label='Yes']
  alt_4   -> no_r     [label='No']

  subgraph cluster_analysed {
    label='Analysed Participants'
    style=dashed; color='#666666'; fontsize=12; fontname='Helvetica-Bold'
    hai_4; hai_r; hai_nr; alt_4; alt_r; no_r
  }
}
")


# ═══════════════════════════════════════════════════════════════
# Figure S2 - Fold-change by age comparing HAI responders and Non-responders --------------------------------------
# ═══════════════════════════════════════════════════════════════

# Scatter + LOESS smoothers showing fold-change vs age, faceted by assay, colored by HAI response group
# Reshape fold-change data to long format for plotting across antibody types
rs2 <- c05 %>%
  select(codigo, group1, impage, ends_with("_fold"), -contains("comb"), hk14_foldchange) %>%
  rename(hk14_fold = hk14_foldchange) %>%
  pivot_longer(cols = contains("_fold"), names_to = "antibody_type", values_to = "fold_change") %>%
  mutate(
    antibody_type = recode(antibody_type,
                           hk14_fold = "HAI", flha_fold = "Full-length HA",
                           stalk_fold = "HA Stalk", na_fold = "NA")
  )

rs2_p <- rs2 %>%
  ggplot(aes(x = impage, y = fold_change, color = group1)) +
  geom_point() +
  geom_smooth(aes(fill = after_scale(color)), alpha = 0.3, method = "loess", formula = y ~ x) +
  scale_y_continuous(trans = "log2", breaks = c(0.0625, 1, 16, 256, 4096),
                     labels = c("0.0625", "1", "16", "256", "4096")) +
  scale_color_brewer(palette = "Set2") +
  labs(
    x = "Age",
    y = "Fold Change (Log Scale)",
    color = "HAI Responding",
    fill = "HAI Responding",
    title = "Fold Change by Age Comparing HAI Responders and Non-responders"
  ) +
  facet_wrap(~antibody_type) +
  theme_minimal() +
  theme(
    plot.title = element_text(size = 16, face = "bold"),
    legend.position = "bottom",
    strip.text = element_text(size = 12, face = "bold")
  )
rs2_p


# Fig S3 - Pre-exposure antibody titers by assay among HAI responders, alternate responders, and no response individuals - pairwise binomial analysis ----------

# Combine sensitivity analyses (each excluding one response group) and prepare for forest plot
plot_data_stratified <- bind_rows(res_2.2,   # excludes HAI Responders
                                  res_2.3) %>% # excludes Alternate Response
  # Keep only the assay titer terms
  filter(grepl("^log2_", term)) %>%
  # FDR correction across the pooled sensitivity results
  mutate(adj.p = p.adjust(p, method = "BH")) %>%
  mutate(
    assay_label = case_when(
      term == "log2_pre_hai"   ~ "HAI",
      term == "log2_flha_pre"  ~ "FLHA",
      term == "log2_stalk_pre" ~ "HA Stalk",
      term == "log2_na_pre"    ~ "NA"
    ),
    assay_label = factor(assay_label, levels = c("NA", "HA Stalk", "FLHA", "HAI")),
    sig = ifelse(adj.p < 0.05, "p < 0.05", "NS"),
    or_label = paste0(round(OR, 2), " (", round(lower, 2), "-", round(upper, 2), ")"),
    p_label = ifelse(adj.p < 0.001, "p<0.001", paste0("p=", format(round(adj.p, 3), nsmall = 3)))
  )

# Forest plot: each pairwise sensitivity comparison side by side, annotated with OR (95% CI) and adjusted p
ggplot(plot_data_stratified, aes(x = OR, y = assay_label, color = comparison)) +
  geom_vline(xintercept = 1, linetype = "dashed", color = "grey50") +
  geom_pointrange(
    aes(xmin = lower, xmax = upper),
    position = position_dodge(width = 0.5),
    size = 0.8, linewidth = 0.8
  ) +
  geom_text(
    aes(label = paste0(round(OR, 2), " (", round(lower, 2), "-", round(upper, 2), ") ", p_label)),
    position = position_dodge(width = 0.5),
    size = 3.5, vjust = -1.1, show.legend = FALSE
  ) +
  # Remap internal comparison labels to reader-friendly names in the legend
  scale_color_manual(
    values = c(
      "No Response vs HAI Responder" = "#2166AC",
      "No Response vs Alternate Response" = "#B2182B"
    ),
    breaks = c(
      "No Response vs HAI Responder",
      "No Response vs Alternate Response"
    ),
    labels = c(
      "HAI Responder vs\nNo Response",
      "Alternate Response vs\nNo Response"
    ),
    name = "Comparison"
  ) +
  scale_x_continuous(
    trans = "log2",
    breaks = c(0.5, 0.625, 0.75, 1),
    labels = c("0.50", "0.625", "0.75", "1.00"),
    limits = c(0.35, 1.15)
  ) +
  labs(
    x = expression("OR per "*log[2]*" unit increase (95% CI)"),
    y = NULL
  ) +
  theme_minimal(base_size = 14) +
  theme(
    panel.grid.major.y  = element_blank(),
    panel.grid.minor    = element_blank(),
    legend.position     = "bottom",
    legend.box          = "vertical",
    plot.title          = element_text(face = "bold", size = 15),
    axis.text.y         = element_text(size = 13, face = "bold")
  )



# Figure S4 - Linear spline analysis of HAI response probabilities by pre-exposure titer ----------------------

# This plot predicted marginal probability of classification into each immune response  
# Plots predicted probabilities from a piecewise-linear spline model across assay titers
# Each panel shows how the probability of each response group 
# changes as one assay's pre-exposure titer increases, holding 
# other covariates at reference levels (age 0-14, female, non-index, non-ILI).


plot_spline_predictions <- function(data,
                                    outcome = "group2",
                                    ref_level = "No Response",
                                    spline_knots = c(log2(10), log2(80))) {
  
  assays <- c("log2_pre_hai", "log2_flha_pre", "log2_stalk_pre", "log2_na_pre")
  assay_labels <- c(log2_pre_hai = "HAI", log2_flha_pre = "FLHA",
                    log2_stalk_pre = "HA Stalk", log2_na_pre = "NA")
  covariates <- c("age_group", "male", "index")
  symptom <- "ili"
  
  # Loop over each assay, fit a separate spline model, and generate predicted probabilities
  all_preds <- map_dfr(assays, function(a) {
    df <- data %>%
      select(all_of(c(outcome, covariates, a, symptom))) %>%
      drop_na() %>%
      mutate(!!outcome := factor(.data[[outcome]],
                                 levels = c(ref_level, setdiff(unique(.data[[outcome]]), ref_level))))
    
    # Use linear B-splines (degree=1) for piecewise-linear fit at specified knots
    knots_str <- paste(spline_knots, collapse = ", ")
    rhs <- paste(c(covariates,
                   paste0("splines::bs(", a, ", knots = c(", knots_str, "), degree = 1)"),
                   symptom), collapse = " + ")
    fml <- as.formula(paste(outcome, "~", rhs))
    fit <- multinom(fml, data = df, trace = FALSE)
    
    # Prediction grid: vary assay across its range, hold covariates at reference levels
    newdata <- tibble(
      !!sym(a)   := seq(min(df[[a]]), max(df[[a]]), length.out = 100),
      age_group  = factor("0-14", levels = levels(df$age_group)),
      male       = FALSE,
      index      = FALSE,
      ili        = FALSE
    )
    
    # Get predicted probabilities for each outcome level
    probs_mat <- predict(fit, newdata = newdata, type = "probs")
    # Handle binary case where predict returns a vector instead of a matrix
    if (is.null(dim(probs_mat))) {
      lvls <- levels(df[[outcome]])
      probs_mat <- cbind(1 - probs_mat, probs_mat)
      colnames(probs_mat) <- lvls
    }
    
    # Reshape to long format for ggplot
    as.data.frame(probs_mat) %>%
      bind_cols(newdata) %>%
      pivot_longer(cols = all_of(colnames(probs_mat)), names_to = "comparison", values_to = "probability") %>%
      mutate(assay_label = assay_labels[a], titer = .data[[a]])
  })
  
  # Plot predicted probability curves faceted by assay, colored by response group
  all_preds %>%
    mutate(assay_label = factor(assay_label, levels = c("HAI", "FLHA", "HA Stalk", "NA"))) %>%
    ggplot(aes(x = titer, y = probability, color = comparison)) +
    geom_line(linewidth = 1) +
    facet_wrap(~ assay_label, scales = "free_x") +
    scale_color_brewer(palette = "Set2") +
    theme_minimal(base_size = 14) +
    theme(
      panel.grid.major.y = element_blank(),
      panel.grid.minor   = element_blank(),
      legend.position    = "bottom",
      legend.box         = "vertical",
      strip.text         = element_text(face = "bold")
    ) +
    labs(
      x = "Pre-exposure titer (log2)",
      y = "Predicted probability",
      color = "Response group"
    )
}

# Primary: predicted probabilities for 3-level response outcome across all assays
Figs4 = plot_spline_predictions(c05, ref_level = "No Response")

# ═══════════════════════════════════════════════════════════════
# Figure S5 - Pre-exposure antibody titers by assay among HAI responders, alternate responders, and no response individuals and age quintiles --------------------
# ═══════════════════════════════════════════════════════════════

# Split participants into 5 age quintiles for age-stratified analysis
c05 <- c05 %>%
  mutate(age_q = ntile(impage, 5),
         age_q_label = paste0("Q", age_q))

# Verify quintile boundaries
c05 %>% group_by(age_q_label) %>%
  summarise(n = n(), min_age = min(impage), max_age = max(impage))

# Run multinomial regression within each age quintile (dropping age_group from covariates since quintiles handle it)
res_by_age <- map_dfr(sort(unique(c05$age_q)), function(q) {
  df_q <- c05 %>% filter(age_q == q)
  label <- unique(df_q$age_q_label)
  age_range <- paste0(round(min(df_q$impage), 1), "-", round(max(df_q$impage), 1), "y")
  run_all_models(
    data = df_q, outcome = "group2", assays = cont_assays,
    symptoms = list("ili"), covariates = c("male", "index"),
    ref_level = "No Response"
  ) %>%
    mutate(age_quintile = label, age_range = age_range)
})

# Prepare titer results for forest plot display
plot_data_age <- res_by_age %>%
  # Keep only the assay titer terms
  filter(grepl("^log2_", term)) %>%
  mutate(
    assay_label = case_when(
      term == "log2_pre_hai"   ~ "HAI",
      term == "log2_flha_pre"  ~ "FLHA",
      term == "log2_stalk_pre" ~ "HA Stalk",
      term == "log2_na_pre"    ~ "NA"
    ),
    assay_label = factor(assay_label, levels = c("HAI", "FLHA", "HA Stalk", "NA")),
    # Create display labels combining quintile number and age range
    age_label = paste0(age_quintile, " (",
                       round(as.numeric(sub("-.*", "", age_range))), "-",
                       round(as.numeric(sub(".*-", "", sub("y", "", age_range)))), "y)"),
    age_label = factor(age_label, levels = unique(age_label[order(age_quintile)])),
    # Mark significance with stars and as a shape flag
    sig_label = case_when(p < 0.001 ~ "***", p < 0.01 ~ "**", p < 0.05 ~ "*", TRUE ~ ""),
    sig = p < 0.05,
    # Clamp extreme confidence intervals for cleaner axis limits
    lower = pmax(lower, 0.1),
    upper = pmin(upper, 10)
  )

# Reusable forest plot function: one row per age quintile, colored by assay
make_forest <- function(df, title) {
  df %>%
    # Jitter assay points vertically within each quintile row to avoid overlap
    mutate(y_pos = as.numeric(age_label) + as.numeric(assay_label) * 0.12 - 0.3) %>%
    ggplot(aes(x = OR, y = y_pos, color = assay_label, shape = sig)) +
    geom_vline(xintercept = 1, linetype = "dashed", color = "grey60") +
    geom_pointrange(aes(xmin = lower, xmax = upper), size = 0.5, linewidth = 0.6) +
    geom_text(aes(x = upper + 0.15, label = sig_label),
              size = 4, show.legend = FALSE, hjust = 0) +
    scale_x_continuous(trans = "log2", breaks = c(0.25, 0.5, 1, 2, 4),
                       labels = c("0.25", "0.5", "1", "2", "4"), limits = c(0.1, 8)) +
    scale_y_continuous(breaks = seq_along(levels(df$age_label)),
                       labels = levels(df$age_label), expand = expansion(mult = 0.1)) +
    scale_color_manual(values = c("HAI" = "#2166AC", "FLHA" = "#67A9CF",
                                  "HA Stalk" = "#EF8A62", "NA" = "#B2182B")) +
    scale_shape_manual(values = c(`TRUE` = 16, `FALSE` = 1),
                       labels = c(`TRUE` = "p < 0.05", `FALSE` = "NS"),
                       name = "Significance") +
    labs(title = title, x = "OR per log2 unit (95% CI)", y = NULL, color = "Assay") +
    theme_minimal(base_size = 13) +
    theme(panel.grid.major.y = element_blank(), panel.grid.minor = element_blank(),
          legend.position = "bottom", legend.box = "horizontal",
          plot.title = element_text(face = "bold", size = 13))
}

# Create one forest plot per multinomial comparison, then stack them vertically
p1 <- plot_data_age %>% filter(comparison == "HAI Responder") %>%
  make_forest("HAI Responder vs No Response")
p2 <- plot_data_age %>% filter(comparison == "Alternate Response") %>%
  make_forest("Alternate Response vs No Response")

# Figure s5
p1 / p2 + plot_layout(guides = "collect") & theme(legend.position = "bottom")


# ═══════════════════════════════════════════════════════════════
# Figure S6 - Pre-exposure antibody titers by assay among responders, alternate responders, and no response individuals of different HAI strains -------------------------
# ═══════════════════════════════════════════════════════════════


# Define 4-fold seroconversion for the two additional H3N2 strains (SI16 and SW13)
c05 <- c05 %>%
  mutate(
    si16_4fold = si16_foldchange >= 4,
    sw13_4fold = sw13_foldchange >= 4,
    # Classify response groups using the same hierarchy as HK14: strain responder > alternate > no response
    group2_si16 = case_when(
      si16_4fold ~ "SI16 Responder",
      !si16_4fold & alternate_4fold ~ "Alternate Response",
      !si16_4fold & !alternate_4fold ~ "No Response"
    ),
    group2_sw13 = case_when(
      sw13_4fold ~ "SW13 Responder",
      !sw13_4fold & alternate_4fold ~ "Alternate Response",
      !sw13_4fold & !alternate_4fold ~ "No Response"
    )
  )

# Check group distributions for each strain definition
table(c05$group2, useNA = "ifany")
table(c05$group2_si16, useNA = "ifany")
table(c05$group2_sw13, useNA = "ifany")

# Run the same multinomial model separately for each strain's response grouping
res_hk14 <- c05 %>%
  filter(!is.na(group2)) %>%
  mutate(group2 = relevel(factor(group2), ref = "No Response")) %>%
  {run_all_models(data = ., outcome = "group2", assays = cont_assays,
                  symptoms = list("ili"), ref_level = "No Response")} %>%
  mutate(strain = "HK14")

res_sw13 <- c05 %>%
  filter(!is.na(group2_sw13)) %>%
  mutate(group2_sw13 = relevel(factor(group2_sw13), ref = "No Response")) %>%
  {run_all_models(data = ., outcome = "group2_sw13", assays = cont_assays,
                  symptoms = list("ili"), ref_level = "No Response")} %>%
  mutate(strain = "SW13")

res_si16 <- c05 %>%
  filter(!is.na(group2_si16)) %>%
  mutate(group2_si16 = relevel(factor(group2_si16), ref = "No Response")) %>%
  {run_all_models(data = ., outcome = "group2_si16", assays = cont_assays,
                  symptoms = list("ili"), ref_level = "No Response")} %>%
  mutate(strain = "SI16")

# Stack all strain results into one data frame for cross-strain comparison
res_all_strains <- bind_rows(res_hk14, res_sw13, res_si16)

# Prepare titer-only results for a forest plot comparing ORs across strains
plot_data_strain <- res_all_strains %>%
  filter(grepl("^log2_", term)) %>%
  mutate(
    assay_label = case_when(
      term == "log2_pre_hai"   ~ "HAI",
      term == "log2_flha_pre"  ~ "FLHA",
      term == "log2_stalk_pre" ~ "HA Stalk",
      term == "log2_na_pre"    ~ "NA Assay"
    ),
    assay_label = factor(assay_label, levels = c("NA Assay", "HA Stalk", "FLHA", "HAI")),
    # Collapse strain-specific comparison names into generic labels for faceting
    comp_type = case_when(
      grepl("Responder", comparison) ~ "Strain Responder\nvs No Response",
      grepl("Alternate", comparison) ~ "Alternate Response\nvs No Response"
    ),
    sig = ifelse(p < 0.05, "p < 0.05", "NS"),
    strain = factor(strain, levels = c("HK14", "SW13", "SI16"))
  )

# Add star-coded significance labels for annotation on the forest plot
plot_data_strain <- plot_data_strain %>%
  mutate(
    sig_label = case_when(
      p < 0.001 ~ "***",
      p < 0.01  ~ "**",
      p < 0.05  ~ "*",
      TRUE      ~ ""
    )
  )

# Forest plot: one panel per strain, showing ORs for each assay colored by comparison type
ggplot(plot_data_strain, aes(x = OR, y = assay_label, color = comp_type, shape = sig)) +
  # Dashed line at OR=1 (null effect)
  geom_vline(xintercept = 1, linetype = "dashed", color = "grey50") +
  # Point estimates with 95% CI whiskers, dodged so two comparisons don't overlap
  geom_pointrange(
    aes(xmin = lower, xmax = upper),
    position = position_dodge(width = 0.5),
    size = 0.7, linewidth = 0.7
  ) +
  # Significance stars placed just to the right of each CI
  geom_text(
    aes(x = upper * 1.08, label = sig_label, group = comp_type),
    position = position_dodge(width = 0.5),
    hjust = 0, size = 4, show.legend = FALSE
  ) +
  # One panel per H3N2 strain
  facet_wrap(~ strain, ncol = 3) +
  scale_color_manual(
    values = c("Strain Responder\nvs No Response" = "#2166AC",
               "Alternate Response\nvs No Response" = "#B2182B"),
    name = "Comparison"
  ) +
  scale_shape_manual(
    values = c("p < 0.05" = 16, "NS" = 1),
    name = "Significance"
  ) +
  # Log2-scaled x-axis so equal distances represent equal fold-changes in OR
  scale_x_continuous(
    trans = "log2",
    breaks = c(0.25, 0.5, 0.75, 1, 1.5, 2),
    labels = c("0.25", "0.5", "0.75", "1", "1.5", "2")
  ) +
  labs(
    x = "Odds Ratio per log2 unit increase (95% CI)",
    y = NULL
  ) +
  theme_minimal(base_size = 13) +
  theme(
    panel.grid.major.y = element_blank(),
    panel.grid.minor   = element_blank(),
    legend.position    = "bottom",
    legend.box         = "vertical",
    plot.title         = element_text(face = "bold", size = 14),
    plot.subtitle      = element_text(color = "grey40", size = 10),
    axis.text.y        = element_text(size = 12, face = "bold"),
    strip.text         = element_text(face = "bold", size = 13)
  )


# Figure S7 - Linear spline analysis of NA response probabilities by pre-exposure titer --------------------------------

# binary NA seroconversion outcome
c05 <- c05 %>%
  mutate(na_4fold_grp = ifelse(na_4fold, "Responder", "Non-Responder"))
plot_spline_predictions(c05, outcome = "na_4fold_grp", ref_level = "Non-Responder")