#' Perform Pathway Enrichment Analysis for Pre-Categorized DEGs
#'
#' @description
#' Performs pathway enrichment analysis for categorized DEGs using GO, KEGG, Reactome, or a custom TERM2GENE database.
#'
#' @param DEGs_category Data frame with columns 'DEG_Type' and 'Gene_Symbols'.
#' @param enrichment_db "KEGG", "GO", "REACTOME", or "CUSTOM" (default: "KEGG").
#' @param organism Organism code for KEGG (default: "hsa").
#' @param gene_type Gene identifier type supported by OrgDb
#' (e.g., "SYMBOL", "ENTREZID", "ENSEMBL", "UNIPROT").
#' @param org_db OrgDb object for gene annotation (default: org.Hs.eg.db).
#' @param custom_db TERM2GENE data frame, required if enrichment_db = "CUSTOM".
#' @param pvalueCutoff Numeric p-value cutoff (default: 0.05).
#' @param qvalueCutoff Numeric q-value cutoff (default: 0.2).
#' @param return_df Logical, if TRUE returns data.frames, else enrichResult objects (default: FALSE).
#'
#' @return Named list of enrichment results per DEG_Type.
#' @details Gene IDs converted only when needed; KEGG/Reactome require ENTREZID; GO/CUSTOM use provided gene_type.
#' @importFrom dplyr %>%
#' @importFrom AnnotationDbi keytypes
#' @export

categorized_enrich <- function(
    DEGs_category,
    enrichment_db = "KEGG",
    organism = "hsa",
    gene_type = "SYMBOL",
    org_db = NULL,
    custom_db = NULL,
    pvalueCutoff = 0.05,
    qvalueCutoff = 0.2,
    return_df = FALSE
) {
  if (is.null(org_db)) {
    if (!requireNamespace("org.Hs.eg.db", quietly = TRUE)) {
      stop(
        "Please provide an OrgDb or install 'org.Hs.eg.db' via:\n",
        "BiocManager::install('org.Hs.eg.db')"
      )
    }
    org_db <- org.Hs.eg.db::org.Hs.eg.db
  }
  db <- toupper(enrichment_db)
  gene_type <- toupper(gene_type)
  
  # Input validation
  if (!all(c("DEG_Type", "Gene_Symbols") %in% colnames(DEGs_category))) {
    stop("DEGs_category must contain columns 'DEG_Type' and 'Gene_Symbols'.")
  }
  
  if (!gene_type %in% AnnotationDbi::keytypes(org_db)) {
    stop(
      paste(
        "Unsupported gene_type. Supported types include:",
        paste(AnnotationDbi::keytypes(org_db), collapse = ", ")
      )
    )
  }
  
  if (!is.null(custom_db)) {
    if (!is.data.frame(custom_db) ||
        !all(c("term", "gene") %in% colnames(custom_db))) {
      stop("custom_db must be a data frame with columns: term, gene.")
    }
  }
  
  # Standardize DEG_Type
  all_categories <- unique(DEGs_category$DEG_Type)
  
  results <- list()
  
  # ---- NEW: helper-based conversion ----
  
  # For KEGG / Reactome → convert once to ENTREZ
  if (db %in% c("KEGG", "REACTOME") && gene_type != "ENTREZID") {
    DEGs_category <- convert_gene_ids(
      DEGs_category,
      fromType = gene_type,
      toType = "ENTREZID",
      org_db = org_db,
      gene_col = "Gene_Symbols"
    )
    gene_type <- "ENTREZID"
  }
  
  # -------------------------------------
  
  for (type in all_categories) {
    
    genes <- DEGs_category[
      DEGs_category$DEG_Type == type,
      "Gene_Symbols"
    ] %>% na.omit() %>% unique()
    
    if (length(genes) == 0) next
    
    # GO
    if (db == "GO") {
      
      enrichment_result <- clusterProfiler::enrichGO(
        gene = genes,
        OrgDb = org_db,
        keyType = gene_type,
        ont = "BP",   # lighter than ALL
        pvalueCutoff = pvalueCutoff,
        qvalueCutoff = qvalueCutoff
      )
    }
    
    # KEGG
    else if (db == "KEGG") {
      
      enrichment_result <- clusterProfiler::enrichKEGG(
        gene = genes,
        organism = organism,
        pvalueCutoff = pvalueCutoff,
        qvalueCutoff = qvalueCutoff
      )
    }
    
    # Reactome
    else if (db == "REACTOME") {
      
      enrichment_result <- ReactomePA::enrichPathway(
        gene = genes,
        organism = "human",
        pvalueCutoff = pvalueCutoff,
        qvalueCutoff = qvalueCutoff
      )
    }
    
    # Custom
    else if (db == "CUSTOM") {
      
      if (is.null(custom_db)) {
        stop("custom_db must be provided when enrichment_db = 'CUSTOM'.")
      }
      
      enrichment_result <- clusterProfiler::enricher(
        gene = genes,
        TERM2GENE = custom_db,
        pvalueCutoff = pvalueCutoff,
        qvalueCutoff = qvalueCutoff
      )
    }
    
    else {
      stop("Unsupported enrichment_db.")
    }
    
    # Store
    results[[type]] <- if (return_df) {
      as.data.frame(enrichment_result)
    } else {
      enrichment_result
    }
  }
  
  return(results)
}


#' Convert gene identifiers in a DEG table
#'
#' @param df Data frame containing a gene column
#' @param fromType Original ID type (e.g. "UNIPROT", "ENSEMBL")
#' @param toType Target ID type (e.g. "SYMBOL", "ENTREZID")
#' @param org_db OrgDb object (default: org.Hs.eg.db)
#' @param gene_col Column name containing gene IDs (default: "Gene_Symbols")
#' @param remove_na Remove unmapped genes (default: TRUE)
#' @param remove_duplicates Remove duplicated mapped IDs (default: TRUE)
#'
#' @return Updated data frame with converted gene IDs
#' @export
convert_gene_ids <- function(
    df,
    fromType,
    toType = "SYMBOL",
    org_db = NULL,
    gene_col = "Gene_Symbols",
    remove_na = TRUE,
    remove_duplicates = TRUE
) {
  
  if (!gene_col %in% colnames(df)) {
    stop(paste("Column", gene_col, "not found in data frame"))
  }
  
  supported <- AnnotationDbi::keytypes(org_db)
  
  if (is.null(org_db)) {
    if (!requireNamespace("org.Hs.eg.db", quietly = TRUE)) {
      stop(
        "Please provide an OrgDb or install 'org.Hs.eg.db' via:\n",
        "BiocManager::install('org.Hs.eg.db')"
      )
    }
    org_db <- org.Hs.eg.db::org.Hs.eg.db
  }
  if (!fromType %in% supported) {
    stop(paste("fromType not supported. Choose from:", paste(supported, collapse = ", ")))
  }
  
  if (!toType %in% supported) {
    stop(paste("toType not supported. Choose from:", paste(supported, collapse = ", ")))
  }
  
  genes <- unique(df[[gene_col]])
  n_input <- length(genes)
  
  mapping <- clusterProfiler::bitr(
    genes,
    fromType = fromType,
    toType = toType,
    OrgDb = org_db
  )
  
  n_mapped <- length(unique(mapping[[fromType]]))
  unmapped <- setdiff(genes, mapping[[fromType]])
  
  message("Gene ID conversion summary:")
  message(" - Input genes: ", n_input)
  message(" - Mapped genes: ", n_mapped)
  message(" - Unmapped genes: ", length(unmapped))
  message(" - Mapping rate: ", round(n_mapped / n_input * 100, 2), "%")
  
  df <- merge(
    df,
    mapping,
    by.x = gene_col,
    by.y = fromType,
    all.x = TRUE
  )
  
  df[[gene_col]] <- df[[toType]]
  
  if (remove_na) {
    df <- df[!is.na(df[[gene_col]]), ]
  }
  
  if (remove_duplicates) {
    df <- df[!duplicated(df[[gene_col]]), ]
  }
  
  df[[toType]] <- NULL
  
  return(df)
}
