####################################
# Fig. 2b (Landscape of cross-tissue CMs in the human body)
####################################

# 1) prepare and filter data
# 2) freq mat (raw)
# 3) freq mat (normalize to [0,1] globally)
# 4) nmf
# 5) sample CMT annotation
# 6) subsets weights for each CM
# 7) identify specifically correlated subset pairs
# 8) plot network


rm(list = ls())

library(NMF)
library(cluster)
library(dplyr)
library(ggplot2)
library(tidytext)
library(patchwork)
library(ComplexHeatmap)
library(RColorBrewer)
library(circlize)
library(ggsci)
library(viridis)
library(extrafont)
library(igraph)
library(tidyverse)
library(psych)
library(mgsub)

color <- readRDS("../resource/color_list.rds")

# 1) prepare and filter data
meta <- read.csv2("../data/meta_obs_ann_s2_trim.csv2", row.names = 1)
meta <- meta[meta$ann1 != "Epi", ]
nrow(meta) # 1702037
# only Mix and Total
meta <- meta[meta$cellSort %in% c("Mix", "Total"), ]
nrow(meta) # 1448609

sort(table(meta$sampleID), decreasing = TRUE) %>% length() # 545
# remove samples < 100 cells
rt_sp <- names(table(meta$sampleID))[table(meta$sampleID) >= 100]
meta <- meta[meta$sampleID %in% rt_sp, ]
nrow(meta) # 1446646
sort(table(meta$sampleID), decreasing = TRUE) %>% length() # 510
rm(rt_sp)



# 2) freq mat (raw)
mat_ct <- table(meta$subCluster, meta$sampleID)
mat_ct <- matrix(
  data = mat_ct,
  ncol = ncol(mat_ct),
  dimnames = dimnames(mat_ct)
) %>% as.data.frame()
mat_ct[1:3, 1:5]
dim(mat_ct) # 76 510

# add majorCluster
meta_sc <- readRDS("../prepare/meta_subCluster.rds")
id <- match(rownames(mat_ct), meta_sc$subCluster)
mat_ct <- data.frame(
  majorCluster = factor(meta_sc$majorCluster[id]),
  mat_ct,
  check.names = FALSE
)
mat_ct[1:3, 1:5]
dim(mat_ct) # 76 511 # 510 samples

# frequency matrix
mat_fq <- mat_ct %>%
  group_by(majorCluster) %>%
  mutate_if(is.numeric, .funs = list(~ . / sum(.))) %>%
  as.data.frame()
mat_fq[is.na(mat_fq)] <- 0
rownames(mat_fq) <- rownames(mat_ct)

# remove majorCluster
mat_fq <- mat_fq[, -1]
mat_fq[1:3, 1:5]
dim(mat_fq) # 76 510
colSums(mat_fq) %>% table()

saveRDS(mat_fq, "../prepare/mat_freq_raw.rds")

# 3) freq mat (normalize to [0,1] globally)
mat_gb_norm <- mat_fq
mat_gb_norm <- mat_gb_norm - apply(mat_gb_norm, MARGIN = 1, FUN = min)
mat_gb_norm <- mat_gb_norm / apply(mat_gb_norm, MARGIN = 1, FUN = max)
sum(is.na(mat_gb_norm)) # 0
dim(mat_gb_norm) # 76 510
range(mat_gb_norm)
colSums(mat_gb_norm) %>% range()

# 4) nmf
mat_gb_norm[1:5, 1:5]
sum(is.na(mat_gb_norm)) # 0
dim(mat_gb_norm) # 76 510
range(mat_gb_norm)

K <- 12
res <- nmf(mat_gb_norm, rank = K, method = "nsNMF", seed = rep(77, 6), nrun = 30)
saveRDS(res, "../prepare/CoVarNet/NMF_K12.rds")

sum(rownames(res) != rownames(mat_gb_norm)) # 0
sum(colnames(res) != colnames(mat_gb_norm)) # 0

pdf("./consensus_K12.pdf", width = 8, height = 7)
consensusmap(res,
             labCol = NA, labRow = NA,
             annColors = list(
               basis = pal_simpsons()(15),
               consensus = pal_simpsons()(15)
             )
)
dev.off()

png("./consensus_K12.png", width = 800, height = 700)
consensusmap(res,
             labCol = NA, labRow = NA,
             annColors = list(
               basis = pal_simpsons()(15),
               consensus = pal_simpsons()(15)
             )
)
dev.off()

# 5) nmf coef
dim(res) # 76 510  12
meta_sp <- readRDS("../prepare/meta_sample.rds")
meta_sp <- meta_sp[meta_sp$sampleID %in% colnames(res), ]
head(meta_sp)

h <- scoef(res)
rownames(h) <- sprintf("C%02d", 1:nrow(h))
sum(is.na(h)) # 0
sum(colSums(h) != 1) # 74
sum(round(colSums(h)) != 1) # 0

# add CMT
sum(rownames(meta_sp) != colnames(h))
venn::venn(list(A = rownames(meta_sp), B = colnames(h)))
h <- h[, rownames(meta_sp)]
id <- apply(h, MARGIN = 2, FUN = which.max)
meta_sp$CMT <- rownames(h)[id]
head(meta_sp)

# shannon
mat_tissue_CMT <- table(meta_sp$tissue, meta_sp$CMT)
mat_tissue_CMT <- t(t(mat_tissue_CMT) / colSums(mat_tissue_CMT))
colSums(mat_tissue_CMT)
shannon <- apply(mat_tissue_CMT,
                 MARGIN = 2,
                 FUN = function(x) {
                   x <- x[x != 0]
                   sum(x * log2(x)) * (-1)
                 }
)
sort(shannon, decreasing = TRUE)

# rename
name_map <- data.frame(
  int = names(sort(shannon, decreasing = TRUE)),
  entropy = sort(shannon, decreasing = TRUE) %>% unname(),
  CM = sprintf("CM%02d", 1:length(shannon)),
  CMT = sprintf("CMT%02d", 1:length(shannon))
)
# meta_sp
id <- match(meta_sp$CMT, name_map$int)
meta_sp$CMT <- name_map$CMT[id]
meta_sp <- meta_sp[
  order(meta_sp$CMT, meta_sp$system, meta_sp$tissue),
]
# h
id <- match(rownames(h), name_map$int)
rownames(h) <- name_map$CM[id]
h <- h[sort(rownames(h)), ]

saveRDS(h, "../prepare/CoVarNet/NMF_h.rds")
saveRDS(meta_sp, "../prepare/meta_sample_CMT.rds")


# 6) subsets weights for each CM
# basis
w <- basis(res)
colnames(w) <- sprintf("C%02d", 1:ncol(w))
id <- match(colnames(w), name_map$int)
colnames(w) <- name_map$CM[id]
w <- w[, sort(colnames(w))]

sum(is.na(w)) # 0
sum(colSums(w) != 1) # 0  ??
sum(round(colSums(w)) != 1) # 0

w_df <- reshape2::melt(w)
head(w_df)
colnames(w_df) <- c("subCluster", "cm", "weight")

w_df <- w_df %>%
  group_by(cm) %>%
  arrange(desc(weight), .by_group = TRUE) %>%
  as.data.frame()
write.csv(w_df, "../prepare/CoVarNet/NMF_weight_all.csv", quote = FALSE)

w_top <- w_df %>%
  group_by(cm) %>%
  arrange(desc(weight), .by_group = TRUE) %>%
  top_n(15, weight) %>%
  as.data.frame()
# saveRDS(w_top, "./weight_top.rds")


# 7) identify specifically correlated subset pairs
# correlation
cor_ls <- t(mat_fq) %>% psych::corr.test(method = "pearson")
cor_mat <- cor_ls$r
cor_mat[lower.tri(cor_mat, diag = TRUE)] <- NA
cor_df <- reshape2::melt(cor_mat, na.rm = TRUE)
colnames(cor_df) <- c("node1", "node2", "correlation")

# pval
pval_mat <- cor_ls$p
pval_mat[lower.tri(pval_mat, diag = TRUE)] <- NA
pval_df <- reshape2::melt(pval_mat, na.rm = TRUE)
colnames(pval_df) <- c("node1", "node2", "pval")

# fdr
pval_df$pval_fdr <- p.adjust(pval_df$pval, method = "fdr")

# spe
cor_mat <- cor_ls$r
spe_mat <- matrix(NA, nrow(cor_mat), ncol(cor_mat))
colnames(spe_mat) <- rownames(spe_mat) <- rownames(cor_mat)
N <- nrow(spe_mat)
for (i in 1:(N - 1)) {
  for (j in (i + 1):N) {
    bg <- c(cor_mat[i, -i], cor_mat[-c(i, j), j]) %>% unname()
    # spe
    spe_mat[i, j] <- mean(bg <= cor_mat[i, j])
  }
}
spe_df <- reshape2::melt(spe_mat, na.rm = TRUE)
colnames(spe_df) <- c("node1", "node2", "spe")

# merge
cor_df <- Reduce(f = merge, x = list(cor_df, pval_df, spe_df))
table(cor_df$spe + cor_df$pmt)

# majorCluster
cor_df$major1 <- meta_sc$majorCluster[match(cor_df$node1, meta_sc$subCluster)]
cor_df$major2 <- meta_sc$majorCluster[match(cor_df$node2, meta_sc$subCluster)]
head(cor_df)

########## saveRDS(cor_df, sprintf("cor_df_%s.rds", method))


# 8) plot network
# top 10 subsets of each CM
w_top10 <- w_df %>%
  group_by(cm) %>%
  arrange(desc(weight), .by_group = TRUE) %>%
  slice(1:10) %>%
  ungroup()

## network
# filter 1
cdt1 <- cor_df$correlation > 0.2
cdt2 <- cor_df$pval_fdr <= 0.05
pl_df <- cor_df[cdt1 & cdt2, ]
dim(pl_df) # 213   8

# filter 2
n_in_cm <- 10 # node number in each CM
n_all <- nrow(meta_sc) # subset number
spe_cutoff <- 1 - ((n_in_cm - 1) * 2 - 1) / ((n_all - 1) * 2 - 1)
cdt3 <- pl_df$spe >= spe_cutoff
pl_df <- pl_df[cdt3, ]
dim(pl_df) # 197  8

quantile(pl_df$correlation)
graph <- graph_from_data_frame(pl_df, directed = FALSE)

# vertex attr
V(graph)$majorCluster <- meta_sc[V(graph)$name, "majorCluster"]
V(graph)$frame.color <- V(graph)$color <- color$majorCluster[V(graph)$majorCluster]
tmp <- strsplit(V(graph)$name, split = "_") %>% unlist()
V(graph)$short_name <- tmp[seq(1, length(V(graph)$name) * 3, 3)]
V(graph)$short_name <- gsub(pattern = "^CD", replacement = "", x = V(graph)$short_name)

# edge attr
quantile(-log10(E(graph)$pval_fdr), probs = seq(0, 1, 0.1))
quantile(E(graph)$spe)
col_fun <- colorRamp2(
  # breaks = quantile(-log10(E(graph)$pval_fdr), probs = c(0, 0.7)),
  breaks = quantile(E(graph)$spe, probs = seq(0, 1, length = 6)),
  # colors = brewer.pal(9, "Reds")[2:9],
  colors = pal_material(palette = "grey", n = 10)(10)[3:8]
)
# E(graph)$color <- col_fun(-log10(E(graph)$pval_fdr))
E(graph)$color <- col_fun(E(graph)$spe)
E(graph)$width <- 1

vertex_attr(graph)
edge_attr(graph)

## global plot
# node and edge
node <- matrix(NA, nrow = 0, ncol = 2) %>% as.data.frame()
colnames(node) <- c("cm", "node")
node <- V(graph)$name
node <- data.frame(
  node = node,
  major <- meta_sc$majorCluster[match(node, meta_sc$subCluster)]
)
node <- node[order(node$node), ]

# legend
lgd <- Legend(
  col_fun = col_fun,
  at = c(min(E(graph)$spe), max(E(graph)$spe)),
  labels = c("Low", "High"),
  title = expression("Specificity"),
  direction = "horizontal",
  grid_height = unit(3, "mm"),
  legend_width = unit(10, "mm"),
  labels_gp = gpar(fontsize = 6),
  title_gp = gpar(fontsize = 6),
  title_position = c("topcenter")
)

# global full name fr
pdf("global_full.pdf", width = 10, height = 10)
par(mar = c(0, 0, 0, 0))
set.seed(42)
plot.igraph(
  graph,
  layout = layout_with_fr,
  xlim = c(-1.2, 1.2),
  ylim = c(-1.2, 1.2),
  # vertex
  vertex.size = 3,
  # vertex.label = V(graph)$short_name,
  vertex.label.cex = 6 / 12,
  vertex.label.color = "black",
  vertex.label.family = "Arial"
)
draw(lgd, x = unit(0.1, "npc"), y = unit(0.05, "npc"))
dev.off()

# global short name fr
pdf("global_short.pdf", width = 8, height = 8)
par(mar = c(0, 0, 0, 0))
set.seed(42)
plot.igraph(
  graph,
  layout = layout_with_fr,
  # vertex
  vertex.size = 6,
  vertex.label = V(graph)$short_name,
  vertex.label.cex = 6 / 12,
  vertex.label.color = "black",
  vertex.label.family = "Arial"
)
draw(lgd, x = unit(0.1, "npc"), y = unit(0.05, "npc"))
dev.off()


### each cm
# node of each cm
node <- matrix(NA, nrow = 0, ncol = 2) %>% as.data.frame()
colnames(node) <- c("cm", "node")
edge <- matrix(NA, nrow = 0, ncol = ncol(pl_df) + 1) %>% as.data.frame()
colnames(edge) <- c("cm", colnames(pl_df))
index <- matrix(NA, nrow = 0, ncol = 3) %>% as.data.frame()
colnames(index) <- c("cm", "mean_distance", "global_efficiency")
for (cm in unique(w_top10$cm)) {
  sub_graph <- subgraph(
    graph,
    vids = intersect(V(graph)$name, w_top10$subCluster[w_top10$cm == cm])
  )
  sub_graph <- delete_vertices(sub_graph, V(sub_graph)[igraph::degree(sub_graph) == 0])
  # stat
  tmp1 <- data.frame(cm = cm, node = V(sub_graph)$name)
  node <- rbind(node, tmp1)
  tmp2 <- pl_df[(pl_df$node1 %in% tmp1$node) & (pl_df$node2 %in% tmp1$node), ]
  edge <- rbind(edge, data.frame(cm = cm, tmp2))
  tmp3 <- data.frame(
    cm = cm,
    mean_distance = mean_distance(sub_graph, directed = FALSE),
    global_efficiency = global_efficiency(sub_graph, directed = FALSE)
  )
  index <- rbind(index, tmp3)
}
node$major <- meta_sc$majorCluster[match(node$node, meta_sc$subCluster)]
rownames(edge) <- NULL
write.csv(node, "../prepare/CoVarNet/CoVarNet_node_each.csv", quote = FALSE)

# each short name circle
pdf("Figure2b.pdf", width = 6, height = 2) # width = 4, height = 3.5
par(mfrow = c(2, 6), mar = c(0, 0, 0, 0) + 0.5) # 3，4
for (cm in unique(w_top10$cm)) {
  sub_graph <- subgraph(
    graph,
    vids = intersect(V(graph)$name, w_top10$subCluster[w_top10$cm == cm])
  )
  sub_graph <- delete_vertices(sub_graph, V(sub_graph)[igraph::degree(sub_graph) == 0])
  plot.igraph(
    sub_graph,
    layout = layout_in_circle, # layout_with_fr
    xlim = c(-1.2, 1.2),
    ylim = c(-1.2, 1.2),
    # vertex
    vertex.size = 50,
    vertex.label = V(sub_graph)$short_name,
    vertex.label.cex = 5 / 8,
    vertex.label.color = "black",
    vertex.label.family = "Arial"
  )
  title(cm, cex.main = 7 / 8, line = -0.5, family = "Arial")
}
dev.off()