#' Perform sex-stratified differential expression analysis for single-cell Seurat object
#'
#' @param seurat_obj A Seurat object containing single-cell RNA-seq data.
#' @param sex Character. Name of the column in metadata indicating sex.
#' @param phenotype Character. Name of the column indicating phenotype/condition.
#' @param celltype_col Character. Name of the column indicating cell type.
#' @param sex_labels_vector A character vector of length two defining the labels for sex (e.g., c("F", "M")). The first element must correspond to "female" and the second to "male". 
#' @param phenotype_labels_vector Character vector of length 2 specifying the reference and test phenotypes (e.g., c("WT", "TG")).
#' @param test.use Statistical test to use in FindMarkers.
#' @param min_samples Integer. Minimum number of cells per group (default 3).
#' @return A list with two elements: `male_DEGs` and `female_DEGs`, each containing a list of DEGs per cell type.
#' 
#' @importFrom Seurat GetAssayData FindMarkers 
#' @importFrom SeuratObject DefaultAssay DefaultAssay<- Idents Idents<-
#' @export
sex_stratified_analysis_sc <- function(seurat_obj,
                                       sex = "sex",
                                       phenotype = "status",
                                       celltype_col = "cell_type",
                                       sex_labels_vector = c("F", "M"),
                                       phenotype_labels_vector = c("WT", "TG"),
                                       test.use = "wilcox",
                                       min_samples = 3){
  
  # Use SC validation function 
  validation <- validate_input_sc(
    seurat_obj = seurat_obj,
    sex = sex,
    phenotype = phenotype,
    celltype_col = celltype_col,
    sex_labels_vector = sex_labels_vector,
    phenotype_labels_vector = phenotype_labels_vector,
    min_cells = min_samples
  )
  
  #Prepare metadata factors
  metadata <- seurat_obj@meta.data
  metadata[[sex]] <- factor(metadata[[sex]], levels = sex_labels_vector)
  metadata[[phenotype]] <- factor(metadata[[phenotype]], levels = phenotype_labels_vector)
  metadata[[celltype_col]] <- as.factor(metadata[[celltype_col]])
  
  # Create combined group
  metadata$Group <- factor(paste(metadata[[phenotype]], metadata[[sex]], sep="_"))
  
  # Initialize result lists
  male_results <- list()
  female_results <- list()
  
  # Loop over cell types
  cell_types <- unique(metadata[[celltype_col]])
  for (cell in cell_types) {
    message("Processing cell type: ", cell)
    
    # Subset by cell type
    Idents(seurat_obj) <- celltype_col
    cell_subset <- subset(seurat_obj, idents = cell)

    for (sex_value in sex_labels_vector) {
      # Subset by sex
      Idents(cell_subset) <- sex
      sex_subset <- subset(cell_subset, idents = sex_value)
      
      # Skip if not enough samples per phenotype
      if (sum(sex_subset@meta.data[[phenotype]] == phenotype_labels_vector[1]) < min_samples ||
          sum(sex_subset@meta.data[[phenotype]] == phenotype_labels_vector[2]) < min_samples) {
        next
      }
      
      # Try-catch for FindMarkers
      tryCatch({
        markers <- FindMarkers(
          object = sex_subset,
          group.by = phenotype,
          ident.1 = phenotype_labels_vector[2],
          ident.2 = phenotype_labels_vector[1],
          only.pos = FALSE,
          logfc.threshold = 0,
          test.use = test.use
        )
        markers[[celltype_col]] <- cell
        markers[[sex]] <- sex_value
        markers$gene <- rownames(markers)
        markers$direction <- ifelse(markers$avg_log2FC < 0, "down", "up")
        
        # Store results
        if (sex_value == sex_labels_vector[2]) {
          male_results[[cell]] <- markers
        } else {
          female_results[[cell]] <- markers
        }
      }, error = function(e) {
        message("Error processing cell type ", cell, " and sex ", sex_value, ": ", e$message)
      })
    }
  }
  
  return(list(
    male_DEGs = male_results,
    female_DEGs = female_results,
    validation = validation 
  ))
}

#' Run sex stratified differential expression pipeline for single cell data
#'
#' @description Performs sex specific differential expression analysis per cell type
#' and categorizes genes into male specific, female specific, sex dimorphic,
#' and sex neutral groups.
#'
#' @param seurat_obj Seurat object containing single cell data.
#' @param target_cell_type_flag Character. Cell type to analyze. 
#' @param sex Character. Metadata column for sex.
#' @param phenotype Character. Metadata column for phenotype.
#' @param celltype_col Character. Metadata column for cell type.
#' @param sex_labels_vector Character vector specifying sex levels.
#' @param phenotype_labels_vector Character vector specifying phenotype levels.
#' @param test.use Character. Statistical test used in FindMarkers.
#' @param alpha Numeric. FDR threshold for significance.
#' @param beta Numeric. P value threshold for exclusion in opposite sex.
#' @param min_abs_logfc Numeric. Minimum absolute log2 fold change for categorization.
#' @param min_samples Integer. Minimum number of cells per group (default 3).
#'
#' @return A named list of categorized differential expression results per cell type.
#'
#' @export

sex_stratified_pipeline_sc <- function(seurat_obj,
                                       target_cell_type_flag = "all",
                                       sex = "sex",
                                       phenotype = "status",
                                       celltype_col = "cell_type",
                                       sex_labels_vector = c("F", "M"),
                                       phenotype_labels_vector = c("WT", "TG"),
                                       test.use = "wilcox", 
                                       alpha = 0.05,
                                       beta = 0.5,
                                       min_abs_logfc = 0.25,
                                       min_samples = 3) {
  
  # Validate metadata columns
  required_columns <- c(sex, phenotype, celltype_col)
  if (!all(required_columns %in% colnames(seurat_obj@meta.data))) {
    missing_cols <- required_columns[!required_columns %in% colnames(seurat_obj@meta.data)]
    stop(paste("Missing metadata columns:", paste(missing_cols, collapse = ", ")))
  }
  
  # Optional cell type subsetting
  if (!("all" %in% target_cell_type_flag)) {
    
    valid_celltypes <- unique(seurat_obj@meta.data[[celltype_col]])
    
    if (!all(target_cell_type_flag %in% valid_celltypes)) {
      missing_ct <- target_cell_type_flag[!target_cell_type_flag %in% valid_celltypes]
      stop(paste("Cell type(s) not found:", paste(missing_ct, collapse = ", ")))
    }
    
    # Subset metadata
    cell_meta <- seurat_obj@meta.data[
      seurat_obj@meta.data[[celltype_col]] %in% target_cell_type_flag, ,
      drop = FALSE
    ]
    
    # Subset Seurat object
    seurat_obj <- subset(seurat_obj, cells = rownames(cell_meta))
  }
  
  # Run stratified DE analysis
  results <- sex_stratified_analysis_sc(
    seurat_obj = seurat_obj,
    sex = sex,
    phenotype = phenotype,
    celltype_col = celltype_col,
    sex_labels_vector = sex_labels_vector,
    phenotype_labels_vector = phenotype_labels_vector,
    test.use = test.use,
    min_samples = min_samples
  )
  
  # Extract cell types AFTER subsetting
  cell_types <- unique(seurat_obj@meta.data[[celltype_col]])
  
  # Build DEG list per cell type
  sex_degs <- lapply(cell_types, function(cell) {
    list(
      male   = results$male_DEGs[[cell]],
      female = results$female_DEGs[[cell]]
    )
  })
  
  names(sex_degs) <- cell_types
  
  # Categorize DEGs
  result_categories <- lapply(sex_degs, function(degs) {
    
    if (is.null(degs$male) || is.null(degs$female)) {
      message("Skipping cell type due to missing DEG results")
      return(NULL)
    }
    
    categorize_sex(
      degs$male,
      degs$female,
      alpha = alpha,
      beta = beta,
      min_abs_logfc = min_abs_logfc
    )
  })
  
  # Remove NULL entries
  result_categories <- Filter(Negate(is.null), result_categories)
  
  return(result_categories)
}


#' Validate single-cell input data for sex-stratified analysis
#'
#' @param seurat_obj Seurat object
#' @param sex Column name for sex
#' @param phenotype Column name for phenotype
#' @param celltype_col Column name for cell type
#' @param sex_labels_vector Expected sex labels (length 2)
#' @param phenotype_labels_vector Expected phenotype labels (length 2)
#' @param min_cells Minimum number of cells per group
#'
#' @return A list with validation status, ratios, imbalance flag, and group counts
#' @export
validate_input_sc <- function(seurat_obj,
                              sex = "sex",
                              phenotype = "status",
                              celltype_col = "cell_type",
                              sex_labels_vector = c("F", "M"),
                              phenotype_labels_vector = c("WT", "TG"),
                              min_cells = 3) {
  
  meta <- seurat_obj@meta.data
  

  # Check required columns

  required_columns <- c(sex, phenotype, celltype_col)
  if (!all(required_columns %in% colnames(meta))) {
    missing_cols <- required_columns[!required_columns %in% colnames(meta)]
    stop(paste("Missing metadata columns:", paste(missing_cols, collapse = ", ")))
  }
  

  # Factor conversion 

  sex_vec <- factor(meta[[sex]], levels = sex_labels_vector)
  phenotype_vec <- factor(meta[[phenotype]], levels = phenotype_labels_vector)
  
  if (any(is.na(sex_vec))) {
    stop("Some sex values are not in sex_labels_vector.")
  }
  
  if (any(is.na(phenotype_vec))) {
    stop("Some phenotype values are not in phenotype_labels_vector.")
  }
  
  # Sex ratio

  sex_counts <- table(sex_vec)
  total_cells <- sum(sex_counts)
  sex_ratio <- sex_counts / total_cells
  
  ratio_diff <- if (length(sex_ratio) == 2) abs(diff(sex_ratio)) else 0
  is_imbalanced <- ratio_diff > 0.30
  
 
  # Group counts (robust to missing groups)

  # Ensures all phenotype-sex combinations exist (even with 0 cells)
  group <- interaction(phenotype_vec, sex_vec, drop = FALSE)
  group_counts <- table(group)
  
  low_cell_groups <- group_counts[group_counts < min_cells]
  
  # Sparsity calculation

  assay <- seurat_obj[["RNA"]]
  
  counts <- if (inherits(assay, "Assay5")) {
    # Seurat v5
    if ("counts" %in% names(assay@layers)) {
      assay@layers[["counts"]]
    } else {
      stop("No 'counts' layer found in Assay5 object.")
    }
  } else if (inherits(assay, "Assay")) {
    # Seurat v4
    assay@counts
  } else {
    stop("Unsupported assay type.")
  }
  
  if (inherits(counts, "dgCMatrix")) {
    zero_entries <- length(counts) - length(counts@x[counts@x != 0]) 
    total_entries <- prod(dim(counts)) 
    
  } else {
    
    zero_entries <- sum(counts == 0) 
    total_entries <- length(counts)
    
  }
  
  sparsity <- (zero_entries / total_entries) * 100

  # Warnings

  if (is_imbalanced) {
    warning(
      paste(
        "Sex imbalance detected (>30%). Ratios:",
        paste(names(sex_ratio), round(sex_ratio, 2), collapse = ", "),
        "Proceeding despite imbalance..."
      )
    )
  }
  
  if (length(low_cell_groups) > 0) {
    warning(
      paste(
        "Some groups have fewer than", min_cells, "cells:",
        paste(names(low_cell_groups), low_cell_groups, sep = ":", collapse = ", "),
        "Sex-stratified analysis may be unreliable."
      )
    )
  }
  
  # Return results
  return(list(
    valid = TRUE,
    sex_counts = sex_counts,
    sex_ratio = round(sex_ratio, 3),
    imbalance = is_imbalanced,
    ratio_difference = round(ratio_diff, 3),
    group_counts = group_counts,
    low_cell_groups = low_cell_groups,
    sparsity_percent = round(sparsity, 2)
  ))
}
