plot_clustering_heatmap_wrapper2 <- function(expr, expr01,
  lineage_markers, functional_markers = NULL, sample_ids = NULL,
  cell_clustering, color_clusters, cluster_merging = NULL, 
  plot_cluster_annotation = TRUE){
  
  # Calculate the median expression of lineage markers
  expr_median <- data.frame(expr[, lineage_markers], 
    cell_clustering = cell_clustering) %>%
    group_by(cell_clustering) %>%  summarize_all(funs(median))
  expr01_median <- data.frame(expr01[, lineage_markers], 
    cell_clustering = cell_clustering) %>% 
    group_by(cell_clustering) %>%  summarize_all(funs(median))
  
  # Calculate cluster frequencies
  clustering_table <- as.numeric(table(cell_clustering))
  clustering_prop <- round(clustering_table / sum(clustering_table) * 100, 2)
  
  # Sort the cell clusters with hierarchical clustering
  d <- dist(expr_median[, lineage_markers], method = "euclidean")
  cluster_rows <- hclust(d, method = "average")
  
  expr_heat <- as.matrix(expr01_median[, lineage_markers])
  
  # Median expression of functional markers in each sample per cluster
  expr_median_sample_cluster_tbl <- data.frame(expr01[, functional_markers, 
    drop = FALSE], sample_id = sample_ids, cluster = cell_clustering) %>%
    group_by(sample_id, cluster) %>% summarize_all(funs(median))
  
  # Colors for the heatmap
  color_heat <- colorRampPalette(rev(brewer.pal(n = 9, name = "RdYlBu")))(100)
  legend_breaks = seq(from = 0, to = 1, by = 0.2)
  labels_row <- paste0(expr01_median$cell_clustering, " (", clustering_prop , 
    "%)")
  
  ### Annotation for the original clusters
  annotation_row1 <- data.frame(Cluster = factor(expr01_median$cell_clustering))
  color_clusters1 <- color_clusters[1:nlevels(annotation_row1$Cluster)]
  names(color_clusters1) <- levels(annotation_row1$Cluster)
  
  ### Annotation for the merged clusters
  if(!is.null(cluster_merging)){
    mm <- match(annotation_row1$Cluster, cluster_merging$original_cluster)
    annotation_row2 <- data.frame(Cluster_merging = 
        factor(cluster_merging$new_cluster[mm]))
    color_clusters2 <- color_clusters[1:nlevels(annotation_row2$Cluster_merging)]
    names(color_clusters2) <- levels(annotation_row2$Cluster_merging)
  }
  
  ### Heatmap annotation for the original clusters
  ha1 <- Heatmap(annotation_row1, name = "Cluster", 
    col = color_clusters1, cluster_columns = FALSE, 
    cluster_rows = cluster_rows, row_dend_reorder = FALSE, 
    show_row_names = FALSE, width = unit(0.5, "cm"), 
    rect_gp = gpar(col = "grey"))
  ### Heatmap annotation for the merged clusters
  if(!is.null(cluster_merging)){
    ha2 <- Heatmap(annotation_row2, name = "Cluster \nmerging", 
      col = color_clusters2, cluster_columns = FALSE, 
      cluster_rows = cluster_rows, row_dend_reorder = FALSE, 
      show_row_names = FALSE, width = unit(0.5, "cm"), 
      rect_gp = gpar(col = "grey"))
  }
  ### Cluster names and sizes - text
  ha_text <- rowAnnotation(text = row_anno_text(labels_row, 
    gp = gpar(fontsize = 6)), width = max_text_width(labels_row))
  ### Cluster sizes - barplot
  ha_bar <- rowAnnotation("Frequency (%)" = row_anno_barplot(
    x = clustering_prop, border = FALSE, axis = TRUE, 
    axis_gp = gpar(fontsize = 5), gp = gpar(fill = "#696969", col = "#696969"), 
    bar_width = 0.9), width = unit(0.7, "cm"), show_annotation_name = TRUE, 
    annotation_name_rot = 0, annotation_name_offset = unit(5, "mm"), 
    annotation_name_gp = gpar(fontsize = 7))
  ### Heatmap for the lineage markers
  ht1 <- Heatmap(expr_heat, name = "Expr",  column_title = "Lineage markers", 
    col = color_heat, cluster_columns = FALSE, cluster_rows = cluster_rows, 
    row_dend_reorder = FALSE, heatmap_legend_param = list(at = legend_breaks, 
      labels = legend_breaks, color_bar = "continuous"), 
    show_row_names = FALSE, row_dend_width = unit(2, "cm"), 
    rect_gp = gpar(col = "grey"), column_names_gp = gpar(fontsize = 8))
  
  if(plot_cluster_annotation){
    draw_out <- ha1
  }else{
    draw_out <- NULL  
  }
  if(!is.null(cluster_merging)){
    draw_out <- draw_out + ha2 + ht1 + ha_bar + ha_text
  }else{
    draw_out <- draw_out + ht1 + ha_bar + ha_text
  }
  
  ### Heatmaps for the signaling markers
  if(!is.null(functional_markers)){
    for(i in 1:length(functional_markers)){
      ## Rearange so the rows represent clusters
      expr_heat_fun <- as.matrix(dcast(expr_median_sample_cluster_tbl[, 
        c("sample_id", "cluster", functional_markers[i])], 
        cluster ~ sample_id, value.var = functional_markers[i])[, -1])
      
      draw_out <- draw_out + Heatmap(expr_heat_fun, 
        column_title = functional_markers[i], col = color_heat, 
        cluster_columns = FALSE, cluster_rows = cluster_rows, 
        row_dend_reorder = FALSE, show_heatmap_legend = FALSE, 
        show_row_names = FALSE, rect_gp = gpar(col = "grey"), 
        column_names_gp = gpar(fontsize = 8))
    }
  }
  draw(draw_out, row_dend_side = "left")
}