####

# ROI DE Analysis ####

####

#' getDiffExp

#'

#' Get differential expression with DESeq2

#'

#' @param object a VoltRon object

#' @param assay assay name (exp: Assay1) or assay class (exp: Visium, Xenium), see \link{SampleMetadata}. 

#' if NULL, the default assay will be used, see \link{vrMainAssay}.

#' @param group.by the categorical variable from metadata to get differentially expressed features across

#' @param group.base Optional, the base category in \code{group.by} which is used as control group

#' @param covariates the covariate variable for the design formula

#' @param method the method for DE analysis, e.g. DESeq2

#'

#' @export

#'

getDiffExp <- function(object, assay = NULL, group.by, group.base = NULL, covariates = NULL, method = "DESeq2"){

  # get data and metadata

  data <- vrData(object, assay = assay)

  metadata <- Metadata(object, assay = assay)

  # check groups

  if(group.by %in% colnames(metadata)){

    if(!is.null(group.base)){

      if(!group.base %in% unique(as.vector(metadata[[group.by]])))

        stop("Please specify a group that is included in the group.by column of metadata to define the base group!")

    } else{

      group.base <- levels(factor(as.vector(metadata[[group.by]])))[1]

    }

  } else {

    stop("Column ", group.by, " cannot be found in metadata!")

  }

  # select a method

  results <-

    switch(method,

           DESeq2 = {

             getDiffExpDESeq2(data, metadata, group.by = group.by, group.base = group.base, covariates = covariates)

           })

  # return

  return(results)

}

#' getDiffExpDESeq2

#'

#' @param data the raw data set

#' @param metadata the metadata

#' @param group.by the categorical variable from metadata to get differentially expressed features across

#' @param group.base Optional, the base category in \code{group.by} which is used as control group

#' @param covariates the covariate variable for the design formula

#'

#' @importFrom stats as.formula

#'

#' @noRd

getDiffExpDESeq2 <- function(data, metadata, group.by, group.base = NULL, covariates){

  if (!requireNamespace('DESeq2'))

    stop("Please install DESeq2 package to find differentially expressed genes with DESeq2 method")

  # experimental design for deseq2

  # make formula

  if(is.null(covariates)){

    group.by.data <- as.vector(metadata[[group.by]])

    uniq_groups <- unique(group.by.data)

    if(!is.null(group.base))

      uniq_groups <- c(group.base, uniq_groups[!uniq_groups %in% group.base])

    group.by.data <- factor(group.by.data, levels = uniq_groups)

    colData <- data.frame(group.by.data)

    colnames(colData) <- c(group.by, covariates)

    deseq2.formula <- stats::as.formula(paste0("~", group.by))

  } else {

    if(all(covariates %in% colnames(metadata))){

      design.data <- metadata[,c(group.by, covariates)]

      uniq_groups <- unique(design.data[[group.by]])

      if(is.null(group.base))

        uniq_groups <- c(group.base, uniq_groups[!uniq_groups %in% group.base])

      group.by.data <- factor(design.data[[group.by]], levels = uniq_groups)

      design.data[[group.by]] <- group.by.data

      colData <- data.frame(design.data)

      colnames(colData) <- c(group.by, covariates)

      deseq2.formula <- stats::as.formula(paste0("~", group.by, " + ", paste(covariates, collapse = " + ")))

    } else {

      stop("Columns ", paste(covariates, collapse = ", "), " cannot be found in metadata!")

    }

  }

  # run DESeq2

  data <- as.matrix(as(data, "dgCMatrix"))

  dds <- DESeq2::DESeqDataSetFromMatrix(countData = data, colData = colData, design = deseq2.formula)

  dds <- DESeq2::DESeq(dds)

  all_results <- NULL

  for(i in seq_len(length(uniq_groups)-1)){

    for(j in (i+1):length(uniq_groups)){

      comparison <- c(group.by, uniq_groups[i], uniq_groups[j])

      cur_results <- as.data.frame(DESeq2::results(dds, comparison))

      cur_results <- data.frame(cur_results, gene = rownames(cur_results), comparison = paste(comparison, collapse = "_"))

      all_results <- rbind(all_results, cur_results)

    }

  }

  # return

  return(all_results)

}