####
# Data Transfer ####
####

#' transferData
#'
#' transfer data across assays
#'
#' @param object a VoltRon object
#' @param from the name or class of assay whose data transfered to the second assay
#' @param to the name or class of target assay where data is transfered to
#' @param features the set of features from \link{vrFeatures} or metadata columns from \link{Metadata} that are transferred. 
#' Only one metadata feature can be transfered at a time.
#' @param new_feature_name the name of the new feature set created from the source assay defined in \code{from} argument.
#'
#' @export
transferData <- function(object, from = NULL, to = NULL, features = NULL, new_feature_name = NULL){

  # assay list
  assaytypes <- c("ROI", "spot", "cell", "molecule", "tile")
  
  # get Assay IDs from Names and IDs
  from <- vrAssayNames(object, assay = from)
  to <- vrAssayNames(object, assay = to)
  
  # check assay names
  if(length(from) > 1 | length(to) > 1){
    stop("For now, label transfer can only be accomplished across two assays")
  }
  
  # check if assays are in the same block
  sample.metadata <- SampleMetadata(object)
  from_assayclass <- sample.metadata[from, "Assay"]
  to_assayclass <- sample.metadata[to, "Assay"]
  samples <- sample.metadata[c(from, to), "Sample"]
  
  if(length(unique(samples)) > 1)
    stop("Selected assays have to be within the same sample block!")
  
  # get assay types
  to_object_type <- vrAssayTypes(object[[to]])
  from_object_type <- vrAssayTypes(object[[from]])
  
  if(which(assaytypes == to_object_type) > which(assaytypes == from_object_type) && from_object_type == "ROI"){
    return(transferLabels(object = object, from = from, to = to, features = features))
  } else {
    return(transferFeatureData(object = object, from = from, to = to, features = features, new_feature_name = new_feature_name))
  }
}

#' transferFeatureData
#'
#' transfer feature data across assays
#'
#' @param object a VoltRon object
#' @param from The ID of assay whose data transfer to the second assay
#' @param to The ID of the target assay where data is transfered to
#' @param features The set of data or metadata features that are transfered. Only one metadata feature can be transfered at a time.
#' @param new_feature_name the name of the new feature set created from the source assay defined in \code{from}.
#'
#' @noRd
transferFeatureData <- function(object, from = NULL, to = NULL, features = NULL, new_feature_name = NULL){

  # get assays and metadata
  from_object <- object[[from]]
  from_metadata <- Metadata(object, assay = from)
  to_object <- object[[to]]
  to_metadata <- Metadata(object, assay = to)

  # get assay types
  to_object_type <- vrAssayTypes(to_object)
  from_object_type <- vrAssayTypes(from_object)

  # get transfer data type
  if(to_object_type == "spot"){
    if(from_object_type == "cell"){
      new_assay <- getSpotsFromCells(from_object, from_metadata, to_object, features = features)
    }
  } else if(to_object_type == "cell"){
    if(from_object_type == "tile"){
      new_assay <- getCellsFromTiles(from_object, from_metadata, to_object, features = features)
    } else if(from_object_type == "spot"){
      new_assay <- getCellsFromSpots(from_object, from_metadata, to_object, features = features)
    }
  } else if(to_object_type == "ROI"){
    if(from_object_type == "cell"){
      new_assay <- getROIsFromCells(from_object, from_metadata, to_object, features = features)
    }
  }

  # add new feature set
  if(is.null(new_feature_name)){
    new_feature_name <- paste(vrMainFeatureType(object, assay = from)$Feature, "pseudo", sep = "_")
  }
  object <- addFeature(object, assay = to, data = new_assay, feature_name = new_feature_name)

  # return
  object
}

#' transferLabels
#'
#' transfer labels across assays
#'
#' @param object a VoltRon object
#' @param from The ID of assay whose data transfer to the second assay
#' @param to The ID of the target assay where data is transfered to
#' @param features The set of data or metadata features that are transferred. Only one metadata feature can be transferred at a time.
#'
#' @noRd
transferLabels <- function(object, from = NULL, to = NULL, features = NULL){
  
  # get assays and metadata
  from_object <- object[[from]]
  from_metadata <- Metadata(object, assay = from)
  to_object <- object[[to]]
  to_metadata <- Metadata(object, assay = to)
  
  # get assay types
  to_object_type <- vrAssayTypes(to_object)
  from_object_type <- vrAssayTypes(from_object)
  
  # get transfer data type
  if(from_object_type == "ROI" & to_object_type != "ROI"){
    
    # transfer labels
    transferedLabelsMetadata <- transferLabelsFromROI(from_object, from_metadata, to_object, to_metadata, features = features)

    # update metadata
    # metadata <- Metadata(object, assay = to)
    # entities <- vrSpatialPoints(to_object)
    # if(is.numeric(value)){
    #   metadata[[features]] <- as.numeric(NA)
    # } else {
    #   metadata[[features]] <- ""
    # }
    # if(is.null(rownames(metadata))){
    #   metadata[[features]][match(entities, as.vector(metadata$id))] <- value
    # } else {
    #   metadata[entities,][[features]] <- value
    # }
    # 
    # # transfer labels
    # Metadata(object, assay = to) <- metadata
    object <- addMetadata(object, assay = to, value = transferedLabelsMetadata[[features]], label = features)
  }
    
  # return
  object
}


#' getSpotsFromCells
#'
#' Generate Psuedo counts per spots and insert to as a separate image assay to the Visium object
#'
#' @param from_object The vrAssay object whose data transfer to the second assay
#' @param from_metadata the metadata associated with \code{from_object}
#' @param to_object The ID of the target vrAssay object where data is transfered to
#' @param features the name of the metadata feature to transfer, if NULL, the rawdata will be transfered
#'
#' @importFrom dplyr %>% right_join
#' @importFrom stats aggregate
#' @importFrom magick image_data
#'
#' @noRd
#'
getSpotsFromCells <- function(from_object, from_metadata = NULL, to_object, features = NULL) {

  # get the spot radius of Visium spots
  Vis_spotradius <- vrAssayParams(to_object, param = "spot.radius")

  # get cell and spot coordinates
  message("Cell to Spot Distances \n")
  coords_spots <- vrCoordinates(to_object)
  coords_cells <- vrCoordinates(from_object)

  # get distances from cells to spots
  # alldist <- flexclust::dist2(coords_cells, coords_spots)
  # cell_to_spot <- FNN::get.knnx(coords_spots, coords_cells, k = 1)
  cell_to_spot <- knn_annoy(coords_spots, coords_cells, k = 1)
  names(cell_to_spot) <- c("nn.index", "nn.dist")
  cell_to_spot_nnid <- vrSpatialPoints(to_object)[cell_to_spot$nn.index[,1]]
  names(cell_to_spot_nnid) <- rownames(coords_cells)
  cell_to_spot_nndist <- cell_to_spot$nn.dist[,1]
  names(cell_to_spot_nndist) <- rownames(coords_cells)
  cell_to_spot_nnid <- cell_to_spot_nnid[cell_to_spot_nndist < Vis_spotradius]

  # find associated spot for each cell
  message("Find associated spots for each cell \n")
  cell_to_spot_id <- names(cell_to_spot_nnid)

  # get data
  if(is.null(features)){
    raw_counts <- vrData(from_object, norm = FALSE)
  } else {
    data_features <- features[features %in% vrFeatures(from_object)]
    metadata_features <- features[features %in% colnames(from_metadata)]
    if(length(data_features) > 0){
      if(length(metadata_features) > 0){
        stop("Data and metadata features cannot be transfered in the same time!")
      } else {
        raw_counts <- vrData(from_object, norm = FALSE)
        raw_counts <- raw_counts[features,]
        message("There are ", length(setdiff(features, data_features)), " unknown features!")
      }
    } else {
      if(length(metadata_features) > 1){
        stop("Only one metadata column can be transfered at a time")
      } else if(length(metadata_features) == 1) {
        raw_counts <- from_metadata[,metadata_features, drop = FALSE]
        rownames_raw_counts <- rownames(raw_counts)
        raw_counts <- dummy_cols(raw_counts, remove_first_dummy = FALSE)
        raw_counts <- raw_counts[,-1]
        raw_counts <- t(raw_counts)
        colnames(raw_counts) <- rownames_raw_counts
        rownames(raw_counts) <- gsub(paste0("^", metadata_features, "_"), "", rownames(raw_counts))
      } else {
        stop("Features cannot be found in data and metadata!")
      }
    }
  }
  raw_counts <- raw_counts[,cell_to_spot_id, drop = FALSE]

  # pool cell counts to Spots
  message("Aggregating cell profiles in spots \n")
  aggregate_raw_counts <- stats::aggregate(t(as.matrix(raw_counts)), list(cell_to_spot_nnid), sum)
  aggregate_raw_counts <- data.frame(barcodes = vrSpatialPoints(to_object)) %>% dplyr::right_join(aggregate_raw_counts, by = c("barcodes" = "Group.1"))
  rownames(aggregate_raw_counts) <- aggregate_raw_counts$barcodes
  aggregate_raw_counts <- t(aggregate_raw_counts[,-1])
  aggregate_raw_counts[is.na(aggregate_raw_counts)] <- 0

  # return
  return(aggregate_raw_counts)
}

#' getSpotsFromCells
#'
#' Generate Psuedo counts per spots and insert to as a separate image assay to the Visium object
#'
#' @param from_object The vrAssay object whose data transfer to the second assay
#' @param from_metadata the metadata associated with \code{from_object}
#' @param to_object The ID of the target vrAssay object where data is transfered to
#' @param features the name of the metadata feature to transfer, if NULL, the rawdata will be transfered
#'
#' @importFrom dplyr %>% right_join
#' @importFrom stats aggregate
#' @importFrom magick image_data
#'
#' @noRd
#'
getCellsFromSpots <- function(from_object, from_metadata = NULL, to_object, features = NULL) {
  
  # get the spot radius of Visium spots
  radius <- vrAssayParams(from_object, param = "nearestpost.distance")/2
  
  # get cell and spot coordinates
  message("Spot to Cell Distances \n")
  coords_spots <- vrCoordinates(from_object)
  coords_cells <- vrCoordinates(to_object)
  
  # get distances from cells to spots
  spot_to_cell <- knn_annoy(coords_spots, coords_cells, k = 1)
  names(spot_to_cell) <- c("nn.index", "nn.dist")
  nnindex <- spot_to_cell$nn.index[,1]
  names_nnindex <- names(nnindex)
  nnindex <- vrSpatialPoints(from_object)[nnindex]
  names(nnindex) <- names_nnindex
  nndist <- spot_to_cell$nn.dist[,1]
  nnindex <- nnindex[nndist < radius]

  # find associated spot for each cell
  message("Find associated spot for each cell \n")

  # get data
  if(is.null(features)){
    raw_counts <- vrData(from_object, norm = FALSE)
  } else {
    data_features <- features[features %in% vrFeatures(from_object)]
    metadata_features <- features[features %in% colnames(from_metadata)]
    if(length(data_features) > 0){
      if(length(metadata_features) > 0){
        stop("Data and metadata features cannot be transfered in the same time!")
      } else {
        raw_counts <- vrData(from_object, norm = FALSE)
        raw_counts <- raw_counts[features,]
        message("There are ", length(setdiff(features, data_features)), " unknown features!")
      }
    } else {
      if(length(metadata_features) > 1){
        stop("Only one metadata column can be transfered at a time")
      } else if(length(metadata_features) == 1) {
        raw_counts <- from_metadata[,metadata_features, drop = FALSE]
        rownames_raw_counts <- rownames(raw_counts)
        raw_counts <- dummy_cols(raw_counts, remove_first_dummy = FALSE)
        raw_counts <- raw_counts[,-1]
        raw_counts <- t(raw_counts)
        colnames(raw_counts) <- rownames_raw_counts
        rownames(raw_counts) <- gsub(paste0("^", metadata_features, "_"), "", rownames(raw_counts))
      } else {
        stop("Features cannot be found in data and metadata!")
      }
    }
  }
  raw_counts <- raw_counts[,nnindex, drop = FALSE]
  colnames(raw_counts) <- names(nnindex)
  
  # return
  return(raw_counts)
}

#' getROIsFromCells
#'
#' Generate Psuedo counts per ROIs and insert to as a separate image assay to the Visium object
#'
#' @param from_object The vrAssay object whose data transfer to the second assay
#' @param from_metadata the metadata associated with \code{from_object}
#' @param to_object The ID of the target vrAssay object where data is transfered to
#' @param features the name of the metadata feature to transfer, if NULL, the rawdata will be transfered
#'
#' @importFrom dplyr %>% right_join
#' @importFrom stats aggregate
#' @importFrom magick image_data
#'
#' @noRd
#'
getROIsFromCells <- function(from_object, from_metadata = NULL, to_object, features = NULL) {

  # get cell and ROIs coordinates
  message("Cell to ROI Distances \n")
  segments_rois <- vrSegments(to_object)
  coords_cells <- vrCoordinates(from_object)
  
  # find associated spot for each cell
  message("Find associated ROIs for each cell \n")
  cell_to_roi_id <- NULL
  cell_to_roi_labelid <- NULL
  names_segments_rois <- names(segments_rois)
  for(i in seq_len(length(segments_rois))){
    cur_segt <- segments_rois[[i]]
    if(ncol(cur_segt) > 3){
      in.list <- point.in.circle(coords_cells[,1], coords_cells[,2], cur_segt$x, cur_segt$y, cur_segt$rx)
    } else {
      in.list <- sp::point.in.polygon(coords_cells[,1], coords_cells[,2], cur_segt$x, cur_segt$y)
    }
    in.list.cells <- rownames(coords_cells)[!!in.list]
    cell_to_roi_id <- c(cell_to_roi_id, in.list.cells)
    cell_to_roi_labelid <- c(cell_to_roi_labelid, rep(names_segments_rois[i], length(in.list.cells)))
  }
  
  # get data
  if(is.null(features)){
    raw_counts <- vrData(from_object, norm = FALSE)
  } else {
    data_features <- features[features %in% vrFeatures(from_object)]
    metadata_features <- features[features %in% colnames(from_metadata)]
    if(length(data_features) > 0){
      if(length(metadata_features) > 0){
        stop("Data and metadata features cannot be transfered in the same time!")
      } else {
        raw_counts <- vrData(from_object, norm = FALSE)
        raw_counts <- raw_counts[features,]
        message("There are ", length(setdiff(features, data_features)), " unknown features!")
      }
    } else {
      if(length(metadata_features) > 1){
        stop("Only one metadata column can be transfered at a time")
      } else if(length(metadata_features) == 1) {
        raw_counts <- from_metadata[,metadata_features, drop = FALSE]
        rownames_raw_counts <- rownames(raw_counts)
        raw_counts <- dummy_cols(raw_counts, remove_first_dummy = FALSE)
        raw_counts <- raw_counts[,-1]
        raw_counts <- t(raw_counts)
        colnames(raw_counts) <- rownames_raw_counts
        rownames(raw_counts) <- gsub(paste0("^", metadata_features, "_"), "", rownames(raw_counts))
      } else {
        stop("Features cannot be found in data and metadata!")
      }
    }
  }
  raw_counts <- raw_counts[,cell_to_roi_id, drop = FALSE]
  
  # pool cell counts to Spots
  message("Aggregating cell profiles in spots \n")
  aggregate_raw_counts <- stats::aggregate(t(as.matrix(raw_counts)), list(cell_to_roi_labelid), sum)
  aggregate_raw_counts <- data.frame(barcodes = vrSpatialPoints(to_object)) %>% dplyr::right_join(aggregate_raw_counts, by = c("barcodes" = "Group.1"))
  rownames(aggregate_raw_counts) <- aggregate_raw_counts$barcodes
  aggregate_raw_counts <- t(aggregate_raw_counts[,-1])
  aggregate_raw_counts[is.na(aggregate_raw_counts)] <- 0
  
  # return
  return(aggregate_raw_counts)
}

getCellsFromTiles <- function(from_object, from_metadata = NULL, to_object, features = NULL, k = 5) {

  # get cell and spot coordinates
  message("Tile to Cell Distances \n")
  coords_cells <- vrCoordinates(to_object)
  coords_tiles <- vrCoordinates(from_object)

  # get distances from cells to spots
  # tile_to_cell <- FNN::get.knnx(coords_tiles, coords_cells, k = k)
  tile_to_cell <- knn_annoy(coords_tiles, coords_cells, k = k)
  names(tile_to_cell) <- c("nn.index", "nn.dist")
  tile_to_cell_nnid <- data.frame(id = rownames(coords_cells), tile_to_cell$nn.index)
  # tile_to_cell_nnid <- reshape2::melt(tile_to_cell_nnid, id.vars = "id")
  tile_to_cell_nnid <- data.table::melt(tile_to_cell_nnid, id.vars = "id")
  tile_id <- vrSpatialPoints(from_object)[tile_to_cell_nnid$value]
  tile_to_cell_nnid <- tile_to_cell_nnid$id

  # get data
  raw_counts <- vrData(from_object, norm = FALSE)
  raw_counts <- raw_counts[,tile_id]

  # pool cell counts to Spots
  message("Aggregating tile profiles in cells \n")
  aggregate_raw_counts <- stats::aggregate(t(as.matrix(raw_counts)), list(tile_to_cell_nnid), mean)
  aggregate_raw_counts <- data.frame(barcodes = vrSpatialPoints(to_object)) %>% dplyr::right_join(aggregate_raw_counts, by = c("barcodes" = "Group.1"))
  rownames(aggregate_raw_counts) <- aggregate_raw_counts$barcodes
  aggregate_raw_counts <- t(aggregate_raw_counts[,-1])
  aggregate_raw_counts[is.na(aggregate_raw_counts)] <- 0

  # return
  return(aggregate_raw_counts)
}

transferLabelsFromROI <- function(from_object, from_metadata = NULL, to_object, to_metadata = NULL, features = NULL) {
  
  # get ROI and other coordinates
  segments_roi <- vrSegments(from_object)
  coords <- vrCoordinates(to_object)
  spatialpoints <- rownames(coords)
  
  # check if all features are in from_metadata
  if(!all(features %in% colnames(from_metadata))){
    stop("Some features are not found in the ROI metadata!")
  }
  
  # annotate points in the to object
  for(feat in features){
    
    # get from metadata labels
    feat_labels <- from_metadata[,feat]
    
    # get to metadata
    new_label <- rep("undefined", length(spatialpoints))
    names(new_label) <- spatialpoints
    
    for(i in seq_len(length(segments_roi))){
      cur_poly <- segments_roi[[i]][,c("x","y")]
      in.list <- sp::point.in.polygon(coords[,1], coords[,2], cur_poly[,1], cur_poly[,2])
      new_label[rownames(coords)[!!in.list]] <- feat_labels[i]
    }
    
    to_metadata[[feat]] <- new_label
  }

  
  # return label
  return(to_metadata)
}

####
# Embedding ####
####

getJointPCA <- function(object, assay.1 = NULL, assay.2 = NULL, dims = 30, seed = 1, ...){

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

  # if there are features of a VoltRon object, then get variable features too
  assay_features <- vrFeatures(object, assay = assay)
  if(length(assay_features) > 0) {
    features <- getVariableFeatures(object, assay = assay)
    vrMainAssay(object) <- object@sample.metadata[assay, "Assay"]
    object_subset <- subsetVoltRon(object, features = features)
    vrMainAssay(object_subset) <- vrMainAssay(object)
    if(dims > length(features)){
      message("Requested more PC dimensions than existing features: dims = length(features) now!")
      dims <- length(features)
    }
  } else {
    object_subset <- object
  }

  # set seed
  set.seed(seed)
  
  # get PCA embedding
  normdata.1 <- vrData(object_subset, assay = assay, norm = TRUE)
  scale.data <- apply(normdata.1, 1, scale)
  pr.data <- irlba::prcomp_irlba(scale.data, n=dims, center=colMeans(scale.data))
  pr.data.1 <- pr.data$x
  colnames(pr.data.1) <- paste0("PC", seq_len(dims))
  rownames(pr.data.1) <- colnames(normdata.1)
  normdata.1 <- normdata.1/(pr.data$sdev[1]^2)

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

  # if there are features of a VoltRon object, then get variable features too
  assay_features <- vrFeatures(object, assay = assay)
  if(length(assay_features) > 0) {
    features <- getVariableFeatures(object, assay = assay)
    vrMainAssay(object) <- object@sample.metadata[assay, "Assay"]
    object_subset <- subsetVoltRon(object, features = features)
    vrMainAssay(object_subset) <- vrMainAssay(object)
    if(dims > length(features)){
      message("Requested more PC dimensions than existing features: dims = length(features) now!")
      dims <- length(features)
    }
  } else {
    object_subset <- object
  }

  # get PCA embedding
  normdata.2 <- vrData(object_subset, assay = assay, norm = TRUE)
  scale.data <- apply(normdata.2, 1, scale)
  pr.data <- irlba::prcomp_irlba(scale.data, n=dims, center=colMeans(scale.data))
  pr.data.2 <- pr.data$x
  colnames(pr.data.2) <- paste0("PC", seq_len(dims))
  rownames(pr.data.2) <- colnames(normdata.2)
  normdata.2 <- normdata.2/(pr.data$sdev[1]^2)

  # get joint PCA
  normdata <- rbind(normdata.1, normdata.2)
  scale.data <- apply(normdata, 1, scale)
  pr.data <- irlba::prcomp_irlba(scale.data, n=dims, center=colMeans(scale.data))
  prsdev <- pr.data$sdev
  pr.data <- pr.data$x
  colnames(pr.data) <- paste0("PC", seq_len(dims))
  rownames(pr.data) <- colnames(normdata)
  normdata <- normdata/(prsdev[1]^2)

  # set Embeddings
  vrEmbeddings(object, type = "pca_joint", assay = assay.1, ...) <- pr.data
  vrEmbeddings(object, type = "pca_joint", assay = assay.2, ...) <- pr.data

  # return
  return(object)
}