# All plots for differential expression analysis results
#
# Author: vidal
###############################################################################

utils::globalVariables(c("Average Expression", "logFC", "log2 Fold-Change", "adj. P-Value", "label_id", "P-Value"))

#' Function to generate a M-A plot using `ggplot2`
#' @param dat data frame; annotated output table generated by performing pair-wise comparisons and
#'   corresponding statistical tests.
#' @param contr Character; contrast to be processed. This is either a column name of a contrast matrix or a character
#'   conforming to the same format.
#' @param id \code{character}. Name of the column with the IDs used for labeling. Used as fallback for \code{sym.col}.
#' @param sym.col \code{character}. Name of the column with the IDs used for labeling. Takes priority over \code{id}.
#' @param p.thr Numeric; P-Value threshold for filtering gene labels. Only genes (points) with a value below this threshold will be labeled.
#' @param fdr.thr Numeric; FDR threshold for filtering gene labels (see \code{p.thr}).
#' @param logfc.thr Numeric; FC threshold on the log2-scale for filtering gene labels.
#' @param numlab \code{numeric}. Maximum number of labels per plot. Overrides numbers calculated based on `p.thr` and `fdr.thr`.
#' @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 point.lab \code{logical}. Should points be labelled, at all?
#' @param biom.attributes \code{character}. Vector of column names to be retrieved from biomart.
#' @param font.size Size of point labels in M-A plots.
#' @param lists Logical; should output tables be written to files?
#' @param base.size \code{numeric}. Overall text scale (`ggplot2` \code{base_size}, default 16) applied via
#'   \code{theme_gray}; controls legend, axis and title text. Independent of \code{font.size}, which sets the
#'   point-label size. Raise it for readability on a large device (e.g. the 15x15 inch pipeline PDF).
#' @return A named \code{list} of \code{ggplot} objects with the elements \code{FDR} and \code{Pval},
#'   holding the M-A plot with points highlighted by false discovery rate and by p-value, respectively.
#' @examples
#' \donttest{
#' set.seed(1)
#' d3 <- data.frame(gene_symbol = paste0("G", 1:100),
#'                  logFC = rnorm(100), PValue = runif(100), FDR = runif(100),
#'                  AveExpr = rnorm(100, 5))
#' rownames(d3) <- d3$gene_symbol
#' g <- diff_expr_ma_plot(d3, contr = "treated-control", sym.col = "gene_symbol",
#'                        out.dir = tempdir(), lists = FALSE)
#' names(g)
#' }
#' @export
diff_expr_ma_plot <-
		function(dat, contr, id=NULL, sym.col="gene_symbol", p.thr=0.05, fdr.thr=0.05, logfc.thr=1, numlab=25,
		         out.dir, analysis.name=NULL, point.lab=TRUE, biom.attributes=c("ensembl_gene_id","hgnc_symbol","description"),
		         font.size=5, lists=TRUE, base.size=16)
{
	if (is.null(id) && !sym.col %in% names(dat)) {
		stop("Need one of 'id' or 'sym.col'.")
	}
	if (sym.col %in% names(dat)) {
		dat$label_id <- dat[[sym.col]]
	} else if (id %in% names(dat)) {
		dat$label_id <- dat[[id]]
	} else {
		stop("No ID columns found.")
	}
	rn <- rownames(dat)
	pv.col <- dw_find_col(names(dat), "^p\\.{0,1}val[e-u]{0,2}$", "p-value")
	fdr.col <- dw_find_col(names(dat), "^fdr$|^adj*\\.{0,1}p\\.{0,1}val[e-u]{0,2}$", "FDR")
	numlab <- ceiling(1.25*numlab)
	# average-expression column: 'logCPM' from edgeR::topTags, 'AveExpr' from limma::topTable
	A <- dw_find_col(names(dat), "^logcpm$|^aveexpr$", "average expression")
	degFDR <- rn[dat[[fdr.col]] < fdr.thr & abs(dat$logFC) >= logfc.thr]
	degPval <- rn[dat[[pv.col]] < p.thr & abs(dat$logFC) >= logfc.thr]
	dat$`P-Value` <- dat[[pv.col]]
	dat$`adj. P-Value` <- dat[[fdr.col]]
	dat$`log2 Fold-Change` <- abs(dat$logFC)
	dat$`Average Expression` <- dat[[A]]
	# list to collect ggplot2 objects for output
	g.l <- list()
	dw_log("  M-A plot...\n")
	if (length(degFDR)>0) {
		dw_log(" FDR filtered values...\n")

		## ggplot2 M-A plot
		g <- ggplot(data=dat, aes(x=`Average Expression`, y=logFC))
		g <- g + geom_point(aes(size=`log2 Fold-Change`, color=`adj. P-Value`))
		g <- g + scale_colour_gradient2(low="#00106B", high="#A4B1FF", mid="#F3F5FF", midpoint=0.2)
		if (point.lab) {
			gene.lab <- dat[rn %in% degFDR, ]
			dw_log("    ", nrow(gene.lab), "point(s) labelled...\n")
			if (nrow(gene.lab)>numlab) {
				dw_log("     Restricting to", numlab, "...\n")
				ix <- sort(gene.lab$`adj. P-Value`, index=TRUE)$ix
				gene.lab <- gene.lab[ix[1:numlab], ][order(ix[1:numlab]), ]
			}
			dw_log("   Highlighted features:\n")
			if (length(degFDR)>8) {
				dw_log_obj(gene.lab[1:8, 1:2])
				dw_log("_truncated_ (", length(degFDR), "features)\n")
			} else {
				dw_log_obj(gene.lab[1:8, 1:2])
			}
			g <- g + ggtitle(paste0("M-A plot for ", contr, " (highl.: FDR < ", fdr.thr, "; FC >= ", 2^logfc.thr, "-fold)"))
			g <- g + ggrepel::geom_text_repel(data = gene.lab, aes(label = label_id), size = font.size, box.padding = unit(0.35, "lines"), point.padding = unit(0.3, "lines"), max.overlaps = 15, show.legend = FALSE)
			if (lists) {
				write.table(gene.lab, file.path(out.dir, paste(analysis.name, contr, "labelledPointsSmearPlot.tsv", sep="_")), sep="\t", quote=FALSE, row.names=FALSE)
			}
		} else {
			g <- g + ggtitle(paste("M-A plot for", contr))
		}
		g <- g + theme_gray(base_size = base.size)
		suppressWarnings(print(g))
		g.l[["FDR"]] <- g
	}
	if (length(degPval)>0) {
		dw_log("    for P-value filtered values...\n")

		## ggplot implementation of MA plot
		g <- ggplot(data=dat, aes(x=`Average Expression`, y=logFC))
		g <- g + geom_point(aes(size=`log2 Fold-Change`, color=`P-Value`))
		g <- g + scale_colour_gradient2(low="#00106B", high="#A4B1FF", mid="#F3F5FF", midpoint=0.2)
		if (point.lab) {
			gene.lab <- dat[rn %in% degPval, ]
			dw_log("    ", nrow(gene.lab), "point(s) labelled...\n")
			if (nrow(gene.lab)>numlab) {
				dw_log("     Restricting to", numlab, "...\n")
				ix <- sort(gene.lab$`P-Value`, index=TRUE)$ix
				gene.lab <- gene.lab[ix[1:numlab], ][order(ix[1:numlab]), ]
			}
			dw_log("   Highlighted features:\n")
			if (length(degPval)>8) {
				dw_log_obj(gene.lab[1:8, 1:2])
				dw_log(paste0("   _truncated_ (", length(degPval), " features)\n"))
			} else {
				dw_log_obj(gene.lab[1:8, 1:2])
			}
			g <- g + ggtitle(paste0("M-A plot for ", contr, " (highl.: P-value < ", p.thr, "; FC >= ", 2^logfc.thr, "-fold)"))
			g <- g + ggrepel::geom_text_repel(data = gene.lab, aes(label = label_id), size = font.size, box.padding = unit(0.35, "lines"), point.padding = unit(0.3, "lines"), max.overlaps = 15, show.legend = FALSE)
			if (lists) {
				write.table(gene.lab, file.path(out.dir, paste(analysis.name, contr, "labelledPointsSmearPlot.tsv", sep="_")), sep="\t", quote=FALSE, row.names=FALSE)
			}
		} else {
			g <- g + ggtitle(paste("M-A plot for", contr))
		}
		g <- g + theme_gray(base_size = base.size)
		print(g)
		g.l[["Pval"]] <- g
	}
	return(g.l)
}


#' Helper function that returns a volcano plot.
#' @description Builds a single volcano plot for one significance measure. Points are coloured by an
#'   'EnhancedVolcano'-style, colour-blind safe (Okabe-Ito) four-level scheme distinguishing genes that
#'   pass neither threshold ("not sign."), the fold-change threshold only, the significance threshold only, or
#'   both. Only the doubly-significant genes are eligible for text labels.
#' @param data.df \code{data.frame}. Data frame containing all necessary columns to generate a Volcano plot with
#'   gene labels (at least p-values, FDR values, log-ratios and gene symbols or other IDs)
#' @param property.to.plot \code{character}. Statistic to be plotted, one of "fdr" or "p". Matched via \code{match.arg}.
#' @param property.column \code{character}. Name of the actual column holding the statistic to be plotted.
#' @param property.thr \code{numeric}. Significance threshold; genes below it (in \option{property.column})
#'   are considered significant. Defaults to 0.05.
#' @param logfc.thr \code{numeric}. Absolute log2 fold-change threshold; genes beyond it, together with the
#'   significance threshold, form the highlighted ("both") set that is eligible for labelling. Defaults to 1.
#' @param main \code{character}. Main plot title. Defaults to \code{NULL}.
#' @param numlab \code{numeric}. Maximum total number of gene labels per plot (default 25). Labels are taken
#'   from the doubly-significant genes and split between up- and down-regulated in proportion to the skew,
#'   with a guaranteed minimum per side so the minor side is never left empty. The chosen split is logged.
#' @param point.lab \code{logical}. Should points be labelled, at all?
#' @param sym.col \code{character}. Name of column with gene symbols, e.g., HGNC Symbols.
#' @param pretty.breaks \code{logical}. Should the breaks on the axes be pretty? Uses `scales::breaks_extended`.
#' @param base.size \code{numeric}. Overall text/point scale (the `ggplot2` \code{base_size}, default 16).
#'   Directly used by legend, axis and title text; point-label, threshold-annotation and point sizes are derived
#'   from it. Useful when the plot is printed to a large device (e.g. the 15x15 inch pipeline PDF), where the
#'   default `ggplot2` sizes are too small.
#' @return A \code{ggplot} object containing the volcano plot for the requested significance measure.
#' @examples
#' set.seed(1)
#' d3 <- data.frame(gene_symbol = paste0("G", 1:100),
#'                  logFC = rnorm(100), PValue = runif(100))
#' g <- prepare_volcano_of_given_property(d3, property.to.plot = "p",
#'                                        property.column = "PValue",
#'                                        property.thr = 0.05, logfc.thr = 1)
#' class(g)
#' @export
prepare_volcano_of_given_property <-
  function(data.df, property.to.plot = c("fdr", "p"), property.column,
           property.thr = 0.05, logfc.thr = 1, main = NULL, numlab = 25,
           point.lab = TRUE, sym.col = "gene_symbol", pretty.breaks = FALSE,
           base.size = 16) {

    property.to.plot <- match.arg(property.to.plot)
    if (pretty.breaks && !requireNamespace("scales", quietly = TRUE)) {
      stop("Package ", sQuote("scales"),
           " must be installed to use the 'pretty.breaks' option.", call. = FALSE)
    }

    # --- significance categories (EnhancedVolcano-style, colour-blind safe) ------
    prop    <- data.df[[property.column]]
    data.df[["neg.log10"]] <- -log10(prop)
    measure <- toupper(property.to.plot)                       # "FDR" or "P"

    pass.p  <- !is.na(prop) & prop < property.thr
    pass.fc <- !is.na(data.df$logFC) & abs(data.df$logFC) > logfc.thr

    lab.ns   <- "not sign."
    lab.fc   <- sprintf("logFC > %s", logfc.thr)
    lab.sig  <- sprintf("%s < %s", measure, property.thr)
    lab.both <- sprintf("%s < %s & logFC > %s", measure, property.thr, logfc.thr)
    levs     <- c(lab.ns, lab.fc, lab.sig, lab.both)
    # Okabe-Ito palette: grey / bluish-green / sky-blue / vermillion
    pal      <- stats::setNames(c("#999999", "#009E73", "#56B4E9", "#D55E00"), levs)

    cat.vec <- ifelse(pass.p & pass.fc, lab.both,
               ifelse(pass.p, lab.sig, ifelse(pass.fc, lab.fc, lab.ns)))
    data.df[["Significance"]] <- factor(cat.vec, levels = levs)

    # --- choose genes to label: the doubly-significant set, budget 'numlab' ------
    #     total <= numlab, split between up- and down-regulated in proportion to the
    #     skew, but with a guaranteed minimum so the minor side is never empty
    lab <- data.df[0, , drop = FALSE]
    if (isTRUE(point.lab)) {
      cand <- data.df[pass.p & pass.fc, , drop = FALSE]
      cand <- cand[order(cand[[property.column]]), , drop = FALSE]  # most significant first
      n.dn <- sum(cand$logFC < 0)
      n.up <- sum(cand$logFC > 0)
      tot  <- min(numlab, n.dn + n.up)
      if (tot > 0) {
        avail <- c(down = n.dn, up = n.up)
        alloc <- pmin(avail, round(tot * avail / sum(avail)))       # skew-weighted
        floor.each <- max(1L, tot %/% 4L)                           # never crush the minor side
        alloc <- pmax(alloc, ifelse(avail > 0L, pmin(avail, floor.each), 0L))
        while (sum(alloc) > tot) alloc[which.max(alloc)] <- alloc[which.max(alloc)] - 1L
        while (sum(alloc) < tot) {
          room <- avail - alloc
          if (max(room) <= 0L) break
          alloc[which.max(room)] <- alloc[which.max(room)] + 1L
        }
        lab <- rbind(utils::head(cand[cand$logFC < 0, , drop = FALSE], alloc[["down"]]),
                     utils::head(cand[cand$logFC > 0, , drop = FALSE], alloc[["up"]]))
        dw_log(sprintf("    Labelling %d significant gene(s): %d down / %d up",
                       nrow(lab), alloc[["down"]], alloc[["up"]]), "\n")
        if (n.dn + n.up > numlab)
          dw_log(sprintf("     (%d significant, capped at numlab = %d; split weighted by skew)\n",
                         n.dn + n.up, numlab))
      }
    }

    # --- assemble the plot ------------------------------------------------------
    line.tag <- paste0(if (property.to.plot == "fdr") "q" else "p", " = ", property.thr)

    g <- ggplot(data.df, aes(x = .data[["logFC"]], y = .data[["neg.log10"]])) +
      geom_point(aes(colour = .data[["Significance"]]), size = base.size / 8, alpha = 0.7) +
      scale_colour_manual(values = pal, drop = FALSE) +
      geom_hline(yintercept = -log10(property.thr), colour = "red", linetype = "dashed") +
      geom_vline(xintercept =  logfc.thr, colour = "red", linetype = "dashed") +
      geom_vline(xintercept = -logfc.thr, colour = "red", linetype = "dashed") +
      annotate("text", x = min(data.df$logFC, na.rm = TRUE),
               y = -log10(property.thr), label = line.tag,
               hjust = 0, vjust = -0.4, size = base.size / 3.5, colour = "red") +
      labs(x = "Logarithmic fold change", y = paste0("-log10(", measure, ")"),
           colour = NULL, title = main, subtitle = paste0("(", measure, "-values)")) +
      theme_bw(base_size = base.size) +
      theme(panel.border = element_blank(),
            axis.line = element_line(colour = "black"),
            panel.grid.major = element_line(linewidth = 0.2),
            panel.grid.minor = element_line(linewidth = 0.2),
            plot.title = element_text(hjust = 0.5),
            plot.subtitle = element_text(hjust = 0.5))

    if (pretty.breaks)
      g <- g + scale_x_continuous(breaks = scales::breaks_extended(6)) +
               scale_y_continuous(breaks = scales::breaks_extended(8))

    if (isTRUE(point.lab) && nrow(lab) > 0)
      g <- g + ggrepel::geom_text_repel(
        data = lab, inherit.aes = FALSE,
        aes(x = .data[["logFC"]], y = .data[["neg.log10"]], label = .data[[sym.col]]),
        size = base.size / 3, colour = "gray30", max.overlaps = 15,
        box.padding = grid::unit(0.35, "lines"),
        point.padding = grid::unit(0.3, "lines"), show.legend = FALSE)

    g
  }

#' Function to generate a Volcano plot using `ggplot2`
#' @description `diff_expr_volcano_plot` generates two Volcano plots highlighting genes that are differentially expressed beyond custom thresholds for significance (set by parameters `p.thr` and fdr.thr`) and differential expression level (set by parameter `logfc.thr`).
#' @param d3 \code{data.frame}. Data frame containing all necessary columns to generate a Volcano plot with gene labels (at least p-values, FDR values, log-ratios and gene symbols or other IDs)
#' @param id \code{character}. Name of the gene ID column. Can be the same as `sym.col` but usually refers to an additional column with, e.g., Ensembl Gene IDs.
#' @param sym.col \code{character}. Name of column with gene symbols, e.g., HGNC Symbols.
#' @param main \code{character}. Main plot title. (Will be complemented with additional information, e.g., 'FDR' when labelling according to and FDR threshold.)
#' @param p.thr \code{numeric}. Plotted values with a P-Value below this threshold will be labelled in the P-Value plot.
#' @param fdr.thr \code{numeric}. Plotted values with a FDR below this threshold will be labelled in the FDR plot.
#' @param logfc.thr \code{numeric}. Plotted (`abs`olute) values above this threshold will have bigger dots.
#' @param numlab \code{numeric}. Maximum number of labels per plot. Takes precedence to numbers calculated based on `p.thr` and `fdr.thr`.
#' @param point.lab \code{logical}. Should points be labelled, at all?
#' @param base.size \code{numeric}. Overall text/point scale passed on to \code{prepare_volcano_of_given_property}
#'   (the `ggplot2` \code{base_size}, default 16). Raise it to enlarge legend, axes, labels and points together
#'   when printing to a large device (e.g. the 15x15 inch pipeline PDF).
#' @return A named \code{list} of \code{ggplot} objects with the elements \code{FDR} and \code{Pval},
#'   holding the volcano plot with points highlighted by false discovery rate and by p-value,
#'   respectively. Both plots are additionally printed, so that they are captured when the function is
#'   called with an open plotting device.
#' @examples
#' \donttest{
#' set.seed(1)
#' d3 <- data.frame(gene_symbol = paste0("G", 1:100),
#'                  logFC = rnorm(100), PValue = runif(100), FDR = runif(100))
#' g <- diff_expr_volcano_plot(d3, id = "gene_symbol", sym.col = "gene_symbol")
#' names(g)
#' }
#' @export
diff_expr_volcano_plot <-
		function(d3, id, sym.col="gene_symbol", main=NULL, p.thr=0.05, fdr.thr=0.05, logfc.thr=1, numlab=25, point.lab=TRUE, base.size=16)
{
	pv.col <- dw_find_col(names(d3), "^p\\.{0,1}val[e-u]{0,2}$", "p-value")
	fdr.col <- dw_find_col(names(d3), "^fdr$|^adj*\\.{0,1}p\\.{0,1}val[e-u]{0,2}$", "FDR")

	if (point.lab) {
		cont.dat <- d3[, c(id, sym.col, "logFC", pv.col, fdr.col)]
		if (any(cont.dat[[sym.col]]=="") || any(is.na(cont.dat[[sym.col]]))) {
			repl <- which(cont.dat[[sym.col]]==""|is.na(cont.dat[[sym.col]]))
			cont.dat[repl, sym.col] <- cont.dat[repl, as.character(id)]
		}
	} else {
		cont.dat <- d3[, c(id, "logFC", pv.col, fdr.col)]
	}


	g.l <- list()
	g.l[["FDR"]] = prepare_volcano_of_given_property(data.df = cont.dat,
			property.to.plot = 'fdr',
			property.column = fdr.col,
			property.thr = fdr.thr,
			logfc.thr = logfc.thr,
			main = main,
			numlab = numlab,
			point.lab = TRUE,
			sym.col = sym.col,
			base.size = base.size)

	g.l[["Pval"]] = prepare_volcano_of_given_property(data.df = cont.dat,
			property.to.plot = 'p',
			property.column = pv.col,
			property.thr = p.thr,
			logfc.thr = logfc.thr,
			main = main,
			numlab = numlab,
			point.lab = TRUE,
			sym.col = sym.col,
			base.size = base.size)

	#Printing the plots so they can be stored into pdf-file when this function is called
	suppressWarnings(print(g.l[["Pval"]]))
	suppressWarnings(print(g.l[["FDR"]]))

	return(g.l)

}

#' Function to generate a histogram of the P-Value distribution
#' @param d3 \code{data.frame}. Data frame containing a P-Value column to generate a histogram. The column name will be determined
#'   by pattern matching.
#' @return Invisibly, an object of class \code{histogram} as returned by \code{\link[graphics]{hist}}.
#'   Called for its side effect of drawing the p-value distribution.
#' @examples
#' d3 <- data.frame(ID = paste0("g", 1:200), PValue = runif(200))
#' diff_expr_pval_hist_plot(d3)
#' @export
diff_expr_pval_hist_plot <-
		function(d3)
{
	dw_log("   Histogram of P-value distribution...\n")
	pv.col <- dw_find_col(names(d3), "^p\\.{0,1}val[e-u]{0,2}$", "p-value")
	hist(d3[[pv.col]],breaks=20, xlab="P Value", ylab="Frequency", main="P-value distribution")
}


