# ==============================================================================
# Core Analysis Pipeline for Malignant-Myeloid Interaction Signature (MMIS)
# ==============================================================================
# This script provides the core logic for Reference Mapping, WGCNA, and Meta-Analysis.
# Libraries used: Seurat, WGCNA, meta, survival.

# ------------------------------------------------------------------------------
# 1. Single-Cell Reference Mapping
# ------------------------------------------------------------------------------
# Input Data Format: 
#   - query: A Seurat object containing raw counts from target scRNA-seq datasets.
#   - reference: A well-annotated Seurat object (e.g., GBmap).
# ------------------------------------------------------------------------------

run_reference_mapping <- function(query_obj, reference_obj) {
  library(Seurat)
  
  # Standard Query Processing
  query_obj <- NormalizeData(query_obj) %>% 
    FindVariableFeatures() %>% 
    ScaleData() %>% 
    RunPCA()
  
  # Find Transfer Anchors
  anchors <- FindTransferAnchors(
    reference = reference_obj,
    query = query_obj,
    dims = 1:50,
    reference.reduction = "pca"
  )
  
  # Transfer Cell Type Labels
  query_obj <- TransferData(
    anchorset = anchors,
    reference = reference_obj,
    query = query_obj,
    refdata = list(celltype = "annotation_level_4")
  )
  
  return(query_obj)
}

# ------------------------------------------------------------------------------
# 2. Weighted Correlation Network Analysis (WGCNA) - Core Logic
# ------------------------------------------------------------------------------
# Input Data Format:
#   - exp_matrix: A log2-transformed expression matrix (Samples as rows, Genes as columns).
#   - trait_data: A data frame containing numeric myeloid scores for the same samples.
# ------------------------------------------------------------------------------

run_core_wgcna <- function(exp_matrix, trait_data) {
  library(WGCNA)
  
  # Pick Soft Thresholding Power
  powers = c(1:20)
  sft = pickSoftThreshold(exp_matrix, powerVector = powers, verbose = 5)
  softPower = sft$powerEstimate # e.g., β selected to reach R^2 > 0.9
  
  # Calculate Adjacency and TOM
  adjacency = adjacency(exp_matrix, power = softPower)
  TOM = TOMsimilarity(adjacency)
  dissTOM = 1 - TOM
  
  # Hierarchical Clustering and Module Identification
  geneTree = hclust(as.dist(dissTOM), method = "average")
  dynamicMods = cutreeDynamic(dendro = geneTree, distM = dissTOM,
                              deepSplit = 2, pamRespectsDendro = FALSE,
                              minClusterSize = 50)
  
  # Correlate Modules with Myeloid Scores
  MEs = moduleEigengenes(exp_matrix, colors = dynamicMods)$eigengenes
  moduleTraitCor = cor(MEs, trait_data, use = "p")
  
  return(list(modules = dynamicMods, correlation = moduleTraitCor))
}

# ------------------------------------------------------------------------------
# 3. Prognostic Meta-Analysis
# ------------------------------------------------------------------------------
# Input Data Format:
#   - cohort_list: A list of data frames. Each data frame must contain:
#     "overall_survival" (numeric), "deceased" (0/1), and Gene Expression (Z-score).
#   - gene_symbol: The gene name to be analyzed.
# ------------------------------------------------------------------------------

run_prognostic_meta <- function(cohort_list, gene_symbol) {
  library(survival)
  library(meta)
  
  meta_results <- data.frame()
  
  # Step 1: Calculate Univariate Cox for each cohort
  for (cohort_name in names(cohort_list)) {
    df <- cohort_list[[cohort_name]]
    formula <- as.formula(paste0("Surv(overall_survival, deceased) ~ ", gene_symbol))
    fit <- coxph(formula, data = df)
    
    # Store HR and SE
    sum_fit <- summary(fit)
    meta_results <- rbind(meta_results, data.frame(
      dataset = cohort_name,
      logHR = sum_fit$coefficients[1, 1],
      seLogHR = sum_fit$coefficients[1, 3]
    ))
  }
  
  # Step 2: Random-effects Meta-Analysis
  meta_integration <- metagen(
    TE = logHR, 
    seTE = seLogHR, 
    data = meta_results, 
    studlab = dataset,
    sm = "HR"
  )
  
  return(meta_integration)
}