# Functions for performing differential expression analysis
#
# Author: vidal
###############################################################################

#' Function to compute linear model fit and optionally apply 'voom' beforehand
#' @param counts Count matrix.
#' @param d Passed to \code{voom.fun}: a numeric matrix containing raw counts, or an ExpressionSet containing raw counts,
#'   or a DGEList object.
#' @param design Numeric design matrix.
#' @param do.voom Logical; should data be "voomed"?
#' @param voom.fun Voom function: one of "voom", "voomWithQualityWeights" or voomLmFit.
#' @param use_weights \code{logical}. Should sample-specific quality weights be estimated?
#' @param norm.method Character; one of "quantile", "tmm".
#' @param quasi.likelihood Logical; should quasi-likelihood methods be used? See \emph{Details} section.
#'   If \code{TRUE}, the default, then the quasi-likelihood methods are used, if \code{FALSE}, then the
#'   likelihood ratio test is used. Please note the main wrapper function \code{diffExpr()} will set this parameter.
#' @param bayes.trend Logical passed to \code{eBayes()}: should an intensity-dependent trend be allowed for the prior variance?
#'   If FALSE then the prior variance is constant.
#'   Alternatively, trend can be a row-wise numeric vector, which will be used as the covariate for the prior variance.
#' @param bayes.robust Logical passed to \code{eBayes()}: should the estimation of df.prior and var.prior be robustified
#'   against outlier sample variances?
#' @param pairs Factor of identifiers specifying paired samples for paired or other block designs, or batch effects.
#' @param pairs_col \code{character}. Name of the column in 'samp.info' containing paired sample information.
#' @param block \code{logical}. Are the comparisons to be made within AND between subjects? See Details section.
#' @param contrasts Contrast matrix as generated by \code{makeContrasts()}. If \code{NULL}, the default, the design matrix
#'   is expected to describe simple pair-wise comparisons, i.e., with the contrasts already included in the matrix.
#'   Consequently, in that case, no contrasts will be fitted.
#' @param disp Character; one of "tagwise.dispersion", "trended.dispersion", "bin.dispersion"
#' @details For experimental designs involving comparisons within as well as between subjects inter-subject needs to be computed.
#'     In this case, the column specified in the 'pairs' argument must assign the subjects to the treatment/tissue/etc groups.
#'     For example, if we have two treatments the effects of which are to be observed in each two tissues, this design would apply.
#'     The 'pairs' factor is passed to the functions \code{duplicateCorrelation()} and \code{lmFit()}.
#'     Quasi-likelihood pipeline:
#'     While the likelihood ratio test is a more obvious choice for inferences with GLMs, the QL
#'     F-test is preferred as it reflects the uncertainty in estimating the dispersion for each gene. It
#'     provides more robust and reliable error rate control when the number of replicates is small.
#'     The QL dispersion estimation and hypothesis testing is done by using the functions
#'     glmQLFit() and glmQLFTest().
#' @return A named \code{list}. If \code{do.voom=TRUE} the elements are \code{v} (the \code{EList} of
#'   voom-transformed counts), \code{fit} (the \code{MArrayLM} linear model fit) and \code{fit2} (the
#'   \code{MArrayLM} fit after \code{eBayes()}, with contrasts applied if these were supplied). Otherwise
#'   they are \code{d} (the input \code{DGEList}), \code{d2} (the \code{DGEList} with estimated
#'   dispersions) and \code{fit} (the \code{DGEGLM} generalised linear model fit).
#' @examples
#' si <- diff_expr_get_samp_info(diffwrap_samp_info, "SampleName", "Group")
#' counts <- diff_expr_filter_counts(diff_expr_read_counts(diffwrap_counts, si), si)
#' groups <- stats::relevel(si$Groups, ref = "control")
#' d <- edgeR::calcNormFactors(edgeR::DGEList(counts, group = groups))
#' design <- diff_expr_make_design(si, groups)
#' fit <- diff_expr_fit(counts, d, design, do.voom = FALSE, quasi.likelihood = TRUE)
#' class(fit$fit)
#' @export
diff_expr_fit <-
  function(counts, d, design, do.voom=TRUE, voom.fun=edgeR::voomLmFit, norm.method=c("quantile", "tmm"), use_weights=FALSE, quasi.likelihood=TRUE,
           bayes.trend=FALSE, bayes.robust=FALSE, pairs=NULL, pairs_col=NULL, block=FALSE, contrasts=NULL, disp="tagwise.dispersion")
  {
    if (do.voom) {
      if (norm.method=="tmm") {
        cts <- d
        norm.method="none"
        # TMM normalisation factors are calculated and added to the DGEList object. If voom is used then counts should usually not be renormalised,
        # so norm.method is set to "none", but "any normalization factors found in 'cts' will still be used even if normalize.method="none".
        # Gordon Smyth notes:
        # I wouldn't recommend TMM before quantile normalization, but it won't make much difference. If you apply scale normalization (such as TMM)
        # then follow it up with quantile normalization of the logCPM values, then the scale normalization will be overwritten. Scale normalization
        # becomes essentially an additive constant on the log-scale for each library, so it will be removed by quantile normalization.
        # We have never recommended using multiple normalization methods together as it confuses the issue somewhat.
      } else {
        cts <- counts
      }

      if (block && !is.null(pairs) && !is.factor(pairs)) {
        dw_log("  Attempting to coerce the pairs vector to a factor...", "\n")
        pairs <- try(as.factor(pairs))
        if (is(pairs, "try-error")) stop("For block designs 'pairs' needs to be a factor.")
      }

      vf <- voom.fun
      if (length(grep("Weights", vf))) {
        dw_log("  'voomWithQualityWeights' is set as voom function; using settings for sample weights...\n")
        use_weights <- TRUE
        voom.fun <- limma::voomWithQualityWeights
      } else if (vf == "voom") {
        dw_log("  Using standard voom...\n")
        voom.fun <- limma::voom
      } else if (vf == "voomLmFit") {
        dw_log("  Using voom wrapper 'voomLmFit'\n")
        voom.fun <- edgeR::voomLmFit
      }

      if (vf == "voomLmFit") {
        if (block && !is.null(pairs)) {
          dw_log("  Samples are not independent. Using column", sQuote(pairs_col), "as block variable...\n")
          dw_log("    ", levels(pairs), "\n")
        } else {
          dw_log("  Performing standard linear model fit...\n")
          pairs <- NULL
        }
        if (use_weights) dw_log("    Empirical sample quality weights will be estimated...\n")
        fit <- voomLmFit(cts,
                       design = design,
                       plot = FALSE,
                       block = pairs,
                       sample.weights = use_weights,
                       normalize.method = norm.method,
                       keep.EList = TRUE)
        v <- fit$EList
      } else {
        v <- voom.fun(cts, design=design, plot=FALSE, normalize.method=norm.method)

        # For block designs (here, the pairs column has the blocking information):
        if (block && !is.null(pairs)) {
          dw_log("  Samples are not independent. Using column", sQuote(pairs_col), "as block variable...\n")
          dw_log("    ", levels(pairs), "\n")
          dw_log("  Calculating consensus correlation...\n")
          corfit <- limma::duplicateCorrelation(v$E , design, block=pairs)
          dw_log("  Fitting linear model using consensus correlation...\n")
          fit <- limma::lmFit(v, design, block=pairs, correlation=corfit$consensus)
        } else {
          dw_log("  Fitting linear model...\n")
          fit <- limma::lmFit(v, design)
        }
      }
      if (!is.null(contrasts)) {
        # Contrasts have been provided in a character vector and a contrast matrix was created, so pair-wise comparisons
        # are calculated here. Since we used voom, limma function are applied.
        if (length(grep("^Intercept$", rownames(contrasts)))) {
          rownames(contrasts)[grep("^Intercept$", rownames(contrasts))] <- "(Intercept)"
        }
        dw_log("  Fitting contrasts...\n")
        fit2 <- limma::contrasts.fit(fit, contrasts)
        dw_log("  eBayes fit...\n")
        fit2 <- limma::eBayes(fit2, trend = bayes.trend, robust = bayes.robust)
      } else {
        # Contrasts are "included" in the design, i.e., pair-wise comparisons have already been calculated in the lmFit() call.
        # Here, only the Bayes fit will be added.
        dw_log("  eBayes fit...\n")
        fit2 <- limma::eBayes(fit, trend = bayes.trend, robust = bayes.robust)
      }
      return(list(v=v, fit=fit, fit2=fit2))
    } else {
      ## Estimate dispersion values, relative to the design matrix, using the Cox-Reid (CR)-adjusted likelihood
      # Note, that batch effects or paired or blocked designs are taken into account in the design matrix and will be considered in the respective model.
      dw_log("  Estimating dispersion values...\n")
      d2 <- edgeR::estimateDisp(d, design)
      # Given the design matrix and dispersion estimates, fit a GLM to each feature.
      # If quasi-likelihood methods are used, estimate the gene-wise quasi-dispersions with empirical Bayes moderation.
      dw_log("  Fitting linear model...\n")
      if (quasi.likelihood) {
        fit <- edgeR::glmQLFit(d2, design, dispersion=d2[[disp]], robust=bayes.robust, abundance.trend=bayes.trend)
      } else {
        fit <- edgeR::glmFit(d2, design, dispersion=d2[[disp]])
      }
      return(list(d=d, d2=d2, fit=fit))
    }
  }



#' Helper function to generate an output table with only most relevant columns
#' @param contrast Character; a contrast group name pair in the typical format, i.e., "group2-group1"
#' @param annotated.normcnt data frame; annotated output table generated by performing pair-wise comparisons and
#'   corresponding statistical tests. From this the normalised counts are extracted.
#' @param samp.name.and.group.key \code{data.frame} with sample names (matching the count columns of
#'   \code{annotated.normcnt}) as row names and a \code{group} column giving each sample's group. Used to
#'   drop count columns of samples outside the contrast in multi-group designs; may be \code{NULL}.
#' @param out.dir Character; output directory for final tables. Required; no default is used so that nothing is written to the working directory unintentionally.
#' @param analysis.name Character used in the output file name.
#' @param filtered.lists Logical; should the DGE table be filtered before saving it to a file?
#' @param fdr.thr Numeric; the FDR threshold for filtering.
#' @param logfc.thr Numeric; the fold-change threshold on the log2-scale for filtering.
#' @return No return value. Called for its side effect of writing a reduced and optionally
#'   significance-filtered differential expression table to a tab-separated file below \option{out.dir}.
#' @examples
#' \donttest{
#' set.seed(1)
#' key <- data.frame(group = rep(c("control", "treated"), each = 4),
#'                   row.names = paste0("S", 1:8))
#' d3 <- data.frame(ID = paste0("g", 1:50), gene_symbol = paste0("G", 1:50),
#'                  logFC = rnorm(50), FDR = runif(50))
#' diffr_expr_generate_cleaned_de_table_output(
#'   contrast = "treated-control", annotated.normcnt = d3,
#'   samp.name.and.group.key = key, out.dir = tempdir(),
#'   analysis.name = "demo")
#' }
#' @export
diffr_expr_generate_cleaned_de_table_output <-
  function(contrast, annotated.normcnt, out.dir, samp.name.and.group.key, analysis.name=NULL, filtered.lists = TRUE,  fdr.thr=0.05, logfc.thr=1 )
  {
    dw_log("Generating a cleaned version of expression table...\n")

    # output file name; 'analysis.name' is optional, so drop it when empty/NULL
    nm <- c(analysis.name, contrast, "differential_expression_clean.tsv")
    DE.out <- paste(nm[nzchar(nm)], collapse = "_")

    # For multi-group designs, drop the count columns of samples that are not part of this
    # contrast. This is only attempted when the contrast maps cleanly onto two real group
    # names present in the key. For coefficient-style contrasts (e.g. "groupsX"), a NULL key,
    # or group names containing "-", the parsing is ambiguous, so all sample columns are kept.
    contrast.groups <- unlist(strsplit(contrast, split = "-", fixed = TRUE))
    if (!is.null(samp.name.and.group.key) && "group" %in% names(samp.name.and.group.key)) {
      grp <- as.character(samp.name.and.group.key$group)
      if (length(contrast.groups) == 2L && all(contrast.groups %in% grp)) {
        not.contrast.samples <- rownames(samp.name.and.group.key)[!(grp %in% contrast.groups)]
        drop.cols <- intersect(colnames(annotated.normcnt), not.contrast.samples)
        if (length(drop.cols)) {
          annotated.normcnt <- annotated.normcnt[, !colnames(annotated.normcnt) %in% drop.cols, drop = FALSE]
          dw_log("   Samples outside the contrast were removed...\n")
        }
      } else {
        dw_log("   NOTE: could not map the contrast onto the sample groups; keeping all sample columns...\n")
      }
    }

    # tidy up aligner-specific column-name suffixes (e.g. STAR 'ReadsPerGene.out')
    colnames(annotated.normcnt) <- gsub("ReadsPerGene.out", "", colnames(annotated.normcnt))

    if (filtered.lists) {
      dw_log("     Filtering the DE list to be saved by fdr <", fdr.thr, "and abs(logfc) >=", logfc.thr, "...\n")
      fdr.col <- names(annotated.normcnt)[grep("^fdr$|^adj*\\.{0,1}p\\.{0,1}val[e-u]{0,2}$", tolower(names(annotated.normcnt)))]
      fc.col  <- names(annotated.normcnt)[grep("^logfc$|fold$", tolower(names(annotated.normcnt)))]

      if (length(fdr.col) == 0L || length(fc.col) == 0L) {
        stop("Could not locate the FDR and/or fold-change column needed for filtering ",
             "(looked for an FDR/adjusted-p column and a logFC/fold column). ",
             "Set 'filtered.lists = FALSE' to save the unfiltered table.", call. = FALSE)
      }
      if (length(fdr.col) > 1L) {
        dw_log("   NOTE: several FDR-like columns found; using", sQuote(fdr.col[1]), "\n")
        fdr.col <- fdr.col[1]
      }
      if (length(fc.col) > 1L) {
        dw_log("   NOTE: several fold-change columns found; using", sQuote(fc.col[1]), "\n")
        fc.col <- fc.col[1]
      }

      nrow.before <- nrow(annotated.normcnt)
      # NA-safe: without the is.na() guards, rows with a missing FDR/logFC would be turned
      # into rows of all-NA in the output rather than being dropped
      keep <- !is.na(annotated.normcnt[[fdr.col]]) & !is.na(annotated.normcnt[[fc.col]]) &
              annotated.normcnt[[fdr.col]] < fdr.thr &
              abs(annotated.normcnt[[fc.col]]) >= logfc.thr
      annotated.normcnt <- annotated.normcnt[keep, , drop = FALSE]

      dw_log("       ", nrow.before - nrow(annotated.normcnt), "unsignificant genes were filtered out...\n")
      dw_log("        The dimensions of the table to be saved:", dim(annotated.normcnt), "...\n")

      DE.out <- paste("filtered", DE.out, sep = "_")
    }

    dw_log("   Saving cleaned DE-list to", DE.out, "...\n")
    write.table(annotated.normcnt, file.path(out.dir, DE.out), sep = "\t", quote = FALSE, row.names = FALSE)

    invisible(NULL)
  }

# All per-contrast plotting for diff_expr_extract_contrasts(), factored out so that the PDF device
# and the graphical parameters set on it share a single function lifetime. That is what makes the
# graphics state safe:
#   * the device is opened here, so par() can be captured *after* pdf() - capturing
#     par(no.readonly = TRUE) while no device exists forces one open and leaves the grid state such
#     that the pheatmap output never reaches the PDF (the blank-heatmap bug);
#   * par() called with settings returns the *previous* values, so 'oldpar' holds the margins of our
#     own PDF device, not the caller's;
#   * the immediate on.exit() restores those margins and closes the device, but only while our own
#     PDF is still the current device. On the normal path dev.off() has already run, so the handler
#     is a no-op and the caller's device and par() settings are never touched. On the error path it
#     is the cleanup.
# Keeping this out of the per-contrast loop in the caller also avoids accumulating one redundant
# on.exit() handler per contrast. Returns the three plot objects for this contrast.
dw_contrast_plots <- function(d3, contr, id.col, sym.col, contr.out.dir, analysis.name,
                              p.thr, fdr.thr, logfc.thr, numlab, point.lab, biom.attributes,
                              font.size, de.plot.base.size, lists,
                              samp.info, samples, groups, sample.plot.names,
                              color.blind.pal, n.pal.cols, color.extremes, palette.length,
                              anno.color, anno.name, heatmap.main,
                              hm.p.thr, hm.fdr.thr, hm.logfc.thr, heatmap.topn, heatmap.split.expr) {

  out <- list()

  grDevices::pdf(file.path(contr.out.dir, paste(analysis.name, contr, "_plots.pdf", sep = "_")),
                 width = 15, height = 15)
  own.dev <- grDevices::dev.cur()          # the device this function owns; updated on any swap below
  oldpar  <- graphics::par(mar = c(6, 6, 5, 3))
  on.exit({
    if (grDevices::dev.cur() == own.dev) { # only ever our own PDF, never the caller's device
      graphics::par(oldpar)
    }
    dw_dev_off(own.dev)                    # no-op if it was already closed on the normal path
  }, add = TRUE)

  dw_log(" MA-plot...\n")
  out$MAplot <- diff_expr_ma_plot(d3, contr, id.col, sym.col, p.thr, fdr.thr, logfc.thr, numlab,
                                  out.dir = contr.out.dir, analysis.name, point.lab, biom.attributes,
                                  font.size, lists, base.size = de.plot.base.size)

  ## Volcano plot
  dw_log(" Volcano plot...\n")
  volcano.name <- gsub(".", " ", contr, fixed = TRUE)
  volcano.name <- gsub("_", " ", volcano.name, fixed = TRUE)
  volcano.name <- gsub("-", " vs. ", volcano.name, fixed = TRUE)
  out$volcanoPlot <- diff_expr_volcano_plot(d3, id.col, sym.col = "gene_symbol", main = volcano.name,
                                            p.thr = p.thr, fdr.thr = fdr.thr, logfc.thr = logfc.thr,
                                            numlab = numlab, point.lab = point.lab,
                                            base.size = de.plot.base.size)

  ## Histogram of P-value distribution
  dw_log(" Dendrogram plot...\n")
  diff_expr_pval_hist_plot(d3)

  ## Heatmaps. Large ones get their own, bigger, PDF; the handler above then guards that device.
  if (heatmap.topn > 100) {
    dw_dev_off(own.dev)
    grDevices::pdf(file.path(contr.out.dir, paste(analysis.name, contr, "_heatmaps.pdf", sep = "_")),
                   width = 25, height = 35)
    own.dev <- grDevices::dev.cur()
    # Same capture-and-restore for the second device. The handler above already covers it (it
    # reads the reassigned 'own.dev'/'oldpar' at exit time), but the pattern is repeated here so
    # that every par() call is followed directly by its on.exit(). Both handlers are idempotent:
    # the restore is guarded by the device check and dw_dev_off() no-ops on a closed device.
    oldpar  <- graphics::par(mar = c(6, 6, 5, 3))
    on.exit({
      if (grDevices::dev.cur() == own.dev) graphics::par(oldpar)
      dw_dev_off(own.dev)
    }, add = TRUE)
  }
  dw_log("Heatmap plots...\n")
  out$heatmapPlot <- pheatmap_plots(d3 = d3,
                                    id = id.col,
                                    sym.col = "gene_symbol",
                                    samp.info = samp.info,
                                    samples = samples,
                                    groups = groups,
                                    sample.plot.names = sample.plot.names,
                                    color.blind.pal = color.blind.pal,
                                    n.pal.cols = n.pal.cols,
                                    color.extremes = color.extremes,
                                    palette.length = palette.length,
                                    anno.color = anno.color,
                                    anno.name = anno.name,
                                    main = NULL,
                                    add.main = heatmap.main,
                                    p.thr = hm.p.thr,
                                    fdr.thr = hm.fdr.thr,
                                    logfc.thr = hm.logfc.thr,
                                    topn = heatmap.topn,
                                    split.expr = heatmap.split.expr)
  # restore on our own device, before it is destroyed; both steps are no-ops if a plotting helper
  # closed the device itself (dev.off() would then hit the null device and error)
  if (grDevices::dev.cur() == own.dev) graphics::par(oldpar)
  dw_dev_off(own.dev)

  out
}

#' Function to extract contrasts and generate top tables and plots
#' @param contrasts Contrast matrix as generated by \code{makeContrasts()}. If \code{NULL}, the default, the design matrix
#'   is expected to describe simple pair-wise comparisons, i.e., with the contrasts already included in the matrix.
#'   Consequently, in that case, no contrasts will be fitted.
#' @param fit An MArrayLM object produced by \code{lmFit}.
#' @param fit2 An MArrayLM object produced by \code{eBayes}.
#' @param normcnt Matrix; depth-adjusted reads per million or "voomed" counts, i.e., a numeric matrix of
#'   normalized expression values on the log2 scale.
#' @param do.voom Logical; should data be "voomed"?
#' @param quasi.likelihood Logical; should quasi-likelihood methods be used? See \emph{Details} section.
#'   If \code{TRUE}, the default, then the quasi-likelihood methods are used, if \code{FALSE}, then the
#'   likelihood ratio test is used. Please note the main wrapper function \code{diffExpr()} will set this parameter.
#' @param out.l List; list of result objects generated upstream in the workflow to add output from this function to.
#' @param out.dir Character; output directory for final tables. Required; no default is used so that nothing is written to the working directory unintentionally.
#' @param analysis.name Character used in the output file name.
#' @param biomart Logical; should biomart annotation be added?
#' @param biom.data.set \code{character} of length one. Biomart data set to use.
#' @param biom.mart \code{character} vector. Biomart to use (uses the first element of the vector), defaults to "ensembl".
#' @param host \code{character} of length one. Host URL.
#' @param biom.filter \code{character} of length one. Name of biomart filter, i.e., type of query ids, defaults to "ensembl_gene_id".
#' @param biom.attributes \code{character} vector. Biomart attributes, i.e., type of desired result(s); make sure query id type is included!
#' @param biom.force.ensg \code{logical}. Should Ensembl Gene IDs be checked for (and stripped of) Ensembl version numbers? Defaults to \code{FALSE}.
#' @param biom.cache \code{character}. Path name giving the location of the cache \command{getBM()} uses if \code{use.cache=TRUE}. Defaults to the value in the \emph{BIOMART_CACHE} environment variable.
#' @param use.cache (\code{logical}). Should \command{getBM()} use the cache? Defaults to \code{TRUE} as in the \command{getBM()} function and is passed on to that.
#' @param sym.col \code{character}. Name of the column in the query result with gene symbols.
#' @param rm.dups \code{logical}. Should duplicated input IDs (\option{biom.filter}) be removed from the result?
#' @param p.thr Numeric; P-Value threshold used in plots.
#' @param fdr.thr Numeric; FDR threshold used in plots and for generating cleaned output tables.
#' @param logfc.thr Numeric; FC threshold on the log2-scale used in plots and for generating cleaned output tables.
#' @param numlab \code{numeric}. Maximum number of labels per plot. Overrides numbers calculated based on `p.thr` and `fdr.thr`.
#' @param point.lab \code{logical}. Should points be labelled, at all?
#' @param heatmap.topn \code{numeric}. Number of top values to be plotted. Defaults to 100.
#' @param hm.p.thr,hm.fdr.thr,hm.logfc.thr \code{numeric}. Significance and fold-change thresholds used
#'   specifically when selecting genes for the heatmaps. These are kept separate from the plot/table
#'   thresholds (\option{p.thr}, \option{fdr.thr}, \option{logfc.thr}) because heatmaps usually read best
#'   with a stricter gene set. Default to 0.05, 0.05 and 1.
#' @param heatmap.split.expr \code{logical}. Should the top up- and top down-regulated genes be displayed at equal numbers (50/50),
#' if they meet the significance threshold (regardless of the actual significance)? Defaults to \code{FALSE}.
#' @param color.blind.pal string determining the RColorBrewer color blind palette (default = "PuOr");
#' other option can be visualized with the following command: brewer.pal.info[brewer.pal.info$colorblind,]
#' @param n.pal.cols desired length of the number of different colours in 'color.blind.pal'. Will also be used
#' as length of the numeric vector of probabilities in 'quantile_breaks()' (see ?quantile); defaults to 11
#' @param color.extremes character vector of length 2 giving the two extremes of a user-defined colour palette
#' varying from the first hue to the second via white.
#' @param palette.length integer setting the desired length of the colour palette to be used in the heatmap
#' @param anno.color list of named character vectors giving the colours used in the heatmap annotation bars. See 'annotation_colors'
#' in [pheatmap()]. Automatically generated if NULL (default).
#' @param anno.name character string used as the column annotation legend title. If 'anno.color' is not NULL and of length 1 the slot name will be used if existing.
#' @param heatmap.main \code{character}. Additional information printed in the heatmap title.
#' @param font.size \code{numeric}. Size of point labels in M-A plots.
#' @param plots Logical; should plots be generated?
#' @param lists Logical; should output tables be written to files?
#' @param filtered.lists Logical; should the DGE table be filtered before saving it to a file?
#' @param samp.info \code{data.frame}. \code{samp.info} object containing information of the project's sample sheet.
#' @param samples \code{character}. Name of the column in 'samp.info' containing sample names. If 'samp.info' is not supplied
#'     vector of sample names.
#' @param groups \code{character}. Name of the column in 'samp.info' containing grouping information. If 'samp.info' is not supplied
#'     vector of groups.
#' @param sample.plot.names \code{character}. Optional name of a column with "nice" sample names for plotting.
#' @param de.plot.base.size \code{numeric}. Overall text/point scale (`ggplot2` \code{base_size}, default 16)
#'   passed to the M-A and volcano plots. Raise it so labels and legends stay readable on the large
#'   per-contrast PDF (opened at 15x15 inches).
#' @return The input list \option{out.l}, extended with a \code{contrasts} element holding one annotated
#'   result \code{data.frame} per contrast and, if plots were requested, with the \code{MAplots},
#'   \code{volcanoPlots} and \code{heatmapPlots} elements. Called also for its side effects of writing
#'   result tables, plot files and Venn section spreadsheets below \option{out.dir}.
#' @examples
#' \donttest{
#' si <- diff_expr_get_samp_info(diffwrap_samp_info, "SampleName", "Group")
#' counts <- diff_expr_filter_counts(diff_expr_read_counts(diffwrap_counts, si), si)
#' groups <- stats::relevel(si$Groups, ref = "control")
#' d <- edgeR::calcNormFactors(edgeR::DGEList(counts, group = groups))
#' design <- diff_expr_make_design(si, groups)
#' contrasts <- diff_expr_make_contrasts(design, groups)
#' fit <- diff_expr_fit(counts, d, design, do.voom = FALSE, quasi.likelihood = TRUE)
#' normcnt <- edgeR::cpm(fit$d2, log = TRUE)
#' out <- diff_expr_extract_contrasts(contrasts = contrasts, fit = fit$fit,
#'                                    normcnt = normcnt, out.l = list(),
#'                                    do.voom = FALSE, out.dir = tempdir(),
#'                                    analysis.name = "demo", plots = FALSE,
#'                                    samp.info = si, samples = "SampleNames",
#'                                    groups = groups)
#' names(out$contrasts)
#' }
#' @export
diff_expr_extract_contrasts <-
  function(contrasts=NULL, fit, fit2=NULL, normcnt, out.l, do.voom=TRUE, quasi.likelihood=TRUE, out.dir,
           analysis.name=NULL, biomart=FALSE, biom.data.set="hsapiens_gene_ensembl", biom.mart="ensembl",
           host="https://www.ensembl.org", biom.filter="ensembl_gene_id", biom.attributes=c("ensembl_gene_id","hgnc_symbol","description"),
           biom.force.ensg = FALSE, biom.cache = NULL, use.cache = FALSE, sym.col="hgnc_symbol",
           rm.dups=FALSE, p.thr=0.05, fdr.thr=0.05, logfc.thr=1, numlab=25, point.lab=TRUE, heatmap.topn = 100,
           hm.p.thr = 0.05, hm.fdr.thr = 0.05, hm.logfc.thr = 1,
           heatmap.split.expr = FALSE, color.blind.pal = "PuOr", n.pal.cols = 11, color.extremes = c("#3182BD", "#E6550D"),
           palette.length = NULL, anno.color = NULL, anno.name = "Sample Class", heatmap.main = NULL, font.size=5, plots=TRUE,
           lists=TRUE, filtered.lists = TRUE, samp.info = NULL, samples = NULL, groups = NULL, sample.plot.names = NULL,
           de.plot.base.size = 16)
  {
    # initial checks
    if (is.null(samp.info)) stop("Need sample sheet! Provide as data frame to 'samp.info' argument.")
    if (is.null(samples)) stop("Need name of column in sample sheet containing sample names!")
    if (is.null(groups)) stop("Need name of column in sample sheet containing grouping information!")
    if (is.null(analysis.name)) {
      analysis.name <- paste0(paste(levels(groups)[1:2], collapse="_"), "_")
      dw_log("Using default settings for the 'analysis name':", sQuote(analysis.name), "\n")
      warning("A unique and descriptive name for the analysis should always be provided!")
    }

    # exact name match, not grep(): a substring test would also be satisfied by any future
    # element whose name merely contains "contrasts" and would then skip the initialisation
    if (!"contrasts" %in% names(out.l)) {
      out.l$contrasts <- list()
    }
    if (is.null(contrasts)) {
      dw_log("  Extracting 'groups' levels from design...\n")
      cn <- grep("^groups.+", colnames(fit$design), value=TRUE)
    } else if (is.matrix(contrasts) && !is.null(colnames(contrasts))) {
      # if do.voom=TRUE contrasts have been calculated already in the fit wrapper function above and only need to be extracted
      # in the case of paired samples, contrasts are inherent to the design and only need to be extracted, i.e., we provide the coefficient of interest to the test function.
      # For GL models, we need to feed the contrast of interest in form of a matrix to glmLRT or glmQLFTest to instruct it to calculate the desired contrasts before doing the test
      # see the different behaviour of glmLRT/glmQLFTest below for do.voom=FALSE
      cn <- colnames(contrasts)
    } else {
      # contrasts is a character vector
      # this will only be called for the non-contrasted columns of the design matrix for paired samples
      # which means contrasts are inherent to the design and only need to be extracted
      cn <- contrasts
    }

    # Create the collections that hold the returnable plot objects (ggplots for MA and volcano,
    # pheatmap objects for the heatmaps), but only if they do not exist yet: diffExpr() calls this
    # function twice for designs where some contrasts are inherent to the design and the remaining
    # ones (e.g. all-vs-all or explicitly requested contrasts) are fitted separately. Both calls
    # share one 'out.l', so resetting here unconditionally would discard the plots of the first
    # call. Same idiom as for 'contrasts' above. Skipped entirely when no plots are produced, so a
    # 'plots = FALSE' run does not return three collections that stay empty.
    if (plots) {
      for (el in c("MAplots", "volcanoPlots", "heatmapPlots")) {
        if (!el %in% names(out.l)) out.l[[el]] <- list()
      }
    }
    for (contr in cn) {
      dw_log("Calculating differential expression for", contr, "\n")
      contr.out.dir <- file.path(out.dir, contr)
      dir.create(contr.out.dir, showWarnings = FALSE)
      dw_log("Storing all results under", contr.out.dir, "\n")

      if (do.voom) {
        if (is.null(fit2)) stop("Need 'fit2' object!")
        dw_log("Generating output table of differentially expressed features...\n")
        tt <- limma::topTable(fit2, coef=contr, number=nrow(fit2), sort.by="P")
      } else {
        test.fun <- ifelse(quasi.likelihood, edgeR::glmQLFTest, edgeR::glmLRT)
        # Perform a likelihood ratio test or a quasi-likelihood F-Test, specifying the contrast of interest.
        # Needs to be done per-contrast as otherwise it will perform the LR test of all contrasts together while we want this done individually.
        dw_log("Performing likelihood ratio test on this contrast...\n")
        if (is.matrix(contrasts)) {
          # calculating the contrast
          de <- test.fun(fit, contrast = contrasts[, contr])
        } else {
          # only extracting the contrast (for paired samples)
          de <- test.fun(fit, coef=contr)
        }
        # Use the topTags function to present a tabular summary of the differential expression statistics (note that topTags
        # operates on the output of exactTest or glmLRT, but only the latter is shown here)
        dw_log("Generating tablular summary of differential expression statistics ('top table')\n")
        tt <- edgeR::topTags(de, n = nrow(de))
        #			} else {
        #				stop("Need contrast matrix as input.")
        #			}
      }
      if (is.null(contrasts)) {
        fitcnt <- fit$coefficients[, contr]
        addinfo <- cbind(normcnt, fitcnt)
        if (do.voom) {
          names(addinfo)[names(addinfo) %in% "fitcnt"] <- "Fitted Coefficients (log2)"
        } else {
          names(addinfo)[names(addinfo) %in% "fitcnt"] <- "Fitted Coefficients (ln)"
        }
      } else if (is.matrix(contrasts)) {
        fitcnt <- fit$coefficients[, rownames(contrasts)[contrasts[, contr]!=0]]
        addinfo <- cbind(normcnt, fitcnt)
      } else {
        addinfo <- normcnt
      }
      d3 <- merge(addinfo, tt, by="row.names", sort=FALSE)
      dw_log("Renaming ID column...\n")
      names(d3)[1] <- "ID"
      if (biomart) {
        dw_step(paste("  |--> Retrieving gene annotation from BioMart for"), contr, "...\n\n")
        # diffExpr() probes BioMart up front, but keep a guard here in case a host drops out
        # mid-run: on failure fall back to offline convertId2() symbols for the remaining contrasts.
        # verbose = TRUE turns on convert.bm()/getBM() progress output (shows if the session surfaces it).
        d3.bm <- tryCatch(
          diff_expr_biomart(d3, biom.data.set, biom.mart, host, biom.filter, biom.attributes, biom.cache, use.cache, sym.col, rm.dups, force.ensg = biom.force.ensg, verbose = TRUE),
          error = function(e) {
            warning("biomart query failed (", conditionMessage(e), "); falling back to offline ",
                    "convertId2() gene symbols for the remaining contrasts.", call. = FALSE)
            dw_log("  !NOTE! biomart query failed:", conditionMessage(e), "- falling back to convertId2()...\n")
            NULL
          })
        if (is.null(d3.bm)) {
          biomart <- FALSE          # stop retrying biomart for the remaining contrasts
        } else {
          d3 <- d3.bm
          id.col <- names(d3)[names(d3) %in% biom.filter]
        }
      }
      if (!biomart) {
        dw_step("\n")
        # convertId2() defaults to human; pass the species matching 'biom.data.set' so mouse
        # datasets are not annotated with human symbols. Also the offline fallback used above.
        conv.species <- switch(biom.data.set,
                               mmusculus_gene_ensembl = "Mouse",
                               "Human")
        syms <- convertid::convertId2(as.character(d3$ID), conv.species)
        syms <- data.frame(ID=names(syms), gene_symbol=as.character(syms), stringsAsFactors=FALSE)
        d3 <- merge(syms, d3, by="ID", all.y=TRUE, all.x=FALSE, sort=TRUE)
        if (any(d3$gene_symbol=="") || any(is.na(d3$gene_symbol))) {
          dw_log("  Replacing", length(which(d3$gene_symbol=="" | is.na(d3$gene_symbol))), "missing Gene Symbols by Ensembl IDs...\n")
          d3$gene_symbol[d3$gene_symbol=="" | is.na(d3$gene_symbol)] <- as.character(d3$ID[d3$gene_symbol=="" | is.na(d3$gene_symbol)])
        }
        id.col <- "ID"
      }
      dw_log("Setting unique row names...\n")
      rownames(d3) <- make.unique(as.character(d3[[id.col]]))
      out.l$contrasts[[contr]] <- d3
      if (lists) {
        DE.out <- paste(analysis.name, contr, "differential_expression.tsv", sep="_")
        dw_log("Saving list to", DE.out, "...\n")
        pv.col <- dw_find_col(names(d3), "^p\\.{0,1}val[e-u]{0,2}$", "p-value")
        dw_log("The result table is ordered in increasing order by column ", pv.col, "...\n")
        d3 = d3[order(d3[[pv.col]]),]
        write.table(d3, file.path(contr.out.dir, DE.out), sep="\t", quote=FALSE, row.names=FALSE)

        diffr_expr_generate_cleaned_de_table_output(contrast=contr, annotated.normcnt=d3, samp.name.and.group.key = fit$samples, out.dir=contr.out.dir,
                                                    analysis.name = analysis.name, filtered.lists = filtered.lists, fdr.thr=fdr.thr, logfc.thr=logfc.thr)
      }

      # Create a graphical summary, such as an M (log-fold change) versus A (log-average expression) plot, here showing the
      # genes selected as differentially expressed (with a 5% false discovery rate)
      if (plots) {
        dw_log("Plotting...\n")
        # The PDF device and the par() settings applied to it are owned by this helper, so they
        # share one function lifetime and are cleaned up together (see its comment block).
        plots.l <- dw_contrast_plots(d3 = d3, contr = contr, id.col = id.col, sym.col = sym.col,
                                     contr.out.dir = contr.out.dir, analysis.name = analysis.name,
                                     p.thr = p.thr, fdr.thr = fdr.thr, logfc.thr = logfc.thr,
                                     numlab = numlab, point.lab = point.lab,
                                     biom.attributes = biom.attributes, font.size = font.size,
                                     de.plot.base.size = de.plot.base.size, lists = lists,
                                     samp.info = samp.info, samples = samples, groups = groups,
                                     sample.plot.names = sample.plot.names,
                                     color.blind.pal = color.blind.pal, n.pal.cols = n.pal.cols,
                                     color.extremes = color.extremes, palette.length = palette.length,
                                     anno.color = anno.color, anno.name = anno.name,
                                     heatmap.main = heatmap.main,
                                     hm.p.thr = hm.p.thr, hm.fdr.thr = hm.fdr.thr,
                                     hm.logfc.thr = hm.logfc.thr, heatmap.topn = heatmap.topn,
                                     heatmap.split.expr = heatmap.split.expr)
        out.l$MAplots[[contr]]      <- plots.l$MAplot
        out.l$volcanoPlots[[contr]] <- plots.l$volcanoPlot
        out.l$heatmapPlots[[contr]] <- plots.l$heatmapPlot
      }
    }
    dw_log("=============================================================================\n")
    dw_log("Venn sections list\n")
    cols.interest.de.table = c("ensembl","symbol|hgnc","description","entrez","^average$|^ave[a-r]{0,4}expr[e-s]{0,6}$","^fdr$|^adj*\\.{0,1}p\\.{0,1}val[e-u]{0,2}$","^p\\.{0,1}val[e-u]{0,2}$")
    pv.col = dw_find_col(names(out.l$contrasts[[1]]), "^p\\.{0,1}val[e-u]{0,2}$", "p-value")
    fdr.col = dw_find_col(names(out.l$contrasts[[1]]), "^fdr$|^adj*\\.{0,1}p\\.{0,1}val[e-u]{0,2}$", "FDR")

    sign.de.tables <- lapply(out.l$contrasts, function(x) {
      sel.col.de.table <- names(x)[grep(paste(cols.interest.de.table,collapse="|"), tolower(names(x)))]
      new.de.table <- x %>% dplyr::select(dplyr::all_of(sel.col.de.table))
      return(new.de.table)
    })

    join.interest <- paste(c("ensembl", "symbol|hgnc", "description", "entrez"), collapse = "|")
    join.vec <- grep(join.interest, tolower(names(out.l$contrasts[[1]])), value = TRUE)

    if (all(lapply(sign.de.tables, function(x) nrow(x[x[[fdr.col]] < 0.05,])) > 0)) {
      sign.de.tables = lapply(sign.de.tables, function(x) x[x[[fdr.col]] < 0.05,])
    } else {
      sign.de.tables = lapply(sign.de.tables, function(x) x[x[[pv.col]] < 0.05,])
    }

    # v = list with the venn.diagram plot and all venn sections tables
    v <- diffr_venn(list.comp.tables = sign.de.tables, join.vec = join.vec, .log = FALSE)
    #print(class(v))

    if (plots && !is(v[["venn.diagram"]], "try-error")) {
      dw_log(">>> Plotting Venn diagram...", "\n")
      pdf(file.path(out.dir,"Venn_Diagram.pdf"), width = 15, height= 15)
      venn.dev <- grDevices::dev.cur()
      grid::grid.draw(v[["venn.diagram"]])
      dw_dev_off(venn.dev)
    }

    ifelse(!dir.exists(file.path(out.dir, "Venn sections")), dir.create(file.path(out.dir, "Venn sections")), FALSE)
    for(j in seq_along(v[["venn.sections"]])){
      dw_log_obj(names(v[["venn.sections"]])[j])
      openxlsx::write.xlsx(v[["venn.sections"]][[j]], file = file.path(out.dir, "Venn sections", paste0(names(v[["venn.sections"]])[j],".xlsx")))
    }
    dw_log("Venn sections list end", "\n")
    return(out.l)
  }

# Is a BioMart server reachable? Establishes a session for the requested mart and dataset (no gene
# query needed). Internal availability gate used by diffExpr() to decide whether to fall back to
# offline convertId2() annotation; modelled on the 'convertid' package's availability check.
# Returns TRUE/FALSE. 'biomaRt' is in Suggests, so the call is requireNamespace()-guarded.
#
# The configured host is probed first, because that is the server the run itself will query.
# The standard mirrors are only tried when that host is the public Ensembl site: for a custom
# host (a versioned release, an in-house mart) a working ensembl.org mirror says nothing about
# whether the run will succeed. Note that useEnsembl() treats 'host' and 'mirror' as mutually
# exclusive - a given host is used as is - so the two are probed in separate calls.
dw_biomart_available <- function(biom.data.set = "hsapiens_gene_ensembl",
                                 biom.mart = "ensembl",
                                 host = NULL,
                                 timeout = 10) {
  if (!requireNamespace("biomaRt", quietly = TRUE)) return(FALSE)
  mart <- biom.mart[1]

  # Each attempt gets a time budget: without a budget every dead candidate costs a full TCP/HTTP
  # timeout. setTimeLimit() raises an error once the budget is exceeded, which the tryCatch turns
  # into FALSE like any other failure, and the limit is reset on exit so nothing leaks into the
  # caller's session. It takes effect at R's interrupt checkpoints, so it bounds the wait in
  # practice rather than being a hard kill; raise 'timeout' if a legitimate mart needs longer.
  attempt <- function(expr) {
    setTimeLimit(elapsed = timeout, transient = TRUE)
    on.exit(setTimeLimit(cpu = Inf, elapsed = Inf, transient = FALSE), add = TRUE)
    tryCatch({ force(expr); TRUE }, error = function(e) FALSE)
  }

  probe.host <- function(h) {
    attempt(biomaRt::useEnsembl(mart, dataset = biom.data.set, host = h))
  }
  probe.mirror <- function(m) {
    if (is.null(m)) {
      attempt(biomaRt::useEnsembl(mart, dataset = biom.data.set))
    } else {
      attempt(biomaRt::useEnsembl(mart, dataset = biom.data.set, mirror = m))
    }
  }

  have.host <- !is.null(host) && length(host) && nzchar(host[1])
  if (have.host) {
    if (isTRUE(probe.host(host[1]))) return(TRUE)
    # a custom host has no mirrors to fall back on
    if (!grepl("(^|[./])ensembl\\.org", host[1])) return(FALSE)
  }
  for (m in list(NULL, "useast", "uswest", "asia")) {
    if (isTRUE(probe.mirror(m))) return(TRUE)
  }
  FALSE
}

#' Function to retrieve additional information from biomart
#' @param d3 data frame to be annotated with biomart; annotated output table generated by performing pair-wise comparisons and
#'   corresponding statistical tests.
#' @param biom.data.set \code{character} of length one. Biomart data set to use.
#' @param biom.mart \code{character} vector. Biomart to use (uses the first element of the vector), defaults to "ensembl".
#' @param host \code{character} of length one. Host URL.
#' @param biom.filter \code{character} of length one. Name of biomart filter, i.e., type of query ids, defaults to "ensembl_gene_id".
#' @param biom.attributes \code{character} vector. Biomart attributes, i.e., type of desired result(s); make sure query id type is included!
#' @param biom.cache \code{character}. Path name giving the location of the cache \command{getBM()} uses if \code{use.cache=TRUE}. Defaults to the value in the \emph{BIOMART_CACHE} environment variable.
#' @param use.cache (\code{logical}). Should \command{getBM()} use the cache? Defaults to \code{TRUE} as in the \command{getBM()} function and is passed on to that.
#' @param sym.col \code{character}. Name of the column in the query result with gene symbols.
#' @param rm.dups \code{logical}. Should duplicated input IDs (\option{biom.filter}) be removed from the result?
#' @param force.ensg \code{logical}. Should Ensembl Gene IDs be checked for (and stripped of) Ensembl version numbers? Defaults to \code{FALSE}.
#' @param verbose \code{logical}. Passed to \code{convert.bm()}/\code{getBM()} to show query and
#'   chunk progress for large BioMart queries. Defaults to \code{FALSE}.
#' @return A \code{data.frame}: the input table \option{d3} merged with the annotation retrieved from
#'   biomart. The column holding gene symbols is renamed to \code{gene_symbol}.
#' @examples
#' \donttest{
#' # Needs network access to Ensembl BioMart. The query is wrapped in tryCatch() so that an
#' # unreachable or slow server reports the problem instead of failing the example.
#' if (requireNamespace("biomaRt", quietly = TRUE)) {
#'   d3 <- data.frame(ID = c("ENSG00000141510", "ENSG00000012048"))
#'   ann <- tryCatch(diff_expr_biomart(d3, biom.data.set = "hsapiens_gene_ensembl"),
#'                   error = function(e) {
#'                     message("BioMart not reachable: ", conditionMessage(e))
#'                     NULL
#'                   })
#'   if (!is.null(ann)) nrow(ann)
#' }
#' }
#' @export
diff_expr_biomart <-
  function(d3, biom.data.set="hsapiens_gene_ensembl", biom.mart="ensembl",
           host="https://www.ensembl.org", biom.filter="ensembl_gene_id", biom.attributes=c("ensembl_gene_id","hgnc_symbol","description"),
           biom.cache = NULL, use.cache = FALSE, sym.col="hgnc_symbol", rm.dups=FALSE, force.ensg=FALSE, verbose = FALSE)
  {
    # test if needed packages are installed
    if (use.cache && !requireNamespace("rappdirs", quietly = TRUE)) {
      stop(
        paste("Package", sQuote("rappdirs"), "must be installed to use the biomart cache."),
        call. = FALSE
      )
    }
    if (use.cache && is.null(biom.cache)) {
      biom.cache <- rappdirs::user_cache_dir("biomaRt")
    }
    # make sure the filter column is also retrieved with the query
    if (!biom.filter %in% biom.attributes) {
      biom.attributes <- c(biom.attributes, biom.filter)
    }
    if (force.ensg) {
      dw_log(" Enforcing standard Ensembl Gene IDs...\n")
      if (length(grep("\\.[[:digit:]]{1,}$", d3$ID, perl = TRUE))) {
        d3$ID <- sub("\\.[[:digit:]]{1,}$", "", d3$ID)
      } else {
        dw_log("  -> Nothing to do\n")
      }
    }
    # NB: named arguments -- the newer convert.bm() inserted 'biomart.fallback' and 'chunk.size'
    # (which keep their defaults here: Ensembl mirrors + 500-ID chunking for large queries).
    gene.lab <- convertid::convert.bm(dat = d3, id = "ID",
                                      biom.data.set = biom.data.set, biom.mart = biom.mart,
                                      host = host, biom.filter = biom.filter,
                                      biom.attributes = biom.attributes, biom.cache = biom.cache,
                                      use.cache = use.cache, sym.col = sym.col, rm.dups = rm.dups,
                                      verbose = verbose)
    names(gene.lab)[names(gene.lab)==sym.col] <- "gene_symbol"
    dw_log("  Extended annotation:\n")
    biom.attributes[biom.attributes==sym.col] <- "gene_symbol"
    if (length(d3$ID)>8) {
      dw_log_obj(gene.lab[1:8, biom.attributes])
      dw_log("_truncated_ (", length(d3$ID), "features)\n")
    } else {
      dw_log_obj(gene.lab[, biom.attributes])
    }
    return(gene.lab)
  }


