####################################
# Fig. 2c (CM abundance prevalence across 34 tissues)
####################################

rm(list = ls())
library(dplyr)
library(ggplot2)
library(patchwork)
library(ggsci)
library(extrafont)
library(ComplexHeatmap)
library(RColorBrewer)
library(circlize)
# font_import()
# fonts()
color <- readRDS("../resource/color_list.rds")

# 1) load data
h <- readRDS("../prepare/CoVarNet/NMF_h.rds")
meta_sp <- readRDS("../prepare/meta_sample_CMT.rds")
range(colSums(h)) # 1 1

# 2) tissue-level abundance (h_ts)
h <- h[, rownames(meta_sp)]
sum(colnames(h) != rownames(meta_sp)) # 0
meta_sp <- cbind(meta_sp, t(h))
h_ts <- meta_sp %>%
  dplyr::group_by(tissue) %>%
  summarise_at(sprintf("CM%02d", 1:12), mean) %>%
  data.frame(row.names = 1) %>%
  t()
range(colSums(h_ts)) # 1 1

# 3) roe
colSums(h_ts)
rowSums(h_ts)
roe <- h_ts / rowSums(h_ts) * sum(h_ts)


# 4) plot (CM abundance prevalence)
# ann
ann <- meta_sp[!duplicated(meta_sp$tissue), c("system", "tissue")]
rownames(ann) <- ann$tissue
ann <- ann[colnames(roe), ]
sum(colnames(roe) != rownames(ann)) # 0

ann_top <- HeatmapAnnotation(
  System = ann$system,
  col = list(System = color$system),
  simple_anno_size = unit(3, "mm"),
  border = TRUE, annotation_name_gp = gpar(fontsize = 6),
  annotation_legend_param = list(
    labels_gp = gpar(fontsize = 6),
    title_gp = gpar(fontsize = 6),
    grid_width = unit(3, "mm"),
    grid_height = unit(2, "mm")
  )
)

# roe
Heatmap(
  matrix = roe,
  # col = viridis_pal(begin = 0, end = 0.5)(7),
  col = colorRamp2(seq(0, 3, length.out = 4), brewer.pal(n = 7, name = "RdBu")[4:1]),
  border = TRUE, # name = "Ro/e",
  clustering_method_rows = "ward.D", clustering_method_columns = "ward.D",
  width = unit(10, "cm"), height = unit(3.5, "cm"),
  column_title = "CM abundance prevalence across 34 tissues",
  column_title_gp = gpar(fontsize = 8),
  column_names_rot = 90, top_annotation = ann_top,
  row_names_gp = gpar(fontsize = 6), column_names_gp = gpar(fontsize = 6),
  column_dend_height = unit(5, "mm"), row_dend_width = unit(5, "mm"),
  cell_fun = function(j, i, x, y, width, height, fill) {
    if (roe[i, j] > 3) {
      grid.text(print("+++"), x, y, gp = gpar(fontsize = 5))
    } else if (roe[i, j] > 2) {
      grid.text(print("++"), x, y, gp = gpar(fontsize = 5))
    } else if (roe[i, j] > 1) {
      grid.text(print("+"), x, y, gp = gpar(fontsize = 5))
    }
  },
  heatmap_legend_param = list(
    at = seq(0, 3.5, 0.5),
    labels = c("0", "0.5", "1", "(+)", "2", "(++)", "3", "(+++)"),
    title = "Ro/e",
    # legend_direction = "horizontal",
    grid_width = unit(3, "mm"), # legend_height = unit(1, "cm"),
    labels_gp = gpar(fontsize = 6), title_gp = gpar(fontsize = 6)
    # title_position = "topcenter"
  )
) -> p2

pdf("Figure2c.pdf", width = 9, height = 4, family = "Arial")
draw(p2,
     heatmap_legend_side = "right", annotation_legend_side = "right",
     merge_legends = FALSE # adjust_annotation_extension = TRUE,
)
dev.off()