mk_spatial_tile <- function(dat, square = TRUE) {
  if (nrow(dat) == 0) {
    dat <- tibble::tibble(
      Location_Center_X = numeric(),
      Location_Center_Y = numeric(),
      draw = character()
    )
  }
  tile <- ggplot(dat, aes(x = Location_Center_X, y = Location_Center_Y)) +
    geom_point(data = filter(dat, draw == "other"), colour = "#39FF14", size = 0.10) +
    geom_point(data = filter(dat, draw == "pdl1"), colour = "#FFFF00", size = 0.16) +
    geom_point(data = filter(dat, draw == "m2"), colour = "#FF00FF", size = 0.16) +
    scale_x_continuous(labels = NULL, guide = "none", expand = c(0, 0)) +
    scale_y_reverse(labels = NULL, guide = "none", expand = c(0, 0)) +
    theme_void(base_size = 8) +
    theme(
      plot.background = element_rect(fill = "white", colour = NA),
      panel.background = element_rect(fill = "black", colour = NA),
      legend.position = "none",
      plot.margin = margin(
        if (square) 1 else 0, if (square) 1 else 0,
        if (square) 1 else 0, if (square) 1 else 0,
        unit = "pt"
      )
    )
  if (square) tile <- tile + theme(aspect.ratio = 1)
  tile
}

roi <- read_csv(
  file.path(results_dir, "sccomp_long_tme_immune_roi.csv"),
  col_select = c(PatientID, BiopsyPhase, ImageID),
  show_col_types = FALSE
) |>
  filter(PatientID == "NT193") |>
  distinct() |>
  mutate(BiopsyPhase = factor(BiopsyPhase, phase_levels)) |>
  arrange(BiopsyPhase, ImageID) |>
  mutate(i = row_number(), .by = BiopsyPhase)

n_phase <- roi |>
  summarise(n = n(), .by = BiopsyPhase) |>
  mutate(x0 = cumsum(n + 2) - (n + 2))
roi <- roi |>
  left_join(n_phase, by = "BiopsyPhase") |>
  mutate(x = i + x0)

cells_nt <- read_fst(file.path(post_dir, "nt193-cells.fst")) |>
  as_tibble()
dx <- 0.46
pad <- 0.62
xmin <- min(roi$x) - pad
xmax <- max(roi$x) + pad
fig_w <- 12
fig_h <- 5.4
target_in <- (2 * dx) / (xmax - xmin) * fig_w
headroom <- 1.9
tile_h <- target_in * headroom / (fig_h - target_in)

phase <- roi |>
  summarise(x = mean(x), .by = BiopsyPhase) |>
  mutate(
    y = tile_h + 0.62,
    label = case_match(
      as.character(BiopsyPhase),
      "Baseline" ~ "Baseline\n10 ROIs",
      "On-treatment" ~ "On-treatment\n10 ROIs",
      "Post-treatment" ~ "Post-treatment\n3 ROIs"
    )
  )
patient <- tibble::tibble(
  x = mean(phase$x), y = phase$y[1] + 0.95, label = "Donor"
)
edges <- bind_rows(
  phase |> mutate(x, y, xend = patient$x, yend = patient$y),
  roi |> mutate(
    x,
    y = tile_h,
    xend = phase$x[match(BiopsyPhase, phase$BiopsyPhase)],
    yend = phase$y[1]
  )
)
phase_cols <- c(
  Baseline = "#1B9E77",
  `On-treatment` = "#D95F02",
  `Post-treatment` = "#7570B3"
)

tree <- ggplot() +
  geom_segment(
    data = edges, aes(x = x, y = y, xend = xend, yend = yend),
    colour = "grey70", linewidth = 0.4
  ) +
  geom_label(
    data = phase,
    aes(x, y, label = label, fill = BiopsyPhase),
    colour = "white", fontface = "bold", size = 3, label.size = 0
  ) +
  geom_label(
    data = patient, aes(x, y, label = label),
    fill = "grey15", colour = "white", fontface = "bold", size = 3.6, label.size = 0
  ) +
  scale_fill_manual(values = phase_cols, guide = "none") +
  coord_cartesian(
    xlim = c(xmin, xmax),
    ylim = c(-0.04, patient$y + 0.4),
    clip = "off"
  ) +
  labs(title = "Patient NT193") +
  theme_void(base_size = 11) +
  theme(
    plot.title = element_text(face = "bold", hjust = 0.5, size = 12),
    plot.margin = margin(6, 6, 6, 6)
  )

for (i in seq_len(nrow(roi))) {
  tree <- tree + annotation_custom(
    ggplotGrob(mk_spatial_tile(filter(cells_nt, ImageID == roi$ImageID[i]), square = FALSE)),
    xmin = roi$x[i] - dx,
    xmax = roi$x[i] + dx,
    ymin = 0,
    ymax = tile_h
  )
}
tree

library(readr)
library(dplyr)

long <- read_csv(file.path(results_dir, "sccomp_long_tme_immune_roi.csv"),
                 show_col_types = FALSE) |>
  rename(ROI_id = sample) |>
  mutate(
    count = as.integer(count),
    BiopsyPhase = factor(BiopsyPhase, c("Baseline", "On-treatment", "Post-treatment")),
    Arm = factor(Arm, c("C", "C&I")),
    pCR = factor(pCR, c("RD", "pCR"))
  )

long |>
  select(ROI_id, cell_group, count, PatientID, BiopsyPhase, Arm, pCR) |>
  gt::gt_preview() |>
  gt::tab_caption("Long-format sccomp input (1,740 ROIs x 15 immune phenotypes)")

library(sccomp)

fit <- sccomp_estimate(
  long,
  formula_composition = ~ Arm * BiopsyPhase * pCR + (1 | PatientID),
  formula_variability = ~ 1,
  sample = "ROI_id",
  cell_group = "cell_group",
  abundance = "count",
  inference_method = "hmc",
  cores = parallel::detectCores(),
  verbose = FALSE
) |>
  sccomp_test()

fit |>
  sccomp_proportional_fold_change(
    formula_composition = ~ Arm * BiopsyPhase * pCR,
    from = "C&I:Post-treatment:RD",  # chemoimmunotherapy, surgery, residual disease
    to = "C&I:Post-treatment:pCR"    # same arm and visit, pathological complete response
  ) |>
  gt::gt_preview() |>
  gt::tab_caption(
    "Proportional fold change: C&I post-treatment, pCR versus residual disease"
  )

fit |>
  filter(is.finite(c_ess_bulk)) |>
  slice_min(c_ess_bulk, n = 6, with_ties = FALSE) |>
  gt::gt_preview() |>
  gt::tab_caption(
    "Composition parameters with the smallest bulk effective sample size"
  )

contrasts <- c(
  marginal_CI_post = paste(
    "`pCRpCR`",                                  # pCR versus RD in chemotherapy at baseline
    "`ArmC&I:pCRpCR`",                           # additional pCR versus RD gap in the C&I arm
    "`BiopsyPhasePost-treatment:pCRpCR`",        # additional pCR versus RD gap at post-treatment
    "`ArmC&I:BiopsyPhasePost-treatment:pCRpCR`", # additional pCR versus RD gap for C&I at post-treatment
    sep = " + "
  ),
  did_CI_post = paste(
    "`BiopsyPhasePost-treatment:pCRpCR`",        # change from baseline in the pCR contrast
    "`ArmC&I:BiopsyPhasePost-treatment:pCRpCR`", # further change from baseline in the C&I arm
    sep = " + "
  )
)

res <- sccomp_test(fit, contrasts = contrasts)

hmc_did |>
  filter(contrast_family == "did") |>
  arrange(c_FDR) |>
  gt::gt_preview() |>
  gt::tab_caption("Change from baseline (strongest by FDR)")

# One row per arm, timepoint, and response group.
pop <- tibble::tribble(
  ~Arm,  ~BiopsyPhase,     ~pCR,
  "C",   "Baseline",       "RD",
  "C",   "Baseline",       "pCR",
  "C",   "On-treatment",   "RD",
  "C",   "On-treatment",   "pCR",
  "C",   "Post-treatment", "RD",
  "C",   "Post-treatment", "pCR",
  "C&I", "Baseline",       "RD",
  "C&I", "Baseline",       "pCR",
  "C&I", "On-treatment",   "RD",
  "C&I", "On-treatment",   "pCR",
  "C&I", "Post-treatment", "RD",
  "C&I", "Post-treatment", "pCR"
) |>
  mutate(
    Arm = factor(Arm, c("C", "C&I")),
    BiopsyPhase = factor(BiopsyPhase, c("Baseline", "On-treatment", "Post-treatment")),
    pCR = factor(pCR, c("RD", "pCR")),
    ROI_id = paste0("pred_", row_number()),
    PatientID = "NT031"
  )

gt::gt_preview(pop) |>
  gt::tab_caption("Population grid passed to sccomp_predict()")

pred_pop <- sccomp_predict(fit, new_data = pop)

pred_pop |>
  gt::gt_preview() |>
  gt::tab_caption("Population predicted proportions")

# Per-patient grid: one row per patient and timepoint (uses the patient random intercept)
pat <- long |>
  distinct(PatientID, BiopsyPhase, Arm, pCR) |>
  mutate(ROI_id = paste0(PatientID, "_", BiopsyPhase))

pat |>
  gt::gt_preview() |>
  gt::tab_caption("Per-patient prediction grid (624 rows)")

pred_pat <- sccomp_predict(fit, new_data = pat)

pred_pat

# Plot shift. Within each arm and cell type, move pCR and RD to the mean of
# their baseline predictions, on the posterior draws.
draws <- sccomp_predict(
  fit, new_data = pop, number_of_draws = 500, summary_instead_of_draws = FALSE
) |>
  as_tibble() |>
  mutate(across(c(Arm, BiopsyPhase, pCR), as.character))

# Shift pCR and RD to their average baseline within each Arm x cell_group, per draw
shifts <- draws |>
  filter(BiopsyPhase == "Baseline") |>
  select(Arm, cell_group, .draw, pCR, proportion) |>
  pivot_wider(names_from = pCR, values_from = proportion) |>
  mutate(shift_pCR = (pCR + RD) / 2 - pCR,
         shift_RD  = (pCR + RD) / 2 - RD) |>
  select(Arm, cell_group, .draw, shift_pCR, shift_RD)

aligned_summary <- draws |>
  left_join(shifts, by = c("Arm", "cell_group", ".draw")) |>
  mutate(aligned_prop = proportion + if_else(pCR == "pCR", shift_pCR, shift_RD)) |>
  summarise(
    aligned_mean  = mean(aligned_prop),
    aligned_lower = quantile(aligned_prop, 0.025),
    aligned_upper = quantile(aligned_prop, 0.975),
    .by = c(Arm, BiopsyPhase, pCR, cell_group)
  )

cell_order <- c(
  "CD8^+T", "CD8^+GZMB^+T", "CD8^+TCF1^+T", "CD8^+PD1^+T_{Ex}",
  "CD4^+TCF1^+T", "CD4^+PD1^+T", "Treg", "CD20^+B", "CD79a^+Plasma",
  "CD56^+NK", "M2 Mac", "DCs", "Neutrophils", "PD-L1^+APCs", "PD-L1^+IDO^+APCs"
)
cols_response <- c(RD = "#377EB8", pCR = "#E41A1C")

pred_main <- aligned_summary |>
  mutate(
    BiopsyPhase = factor(BiopsyPhase, phase_levels),
    Arm = factor(Arm, arm_levels),
    pCR = factor(pCR, c("RD", "pCR")),
    cell_group = factor(cell_group, levels = cell_order),
    # A few intervals sit just below zero after the baseline shift.
    aligned_mean = pmax(aligned_mean, 0),
    aligned_lower = pmax(aligned_lower, 0)
  )

# Asterisk just above the ribbon where the DiD FDR is below 0.05.
did_sig_fig1 <- hmc_did |>
  filter(contrast_family == "did", is.finite(c_FDR), c_FDR < 0.05) |>
  distinct(cell_group, Arm, BiopsyPhase) |>
  mutate(cell_group = factor(cell_group, levels = cell_order)) |>
  inner_join(
    pred_main |>
      summarise(ymax = max(aligned_upper), .by = c(cell_group, Arm, BiopsyPhase)),
    by = c("cell_group", "Arm", "BiopsyPhase")
  )

ggplot(pred_main, aes(x = BiopsyPhase, y = aligned_mean, colour = pCR, group = pCR)) +
  geom_ribbon(
    aes(ymin = aligned_lower, ymax = aligned_upper, fill = pCR),
    alpha = 0.10, colour = NA
  ) +
  geom_line(linewidth = 0.75) +
  geom_point(size = 1.2) +
  geom_text(
    data = did_sig_fig1,
    aes(x = BiopsyPhase, y = ymax, label = "*"),
    inherit.aes = FALSE,
    size = 4.5, colour = "black", fontface = "bold", vjust = -0.2
  ) +
  facet_grid(cell_group ~ Arm, scales = "free_y") +
  scale_y_continuous(labels = scales::label_number(accuracy = 0.01)) +
  coord_cartesian(clip = "off") +
  scale_colour_manual(values = cols_response, name = "Response") +
  scale_fill_manual(values = cols_response, name = "Response") +
  labs(
    title = "Baseline-aligned immune compositional trajectories",
    x = NULL,
    y = "Baseline-aligned predicted proportion"
  ) +
  theme_pub() +
  theme(strip.text.y = element_text(size = 7))

did_hm <- hmc_did |>
  filter(contrast_family == "did") |>
  mutate(
    cell_group = factor(cell_group, levels = rev(cell_order)),
    BiopsyPhase = factor(BiopsyPhase, levels = c("On-treatment", "Post-treatment")),
    Arm = factor(Arm, arm_levels)
  ) |>
  filter(!is.na(cell_group)) |>
  mutate(
    sig = is.finite(c_FDR) & c_FDR < 0.05,
    facet_lab = factor(
      paste(Arm, BiopsyPhase, sep = " · "),
      levels = c(
        "C · On-treatment", "C · Post-treatment",
        "C&I · On-treatment", "C&I · Post-treatment"
      )
    )
  )

lim <- max(abs(did_hm$c_effect), na.rm = TRUE)

ggplot(did_hm, aes(x = facet_lab, y = cell_group, fill = c_effect)) +
  geom_tile(colour = "grey90", linewidth = 0.3) +
  geom_tile(
    data = filter(did_hm, sig),
    colour = "black", linewidth = 0.7, fill = NA
  ) +
  geom_text(
    aes(label = if_else(sig, sprintf("%.2f", c_effect), "")),
    size = 2.6
  ) +
  scale_fill_gradient2(
    low = cols_response[["RD"]], mid = "white", high = cols_response[["pCR"]],
    midpoint = 0, limits = c(-lim, lim), name = "center_log_ratio"
  ) +
  labs(
    title = "Response-associated change from baseline (difference-in-differences)",
    x = NULL, y = NULL
  ) +
  theme_pub() +
  theme(
    axis.text.x = element_text(angle = 30, hjust = 1),
    panel.grid = element_blank()
  )

m2_volcano <- bind_rows(
  hmc_marginal_pcr |>
    filter(cell_group == "M2 Mac") |>
    mutate(Arm, BiopsyPhase, c_effect, c_FDR, family = "Marginal pCR vs RD"),
  hmc_did |>
    filter(contrast_family == "did", cell_group == "M2 Mac") |>
    mutate(Arm, BiopsyPhase, c_effect, c_FDR, family = "DiD (change from T0)")
) |>
  mutate(label = paste(Arm, BiopsyPhase, sep = " · "))

ggplot(m2_volcano, aes(x = c_effect, y = -log10(pmax(c_FDR, 1e-12)))) +
  geom_vline(xintercept = 0, linetype = 2, colour = "grey70") +
  geom_hline(yintercept = -log10(0.05), linetype = 2, colour = "grey70") +
  geom_point(aes(colour = family, shape = BiopsyPhase), size = 2.8) +
  geom_text(aes(label = label), size = 2.3, vjust = -0.9, colour = "grey25") +
  scale_colour_manual(values = c(
    "Marginal pCR vs RD" = "#999999",
    "DiD (change from T0)" = "#d73027"
  )) +
  labs(
    title = "M2 Mac: visit contrast and change from baseline",
    x = "center_log_ratio",
    y = expression(-log[10](FDR)),
    colour = NULL, shape = NULL
  ) +
  theme_pub()

pred_pat <- hmc_pred_patient |>
  filter(cell_group == "M2 Mac") |>
  mutate(cell_group = factor(cell_group))

pcr_cols <- c("RD" = "#4575b4", "pCR" = "#d73027")

ggplot(pred_pat, aes(x = BiopsyPhase, y = proportion_mean, group = PatientID, colour = pCR)) +
  geom_line(linewidth = 0.4, alpha = 0.5) +
  facet_grid(cell_group ~ Arm + pCR, scales = "free_y") +
  scale_colour_manual(values = pcr_cols, name = "Response") +
  labs(
    title = "Between-patient variation in predicted M2 Mac proportion",
    x = NULL,
    y = "Predicted proportion"
  ) +
  theme_pub() +
  theme(legend.position = "bottom", strip.text = element_text(size = 8))

cores <- read_csv(file.path(post_dir, "spatial-cores.csv"), show_col_types = FALSE)
cells <- read_fst(file.path(post_dir, "spatial-cells.fst")) |> as_tibble()

mk_roi <- function(image_id) {
  if (is.na(image_id)) return(patchwork::plot_spacer())
  mk_spatial_tile(filter(cells, ImageID == image_id))
}

mk_row <- function(pid) {
  image_ids <- cores |>
    filter(PatientID == pid) |>
    arrange(roi) |>
    pull(ImageID)
  ids <- c(image_ids, rep(NA_character_, 3L - length(image_ids)))
  patchwork::wrap_plots(
    ggplot() +
      annotate("text", x = 0.5, y = 0.5, label = pid, fontface = "bold", size = 2.8) +
      scale_x_continuous(limits = c(0, 1), expand = c(0, 0)) +
      scale_y_continuous(limits = c(0, 1), expand = c(0, 0)) +
      theme_void(),
    patchwork::wrap_plots(lapply(ids, mk_roi), nrow = 1, ncol = 3),
    ncol = 1,
    heights = c(0.16, 1)
  )
}

mk_block <- function(response, title) {
  one_column <- function(which_side) {
    pids <- cores |>
      filter(pCR == response, side == which_side) |>
      arrange(patient) |>
      distinct(PatientID) |>
      pull(PatientID)
    patchwork::wrap_plots(lapply(pids, mk_row), ncol = 1)
  }
  patchwork::wrap_plots(one_column("left"), one_column("right"), ncol = 2) +
    patchwork::plot_annotation(
      title = title,
      theme = theme(
        plot.title = element_text(face = "bold", size = 12, hjust = 0.5, margin = margin(b = 4))
      )
    )
...