#' Merge probability matrix based on annotations
#'
#' @param pm A numeric matrix. Probability matrix generated by the soft_max function.
#' @param group_df A character matrix. Annotation of the neighboring cells to be used.
#'
#' @return A probability matrix, describing the probability of each cell being
#' in each cellular neighborhood.
#' @export
#'
#' @examples
#'
#' data("spe_test")
#'
#' spe <- readHoodData(spe, anno_col = "celltypes")
#'
#' fnc <- findNearCells(spe, k = 100)
#'
#' pm <- scanHoods(fnc$distance)
#'
#' pm2 <- mergeByGroup(pm, fnc$cells)
#'
mergeByGroup <- function(pm, group_df) {
  if (!is(pm, "matrix")){
    stop("The input pm must be a numeric matrix.")
  }
  df <- as.matrix(pm)
  group_df <- as.matrix(group_df)

  if (!all(dim(df) == dim(group_df))) {
    stop("df and group_df should have the same dimensions.")
  }

  colnames(df) <- colnames(group_df)
  rownames(df) <- rownames(group_df)

  uniquegroups <- sort(unique(as.vector((group_df))))

  out_df <- matrix(nrow = nrow(pm), ncol = length(uniquegroups))
  colnames(out_df) <- uniquegroups
  rownames(out_df) <- rownames(df)

  out_df <- vapply(uniquegroups, \(g) {
    mask <- 1 * (group_df == g)
    rowSums(mask * pm)[rownames(df)]
  }, numeric(nrow(df)))

  out_df[is.na(out_df)] <- 0

  return(out_df)
}
