####
# Spot Deconvolution ####
####

#' getDeconvolution
#'
#' Calculate deconvolution of spots and ROIs
#'
#' @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 features features
#' @param sc.object Seurat Object
#' @param sc.assay assay of the Seurat Object used for the single cell data reference
#' @param sc.cluster metadata column variable used for the single cell data reference
#' @param method Deconvolution method, RCTD (spot), SPOTlight (spot), MuSiC (ROI)
#' @param ... additional parameters passed to method specific functions, e.g. RCTD, MuSiC.
#'
#' @export
getDeconvolution <- function(object, assay = NULL, features = NULL, sc.object, sc.assay = "RNA", sc.cluster = "seurat_clusters", method = "RCTD", ...){

  # sample metadata
  sample.metadata <- SampleMetadata(object)

  # get assay names
  assay_names <- vrAssayNames(object, assay = assay)

  # check assay type
  assay.types <- unique(vrAssayTypes(object, assay = assay))

  if(length(assay.types) > 1){
    stop("Please make sure that only assays of one assay type (cell/spot/ROI) are being deconvoluted at a time!")
  } else {

    # make single cell reference
    reference <- getDeconReference(sc.object = sc.object, sc.assay = sc.assay, sc.cluster = sc.cluster,
                                   method = method, assay.type = assay.types)

    # run a list of assays
    for(assy in assay_names){

      # get assay
      cur_assay <- object[[assy]]

      # RCTD
      rawdata <- getDeconSingle(object = cur_assay, features = features, reference = reference, method = method, ...)
      
      # add cell type mixtures as new featureset
      object <- addFeature(object, assay = assy, data = rawdata, feature_name = "Decon")
    }
  }

  return(object)
}

#' getDeconReference
#'
#' Establish and produce the single cell reference for deconvolution
#'
#' @param sc.object Seurat Object
#' @param sc.assay assay of the Seurat Object used for the single cell data reference
#' @param sc.cluster metadata column variable used for the single cell data reference
#' @param method Deconvolution method, RCTD (spot), SPOTlight (spot), MuSiC (ROI)
#' @param assay.type the assay type associated with the single cell deconvolution reference
#'
#' @noRd
getDeconReference <- function(sc.object, sc.assay = "RNA", sc.cluster = "seurat_clusters", method = "RCTD", assay.type = NULL){

  # Deconvolute for spots
  if(assay.type == "spot"){

    # check method
    if(!method %in% c("RCTD")){
      message("The selected method is not provided for spot deconvolution. Switching to RCTD")
      method <- "RCTD"
    }

    # deconvolution with RCTD
    if(method == "RCTD"){

      if (!requireNamespace('spacexr'))
        stop("Please install spacexr package to use the RCTD algorithm")
      if (!requireNamespace('Seurat'))
        stop("Please install Seurat package for using Seurat objects")

      message("Configuring Single Cell Assay (reference) ...\n")
      sccounts <- Seurat::GetAssayData(sc.object[[sc.assay]], slot = "counts")
      # sccounts <- as.matrix(apply(sccounts,2,ceiling))
      rownames(sccounts) <- rownames(sc.object[[sc.assay]])
      cell_types <- as.factor(sc.object@meta.data[[sc.cluster]])
      names(cell_types) <- colnames(sc.object)
      sc.nUMI <- colSums(sccounts)
      names(sc.nUMI) <- colnames(sc.object)
      reference <- spacexr::Reference(sccounts, cell_types, sc.nUMI)
    }

  # Deconvolute for ROIs
  } else if(assay.type == "ROI"){

    # check method
    if(!method %in% c("MuSiC")){
      message("The selected method is not provided for ROI deconvolution. Switching to MuSiC")
      method <- "MuSiC"
    }

    # deconvolution with MuSiC
    if(method == "MuSiC"){

      message("Configuring Single Cell Assay (reference) ...\n")
      if(inherits(sc.object, "SingleCellExperiment")){
        sc.object$music_decon_clusters <- sc.object[[sc.cluster]]
        reference <- sc.object
      } else if(inherits(sc.object, "Seurat")){
        sc.object$music_decon_clusters <- sc.object@meta.data[[sc.cluster]]
        sccounts <- Seurat::GetAssayData(sc.object[[sc.assay]], slot = "counts")
        sccounts <- as.matrix(apply(sccounts,2,ceiling))
        rownames(sccounts) <- rownames(sc.object[[sc.assay]])
        reference <- Seurat::as.SingleCellExperiment(Seurat::CreateSeuratObject(sccounts, meta.data = sc.object@meta.data))
      } else{
        stop("'sc.object' should either be of a Seurat or SingleCellExperiment class!")
      }

    }
  }

  # return
  return(reference)
}

#' getDeconSingle
#'
#' Calculate deconvolution of spots and ROIs of a single vrAssay object
#'
#' @param object a vrAssay object
#' @param features features
#' @param reference the single cell deconvolution reference, generated by \code{getDeconReference}
#' @param method Deconvolution method, RCTD (spot), SPOTlight (spot), MuSiC (ROI)
#' @param ... additional parameters passed to method specific functions
#'
#' @noRd
getDeconSingle <- function(object, features = features, reference, method = "RCTD", ...){

  # get assay type
  assay.type <- vrAssayTypes(object)

  if(assay.type == "spot"){

    # check method
    if(!method %in% c("RCTD")){
      message("The selected method is not provided for spot deconvolution. Switching to RCTD")
      method <- "RCTD"
    }

    if(method == "RCTD"){
      message("Running RCTD for spot deconvolution ...\n")
      rawdata <- getRCTD(object = object, features = features, reference = reference, ...)
    }

  } else if(assay.type == "ROI"){

    # check method
    if(!method %in% c("MuSiC")){
      message("The selected method is not provided for ROI deconvolution. Switching to MuSiC")
      method <- "MuSiC"
    }

    if(method == "MuSiC"){
      message("Running MuSiC for ROI deconvolution ...\n")
      rawdata <- getMuSiC(object = object, features = features, reference = reference, ...)
    }

  }

  # return
  return(rawdata)
}

#' getRCTD
#'
#' Calculate RCTD deconvolution for spot transcriptomics
#'
#' @param object a VoltRon object
#' @param features features
#' @param reference the single cell deconvolution reference, generated \code{getDeconReference}
#' @param ... additional parameters passed to \code{create.RCTD} function
#'
#' @noRd
getRCTD <- function(object, features = NULL, reference, ...){

  if (!requireNamespace('spacexr'))
    stop("Please install spacexr package to use the RCTD algorithm: devtools::install_github('dmcable/spacexr')")
  if (!requireNamespace('Seurat'))
    stop("Please install Seurat package for using Seurat objects: install.packages('Seurat')")

  # create spatial data
  message("Configuring Spatial Assay ...\n")
  spatialcounts <- vrData(object, norm = FALSE)
  coords <- as.data.frame(as(vrCoordinates(object), "dgCMatrix"))[,c("x", "y")]
  spatialnUMI <- colSums(spatialcounts)
  spatialdata <- spacexr::SpatialRNA(coords, spatialcounts, spatialnUMI)

  # Run RCTD
  myRCTD <- spacexr::create.RCTD(spatialdata, reference, ...)
  message("Calculating Cell Type Compositions of spots with RCTD ...\n")
  myRCTD <- quiet(spacexr::run.RCTD(myRCTD, doublet_mode = 'full'))
  results <- as.matrix(myRCTD@results$weights)
  norm_weights <- t(sweep(results, 1, rowSums(results), "/"))

  # return
  return(norm_weights)
}

#' getMuSiC
#'
#' Calculate MuSiC deconvolution for ROIs
#'
#' @param object a vrAssay object
#' @param features features
#' @param reference the single cell deconvolution reference, generated \code{getDeconReference}
#' @param sc.samples metadata column in Seurat that provides the samples in the single cell data
#'
#' @noRd
getMuSiC <- function(object, features = NULL, reference, sc.samples = NULL){

  if (!requireNamespace('Seurat'))
    stop("Please install Seurat package for using Seurat objects: install.packages('Seurat')")
  if (!requireNamespace('MuSiC'))
    stop("Please install MuSiC package for ROI deconvolution: devtools::install_github('xuranw/MuSiC')")
  if (!requireNamespace('SingleCellExperiment'))
    stop("Please install SingleCellExperiment package for ROI deconvolution: BiocManager::install('SingleCellExperiment')")

  if(is.null(sc.samples))
    stop("Please provide a metadata column for samples for MuSiC algorithm to work, e.g. sc.samples = Sample")

  if(is.null(features)) {
    features <- vrFeatures(object)
  }

  # Single cell reference data
  reference <- reference[rownames(reference) %in% features,]

  # data
  datax <- as.matrix(vrData(object))

  # common features
  common_features <- intersect(rownames(reference), rownames(datax))
  common_features <- intersect(common_features, features)
  if(length(common_features) < 5)
    stop("The number of common or selected features are less than 5!")
  reference <- reference[rownames(reference) %in% common_features,]
  datax <- datax[common_features,]

  # deconvolute
  message("Calculating Cell Type Compositions of ROIs with MuSiC ...\n")
  results <- MuSiC::music_prop(bulk.mtx = datax,
                        sc.sce = reference,
                        clusters = "music_decon_clusters",
                        samples = sc.samples,
                        verbose = T)
  t(results$Est.prop.weighted)
}