#####################################
# Fig. 4a
#####################################

library(dplyr)
library(tidytext)
library(ggplot2)
library(patchwork)
library(ggpubr)
library(extrafont)
color <- readRDS("../resource/color_list.rds")

# 1) prepare phenotype meta
# load meta
meta <- readRDS("../prepare/meta_sample_CMT.rds")

# age group
table(meta$age)
meta$age_gp <- meta$age
meta$age_gp[meta$age %in% c("18-30", "18-34", "20-25", "25-30", 17:34)] <- "<35"
meta$age_gp[meta$age %in% c(35:39)] <- "35-39"
meta$age_gp[meta$age %in% c("40-45", "45-50", "42.5", 40:49)] <- "40-49"
meta$age_gp[meta$age %in% c("50-54", "55-60", "57.5", 50:59)] <- "50-59"
meta$age_gp[meta$age %in% c("60-64", "60-65", "65-70", "67.5", 60:69)] <- "60-69"
meta$age_gp[meta$age %in% c("70-74", "70-75", 70:85)] <- "70-85"
meta$age <- as.numeric(meta$age)
table(meta$age_gp)
table(meta$age, meta$age_gp)
table(meta$tissue, meta$age_gp)

# merge
h <- readRDS("../prepare/CoVarNet/NMF_h.rds")
freq <- readRDS("../prepare/mat_freq_raw.rds")

venn::venn(
  x = list(
    meta = rownames(meta),
    h = colnames(h),
    freq = colnames(freq)
  ),
  ilabels = "counts"
)

# unify sample order
sum(rownames(meta) != colnames(h))
sum(rownames(meta) != colnames(freq))
h <- h[, rownames(meta)]
freq <- freq[, rownames(meta)]
# merge
meta <- cbind(meta, t(h))
meta <- cbind(meta, t(freq))
dim(meta) # 510 104
head(meta[, 1:28])
table(meta$sex)

# only tissue with >=5 samples
ts <- table(meta$tissue)
ts <- names(ts)[ts >= 5] # 21 tissues
meta <- meta[meta$tissue %in% ts, ]
dim(meta) # 486 104

# add majorCluster freq
meta_cell <- read.csv2("../data/meta_obs_ann_s2_trim.csv2", row.names = 1)
nrow(meta_cell) # 2293951
meta_cell <- meta_cell[meta_cell$cellSort == "Total", ]
meta_cell <- meta_cell[meta_cell$sampleID %in% meta$sampleID, ]
# only 413 samples without sorting have majorCluster freq
venn::venn(list(meta$sampleID, meta_cell$sampleID), ilabels = "counts")

# create table
mat <- table(meta_cell$sampleID, meta_cell$majorCluster)
mat <- mat / rowSums(mat)

tmp <- matrix(
  data = NA,
  nrow = nrow(meta), ncol = ncol(mat),
  dimnames = list(rownames(meta), colnames(mat))
)
tmp[rownames(mat), colnames(mat)] <- mat

if (all(rownames(meta) == rownames(tmp))) {
  meta <- cbind(meta, tmp)
}

colnames(meta)
meta <- meta[meta$tissue == "Spleen", ]
meta <- meta[!meta$age_gp %in% c("Unknown"), ]
meta[, c(1:16)]


# 2) plot
node <- read.csv("../prepare/CoVarNet/CoVarNet_node_each.csv",row.names = 1)

# CM05
ggplot(meta, aes(age_gp, CM05, col = age_gp)) +
  geom_boxplot(outlier.shape = NA, linewidth = 0.3, alpha = 0) +
  geom_jitter(shape = 16, size = 0.5, width = 0.3) +
  scale_color_viridis_d(option = "cividis") +
  stat_compare_means(
    size = 2, method = "kruskal.test",
    label.x.npc = 0.1, label.y = 1.3
  ) +
  stat_compare_means(
    aes(label = paste0("P = ", after_stat(p.format))),
    comparisons = list( # only label signif. pair
      # c("<35", "40-49"),
      c("40-49", "50-59")
      # c("50-59", "60-69"),
      # c("60-69", "70-85")
    ), size = 2
  ) +
  scale_y_continuous(expand = expansion(mult = c(0.05, 0.15))) +
  labs(x = "", y = "CM05") +
  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_blank(),
    axis.text.x = element_text(angle = 45, hjust = 1),
    plot.title = element_text(size = 6, face = "italic", hjust = 0.5),
    legend.position = "none",
    aspect.ratio = 1.2 / 2
  ) -> p1

# CM06
ggplot(meta, aes(age_gp, CM06, col = age_gp)) +
  geom_boxplot(outlier.shape = NA, linewidth = 0.3, alpha = 0) +
  geom_jitter(shape = 16, size = 0.5, width = 0.3) +
  scale_color_viridis_d(option = "cividis") +
  stat_compare_means(
    size = 2, method = "kruskal.test",
    label.x.npc = 0.3, label.y = 1.5
  ) +
  stat_compare_means(
    aes(label = paste0("P = ", after_stat(p.format))),
    comparisons = list(
      # c("60-69", "70-85"),
      # c("50-59", "60-69"),
      c("40-49", "50-59"),
      c("<35", "40-49")
    ), size = 2
  ) +
  scale_y_continuous(expand = expansion(mult = c(0.05, 0.15))) +
  labs(x = "", y = "CM06") +
  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 = 6, face = "italic", hjust = 0.5),
    legend.position = "none",
    aspect.ratio = 1.2 / 2
  ) -> p2
plt <- p1 + p2 + plot_layout(ncol = 2)
ggsave("./Figure4a.pdf", plot = plt, width = 7, height = 1.7, family = "Arial")




#####################################
# Fig. 4b
#####################################

dim(meta)

# 1) significant subsets
#  p value
meta_sc <- readRDS("../prepare/meta_subCluster.rds")
meta_sc <- meta_sc[meta_sc$majorCluster %in% c("B", "CD4T", "CD8T", "ILC", "Myeloid"), ]
sc_ls <- meta_sc$subCluster
temp <- setNames(object = rep(NA, length(sc_ls)), nm = sc_ls)

for (sc in sc_ls) {
  meta$sc <- meta[, sc]
  if (length(unique(meta$age_gp)) == 2) { # two groups
    tmp <- wilcox.test(sc ~ age_gp, data = meta)$p.value
  } else if (length(unique(meta$age_gp)) > 2) { # more groups
    tmp <- kruskal.test(
      sc ~ age_gp,
      data = meta
    )$p.value
  }
  temp[sc] <- tmp
}

# check NA
sum(is.na(temp)) # 2
temp <- temp[!is.na(temp)]
# fdr
mat2 <- p.adjust(temp, method = "fdr")
# -log10(fdr)
mat3 <- -log10(mat2)

# 2) plot
df <- data.frame(sc = names(mat3), sig = unname(mat3))

# heatmap, mean freq of subsets in each age group
sc_sig <- df$sc[df$sig > (-1) * log10(0.05)]
meta2 <- meta[, c("age_gp", "B08_ABC_FCRL5", sc_sig)]

mean_fq <- aggregate(. ~ age_gp, data = meta2, FUN = mean)
rownames(mean_fq) <- mean_fq[, 1]
mean_fq <- mean_fq[, -1]
mean_fq <- t(mean_fq)

Heatmap(
  matrix = mean_fq,
  # col = viridis_pal(begin = 0, end = 0.5)(7),
  col = colorRamp2(seq(0, 0.7, length.out = 4), brewer.pal(n = 7, name = "RdBu")[4:1]),
  border = TRUE, row_km = 2, # name = "Ro/e",
  clustering_method_rows = "ward.D", cluster_columns = FALSE,
  width = unit(2, "cm"), height = unit(4.5, "cm"),
  column_title = "Chronological dynamics",
  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),
  column_dend_height = unit(5, "mm"), row_dend_width = unit(5, "mm"),
  heatmap_legend_param = list(
    title = "Cell freq.", at = c(0, 0.4, 0.8),
    legend_direction = "horizontal", title_position = "leftcenter",
    grid_height = unit(3, "mm"), legend_width = unit(1, "cm"),
    labels_gp = gpar(fontsize = 6), title_gp = gpar(fontsize = 6)
  )
) -> p2

pdf("Figure4b.pdf", width = 4, height = 4, family = "Arial")
draw(p2, heatmap_legend_side = "bottom")
dev.off()