library(dplyr)
library(ggplot2)
library(ggpubr)
library(ggsci)
library(patchwork)
library(extrafont)
library(ComplexHeatmap)
library(circlize)
library(RColorBrewer)
library(viridis)
color <- readRDS("/raid1/shiq/resource/color/color_list.rds")


###################################
# calculate CM activity
################################

# node weight
node <- readRDS("/raid1/shiq/cellModule/cancer/05_dynamics/01_normal/norm_cm/04_network/node_each.rds")
weight <- readRDS("/raid1/shiq/cellModule/cancer/05_dynamics/01_normal/norm_cm/03_basis/weight_top.rds")
colnames(weight)[1] <- "node"
weight <- merge(node, weight)
rm(node)

# subset freq
meta_sp <- readRDS("/raid1/shiq/cellModule/cancer/04_cm_pair/02_coef/meta_sample_cCMT.rds")
mat_fq <- readRDS("/raid1/shiq/cellModule/cancer/05_dynamics/01_normal/04_cm_pair/00_prepare/mat_freq_global_norm.rds")

venn::venn(list(rownames(meta_sp), colnames(mat_fq)), ilabels = "counts")
mat_fq <- mat_fq[, rownames(meta_sp)]
sum(rownames(meta_sp) != colnames(mat_fq)) # 0

# cm score
meta_sp <- cbind(meta_sp, t(mat_fq))

for (cm in unique(weight$cm)) {
  node <- weight$node[weight$cm == cm]
  a <- meta_sp[, node] %>% as.matrix()
  b <- weight[weight$cm == cm, "weight", drop = FALSE] %>% as.matrix()
  meta_sp[, cm] <- c(a %*% b)
}



###################################
# Fig5d
###################################

#### prepare
meta_sp <- readRDS("/raid1/shiq/cellModule/cancer/05_dynamics/01_normal/meta_sample.rds")

# healthy CMs
pl_ls <- list(
  "Colon-Rectum_CRC" = c("CM03"),
  "Kidney_KIRC" = c("CM04"),
  "Liver_HCC" = c("CM11"),
  "Lung_LUAD" = c("CM11"),
  "Oral_HNSC" = c("CM08"),
  "Ovary_OV" = c("CM07"),
  "Skin_cSCC" = c("CM08"),
  "Uterus_CESC" = c("CM07")
)

# add two columns: health_cm and health_act
meta_sp$health_cm <- sapply(meta_sp$tissue_u, function(x) pl_ls[[x]])
table(meta_sp$tissue_u, meta_sp$health_cm)
meta_sp <- meta_sp %>%
  mutate(health_act = case_when(
    tissue_u == "Colon-Rectum_CRC" ~ CM03,
    tissue_u == "Kidney_KIRC" ~ CM04,
    tissue_u == "Liver_HCC" ~ CM11,
    tissue_u == "Lung_LUAD" ~ CM11,
    tissue_u == "Oral_HNSC" ~ CM08,
    tissue_u == "Ovary_OV" ~ CM07,
    tissue_u == "Skin_cSCC" ~ CM08,
    tissue_u == "Uterus_CESC" ~ CM07,
    TRUE ~ NA_real_
  ))

# add one column: health_act_scale
meta_sp <- meta_sp %>%
  group_by(tissue_u) %>%
  mutate(health_act_scale = scales::rescale(health_act)) %>%
  data.frame()

mat <- meta_sp %>%
  group_by(tissue_u, condition) %>%
  summarise_at(.vars = vars(health_act_scale), .funs = mean)
mat <- reshape2::dcast(data = mat, tissue_u ~ condition, value.var = "health_act_scale")
rownames(mat) <- mat$tissue_u
mat <- mat[, -1]

Heatmap(
  matrix = mat, border = TRUE,
  col = colorRamp2(seq(0.2, 0.7, length.out = 5), viridis(5)),
  width = unit(1.5, "cm"), height = unit(2.5, "cm"),
  cluster_columns = FALSE, cluster_rows = FALSE,
  # row_names_gp = gpar(fontsize = 6), column_names_gp = gpar(fontsize = 6)
  column_title = "Healthy CMs", column_title_gp = gpar(fontsize = 7),
  column_names_rot = 45, # top_annotation = ann_top,
  row_names_gp = gpar(fontsize = 6), column_names_gp = gpar(fontsize = 6),
  heatmap_legend_param = list(
    title = "CM activity", at = c(0, 0.4, 0.8),
    # legend_direction = "horizontal", title_position = "leftcenter",
    grid_width = unit(2, "mm"), legend_height = unit(1, "cm"),
    labels_gp = gpar(fontsize = 6), title_gp = gpar(fontsize = 6)
  )
) -> ht

pdf("../fig/Figure5d.pdf", width = 4, height = 4, family = "Arial")
draw(ht)
dev.off()




###################################
# Figure5e
################################
ggplot(meta_sp, aes(x = condition, y = health_act_scale, col = condition)) +
  geom_boxplot(outlier.shape = NA, linewidth = 0.3, alpha = 0) +
  geom_jitter(size = 0.3, width = 0.3, shape = 16) +
  facet_wrap(~tissue_u, scales = "free", nrow = 2) +
  scale_color_manual(values = color$condition) +
  stat_compare_means(
    # Use only p.format as label. Remove method name.
    aes(label = paste0("P = ", after_stat(p.format))),
    size = 1.6, # label = "p.signif",
    comparisons = list(
      c("Healthy", "Adjacent"),
      c("Adjacent", "Tumor"),
      c("Healthy", "Tumor")
    )
  ) +
  labs(x = "", y = "CM activity") +
  theme_classic() +
  theme(
    line = element_line(linewidth = 0.3),
    text = element_text(size = 6, family = "Arial", color = "black"),
    axis.text = element_text(size = 6, family = "Arial", color = "black"),
    axis.text.x = element_text(angle = 45, hjust = 1),
    plot.title = element_text(size = 7, hjust = 0.5),
    legend.position = "none",
    strip.text.x = element_text(size = 6),
    strip.background = element_blank(),
    aspect.ratio = 1 / 1
  ) -> p1

ggsave(filename = "../fig/Figure5e.pdf", plot = p1, width = 6, height = 3, family = "Arial")