DZ_Script_4_CellChat_Analysis_Visualization_Final.R

# Current as of 5/22/26 11:45am
# =========================================================
# SCRIPT 4: CELLCHAT ANALYSIS AND VISUALIZATION
# =========================================================
#
# Purpose:
#   Load the aggregated CellChat outputs produced by Script 3 and
#   generate a comprehensive set of per-patient and cohort-level
#   figures and summary tables.
#
# Inputs (from Script 3):
#   cellchat_aggregated/
#     <prefix>_aggregated_TAM_TME.csv  — per-patient TAM-TME aggregated LR
#     <prefix>_aggregated_full_TME.csv — per-patient full-TME aggregated LR
#     COHORT_aggregated_TAM_TME_all_patients.csv
#     COHORT_aggregated_full_TME_all_patients.csv
#   composition_exports/
#     combined_TAM_composition_by_sample_type.csv
#
# NOTE: The per-sample raw LR CSVs in cellchat_TAM_full_TME/ and
#   cellchat_full_TME/ are not read by this script. All visualization
#   and analysis uses the patient-level aggregated tables produced by
#   Script 3's aggregate_sample_results() function.
#
# Outputs (all .png):
#   Per-patient figures (saved in outdir_figures/<prefix>/):
#     1. TAM composition bar chart — proportion of each TAM state by
#        sample_type (tumor vs background vs healthy).
#     2. Interaction dot plot — x = mean_prob_norm_all_samples (all-sample
#        normalized strength), y = consistency (proportion of samples
#        where detected), size = mean_n_cells_source, color = pathway_name.
#        Quadrant labels distinguish core, heterogeneous, and background
#        interactions using Script 3's two-metric framework.
#     3. Circular chord diagram — source-to-target interaction strengths
#        summed across LR pairs; TAM states colored with tam_colors.
#     4. Interaction heatmap — source x target matrix of summed
#        mean_prob_norm_all_samples, clustered by hierarchical clustering.
#
#   Cohort-level figures (saved in outdir_figures/cohort/):
#     5. TAM composition stacked bar — all patients side by side,
#        colored by TAM state, faceted by sample_type.
#     6. Top interactions dot plot — top N LR pairs ranked by
#        mean_prob_norm_all_samples averaged across the complete patient
#        grid (absent patients count as zero, not missing).
#     7. Interaction consistency heatmap — patients x LR pairs,
#        colored by consistency score.
#     8. Core vs heterogeneous interaction bar chart — counts per patient
#        per context, faceted by context.
#
#   Summary tables (saved in outdir_summary/):
#     9.  Per-patient TAM composition summary CSV.
#    10.  Per-patient top interactions CSV (top N per patient per context),
#         ranked by mean_prob_norm_all_samples.
#    11.  Cohort-level core interaction table (present in >= 50% of patients).
#    12.  Cohort-level heterogeneous interaction table (spatially restricted
#         but high probability in at least one patient).
#
# =========================================================
#
# Features:
# =========================================================
#
# [Feature 1] Dot plot (Figure 2) now uses mean_prob_norm_all_samples on x-axis
#   Figure 2 correctly represents the two-dimensional interaction landscape that Script 3 was
#   designed to produce. 
#
# [Feature 2] Cohort dot plot (Figure 6) uses a complete patient grid
#   Previously, the cohort-level mean was computed only over patients that
#   had a row for a given LR pair, inflating rankings of patient-specific
#   interactions. Now a full patient x LR pair grid is constructed and
#   absent interactions are assigned mean_prob_norm_all_samples = 0 before
#   averaging. Rankings now reflect true cohort-level prevalence.
#
# [Feature 3] write_top_interactions() reads patient_key from the data column
#   Previously used stringr::str_extract(basename(f), "^[^_]+") which
#   extracted only up to the first underscore and was fragile for any key
#   containing underscores. Now reads patient_key directly from the CSV's
#   patient_key column, which Script 3 guarantees to be present.
#
# [Feature 4] dplyr::cur_data() replaced with dplyr::pick(everything())
#   cur_data() was deprecated in dplyr 1.1.0 and removed in 1.2.0,
#   causing an error in write_cohort_heterogeneous(). Replaced with the
#   current API: pick(everything()).
#
# [Feature 5] Header inputs section corrected
#   Removed references to per-sample raw LR CSVs which are listed as
#   inputs but never actually read by this script.
#
# [Feature 6] Summary tables ranked by mean_prob_norm_all_samples
#   write_top_interactions() previously ranked by mean_prob_norm. Now
#   explicitly ranks by mean_prob_norm_all_samples so the ranking metric
#   matches the primary visualization axis.
#
# [Feature 7] Cohort core and heterogeneous tables use mean_prob_norm_all_samples
#   All threshold comparisons in write_cohort_core() and
#   write_cohort_heterogeneous() updated from mean_prob_norm to
#   mean_prob_norm_all_samples for consistency with the new framework.
#
# =========================================================

suppressPackageStartupMessages({
  library(tidyverse)      # Data wrangling and ggplot2 visualization.
  library(circlize)       # Circular chord diagrams for interaction plots.
  library(ComplexHeatmap) # Heatmap visualization.
  library(RColorBrewer)   # Color palettes.
  library(scales)         # Axis formatting helpers.
  library(ggrepel)        # Non-overlapping labels on dot plots.
  library(Seurat)         # For visualization with the tile plot
})

set.seed(777)
gc()
rm(list = ls())

# =========================================================
# USER SETTINGS
# =========================================================

base_outdir <- "/Users/rileyjones/Desktop/BIOINFORMATICS/Thesis/Data/scRNAseq/De_Zuani_NatComm_2024_EMTAB13530/DZ_scRNAseq"

# Input directories from Script 3.
indir_aggregated  <- file.path(base_outdir, "cellchat_aggregated")
indir_composition <- file.path(base_outdir, "composition_exports")

# Output directories.
outdir_figures <- file.path(base_outdir, "figures_script4")
outdir_summary <- file.path(base_outdir, "summary_tables_script4")

# Create output directories. recursive = TRUE is safe if they already exist.
dir.create(outdir_figures,                     showWarnings = FALSE, recursive = TRUE)
dir.create(file.path(outdir_figures, "cohort"), showWarnings = FALSE, recursive = TRUE)
dir.create(outdir_summary,                     showWarnings = FALSE, recursive = TRUE)

# ---------------------------------------------------------
# Figure settings
# ---------------------------------------------------------

fig_dpi         <- 300  # PNG resolution. 300 dpi is publication quality.
fig_width       <- 10   # Default figure width in inches.
fig_height      <- 8    # Default figure height in inches.
fig_width_wide  <- 16   # Wide format for cohort-level multi-patient figures.
fig_height_tall <- 12   # Tall format for large heatmaps.

# If TRUE, skip saving a figure when its output file already exists.
# Set FALSE to regenerate all figures from scratch.
skip_existing_figures <- FALSE

# ---------------------------------------------------------
# Analysis thresholds
# ---------------------------------------------------------

# Minimum mean_prob_norm_all_samples to label an LR pair on dot plots.
# Points below this threshold are plotted but not labeled, keeping
# the figure readable when thousands of interactions are present.
label_prob_threshold <- 0.3

# Top N LR pairs to show in cohort-level dot and consistency heatmap figures.
top_n_interactions <- 30

# Consistency threshold separating "core" interactions from "heterogeneous".
# An interaction with consistency >= this value is detected in at least
# this proportion of a patient's samples and is considered broadly present.
core_consistency_threshold <- 0.75

# Minimum mean_prob_norm_all_samples for a spatially restricted interaction
# to be classified as "heterogeneous" rather than weak/background.
# Prevents reporting very low-probability interactions as scientifically
# meaningful just because they appear in only one sample.
heterogeneous_prob_threshold <- 0.10

# Minimum fraction of cohort patients an interaction must appear in to
# qualify for the cohort-level core interaction summary table.
cohort_core_patient_fraction <- 0.50

# Number of top interactions written to the per-patient summary CSV.
top_n_per_patient <- 20


# =========================================================
# COLOR PALETTES
# =========================================================
# Defined once here and used consistently across all figures so that
# the same TAM state always maps to the same color.

# TAM state colors — seven functional programs plus Ambiguous.
# Keys match the raw TAM state names stored in TAM_label_for_analysis
# (no "TAM_" prefix). get_group_colors() strips the prefix when mapping
# cellchat_group labels that do carry it.
tam_colors <- c(
  "IFN"       = "#E41A1C",
  "INFLAM"    = "#FF7F00",
  "LA"        = "#4DAF4A",
  "ANGIO"     = "#984EA3",
  "REG"       = "#377EB8",
  "PROLIF"    = "#A65628",
  "TRM"       = "#F781BF",
  "Ambiguous" = "#999999"
)

# Sample type colors for composition figures.
sample_type_colors <- c(
  "tumor"      = "#D73027",
  "background" = "#4575B4",
  "healthy"    = "#1A9850"
)

# General cell type palette for non-TAM CellChat identity groups.
# colorRampPalette interpolates to however many unique groups are present.
celltype_palette <- colorRampPalette(
  RColorBrewer::brewer.pal(12, "Set3")
)


# =========================================================
# HELPER FUNCTIONS
# =========================================================

# ---------------------------------------------------------
# save_png()
# ---------------------------------------------------------
# Saves a ggplot object or a base-R graphics expression to a PNG file.
# When skip_existing_figures = TRUE, the function returns early without
# re-rendering if the output file already exists.
#
# Arguments:
#   outfile  — full path to the output PNG.
#   plot_obj — a ggplot2 object (use this for all ggplot figures).
#   expr     — an unevaluated base-R graphics expression wrapped in
#              quote({}) (use this for circlize and ComplexHeatmap).
#   width, height — figure dimensions in inches.
save_png <- function(outfile, plot_obj = NULL, expr = NULL,
                     width = fig_width, height = fig_height) {
  
  # Skip if the file already exists and we are not forcing a rebuild.
  if (skip_existing_figures && file.exists(outfile)) {
    message("Skipping (exists): ", basename(outfile))
    return(invisible(NULL))
  }
  
  # Create the output directory if it does not yet exist.
  dir.create(dirname(outfile), showWarnings = FALSE, recursive = TRUE)
  
  if (!is.null(plot_obj)) {
    # ggplot2 path: ggsave handles device open/close internally.
    ggplot2::ggsave(
      filename = outfile,
      plot     = plot_obj,
      width    = width,
      height   = height,
      dpi      = fig_dpi,
      bg       = "white"
    )
  } else if (!is.null(expr)) {
    # Base-R graphics path: open PNG device, evaluate the expression,
    # then close the device in a finally block so it always closes.
    #
    # IMPORTANT: eval() must use the CALLER'S environment (parent.frame()),
    # not save_png()'s own local environment. Variables like sq_mat and ht
    # are defined in the for-loop or function that calls save_png(); if
    # eval() runs in save_png()'s frame, those names are not found and R
    # throws "object 'sq_mat' not found" (or similar). parent.frame()
    # resolves to the environment where save_png() was called, which is
    # where those variables live.
    caller_env <- parent.frame()
    png(outfile,
        width  = width  * fig_dpi,
        height = height * fig_dpi,
        res    = fig_dpi,
        bg     = "white")
    tryCatch(eval(expr, envir = caller_env), finally = dev.off())
  }
  
  message("Saved: ", basename(outfile))
}

# ---------------------------------------------------------
# get_group_colors()
# ---------------------------------------------------------
# Assigns a color to each unique cell identity group in a CellChat LR table.
# TAM functional states (groups starting with "TAM_") are assigned their
# canonical color from tam_colors by stripping the "TAM_" prefix. All
# other groups (T cells, tumor, fibroblasts, etc.) receive colors from
# celltype_palette so that TAM states are visually distinct in every figure.
get_group_colors <- function(groups) {
  unique_groups <- sort(unique(as.character(groups)))
  
  # Pre-allocate with NA; filled in two passes below.
  colors <- setNames(rep(NA_character_, length(unique_groups)), unique_groups)
  
  # Pass 1: assign tam_colors to any group whose name (minus "TAM_") matches.
  for (g in unique_groups) {
    tam_key <- sub("^TAM_", "", g)       # Strip prefix: "TAM_IFN" -> "IFN".
    if (tam_key %in% names(tam_colors)) {
      colors[g] <- tam_colors[tam_key]
    }
  }
  
  # Pass 2: assign interpolated palette colors to all remaining (non-TAM) groups.
  non_tam   <- unique_groups[is.na(colors)]
  palette_n <- max(length(non_tam), 1)   # Avoid colorRampPalette(0) edge case.
  pal       <- celltype_palette(palette_n)
  colors[non_tam] <- pal[seq_along(non_tam)]
  
  colors
}

# ---------------------------------------------------------
# make_patient_dir()
# ---------------------------------------------------------
# Creates (if needed) and returns the per-patient figure subdirectory.
# All four per-patient figures for a given patient are saved here.
make_patient_dir <- function(prefix) {
  d <- file.path(outdir_figures, prefix)
  dir.create(d, showWarnings = FALSE, recursive = TRUE)
  d
}

# ---------------------------------------------------------
# make_complete_cohort_grid()
# ---------------------------------------------------------
# [Feature 2] Builds a complete patient x LR-pair grid from a cohort data.frame.
# For interactions absent from a patient's aggregated table (filtered out
# by Script 3's noise filter or simply not detected), the returned grid
# fills mean_prob_norm_all_samples = 0 and consistency = 0. This ensures
# that cohort-level averages reflect true absence rather than missing data,
# correcting the bias that inflated rankings of patient-specific interactions.
#
# Arguments:
#   cohort_df — the COHORT_aggregated_*.csv data.frame (all patients stacked).
#   lr_key_cols — character vector of columns that define a unique LR interaction.
make_complete_cohort_grid <- function(cohort_df,
                                      lr_key_cols = c("source", "target",
                                                      "ligand", "receptor",
                                                      "pathway_name")) {
  
  all_patients <- sort(unique(cohort_df$patient_key))
  
  # Build the universe of unique LR interactions detected in any patient.
  interaction_keys <- cohort_df %>%
    dplyr::distinct(dplyr::across(dplyr::all_of(lr_key_cols)))
  
  # Cross patient list with interaction universe to get the complete grid.
  complete_grid <- tidyr::crossing(patient_key = all_patients, interaction_keys)
  
  # Left join observed data; rows absent from cohort_df get NA filled as zero.
  complete_grid %>%
    dplyr::left_join(cohort_df, by = c("patient_key", lr_key_cols)) %>%
    dplyr::mutate(
      # Absent interactions are true zeroes for both probability and consistency.
      mean_prob_norm_all_samples = dplyr::coalesce(mean_prob_norm_all_samples, 0),
      mean_prob_norm_detected    = dplyr::coalesce(mean_prob_norm_detected,    0),
      mean_prob_norm             = dplyr::coalesce(mean_prob_norm,             0),
      consistency                = dplyr::coalesce(consistency,                0)
    )
}


# =========================================================
# LOAD DATA
# =========================================================

message("Loading aggregated CellChat tables...")

# Per-patient aggregated tables. These are the primary input for all
# per-patient figures and for building the summary tables.
tam_tme_files <- list.files(
  indir_aggregated,
  pattern    = "_aggregated_TAM_TME\\.csv$",
  full.names = TRUE
)
full_tme_files <- list.files(
  indir_aggregated,
  pattern    = "_aggregated_full_TME\\.csv$",
  full.names = TRUE
)

# Exclude COHORT files from the per-patient file lists.
# COHORT files are loaded separately below.
tam_tme_files  <- tam_tme_files[!grepl("^COHORT_",  basename(tam_tme_files))]
full_tme_files <- full_tme_files[!grepl("^COHORT_", basename(full_tme_files))]

# Cohort-level tables: all patients stacked into one data.frame per context.
cohort_tam_tme_file  <- file.path(indir_aggregated, "COHORT_aggregated_TAM_TME_all_patients.csv")
cohort_full_tme_file <- file.path(indir_aggregated, "COHORT_aggregated_full_TME_all_patients.csv")

cohort_tam_tme  <- if (file.exists(cohort_tam_tme_file))
  read.csv(cohort_tam_tme_file,  stringsAsFactors = FALSE) else NULL
cohort_full_tme <- if (file.exists(cohort_full_tme_file))
  read.csv(cohort_full_tme_file, stringsAsFactors = FALSE) else NULL

# TAM composition table from Script 3's incremental composition export.
composition_file <- file.path(indir_composition, "combined_TAM_composition_by_sample_type.csv")
composition_df   <- if (file.exists(composition_file))
  read.csv(composition_file, stringsAsFactors = FALSE) else NULL

# Stop early if no cohort data was found — nothing useful can be produced.
if (is.null(cohort_tam_tme) && is.null(cohort_full_tme)) {
  stop(
    "No cohort-level aggregated CellChat tables found in:\n  ", indir_aggregated,
    "\nRun Script 3 first."
  )
}

message("Loaded ", length(tam_tme_files),  " per-patient TAM-TME files.")
message("Loaded ", length(full_tme_files), " per-patient full-TME files.")
if (!is.null(cohort_tam_tme))  message("Cohort TAM-TME:  ", nrow(cohort_tam_tme),  " rows, ",
                                       length(unique(cohort_tam_tme$patient_key)),  " patients.")
if (!is.null(cohort_full_tme)) message("Cohort full-TME: ", nrow(cohort_full_tme), " rows, ",
                                       length(unique(cohort_full_tme$patient_key)), " patients.")


# =========================================================
# FIGURE 1 (per patient): TAM COMPOSITION BAR CHART
# =========================================================
# Stacked bar chart showing the proportion of each TAM functional state
# within each sample_type (tumor, background, healthy) for one patient.
#
# x-axis    : sample_type
# y-axis    : proportion of cells (0–100%)
# fill color: TAM_label_for_analysis (raw TAM state name, no "TAM_" prefix)
#
# This is an early QC check: if all bars are 100% Ambiguous, the TAM
# label transfer from Script 1 likely failed for this patient.

message("\n--- Generating per-patient TAM composition bar charts ---")

if (!is.null(composition_df)) {
  
  for (prefix in sort(unique(composition_df$patient_key))) {
    
    pdir    <- make_patient_dir(prefix)
    outfile <- file.path(pdir, paste0(prefix, "_TAM_composition_bar.png"))
    
    pat_data <- composition_df %>%
      dplyr::filter(patient_key == prefix) %>%
      dplyr::mutate(
        # Fix factor levels so bars are always ordered tumor > background > healthy.
        sample_type = factor(sample_type,
                             levels = c("tumor", "background", "healthy")),
        # Put Ambiguous last so the most informative states are at the base.
        TAM_label_for_analysis = factor(
          TAM_label_for_analysis,
          levels = c(names(tam_colors)[names(tam_colors) != "Ambiguous"], "Ambiguous")
        )
      )
    
    if (nrow(pat_data) == 0) next
    
    # Aggregate to sample_type x TAM state and compute proportions.
    pat_props <- pat_data %>%
      dplyr::group_by(sample_type, TAM_label_for_analysis) %>%
      dplyr::summarise(n_cells = sum(n_cells), .groups = "drop") %>%
      dplyr::group_by(sample_type) %>%
      dplyr::mutate(prop = n_cells / sum(n_cells)) %>%
      dplyr::ungroup()
    
    # Use only colors for TAM states that are actually present in this patient
    # so scale_fill_manual does not warn about unused levels.
    present_states <- intersect(
      levels(pat_props$TAM_label_for_analysis),
      unique(as.character(pat_props$TAM_label_for_analysis))
    )
    fill_colors <- tam_colors[present_states]
    fill_colors <- fill_colors[!is.na(fill_colors)]
    
    p <- ggplot(pat_props,
                aes(x = sample_type, y = prop,
                    fill = TAM_label_for_analysis)) +
      geom_bar(stat = "identity", position = "stack",
               width = 0.7, color = "white", linewidth = 0.3) +
      scale_fill_manual(values = fill_colors, drop = FALSE, name = "TAM State") +
      scale_y_continuous(
        labels = scales::percent_format(accuracy = 1),
        limits = c(0, 1)
      ) +
      labs(
        title    = paste0(prefix, " — TAM State Composition by Sample Type"),
        subtitle = paste0(
          "Total macrophages: ",
          format(sum(pat_data$n_cells), big.mark = ",")
        ),
        x = "Sample Type",
        y = "Proportion of Macrophages"
      ) +
      theme_classic(base_size = 13) +
      theme(
        plot.title      = element_text(face = "bold", size = 14),
        plot.subtitle   = element_text(color = "grey40", size = 11),
        legend.position = "right",
        axis.text.x     = element_text(angle = 30, hjust = 1)
      )
    
    save_png(outfile, plot_obj = p, width = 7, height = 6)
    rm(p, pat_data, pat_props)
  }
}


# =========================================================
# FIGURES 2-4 (per patient): INTERACTION DOT PLOT,
#             CIRCULAR CHORD DIAGRAM, INTERACTION HEATMAP
# =========================================================

run_per_patient_figures <- function(lr_files, context_label) {
  
  message("\n--- Generating per-patient figures: ", context_label, " ---")
  
  for (f in lr_files) {
    
    # [Feature 3] Read patient_key directly from the data column rather than
    # parsing the filename. Script 3 guarantees patient_key is present in
    # every aggregated CSV, so this is more robust than regex on filenames.
    lr_df  <- read.csv(f, stringsAsFactors = FALSE)
    
    if (nrow(lr_df) == 0) {
      message("No interactions in: ", basename(f), " (", context_label, "). Skipping.")
      next
    }
    
    prefix <- dplyr::first(lr_df$patient_key)   # [Feature 3] Read from column, not filename.
    pdir   <- make_patient_dir(prefix)
    
    # Build a named color vector for all unique identity groups in this patient.
    all_groups <- sort(unique(c(lr_df$source, lr_df$target)))
    group_cols <- get_group_colors(all_groups)
    
    # ---------------------------------------------------------
    # FIGURE 2: Interaction Dot Plot
    # ---------------------------------------------------------
    # [Feature 1] x-axis uses mean_prob_norm_all_samples (the all-sample
    # normalized strength metric from Script 3) rather than mean_prob_norm
    # (its alias). This correctly represents the two-dimensional interaction
    # landscape: x = how strong at the patient level (penalized for absent
    # samples), y = how broadly detected across the patient's samples.
    #
    # Quadrant interpretation:
    #   Top-right    : core (strong + broadly detected across samples)
    #   Bottom-right : heterogeneous (strong in ≥1 sample, spatially restricted)
    #   Top-left     : consistently detected but weak
    #   Bottom-left  : noise (already filtered by Script 3)
    
    dot_outfile <- file.path(pdir,
                             paste0(prefix, "_", context_label, "_dot_plot.png"))
    
    dot_data <- lr_df %>%
      dplyr::mutate(
        lr_pair = paste0(ligand, " \u2192 ", receptor),   # Unicode arrow for readability.
        # Label only interactions above the probability threshold to keep the
        # plot readable when thousands of LR pairs are present.
        label_txt = dplyr::if_else(
          mean_prob_norm_all_samples >= label_prob_threshold,
          paste0(lr_pair, "\n(", source, "\u2192", target, ")"),
          ""
        )
      )
    
    # Cap point size at 2000 cells so one dominant group does not visually
    # overwhelm all others; the legend uses the capped breaks.
    dot_data$size_capped <- pmin(dot_data$mean_n_cells_source, 2000)
    
    p_dot <- ggplot(dot_data,
                    aes(x     = mean_prob_norm_all_samples,  # [Feature 1]
                        y     = consistency,
                        size  = size_capped,
                        color = pathway_name)) +
      # Reference lines marking the core/heterogeneous quadrant boundaries.
      geom_hline(yintercept = core_consistency_threshold,
                 linetype = "dashed", color = "grey60", linewidth = 0.5) +
      geom_vline(xintercept = heterogeneous_prob_threshold,
                 linetype = "dashed", color = "grey60", linewidth = 0.5) +
      geom_point(alpha = 0.65) +
      ggrepel::geom_text_repel(
        aes(label = label_txt),
        size          = 2.8,
        color         = "black",
        max.overlaps  = 20,
        box.padding   = 0.3,
        segment.color = "grey50",
        segment.size  = 0.3
      ) +
      scale_size_continuous(
        name   = "Mean source\ncell count",
        range  = c(1, 8),
        breaks = c(50, 200, 500, 1000, 2000),
        labels = c("50", "200", "500", "1,000", "\u22652,000")
      ) +
      scale_x_continuous(
        name   = "Mean Normalized Probability (all samples)",   # [Feature 1]
        limits = c(0, NA),
        expand = expansion(mult = c(0.02, 0.1))
      ) +
      scale_y_continuous(
        name   = "Consistency (proportion of samples detected)",
        limits = c(0, 1),
        labels = scales::percent_format(accuracy = 1)
      ) +
      # Quadrant corner annotations for reader orientation.
      annotate("text", x = Inf, y = 1.0,
               label = "Core", hjust = 1.1, vjust = 1.2,
               color = "grey30", size = 3.5, fontface = "italic") +
      annotate("text", x = Inf, y = 0.0,
               label = "Heterogeneous", hjust = 1.1, vjust = -0.2,
               color = "grey30", size = 3.5, fontface = "italic") +
      labs(
        title = paste0(prefix, " \u2014 Interaction Landscape (", context_label, ")"),
        subtitle = paste0(
          nrow(dot_data), " LR pairs | ",
          sum(dot_data$consistency >= core_consistency_threshold), " core | ",
          sum(dot_data$consistency < core_consistency_threshold &
                dot_data$mean_prob_norm_all_samples >= heterogeneous_prob_threshold),  # [Feature 1]
          " heterogeneous"
        ),
        color = "Pathway"
      ) +
      theme_classic(base_size = 12) +
      theme(
        plot.title      = element_text(face = "bold", size = 13),
        plot.subtitle   = element_text(color = "grey40", size = 10),
        legend.position = "right",
        legend.key.size = unit(0.4, "cm")
      )
    
    save_png(dot_outfile, plot_obj = p_dot,
             width = fig_width, height = fig_height)
    rm(p_dot, dot_data)
    
    # ---------------------------------------------------------
    # FIGURE 3: Circular Chord Diagram
    # ---------------------------------------------------------
    # Directed chord diagram summarizing total interaction strength between
    # identity groups. Each chord goes FROM source TO target; chord width
    # represents the sum of mean_prob_norm_all_samples across all LR pairs
    # between that source-target pair. TAM states use their canonical colors.
    
    chord_outfile <- file.path(pdir,
                               paste0(prefix, "_", context_label, "_chord.png"))
    
    # Aggregate LR-level probabilities to source-target pairs.
    chord_data <- lr_df %>%
      dplyr::group_by(source, target) %>%
      dplyr::summarise(
        total_prob = sum(mean_prob_norm_all_samples, na.rm = TRUE),  # [Feature 1]
        n_pairs    = dplyr::n(),
        .groups    = "drop"
      ) %>%
      dplyr::filter(total_prob > 0)  # Remove pairs with no detectable interaction.
    
    if (nrow(chord_data) >= 2) {
      
      # Pivot to a source x target probability matrix.
      chord_mat <- chord_data %>%
        tidyr::pivot_wider(
          id_cols     = source,
          names_from  = target,
          values_from = total_prob,
          values_fill = 0
        ) %>%
        tibble::column_to_rownames("source") %>%
        as.matrix()
      
      # Make the matrix square: all identity groups appear as both rows and columns.
      all_g  <- sort(union(rownames(chord_mat), colnames(chord_mat)))
      sq_mat <- matrix(0,
                       nrow     = length(all_g),
                       ncol     = length(all_g),
                       dimnames = list(all_g, all_g))
      sq_mat[rownames(chord_mat), colnames(chord_mat)] <- chord_mat
      
      chord_colors <- group_cols[all_g]
      chord_colors[is.na(chord_colors)] <- "grey70"  # Fallback for unmapped groups.
      
      save_png(
        chord_outfile,
        expr = quote({
          par(mar = c(1, 1, 3, 1))   # Extra top margin for the title.
          
          circlize::chordDiagram(
            sq_mat,
            grid.col          = chord_colors,
            transparency      = 0.4,          # Partial transparency to show overlaps.
            directional       = 1,            # Show direction: source → target.
            direction.type    = c("diffHeight", "arrows"),
            link.arr.type     = "big.arrow",
            annotationTrack   = "grid",
            preAllocateTracks = list(track.height = 0.08),
            self.link         = 1             # Allow self-interactions (TAM → TAM).
          )
          
          # Rotate sector labels for readability across all arc positions.
          circlize::circos.trackPlotRegion(
            track.index = 1,
            panel.fun   = function(x, y) {
              xlim   <- circlize::get.cell.meta.data("xlim")
              ylim   <- circlize::get.cell.meta.data("ylim")
              sector <- circlize::get.cell.meta.data("sector.index")
              circlize::circos.text(
                mean(xlim), ylim[1],
                sector,
                facing     = "clockwise",
                niceFacing = TRUE,
                adj        = c(0, 0.5),
                cex        = 0.7
              )
            },
            bg.border = NA
          )
          
          title(
            paste0(prefix, " \u2014 Cell Interaction Chord (", context_label, ")"),
            cex.main = 1.1
          )
          
          circlize::circos.clear()   # Reset circlize state for the next figure.
        }),
        width  = 9,
        height = 9
      )
      
    } else {
      message("Too few source-target pairs for chord diagram: ",
              prefix, " (", context_label, ")")
    }
    
    rm(chord_data)
    
    # ---------------------------------------------------------
    # FIGURE 4: Interaction Strength Heatmap
    # ---------------------------------------------------------
    # Source (rows) x Target (columns) matrix of total summed
    # mean_prob_norm_all_samples. Cells are colored by interaction strength.
    # Rows and columns are hierarchically clustered to reveal which cell
    # type pairs have the strongest overall communication and which are silent.
    # A row annotation bar distinguishes TAM vs non-TAM source groups.
    
    heatmap_outfile <- file.path(pdir,
                                 paste0(prefix, "_", context_label, "_heatmap.png"))
    
    heatmap_data <- lr_df %>%
      dplyr::group_by(source, target) %>%
      dplyr::summarise(
        total_prob = sum(mean_prob_norm_all_samples, na.rm = TRUE),  # [Feature 1]
        .groups    = "drop"
      )
    
    if (nrow(heatmap_data) >= 2) {
      
      heatmap_mat <- heatmap_data %>%
        tidyr::pivot_wider(
          id_cols     = source,
          names_from  = target,
          values_from = total_prob,
          values_fill = 0
        ) %>%
        tibble::column_to_rownames("source") %>%
        as.matrix()
      
      # Row annotation: TAM vs Other — allows quick visual check of
      # whether TAM states are the dominant senders or receivers.
      row_ann_df <- data.frame(
        Type      = ifelse(grepl("^TAM_", rownames(heatmap_mat)), "TAM", "Other"),
        row.names = rownames(heatmap_mat)
      )
      row_ann_colors <- list(Type = c("TAM" = "#E41A1C", "Other" = "#4575B4"))
      
      row_ann <- ComplexHeatmap::rowAnnotation(
        df          = row_ann_df,
        col         = row_ann_colors,
        width       = unit(4, "mm"),
        show_legend = TRUE
      )
      
      ht <- ComplexHeatmap::Heatmap(
        heatmap_mat,
        name     = "Sum\nprob_norm",
        col      = circlize::colorRamp2(
          c(0, max(heatmap_mat) / 2, max(heatmap_mat)),
          c("white", "#FEB24C", "#BD0026")   # White → orange → red gradient.
        ),
        cluster_rows         = nrow(heatmap_mat) > 2,  # No clustering if only 2 rows.
        cluster_columns      = ncol(heatmap_mat) > 2,
        show_row_dend        = TRUE,
        show_column_dend     = TRUE,
        row_names_gp         = grid::gpar(fontsize = 9),
        column_names_gp      = grid::gpar(fontsize = 9),
        column_names_rot     = 45,
        right_annotation     = row_ann,
        column_title         = paste0(prefix, " \u2014 Interaction Heatmap (",
                                      context_label, ")"),
        column_title_gp      = grid::gpar(fontsize = 12, fontface = "bold"),
        heatmap_legend_param = list(
          title_gp  = grid::gpar(fontsize = 9),
          labels_gp = grid::gpar(fontsize = 8)
        )
      )
      
      save_png(
        heatmap_outfile,
        expr = quote({
          ComplexHeatmap::draw(ht, padding = unit(c(5, 10, 5, 5), "mm"))
        }),
        width  = fig_width,
        height = fig_height
      )
      
    } else {
      message("Too few source-target pairs for heatmap: ",
              prefix, " (", context_label, ")")
    }
    
    rm(lr_df, heatmap_data, all_groups, group_cols)
    gc()
  }
}

# Run per-patient figures for both CellChat contexts.
run_per_patient_figures(tam_tme_files,  "TAM_TME")
run_per_patient_figures(full_tme_files, "full_TME")


# =========================================================
# FIGURE 5 (cohort): TAM COMPOSITION STACKED BAR
# =========================================================
# All patients side by side on the x-axis, faceted by sample_type.
# Colored by TAM state. Allows rapid visual comparison of TAM state
# distributions across the cohort and across tissue compartments.

message("\n--- Generating cohort-level TAM composition figure ---")

if (!is.null(composition_df) && nrow(composition_df) > 0) {
  
  comp_outfile <- file.path(outdir_figures, "cohort",
                            "COHORT_TAM_composition_stacked_bar.png")
  
  comp_cohort <- composition_df %>%
    dplyr::mutate(
      sample_type = factor(sample_type,
                           levels = c("tumor", "background", "healthy")),
      TAM_label_for_analysis = factor(
        TAM_label_for_analysis,
        levels = c(names(tam_colors)[names(tam_colors) != "Ambiguous"], "Ambiguous")
      ),
      patient_key = factor(patient_key, levels = sort(unique(patient_key)))
    ) %>%
    dplyr::group_by(patient_key, sample_type, TAM_label_for_analysis) %>%
    dplyr::summarise(n_cells = sum(n_cells), .groups = "drop") %>%
    dplyr::group_by(patient_key, sample_type) %>%
    dplyr::mutate(prop = n_cells / sum(n_cells)) %>%
    dplyr::ungroup()
  
  p_cohort_comp <- ggplot(comp_cohort,
                          aes(x = patient_key, y = prop,
                              fill = TAM_label_for_analysis)) +
    geom_bar(stat = "identity", position = "stack",
             width = 0.8, color = "white", linewidth = 0.2) +
    scale_fill_manual(values = tam_colors, drop = FALSE, name = "TAM State") +
    scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
    facet_wrap(~ sample_type, ncol = 1, scales = "free_y") +
    labs(
      title = "Cohort-Level TAM State Composition",
      x     = "Patient",
      y     = "Proportion of Macrophages"
    ) +
    theme_classic(base_size = 12) +
    theme(
      plot.title       = element_text(face = "bold", size = 14),
      axis.text.x      = element_text(angle = 45, hjust = 1, size = 9),
      strip.background = element_rect(fill = "grey92", color = NA),
      strip.text       = element_text(face = "bold", size = 11),
      legend.position  = "right"
    )
  
  save_png(comp_outfile, plot_obj = p_cohort_comp,
           width = fig_width_wide, height = fig_height_tall)
  rm(p_cohort_comp, comp_cohort)
}


# =========================================================
# FIGURE 6 (cohort): TOP INTERACTIONS DOT PLOT
# =========================================================
# Top N LR pairs ranked by mean mean_prob_norm_all_samples across the
# complete cohort grid (absent patients contribute zero, not NA).
#
# x-axis : patient_key
# y-axis : LR pair (ligand → receptor, source → target)
# size   : mean_prob_norm_all_samples (interaction strength)
# color  : consistency (proportion of a patient's samples where detected)
#
# This plot identifies both cohort-wide interactions (present and strong
# across all patients) and patient-enriched interactions (strong in a few
# but absent in most).

run_cohort_dot_plot <- function(cohort_df, context_label) {
  
  if (is.null(cohort_df) || nrow(cohort_df) == 0) return(invisible(NULL))
  
  message("Generating cohort dot plot: ", context_label)
  
  outfile <- file.path(outdir_figures, "cohort",
                       paste0("COHORT_top_interactions_dot_", context_label, ".png"))
  
  # [Feature 2] Build a complete patient x LR pair grid before ranking.
  # This ensures the cohort-level mean includes zeros for absent interactions
  # rather than computing a mean over only the patients that detected it.
  complete_df <- make_complete_cohort_grid(cohort_df)
  
  # Construct a LR pair label for the y-axis.
  complete_df <- complete_df %>%
    dplyr::mutate(
      lr_pair = paste0(ligand, " \u2192 ", receptor,
                       "\n(", source, " \u2192 ", target, ")")
    )
  
  # Rank LR pairs by their mean all-sample normalized probability across
  # the complete patient grid. High-ranking pairs are both strong and common.
  top_pairs <- complete_df %>%
    dplyr::group_by(lr_pair) %>%
    dplyr::summarise(
      mean_across_patients = mean(mean_prob_norm_all_samples, na.rm = TRUE),  # [Feature 2]
      .groups = "drop"
    ) %>%
    dplyr::arrange(dplyr::desc(mean_across_patients)) %>%
    dplyr::slice_head(n = top_n_interactions) %>%
    dplyr::pull(lr_pair)
  
  plot_df <- complete_df %>%
    dplyr::filter(lr_pair %in% top_pairs) %>%
    dplyr::mutate(
      lr_pair     = factor(lr_pair, levels = rev(top_pairs)),  # Rev so top is at top of y.
      patient_key = factor(patient_key, levels = sort(unique(patient_key)))
    )
  
  p <- ggplot(plot_df,
              aes(x     = patient_key,
                  y     = lr_pair,
                  size  = mean_prob_norm_all_samples,   # [Feature 1]
                  color = consistency)) +
    geom_point(alpha = 0.8) +
    scale_size_continuous(
      name   = "All-sample norm.\nprobability",
      range  = c(0.5, 8),
      breaks = c(0.1, 0.3, 0.5, 0.75, 1.0)
    ) +
    scale_color_gradientn(
      name   = "Consistency",
      colors = c("#4575B4", "#FEE090", "#D73027"),   # Blue → yellow → red.
      limits = c(0, 1),
      labels = scales::percent_format(accuracy = 1)
    ) +
    labs(
      title    = paste0("Cohort \u2014 Top ", top_n_interactions,
                        " Interactions (", context_label, ")"),
      subtitle = paste0(
        "Size = all-sample interaction strength | ",
        "Color = proportion of samples where detected\n",
        "Absent patients shown as size 0 (complete grid)"   # [Feature 2]
      ),
      x = "Patient",
      y = "LR Pair (Source \u2192 Target)"
    ) +
    theme_classic(base_size = 11) +
    theme(
      plot.title      = element_text(face = "bold", size = 13),
      plot.subtitle   = element_text(color = "grey40", size = 9),
      axis.text.x     = element_text(angle = 45, hjust = 1, size = 9),
      axis.text.y     = element_text(size = 7),
      legend.position = "right"
    )
  
  save_png(outfile, plot_obj = p,
           width = fig_width_wide, height = fig_height_tall)
  rm(p, plot_df, top_pairs, complete_df)
}

run_cohort_dot_plot(cohort_tam_tme,  "TAM_TME")
run_cohort_dot_plot(cohort_full_tme, "full_TME")


# =========================================================
# FIGURE 7 (cohort): CONSISTENCY HEATMAP
# =========================================================
# Patients (columns) x LR pairs (rows), colored by consistency score.
# Restricted to the top N interactions by mean consistency across patients.
# Reveals which interactions are universally present vs patient-specific.
# Numeric consistency values are printed inside each cell.

run_consistency_heatmap <- function(cohort_df, context_label) {
  
  if (is.null(cohort_df) || nrow(cohort_df) == 0) return(invisible(NULL))
  
  message("Generating cohort consistency heatmap: ", context_label)
  
  outfile <- file.path(outdir_figures, "cohort",
                       paste0("COHORT_consistency_heatmap_", context_label, ".png"))
  
  # Build LR pair label and select top interactions by mean consistency.
  top_pairs <- cohort_df %>%
    dplyr::mutate(
      lr_pair = paste0(ligand, " \u2192 ", receptor,
                       " (", source, "\u2192", target, ")")
    ) %>%
    dplyr::group_by(lr_pair) %>%
    dplyr::summarise(
      mean_consistency = mean(consistency, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::arrange(dplyr::desc(mean_consistency)) %>%
    dplyr::slice_head(n = top_n_interactions) %>%
    dplyr::pull(lr_pair)
  
  # Pivot to a rows (LR pairs) x columns (patients) matrix.
  # Fill absent combinations with 0 — consistent with the Script 3 framework
  # where non-detection is a true zero, not missing data.
  hm_df <- cohort_df %>%
    dplyr::mutate(
      lr_pair = paste0(ligand, " \u2192 ", receptor,
                       " (", source, "\u2192", target, ")")
    ) %>%
    dplyr::filter(lr_pair %in% top_pairs) %>%
    dplyr::select(patient_key, lr_pair, consistency) %>%
    tidyr::pivot_wider(
      id_cols     = lr_pair,
      names_from  = patient_key,
      values_from = consistency,
      values_fill = 0   # Non-detected patient-LR combinations are truly zero.
    ) %>%
    tibble::column_to_rownames("lr_pair") %>%
    as.matrix()
  
  # Sort rows by mean consistency descending so the most consistent
  # interactions are at the top of the heatmap.
  hm_df <- hm_df[order(rowMeans(hm_df, na.rm = TRUE), decreasing = TRUE), , drop = FALSE]
  
  ht <- ComplexHeatmap::Heatmap(
    hm_df,
    name  = "Consistency",
    col   = circlize::colorRamp2(
      c(0, 0.5, 1),
      c("#4575B4", "#FFFFBF", "#D73027")   # Blue (absent) → yellow → red (universal).
    ),
    cluster_rows         = FALSE,              # Rows already ordered by mean consistency.
    cluster_columns      = ncol(hm_df) > 2,   # Cluster patients to group similar profiles.
    row_names_gp         = grid::gpar(fontsize = 7),
    column_names_gp      = grid::gpar(fontsize = 9),
    column_names_rot     = 45,
    column_title         = paste0("Cohort Consistency Heatmap (", context_label, ")"),
    column_title_gp      = grid::gpar(fontsize = 12, fontface = "bold"),
    # Print numeric consistency values inside cells for precise reading.
    cell_fun = function(j, i, x, y, width, height, fill) {
      val <- round(hm_df[i, j], 2)
      if (val > 0) {
        grid::grid.text(
          val, x, y,
          gp = grid::gpar(
            fontsize = 6,
            col      = ifelse(val > 0.6, "white", "black")  # White text on dark cells.
          )
        )
      }
    },
    heatmap_legend_param = list(
      title_gp  = grid::gpar(fontsize = 9),
      labels_gp = grid::gpar(fontsize = 8),
      at        = c(0, 0.25, 0.5, 0.75, 1.0),
      labels    = c("0%", "25%", "50%", "75%", "100%")
    )
  )
  
  # Scale figure height to the number of rows so labels are never cramped.
  save_png(
    outfile,
    expr   = quote(ComplexHeatmap::draw(ht, padding = unit(c(5, 10, 5, 5), "mm"))),
    width  = fig_width_wide,
    height = max(fig_height, nrow(hm_df) * 0.28 + 3)
  )
  
  rm(ht, hm_df, top_pairs)
}

run_consistency_heatmap(cohort_tam_tme,  "TAM_TME")
run_consistency_heatmap(cohort_full_tme, "full_TME")


# =========================================================
# FIGURE 8 (cohort): CORE VS HETEROGENEOUS INTERACTION BAR CHART
# =========================================================
# For each patient, count the number of:
#   Core interactions        : consistency >= core_consistency_threshold
#   Heterogeneous interactions: mean_prob_norm_all_samples >= threshold
#                               AND consistency < core_consistency_threshold
# Shown as a grouped bar chart, faceted by CellChat context.
# Summarizes whether each patient's TME has predominantly stable (core)
# or spatially variable (heterogeneous) interaction patterns.

run_core_het_bar <- function(cohort_tam, cohort_full) {
  
  message("Generating cohort core vs heterogeneous bar chart...")
  
  outfile <- file.path(outdir_figures, "cohort",
                       "COHORT_core_vs_heterogeneous_bar.png")
  
  # Classify each interaction row as Core, Heterogeneous, or Weak/Background.
  classify_interactions <- function(df, ctx) {
    df %>%
      dplyr::mutate(
        interaction_class = dplyr::case_when(
          consistency >= core_consistency_threshold
          ~ "Core",
          mean_prob_norm_all_samples >= heterogeneous_prob_threshold &   # [Feature 1]
            consistency < core_consistency_threshold
          ~ "Heterogeneous",
          TRUE
          ~ "Weak/Background"   # Retained by noise filter but not analytically focal.
        ),
        context = ctx
      )
  }
  
  combined <- dplyr::bind_rows(
    if (!is.null(cohort_tam))  classify_interactions(cohort_tam,  "TAM-TME")  else NULL,
    if (!is.null(cohort_full)) classify_interactions(cohort_full, "Full TME") else NULL
  )
  
  if (nrow(combined) == 0) return(invisible(NULL))
  
  # Count Core and Heterogeneous per patient per context (exclude Weak/Background).
  bar_df <- combined %>%
    dplyr::filter(interaction_class != "Weak/Background") %>%
    dplyr::count(patient_key, context, interaction_class) %>%
    dplyr::mutate(
      patient_key       = factor(patient_key, levels = sort(unique(patient_key))),
      interaction_class = factor(interaction_class, levels = c("Core", "Heterogeneous"))
    )
  
  p <- ggplot(bar_df,
              aes(x = patient_key, y = n, fill = interaction_class)) +
    geom_bar(stat = "identity", position = "dodge",
             width = 0.7, color = "white", linewidth = 0.2) +
    scale_fill_manual(
      values = c("Core" = "#2166AC", "Heterogeneous" = "#D6604D"),
      name   = "Interaction Class"
    ) +
    facet_wrap(~ context, ncol = 1, scales = "free_y") +
    labs(
      title    = "Core vs Heterogeneous Interactions per Patient",
      subtitle = paste0(
        "Core: consistency \u2265 ", scales::percent(core_consistency_threshold, accuracy = 1),
        " | Heterogeneous: all-sample prob \u2265 ", heterogeneous_prob_threshold,   # [Feature 1]
        " and consistency < ", scales::percent(core_consistency_threshold, accuracy = 1)
      ),
      x = "Patient",
      y = "Number of LR Interactions"
    ) +
    theme_classic(base_size = 12) +
    theme(
      plot.title       = element_text(face = "bold", size = 13),
      plot.subtitle    = element_text(color = "grey40", size = 9),
      axis.text.x      = element_text(angle = 45, hjust = 1, size = 9),
      strip.background = element_rect(fill = "grey92", color = NA),
      strip.text       = element_text(face = "bold", size = 11),
      legend.position  = "right"
    )
  
  save_png(outfile, plot_obj = p, width = fig_width_wide, height = 10)
  rm(p, bar_df, combined)
}

run_core_het_bar(cohort_tam_tme, cohort_full_tme)


# =========================================================
# SUMMARY TABLES
# =========================================================

message("\n--- Writing summary tables ---")

# ---------------------------------------------------------
# Table 1: Per-patient TAM composition summary
# ---------------------------------------------------------
# One row per patient + sample_type + TAM state with cell counts and
# proportions. Consolidates numerically what the bar charts show visually.

if (!is.null(composition_df)) {
  
  comp_summary <- composition_df %>%
    dplyr::group_by(patient_key, patient_id, cancer_type, sample_type,
                    TAM_label_for_analysis) %>%
    dplyr::summarise(n_cells = sum(n_cells), .groups = "drop") %>%
    dplyr::group_by(patient_key, sample_type) %>%
    dplyr::mutate(
      total_TAM_cells = sum(n_cells),
      prop_cells      = round(n_cells / total_TAM_cells, 4)
    ) %>%
    dplyr::ungroup() %>%
    dplyr::arrange(patient_key, sample_type, dplyr::desc(prop_cells))
  
  write.csv(
    comp_summary,
    file.path(outdir_summary, "TAM_composition_summary_all_patients.csv"),
    row.names = FALSE
  )
  message("Written: TAM_composition_summary_all_patients.csv")
  rm(comp_summary)
}

# ---------------------------------------------------------
# Table 2: Per-patient top interactions
# ---------------------------------------------------------
# Top N interactions per patient per context, ranked by
# mean_prob_norm_all_samples. Useful for quickly identifying the most
# prominent interactions per patient without opening the full cohort tables.

write_top_interactions <- function(lr_files, context_label) {
  
  if (length(lr_files) == 0) return(invisible(NULL))
  
  top_list <- lapply(lr_files, function(f) {
    
    df <- read.csv(f, stringsAsFactors = FALSE)
    if (nrow(df) == 0) return(NULL)
    
    # [Feature 3] Read patient_key from the data column; do not parse from filename.
    prefix <- dplyr::first(df$patient_key)
    
    df %>%
      dplyr::mutate(patient_key = prefix, context = context_label) %>%
      dplyr::arrange(dplyr::desc(mean_prob_norm_all_samples)) %>%   # [Feature 6]
      dplyr::slice_head(n = top_n_per_patient)
  })
  
  top_df <- dplyr::bind_rows(Filter(Negate(is.null), top_list))
  
  if (nrow(top_df) > 0) {
    outfile <- file.path(
      outdir_summary,
      paste0("top_", top_n_per_patient, "_interactions_per_patient_",
             context_label, ".csv")
    )
    write.csv(top_df, outfile, row.names = FALSE)
    message("Written: ", basename(outfile))
  }
}

write_top_interactions(tam_tme_files,  "TAM_TME")
write_top_interactions(full_tme_files, "full_TME")

# ---------------------------------------------------------
# Table 3: Cohort-level core interaction table
# ---------------------------------------------------------
# Interactions present (consistency > 0) in >= cohort_core_patient_fraction
# of patients. These are the most reproducible interactions in the cohort
# and the strongest candidates for follow-up functional validation.

write_cohort_core <- function(cohort_df, context_label) {
  
  if (is.null(cohort_df) || nrow(cohort_df) == 0) return(invisible(NULL))
  
  n_patients_total <- length(unique(cohort_df$patient_key))
  
  core_table <- cohort_df %>%
    dplyr::group_by(source, target, ligand, receptor, pathway_name) %>%
    dplyr::summarise(
      # Count patients where this interaction was detected (consistency > 0).
      n_patients_with_interaction = sum(consistency > 0, na.rm = TRUE),
      patient_fraction            = n_patients_with_interaction / n_patients_total,
      # Cohort-level summary statistics using the all-sample metric.  [Feature 7]
      mean_prob_norm_all_samples_cohort = mean(mean_prob_norm_all_samples, na.rm = TRUE),
      mean_consistency_cohort           = mean(consistency,                na.rm = TRUE),
      sd_prob_norm_all_samples_cohort   = sd(mean_prob_norm_all_samples,   na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::filter(patient_fraction >= cohort_core_patient_fraction) %>%
    dplyr::arrange(dplyr::desc(mean_prob_norm_all_samples_cohort))   # [Feature 7]
  
  outfile <- file.path(outdir_summary,
                       paste0("cohort_core_interactions_", context_label, ".csv"))
  write.csv(core_table, outfile, row.names = FALSE)
  message("Written: ", basename(outfile), " (", nrow(core_table), " core interactions)")
  rm(core_table)
}

write_cohort_core(cohort_tam_tme,  "TAM_TME")
write_cohort_core(cohort_full_tme, "full_TME")

# ---------------------------------------------------------
# Table 4: Cohort-level heterogeneous interaction table
# ---------------------------------------------------------
# Interactions that are strong (mean_prob_norm_all_samples >= threshold) in
# at least one patient but have low mean consistency across patients.
# These are candidate spatially restricted or patient-specific interactions
# worth reporting for intratumor heterogeneity analysis.

write_cohort_heterogeneous <- function(cohort_df, context_label) {
  
  if (is.null(cohort_df) || nrow(cohort_df) == 0) return(invisible(NULL))
  
  n_patients_total <- length(unique(cohort_df$patient_key))
  
  het_table <- cohort_df %>%
    dplyr::group_by(source, target, ligand, receptor, pathway_name) %>%
    dplyr::summarise(
      # Count patients where all-sample probability meets the threshold.
      n_patients_detected       = sum(mean_prob_norm_all_samples >= heterogeneous_prob_threshold,  # [Feature 7]
                                      na.rm = TRUE),
      patient_fraction          = n_patients_detected / n_patients_total,
      # Maximum all-sample probability in any single patient.
      max_prob_any_patient      = max(mean_prob_norm_all_samples, na.rm = TRUE),   # [Feature 7]
      mean_prob_norm_all_samples_cohort = mean(mean_prob_norm_all_samples, na.rm = TRUE),  # [Feature 7]
      mean_consistency_cohort   = mean(consistency,                na.rm = TRUE),
      sd_prob_norm_all_samples_cohort   = sd(mean_prob_norm_all_samples,   na.rm = TRUE),  # [Feature 7]
      # Top 3 patients by all-sample probability for annotation.
      # [Feature 4] dplyr::cur_data() removed (deprecated in dplyr ≥1.1.0, errors in ≥1.2.0).
      # Replaced with dplyr::pick(everything()) which is the current API.
      top_patients = paste(
        dplyr::pick(everything()) %>%
          dplyr::arrange(dplyr::desc(mean_prob_norm_all_samples)) %>%   # [Feature 4, Feature 7]
          dplyr::slice_head(n = 3) %>%
          dplyr::pull(patient_key),
        collapse = "; "
      ),
      .groups = "drop"
    ) %>%
    # Heterogeneous: strong in at least one patient, but not core across the cohort.
    dplyr::filter(
      max_prob_any_patient  >= heterogeneous_prob_threshold &    # [Feature 7]
        patient_fraction    <  cohort_core_patient_fraction
    ) %>%
    dplyr::arrange(dplyr::desc(max_prob_any_patient))   # [Feature 7]
  
  outfile <- file.path(outdir_summary,
                       paste0("cohort_heterogeneous_interactions_",
                              context_label, ".csv"))
  write.csv(het_table, outfile, row.names = FALSE)
  message("Written: ", basename(outfile),
          " (", nrow(het_table), " heterogeneous interactions)")
  rm(het_table)
}

write_cohort_heterogeneous(cohort_tam_tme,  "TAM_TME")
write_cohort_heterogeneous(cohort_full_tme, "full_TME")

gc()
message("\nScript 4 complete.")
message("Figures saved to:       ", outdir_figures)
message("Summary tables saved to: ", outdir_summary)