require(Seurat)
require(magrittr)
require(decoupleR)
require(reshape2)
require(tidyverse)
require(tidyr)
require(broom)
require(limma)
require(ggsci)
require(fossil)
require(ggvenn)
require(pryr)
require(stringr)
source("script/20240910.R")
net <- get_collectri(organism='human', split_complexes=FALSE)
obj <- readRDS("result/obj.pbmc3k.expression.rds")

################################################################################
#                                 NNet
################################################################################

tfs <- unique(net$source)
targets <- unique(net$target)
rt.ppr <- get.ppr()
genes <- select.gene(obj, tfs = tfs, targets = targets, min.cells = 0)
tfs <- genes$tfs
n.tfs <- length(tfs)

# Parameters
pars <- data.frame(truncated = FALSE, k = c(30,50,100), knn = 30, t = 3)
pars <- pars %>% rbind(
  data.frame(truncated = TRUE, k = 100, knn = c(15,50,100), t = 3)
)
pars <- pars %>% rbind(
  data.frame(truncated = TRUE, k = 100, knn = 30, t = c(1,5,10))
)

test.fun <- function(truncated, k, knn, t){
  obj <- prepare.seurat(obj, genes = genes$genes, k = k)
  obj <- prepare.graph(obj, truncated = truncated, knn = knn)
  obj <- prepare.reg(obj, predictors = tfs,
                     responses = tfs)
  obj <- run.nn.reg(obj, tfs, t=t)

  k <- ncol(obj@misc$setting$pcs)
  result <- list(
    degree.mat = apply(2*obj@misc$mod$effect^2, 1, colSums),
    self.rank.mat = sapply(1:n.tfs, function(x) apply(2*obj@misc$mod$effect[x,,]^2, 2, rank)[x,]/n.tfs)
  )
  saveRDS(result, paste0("response/","nPCs-",k,"-knn-",knn,"-t-",t,".rds"))
}

apply(pars, 1,
      function(x){
        do.call(test.fun,as.list(x))
      }
)

################################################################################
#                           Analysis (Self-ranking)
################################################################################
self.rank <- data.frame(matrix(nrow = n.tfs, ncol = 9))
files <- list.files("response", "nPCs")
file.names <- str_split(files, "\\.", simplify =TRUE)[,1]
colnames(self.rank) <- file.names
for(f in file.names){
  self.rank[,f] <- readRDS(paste0("response/", f, ".rds"))[[2]] %>% colMeans()
}
self.rank <- melt(self.rank, variable.name = "pars") %>%
  mutate(nPCs = str_split(pars, "-", simplify =TRUE )[,2],
         KNN = str_split(pars, "-", simplify =TRUE )[,4],
         t = str_split(pars, "-", simplify =TRUE )[,6]
  ) %>%
  mutate(nPCs = factor(nPCs, levels = c(30,50,57,100)),
         KNN = factor(KNN, levels = c(15,50,100)),
         t = factor(t, levels = c(1,5,10))
  ) %>%
  mutate(setting = case_when(nPCs != 57 ~ "Alternating PCs",
                             t != 3 ~ "Alternating t",
                             KNN != 30 ~ "Alternating KNN")
  )

# Plot self ranking
fill.col <- RColorBrewer::brewer.pal(11, "Spectral")[c(1,6,11)]

p1 <- subset(self.rank, setting == "Alternating PCs") %>%
  ggplot() +
  geom_violin(aes(nPCs, value, fill = nPCs)) +
  geom_jitter(aes(nPCs, value), alpha = 1, size = 0.8) +
  theme_classic() +
  theme(text = element_text(size = 13, face = "bold"),
        legend.position = "none") +
  scale_color_manual(values = fill.col)+
  labs(y = "Self-ranking percentile",
       title = "Alternating PCs")

p2 <- subset(self.rank, setting == "Alternating KNN") %>%
  ggplot() +
  geom_violin(aes(KNN, value, fill = KNN)) +
  geom_jitter(aes(KNN, value), alpha = 1, size = 0.8) +
  theme_classic() +
  theme(text = element_text(size = 13, face = "bold"),
        legend.position = "none") +
  scale_color_manual(values = fill.col)+
  labs(y = "",
       title = "Alternating KNN")

p3 <- subset(self.rank, setting == "Alternating t") %>%
  ggplot() +
  geom_violin(aes(t, value, fill = t)) +
  geom_jitter(aes(t, value), alpha = 1, size = 0.8) +
  theme_classic() +
  theme(text = element_text(size = 13, face = "bold"),
        legend.position = "none") +
  scale_color_manual(values = fill.col)+
  labs(y = "",
       title = "Alternating t")

p1 + p2 + p3


################################################################################
#                           Analysis (Clustering)
################################################################################
obj.list <- list()
files <- list.files("response", ".rds")
file.names <- str_split(files, "\\.", simplify =TRUE)[,1]
for(f in file.names){
  obj.f<- CreateSeuratObject(t(readRDS(paste0("response/", f, ".rds"))[[1]]))
  LayerData(obj.f, "data") <- obj.f@assays$RNA$counts
  VariableFeatures(obj.f) <- rownames(obj.f)
  obj.f <- ScaleData(obj.f)
  obj.f <- RunPCA(obj.f)
  obj.f <- FindNeighbors(obj.f, dims = 1:10)
  obj.f <- FindClusters(obj.f)
  obj.f <- RunUMAP(obj.f, dims = 1:10)
  obj.f$celltype <- Idents(obj)
  obj.list[[f]] <- obj.f
}

plot.list <- list()
for(f in file.names){
  sil <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                             dist(obj.list[[f]]@reductions$pca@cell.embeddings[,1:10]))[,3] %>%
    median %>% round(.,2)
  plot.list[[f]] <- DimPlot(obj.list[[f]], group.by = "celltype", pt.size = 0.2) +
    theme_classic()+
    theme(text = element_text(size = 13,face = "bold"),
          axis.line = element_blank(),
          legend.position = "none",
          axis.text = element_blank(),
          axis.ticks = element_blank(),
    )+ xlab("")+ylab("")+
    scale_color_bmj()+
    ggtitle(f, paste0("MSI=",sil))
}
file.names.ordered <- c("nPCs-30-knn-30-t-3", "nPCs-50-knn-30-t-3", "nPCs-100-knn-30-t-3",
                        "nPCs-57-knn-15-t-3", "nPCs-57-knn-50-t-3", "nPCs-57-knn-100-t-3",
                        "nPCs-57-knn-30-t-1", "nPCs-57-knn-30-t-5", "nPCs-57-knn-30-t-10")

ggarrange(plotlist = plot.list[file.names.ordered])
