#' Prepare Term Data for Enrichment Analysis

 #'

 #' Process term data downloaded with the \code{fetch_*} functions, preparing it
 #' for fast enrichment analysis using \code{functional_enrichment}.

 #'

 #' @details
 #'

 #' This function takes two tibbles containing functional term information
 #' (\code{terms}) and feature mapping (\code{mapping}), and converts them into
 #' an object required by \code{functional_enrichment} for efficient analysis.
 #' Terms and mapping can be generated with the database access functions
 #' included in this package, such as \code{fetch_reactome} or
 #' \code{fetch_go_from_go}.

 #'

 #' @param terms A tibble with at least two columns: \code{term_id} and

 #'   \code{term_name}. This tibble contains information about functional term
 #'   names and descriptions.
 #' @param mapping A tibble with at least two columns, containing the mapping
 #'   between functional terms and features. One column must be named
 #'   \code{term_id}, while the other column should have a name specified by the
 #'   \code{feature_name} argument. For example, if \code{mapping} contains
 #'   columns \code{term_id}, \code{accession_number}, and \code{gene_symbol},
 #'   setting \code{feature_name = "gene_symbol"} indicates that gene symbols
 #'   will be used for enrichment analysis.
 #' @param all_features A vector with all feature IDs used as the background for

 #'   enrichment. If not specified, all features found in \code{mapping} will be
 #'   used, resulting in a larger object size.

 #' @param feature_name The name of the column in the \code{mapping} tibble to be
 #'   used as the feature identifier. For example, if \code{mapping} contains
 #'   columns \code{term_id}, \code{accession_number}, and \code{gene_symbol},
 #'   setting \code{feature_name = "gene_symbol"} indicates that gene symbols
 #'   will be used for enrichment analysis.

 #'

 #' @return An object of class \code{fenr_terms} required by

 #'   \code{functional_enrichment}.

 #' @importFrom assertthat assert_that

 #' @export

 #' @examples

 #' \dontrun{

 #' data(exmpl_all)

 #' go <- fetch_go(species = "sgd")

 #' go_terms <- prepare_for_enrichment(go$terms, go$mapping, exmpl_all,
 #'                                    feature_name = "gene_symbol")

 #' }

 prepare_for_enrichment <- function(terms, mapping, all_features = NULL,
                                    feature_name = "gene_id") {

   # Binding variables from non-standard evaluation locally
   feature_id <- term_id <- NULL
 

   # Argument checks

   assert_that(is.data.frame(terms) || tibble::is_tibble(terms),
               msg = "'terms' must be a data frame or tibble.")

   assert_that(is.data.frame(mapping) || tibble::is_tibble(mapping),
               msg = "'mapping' must be a data frame or tibble.")

   assert_that(is.null(all_features) || is.vector(all_features),
               msg = "'all_features' must be a vector or NULL.")

   assert_that(is.character(feature_name) && length(feature_name) == 1,
               msg = "'feature_name' must be a single string.")

   # Check terms

   assert_that(all(c("term_id", "term_name") %in% colnames(terms)),
               msg = "Column names in 'terms' should be 'term_id' and 'term_name'.")
 
   assert_that(anyDuplicated(terms$term_id) == 0,
               msg = "Duplicated term_id detected in 'terms'.")

 
   # Check mapping

   assert_that("term_id" %in% colnames(mapping),
               msg = "'mapping' should contain a column named 'term_id'.")

 
   # Check for feature name

   assert_that(feature_name %in% colnames(mapping),

               msg = paste0(feature_name, " column not found in mapping table.
                            Check 'feature_name' argument."))

 
   # Replace empty all_features with everything from mapping
   map_features <- mapping[[feature_name]] |>
     unique()
   if (is.null(all_features)) {
     all_features <- map_features
   } else {
     # Check if mapping is contained in all features
     if (length(intersect(all_features, map_features)) == 0)

       stop("No overlap between 'all_features' and features found in 'mapping'.
            Did you provide correct 'all_features'?")

   }
 
   # Check for missing term descriptions

   mis_term <- setdiff(mapping$term_id, terms$term_id)

   if (length(mis_term) > 0) {
     dummy <- tibble::tibble(
       term_id = mis_term,
       term_name = rep(NA_character_, length(mis_term))
     )

     terms <- dplyr::bind_rows(terms, dummy)

   }
 

   # Hash to select term name

   term2name <- new.env(hash = TRUE)

   for (i in seq_len(nrow(terms))) {

     r <- terms[i, ]
     term2name[[r$term_id]] <- r$term_name
   }

 
   # feature-term tibble
   feature_term <- mapping |>
     dplyr::rename(feature_id = !!feature_name) |>

     dplyr::filter(feature_id %in% all_features) |>

     dplyr::select(feature_id, term_id) |>
     dplyr::distinct()

   # Feature to terms hash
   f2t <- feature_term |>

     dplyr::group_by(feature_id) |>

     dplyr::summarise(terms = list(term_id)) |>
     tibble::deframe()

   feature2term <- new.env(hash = TRUE)

   for(feat in names(f2t))

     feature2term[[feat]] <- f2t[[feat]]

 
   # Term to feature hash
   t2f <- feature_term |>

     dplyr::group_by(term_id) |>

     dplyr::summarise(features = list(feature_id)) |>
     tibble::deframe()

   term2feature <- new.env(hash = TRUE)

   for(term in names(t2f))

     term2feature[[term]] <- t2f[[term]]

 
   list(
     term2name = term2name,
     term2feature = term2feature,
     feature2term = feature2term

   ) |>

     structure(class = "fenr_terms")

 }
 
 

 
 #' Fast Functional Enrichment

 #'

 #' Perform fast functional enrichment analysis based on the hypergeometric
 #' distribution. Designed for use in interactive applications.

 #'

 #' @details This function carries out functional enrichment analysis on a
 #'   selection of features (e.g., differentially expressed genes) using the
 #'   hypergeometric probability distribution (Fisher's exact test). Features can
 #'   be genes, proteins, etc. The \code{term_data} object contains functional
 #'   term information and feature-term mapping.

 #'

 #' @param feat_all A character vector with all feature identifiers, serving as
 #'   the background for enrichment.

 #' @param feat_sel A character vector with feature identifiers in the selection.

 #' @param term_data An object of class \code{fenr_terms}, created by

 #'   \code{prepare_for_enrichment}.

 #' @param feat2name An optional named list to convert feature IDs into feature

 #'   names.

 #'

 #' @return A tibble with enrichment results, providing the following information
 #'   for each term:
 #'   \itemize{
 #'     \item{\code{N_with} - number of features with this term among all features}
 #'     \item{\code{n_with_sel} - number of features with this term in the selection}
 #'     \item{\code{n_expect} - expected number of features with this term in the selection,
 #'       under the null hypothesis that terms are mapped to features randomly}
 #'     \item{\code{enrichment} - ratio of n_with_sel / n_expect}
 #'     \item{\code{odds_ratio} - odds ratio for enrichment; is infinite when all
 #'       features with the given term are in the selection}
 #'     \item{\code{p_value} - p-value from a single hypergeometric test}

 #'     \item{\code{p_adjust} - p-value adjusted for multiple tests using the
 #'       Benjamini-Hochberg approach}

 #'   }.

 #'

 #' @importFrom assertthat assert_that

 #' @importFrom methods is
 #' @export

 #' @examples

 #' \dontrun{

 #' data(exmpl_all, exmpl_sel)

 #' go <- fetch_go(species = "sgd")
 #' go_terms <- prepare_for_enrichment(go$terms, go$mapping, exmpl_all, feature_name = "gene_symbol")
 #' enr <- functional_enrichment(exmpl_all, exmpl_sel, go_terms)

 #' }

 functional_enrichment <- function(feat_all, feat_sel, term_data, feat2name = NULL) {

   # Binding variables from non-standard evaluation locally
   N_with <- n_with_sel <- n_expect <- enrichment <- odds_ratio <- NULL
   desc <- p_value <- p_adjust <- NULL
 

   # Check for character vectors
   assert_that(is.character(feat_all))
   assert_that(is.character(feat_sel))
   assert_that(length(feat_all) > 1)
   assert_that(length(feat_sel) > 1)
 

   # Check term_data class

   assert_that(is(term_data, "fenr_terms"))

   # If no overlap between selection and all, return NULL
   if(!any(feat_sel %in% feat_all))
     return(NULL)
 

   # all terms present in the selection
   our_terms <- feat_sel |>

     purrr::map(~term_data$feature2term[[.x]]) |>

     unlist() |>
     unique()

   # number of features in selection
   N_sel <- length(feat_sel)
   # total number of features
   N_tot <- length(feat_all)
 
   res <- purrr::map_dfr(our_terms, function(term_id) {
     # all features with the term

     # term_data$term2feature is a hash environment
     tfeats <- term_data$term2feature[[term_id]]

     # necessary if term data contain features not present in feat_all
     tfeats <- tfeats[tfeats %in% feat_all]

 
     # features from selection with the term

     # this is faster than intersect(tfeats, feat_sel)
     tfeats_sel <- tfeats[tfeats %in% feat_sel]

 
     N_with <- length(tfeats)
     N_without <- N_tot - N_with
 
     # building contingency table
     n_with_sel <- length(tfeats_sel)
     n_without_sel <- N_sel - n_with_sel
     n_with_nsel <- N_with - n_with_sel
     n_without_nsel <- N_tot - (n_with_sel + n_without_sel + n_with_nsel)
 
     # contingency table is
     #               | In selection  | Not in selection
     #-------------------------------------------------
     #  With term    | n_with_sel    | n_with_nsel
     #  Without term | n_without_sel | n_without_nsel
 

     if (n_with_sel < 2) return(NULL)

 
     # Expected number of features in selection, if random
     n_expect <- N_with * N_sel / N_tot
 
     # Odds ratio
     odds_ratio <- (n_with_sel / n_without_sel) / (n_with_nsel / n_without_nsel)
 
     # Hypergeometric function much faster than fisher.test
     p <- 1 - stats::phyper(n_with_sel - 1, N_with, N_without, N_sel)
 

     # Convert feature IDs to feature names;

     if (!is.null(feat2name))

       tfeats_sel <- feat2name[tfeats_sel] |> unname()

     term_name <- term_data$term2name[[term_id]]

     # returns NAs if no term found
     if (is.null(term_name)) term_name <- NA_character_
 
     # constructing a tibble in every iteration is more time expensive than `c`,
     # even with overhead of converting types afterwards. Not elegant, but fast.
     c(
       term_id = term_id,
       term_name = term_name,
       N_with = N_with,
       n_with_sel = n_with_sel,
       n_expect = n_expect,
       enrichment = n_with_sel / n_expect,
       odds_ratio = odds_ratio,
       ids = paste(tfeats_sel, collapse = ", "),
       p_value = p
     )
   })

   # Drawback - if all selections below minimum, res is tibble 0 x 0, need to
   # catch it

   if (nrow(res) == 0) {
     res <- NULL
   } else {
     res <- res |>

       dplyr::mutate(

         dplyr::across(c(N_with, n_with_sel), as.integer),
         dplyr::across(c(n_expect, enrichment, odds_ratio, p_value), as.numeric),
         p_adjust = stats::p.adjust(p_value, method = "BH"),
         dplyr::across(c(enrichment, odds_ratio, p_value, p_adjust), ~signif(.x, 3)),

         n_expect = round(n_expect, 2)
       ) |>

       dplyr::arrange(desc(odds_ratio))

   }

   res

 }