#' Perform differential expression analysis within each sex
#'
#' This function identifies differentially expressed genes between phenotype groups
#' separately for each sex using limma (default), DESeq2, or edgeR.
#'
#' @param expression_data A numeric matrix of expression data with features in rows
#'   and samples in columns. For \code{method = "limma"}, this can be a normalized
#'   log-expression matrix or a raw count matrix. For \code{method = "deseq2"} or
#'   \code{method = "edger"}, this must be a raw (non-negative integer) count matrix.
#' @param sex A vector indicating sex for each sample.
#' @param phenotype A vector indicating phenotype labels for each sample.
#' @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 specifying phenotype levels. The
#'   first element is the reference, the second the comparison; contrast = 2 - 1
#'   (e.g. c("WT","TG")).
#' @param min_samples Integer. Minimum number of samples per group (default 3).
#' @param method Character. Differential expression method to use. One of
#'   \code{"limma"} (default), \code{"deseq2"}, or \code{"edger"}.
#'
#' @return A list with three elements:
#'   \describe{
#'     \item{male_DEGs}{A data.frame of DE results for males.}
#'     \item{female_DEGs}{A data.frame of DE results for females.}
#'     \item{validation}{Output from \code{validate_input_bulk()}.}
#'   }
#'   Regardless of the method, each DEG data.frame contains the standardized
#'   columns \code{avg_log2FC}, \code{p_val}, \code{p_val_adj}, and \code{gene},
#'   with gene names as rownames. Additional method-specific columns (e.g.
#'   \code{baseMean} for DESeq2, \code{logCPM} for edgeR, \code{AveExpr} for limma)
#'   are retained for reference.
#'
#' @details For each sex, the function subsets the data and fits the appropriate
#'   model:
#'   \itemize{
#'    \item \strong{limma}: Fits a linear model via \code{lmFit} + \code{eBayes}
#'       on the input matrix directly. Suitable for microarray, proteomics, or
#'       pre-normalized RNA-seq data. If raw RNA-seq counts are provided, the
#'       function automatically detects integer matrices and applies
#'       \code{limma-voom} internally with a warning. For full control, consider
#'       using \code{"edger"} or \code{"deseq2"} instead, or apply \code{voom}
#'       externally and pass the transformed data with
#'       \code{method = "limma"}.
#'     \item \strong{DESeq2}: Fits a negative binomial GLM via
#'       \code{DESeq2::DESeq()}. Requires raw integer counts. Performs internal
#'       normalization (median-of-ratios) and dispersion estimation.
#'     \item \strong{edgeR}: Fits a negative binomial GLM via
#'       \code{edgeR::glmQLFit()} + \code{edgeR::glmQLFTest()}. Requires raw
#'       integer counts. Performs TMM normalization internally.
#'   }
#'   Positive logFC values indicate higher expression in the test phenotype
#'   (second element of \code{phenotype_labels_vector}) compared to the reference
#'   (first element).
#'
#' @importFrom stats model.matrix
#' @importFrom limma lmFit contrasts.fit eBayes topTable makeContrasts
#'
#' @export
sex_stratified_analysis_bulk <- function(expression_data,
                                         sex,
                                         phenotype,
                                         sex_labels_vector = c("F", "M"),
                                         phenotype_labels_vector = c("WT", "TG"),
                                         min_samples = 3,
                                         method = c("limma", "deseq2", "edger")) {
  
  method <- match.arg(method)
  
  # Validate input
  validation <- validate_input_bulk(
    expression_data = expression_data,
    sex = sex,
    phenotype = phenotype,
    sex_labels_vector = sex_labels_vector,
    phenotype_labels_vector = phenotype_labels_vector,
    min_samples = min_samples
  )
  
  # Check count matrix requirement for DESeq2 and edgeR
  if (method %in% c("deseq2", "edger")) {
    if (any(expression_data < 0, na.rm = TRUE)) {
      stop(
        "method = '", method, "' requires non-negative count data. ",
        "The input matrix contains negative values. ",
        "Use method = 'limma' for log-transformed or normalized data."
      )
    }
    if (!all(expression_data == floor(expression_data), na.rm = TRUE)) {
      warning(
        "method = '", method, "' expects raw integer counts, but non-integer ",
        "values were detected. This usually means the data has already been ",
        "normalized (e.g. log-transformed, TPM, FPKM). Rounding to nearest ",
        "integer, but results may be unreliable. Consider using ",
        "method = 'limma' for pre-normalized data instead."
      )
      expression_data <- round(expression_data)
    }
    storage.mode(expression_data) <- "integer"
  }
  
  # Convert to factors
  sex <- factor(sex, levels = sex_labels_vector)
  phenotype <- factor(phenotype, levels = phenotype_labels_vector)
  
  male_results <- NULL
  female_results <- NULL
  
  for (sex_label in sex_labels_vector) {
    
    selected_samples <- which(as.character(sex) == sex_label)
    
    if (length(selected_samples) == 0) {
      message("No samples for sex: ", sex_label)
      next
    }
    
    sex_subset <- expression_data[, selected_samples, drop = FALSE]
    phenotype_subset <- phenotype[selected_samples]
    
    if (length(unique(phenotype_subset)) < 2) {
      message("Only one phenotype level found for sex: ", sex_label)
      next
    }
    
    # Dispatch to method-specific backend
    results <- switch(method,
                      limma  = .de_limma(sex_subset, phenotype_subset, phenotype_labels_vector),
                      deseq2 = .de_deseq2(sex_subset, phenotype_subset, phenotype_labels_vector),
                      edger  = .de_edger(sex_subset, phenotype_subset, phenotype_labels_vector)
    )
    

    
    # Assign to sex
    if (sex_label == sex_labels_vector[2]) {
      male_results <- results
    } else {
      female_results <- results
    }
  }
  
  return(list(
    male_DEGs = male_results,
    female_DEGs = female_results,
    validation = validation
  ))
}


# --- Internal DE backends ------------------------------------------------

#' limma backend for sex-stratified DE
#' @noRd
#' @keywords internal
.de_limma <- function(counts, phenotype, phenotype_labels_vector) {
  
  design <- model.matrix(~ 0 + phenotype)
  colnames(design) <- levels(phenotype)
  
  # --- Detect raw counts and apply voom if needed ---
  is_likely_counts <- mean(counts %% 1 == 0, na.rm = TRUE) > 0.99
  
  if (is_likely_counts) {
    warning("Expression data appear to be raw counts. Applying limma-voom.")
    v <- limma::voom(counts, design, plot = FALSE)
    fit <- limma::lmFit(v, design)
  } else {
    fit <- limma::lmFit(counts, design)
  }
  
  # --- Contrast setup ---
  contrast_formula <- paste0(
    phenotype_labels_vector[2], " - ", phenotype_labels_vector[1]
  )
  
  contrast_matrix <- limma::makeContrasts(
    contrasts = contrast_formula,
    levels = design
  )
  
  fit <- limma::contrasts.fit(fit, contrast_matrix)
  fit <- limma::eBayes(fit)
  
  res <- limma::topTable(fit, coef = 1, number = Inf)
  
  # --- Standardize column names ---
  res$avg_log2FC <- res$logFC
  res$p_val      <- res$P.Value
  res$p_val_adj  <- res$adj.P.Val
  res$gene       <- rownames(res)
  
  return(res)
}

#' DESeq2 backend for sex-stratified DE
#' 
#' @noRd
#' @keywords internal
.de_deseq2 <- function(counts, phenotype, phenotype_labels_vector) {
  
  if (!requireNamespace("DESeq2", quietly = TRUE)) {
    stop("Package 'DESeq2' is required for method = 'deseq2'. ",
         "Install it from Bioconductor: BiocManager::install('DESeq2')")
  }
  
  col_data <- data.frame(
    condition = phenotype,
    row.names = colnames(counts)
  )
  
  dds <- DESeq2::DESeqDataSetFromMatrix(
    countData = counts,
    colData   = col_data,
    design    = ~ condition
  )
  
  # Ensure the reference level is set correctly
  dds$condition <- stats::relevel(dds$condition, ref = phenotype_labels_vector[1])
  
  # Run DESeq2 pipeline with fallback for dispersion estimation failures.
  # Small or atypical datasets can cause the parametric dispersion curve fit
  # to fail. In that case we fall back to gene-wise estimates, which is the
  # approach recommended by the DESeq2 authors in the error message itself.
  dds <- tryCatch(
    suppressMessages(DESeq2::DESeq(dds)),
    error = function(e) {
      if (grepl("dispersion", e$message, ignore.case = TRUE)) {
        message(
          "DESeq2: parametric dispersion fit failed. ",
          "Falling back to gene-wise dispersion estimates. ",
          "This is expected for small or atypical datasets."
        )
        dds <- suppressMessages(DESeq2::estimateSizeFactors(dds))
        dds <- suppressMessages(DESeq2::estimateDispersionsGeneEst(dds))
        DESeq2::dispersions(dds) <- S4Vectors::mcols(dds)$dispGeneEst
        dds <- suppressMessages(DESeq2::nbinomWaldTest(dds))
        return(dds)
      } else {
        stop(e)
      }
    }
  )
  
  res <- DESeq2::results(
    dds,
    contrast = c("condition",
                 phenotype_labels_vector[2],
                 phenotype_labels_vector[1])
  )
  
  res <- as.data.frame(res)
  
  # Remove genes with NA p-values (e.g. low count outliers filtered by DESeq2)
  res <- res[!is.na(res$pvalue), ]
  
  # Standardize column names
  res$avg_log2FC <- res$log2FoldChange
  res$p_val      <- res$pvalue
  res$p_val_adj  <- res$padj
  # Genes with extreme Cook's distance get NA padj; fill with 1
  res$p_val_adj[is.na(res$p_val_adj)] <- 1
  res$gene       <- rownames(res)
  
  return(res)
}


#' edgeR backend for sex-stratified DE
#' @noRd
#' @keywords internal
.de_edger <- function(counts, phenotype, phenotype_labels_vector) {
  
  if (!requireNamespace("edgeR", quietly = TRUE)) {
    stop("Package 'edgeR' is required for method = 'edger'. ",
         "Install it from Bioconductor: BiocManager::install('edgeR')")
  }
  
  dge <- edgeR::DGEList(counts = counts, group = phenotype)
  
  # Filter lowly expressed genes
  keep <- edgeR::filterByExpr(dge, group = phenotype)
  dge <- dge[keep, , keep.lib.sizes = FALSE]
  
  # TMM normalization
  dge <- edgeR::calcNormFactors(dge, method = "TMM")
  
  # Design matrix
  design <- model.matrix(~ 0 + phenotype)
  colnames(design) <- levels(phenotype)
  
  # Estimate dispersions and fit quasi-likelihood model
  dge <- edgeR::estimateDisp(dge, design)
  fit <- edgeR::glmQLFit(dge, design)
  
  # Contrast: test vs reference
  contrast_vector <- rep(0, ncol(design))
  names(contrast_vector) <- colnames(design)
  contrast_vector[phenotype_labels_vector[2]] <-  1
  contrast_vector[phenotype_labels_vector[1]] <- -1
  
  qlf <- edgeR::glmQLFTest(fit, contrast = contrast_vector)
  
  res <- edgeR::topTags(qlf, n = Inf, sort.by = "PValue")$table
  
  # Standardize column names
  res$avg_log2FC <- res$logFC
  res$p_val      <- res$PValue
  res$p_val_adj  <- res$FDR
  res$gene       <- rownames(res)
  
  return(res)
}

#' Validate bulk input data for sex-stratified analysis
#'
#' @param expression_data Numeric matrix (features x samples)
#' @param sex Vector of sex labels
#' @param phenotype Vector of phenotype labels
#' @param sex_labels_vector Expected sex labels (length 2)
#' @param phenotype_labels_vector Expected phenotype labels (length 2)
#' @param min_samples Minimum samples per group
#'
#' @return A list with validation status, ratios, imbalance flag, and group counts
#' @export
validate_input_bulk <- function(expression_data,
                                sex,
                                phenotype,
                                sex_labels_vector = c("F", "M"),
                                phenotype_labels_vector = c("WT", "TG"),
                                min_samples = 3) {
  
  # -------------------------------
  # Basic checks
  # -------------------------------
  if (!is.matrix(expression_data)) {
    stop("expression_data must be a matrix.")
  }
  
  if (ncol(expression_data) != length(phenotype)) {
    stop("Number of columns in expression_data must match length of phenotype.")
  }
  
  if (length(sex) != length(phenotype)) {
    stop("Length of sex and phenotype must match.")
  }
  
  if (length(sex_labels_vector) != 2) {
    stop("sex_labels_vector must contain exactly two labels.")
  }
  
  if (length(phenotype_labels_vector) != 2) {
    stop("phenotype_labels_vector must contain exactly two labels.")
  }
  
  # -------------------------------
  # Factor conversion (no shadowing)
  # -------------------------------
  sex_vec <- factor(sex, levels = sex_labels_vector)
  phenotype_vec <- factor(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_samples <- sum(sex_counts)
  sex_ratio <- sex_counts / total_samples
  
  ratio_diff <- if (length(sex_ratio) == 2) abs(diff(sex_ratio)) else 0
  is_imbalanced <- ratio_diff > 0.30
  
  # -------------------------------
  # Group counts (robust)
  # -------------------------------
  group <- interaction(phenotype_vec, sex_vec, drop = FALSE)
  group_counts <- table(group)
  
  low_sample_groups <- group_counts[group_counts < min_samples]
  
  # -------------------------------
  # Warnings
  # -------------------------------
  if (is_imbalanced) {
    warning(
      paste(
        "Sex imbalance detected (>30%). Ratios:",
        paste(names(sex_ratio), round(sex_ratio, 2), collapse = ", "),
        "Proceeding despite sex imbalance..."
      )
    )
  }
  
  if (length(low_sample_groups) > 0) {
    warning(
      paste(
        "Some groups have fewer than", min_samples, "samples:",
        paste(names(low_sample_groups), low_sample_groups, sep = ":", collapse = ", "),
        "Sex-analysis may be unreliable."
      )
    )
  }
  
  # -------------------------------
  # Return
  # -------------------------------
  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_sample_groups = low_sample_groups
  ))
}