.makeResultColumns <- function(contrast, resultsColumnNamePrefix) {
  assertthat::assert_that(
    is.matrix(contrast)|is.character(contrast), 
    msg = "'contrast' should be a contrast matrix or 
    a vector with the contrast names")
  
  if(is.matrix(contrast)){
    if (is.null(colnames(contrast))) {
      if (resultsColumnNamePrefix == "")
        resultsColumnNamePrefix <- "msqrobResults"
      if (ncol(contrast) > 1)
        colnames(contrast) <- seq_len(ncol(contrast))
      }
    return(paste0(resultsColumnNamePrefix, colnames(contrast)))
    } else 
      return(paste0(resultsColumnNamePrefix, contrast))
  }

.getMsqrob2Results <- function(object, row_var) {
  if (!row_var %in% colnames(rowData(object)))
    stop(
      sQuote(row_var), "is not found in the rowData of the ",
      "SummarizedExperiment object. Make sure you set the ",
      "same 'contrast' and 'resultsColumnNamePrefix' as for ",
      "'hypothesisTest()'."
    )
  rowData(object)[[row_var]]
}


#' Function to collect the inference tables generated by the msqrob2
#' statistical inference workflow.
#'
#' @param object A `SummarizedExperiment` object containing model
#'    inference performed by `hypothesisTest()`.
#' @param contrast A numeric `matrix` providing the contrast matrix
#'     used during model estimation or a character `vector` providing 
#'     the names of the contrast used during model estimation.
#' @param resultsColumnNamePrefix A `character(1)` providing the
#'     prefix used when running `hypothesisTest()`.
#' @param combine A `logical(1)` indicating whether the result tables
#'     should be combined in a single table (default) or return as a
#'     list of tables. When combined, two new variables are created:
#'     1. `contrast`tells from which contrast (column of the `contrast`
#'     matrix) the results where obtained; 2. `feature` provides the
#'     name of the modelled feature (this information is taken from
#'     the rownames of the tables, but this are made unique upon
#'     combining).
#' @return Result tables for contrasts combined in a single table or list
#' @examples
#'
#' ## Load example data
#' ## The data are a Feature object containing
#' ## a SummarizedExperiment named "peptide" with MaxQuant peptide intensities
#' ## The data are a subset of spike-in the human-ecoli study
#' ## The variable condition in the colData of the Feature object
#' ## contains information on the spike in condition a-e (from low to high)
#' data(pe)
#'
#' ## Aggregate peptide intensities in protein expression values 
#' ## Note that the peptide intensities were already normalised!
#' pe <- aggregateFeatures(pe, 
#'                         i = "peptide", 
#'                         fcol = "Proteins", 
#'                         name = "protein")
#'
#' ## Fit msqrob model
#' pe <- msqrob(pe, i = "protein", formula = ~condition)
#'
#' ## Define contrast
#' design <- model.matrix(~condition, data = colData(pe))
#' 
#' ## Assess log2 fold change between reference condition a and condition b
#' L <- makeContrast(
#'     contrasts = c("conditionb = 0",
#'                   "conditionc = 0",
#'                   "conditionc - conditionb = 0"),
#'     parameterNames = colnames(design))
#' 
#'
#' pe <- hypothesisTest(pe, i = "protein", L)
#' 
#' ## Extract Results
#' inference <- msqrobCollect(pe[["protein"]], L) 
#' head(inference)
#' 
#' @importFrom assertthat assert_that
#' @export

msqrobCollect <- function(object, contrast, resultsColumnNamePrefix = "",
                          combine = TRUE) {
  # check type of arguments
  assertthat::assert_that(
    is(object, "SummarizedExperiment"), 
    msg = "'object' should be of the 'SummarizedExperiment' class")
  assertthat::assert_that(
    is.matrix(contrast)|is.character(contrast), 
    msg = "'contrast' should be a contrast matrix or 
    a vector with the contrast names")
  assertthat::assert_that(is.character(resultsColumnNamePrefix))
  assertthat::assert_that(is.logical(combine))
  
  resultCols <- .makeResultColumns(contrast, resultsColumnNamePrefix)
  out <- lapply(resultCols, .getMsqrob2Results, object = object)
  names(out) <- resultCols
  if (combine) {
    for (i in names(out)) {
      out[[i]]$contrast <- i
      out[[i]]$feature <- rownames(out[[i]])
    }
    out <- do.call(rbind, out)
  }
  out
}

#' Volcano plot
#'
#' \code{plot_volcano} generates a volcano plot for the results of hypothesis 
#' tests generated with the msqrob2 hypothesisTest method. 
#'
#' @param resultsTable data frame returned by the msqrob2 hypothesisTest method
#' @param significanceLevel nominal FDR-level at which features are considered 
#' significant (value between 0 and 1)
#' @param  significance_colors a named vector with two colors, one for 
#' non-significant features named 'FALSE' and one for significant features 
#' named 'TRUE'. The default is 'c('FALSE' = 'black', 'TRUE' = 'red')'. 
#' @param opacity refers to the opacity of the dots in the plot. 
#' Values of opacity range from 0 to 1, with lower values corresponding to more 
#' transparent colors. The default is `opacity = 0.5`.
#' @return A ggplot object with the volcano plot. 
#' 
#' @importFrom assertthat assert_that
#' @importFrom ggplot2 aes ggplot geom_point theme_bw ylab scale_color_manual labs
#' @importFrom dplyr filter
#' @importFrom grDevices col2rgb
#' @importFrom rlang .data

#'
#' @examples
#'
#' # Load example data
#' # The data are a Feature object containing
#' # a SummarizedExperiment named "peptide" with MaxQuant peptide intensities
#' # The data are a subset of spike-in the human-ecoli study
#' # The variable condition in the colData of the Feature object
#' # contains information on the spike in condition a-e (from low to high)
#' data(pe)
#'
#' # Aggregate peptide intensities in protein expression values 
#' # Note that the peptide intensities were already normalised!
#' pe <- aggregateFeatures(pe, 
#'                         i = "peptide", 
#'                         fcol = "Proteins", 
#'                         name = "protein")
#'
#' # Fit msqrob model
#' pe <- msqrob(pe, i = "protein", formula = ~condition)
#'
#' # Define contrast
#' design <- model.matrix(~condition, data = colData(pe))
#' # Assess log2 fold change between reference condition a and condition b
#' L <- makeContrast(
#'     contrasts = c("conditionb = 0",
#'                   "conditionc = 0",
#'                   "conditionc - conditionb = 0"),
#'     parameterNames = colnames(design))
#' 
#'
#' pe <- hypothesisTest(pe, i = "protein", L)
#' 
#' # Volcano plots
#' library(ggplot2)
#' inference <- msqrobCollect(pe[["protein"]], L) 
#' plotVolcano(inference) + facet_wrap(~contrast)
#' plotVolcano(rowData(pe[["protein"]])$`conditionc`)
#' plotVolcano(rowData(pe[["protein"]])$`conditionc - conditionb`)
#' @export

plotVolcano <- function(resultsTable, 
                        significanceLevel  = 0.05, 
                        significance_colors = 
                          c("FALSE" = "black", "TRUE" = "red"), 
                        opacity = .5)
{
  assertthat::assert_that(
    is.data.frame(resultsTable), 
    msg = "'resultsTable' is no proper 'msqrob2::hypothesisTest' object")
  
  if (!all(c("logFC", "adjPval", "pval") %in% names(resultsTable))) 
    stop("'resultsTable' is no proper 'msqrob2::hypothesisTest' object")
  
  assertthat::assert_that(
    min(resultsTable$pval, na.rm = TRUE) >=0 & 
      max(resultsTable$pval, na.rm = TRUE) <=1, 
    msg = "p-values are not between 0 and 1")
  
  assertthat::assert_that(
    min(resultsTable$adjPval, na.rm = TRUE) >=0 & 
      max(resultsTable$adjPval, na.rm = TRUE) <=1, 
    msg = "Adjusted p-values are not between 0 and 1")
  
  assertthat::assert_that(is.numeric(resultsTable$logFC))
  
  assertthat::assert_that(
    significanceLevel  >=0 & significanceLevel  <= 1, 
    msg = "significance level is not between 0 and 1")
  
  assertthat::assert_that(
    opacity >=0 &  opacity <= 1, 
    msg = "Opacity is not between 0 and 1")
  
  assertthat::assert_that(
    all(names(significance_colors) %in% c("FALSE","TRUE")), 
    msg = "significance colors are not named 'FALSE' and 'TRUE'")
  
  assertthat::assert_that(
    tryCatch(
      {grDevices::col2rgb(significance_colors);TRUE}, 
      error = function(e) {FALSE}), 
    msg = "significance_colors are no valid colors")
  
  volcano_out <- 
    resultsTable |> 
    dplyr::filter(is.finite(.data$adjPval)) |> 
    ggplot2::ggplot( 
      ggplot2::aes(x = .data$logFC,
          y = -log10(.data$pval), 
          col = .data$adjPval < significanceLevel )
      ) + 
    ggplot2::geom_point(size = 1, alpha = opacity) + 
    ggplot2::theme_bw() +
    ylab("-log10(p-value)") +
    ggplot2::scale_color_manual(
      values = significance_colors, 
      breaks = c("TRUE", "FALSE")) +
    ggplot2::labs(color=paste0("FDR < ",significanceLevel ))
  return(volcano_out)
}