#' @include allgenerics.R
#'
NULL

####
# Nearest Neighbor graphs ####
####

#' Get profile specific neighbors
#'
#' Get neighbors of spatial points
#'
#' @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 method the method used for graph construction, SNN or kNN
#' @param k number of neighbors for kNN
#' @param data.type the type of embedding used for neighborhood calculation, e.g. raw counts (raw), normalized counts (norm), PCA embeddings (pca), UMAP embeddings (umap) etc.
#' @param dims the set of dimensions of the embedding data
#' @param graph.key the name of the graph
#'
#' @importFrom igraph add_edges simplify make_empty_graph vertices E<- E
#'
#' @export
getProfileNeighbors <- function(object, assay = NULL, method = "kNN", k = 10, data.type = "pca", dims = seq_len(30), graph.key = method){

  # get data
  if(data.type %in% c("raw", "norm")){
    nndata <- vrData(object, assay = assay, norm = (data.type == "norm"))
    nndata <- t(nndata)
  } else {
    embedding_names <- vrEmbeddingNames(object)
    if(data.type %in% vrEmbeddingNames(object)) {
      nndata <- 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")
    }
  }

  # find profile neighbors
  # if(knn.method == "FNN"){
  #   nnedges <- FNN::get.knn(nndata, k = k + 1)
  nnedges <- knn_annoy(nndata, k = k + 1)
  names(nnedges) <- c("nn.index", "nn.dist")
  weights <- NULL
  nnedges <-
    switch(method,
           SNN = {
             g.out <- build_snn_number(nnedges$nn.index)
             nnedges <- g.out[[1]]
             weights <- g.out[[2]]
             weights <- weights/(2 * (k+2) - weights)
             nnedges
           },
           kNN = {
             nnedges <- nnedges$nn.index
             nnedges <- cbind(seq_len(nrow(nndata)), nnedges)
             nnedges <- apply(nnedges, 1, function(x){
               do.call(c,lapply(x[-1], function(y) return(c(x[1],y))))
             })
             nnedges
           })
  nnedges <- rownames(nndata)[nnedges]

  # make graph and add edges
  graph <- make_empty_graph(directed = FALSE) + vertices(rownames(nndata))
  graph <- add_edges(graph, edges = nnedges)
  if(!is.null(weights))
    igraph::E(graph)$weight <- weights
  graph <- simplify(graph, remove.multiple = TRUE, remove.loops = FALSE)
  vrGraph(object, graph.type = graph.key) <- graph

  # return
  return(object)
}

#' knn_annoy
#' 
#' knn engine employed by RcppAnnoy package, adapted from \code{BPCells} package.
#' 
#' @rdname knn
#' 
#' @details **knn_annoy**: Use RcppAnnoy as knn engine
#' 
#' @param data data
#' @param query query data (Default: data)
#' @param k number of neighbors for kNN
#' @param n_trees Number of trees during index build time. More trees gives higher accuracy
#' @param search_k Number of nodes to inspect during the query, or -1 for default value. Higher number gives higher accuracy
#' 
#' @importFrom RcppAnnoy AnnoyEuclidean
knn_annoy <- function(data, query = data, k = 10, n_trees = 50, search_k = -1) {
  annoy <- new(RcppAnnoy::AnnoyEuclidean, ncol(data))
  for (i in seq_len(nrow(data))) {
    annoy$addItem(i - 1, data[i, ])
  }
  annoy$build(n_trees)
  
  idx <- matrix(nrow = nrow(query), ncol = k)
  dist <- matrix(nrow = nrow(query), ncol = k)
  rownames(idx) <- rownames(query)
  rownames(dist) <- rownames(query)
  for (i in seq_len(nrow(query))) {
    res <- annoy$getNNsByVectorList(query[i, ], k, search_k, include_distances = TRUE)
    idx[i, ] <- res$item + 1
    dist[i, ] <- res$dist
  }
  list(idx = idx, dist = dist)
}

####
# Clustering ####
####

#' getClusters
#'
#' Get clustering of the VoltRon object
#'
#' @param object a VoltRon object
#' @param resolution the resolution parameter for leiden clustering.
#' @param nclus The number of cluster centers for K-means clustering.
#' @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 method of clustering. Use 'leiden' to perform graph clustering and 'kmeans' for K-means based clustering
#' @param label the name for the newly created clustering column in the metadata
#' @param graph the graph type to be used
#' @param seed seed
#' @param abundance_limit the minimum number of points for a cluster, hence clusters with abundance lower than this limit will be appointed to other nearby clusters
#'
#' @importFrom igraph cluster_leiden
#' @importFrom stats kmeans
#' 
#' @export
getClusters <- function(object, resolution = 1, nclus = integer(0), assay = NULL, method = "leiden", label = "clusters", graph = "kNN", seed = 1, abundance_limit = 2){

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

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

  # get assays
  object_subset <- subsetVoltRon(object, assays = assay_names)

  # check clustering parameters
  .check_clustering_params(method, resolution, nclus, abundance_limit)
  
  # clustering
  set.seed(seed)
  if(method == "leiden"){
    object_graph <- vrGraph(object_subset, assay = assay, graph.type = graph)
    clusters <- igraph::cluster_leiden(object_graph, objective_function = "modularity", resolution = resolution) 
  } else if(method == "kmeans"){
    vrdata <- vrData(object_subset, norm = TRUE)
    clusters <- stats::kmeans(t(vrdata), centers = nclus)
    clusters <- list(names = names(clusters$cluster), membership = clusters$cluster)
  } else {
    stop("Unrecognized clustering method! Use either 'leiden' for graph clustering or 'kmeans' for K-means clustering")
  }

  # correct clustering
  clusters <- .correct_low_abundant_clusters(object_graph, clusters, abundance_limit)
    
  # update metadata
  spatialpoints <- vrSpatialPoints(object, assay = assay)
  membership <- setNames(rep(NA,length(spatialpoints)), spatialpoints)
  membership[clusters$names] <- clusters$membership
  object <- addMetadata(object, assay = assay, value = membership, label = label)

  # return
  return(object)
}

#' @noRd
.correct_low_abundant_clusters <- function(object_graph, clusters, abundance_limit){

  # cluster abundances
  cluster_abundance <- table(clusters$membership)
  
  # check if some clusters are low in abundance
  ind <- cluster_abundance < abundance_limit
  if(any(ind)){
    low_abundant_clusters <- names(cluster_abundance)[ind]
    clusters$membership[clusters$membership %in% low_abundant_clusters] <- NA
  } 
  return(clusters)
}

#' @noRd
.check_clustering_params <- function(method, resolution, nclus, abundance_limit){
  
  # method related params
  if(method == "leiden"){
    msg <- "Resolution must be a single numeric and above 0"
    if(!is.numeric(resolution))
      stop(msg) 
    if(length(resolution) > 1)
      stop(msg) 
    if(resolution == 0 | resolution < 0)
      stop(msg)
  } else if(method == "kmeans"){
    msg <- "Number of cluster centres (nclus) must be a single integer and should be above 1"
    if(!is.numeric(nclus))
      stop(msg) 
    if(length(nclus) > 1)
      stop(msg) 
    if(nclus %% 1 != 0)
      stop(msg) 
    if(nclus == 0)
      stop(msg) 
  }
  
  # low abundant clusters
  msg <- "Low abundance limit must be a single integer and should be above 0"
  if(!is.numeric(abundance_limit))
    stop(msg) 
  if(length(abundance_limit) > 1)
    stop(msg) 
  if(abundance_limit %% 1 != 0)
    stop(msg) 
  if(abundance_limit == 0 || abundance_limit < 0)
    stop(msg) 
}