require(Seurat)
require(SeuratDisk)
require(SeuratData)
require(magrittr)
require(decoupleR)
require(reshape2)
require(tidyverse)
require(tidyr)
require(broom)
require(limma)
require(ggsci)
require(fossil)
require(ggvenn)
require(pryr)
require(qs)
require(anndataR)
require(patchwork)
source("script/20240910.R")
#remotes::install_github("https://github.com/timoast/seurat-disk/tree/master")
net <- get_collectri(organism='human', split_complexes=FALSE)

# http://bioailab.com/PerturbSeq.db/
# Characterizing the molecular regulation of inhibitory immune checkpoints with multimodal single-cell screens
# Papalexi etal. 2021
obj <- LoadH5Seurat("response/data/filtered_003.h5Seurat")
obj$gene <- droplevels(obj$gene)

perturbed <- paste(levels(obj$gene), collapse = ",") %>%
  str_split(., ",") %>%
  do.call(c,.) %>%
  unique()
# obj <- sapply(perturbed,
#                function(x) {
#                  grepl(x, obj$gene)
#                }
#              ) %>%
#        AddMetaData(obj, .)

# Run regression
tfs <- unique(net$source)
targets <- unique(net$target)
rt.ppr <- get.ppr()
genes <- select.gene(obj, tfs = tfs, targets = targets, min.cells = 0)
obj <- prepare.seurat(obj, genes = genes$genes)

mem_change({
  obj <- prepare.graph(obj)
  obj <- select.cell(obj, all = T)
  t1 <- system.time({
    obj <- prepare.reg(obj, predictors = genes$targets,
                       responses = genes$tfs)
  })

  # Evaluate on TF with sufficient coverage
  good.TF <- c()
  for(tf in genes$tfs){
    targets <- intersect(net[net$source == tf,"target",drop=T],
                         obj@misc$setting$genes)
    if(length(targets) > 10 & length(targets)/sum(net$source == tf) > 0.2){
      good.TF <- c(good.TF,tf)
    }
  }
  good.TF <- intersect(good.TF, perturbed)
  responses <- good.TF

  t2 <- system.time({
    obj <- run.nn.reg(obj, responses)
  })
})

# Run UMAP
obj <- RunUMAP(obj, dims = 1:30)

# Plot
DimPlot(obj, group.by = "gene",reduction = "umap", label = T)

################################################################################
#                              Analaysis decoupleR
################################################################################
# Avoid introducing bias as CRISPRi influence expression level of TF.
auc.results <- data.frame(matrix(nrow = length(responses), ncol=4))
colnames(auc.results) <- c("nnPC.ulm", "ulm", "combined", "expression")
rownames(auc.results) <- responses

# CoETAS
for(tf in responses){
  y <- !obj$gene == tf
  cells <- obj$gene %in% c("CTRL",tf)
  y <- y[cells]

  # FeaturePlot(obj, tf)
  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = grepl(tf, obj$gene)))

  # CoETAS
  acts <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                  .mor='mor', minsize = 1)
  acts <- dcast(acts, source ~ condition, value.var = "score")
  x <- as.numeric(acts[,colnames(obj)])
  x <- x[cells]
  auc.results[tf,1] <- pROC::roc(y, x)$auc

  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = acts)) +
  #   scale_color_gradientn(colours = viridis::inferno(100))
  # ggplot() + geom_boxplot(aes(obj$gene, y = acts, fill = obj$gene == tf))


  # ETAS
  acts <- run_ulm(mat=obj@assays$RNA$counts[genes$genes,], net=net[net$source == tf,], .source='source', .target='target',
                  .mor='mor', minsize = 1)
  acts <- dcast(acts, source ~ condition, value.var = "score")
  x1 <- as.numeric(acts[,colnames(obj)])
  x1 <- x1[cells]
  auc.results[tf,2] <- pROC::roc(y, x1)$auc

  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = acts)) +
  #   scale_color_gradientn(colours = viridis::inferno(100))
  # ggplot() + geom_boxplot(aes(obj$gene, y = acts, fill = obj$gene == tf))

  # Combined
  x2 <- x*x1
  auc.results[tf,3] <- pROC::roc(y, x2)$auc

  # Expression
  x <- obj@assays$RNA$data[tf,]
  x <- x[cells]
  auc.results[tf,4] <- pROC::roc(y, x)$auc

}

auc.results$TF <- rownames(auc.results)
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable), show.legend = F) +
  geom_point(aes(variable, value, color = TF), size=3) +
  scale_fill_uchicago() + ylim(0.4,1) +
  scale_color_bmj()+
  scale_x_discrete(labels = c("CoETAS", "ETAS", "Combined","Expression"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold")) +
  xlab("Measure") + ylab("AUC (Perturbed vs CTRL)") + ggtitle("Activity with decoupleR")

t.test(auc.results[,1] - auc.results[,4], alternative = "greater")

p1 <- ggplot(melt(auc.result)) + geom_boxplot(aes(variable, value, fill = variable), show.legend = F) +
  geom_point(aes(variable, value, color = TF), size=3) +
  scale_fill_uchicago() + ylim(0.4,1) +
  scale_color_bmj()+
  scale_x_discrete(labels = c("CoETAS", "ETAS","Expression"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold")) +
  xlab("Measure") + ylab("AUC (Perturbed vs CTRL)") + ggtitle("Activity with decoupleR")

################################################################################
#                              Analaysis AUCell
################################################################################
auc.results <- data.frame(matrix(nrow = length(responses), ncol=3))
colnames(auc.results) <- c("nnPC.ulm", "ulm", "expression")
rownames(auc.results) <- responses

# CoETAS
for(tf in responses){
  y <- !obj$gene == tf
  cells <- obj$gene %in% c("CTRL",tf)
  y <- y[cells]

  # FeaturePlot(obj, tf)
  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = grepl(tf, obj$gene)))

  # CoETAS
  acts <- run_aucell(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                     minsize = 1)
  acts <- dcast(acts, source ~ condition, value.var = "score")
  x <- as.numeric(acts[,colnames(obj)])
  x <- x[cells]
  auc.results[tf,1] <- pROC::roc(y, x)$auc

  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = acts)) +
  #   scale_color_gradientn(colours = viridis::inferno(100))
  # ggplot() + geom_boxplot(aes(obj$gene, y = acts, fill = obj$gene == tf))


  # ETAS
  acts <- run_aucell(mat=obj@assays$RNA$counts[genes$genes,], net=net[net$source == tf,], .source='source', .target='target',
                     minsize = 1)
  acts <- dcast(acts, source ~ condition, value.var = "score")
  x <- as.numeric(acts[,colnames(obj)])
  x <- x[cells]
  auc.results[tf,2] <- pROC::roc(y, x)$auc

  # ggplot(data.frame(Embeddings(obj, "umap"))) + geom_point(aes(umap_1,umap_2, col = acts)) +
  #   scale_color_gradientn(colours = viridis::inferno(100))
  # ggplot() + geom_boxplot(aes(obj$gene, y = acts, fill = obj$gene == tf))

  # Expression
  x <- obj@assays$RNA$data[tf,]
  x <- x[cells]
  auc.results[tf,3] <- pROC::roc(y, x)$auc
}

auc.results$TF <- rownames(auc.results)
p2 <- ggplot(melt(auc.results[,-3])) + geom_boxplot(aes(variable, value, fill = variable), show.legend = F) +
  geom_point(aes(variable, value, color = TF), size=3) +
  scale_fill_uchicago() + ylim(0.4,1) +
  scale_x_discrete(labels = c("CoETAS", "ETAS", "Expression"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold")) +
  xlab("Measure") + ylab("AUC (Perturbed vs CTRL)") + ggtitle("Activity with AUcell")

################################################################################
#                              Plot
################################################################################
p3 <- ggplot(melt(auc.results[,c(3,4),drop=F])) +
  geom_boxplot(aes(variable, value, fill = variable), show.legend = F, fill = ggsci::pal_uchicago()(3)[3]) +
  geom_point(aes(variable, value, color = TF), size=3) +
  ylim(0.4,1) +
  scale_x_discrete(labels = c("Expression"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold")) +
  xlab("Measure") + ylab("AUC (Perturbed vs CTRL)") + ggtitle("Expression")

p3 + guides(color = "none") + xlab("") +
  p1 + guides(color = "none") + ylab("") + ggtitle("Activity (DecoupleR)") +
  p2 + ylab("") + xlab("") + ggtitle("Activity (AUCell)")+
  plot_layout(widths = c(1, 1.5,1.5)) &
  scale_color_bmj()

p3 + guides(color = "none") + xlab("") +
  p1 + guides(color = "none") + ylab("") + ggtitle("Activity score") +
  theme(axis.title.x = element_text(hjust=-0.2))+
  plot_layout(widths = c(1, 1.5)) &
  scale_color_bmj()


# Dimplot
to.plot <- obj$gene %in% c(responses, "CTRL")
gene.type <- factor(obj$gene[to.plot], levels = c(sort(responses),"CTRL"))

umap <- data.frame(Embeddings(obj, "umap"))
ggplot(data = umap) +
  geom_point(aes(umap_1, umap_2), color = "lightgrey", size = 0.8) +
  geom_point(data = umap[to.plot,], aes(umap_1, umap_2, col= gene.type), size = 0.8)+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        legend.text = element_text(face = "bold.italic")) +
  scale_color_bmj()+
  guides(color = guide_legend("Perturbation")) +
  ylab("UMAP2") +
  xlab("UMAP1") +
  ggtitle("Papalexi et al. 2021", "CRISPRi on T cells")

# Activity plot
tf <- "SMAD4"
acts <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                .mor='mor', minsize = 1)
acts <- dcast(acts, source ~ condition, value.var = "score")
acts <- as.numeric(acts[,colnames(obj)])

ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = acts),size = 0.8) +
  scale_color_gradientn(colours = rev(RColorBrewer::brewer.pal(11,"Spectral")))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        legend.position = "bottom",
        legend.title =  ggtext::element_markdown())+
  ylab("UMAP2") +
  xlab("UMAP1") + ggtitle("") +
  guides(color = guide_colorbar("*SMAD4* CoETAS"))


# SMAD4
tf <- "SMAD4"
p1 <- ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = obj$gene == tf),size = 0.8) +
  theme_classic()+
  scale_color_manual(values = c("lightgrey","black"))+
  theme(text = element_text(size = 13,face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
        axis.text = element_blank(),
        axis.ticks = element_blank(),
        plot.title = element_text(face = "bold.italic")
  )+
  ylab("") +
  xlab("") + ggtitle(tf) +
  guides(color = guide_legend("Perturbed",override.aes = list(size = 2)))

acts1 <- run_ulm(mat=obj@assays$RNA$counts[genes$genes,], net=net[net$source == tf,], .source='source', .target='target',
                 .mor='mor', minsize = 1)
acts1 <- dcast(acts1, source ~ condition, value.var = "score")
acts1 <- as.numeric(acts1[,colnames(obj)])

p2 <- ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = acts1),size = 0.8) +
  scale_color_gradientn(colours = rev(RColorBrewer::brewer.pal(11,"Spectral")))+
  theme_classic()+
  theme(text = element_text(size = 13,face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
        axis.text = element_blank(),
        axis.ticks = element_blank(),
  )+
  ylab("") +
  xlab("") + ggtitle("") +
  guides(color = guide_colorbar("ETAS"))


p1 + p2


# STAT2
tf <- "STAT2"
p1 <- ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = obj$gene == tf),size = 0.8) +
  theme_classic()+
  scale_color_manual(values = c("lightgrey","black"))+
  theme(text = element_text(size = 13,face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
        axis.text = element_blank(),
        axis.ticks = element_blank(),
        plot.title = element_text(face = "bold.italic")
  )+
  ylab("") +
  xlab("") + ggtitle(tf) +
  guides(color = guide_legend("Perturbed",override.aes = list(size = 2)))

acts <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                .mor='mor', minsize = 1)
acts <- dcast(acts, source ~ condition, value.var = "score")
acts <- as.numeric(acts[,colnames(obj)])

p2 <- ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = acts),size = 0.8) +
  scale_color_gradientn(colours = rev(RColorBrewer::brewer.pal(11,"Spectral")))+
  theme_classic()+
  theme(text = element_text(size = 13,face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
        axis.text = element_blank(),
        axis.ticks = element_blank(),
  )+
  ylab("") +
  xlab("") + ggtitle("") +
  guides(color = guide_colorbar("CoETAS"))

acts1 <- run_ulm(mat=obj@assays$RNA$counts[genes$genes,], net=net[net$source == tf,], .source='source', .target='target',
                 .mor='mor', minsize = 1)
acts1 <- dcast(acts1, source ~ condition, value.var = "score")
acts1 <- as.numeric(acts1[,colnames(obj)])

p3 <- ggplot(data.frame(Embeddings(obj, "umap"))) +
  geom_point(aes(umap_1,umap_2, col = acts1),size = 0.8) +
  scale_color_gradientn(colours = rev(RColorBrewer::brewer.pal(11,"Spectral")))+
  theme_classic()+
  theme(text = element_text(size = 13,face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
        axis.text = element_blank(),
        axis.ticks = element_blank(),
  )+
  ylab("") +
  xlab("") + ggtitle("") +
  guides(color = guide_colorbar("ETAS"))


p1 + p2 + p3
