DZ_Script_1_Preprocess_Samples_Final.R

# Most up-to-date as of 3pm 5/15/26
# =========================================================
# SCRIPT 1: PER-PATIENT AND HEALTHY DONOR PROCESSING
# =========================================================
#
# Purpose:
#   This script processes each cancer patient and healthy donor
#   independently (one at a time) to keep RAM usage manageable,
#   then saves the objects needed by downstream scripts.
#
#   For each patient/donor it produces:
#     1. A full annotated Seurat object with TAM labels.
#     2. A tumor-only Seurat object (cancer patients only).
#     3. A background-only Seurat object (cancer patients only).
#     4. A macrophage/TAM-only Seurat object.
#
# Pipeline overview (per patient/donor):
#   Step 1  - Read raw per-sample MTX files.
#   Step 2  - Run DoubletFinder per sample to remove doublets.
#   Step 3  - Merge singlet-filtered samples into one object.
#   Step 4  - Final post-merge QC filtering.
#   Step 5  - SCTransform normalization, PCA, Harmony batch correction,
#             graph-based clustering, and UMAP embedding.
#   Step 6  - Broad marker-based cell type scoring (Immune, Epithelial,
#             Endothelial, Fibroblast, Mast, etc) + SingleR fine annotation
#             using the HLCA centroid reference.
#   Step 7  - Subset macrophage-like cells; score each cell for seven
#             TAM programs based upon Ma Black Qian Trends Immuno 2022
#             (IFN, INFLAM, LA, ANGIO, REG, PROLIF, TRM)
#             and assign a confident TAM_subset label.
#   Step 8  - Transfer TAM labels back to the full object; build
#             cellchat_group column for Script 2.
#   Step 9  - Save final RDS outputs and clean up memory.
#
# Cohort-level integration (Script 1.5) is handled separately.
#
# Key design principles:
#   - skip_existing_outputs = TRUE means already-completed patients
#     are silently skipped, making reruns safe and fast.
#   - DoubletFinder is run per capture library BEFORE merging, which
#     is the biologically correct level for doublet detection.
#   - All intermediate SCT/PCA/graph outputs from DoubletFinder are
#     deleted before returning each sample, so the merged object
#     starts cleanly with only raw RNA counts as these are not needed by DoubletFinder.
#
# =========================================================

suppressPackageStartupMessages({
  library(Seurat)          # Single-cell analysis framework (v5).
  library(Matrix)          # Sparse matrix support for count data.
  library(tidyverse)       # Data wrangling and plotting utilities.
  library(harmony)         # Batch correction across samples/patients.
  library(sctransform)     # Regularized NB normalization (SCTransform).
  library(glmGamPoi)       # Fast GLM backend used by SCTransform.
  library(SingleR)         # Reference-based cell type annotation.
  library(SingleCellExperiment) # SCE container required by SingleR.
  library(future)          # Parallelism control (set to sequential here).
  library(DoubletFinder)   # Doublet detection per capture library.
})

# Run sequentially. Harmony and SCTransform internally use multiple
# threads but future::plan("sequential") prevents nested parallelism
# that can exhaust RAM on large patient objects.
plan("sequential")

# Set future globals limit high enough for large patients (e.g. P20 ~110k cells).
options(future.globals.maxSize = 128 * 1024^3)

set.seed(777)
gc()

# =========================================================
# USER SETTINGS
# =========================================================
# All paths and tunable parameters are collected here so that
# nothing else in the script needs to be edited for a new run.

# Directory containing per-sample MTX files named as:
#   <sample_id>-matrix.mtx
#   <sample_id>-features.tsv
#   <sample_id>-barcodes.tsv
data_dir <- "/Users/rileyjones/Desktop/BIOINFORMATICS/Thesis/Data/scRNAseq/De_Zuani_NatComm_2024_EMTAB13530/DZ_scRNAseq/All_Samples/"

# HLCA centroid reference RDS for SingleR annotation.
# Must contain named slots: $centroids (gene x cell-type matrix),
# $fine_labels (character vector), and optionally $broad_map.
centroid_rds <- "/Users/rileyjones/Desktop/BIOINFORMATICS/Thesis/Data/scRNAseq/Azimuth_HLCA_Ref/hlca_centroids_all_cells.rds"

# Root output directory. Sub-directories are created automatically.
base_outdir <- "/Users/rileyjones/Desktop/BIOINFORMATICS/Thesis/Data/scRNAseq/De_Zuani_NatComm_2024_EMTAB13530/DZ_scRNAseq"

# Output sub-directories.
outdir_full_annot  <- file.path(base_outdir, "full_patient_annotated")
outdir_tumor       <- file.path(base_outdir, "tumor_only")
outdir_background  <- file.path(base_outdir, "background_only")
outdir_macro       <- file.path(base_outdir, "macrophage_only")
outdir_qc          <- file.path(base_outdir, "qc_tables")

dir.create(outdir_full_annot,  showWarnings = FALSE, recursive = TRUE)
dir.create(outdir_tumor,       showWarnings = FALSE, recursive = TRUE)
dir.create(outdir_background,  showWarnings = FALSE, recursive = TRUE)
dir.create(outdir_macro,       showWarnings = FALSE, recursive = TRUE)
dir.create(outdir_qc,          showWarnings = FALSE, recursive = TRUE)

# If TRUE, patients/donors whose final output RDS files already exist
# are silently skipped. Set to TRUE for normal runs and reruns after
# partial failures. Set to FALSE only if you want to force-reprocess
# every patient from scratch.
skip_existing_outputs <- TRUE

# Whether to write per-sample DoubletFinder QC and final cell count
# summary CSVs at the end of the script.
save_doubletfinder_summary_csv      <- TRUE
save_final_cell_count_summary_csv   <- TRUE

# ---------------------------------------------------------
# QC thresholds
# ---------------------------------------------------------
# These are aligned closely to the De Zuani et al. (2024) paper.
# Cells failing any threshold are removed before downstream analysis.
min_features    <- 180      # Minimum detected genes per cell.
max_features    <- 6000     # Maximum detected genes (removes likely doublets).
min_counts      <- 400      # Minimum total UMI counts per cell.
max_counts      <- 100000   # Maximum total UMI counts per cell.
max_percent_mt  <- 20       # Maximum mitochondrial read fraction (%).

# ---------------------------------------------------------
# Analysis settings
# ---------------------------------------------------------
# TAM margin cutoff: the difference between the top and second-best
# TAM program score must exceed this value for a cell to receive a
# confident (non-Ambiguous) label. 0.075 is intentionally conservative
# to reduce false-positive TAM subset assignments.
tam_margin_cutoff   <- 0.075

# Broad label margin cutoff: same concept applied to the five broad
# cell-type categories (Immune, Epithelial, Endothelial, Fibroblast, Mast).
broad_margin_cutoff <- 0.05

# Number of PCs to use for FindNeighbors, Harmony, UMAP, and clustering.
npcs_use    <- 20

# Whether to run UMAP. Set to FALSE for test runs to save time.
run_umap    <- TRUE

# ---------------------------------------------------------
# DoubletFinder settings
# ---------------------------------------------------------
run_doubletfinder               <- TRUE   # Master switch for doublet detection.
doublet_rate                    <- 0.06   # Expected doublet rate (~6% for 10x & author recommendations).
doubletfinder_pcs               <- 1:20   # PCs passed to DoubletFinder.
doubletfinder_pN                <- 0.25   # Proportion of artificial doublets to generate.
doubletfinder_use_param_sweep   <- TRUE   # If TRUE, use paramSweep to estimate optimal pK.
doubletfinder_pK_default        <- 0.09   # Fallback pK if paramSweep fails or is disabled.
doubletfinder_min_cells         <- 100    # Skip DoubletFinder for samples below this cell count.


# =========================================================
# PATIENT METADATA
# =========================================================
# Sample IDs are auto-detected from the data directory by pattern
# matching (P<number>_<T|B><number>). The patient_meta table provides
# clinical annotations that are stamped onto every cell.

# Detect all patient sample IDs from the data directory.
sample_ID <- list.files(data_dir, full.names = FALSE) %>%
  stringr::str_extract("^P\\d+_[TB]\\d+") %>%
  unique() %>%
  .[!is.na(.)] %>%
  sort()

# Clinical metadata for all 24 cancer patients.
# Columns:
#   patient_key    - Short key used for file naming (P1, P2, ...).
#   patient_id     - Full patient label (somewhat redundant).
#   dataset        - Study identifier (all DZ for De Zuani et al.).
#   sex            - Biological sex (M/F).
#   cancer_type    - LC (unspecified or adenosquamous), LUSC (squamous), LUAD (adenocarcinoma).
#   stage          - Pathological T stage at resection.
#   smoking_status - Smoking history (current/ ex (any kind of previous smoking)/ no or never smoker).
#   age            - Age at surgery (years, stored as character for flexibility).
#   treatment      - Treatment status at time of sample (all untreated here).
patient_meta <- tibble::tribble(
  ~patient_key, ~patient_id,   ~dataset, ~sex, ~cancer_type, ~stage, ~smoking_status, ~age, ~treatment,
  "P1",  "Patient 1",  "DZ", "F", "LC",   "T2", "ex",      "73", "untreated",
  "P2",  "Patient 2",  "DZ", "F", "LUSC", "T2", "ex",      "82", "untreated",
  "P3",  "Patient 3",  "DZ", "F", "LUSC", "T4", "ex",      "77", "untreated",
  "P4",  "Patient 4",  "DZ", "M", "LUSC", "NA", "ex",      "69", "untreated",
  "P5",  "Patient 5",  "DZ", "F", "LUAD", "T2", "no",      "72", "untreated",
  "P6",  "Patient 6",  "DZ", "F", "LUAD", "T1", "current", "69", "untreated",
  "P7",  "Patient 7",  "DZ", "M", "LUAD", "T1", "no",      "58", "untreated",
  "P8",  "Patient 8",  "DZ", "M", "LUSC", "T2", "ex",      "75", "untreated",
  "P9",  "Patient 9",  "DZ", "M", "LUAD", "T2", "ex",      "79", "untreated",
  "P10", "Patient 10", "DZ", "F", "LUAD", "T3", "no",      "72", "untreated",
  "P11", "Patient 11", "DZ", "M", "LUSC", "T3", "ex",      "81", "untreated",
  "P12", "Patient 12", "DZ", "F", "LC",   "T1", "ex",      "63", "untreated",
  "P13", "Patient 13", "DZ", "F", "LUAD", "T1", "ex",      "59", "untreated",
  "P14", "Patient 14", "DZ", "M", "LUAD", "T4", "ex",      "81", "untreated",
  "P15", "Patient 15", "DZ", "F", "LUAD", "T3", "current", "73", "untreated",
  "P16", "Patient 16", "DZ", "F", "LUAD", "T2", "current", "87", "untreated",
  "P17", "Patient 17", "DZ", "M", "LC",   "T2", "ex",      "65", "untreated",
  "P18", "Patient 18", "DZ", "M", "LUSC", "T2", "current", "77", "untreated",
  "P19", "Patient 19", "DZ", "M", "LUSC", "T3", "ex",      "78", "untreated",
  "P20", "Patient 20", "DZ", "M", "LUSC", "T4", "current", "72", "untreated",
  "P21", "Patient 21", "DZ", "M", "LUAD", "T1", "ex",      "68", "untreated",
  "P22", "Patient 22", "DZ", "M", "LUAD", "T3", "ex",      "86", "untreated",
  "P23", "Patient 23", "DZ", "F", "LC",   "T4", "ex",      "67", "untreated",
  "P24", "Patient 24", "DZ", "F", "LUAD", "T3", "ex",      "52", "untreated"
)

# Join sample-level metadata (auto-detected IDs) with clinical metadata.
# sample_class: T = tumor, B = background.
# sample_type:  "tumor" or "background" — used throughout for subsetting.
sample_meta <- tibble(sample_id = sample_ID) %>%
  mutate(
    patient_key   = stringr::str_extract(sample_id, "^P\\d+"),
    sample_class  = stringr::str_match(sample_id, "^P\\d+_([TB])\\d+$")[, 2],
    sample_number = as.integer(stringr::str_match(sample_id, "^P\\d+_[TB](\\d+)$")[, 2]),
    sample_type   = dplyr::if_else(sample_class == "T", "tumor", "background")
  ) %>%
  left_join(patient_meta, by = "patient_key")

# Warn if any detected sample IDs could not be matched to patient_meta.
if (any(is.na(sample_meta$patient_id))) {
  warning("Some sample IDs did not match patient_meta:")
  print(sample_meta %>% filter(is.na(patient_id)))
}

# Build named lookup vectors for fast metadata stamping inside
# read_one_patient_sample(). Each vector maps sample_id -> metadata value.
patient_ID_map     <- setNames(sample_meta$patient_id,      sample_meta$sample_id)
patient_key_map    <- setNames(sample_meta$patient_key,     sample_meta$sample_id)
dataset_map        <- setNames(sample_meta$dataset,         sample_meta$sample_id)
sex_map            <- setNames(sample_meta$sex,             sample_meta$sample_id)
cancer_type_map    <- setNames(sample_meta$cancer_type,     sample_meta$sample_id)
stage_map          <- setNames(sample_meta$stage,           sample_meta$sample_id)
smoking_map        <- setNames(sample_meta$smoking_status,  sample_meta$sample_id)
age_map            <- setNames(sample_meta$age,             sample_meta$sample_id)
treatment_map      <- setNames(sample_meta$treatment,       sample_meta$sample_id)
sample_type_map    <- setNames(sample_meta$sample_type,     sample_meta$sample_id)
sample_number_map  <- setNames(sample_meta$sample_number,   sample_meta$sample_id)


# =========================================================
# MARKER PANELS
# =========================================================

# ---------------------------------------------------------
# Broad cell-type markers
# ---------------------------------------------------------
# Used by assign_broad_labels() to score every cell for five major
# compartments before SingleR fine annotation. The top-scoring
# compartment with a margin >= broad_margin_cutoff is assigned as
# the confident broad label. Cells below the margin threshold receive
# "Ambiguous" and fall through to SingleR as the primary annotation.
broad_marker_list <- list(
  Immune             = c("PTPRC", "LST1", "TYROBP", "CD3D", "NKG7", "MS4A1", "FCER1A"),
  Epithelial_Tumor   = c("EPCAM", "KRT8", "KRT18", "KRT19", "KRT17", "TACSTD2", "MUC1"),
  Endothelial        = c("PECAM1", "VWF", "KDR", "EMCN", "RAMP2"),
  Fibroblast_Stromal = c("COL1A1", "COL1A2", "DCN", "LUM", "COL3A1", "TAGLN"),
  Mast               = c("TPSAB1", "TPSB2", "CPA3", "KIT", "HDC")
)

# ---------------------------------------------------------
# TAM program gene signatures (expanded)
# ---------------------------------------------------------
# Seven macrophage polarization programs derived from 
# Ma Black Qian Trends Immuno 2022. Each program
# is represented by a curated gene list. label_tams_expanded() computes
# the mean log-normalized expression of each program's genes per cell
# and assigns the highest-scoring program as the TAM_subset. The
# TAM_margin (top score minus second-best score) must exceed
# tam_margin_cutoff for the label to be considered confident.
#
# Programs:
#   IFN    - Interferon-stimulated macrophages (antiviral/antitumor).
#   INFLAM - Inflammatory/classically activated macrophages.
#   LA     - Lipid-associated macrophages (TREM2+/APOE+).
#   ANGIO  - Angiogenic/SPP1+ pro-tumoral macrophages.
#   REG    - Regulatory/alternatively activated macrophages.
#   PROLIF - Proliferating macrophages.
#   TRM    - Tissue-resident macrophages (alveolar-like, LYVE1+/FABP4+).
tam_gene_list_expanded <- list(
  IFN    = c("ISG15","IFIT1","IFIT2","IFIT3","IFITM1","IFITM3","CXCL9","CXCL10","CXCL11","IDO1","CD274","IRF1","IRF7","STAT1","IL4I1","TNFSF10","LAMP3"),
  INFLAM = c("IL1B","IL1RN","CCL2","CCL3","CCL3L1","CCL4","CCL20","CXCL1","CXCL2","CXCL3","CXCL5","CXCL8","INHBA","S100A8","S100A9","NFKBIA"),
  LA     = c("ACP5","APOE","APOC1","C1QA","C1QB","C1QC","FABP5","GPNMB","LGALS3","LIPA","LPL","CD36","CTSB","CTSD","TREM2","SPP1","CHI3L1","F13A1"),
  ANGIO  = c("VEGFA","SPP1","FN1","VCAN","THBS1","FLT1","AREG","CCL20","CD163","CEBPB","FCN1","PPARG","SLC2A1","IL1B","IL1RN","OLR1","TIMP1","SERPINB2","BNIP3","CLEC5A"),
  REG    = c("ARG1","MRC1","CX3CR1","CD274","CD40","CD80","CD86","ICOSLG","IL10","TGFB2","LGALS9","HLA-DRA","HLA-DRB1","HLA-DQA1","HLA-DQB1","ITGA4","CHIT1"),
  PROLIF = c("MKI67","CDK1","CDC45","STMN1","TOP2A","TYMS","RRM2","CCNA2","TUBA1B","TUBB","HMGN2","HMGB1"),
  TRM    = c("LYVE1","HES1","FOLR2","MARCO","FABP4","CCL18","FBP1","LGALS3","MCEMP1","MRC1","MSR1","PPARG","RBP4","VSIG4","CD5L","MS4A7","SLC40A1","VCAM1","SEPP1","PLTP")
)


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

# Accumulates DoubletFinder results per sample.
# Written to CSV at end of script if save_doubletfinder_summary_csv = TRUE.
doubletfinder_summary_list <- list()

# ---------------------------------------------------------
# recode_smoking_binary()
# ---------------------------------------------------------
# Converts free-text smoking history values into a two-level binary:
#   "nonsmoker"  - never smoked
#   "smoker"     - ever smoked (current or former)
# This is used for cohort-level smoking comparisons in Script 3.
recode_smoking_binary <- function(x) {
  x <- tolower(trimws(as.character(x)))
  out <- rep(NA_character_, length(x))
  out[x %in% c("never","never smoker","never_smoker","nonsmoker",
               "non-smoker","non smoker","no")] <- "nonsmoker"
  out[x %in% c("former","former smoker","former_smoker","ex","ex-smoker",
               "ex smoker","current","current smoker","current_smoker",
               "smoker","ever","ever smoker")] <- "smoker"
  out
}

# ---------------------------------------------------------
# find_existing_file()
# ---------------------------------------------------------
# Returns the first path in `paths` that exists on disk, or
# NA_character_ if none exist. Used to handle naming variants for
# donor MTX files (e.g. -matrix.mtx vs -matrix.tsv).
find_existing_file <- function(paths) {
  hit <- paths[file.exists(paths)]
  if (length(hit) == 0) return(NA_character_)
  hit[1]
}

# ---------------------------------------------------------
# read_one_patient_sample()
# ---------------------------------------------------------
# Reads one cancer patient sample from MTX format and attaches all
# clinical metadata as cell-level columns. Cell barcodes are prefixed
# with the sample ID to ensure uniqueness after merging.
read_one_patient_sample <- function(sid) {
  counts <- ReadMtx(
    mtx      = file.path(data_dir, paste0(sid, "-matrix.mtx")),
    features = file.path(data_dir, paste0(sid, "-features.tsv")),
    cells    = file.path(data_dir, paste0(sid, "-barcodes.tsv"))
  )
  # Prefix barcodes with sample ID to avoid collisions after merge().
  colnames(counts) <- paste0(sid, "_", colnames(counts))
  obj <- CreateSeuratObject(counts = counts, project = sid)
  # Stamp clinical metadata onto every cell using the lookup vectors
  # built from patient_meta above.
  obj$sample_id      <- sid
  obj$patient_key    <- patient_key_map[[sid]]
  obj$patient_id     <- patient_ID_map[[sid]]
  obj$cancer_type    <- cancer_type_map[[sid]]
  obj$dataset        <- dataset_map[[sid]]
  obj$stage          <- stage_map[[sid]]
  obj$age            <- age_map[[sid]]
  obj$treatment      <- treatment_map[[sid]]
  obj$sex            <- sex_map[[sid]]
  obj$smoking_status <- smoking_map[[sid]]
  obj$sample_type    <- sample_type_map[[sid]]
  obj$sample_number  <- sample_number_map[[sid]]
  obj
}

# ---------------------------------------------------------
# get_donor_sample_ids()
# ---------------------------------------------------------
# Auto-detects healthy donor sample IDs by pattern-matching filenames
# in data_dir for the D<number>_<number> naming convention used by
# the De Zuani healthy donor samples.
get_donor_sample_ids <- function(data_dir) {
  list.files(data_dir, full.names = FALSE) %>%
    str_extract("^D\\d+_\\d+") %>%
    unique() %>%
    .[!is.na(.)] %>%
    sort()
}

# ---------------------------------------------------------
# read_one_donor_sample()
# ---------------------------------------------------------
# Reads one healthy donor sample. Handles two filename variants for
# the matrix file (-matrix.mtx or -matrix.tsv) and for the features
# file (-features.tsv or -genes.tsv), since donor data may have been
# exported with slightly different naming conventions.
read_one_donor_sample <- function(sid) {
  matrix_file <- find_existing_file(c(
    file.path(data_dir, paste0(sid, "-matrix.mtx")),
    file.path(data_dir, paste0(sid, "-matrix.tsv"))
  ))
  features_file <- find_existing_file(c(
    file.path(data_dir, paste0(sid, "-features.tsv")),
    file.path(data_dir, paste0(sid, "-genes.tsv"))
  ))
  barcodes_file <- find_existing_file(
    file.path(data_dir, paste0(sid, "-barcodes.tsv"))
  )
  if (any(is.na(c(matrix_file, features_file, barcodes_file)))) {
    warning("Missing matrix/features/barcodes file for donor sample: ", sid)
    return(NULL)
  }
  counts <- ReadMtx(mtx = matrix_file, features = features_file, cells = barcodes_file)
  colnames(counts) <- paste0(sid, "_", colnames(counts))
  donor_key    <- str_extract(sid, "^D\\d+")
  donor_number <- str_match(sid, "^D\\d+_(\\d+)$")[, 2]
  obj <- CreateSeuratObject(counts = counts, project = sid)
  # Healthy donors get fixed metadata; clinical fields unavailable
  # for donors are set to NA rather than left unset.
  obj$sample_id      <- sid
  obj$patient_key    <- donor_key
  obj$patient_id     <- paste("Donor", str_remove(donor_key, "^D"))
  obj$cancer_type    <- "Healthy"
  obj$dataset        <- "DZ_Healthy"
  obj$stage          <- NA_character_
  obj$age            <- NA_character_
  obj$treatment      <- "none"
  obj$sex            <- NA_character_
  obj$smoking_status <- NA_character_
  obj$sample_type    <- "healthy"
  obj$sample_number  <- suppressWarnings(as.integer(donor_number))
  obj
}

# ---------------------------------------------------------
# run_basic_qc()
# ---------------------------------------------------------
# Applies per-cell QC filters using the global thresholds defined
# in USER SETTINGS. Called twice per patient:
#   1. Before DoubletFinder (inside run_doubletfinder_one_sample).
#   2. After merge() on the combined patient object.
# The second call catches any low-quality cells that may have slipped
# through at the sample level due to sample-specific background noise.
run_basic_qc <- function(obj) {
  DefaultAssay(obj) <- "RNA"
  # Compute mitochondrial fraction if not already present.
  # MT- genes are standard human mitochondrial gene prefix.
  if (!("percent.mt" %in% colnames(obj@meta.data))) {
    obj[["percent.mt"]] <- PercentageFeatureSet(obj, pattern = "^MT-", assay = "RNA")
  }
  subset(
    obj,
    subset = nFeature_RNA >= min_features &
      nFeature_RNA <= max_features &
      nCount_RNA   >= min_counts   &
      nCount_RNA   <= max_counts   &
      percent.mt   <  max_percent_mt
  )
}

# ---------------------------------------------------------
# run_doubletfinder_one_sample()
# ---------------------------------------------------------
# Detects and removes doublets for a SINGLE capture library using
# DoubletFinder. Running per sample (before merging) is biologically
# correct because doublets arise within individual capture events,
# not across samples.
#
# Workflow:
#   1. QC filter the sample.
#   2. Temporarily run SCTransform + PCA + graph-based clustering
#      (these are only needed by DoubletFinder and are deleted after).
#   3. Estimate optimal pK via paramSweep (or use a default).
#   4. Estimate expected doublet count, adjusted for homotypic doublets.
#   5. Call DoubletFinder.
#   6. Store classifications in metadata; filter to singlets.
#   7. Delete temporary SCT/PCA/graph objects before returning.
#
# All steps are wrapped in tryCatch so a failure for one sample
# does not abort the entire patient pipeline.
run_doubletfinder_one_sample <- function(obj, sample_id) {
  
  df_start_time <- Sys.time()
  
  # --- Early exit: DoubletFinder disabled globally ---
  if (!run_doubletfinder) {
    log_msg("DoubletFinder disabled by run_doubletfinder = FALSE for sample: ", sample_id)
    obj$DoubletFinder_class <- "Not_run"
    obj$DoubletFinder_pANN  <- NA_real_
    obj$DoubletFinder_nExp  <- NA_integer_
    obj$DoubletFinder_pK    <- NA_real_
    return(obj)
  }
  '
  # --- Early exit: package not installed ---
  if (!requireNamespace("DoubletFinder", quietly = TRUE)) {
    warning("DoubletFinder not installed. Skipping doublet detection for ", sample_id)
    obj$DoubletFinder_class <- "Not_run_package_missing"
    obj$DoubletFinder_pANN  <- NA_real_
    obj$DoubletFinder_nExp  <- NA_integer_
    obj$DoubletFinder_pK    <- NA_real_
    doubletfinder_summary_list[[sample_id]] <<- data.frame(
      sample_id        = sample_id,
      n_cells_after_QC = ncol(obj),
      doublet_rate     = doublet_rate,
      pN               = NA_real_,
      pK               = NA_real_,
      nExp_unadjusted  = NA_integer_,
      homotypic_prop   = NA_real_,
      nExp_adjusted    = NA_integer_,
      n_singlet        = ncol(obj),
      n_doublet        = NA_integer_,
      status           = "Not_run_package_missing",
      stringsAsFactors = FALSE
    )
    return(obj)
  }
  '
  log_msg("DoubletFinder START for sample: ", sample_id)
  
  # Apply QC filter before DoubletFinder. DoubletFinder performs poorly
  # on low-quality cells that can distort the simulated doublet distribution.
  obj <- run_basic_qc(obj)
  
  if (ncol(obj) == 0) {
    warning("No cells remain after QC for ", sample_id)
    return(NULL)
  }
  
  # --- Early exit: sample too small for reliable doublet detection ---
  if (ncol(obj) < doubletfinder_min_cells) {
    warning(
      "Skipping DoubletFinder for ", sample_id,
      " because fewer than ", doubletfinder_min_cells,
      " cells remain after QC."
    )
    obj$DoubletFinder_class <- "Not_run_low_cell_count"
    obj$DoubletFinder_pANN  <- NA_real_
    obj$DoubletFinder_nExp  <- NA_integer_
    obj$DoubletFinder_pK    <- NA_real_
    doubletfinder_summary_list[[sample_id]] <<- data.frame(
      sample_id        = sample_id,
      n_cells_after_QC = ncol(obj),
      doublet_rate     = doublet_rate,
      pN               = NA_real_,
      pK               = NA_real_,
      nExp_unadjusted  = NA_integer_,
      homotypic_prop   = NA_real_,
      nExp_adjusted    = NA_integer_,
      n_singlet        = ncol(obj),
      n_doublet        = NA_integer_,
      status           = "Not_run_low_cell_count",
      stringsAsFactors = FALSE
    )
    return(obj)
  }
  
  # ---------------------------------------------------------
  # Temporary preprocessing for DoubletFinder
  # ---------------------------------------------------------
  # DoubletFinder requires a PCA embedding and cluster assignments.
  # We use SCTransform here for consistency with the main pipeline,
  # but conserve.memory = TRUE and return.only.var.genes = TRUE keep
  # it lightweight. All SCT/PCA/graph outputs are deleted at the end
  # of this function — the returned object contains only RNA counts
  # and DoubletFinder metadata.
  log_msg("DoubletFinder preprocessing START for sample: ", sample_id)
  
  obj <- SCTransform(
    obj,
    method                = "glmGamPoi",
    vars.to.regress       = "percent.mt",
    conserve.memory       = TRUE,
    return.only.var.genes = TRUE,
    verbose               = FALSE
  )
  
  obj <- RunPCA(obj, npcs = max(doubletfinder_pcs), verbose = FALSE)
  
  obj <- FindNeighbors(obj, reduction = "pca", dims = doubletfinder_pcs, verbose = FALSE)
  
  # Broader resolution set here for FindClusters()
  obj <- FindClusters(obj, resolution = 0.5, verbose = FALSE) 
  
  log_msg("DoubletFinder preprocessing DONE for sample: ", sample_id)
  
  # ---------------------------------------------------------
  # Resolve DoubletFinder function names across package versions
  # ---------------------------------------------------------
  # DoubletFinder changed its exported function names between versions
  # (doubletFinder_v3 -> doubletFinder, paramSweep_v3 -> paramSweep).
  # We check for both names so the script works with either version.
  df_ns <- asNamespace("DoubletFinder")
  
  doubletFinder_fun <- if (exists("doubletFinder", envir = df_ns)) {
    get("doubletFinder", envir = df_ns)
  } else if (exists("doubletFinder_v3", envir = df_ns)) {
    get("doubletFinder_v3", envir = df_ns)
  } else {
    stop("Could not find doubletFinder or doubletFinder_v3 in DoubletFinder namespace.")
  }
  
  paramSweep_fun <- if (exists("paramSweep", envir = df_ns)) {
    get("paramSweep", envir = df_ns)
  } else if (exists("paramSweep_v3", envir = df_ns)) {
    get("paramSweep_v3", envir = df_ns)
  } else {
    NULL
  }
  
  summarizeSweep_fun  <- get("summarizeSweep",  envir = df_ns)
  find.pK_fun         <- get("find.pK",         envir = df_ns)
  modelHomotypic_fun  <- get("modelHomotypic",   envir = df_ns)
  
  # ---------------------------------------------------------
  # Estimate optimal pK via paramSweep
  # ---------------------------------------------------------
  # pK is the PC neighbourhood size used by DoubletFinder to define
  # the KNN graph for artificial doublet scoring. The optimal value
  # is estimated by testing a range of pK values and selecting the one
  # that maximises the BCmetric (bimodality coefficient). If paramSweep
  # fails for any reason, we fall back to doubletfinder_pK_default.
  best_pK <- doubletfinder_pK_default
  
  if (doubletfinder_use_param_sweep && !is.null(paramSweep_fun)) {
    best_pK <- tryCatch({
      log_msg("DoubletFinder pK sweep START for sample: ", sample_id)
      sweep_res   <- paramSweep_fun(obj, PCs = doubletfinder_pcs, sct = TRUE)
      log_msg("DoubletFinder pK sweep DONE for sample: ", sample_id)
      log_msg("DoubletFinder pK summary START for sample: ", sample_id)
      sweep_stats <- summarizeSweep_fun(sweep_res, GT = FALSE)
      pK_table    <- as.data.frame(find.pK_fun(sweep_stats))
      pK_values   <- as.numeric(as.character(unlist(pK_table[["pK"]])))
      bc_values   <- as.numeric(as.character(unlist(pK_table[["BCmetric"]])))
      valid       <- is.finite(pK_values) & is.finite(bc_values)
      if (sum(valid) == 0) {
        warning("No valid pK/BCmetric values found for ", sample_id, ". Using default pK.")
        doubletfinder_pK_default
      } else {
        chosen_pK <- pK_values[valid][which.max(bc_values[valid])]
        log_msg("DoubletFinder selected pK for ", sample_id, ": ", chosen_pK)
        chosen_pK
      }
    }, error = function(e) {
      warning(
        "DoubletFinder pK sweep failed for ", sample_id,
        ". Using default pK = ", doubletfinder_pK_default,
        ". Error: ", e$message
      )
      doubletfinder_pK_default
    })
  } else {
    log_msg("DoubletFinder pK sweep skipped for ", sample_id, "; using default pK=", best_pK)
  }
  
  # ---------------------------------------------------------
  # Estimate expected doublets
  # ---------------------------------------------------------
  # nExp_poi: expected number of doublets given the assumed doublet_rate.
  # nExp_poi_adj: adjusted downward by the estimated homotypic proportion,
  #   because homotypic doublets (two cells of the same type) are
  #   transcriptionally similar and harder for DoubletFinder to detect.
  # Using the adjusted value avoids over-filtering singlets.
  nExp_poi <- round(doublet_rate * ncol(obj))
  
  if ("seurat_clusters" %in% colnames(obj@meta.data)) {
    obj$seurat_clusters <- as.factor(as.character(obj@meta.data[["seurat_clusters"]]))
  }
  
  homotypic_prop <- tryCatch({
    modelHomotypic_fun(as.factor(as.character(obj@meta.data[["seurat_clusters"]])))
  }, error = function(e) {
    warning(
      "Homotypic doublet estimate failed for ", sample_id,
      ". Using unadjusted expected doublet count. Error: ", e$message
    )
    0
  })
  
  nExp_poi_adj <- max(round(nExp_poi * (1 - homotypic_prop)), 1)
  
  # ---------------------------------------------------------
  # Force DoubletFinder parameters to simple scalar values
  # ---------------------------------------------------------
  # The pK sweep may return a value derived from a tibble or data.frame,
  # which can cause downstream errors such as:
  #   Error in xtfrm.data.frame(x) : cannot xtfrm data frames
  # Explicitly coercing to numeric scalar [1] prevents this.
  best_pK                <- as.numeric(best_pK)[1]
  doubletfinder_pN_use   <- as.numeric(doubletfinder_pN)[1]
  nExp_poi_adj           <- as.integer(nExp_poi_adj)[1]
  
  if (!is.finite(best_pK)) {
    warning("best_pK was not finite for ", sample_id, ". Using default pK.")
    best_pK <- as.numeric(doubletfinder_pK_default)[1]
  }
  if (!is.finite(doubletfinder_pN_use)) {
    warning("doubletfinder_pN was not finite for ", sample_id, ". Using pN = 0.25.")
    doubletfinder_pN_use <- 0.25
  }
  if (!is.finite(nExp_poi_adj) || nExp_poi_adj < 1) {
    warning("nExp_poi_adj was invalid for ", sample_id, ". Using nExp = 1.")
    nExp_poi_adj <- 1L
  }
  
  # ---------------------------------------------------------
  # Run DoubletFinder
  # ---------------------------------------------------------
  # Wrapped in tryCatch: if DoubletFinder fails for one sample, the
  # cell is marked as "DF_final_call_failed" in metadata and included
  # (unfiltered) in the merged object. The pipeline continues.
  log_msg(
    "DoubletFinder final call START for ", sample_id,
    " | pN=", doubletfinder_pN_use,
    " | pK=", best_pK,
    " | nExp=", nExp_poi_adj
  )
  
  df_call_success <- TRUE
  
  obj <- tryCatch({
    doubletFinder_fun(
      obj,
      PCs       = doubletfinder_pcs,
      pN        = doubletfinder_pN_use,
      pK        = best_pK,
      nExp      = nExp_poi_adj,
      reuse.pANN = NULL,
      sct       = TRUE
    )
  }, error = function(e) {
    df_call_success <<- FALSE
    warning(
      "Final DoubletFinder call failed for ", sample_id,
      ". Returning sample without doublet filtering. Error: ", e$message
    )
    obj$DoubletFinder_class <- "DF_final_call_failed"
    obj$DoubletFinder_pANN  <- NA_real_
    obj$DoubletFinder_nExp  <- nExp_poi_adj
    obj$DoubletFinder_pK    <- best_pK
    obj
  })
  
  log_msg("DoubletFinder final call DONE/EXIT for sample: ", sample_id)
  
  # ---------------------------------------------------------
  # Extract DoubletFinder classification columns safely
  # ---------------------------------------------------------
  # DoubletFinder writes metadata columns with auto-generated names
  # like DF.classifications_0.25_0.15_119 and pANN_0.25_0.15_119.
  # We extract these using obj@meta.data[[col]] (not obj$col) because
  # obj$col can return a data-frame-like object in newer Seurat builds,
  # which would break subsequent vector operations (e.g. %in%, which).
  if (df_call_success) {
    class_cols <- grep("^DF.classifications", colnames(obj@meta.data), value = TRUE)
    pANN_cols  <- grep("^pANN",              colnames(obj@meta.data), value = TRUE)
    
    if (length(class_cols) == 0) {
      warning("No DoubletFinder classification column found for ", sample_id)
      df_class_vec <- rep("Not_found", ncol(obj))
      df_pANN_vec  <- rep(NA_real_,    ncol(obj))
      status       <- "Classification_not_found"
      n_singlet    <- NA_integer_
      n_doublet    <- NA_integer_
    } else {
      # Use the last matching column in case multiple runs left stale columns.
      class_col <- tail(class_cols, 1)
      pANN_col  <- if (length(pANN_cols) > 0) tail(pANN_cols, 1) else NA_character_
      
      df_class_raw <- obj@meta.data[[class_col]]
      if (is.data.frame(df_class_raw)) df_class_raw <- df_class_raw[[1]]
      df_class_vec <- as.character(df_class_raw)
      
      if (!is.na(pANN_col)) {
        df_pANN_raw <- obj@meta.data[[pANN_col]]
        if (is.data.frame(df_pANN_raw)) df_pANN_raw <- df_pANN_raw[[1]]
        df_pANN_vec <- suppressWarnings(as.numeric(df_pANN_raw))
      } else {
        df_pANN_vec <- rep(NA_real_, ncol(obj))
      }
      
      status    <- "Completed"
      n_singlet <- sum(df_class_vec == "Singlet", na.rm = TRUE)
      n_doublet <- sum(df_class_vec == "Doublet", na.rm = TRUE)
    }
  } else {
    df_class_vec <- rep("DF_final_call_failed", ncol(obj))
    df_pANN_vec  <- rep(NA_real_,               ncol(obj))
    status       <- "DF_final_call_failed"
    n_singlet    <- NA_integer_
    n_doublet    <- NA_integer_
  }
  
  # ---------------------------------------------------------
  # Store DoubletFinder results as plain named vectors
  # ---------------------------------------------------------
  # Force all classification outputs into simple named character/numeric
  # vectors before writing to metadata. This prevents Seurat/tibble
  # coercion issues in downstream subset() and match() calls.
  df_class_vec <- as.character(df_class_vec)
  df_pANN_vec  <- suppressWarnings(as.numeric(df_pANN_vec))
  
  # Guard against length mismatches (should never happen, but be safe).
  if (length(df_class_vec) != ncol(obj)) {
    warning(
      "DoubletFinder class vector length mismatch for ", sample_id,
      ". Marking classifications as Not_found."
    )
    df_class_vec <- rep("Not_found", ncol(obj))
    status       <- "Classification_length_mismatch"
    n_singlet    <- NA_integer_
    n_doublet    <- NA_integer_
  }
  if (length(df_pANN_vec) != ncol(obj)) {
    df_pANN_vec <- rep(NA_real_, ncol(obj))
  }
  
  # Name vectors by cell barcode for safe index-based assignment.
  names(df_class_vec) <- colnames(obj)
  names(df_pANN_vec)  <- colnames(obj)
  
  obj@meta.data[["DoubletFinder_class"]] <- df_class_vec
  obj@meta.data[["DoubletFinder_pANN"]]  <- df_pANN_vec
  obj@meta.data[["DoubletFinder_nExp"]]  <- rep(as.integer(nExp_poi_adj), ncol(obj))
  obj@meta.data[["DoubletFinder_pK"]]    <- rep(as.numeric(best_pK),      ncol(obj))
  
  # ---------------------------------------------------------
  # Write compact DoubletFinder summary row
  # ---------------------------------------------------------
  # Appended to the global doubletfinder_summary_list and written
  # to CSV at the end of the script. Useful for QC review and for
  # identifying which patients/samples were affected by doublet issues.
  doubletfinder_summary_list[[sample_id]] <<- data.frame(
    sample_id        = sample_id,
    n_cells_after_QC = ncol(obj),
    doublet_rate     = doublet_rate,
    pN               = doubletfinder_pN_use,
    pK               = best_pK,
    nExp_unadjusted  = nExp_poi,
    homotypic_prop   = homotypic_prop,
    nExp_adjusted    = nExp_poi_adj,
    n_singlet        = n_singlet,
    n_doublet        = n_doublet,
    status           = status,
    stringsAsFactors = FALSE
  )
  
  # ---------------------------------------------------------
  # Keep singlets only
  # ---------------------------------------------------------
  # Direct column indexing (obj[, cells]) is used instead of
  # subset(cells = ...) to avoid additional Seurat coercion behavior
  # that can cause silent failures when cell names contain special chars.
  if (identical(status, "Completed") && any(df_class_vec == "Doublet", na.rm = TRUE)) {
    singlet_cells <- names(df_class_vec)[df_class_vec == "Singlet"]
    singlet_cells <- intersect(singlet_cells, colnames(obj))
    if (length(singlet_cells) == 0) {
      warning("No Singlet cells found after DoubletFinder for ", sample_id, ". Returning unfiltered object.")
    } else {
      obj <- obj[, singlet_cells]
    }
    rm(singlet_cells)
  }
  
  # ---------------------------------------------------------
  # Remove temporary DoubletFinder preprocessing outputs
  # ---------------------------------------------------------
  # SCT assay, PCA reduction, and graph objects were created only
  # for DoubletFinder. Removing them here ensures the returned object
  # contains only the RNA assay with raw counts, so the main pipeline
  # (SCTransform + Harmony) starts from a clean slate.
  if ("RNA" %in% names(obj@assays)) DefaultAssay(obj) <- "RNA"
  if ("SCT" %in% names(obj@assays))  obj[["SCT"]] <- NULL
  if ("pca" %in% names(obj@reductions)) obj[["pca"]] <- NULL
  obj@graphs <- list()
  gc()
  
  df_elapsed <- round(as.numeric(difftime(Sys.time(), df_start_time, units = "mins")), 2)
  log_msg("DoubletFinder DONE for sample: ", sample_id, " | elapsed: ", df_elapsed, " min")
  
  obj
}


# ---------------------------------------------------------
# assign_broad_labels()
# ---------------------------------------------------------
# Scores every cell in `obj` against five broad cell-type marker sets
# using mean log-normalized expression (via AddModuleScore). The
# top-scoring compartment is assigned as broad_label. Cells where the
# margin between the top and second-best score is below
# broad_margin_cutoff receive broad_label_conf = "Ambiguous".
#
# The confident broad label is used as a fallback cell type in
# make_final_cell_type_label() when SingleR does not return a result.
#
# JoinLayers() is called before NormalizeData() to merge per-sample
# split layers into a single unified layer. Without this, AddModuleScore()
# silently reads only the first sample's data, producing near-zero
# expression for all other cells and collapsing every cell to "Ambiguous".
assign_broad_labels <- function(obj) {
  DefaultAssay(obj) <- "RNA"
  
  # Merges per-sample split RNA layers (counts.P1_T1, counts.P1_B1, ...)
  # into a single unified layer before normalization. Without this,
  # NormalizeData() and AddModuleScore() silently read only the first
  # layer in Seurat v5, producing near-zero expression for cells from
  # all other samples. JoinLayers() is a no-op when layers are already joined
  if ("JoinLayers" %in% getNamespaceExports("Seurat")) {
    obj <- JoinLayers(obj, assay = "RNA")
  }
  
  obj <- NormalizeData(obj, verbose = FALSE)
  
  # Only use markers that are actually present in the dataset.
  # Require at least 2 genes per compartment to produce a meaningful score.
  present_marker_list <- lapply(broad_marker_list, function(g) intersect(g, rownames(obj)))
  present_marker_list <- present_marker_list[lengths(present_marker_list) >= 2]
  
  if (length(present_marker_list) == 0) {
    # No usable markers — mark all cells as Ambiguous and return.
    obj$broad_label      <- NA_character_
    obj$broad_margin     <- NA_real_
    obj$broad_label_conf <- "Ambiguous"
    return(obj)
  }
  
  # Score each compartment. AddModuleScore appends "1" to the name
  # (a quirk of the function); the rename step below strips it.
  for (nm in names(present_marker_list)) {
    obj <- AddModuleScore(
      obj,
      features = list(present_marker_list[[nm]]),
      name     = paste0("broad_", nm),
      assay    = "RNA"
    )
  }
  
  # Strip the trailing "1" appended by AddModuleScore.
  colnames(obj@meta.data) <- sub("^broad_(.*)1$", "broad_\\1", colnames(obj@meta.data))
  
  # Extract the score columns and derive the top label and margin.
  score_cols <- intersect(paste0("broad_", names(present_marker_list)), colnames(obj@meta.data))
  score_mat  <- obj@meta.data[, score_cols, drop = FALSE]
  
  # Assign the compartment with the highest score as the broad label.
  obj$broad_label <- gsub("^broad_", "", colnames(score_mat)[max.col(score_mat, ties.method = "first")])
  
  # Margin = top score minus second-best score.
  # A higher margin means the top assignment is more clearly supported.
  obj$broad_margin <- if (ncol(score_mat) >= 2) {
    apply(t(apply(score_mat, 1, sort, decreasing = TRUE)), 1, function(x) x[1] - x[2])
  } else NA_real_
  
  # Confident label: only assigned if margin >= broad_margin_cutoff.
  obj$broad_label_conf <- ifelse(
    !is.na(obj$broad_margin) & obj$broad_margin >= broad_margin_cutoff,
    obj$broad_label,
    "Ambiguous"
  )
  
  obj
}


# ---------------------------------------------------------
# run_singler_with_centroids()
# ---------------------------------------------------------
# Annotates cells using SingleR against pre-computed HLCA centroid
# references. Centroids are stored as a compact gene x cell-type matrix
# rather than a full reference dataset, which is much faster and uses
# less RAM than a full SCE reference.
#
# Uses the SCT-normalized assay (rather than log-normalized RNA) to
# match the expression scale expected by the HLCA reference.
#
# fine_label_broad maps SingleR fine labels back to broad HLCA
# categories via the centroid reference's broad_map (if available).
run_singler_with_centroids <- function(obj, centroid_ref) {
  DefaultAssay(obj) <- "SCT"
  
  # Convert to SingleCellExperiment for SingleR compatibility.
  query_sce <- as.SingleCellExperiment(obj, assay = "SCT")
  
  # Build a minimal SCE reference from the centroid matrix.
  ref_sce <- SingleCellExperiment(assays = list(logcounts = centroid_ref$centroids))
  colData(ref_sce)$label.fine  <- centroid_ref$fine_labels
  colData(ref_sce)$label.broad <- if (!is.null(centroid_ref$broad_map)) {
    unname(centroid_ref$broad_map[centroid_ref$fine_labels])
  } else NA_character_
  
  # Restrict to genes present in both query and reference.
  common_genes <- intersect(rownames(query_sce), rownames(ref_sce))
  query_sce    <- query_sce[common_genes, ]
  ref_sce      <- ref_sce[common_genes,  ]
  
  # Run SingleR annotation.
  pred <- SingleR(test = query_sce, ref = ref_sce, labels = ref_sce$label.fine)
  
  # pruned.labels removes low-confidence assignments (set to NA).
  # Preference for pruned labels but fall back to labels when pruned is NA.
  obj$SingleR_label        <- as.character(pred$labels)
  obj$SingleR_pruned_label <- as.character(pred$pruned.labels)
  obj$fine_label           <- ifelse(
    is.na(obj$SingleR_pruned_label) | obj$SingleR_pruned_label == "",
    obj$SingleR_label,
    obj$SingleR_pruned_label
  )
  
  # Map fine labels back to broad HLCA categories if a broad_map is available.
  obj$fine_label_broad <- if (!is.null(centroid_ref$broad_map)) {
    unname(centroid_ref$broad_map[as.character(obj$fine_label)])
  } else NA_character_
  
  rm(query_sce, ref_sce, pred)
  gc()
  
  obj
}


# ---------------------------------------------------------
# make_final_cell_type_label()
# ---------------------------------------------------------
# Builds a single consolidated cell_type_label column using a
# hierarchical priority:
#   1. SingleR fine label with a valid broad map entry (most specific).
#   2. SingleR fine label without broad map.
#   3. Confident broad marker label (fallback if SingleR failed).
#   4. "Unassigned" (last resort).
#
# cell_type_label_source records which tier the label came from,
# which is useful for downstream QC and filtering.
make_final_cell_type_label <- function(obj) {
  obj$cell_type_label <- case_when(
    !is.na(obj$fine_label_broad) & obj$fine_label_broad != "" ~ as.character(obj$fine_label),
    !is.na(obj$fine_label)       & obj$fine_label != ""       ~ as.character(obj$fine_label),
    !is.na(obj$broad_label_conf) & obj$broad_label_conf != "" ~ as.character(obj$broad_label_conf),
    TRUE ~ "Unassigned"
  )
  obj$cell_type_label_source <- case_when(
    !is.na(obj$fine_label_broad) & obj$fine_label_broad != "" ~ "SingleR_with_broad_map",
    !is.na(obj$fine_label)       & obj$fine_label != ""       ~ "SingleR_no_broad_map",
    !is.na(obj$broad_label_conf) & obj$broad_label_conf != "" ~ "Broad_marker_fallback",
    TRUE ~ "Unassigned"
  )
  obj
}


# ---------------------------------------------------------
# label_tams_expanded()
# ---------------------------------------------------------
# Scores each cell in a macrophage-only Seurat object across the seven
# TAM program signatures defined in tam_gene_list_expanded. For each cell,
# the program with the highest mean log-normalized expression is assigned
# as TAM_subset. If the margin (top minus second-best score) meets or
# exceeds tam_margin_cutoff, the cell receives that confident TAM_subset_conf
# label; otherwise it is labelled "Ambiguous".
#
# Raw per-program mean scores are also stored as *_expanded metadata
# columns for downstream heatmaps and GSVA.

label_tams_expanded <- function(obj) {
  DefaultAssay(obj) <- "RNA"
  
  # After merge(), Seurat v5 stores each sample's counts in a separate
  # named layer (e.g. counts.P4_B1, counts.P4_T1). NormalizeData()
  # normalizes each layer independently, producing fragmented data.*
  # layers. GetAssayData(layer = "data") then returns ONLY the first
  # data layer, so every cell not in the first sample appears to have
  # near-zero expression. With near-zero expression, all seven TAM
  # program scores are ~0, no margin ever exceeds tam_margin_cutoff,
  # and every cell is assigned "Ambiguous" — yielding an all-NA
  # TAM_subset_conf column that causes Script 2 to skip the patient.
  #
  # JoinLayers() collapses all split layers into one before normalization,
  # ensuring every cell receives a valid normalized expression value.
  # It is a no-op for single-sample objects (layers already joined).
  if ("JoinLayers" %in% getNamespaceExports("Seurat")) {
    obj <- JoinLayers(obj, assay = "RNA")
  }
  
  obj <- NormalizeData(obj, verbose = FALSE)
  
  # Filter each program's gene list to genes present in the dataset.
  # Programs with fewer than 2 detectable genes are dropped entirely.
  marker_list <- lapply(tam_gene_list_expanded, function(g) intersect(g, rownames(obj)))
  marker_list <- marker_list[lengths(marker_list) >= 2]
  
  # If no programs have enough detectable genes, return Ambiguous labels.
  if (length(marker_list) == 0 || ncol(obj) == 0) {
    obj$TAM_subset      <- NA_character_
    obj$TAM_margin      <- NA_real_
    obj$TAM_subset_conf <- "Ambiguous"
    return(obj)
  }
  
  # ---------------------------------------------------------
  # Layer-aware expression matrix extraction
  # ---------------------------------------------------------
  # GetAssayData(layer = "data") fails with:
  #   "GetAssayData doesn't work for multiple layers in v5 assay"
  # when any data.* layers are still split — even after JoinLayers()
  # is called — because the macro_obj subset is created AFTER
  # JoinLayers() runs on the full object, and Seurat v5 subsetting
  # can re-fragment layers when the subset spans cells from multiple
  # original layers.
  #
  # The robust fix is to enumerate all data.* layers explicitly using
  # Layers(), extract each one with LayerData(), and cbind them into
  # a single unified matrix. This bypasses GetAssayData() entirely and
  # works correctly regardless of how many layers remain after subsetting.
  #
  # For a properly joined single-layer object, Layers() returns exactly
  # one "data" entry and the cbind is a no-op (returns that one matrix).
  all_layers  <- Layers(obj, assay = "RNA")
  data_layers <- grep("^data", all_layers, value = TRUE)
  
  if (length(data_layers) == 0) {
    # No normalized data layers found — NormalizeData() above should have
    # created them. This should not happen, but if it does, fall back to
    # the counts layer and normalize on the fly.
    warning("No data layers found for TAM scoring in ", file_prefix,
            ". Falling back to counts layer.")
    data_layers <- grep("^counts", all_layers, value = TRUE)
  }
  
  if (length(data_layers) == 1) {
    # Single layer — extract directly, no cbind needed.
    expr_mat <- LayerData(obj, assay = "RNA", layer = data_layers[1])
  } else {
    # Multiple data layers still present after subsetting — cbind all of them.
    # Rows (genes) are identical across layers; columns (cells) are disjoint.
    expr_mat <- do.call(cbind, lapply(data_layers, function(l) {
      LayerData(obj, assay = "RNA", layer = l)
    }))
    # Reorder columns to match Seurat's cell order to avoid any barcode
    # mismatches when indexing into score_mat by colnames(obj) below.
    expr_mat <- expr_mat[, colnames(obj), drop = FALSE]
  }
  
  # Compute per-program mean expression across each program's gene set.
  # This is equivalent to an unweighted module score without background
  # correction, which is appropriate here because we are comparing
  # programs relative to each other rather than to a random background.
  score_mat <- sapply(names(marker_list), function(nm) {
    genes <- marker_list[[nm]]
    vals  <- Matrix::colMeans(expr_mat[genes, , drop = FALSE])
    as.numeric(vals)
  })
  
  # Handle edge case: sapply returns a vector (not matrix) when only
  # one program passes the >= 2 gene filter.
  if (is.null(dim(score_mat))) {
    score_mat         <- matrix(score_mat, ncol = 1)
    colnames(score_mat) <- names(marker_list)
  }
  
  score_mat         <- as.data.frame(score_mat)
  rownames(score_mat) <- colnames(obj)
  
  # Store individual program scores as *_expanded metadata columns.
  # These are used by Script 3 for per-program heatmaps and GSVA scoring.
  for (nm in colnames(score_mat)) {
    obj@meta.data[[paste0(nm, "_expanded")]] <- score_mat[[nm]]
  }
  
  # Edge case: only one program available — cannot compute a meaningful margin.
  if (ncol(score_mat) == 1) {
    obj$TAM_subset      <- colnames(score_mat)[1]
    obj$TAM_margin      <- NA_real_
    obj$TAM_subset_conf <- "Ambiguous"
    return(obj)
  }
  
  # Assign the top-scoring program as TAM_subset.
  best_idx       <- max.col(score_mat, ties.method = "first")
  obj$TAM_subset <- colnames(score_mat)[best_idx]
  
  # Compute margin as top score minus second-best score.
  sorted_scores  <- t(apply(score_mat, 1, sort, decreasing = TRUE))
  obj$TAM_margin <- sorted_scores[, 1] - sorted_scores[, 2]
  
  # Confident label: only assigned when margin >= tam_margin_cutoff.
  # Cells below the cutoff receive "Ambiguous" to flag uncertain assignments.
  obj$TAM_subset_conf <- ifelse(
    !is.na(obj$TAM_margin) & obj$TAM_margin >= tam_margin_cutoff,
    obj$TAM_subset,
    "Ambiguous"
  )
  
  obj
}


# =========================================================
# PROGRESS / TIMING HELPERS
# =========================================================

# Record the script-level start time for cumulative elapsed reporting.
script_start_time <- Sys.time()

# log_msg(): timestamped message printed to console/log.
log_msg <- function(...) {
  message("[", format(Sys.time(), "%Y-%m-%d %H:%M:%S"), "] ", paste0(...))
}

# log_step(): logs a numbered pipeline step with percent-complete.
log_step <- function(step_i, step_n, label, detail = NULL) {
  pct <- round(100 * step_i / step_n, 1)
  if (is.null(detail)) {
    log_msg("STEP ", step_i, "/", step_n, " (", pct, "%): ", label)
  } else {
    log_msg("STEP ", step_i, "/", step_n, " (", pct, "%): ", label, " | ", detail)
  }
}

# time_block(): wraps an expression, logging start/done with step and
# total elapsed times. Returns the expression result transparently.
time_block <- function(label, expr) {
  t0     <- Sys.time()
  log_msg("START: ", label)
  result <- force(expr)
  elapsed       <- round(as.numeric(difftime(Sys.time(), t0,                  units = "mins")), 2)
  total_elapsed <- round(as.numeric(difftime(Sys.time(), script_start_time,   units = "mins")), 2)
  log_msg(
    "DONE: ", label,
    " | step elapsed: ", elapsed, " min",
    " | total elapsed: ", total_elapsed, " min"
  )
  result
}


# =========================================================
# process_group()
# =========================================================
# Core processing unit for one cancer patient or healthy donor.
# Called once per patient/donor from the main loops below.
#
# Arguments:
#   sample_ids  - Character vector of sample IDs belonging to this group.
#   file_prefix - Short string used for output file naming (e.g. "P4", "D1").
#   group_label - Expanded label for logging (e.g. "Patient 4").
#   read_fun    - Function to read one sample: either read_one_patient_sample
#                 or read_one_donor_sample.
#
# Output files saved per patient/donor:
#   outdir_full_annot  / <prefix>_full_annotated_with_TAM_labels.rds
#   outdir_macro       / <prefix>_macrophage_only.rds
#   outdir_tumor       / <prefix>_tumor_only.rds        (cancer patients only)
#   outdir_background  / <prefix>_background_only.rds   (cancer patients only)
process_group <- function(sample_ids, file_prefix, group_label, read_fun) {
  
  group_start_time <- Sys.time()
  step_n           <- 9
  
  log_msg("=========================================================")
  log_msg("Processing group: ", group_label, " / ", file_prefix)
  log_msg("Samples: ", paste(sample_ids, collapse = ", "))
  log_msg("=========================================================")
  
  # Define output file paths for this patient/donor.
  final_full_file       <- file.path(outdir_full_annot,  paste0(file_prefix, "_full_annotated_with_TAM_labels.rds"))
  final_macro_file      <- file.path(outdir_macro,       paste0(file_prefix, "_macrophage_only.rds"))
  final_tumor_file      <- file.path(outdir_tumor,       paste0(file_prefix, "_tumor_only.rds"))
  final_background_file <- file.path(outdir_background,  paste0(file_prefix, "_background_only.rds"))
  
  # Skip if both primary outputs already exist and skip flag is set.
  # The full annotated and macrophage objects are the minimum needed
  # by Script 2; tumor/background objects are regenerated if missing.
  if (skip_existing_outputs && file.exists(final_full_file) && file.exists(final_macro_file)) {
    log_msg("Skipping ", file_prefix, ": final full and TAM outputs already exist.")
    return(invisible(NULL))
  }
  
  # -------------------------------------------------------
  # STEP 1: Read samples and run sample-level DoubletFinder
  # -------------------------------------------------------
  # Each sample in the patient/donor group is read and doublet-filtered
  # independently before merging. This loop accumulates processed
  # per-sample objects; NULLs (failed reads or all-filtered samples)
  # are dropped after the loop.
  
  log_step(1, step_n, "Read samples and run DoubletFinder", file_prefix)
  
  objs <- time_block(paste0("Reading samples + DoubletFinder for group ", file_prefix), {
    objs <- vector("list", length(sample_ids))
    names(objs) <- sample_ids
    for (i in seq_along(sample_ids)) {
      sid        <- sample_ids[i]
      sample_pct <- round(100 * i / length(sample_ids), 1)
      log_msg("Sample ", i, "/", length(sample_ids), " (", sample_pct, "%) for ", file_prefix, ": ", sid)
      sample_obj <- read_fun(sid)
      if (!is.null(sample_obj)) {
        sample_obj <- run_doubletfinder_one_sample(obj = sample_obj, sample_id = sid)
      }
      objs[[i]] <- sample_obj
      rm(sample_obj)
      gc()
    }
    objs
  })
  
  # Drop any samples that returned NULL (read failure or zero cells after QC).
  objs <- objs[!vapply(objs, is.null, logical(1))]
  
  if (length(objs) == 0) {
    log_msg("No readable samples retained for ", file_prefix, ". Skipping group.")
    return(invisible(NULL))
  }
  
  # -------------------------------------------------------
  # STEP 2: Merge singlet-filtered samples
  # -------------------------------------------------------
  # After DoubletFinder, merge all per-sample Seurat objects into one
  # patient-level object. In Seurat v5, merge() creates split layers
  # (counts.P4_B1, counts.P4_T1, ...) which is the upstream cause of
  # the split-layer bug. The JoinLayers() calls in assign_broad_labels()
  # and label_tams_expanded() resolve this before any scoring is done.
  
  log_step(2, step_n, "Merge singlet-filtered samples", file_prefix)
  
  full_obj <- time_block(paste0("Merging samples for ", file_prefix), {
    if (length(objs) == 1) {
      objs[[1]]
    } else {
      merge(objs[[1]], y = objs[2:length(objs)])
    }
  })
  
  rm(objs)
  gc()
  
  # -------------------------------------------------------
  # STEP 3: Final merged-object QC
  # -------------------------------------------------------
  # Apply QC thresholds a second time on the merged object. This catches
  # any borderline cells that may have passed sample-level QC only because
  # per-sample ambient RNA background was lower or higher than the cohort
  # average.
  
  log_step(3, step_n, "Final merged-object QC", file_prefix)
  
  full_obj <- time_block(paste0("Final QC for ", file_prefix), {
    run_basic_qc(full_obj)
  })
  
  if (ncol(full_obj) == 0) {
    log_msg("No cells retained after final QC for ", file_prefix)
    rm(full_obj)
    gc()
    return(invisible(NULL))
  }
  
  # -------------------------------------------------------
  # STEP 4: SCT normalization, PCA, Harmony, clustering, UMAP
  # -------------------------------------------------------
  # SCTransform (glmGamPoi backend) regresses out mitochondrial fraction
  # and produces Pearson residuals used for PCA.
  #
  # Harmony batch-corrects the PCA embedding across samples within the
  # patient (relevant when a patient has both tumor and background samples
  # from the same individual but different captures).
  #
  # For single-sample patients/donors, Harmony is skipped and PCA
  # embeddings are used directly for clustering and UMAP.
  #
  # Memory note:
  #   Large patients (e.g. P20 with ~110k cells across 4 samples) can
  #   exhaust RAM during SCTransform if any intermediate objects from
  #   the merge/QC steps are still allocated. DietSeurat() strips all
  #   non-essential slots (scale.data, existing reductions, graphs)
  #   before SCTransform runs, minimising peak RAM. This is safe because:
  #     - scale.data is not used downstream (SCT produces its own).
  #     - Any reductions/graphs from this point are pre-SCT and would
  #       be replaced by the SCT-based ones immediately below.
  #   For small patients the overhead is negligible; for large ones it
  #   can save 10-20 GB of peak allocation.
  
  log_step(4, step_n, "SCT/PCA/Harmony/clustering/UMAP", file_prefix)
  
  # ---------------------------------------------------------
  # Pre-SCTransform memory diet
  # ---------------------------------------------------------
  # Strip scale.data, all reductions, and all graphs before SCTransform.
  # counts = TRUE and data = TRUE preserve the raw and log-normalised
  # RNA layers needed by label_tams_expanded() and assign_broad_labels()
  # later. SCT assay is excluded here (it doesn't exist yet at this point).
  full_obj <- DietSeurat(
    full_obj,
    counts     = TRUE,
    data       = TRUE,
    scale.data = FALSE,
    assays     = "RNA",
    dimreducs  = NULL,
    graphs     = NULL
  )
  gc()
  
  full_obj <- time_block(paste0("Dimensional reduction and clustering for ", file_prefix), {
    
    full_obj <- SCTransform(
      full_obj,
      method                = "glmGamPoi",
      vars.to.regress       = "percent.mt",
      conserve.memory       = TRUE,
      return.only.var.genes = FALSE,    # Keep all genes in the SCT assay for SingleR.
      verbose               = TRUE
    )
    
    full_obj <- RunPCA(full_obj, npcs = npcs_use, verbose = FALSE)
    
    if (length(unique(full_obj$sample_id)) > 1) {
      # Multi-sample patient: use Harmony to correct per-sample technical variation.
      full_obj <- RunHarmony(
        full_obj,
        group.by.vars  = "sample_id",
        assay.use      = "SCT",
        reduction      = "pca",
        reduction.save = "integrated.harmony",
        max_iter       = 100,
        verbose        = TRUE
      )
      full_obj <- FindNeighbors(full_obj, reduction = "integrated.harmony", dims = 1:npcs_use)
      full_obj <- FindClusters(full_obj, resolution = 0.8, cluster.name = "harmony_clusters")
      if (run_umap) {
        full_obj <- RunUMAP(full_obj, reduction = "integrated.harmony", dims = 1:npcs_use, reduction.name = "umap.harmony")
      }
    } else {
      # Single-sample patient/donor: skip Harmony.
      full_obj <- FindNeighbors(full_obj, reduction = "pca", dims = 1:npcs_use)
      full_obj <- FindClusters( full_obj, resolution = 0.8, cluster.name = "pca_clusters")
      if (run_umap) {
        # Use "umap.harmony" as the name even for PCA-based UMAP so that
        # Script 2 can reference a consistent reduction name.
        full_obj <- RunUMAP(full_obj, reduction = "pca", dims = 1:npcs_use, reduction.name = "umap.harmony")
      }
    }
    
    full_obj
  })
  
  gc()
  
  # -------------------------------------------------------
  # STEP 5: Broad marker labels + SingleR annotation
  # -------------------------------------------------------
  # assign_broad_labels() provides a coarse five-compartment classification
  # based on manually curated markers.
  #
  # run_singler_with_centroids() provides a fine-grained HLCA cell type
  # annotation using the SCT-normalized assay.
  #
  # make_final_cell_type_label() consolidates both into cell_type_label
  # using the hierarchical priority described in that function.
  #
  # smoking_binary is computed here so it is available in all downstream
  # metadata joins and CSVs.
  
  log_step(5, step_n, "Broad marker labels and SingleR annotation", file_prefix)
  
  full_obj <- time_block(paste0("Annotation for ", file_prefix), {
    full_obj <- assign_broad_labels(full_obj)
    full_obj <- run_singler_with_centroids(obj = full_obj, centroid_ref = hlca_centroids)
    full_obj <- make_final_cell_type_label(full_obj)
    full_obj$smoking_binary <- recode_smoking_binary(full_obj$smoking_status)
    full_obj
  })
  
  gc()
  
  # -------------------------------------------------------
  # STEP 6: Macrophage/TAM subsetting and TAM-state labeling
  # -------------------------------------------------------
  # Cells annotated as macrophage-like by cell_type_label (case-insensitive
  # substring match) are extracted into a separate macrophage object.
  # label_tams_expanded() then applies the split-layer fix and scores
  # each macrophage for the seven TAM programs.
  #
  # The TAM scoring uses log-normalized RNA expression, which is computed
  # fresh inside label_tams_expanded() after JoinLayers().
  
  log_step(6, step_n, "Macrophage/TAM subsetting and TAM-state labeling", file_prefix)
  
  macro_obj <- time_block(paste0("TAM subsetting/labeling for ", file_prefix), {
    
    # Identify macrophage-like cells by cell_type_label string matching.
    # The regex covers all HLCA macrophage subtypes expected in lung tissue.
    macrophage_keep <- str_detect(
      tolower(as.character(full_obj@meta.data[["cell_type_label"]])),
      "macrophage|alveolar macrophage|interstitial macrophage|elicited macrophage|lung macrophage"
    )
    macrophage_keep[is.na(macrophage_keep)] <- FALSE
    n_macro <- sum(macrophage_keep)
    log_msg("Macrophage-like cells detected for ", file_prefix, ": ", n_macro)
    
    if (n_macro > 0) {
      macro_cells <- colnames(full_obj)[macrophage_keep]
      macro_obj   <- full_obj[, macro_cells]
      
      # Wrap TAM labeling in tryCatch so a gene-list or expression issue
      # for one patient does not abort the entire pipeline.
      macro_obj <- tryCatch({
        label_tams_expanded(macro_obj)
      }, error = function(e) {
        warning(
          "TAM labeling failed for ", file_prefix,
          ". Returning macrophage object with Ambiguous TAM labels. Error: ", e$message
        )
        macro_obj$TAM_subset      <- NA_character_
        macro_obj$TAM_margin      <- NA_real_
        macro_obj$TAM_subset_conf <- "Ambiguous"
        macro_obj
      })
    } else {
      log_msg("No macrophage cells retained for ", file_prefix)
      macro_obj <- NULL
    }
    
    rm(macrophage_keep)
    macro_obj
  })
  
  gc()
  
  # -------------------------------------------------------
  # STEP 7: Transfer TAM labels back to the full object
  # -------------------------------------------------------
  # TAM_subset, TAM_margin, and TAM_subset_conf are written back to
  # the full patient object by cell barcode match. Non-macrophage cells
  # receive NA for TAM_subset and TAM_margin.
  #
  # cellchat_group is built here as a convenience column for Script 2:
  #   - Macrophages with confident TAM labels -> "TAM_<subset>"
  #   - All other cells -> their cell_type_label
  #   - Unlabeled cells -> "Other"
  
  log_step(7, step_n, "Transfer TAM labels back to full object", file_prefix)
  
  full_obj <- time_block(paste0("Transferring TAM labels for ", file_prefix), {
    
    # Initialise TAM columns for all cells (non-macrophages stay NA).
    full_obj$TAM_subset      <- NA_character_
    full_obj$TAM_margin      <- NA_real_
    full_obj$TAM_subset_conf <- NA_character_
    
    if (!is.null(macro_obj)) {
      # Match by cell barcode to avoid index misalignment.
      shared_cells <- intersect(colnames(full_obj), colnames(macro_obj))
      full_idx     <- match(shared_cells, colnames(full_obj))
      macro_idx    <- match(shared_cells, colnames(macro_obj))
      keep_idx     <- !is.na(full_idx) & !is.na(macro_idx)
      full_idx     <- full_idx[keep_idx]
      macro_idx    <- macro_idx[keep_idx]
      
      full_obj@meta.data[full_idx, "TAM_subset"]      <- as.character(macro_obj@meta.data[macro_idx, "TAM_subset"])
      full_obj@meta.data[full_idx, "TAM_margin"]      <- as.numeric( macro_obj@meta.data[macro_idx, "TAM_margin"])
      full_obj@meta.data[full_idx, "TAM_subset_conf"] <- as.character(macro_obj@meta.data[macro_idx, "TAM_subset_conf"])
      
      rm(full_idx, macro_idx, keep_idx, shared_cells)
    }
    
    # Build cellchat_group column used by Script 2.
    # Priority: confident TAM subset > cell_type_label > "Other".
    full_obj$cellchat_group <- case_when(
      !is.na(full_obj$TAM_subset_conf) &
        full_obj$TAM_subset_conf != "" &
        full_obj$TAM_subset_conf != "Ambiguous" ~ paste0("TAM_", full_obj$TAM_subset_conf),
      !is.na(full_obj$cell_type_label) &
        full_obj$cell_type_label != "" ~ as.character(full_obj$cell_type_label),
      TRUE ~ "Other"
    )
    
    full_obj
  })
  
  gc()
  
  # -------------------------------------------------------
  # STEP 8: Save final output RDS files
  # -------------------------------------------------------
  # Four files are saved per patient/donor (tumor and background files
  # are only written if the relevant sample type is present):
  #   1. Full annotated object (all cells, all annotations).
  #   2. Macrophage-only object (TAM-labelled).
  #   3. Tumor-only object.
  #   4. Background-only object.
  #
  # The full and macrophage objects are the primary inputs to Script 2.
  # Script 1.5 uses the tumor and background objects for cohort integration.
  
  log_step(8, step_n, "Save final RDS objects", file_prefix)
  
  time_block(paste0("Saving final objects for ", file_prefix), {
    
    saveRDS(full_obj, final_full_file)
    
    if (!is.null(macro_obj)) {
      saveRDS(macro_obj, final_macro_file)
    }
    
    if (sum(full_obj$sample_type == "tumor", na.rm = TRUE) > 0) {
      tumor_obj <- subset(full_obj, subset = sample_type == "tumor")
      saveRDS(tumor_obj, final_tumor_file)
      rm(tumor_obj)
      gc()
    }
    
    if (sum(full_obj$sample_type == "background", na.rm = TRUE) > 0) {
      background_obj <- subset(full_obj, subset = sample_type == "background")
      saveRDS(background_obj, final_background_file)
      rm(background_obj)
      gc()
    }
    
    invisible(TRUE)
  })
  
  # -------------------------------------------------------
  # STEP 9: Memory cleanup
  # -------------------------------------------------------
  # Explicitly remove the two large patient objects and force garbage
  # collection before moving to the next patient. In a 24-patient cohort
  # this is important for keeping peak RAM usage bounded.
  
  log_step(9, step_n, "Cleanup", file_prefix)
  
  rm(full_obj, macro_obj)
  gc()
  
  group_elapsed <- round(as.numeric(difftime(Sys.time(), group_start_time, units = "mins")), 2)
  total_elapsed <- round(as.numeric(difftime(Sys.time(), script_start_time,  units = "mins")), 2)
  
  log_msg(
    "FINISHED group: ", group_label, " / ", file_prefix,
    " | group elapsed: ", group_elapsed, " min",
    " | total elapsed: ", total_elapsed, " min"
  )
  
  invisible(NULL)
}


# =========================================================
# LOAD CENTROID REFERENCE
# =========================================================
# The HLCA centroid RDS must contain at minimum:
#   $centroids   - gene x cell-type numeric matrix of log-normalised
#                  mean expression values (the reference "centroids").
#   $fine_labels - character vector of cell type names (one per column).
# Optionally:
#   $broad_map   - named character vector mapping fine labels to broad
#                  HLCA compartment names.

hlca_centroids <- readRDS(centroid_rds)

required_centroid_slots <- c("centroids", "fine_labels")
missing_centroid_slots  <- setdiff(required_centroid_slots, names(hlca_centroids))

if (length(missing_centroid_slots) > 0) {
  stop("Centroid reference missing required slots: ", paste(missing_centroid_slots, collapse = ", "))
}


# =========================================================
# PART 1: PROCESS CANCER PATIENTS
# =========================================================
# Iterates over each unique patient_id in the cohort, collects all
# sample IDs belonging to that patient, derives the file_prefix from
# the patient_key (P1, P2, ...), and calls process_group().

for (pid in sort(unique(patient_ID_map))) {
  patient_samples     <- sort(names(patient_ID_map)[patient_ID_map == pid])
  patient_key_current <- unique(patient_key_map[patient_samples])
  patient_key_current <- patient_key_current[!is.na(patient_key_current)]
  file_prefix         <- if (length(patient_key_current) == 0) gsub(" ", "_", pid) else patient_key_current[1]
  
  process_group(
    sample_ids  = patient_samples,
    file_prefix = file_prefix,
    group_label = pid,
    read_fun    = read_one_patient_sample
  )
}


# =========================================================
# PART 2: PROCESS HEALTHY DONORS
# =========================================================
# Healthy donors use a different sample naming convention (D<n>_<n>)
# and a separate read function. All other processing steps are identical
# to cancer patients. Donors do not produce tumor or background sub-objects
# since all samples are labelled sample_type = "healthy".

donor_sample_ID <- get_donor_sample_ids(data_dir)

if (length(donor_sample_ID) == 0) {
  log_msg("No healthy donor samples detected. Skipping healthy donor pipeline.")
} else {
  donor_meta <- tibble(sample_id = donor_sample_ID) %>%
    mutate(
      donor_key    = str_extract(sample_id, "^D\\d+"),
      donor_number = as.integer(str_match(sample_id, "^D\\d+_(\\d+)$")[, 2])
    )
  
  for (did in sort(unique(donor_meta$donor_key))) {
    donor_samples <- sort(donor_meta$sample_id[donor_meta$donor_key == did])
    process_group(
      sample_ids  = donor_samples,
      file_prefix = did,
      group_label = paste("Healthy donor", did),
      read_fun    = read_one_donor_sample
    )
  }
}


# =========================================================
# SAVE DOUBLETFINDER QC SUMMARY
# =========================================================
# Writes one row per sample to a CSV in outdir_qc.
# Columns include: sample_id, n_cells, pK, nExp, n_singlet, n_doublet,
# and status. Used to verify DoubletFinder ran correctly and to
# cross-reference with the split-layer diagnostic (n_samples per patient).

if (save_doubletfinder_summary_csv && length(doubletfinder_summary_list) > 0) {
  write.csv(
    bind_rows(doubletfinder_summary_list),
    file.path(outdir_qc, "DoubletFinder_summary_all_samples.csv"),
    row.names = FALSE
  )
}


# =========================================================
# SAVE FINAL CELL COUNT SUMMARY TABLE
# =========================================================
# Reads every final full annotated object and tabulates cell counts at
# two levels:
#   - per sample (sample_id + patient metadata)
#   - per patient/donor (patient_key + patient metadata)
#
# For each grouping level, five annotation types are counted:
#   1. total_cells        - all cells regardless of label.
#   2. broad_label_conf   - confident broad compartment labels.
#   3. fine_label         - SingleR fine HLCA labels.
#   4. cell_type_label    - final consolidated working label.
#   5. TAM_subset_conf    - confident TAM program labels.
#
# The resulting CSV is used to:
#   a) Verify TAM labeling success per patient (has_real_TAM diagnostic).
#   b) Identify patients affected by the split-layer bug (all Ambiguous).
#   c) Provide cell-type composition tables for the Methods section.

make_count_rows <- function(meta, group_cols, group_level_name) {
  
  group_cols <- intersect(group_cols, colnames(meta))
  
  # Inner helper: count cells by annotation label within each group.
  count_annotation <- function(annotation_col, annotation_type) {
    if (!(annotation_col %in% colnames(meta))) return(NULL)
    meta %>%
      as.data.frame() %>%
      dplyr::mutate(
        label = as.character(.data[[annotation_col]]),
        label = dplyr::if_else(is.na(label) | label == "", "NA_or_blank", label)
      ) %>%
      dplyr::group_by(dplyr::across(dplyr::all_of(group_cols)), label) %>%
      dplyr::summarise(n_cells = dplyr::n(), .groups = "drop") %>%
      dplyr::mutate(group_level = group_level_name, annotation_type = annotation_type)
  }
  
  # Total cells per group (no annotation breakdown).
  total_counts <- meta %>%
    as.data.frame() %>%
    dplyr::group_by(dplyr::across(dplyr::all_of(group_cols))) %>%
    dplyr::summarise(n_cells = dplyr::n(), .groups = "drop") %>%
    dplyr::mutate(group_level = group_level_name, annotation_type = "total_cells", label = "all_cells")
  
  dplyr::bind_rows(
    total_counts,
    count_annotation("broad_label_conf", "broad_label_conf"),
    count_annotation("fine_label",       "fine_label"),
    count_annotation("cell_type_label",  "cell_type_label"),
    count_annotation("TAM_subset_conf",  "TAM_subset_conf")
  )
}

if (save_final_cell_count_summary_csv) {
  
  final_full_files <- list.files(
    path    = outdir_full_annot,
    pattern = "_full_annotated_with_TAM_labels\\.rds$",
    full.names = TRUE
  )
  
  if (length(final_full_files) == 0) {
    warning("No final full annotated objects found. Skipping final cell count summary.")
  } else {
    
    cell_count_summary_list <- list()
    
    for (f in final_full_files) {
      
      log_msg("Summarizing cell counts for: ", basename(f))
      
      obj  <- readRDS(f)
      meta <- obj@meta.data
      
      # Per-sample grouping columns.
      sample_group_cols <- intersect(
        c("patient_key","patient_id","sample_id","sample_type","dataset",
          "cancer_type","smoking_status","smoking_binary","sex","stage"),
        colnames(meta)
      )
      
      # Per-patient/donor grouping columns (no sample_id).
      patient_group_cols <- intersect(
        c("patient_key","patient_id","dataset","cancer_type",
          "smoking_status","smoking_binary","sex","stage"),
        colnames(meta)
      )
      
      cell_count_summary_list[[basename(f)]] <- dplyr::bind_rows(
        make_count_rows(meta = meta, group_cols = sample_group_cols,  group_level_name = "sample"),
        make_count_rows(meta = meta, group_cols = patient_group_cols, group_level_name = "patient_or_donor")
      )
      
      rm(obj, meta)
      gc()
    }
    
    final_cell_count_summary <- dplyr::bind_rows(cell_count_summary_list)
    
    # Reorder columns for readability.
    preferred_cols <- c(
      "group_level","patient_key","patient_id","sample_id","sample_type",
      "dataset","cancer_type","smoking_status","smoking_binary","sex","stage",
      "annotation_type","label","n_cells"
    )
    
    final_cell_count_summary <- final_cell_count_summary %>%
      dplyr::select(dplyr::any_of(preferred_cols), dplyr::everything()) %>%
      dplyr::arrange(group_level, patient_key, sample_id, annotation_type, dplyr::desc(n_cells))
    
    write.csv(
      final_cell_count_summary,
      file.path(outdir_qc, "Final_cell_count_summary_by_sample_and_patient.csv"),
      row.names = FALSE
    )
    
    rm(final_cell_count_summary, cell_count_summary_list)
    gc()
  }
}

log_msg("Script 1 complete.")

sessionInfo()