#' @include allgenerics.R

#'

NULL

####

# Normalization ####

####

normalizeDataVoltRon <- function(object, assay = NULL, method = "LogNorm", desiredQuantile = 0.9, scale = 0.2, sizefactor = 10000, feat_type = NULL) {

  # get assay names

  assay_names <- vrAssayNames(object, assay = assay)

  # normalize assays

  for(assy in assay_names){

    cur_assay <- object[[assy]]

    object[[assy]] <- normalizeData(cur_assay, method = method, desiredQuantile = desiredQuantile, scale = scale, sizefactor = sizefactor, feat_type = feat_type)

  }

  # return

  return(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 method the normalization method: "LogNorm", "Q3Norm", "LogQ3Norm" or "CLR"

#' @param desiredQuantile the quantile of the data if "QuanNorm" or "LogQuanNorm" is selected as \code{method}.

#' @param scale the scale parameter for the hyperbolic arcsine transformation

#' @param sizefactor size factor if \code{method} is selected as \code{LogNorm}

#' @param feat_type the feature set type

#' 

#' @rdname normalizeData

#' @method normalizeData VoltRon

#'

#' @export

setMethod("normalizeData", "VoltRon", normalizeDataVoltRon)

normalizeDatavrAssay <- function(object, method = "LogNorm", desiredQuantile = 0.9, scale = 0.2, sizefactor = 10000, feat_type = NULL) {

  # size factor

  rawdata <- vrData(object, feat_type = feat_type, norm = FALSE)

  if(!is.numeric(desiredQuantile)){

    stop("desiredQuantile should be numeric")

  } else {

    if(!findInterval(desiredQuantile, c(0,1)) == 1L){

      stop("desiredQuantile should be between [0,1]")

    }

  }

  # normalization method

  if(method == "LogNorm"){

    normdata <- LogNorm(rawdata, colSums(rawdata), sizefactor)

  } else if(method == "Q3Norm") {

    # rawdata[rawdata==0] <- 1

    qs <- getColQuantiles(rawdata, desiredQuantile)

    normdata <- getDivideSweep(rawdata, qs / exp(mean(log(qs))))

  } else if(method == "LogQ3Norm") {

    # rawdata[rawdata==0] <- 1

    qs <- getColQuantiles(rawdata, desiredQuantile)

    normdata <- getDivideSweep(rawdata, qs / exp(mean(log(qs))))

    normdata <- log(normdata + 1)

  } else if(method == "CLR") {

    normdata <- getDivideSweep(rawdata, colSums(rawdata))

    normdata <- apply(normdata, 2, function(x) {

      log1p(x = x / (exp(x = sum(log1p(x = x[x > 0]), na.rm = TRUE) / length(x = x))))

    })

  } else if(method == "hyper.arcsine") {

    normdata <- asinh(rawdata/scale)

  } else {

    stop('Please select one of these methods: "LogNorm", "Q3Norm", "LogQ3Norm" or "CLR"')

  }

  # get normalized data

  catch_connect1 <- try(slot(object, name = "data"), silent = TRUE)

  catch_connect2 <- try(slot(object, name = "rawdata"), silent = TRUE)

  if(!is(catch_connect1, 'try-error') && !methods::is(catch_connect1,'error')){

    if(is.null(feat_type))

      feat_type <- vrMainFeatureType(object)

    object@data[[paste0(feat_type, "_norm")]] <- normdata

  } else if(!is(catch_connect2, 'try-error') && !methods::is(catch_connect2,'error')){

    object@normdata <- normdata

  }

  # return

  return(object)

}

#' @rdname normalizeData

#' @method normalizeData vrAssay

#'

#' @importFrom stats quantile

#'

#' @export

setMethod("normalizeData", "vrAssay", normalizeDatavrAssay)

#' @rdname normalizeData

#' @method normalizeData vrAssayV2

#'

#' @importFrom stats quantile

#'

#' @export

setMethod("normalizeData", "vrAssayV2", normalizeDatavrAssay)

LogNorm <- function(rawdata, coldepth, sizefactor){

  if(inherits(rawdata, "IterableMatrix")){

    if(!requireNamespace("BPCells"))

      stop("You have to install BPCells!")

    normdata <- BPCells::t(BPCells::t(rawdata)/coldepth)

    normdata <- BPCells::log1p_slow(normdata*sizefactor)

  } else {

    normdata <- sweep(rawdata, 2L, coldepth, FUN = "/")

    normdata <- log(normdata*sizefactor + 1)

  }

  return(normdata)

}

getDivideSweep <- function(rawdata, divisor){

  if(inherits(rawdata, "IterableMatrix")){

    if(!requireNamespace("BPCells"))

      stop("You have to install BPCells!")

    return(BPCells::t(BPCells::t(rawdata)/divisor))

  } else {

    return(sweep(rawdata, 2L, divisor, FUN = "/"))

  }

  return(rawdata)

}

####

# Features ####

####

getFeaturesVoltRon <- function(object, assay = NULL, max.count = 1, n = 3000){

  # get assay names

  assay_names <- vrAssayNames(object, assay = assay)

  # get features for all coordinates

  for(assy in assay_names){

    object[[assy]] <- getFeatures(object[[assy]], max.count = max.count, n = n)

  }

  # return

  return(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 max.count maximum count (across spatial points) for low count filtering

#' @param n the top number of variable features 

#' 

#' @rdname getFeatures

#'

#' @export

setMethod("getFeatures", "VoltRon", getFeaturesVoltRon)

getFeaturesvrAssay <- function(object, max.count = 1, n = 3000){

  # get data and coordinates

  rawdata <- vrData(object, norm = FALSE)

  coords <- vrCoordinates(object)

  features <- vrFeatures(object)

  # eliminate genes with low counts

  # keep.genes <- which(apply(rawdata,1,max) > max.count)

  keep.genes <- getMaxCount(rawdata, max.count)

  # vst estimation

  # vst_data <- data.frame(mean = Matrix::rowMeans(rawdata), var = apply(rawdata, 1, stats::var))

  vst_data <- getVstData(rawdata)

  loess_data <- vst_data[keep.genes,]

  loess_results <- stats::loess(var~mean, loess_data, span = 0.3)

  vst_data$adj_var <- 0

  vst_data$rank <- 0

  vst_data[keep.genes,]$adj_var <- stats::predict(loess_results)

  vst_data[keep.genes,]$rank <- order(order(vst_data$adj_var[keep.genes], decreasing = TRUE))

  # set feature data

  vrFeatureData(object) <- vst_data

  # return

  return(object)

}

#' @rdname getFeatures

#'

#' @importFrom stats loess predict var

#' @importFrom Matrix rowMeans

#'

#' @export

setMethod("getFeatures", "vrAssay", getFeaturesvrAssay)

#' @rdname getFeatures

#'

#' @importFrom stats loess predict var

#' @importFrom Matrix rowMeans

#'

#' @export

setMethod("getFeatures", "vrAssayV2", getFeaturesvrAssay)

getVstData <- function(rawdata){

  if(inherits(rawdata, "IterableMatrix")){

    if(!requireNamespace("BPCells"))

      stop("You have to install BPCells!")

    mean_data <- BPCells::rowMeans(rawdata)

    var_data <- BPCells::rowSums(rawdata^2)

    var_data <- (var_data - mean_data^2/nrow(rawdata))/(nrow(rawdata)-1)

    # var_data <- BPCells::matrix_stats(rawdata, row_stats="variance")

  } else {

    mean_data <- Matrix::rowMeans(rawdata)

    var_data <- apply(rawdata, 1, stats::var)

  }

  vst_data <- data.frame(mean = mean_data, var = var_data)

  return(vst_data)

}

getMaxCount <- function(rawdata, max.count){

  if(inherits(rawdata, "IterableMatrix")){

    if(!requireNamespace("BPCells"))

      stop("You have to install BPCells!")

    rawdata <- rawdata > max.count

    keep.genes <- which(BPCells::rowSums(rawdata) > 0)

  } else {

    keep.genes <- which(apply(rawdata,1,max) > max.count)

  }

  return(keep.genes)

}

#' getVariableFeatures

#'

#' get shared variable features across multiple assays

#'

#' @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 n the number of features

#' @param ... additional arguements passed to \link{vrFeatureData}

#'

#' @importFrom dplyr full_join

#' @importFrom utils head

#'

#' @export

getVariableFeatures <- function(object, assay = NULL, n = 3000, ...){

  # get assay names

  assay_names <- vrAssayNames(object, assay = assay)

  # get features for all coordinates

  ranks <- NULL

  for(assy in assay_names){

    feature_data <- vrFeatureData(object[[assy]], ...)

    # if(nrow(feature_data) > 0){

    if(!is.null(feature_data)) {

      if(nrow(feature_data) > 0){

        feature_data$gene <- rownames(feature_data)

      }

    } else {

      feature_data <- data.frame(gene = vrFeatures(object[[assy]]), rank = NA)

    }

    if(is.null(ranks)){

      ranks <- feature_data[,c("gene", "rank")]

    } else {

      ranks <- ranks %>% full_join(feature_data[,c("gene", "rank")], by = c("gene" = "gene"))

    }

  }

  # get geometric mean of ranks, i.e. rank product statistic

  rownames_ranks <- ranks$gene

  ranks <- ranks[,!colnames(ranks) %in% "gene", drop = FALSE]

  ranks <- apply(ranks, 1, function(x) exp(mean(log(x))))

  # names(ranks) <- rownames(feature_data)

  names(ranks) <- rownames_ranks

  ranks <- ranks[ranks != 0]

  # get selected features

  if(length(ranks[!is.na(ranks)]) > 0){

    selected_features <- names(utils::head(sort(ranks, decreasing = FALSE), n))

  } else {

    selected_features <- vrFeatures(object, assay = assay)

  }

  # return

  return(selected_features)

}

####

# vrEmbeddings ####

####

#' getPCA

#'

#' calculate PCA of the VoltRon objects

#'

#' @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 the selected features for PCA reduction

#' @param dims the number of dimensions extracted from PCA

#' @param type the key name for the embedding, default: pca

#' @param overwrite Whether the existing embedding with name 'type' should be overwritten in \link{vrEmbeddings}

#' @param seed seed

#'

#' @importFrom irlba irlba

#'

#' @export

getPCA <- function(object, assay = NULL, features = NULL, dims = 30, type = "pca", overwrite = FALSE, seed = 1){

  # get assay names

  assay_names <- vrAssayNames(object, assay = assay)

  # get shared features and subset

  assay_features <- vrFeatures(object, assay = assay)

  # if there are features of a VoltRon object, then get variable features too

  if(length(assay_features) > 0) {

    if(is.null(features))

      features <- getVariableFeatures(object, assay = assay)

    object_subset <- subsetVoltRon(object, features = features)

    vrMainAssay(object_subset) <- vrMainAssay(object)

    # adjust extraction features length

    if(dims > length(features)){

      message("Requested more PC dimensions than existing features: dims = length(features) now!")

      dims <- length(features)

    }

  # if there are no features in VoltRon object, return the assay as itself

  } else {

    object_subset <- object

  }

  # get data

  normdata <- vrData(object_subset, assay = assay, norm = TRUE)

  # get PCA embedding

  set.seed(seed)

  if(inherits(normdata, "IterableMatrix")){

    if(!requireNamespace("BPCells"))

      stop("You have to install BPCells!")

    svd <- BPCells::svds(normdata, k=dims)

    pr.data <- BPCells::multiply_cols(svd$v, svd$d)

  } else {

    scale.data <- apply(normdata, 1, scale)

    pr.data <- irlba::prcomp_irlba(scale.data, n=dims, center=colMeans(scale.data))

    pr.data <- pr.data$x 

  }

  # change colnames

  colnames(pr.data) <- paste0("PC", seq_len(dims))

  rownames(pr.data) <- colnames(normdata)

  # set Embeddings

  vrEmbeddings(object, assay = assay, type = type, overwrite = overwrite) <- pr.data

  # return

  return(object)

}

#' getUMAP

#'

#' calculate UMAP of the VoltRon objects

#'

#' @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 data.type the type of data used to calculate UMAP from: "pca" (default), "raw" or "norm"

#' @param dims the number of dimensions extracted from PCA

#' @param umap.key the name of the umap embedding, default: umap

#' @param overwrite Whether the existing embedding with name 'type' should be overwritten in \link{vrEmbeddings}

#' @param seed seed

#'

#' @importFrom uwot umap

#' @importFrom Matrix t

#'

#' @export

#'

getUMAP <- function(object, assay = NULL, data.type = "pca", dims = seq_len(30), umap.key = "umap", overwrite = FALSE, seed = 1){

  # get data

  if(data.type %in% c("raw", "norm")){

    data <- vrData(object, assay = assay, norm = (data.type == "norm"))

    data <- as.matrix(as(Matrix::t(data),"dgCMatrix"))

  } else{

    embedding_names <- vrEmbeddingNames(object)

    if(data.type %in% vrEmbeddingNames(object)) {

      data <- vrEmbeddings(object, assay = assay, type = data.type, dims = dims)

    } else {

      stop("Please provide a data type from one of three choices: raw, norm and pca")

    }

  }

  # get umap

  set.seed(seed)

  umap_data <- uwot::umap(data)

  colnames(umap_data) <- c("x", "y")

  vrEmbeddings(object, assay = assay, type = umap.key, overwrite = overwrite) <- umap_data

  # return

  return(object)

}

####

# Image Processing ####

####

#' split_into_tiles

#'

#' split image raster data into tiles

#'

#' @param image_data image raster data

#' @param tile_size tile size

#'

#' @noRd

split_into_tiles <- function(image_data, tile_size = 10) {

  n_rows <- nrow(image_data)

  n_cols <- ncol(image_data)

  # Calculate the number of tiles in rows and columns

  n_row_tiles <- n_rows %/% tile_size

  n_col_tiles <- n_cols %/% tile_size

  # Initialize an empty list to store tiles

  tiles <- list()

  # Loop through the image data matrix to extract tiles

  for (i in seq_len(n_row_tiles)) {

    for (j in seq_len(n_col_tiles)) {

      # Calculate the indices for the current tile

      start_row <- (i - 1) * tile_size + 1

      end_row <- i * tile_size

      start_col <- (j - 1) * tile_size + 1

      end_col <- j * tile_size

      # Extract the current tile from the image data matrix

      tile <- image_data[start_row:end_row, start_col:end_col]

      # Store the tile in the list

      tiles[[length(tiles) + 1]] <- tile

    }

  }

  # Return the list of tiles

  return(tiles)

}