spatialHeatmap: spatialHeatmap

#' Map colours in SHMs to embedding plots through manually-created matching list 
#'
#' @param gg.dim The ggplot of embedding plot.
#' @param gg.shm.all The list of SHM ggplots.
#' @param grob.shm.all The list of SHM grobs.
#' @param col.shm.all The list of SHM colours.
#' @param cell.group A column name in \code{colData} such as \code{cluster} (auto-generated), \code{label}. Cells are divided into clusters by this column name and these clusters are matched to bulk tissues. It is also the legend title in the embedding plots.
#' @param tar.cell The names of target cell clusters to show in embedding plot. The default is \code{matched} and only matching cell clusters have legends in the embedding plot.
#' @param con.na Logical, \code{TRUE} or \code{FALSE}. Default is \code{TRUE}, meaning conditions are available.
#' @param lis.match The maching list of spatial features between data and aSVGs.
#' @param sub.title.size The title size of embedding plots. The default is 11.
#' @param dim.lgd.pos The legend position. The default is \code{bottom}.
#' @param dim.lgd.nrow The number of legend rows. The default is \code{2}.
#' @param dim.lgd.text.size The size of legend text. The default is \code{8}.

#' @return A nested list of embedding and SHM plots.
#' @keywords Internal
#' @noRd

#' @author Jianhai Zhang \email{jzhan067@@ucr.edu} \cr Dr. Thomas Girke \email{thomas.girke@@ucr.edu}

#' @references
#' H. Wickham. ggplot2: Elegant Graphics for Data Analysis. Springer-Verlag New York, 2016

#' @importFrom ggplot2 layer_data ggplot geom_point theme_classic theme element_text element_blank labs scale_color_manual scale_shape_manual margin guide_legend element_rect 

dim_color2cell <- function(gg.dim, gg.shm.all, grob.shm.all, col.shm.all, gg.lgd.all, col.lgd.all, grob.lgd.all, profile=TRUE, cell.group, con.na=TRUE, lis.match=NULL, sub.title.size=11, dim.lgd.pos='bottom', dim.lgd.nrow=2, dim.lgd.key.size=4, dim.lgd.text.size=13, dim.axis.font.size=10, alpha.pt=0.8, shape=NULL) {
  # save(gg.dim, gg.shm.all, grob.shm.all, col.shm.all, gg.lgd.all, col.lgd.all, grob.lgd.all, profile, cell.group, con.na, lis.match, sub.title.size, dim.lgd.pos, dim.lgd.nrow, dim.lgd.key.size, dim.lgd.text.size, alpha.pt, shape, file='dim.color.tocell.arg')
  x <- y <- fill <- feature <- NULL
  # Trasfer svg features to name slots of lis.match.
  lis.match <- lis.match[!unlist(lapply(lis.match, is.null))]
  #lis.match <- lis.match
  #for (i in seq_along(lis.match)) {
  #  names(lis.match)[i] <- unique(subset(df.match, dataBulk==names(lis.match[i]))$SVGBulk)
  #}
  # Ggplots of all reduced dim.
  if (profile==TRUE) {
    n <- length(grob.shm.all); gg.dim.all <- rep(list(gg.dim), n)
    names(gg.dim.all) <- paste0('dim_', names(grob.shm.all))
  } else if (profile==FALSE) {
    n <- length(gg.lgd.all); gg.dim.all <- rep(list(gg.dim), n)
    names(gg.dim.all) <- paste0('dim_', names(gg.lgd.all))
  } 
  # Match colors in SHMs to dim plots. Colour order: data -> svg feature -> embedding plot.
  for (i in seq_along(gg.dim.all)) {
    gg.dim <- gg.dim.all[i]
    if (profile==TRUE) dim.col <- col_dim_tocell(gg.dim=gg.dim, gcol.all=col.shm.all, lis.match=lis.match)
    if (profile==FALSE) dim.col <- col_dim_tocell(gg.dim=gg.dim, gcol.all=col.lgd.all, lis.match=lis.match)

    # Data for plotting embedding plot.
    gg.dim0 <- gg.dim[[1]]
    dat <- gg.dim0$data; lay.dat <- layer_data(gg.dim0) 
    dat.all <- cbind(lay.dat, colour_by=dat$colour_by)
    dat.all$fill <- dat.all$colour <- dim.col[dat.all$colour_by]
    dat.all <- dat.all[, !colnames(dat.all) %in% 'colour']
    colnames(dat.all)[colnames(dat.all)=='colour_by'] <- 'feature'
    # dat.all$tissue <- dat.all$feature
    cell.all <- unique(dat.all$feature)
    # Cells not in lis.match.
    ft.o <- cell.all[!cell.all %in% names(dim.col)]
    col.o <- rep('gray80', length(ft.o))
    names(col.o) <- ft.o; dim.col.all <- c(dim.col, col.o)
    dat.all$fill <- dim.col.all[dat.all$feature]
    dat.all <- rbind(subset(dat.all, fill == 'gray80'), subset(dat.all, fill != 'gray80'))

    if (con.na==TRUE) tit <- gsub('^dim_(.*)_\\d+$', '\\1', names(gg.dim)) else tit <- gsub('^dim_(.*)_con_\\d+$', '\\1', names(gg.dim))
    if (profile==FALSE) tit <- NULL
    # Non-target cell clusters have colour of 'gray80'.
    tar.cell <- unlist(lis.match)
    non.tar <- setdiff(cell.all, tar.cell) 
    dim.col.all[non.tar] <- 'gray80'
    # Legal shapes: c(0:25, 32:127)
    # sp.sel <- c(15:18, 7:14)
    # sp.all <- c(0, 2:25, 32:127)
    # sp.all <- c(sp.sel, setdiff(sp.all, sp.sel))
    # Cell cluster shapes.
    # sp <- sp.all[seq_along(cell.all)]
    # names(sp) <- cell.all
    sp <- shp(shape, cell.all)
    if (length(non.tar) > 0) br <- tar.cell else br <- cell.all
    # Re-plot dimensionlaity plot.
    gg <- ggplot(dat.all, aes(x=x, y=y, text=dat.all$feature)) + geom_point(size=2, alpha=alpha.pt, aes(colour=feature, shape=feature)) + theme_classic() + theme(plot.title=element_text(hjust=0.5, size=sub.title.size), legend.position=dim.lgd.pos, legend.text=element_text(size=dim.lgd.text.size), legend.margin = margin(t=-0.02, l=0.05, r=0.1, unit='npc'), legend.background = element_rect(fill='transparent'), axis.text = element_blank(), axis.ticks = element_blank(), axis.title=element_text(size=dim.axis.font.size), aspect.ratio=1) + scale_color_manual(values=dim.col.all, breaks=br, guide=guide_legend(title=NULL, nrow=dim.lgd.nrow)) + scale_shape_manual(values=sp, breaks=br, guide=guide_legend(title=NULL, nrow=dim.lgd.nrow, override.aes = list(size=dim.lgd.key.size))) + labs(title=tit, x=gg.dim0$labels$x, y=gg.dim0$labels$y, colour=cell.group, shape=cell.group) 
    gg.dim.all[[i]] <- gg
  }
  # Convert all reduced dim of ggplots to grobs.
  grob.dim.all <- grob_shm(gg.dim.all, lgd.pos=NULL)
  # Empty list of all reduced dim and SHMs. 
  dim.shm.gg.lis <- dim.shm.grob.lis <- rep(list(NULL), 2*n)

  # Assign all reduced dims to the empty list.
  dim.shm.gg.lis[seq(1, 2*n, 2)] <- gg.dim.all
  dim.shm.grob.lis[seq(1, 2*n, 2)] <- grob.dim.all
  names(dim.shm.gg.lis)[seq(1, 2*n, 2)] <- names(dim.shm.grob.lis)[seq(1, 2*n, 2)] <- names(grob.dim.all)
  # Assign all SHMs to the empty list.
  if (profile==TRUE) {
    dim.shm.gg.lis[seq(2, 2*n, 2)] <- gg.shm.all
    dim.shm.grob.lis[seq(2, 2*n, 2)] <- grob.shm.all
    names(dim.shm.gg.lis)[seq(2, 2*n, 2)] <- names(dim.shm.grob.lis)[seq(2, 2*n, 2)] <- names(grob.shm.all)
  } else {
    dim.shm.gg.lis[seq(2, 2*n, 2)] <- gg.lgd.all
    dim.shm.grob.lis[seq(2, 2*n, 2)] <- grob.lgd.all
    names(dim.shm.gg.lis)[seq(2, 2*n, 2)] <- names(dim.shm.grob.lis)[seq(2, 2*n, 2)] <- names(grob.lgd.all)
  } 
  return(list(dim.shm.gg.lis=dim.shm.gg.lis, dim.shm.grob.lis=dim.shm.grob.lis))
}


#' Assign colors from SVG features to cell labels in embedding plot through a matching list when mapping bulk to cells.
#' 
#' @return A vector.
#' @keywords Internal
#' @noRd 

col_dim_tocell <- function(gg.dim, dim.na=NULL, gcol.all, lis.match) { 
    if (is.null(dim.na)) na <- sub('^dim_', '', names(gg.dim)) else na <- dim.na
    # Relevant colors in SHM. 
    g.col <- gcol.all[[paste0('col_', na)]]
    svg.ft.na <- names(lis.match)
    # 'gray80' is a reserved color. 
    cell.labs <- unlist(lis.match)    
    dim.col <- rep('gray80', length(cell.labs))
    names(dim.col) <- cell.labs 
    for (j in svg.ft.na) { # Colors: svg to cell.
      matched.col <- g.col[sub('__\\d+', '', names(g.col))==j][1]
      if (matched.col!='NA') dim.col[as.character(lis.match[[j]])] <- matched.col 
    }; return(dim.col)
}