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(ggtext)
source("script/20240910.R")
net <- get_collectri(organism='human', split_complexes=FALSE)

################################# DAY 7 Run regression #################################
load("data/data_zenodo/in_vitro_benchmark/perturbseq/day7/count_matrix.RData")
colnames(matrix) <- strsplit(colnames(matrix), split = "_") %>% sapply(., function(x) x[2])
matrix <- matrix[!duplicated(rownames(matrix)),!duplicated(colnames(matrix))]
obj <- CreateSeuratObject(t(matrix))
load("data/data_zenodo/in_vitro_benchmark/perturbseq/day7/batch_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
load("data/data_zenodo/in_vitro_benchmark/perturbseq/day7/guide_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
obj <- NormalizeData(obj)
obj <- FindVariableFeatures(obj)
obj <- ScaleData(obj)
obj <- RunPCA(obj)
obj <- RunUMAP(obj, dims = 1:10)

# Preprocess meta data
responses <- c("EGR1", "NR2C2", "E2F4", "CREB1", "YY1", "GABPA", "IRF1", "ELK1", "ELF1", "ETS1", "INTERGENIC")
obj <- sapply(responses,
              function(x){
                obj@meta.data %>% select(contains(x)) %>%
                  rowSums() > 0
              }
) %>% AddMetaData(obj, .)

# Dimplot
perturb.type <- apply(obj@meta.data[,responses], 1, function(x){
  paste(responses[x],collapse = "-")
})

p.list <- list()
for(r in responses){
  p.list[[r]] <- ggplot(data = data.frame(Embeddings(obj, "umap")),
                        aes_string(x = "umap_1", y="umap_2", col = obj@meta.data[,r])) +
    geom_point(size = 0.5) +
    theme_classic()+
    theme(text = element_text(size = 13, face = "bold"),
          axis.line = element_blank(),
          axis.text = element_blank(),
          axis.ticks = element_blank()
    )+
    ylab("") +
    xlab("") + ggtitle(r) +
    scale_color_manual("Perturbed",values = c("lightgrey","black"))

}
ggarrange(plotlist = p.list, common.legend = T, nrow = 2, ncol = 6)


# 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)
  })
  t2 <- system.time({
    obj <- run.nn.reg(obj, responses)
  })
})
t1
t2

################################# Figure F decoupleR D7 #################################
# Benchmark D7-decoupleR
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "ulm","ulm.decoupleR")
rownames(auc.results) <- tfs

for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)
  if(length(targets) > 10){
    group <- obj@meta.data[,tf]

    # Get large batch group
    group[rowSums(obj@meta.data[,38:47])==0] <- "control"
    group[rowSums(obj@meta.data[,4:7]) == 0] <- "FALSE"

    # Smaller batch group
    # group[rowSums(obj@meta.data[,4:7]) != 0] <- "FALSE"

    # Expression
    x <- obj@assays$RNA$data[tf,]
    x <- x[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,1] <- pROC::roc(y, x)$auc

    # Expression-activity
    acts <- run_ulm(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                    .source='source', .target='target',
                    .mor='mor', minsize = 5)
    acts <- dcast(acts, source ~ condition, value.var = "score")
    acts <- as.numeric(acts[,-1])
    x <- acts[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,2] <-  pROC::roc(y, x)$auc

    # Co-expression-activity
    acts1 <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                     .mor='mor', minsize = 5)
    acts1 <- dcast(acts1, source ~ condition, value.var = "score")
    acts1 <- as.numeric(acts1[,-1])
    x <- acts1[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,3] <- pROC::roc(y, x)$auc
  }
}

# Save results
# saveRDS(auc.results, "result/d7.auc.results.ulm.rds")
# auc.results <- readRDS("result/d7.auc.results.ulm.rds")

# Figure F1
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.44,0.55) +
  scale_x_discrete(labels = c("TF Expression", "ETAS", "CoETAS"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"), legend.position = "none") +
  xlab("Measure") + ylab("AUC") + ggtitle("PerturbSeq D7")

################################# Figure F AUCell D7 #################################
# Benchmark D7-AUCell
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "AUCell","nnPC.AUCell")
rownames(auc.results) <- tfs
for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)
  if(length(targets) > 10){
    group <- obj@meta.data[,tf]

    # Larger batch group
    group[rowSums(obj@meta.data[,38:47])==0] <- "control"
    group[rowSums(obj@meta.data[,4:7]) != 0] <- "FALSE"

    # Smaller batch group
    # group[rowSums(obj@meta.data[,4:7]) != 0] <- "FALSE"

    # Expression
    x <- obj@assays$RNA$data[tf,]
    x <- x[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,1] <- pROC::roc(y, x)$auc

    # Expression-activity
    acts <- run_aucell(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                       .source='source', .target='target', minsize = 5)
    acts <- dcast(acts, source ~ condition, value.var = "score")
    acts <- as.numeric(acts[,-1])
    x <- acts[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,2] <-  pROC::roc(y, x)$auc

    # Co-expression-activity
    acts1 <- run_aucell(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,],
                        .source='source', .target='target',
                        minsize = 5)
    acts1 <- dcast(acts1, source ~ condition, value.var = "score")
    acts1 <- as.numeric(acts1[,-1])
    x <- acts1[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,3] <- pROC::roc(y, x)$auc
  }
}

# Save results
# saveRDS(auc.results, "result/d7.auc.results.AUCell.rds")
# auc.results <- readRDS("result/d7.auc.results.AUCell.rds")

# Figure F2
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.44,0.55) +
  scale_x_discrete(labels = c("TF Expression", "ETAS", "CoETAS"))+
  theme_classic()+
  theme(text = element_text(size = 18, face = "bold"), legend.position = "none") +
  xlab("Measure") + ylab("AUC") + ggtitle("PerturbSeq D7 AUCell")

################################# Figure deprecated (schemetic diagram) #################################
# Only for generate schematic diagram
x1 <- rnorm(100)
x2 <- rnorm(100, 1)
ggplot()+geom_density(aes(x1, fill = "Unpertubed")) +
  geom_density(aes(x2, fill = "TF i pertubed"))+
  theme_classic()+
  theme(text = element_text(size = 15, face = "bold"), axis.ticks = element_blank(),
        axis.text = element_blank()) +
  xlab("TF i activity scores") + scale_fill_manual("", values = c("darkred", "blue4"))

################################# DAY 13 Run regression #################################
load("data/holland/in_vitro_benchmark/perturbseq/day13/count_matrix.RData")
colnames(matrix) <- strsplit(colnames(matrix), split = "_") %>% sapply(., function(x) x[2])
matrix <- matrix[!duplicated(rownames(matrix)),!duplicated(colnames(matrix))]
obj <- CreateSeuratObject(t(matrix))
load("data/holland/in_vitro_benchmark/perturbseq/day13/batch_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
load("data/holland/in_vitro_benchmark/perturbseq/day13/guide_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
obj <- NormalizeData(obj)
obj <- FindVariableFeatures(obj)
obj <- ScaleData(obj)
obj <- RunPCA(obj)
obj <- RunUMAP(obj, dims = 1:10)

# Preprocess meta data
responses <- c("EGR1", "NR2C2", "E2F4", "CREB1", "YY1", "GABPA", "IRF1", "ELK1", "ELF1", "ETS1", "INTERGENIC")
obj <- sapply(responses,
              function(x){
                obj@meta.data %>% select(contains(x)) %>%
                  rowSums() > 0
              }
) %>% 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)
  })
  t2 <- system.time({
    obj <- run.nn.reg(obj, responses)
  })
})
t1
t2

FeaturePlot(obj, "batchpt214d_G7")

################################# Figure F decoupleR D13 #################################
# Benchmark D13-decoupleR
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "ulm","nnPC.ulm")
rownames(auc.results) <- tfs

for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)
  if(length(targets) > 10){
    group <- obj@meta.data[,tf]
    group[rowSums(obj@meta.data[,36:45])==0] <- "control" # Cells that are not perturbed

    # Larger batch group
    group[obj$batchpt214d_G7==1] <- "FALSE"

    # Smaller batch group
    # group[obj$batchpt214d_G7!=1] <- "FALSE"

    # Expression
    x <- obj@assays$RNA$data[tf,]
    x <- x[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,1] <- pROC::roc(y, x)$auc

    # Expression-activity
    acts <- run_ulm(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                    .source='source', .target='target',
                    .mor='mor', minsize = 5)
    acts <- dcast(acts, source ~ condition, value.var = "score")
    acts <- as.numeric(acts[,-1])
    x <- acts[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,2] <-  pROC::roc(y, x)$auc

    # Co-expression-activity
    acts1 <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                     .mor='mor', minsize = 5)
    acts1 <- dcast(acts1, source ~ condition, value.var = "score")
    acts1 <- as.numeric(acts1[,-1])
    x <- acts1[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,3] <- pROC::roc(y, x)$auc
  }
}

# Save results
# saveRDS(auc.results, "result/d13.auc.results.ulm.rds")
# auc.results <- readRDS("result/d13.auc.results.ulm.rds")

# Figure F1
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.44,0.55) +
  scale_x_discrete(labels = c("TF Expression", "ETAS", "CoETAS"))+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"), legend.position = "none") +
  xlab("Measure") + ylab("AUC") + ggtitle("PerturbSeq D13")

################################# Figure F AUCell D13 #################################
# Benchmark D13-AUCell
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "AUCell","nnPC.AUCell")
rownames(auc.results) <- tfs
for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)
  if(length(targets) > 10){
    group <- obj@meta.data[,tf]
    group[rowSums(obj@meta.data[,36:45])==0] <- "control"

    # Larger batch group
    group[obj$batchpt214d_G7==1] <- "FALSE"

    # Smaller batch group
    # group[obj$batchpt214d_G7!=1] <- "FALSE"

    x <- obj@assays$RNA$data[tf,]
    x <- x[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,1] <- pROC::roc(y, x)$auc

    acts <- run_aucell(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                       .source='source', .target='target', minsize = 5)
    acts <- dcast(acts, source ~ condition, value.var = "score")
    acts <- as.numeric(acts[,-1])
    x <- acts[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,2] <-  pROC::roc(y, x)$auc

    acts1 <- run_aucell(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,],
                        .source='source', .target='target',
                        minsize = 5)
    acts1 <- dcast(acts1, source ~ condition, value.var = "score")
    acts1 <- as.numeric(acts1[,-1])
    x <- acts1[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,3] <- pROC::roc(y, x)$auc
  }
}

# Save results
# saveRDS(auc.results, "result/d13.auc.results.AUCell.rds")
# auc.results <- readRDS("result/d13.auc.results.AUCell.rds")

# Figure F2
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.44,0.55) +
  scale_x_discrete(labels = c("TF Expression", "ETAS", "CoETAS"))+
  theme_classic()+
  theme(text = element_text(size = 18, face = "bold"), legend.position = "none") +
  xlab("Method") + ylab("AUC") + ggtitle("PerturbSeq D13 AUCell")

################################# CRISPRI, potentially supplementary result (deprecated) #################################
load("data/data_zenodo/in_vitro_benchmark/crispri/count_matrix.RData")
matrix <- matrix[!duplicated(rownames(matrix)),!duplicated(colnames(matrix))]
obj <- CreateSeuratObject(t(matrix))
load("data/data_zenodo/in_vitro_benchmark/crispri/batch_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
load("data/data_zenodo/in_vitro_benchmark/crispri/guide_matrix.RData")
obj <- AddMetaData(obj , data.frame(matrix))
obj <- NormalizeData(obj)
obj <- FindVariableFeatures(obj)
obj <- ScaleData(obj)
obj <- RunPCA(obj)
obj <- RunUMAP(obj, dims = 1:10)

# Preprocess meta data
perturbed <- colnames(matrix) %>% strsplit(.,"-") %>% sapply(., head, n = 1) %>%
  unique()
obj <- sapply(perturbed,
              function(x){
                obj@meta.data %>% select(contains(x)) %>%
                  rowSums() > 0
              }
) %>% 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)
obj <- prepare.graph(obj)
obj <- select.cell(obj, all = T)
obj <- prepare.reg(obj, predictors = genes$targets,
                   responses = genes$tfs)

# Evaluate on TF with sufficient coverage
responses <- 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){
    responses <- c(responses,tf)
  }
}
responses <- intersect(responses, perturbed)
responses <- responses[colSums(obj@meta.data[,responses]) > 10]

d.mat <- dist(obj@reductions$pca@cell.embeddings[,1:30])

apply(obj@meta.data[,responses], 2, function(x) {
  cluster::silhouette(as.numeric(x),d.mat)[,3] %>%
    median
}
) %>% sort

DimPlot(obj, group.by = "MYBL2", order = T)

# Avoid introducing bias as CRISPRi influence expression level of TF.
obj <- run.nn.reg(obj, responses)


################################# CRISPRI DecoupleR #################################
# Benchmark CRISPRI decoupleR
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "ulm","nnPC.ulm")
rownames(auc.results) <- tfs

for(tf in tfs){
  y <- !obj@meta.data[,tf]
  #group <- obj@meta.data[,tf]
  #group[obj@meta.data$Scramble] <- "control"

  x <- obj@assays$RNA$data[tf,]
  #x <- x[group!=FALSE]
  #y <- group[group!=FALSE]
  #y <- y == "control"
  auc.results[tf,1] <- pROC::roc(y, x)$auc

  acts <- run_ulm(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                  .source='source', .target='target',
                  .mor='mor', minsize = 5)
  acts <- dcast(acts, source ~ condition, value.var = "score")
  x <- as.numeric(acts[,colnames(obj)])
  #acts <- as.numeric(acts[,colnames(obj)])
  #x <- acts[group!=FALSE]
  #y <- group[group!=FALSE]
  #y <- y == "control"
  auc.results[tf,2] <-  pROC::roc(y, x)$auc

  acts1 <- run_ulm(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,], .source='source', .target='target',
                   .mor='mor', minsize = 5)
  acts1 <- dcast(acts1, source ~ condition, value.var = "score")
  x <- as.numeric(acts1[,colnames(obj)])
  #acts1 <- as.numeric(acts1[,colnames(obj)])
  #x <- acts1[group!=FALSE]
  #y <- group[group!=FALSE]
  #y <- y == "control"
  auc.results[tf,3] <- pROC::roc(y, x)$auc
}
# Save results
# saveRDS(auc.results, "result/crispri.auc.results.ulm.rds")
# auc.results <- readRDS("result/crispri.auc.results.ulm.rds")

# Figure Supp
ggplot(melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.4,1) +
  scale_x_discrete(labels = c("expression", "decoupleR", "NNet-decoupleR"))+
  theme_classic()+
  theme(text = element_text(size = 15, face = "bold"), legend.position = "none", axis.text.x = element_text(angle = 15, vjust = 1, hjust = 1)) +
  xlab("") + ylab("AUC") + ggtitle("CRISPRi decoupleR")

################################# CRISPRI AUCell #################################
# Benchmark CRISPRI AUCell
tfs <- dimnames(obj@misc$mod$effect)[[1]]
auc.results <- data.frame(matrix(nrow = length(tfs), ncol=3))
colnames(auc.results) <- c("expression", "AUCell","nnPC.AUCell")
rownames(auc.results) <- tfs
for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)
  if(length(targets) > 10){
    group <- obj@meta.data[,tf]
    group[obj@meta.data$Scramble] <- "control"

    x <- obj@assays$RNA$data[tf,]
    x <- x[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,1] <- pROC::roc(y, x)$auc

    acts <- run_aucell(mat=obj@assays$RNA$data[obj@misc$mod$gene.sets$genes,], net=net[net$source == tf,],
                       .source='source', .target='target', minsize = 5)
    acts <- dcast(acts, source ~ condition, value.var = "score")
    acts <- as.numeric(acts[,colnames(obj)])
    x <- acts[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,2] <-  pROC::roc(y, x)$auc

    acts1 <- run_aucell(mat=2*obj@misc$mod$effect[tf,,]^2, net=net[net$source == tf,],
                        .source='source', .target='target',
                        minsize = 5)
    acts1 <- dcast(acts1, source ~ condition, value.var = "score")
    acts1 <- as.numeric(acts1[,colnames(obj)])
    x <- acts1[group!=FALSE]
    y <- group[group!=FALSE]
    y <- y == "control"
    auc.results[tf,3] <- pROC::roc(y, x)$auc
  }
}

# Save results
# saveRDS(auc.results, "result/crispri.auc.results.AUCell.rds")
# auc.results <- readRDS("result/crispri.auc.results.AUCell.rds")

# Figure Supp
ggplot( melt(auc.results)) + geom_boxplot(aes(variable, value, fill = variable)) +
  geom_point(aes(variable, value)) +
  scale_fill_uchicago() + ylim(0.4,1) +
  scale_x_discrete(labels = c("TF Expression", "Expression-AUCell", "NNet-AUCell"))+
  theme_classic()+
  theme(text = element_text(size = 15, face = "bold"), legend.position = "none", axis.text.x = element_text(angle = 15, vjust = 1, hjust = 1)) +
  xlab("Method") + ylab("AUC") + ggtitle("CRISPRi AUCell")

################################# PBMC3K process data #################################
pbmc.data <- Read10X(data.dir = "data/pbmc3k/filtered_gene_bc_matrices/hg19/")
obj <- CreateSeuratObject(counts = pbmc.data, project = "pbmc3k", min.cells = 3, min.features = 200)
obj[["percent.mt"]] <- PercentageFeatureSet(obj, pattern = "^MT-")
obj <- subset(obj, subset = nFeature_RNA > 200 & nFeature_RNA < 2500 & percent.mt < 5)
obj <- NormalizeData(obj)
obj <- FindVariableFeatures(obj)
obj <- ScaleData(obj)
obj <- RunPCA(obj)
obj <- FindNeighbors(obj, dims = 1:10)
obj <- FindClusters(obj, resolution = 0.5)
obj <- RunUMAP(obj, dims = 1:10)
new.cluster.ids <- c("Naive CD4 T", "CD14+ Mono", "Memory CD4 T", "B", "CD8 T", "FCGR3A+ Mono",
                     "NK", "DC", "Platelet")
names(new.cluster.ids) <- levels(obj)
obj <- RenameIdents(obj, new.cluster.ids)

# Save the baseline R object
# saveRDS(obj, "result/obj.pbmc3k.expression.rds")
# obj <- readRDS("result/obj.pbmc3k.expression.rds")

################################# Figure B1 2000 VF dimplot #################################
# Figure B1
DimPlot(obj, 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(),
  )+
  scale_color_bmj()+
  xlab("")+ylab("")+
  ggtitle("2000 Variable features")

DimPlot(obj, pt.size = 2)  +
  theme_classic()+
  theme(text = element_text(size=13, face = "bold"),
        axis.line = element_blank(),
        legend.position = "bottom",
  )+
  scale_color_bmj()+
  guides(colour = guide_legend(nrow = 2, title = "Celltype"))+
  xlab("")+ylab("")+
  ggtitle("2000 Variable features")

################################# NNet regression #################################
# Prepare nn 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)
  })
  # Select valid TFs that has enough target coverage.
  responses <- 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){
      responses <- c(responses,tf)
    }
  }
  t2 <- system.time({
    obj <- run.nn.reg(obj, responses, return.prune = T)
  })
  tfs <- responses
})
t1
t2

################################# Figure B3 Infer TF activity based on co-expression #################################
# Calculate TF activity based on co-expression.
acts.ulm <- acts.aucell <- matrix(nrow = length(tfs), ncol = ncol(obj), dimnames = list(tfs, colnames(obj)))
for(tf in tfs){
  targets <- intersect(net[net$source == tf,"target",drop=T],
                       obj@misc$mod$gene.sets$genes)

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

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

################################# Figure B3 Co-expression-decoupleR #################################
# Build Seurat object based on co-expression-decoupleR
obj.ulm <- CreateSeuratObject(acts.ulm)
LayerData(obj.ulm, "data") <- acts.ulm
VariableFeatures(obj.ulm) <- rownames(obj.ulm)
obj.ulm <- ScaleData(obj.ulm)
obj.ulm <- RunPCA(obj.ulm)
obj.ulm <- FindNeighbors(obj.ulm, dims = 1:10)
obj.ulm <- FindClusters(obj.ulm, resolution = 0.28)
obj.ulm <- RunUMAP(obj.ulm, dims = 1:10)
obj.ulm$celltype <- Idents(obj)

# Save results
# saveRDS(obj.ulm, "result/obj.pbmc3k.ulm.rds")
# obj.ulm <- readRDS("result/obj.pbmc3k.ulm.rds")

# Figure B3
DimPlot(obj.ulm, 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("288 TF CoETAS")

################################# Supplementary Figure A Co-expression-AUCell #################################
# Build Seurat object based on co-expression-AUCell
obj.auc <- CreateSeuratObject(acts.aucell)
LayerData(obj.auc, "data") <- acts.aucell
VariableFeatures(obj.auc) <- rownames(obj.auc)
obj.auc <- ScaleData(obj.auc)
obj.auc <- RunPCA(obj.auc)
obj.auc <- FindNeighbors(obj.auc, dims = 1:10)
obj.auc <- FindClusters(obj.auc, resolution = 0.25)
obj.auc <- RunUMAP(obj.auc, dims = 1:10)
obj.auc$celltype <- Idents(obj)

# Save results
# saveRDS(obj.auc, "result/obj.pbmc3k.auc.rds")
# obj.auc <- readRDS("result/obj.pbmc3k.auc.rds")

# Figure B3
DimPlot(obj.auc, 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("288 TF CoETAS (AUCell)")


################################# Figure B4 Expression-decoupleR #################################
# Build Seurat object based on expression-decoupleR
# TF activity
tfs <- rownames(obj.ulm)
genes <- select.gene(obj, tfs = unique(net$source), targets = unique(net$target), min.cells = 0)
genes <- unique(c(tfs,genes$targets))
acts.ulm.expression <- run_ulm(mat=obj@assays$RNA$data[genes,], net=net[net$source %in% tfs,], .source='source', .target='target',
                               .mor='mor', minsize = 5)
acts.ulm.expression <- dcast(acts.ulm.expression, source ~ condition, value.var = "score")
rownames(acts.ulm.expression) <- acts.ulm.expression[,1]
acts.ulm.expression <- acts.ulm.expression[,-1]

# Build Seurat object
obj.ulm.expression <- CreateSeuratObject(acts.ulm.expression)
LayerData(obj.ulm.expression, "data") <- acts.ulm.expression
VariableFeatures(obj.ulm.expression) <- rownames(obj.ulm.expression)
obj.ulm.expression <- ScaleData(obj.ulm.expression)
obj.ulm.expression <- RunPCA(obj.ulm.expression)
obj.ulm.expression <- FindNeighbors(obj.ulm.expression, dims = 1:10)
obj.ulm.expression <- FindClusters(obj.ulm.expression, resolution = 0.99)
obj.ulm.expression <- RunUMAP(obj.ulm.expression, dims = 1:10)
obj.ulm.expression$celltype <- Idents(obj)

# Save results
# saveRDS(obj.ulm.expression, "result/obj.pbmc3k.ulm.expression.rds")
# obj.ulm.expression <- readRDS("result/obj.pbmc3k.ulm.expression.rds")

# Figure B2
DimPlot(obj.ulm.expression, 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("288 TF ETAS")

################################# Figure Response to reviwers (decoupleR) #################################
acts.ulm.expression <- obj.ulm.expression@assays$RNA$counts
rownames(acts.ulm.expression) <- paste0(rownames(acts.ulm.expression), ".expression")
acts.ulm <- obj.ulm@assays$RNA$counts
rownames(acts.ulm) <- paste0(rownames(acts.ulm), ".nnet")

acts.ulm.combined <- rbind(acts.ulm.expression, acts.ulm)
obj.ulm.combined <- CreateSeuratObject(acts.ulm.combined)
LayerData(obj.ulm.combined, "data") <- acts.ulm.combined
VariableFeatures(obj.ulm.combined) <- rownames(obj.ulm.combined)
obj.ulm.combined <- ScaleData(obj.ulm.combined)
obj.ulm.combined <- RunPCA(obj.ulm.combined)
obj.ulm.combined <- FindNeighbors(obj.ulm.combined, dims = 1:10)
obj.ulm.combined <- FindClusters(obj.ulm.combined, resolution = 0.4)
obj.ulm.combined <- RunUMAP(obj.ulm.combined, dims = 1:10)
obj.ulm.combined$celltype <- Idents(obj)

# Save results
# saveRDS(obj.ulm.combined, "result/obj.pbmc3k.ulm.combined.rds")
# obj.ulm.combined <- readRDS("result/obj.pbmc3k.ulm.combined.rds")

# Figure B2
DimPlot(obj.ulm.combined, 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("288 TF ETAS + CoETAS (DecoupleR)")


################################# Figure Response to reviwers (AUCell) #################################
acts.auc.expression <- obj.auc.expression@assays$RNA$counts
rownames(acts.auc.expression) <- paste0(rownames(acts.auc.expression), ".expression")
acts.auc <- obj.auc@assays$RNA$counts
rownames(acts.auc) <- paste0(rownames(acts.auc), ".nnet")

acts.auc.combined <- rbind(acts.auc.expression, acts.auc)
obj.auc.combined <- CreateSeuratObject(acts.auc.combined)
LayerData(obj.auc.combined, "data") <- acts.auc.combined
VariableFeatures(obj.auc.combined) <- rownames(obj.auc.combined)
obj.auc.combined <- ScaleData(obj.auc.combined)
obj.auc.combined <- RunPCA(obj.auc.combined)
obj.auc.combined <- FindNeighbors(obj.auc.combined, dims = 1:10)
obj.auc.combined <- FindClusters(obj.auc.combined, resolution = 0.5)
obj.auc.combined <- RunUMAP(obj.auc.combined, dims = 1:10)
obj.auc.combined$celltype <- Idents(obj)

# Save results
# saveRDS(obj.auc.combined, "result/obj.pbmc3k.auc.combined.rds")
# obj.auc.combined <- readRDS("result/obj.pbmc3k.auc.combined.rds")

# Figure B2
DimPlot(obj.auc.combined, group.by = "celltype", pt.size=0.1)  +
  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("288 TF ETAS + CoETAS (AUCell)")

################################# Supplementary Figure A Expression-AUCell #################################
# Build Seurat object based on expression-AUCell
# TF activity
tfs <- rownames(obj.ulm)
genes <- select.gene(obj, tfs = unique(net$source), targets = unique(net$target), min.cells = 0)
genes <- unique(c(tfs,genes$targets))
acts.auc.expression <- run_aucell(mat=obj@assays$RNA$data[genes,], net=net[net$source %in% tfs,], .source='source', .target='target',
                                  minsize = 5)
acts.auc.expression <- dcast(acts.auc.expression, source ~ condition, value.var = "score")
rownames(acts.auc.expression) <- acts.auc.expression[,1]
acts.auc.expression <- acts.auc.expression[,-1]

# Build Seurat object
obj.auc.expression <- CreateSeuratObject(acts.auc.expression)
LayerData(obj.auc.expression, "data") <- acts.auc.expression
VariableFeatures(obj.auc.expression) <- rownames(obj.auc.expression)
obj.auc.expression <- ScaleData(obj.auc.expression)
obj.auc.expression <- RunPCA(obj.auc.expression)
obj.auc.expression <- FindNeighbors(obj.auc.expression, dims = 1:10)
obj.auc.expression <- FindClusters(obj.auc.expression, resolution = 0.7)
obj.auc.expression <- RunUMAP(obj.auc.expression, dims = 1:10)
obj.auc.expression$celltype <- Idents(obj)

# Save results
# saveRDS(obj.auc.expression, "result/obj.pbmc3k.auc.expression.rds")
# obj.auc.expression <- readRDS("result/obj.pbmc3k.auc.expression.rds")

DimPlot(obj.auc.expression, 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("288 TF ETAS (AUCell)")

################################# Figure B3 TF degree #################################
tf.degree <- matrix(0, ncol = ncol(obj), nrow = length(tfs), dimnames = list(tfs,colnames(obj)))
for(tf in tfs) tf.degree[tf,] <- colSums(2*obj@misc$mod$effect[tf,,]^2)
obj.degree.tfs <- CreateSeuratObject(tf.degree)
LayerData(obj.degree.tfs, "data") <- obj.degree.tfs@assays$RNA$counts
VariableFeatures(obj.degree.tfs) <- rownames(obj.degree.tfs)
obj.degree.tfs <- ScaleData(obj.degree.tfs)
obj.degree.tfs <- RunPCA(obj.degree.tfs)
obj.degree.tfs <- FindNeighbors(obj.degree.tfs, dims = 1:10)
obj.degree.tfs <- FindClusters(obj.degree.tfs, resolution = 0.2)
obj.degree.tfs <- RunUMAP(obj.degree.tfs, dims = 1:10)
obj.degree.tfs$celltype <- Idents(obj)

# Save results
# saveRDS(obj.degree.tfs, "result/obj.pbmc3k.degree.rds")
# obj.degree.tfs <- readRDS("result/obj.pbmc3k.degree.rds")

DimPlot(obj.degree.tfs, 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("288 TF Degree")



################################# Figure C Cluster analysis #################################
# perform LRA
tfs <- rownames(obj.ulm)
genes <- select.gene(obj, tfs = unique(net$source), targets = unique(net$target), min.cells = 0)
lra.genes <- unique(c(tfs,genes$targets))
obj <- prepare.reg(obj, predictors = genes$targets,
                   responses = lra.genes)

obj.lra.tfs <- CreateSeuratObject(t(obj@misc$setting$lra[,tfs]))
LayerData(obj.lra.tfs, "data") <- obj.lra.tfs@assays$RNA$counts
VariableFeatures(obj.lra.tfs) <- rownames(obj.lra.tfs)
obj.lra.tfs <- ScaleData(obj.lra.tfs)
obj.lra.tfs <- RunPCA(obj.lra.tfs)
obj.lra.tfs <- FindNeighbors(obj.lra.tfs, dims = 1:10)
obj.lra.tfs <- FindClusters(obj.lra.tfs, resolution = 1.3)
obj.lra.tfs <- RunUMAP(obj.lra.tfs, dims = 1:10)
obj.lra.tfs$celltype <- Idents(obj)

# Build Seurat object based on TF expression
tfs <- rownames(obj.ulm)
obj.expression.tfs <- CreateSeuratObject(obj@assays$RNA$data[tfs,])
LayerData(obj.expression.tfs, "data") <- obj.expression.tfs@assays$RNA$counts
VariableFeatures(obj.expression.tfs) <- rownames(obj.expression.tfs)
obj.expression.tfs <- ScaleData(obj.expression.tfs)
obj.expression.tfs <- RunPCA(obj.expression.tfs)
obj.expression.tfs <- FindNeighbors(obj.expression.tfs, dims = 1:10)
obj.expression.tfs <- FindClusters(obj.expression.tfs, resolution = 1.2)
obj.expression.tfs <- RunUMAP(obj.expression.tfs, dims = 1:10)
obj.expression.tfs$celltype <- Idents(obj)

# Build Seurat object based on LRA-decoupleR
acts.ulm.lra <- run_ulm(mat=t(obj@misc$setting$lra), net=net[net$source %in% tfs,], .source='source', .target='target',
                        .mor='mor', minsize = 5)
acts.ulm.lra <- dcast(acts.ulm.lra, source ~ condition, value.var = "score")
acts.ulm.lra <- acts.ulm.lra[,-1]

# Build Seurat object
obj.ulm.lra <- CreateSeuratObject(acts.ulm.lra)
LayerData(obj.ulm.lra, "data") <- acts.ulm.lra
VariableFeatures(obj.ulm.lra) <- rownames(obj.ulm.lra)
obj.ulm.lra <- ScaleData(obj.ulm.lra)
obj.ulm.lra <- RunPCA(obj.ulm.lra)
obj.ulm.lra <- FindNeighbors(obj.ulm.lra, dims = 1:10)
obj.ulm.lra <- FindClusters(obj.ulm.lra, resolution = 1.05)
obj.ulm.lra <- RunUMAP(obj.ulm.lra, dims = 1:10)
obj.ulm.lra$celltype <- Idents(obj)

# Build Seurat object based on LRA-AUC
acts.auc.lra <- run_aucell(mat=t(obj@misc$setting$lra), net=net[net$source %in% tfs,], .source='source', .target='target',
                           minsize = 5)
acts.auc.lra <- dcast(acts.auc.lra, source ~ condition, value.var = "score")
acts.auc.lra <- acts.auc.lra[,-1]

# Build Seurat object
obj.auc.lra <- CreateSeuratObject(acts.auc.lra)
LayerData(obj.auc.lra, "data") <- acts.auc.lra
VariableFeatures(obj.auc.lra) <- rownames(obj.auc.lra)
obj.auc.lra <- ScaleData(obj.auc.lra)
obj.auc.lra <- RunPCA(obj.auc.lra)
obj.auc.lra <- FindNeighbors(obj.auc.lra, dims = 1:10)
obj.auc.lra <- FindClusters(obj.auc.lra, resolution = 0.67)
obj.auc.lra <- RunUMAP(obj.auc.lra, dims = 1:10)
obj.auc.lra$celltype <- Idents(obj)

################################# Figure C Quantitative analysis of clusters #################################
# Calculate Rand index
ari <- rep(0, 12)
names(ari) <- c("2000 VF expression", "NNet-decoupleR", "NNet-AUCell", "Expression-decoupleR",
                "Expression-AUCell", "LRA-decoupleR", "LRA-AUCell","288 TF expression",
                "288 TF LRA", "288 TF NNet degree", "Combined-decoupleR","Combined-AUCell")
ari[1] <- adj.rand.index(obj$seurat_clusters, obj$seurat_clusters)
ari[2] <- adj.rand.index(obj$seurat_clusters, obj.ulm$seurat_clusters)
ari[3] <-adj.rand.index(obj$seurat_clusters, obj.auc$seurat_clusters)
ari[4] <-adj.rand.index(obj$seurat_clusters, obj.ulm.expression$seurat_clusters)
ari[5] <-adj.rand.index(obj$seurat_clusters, obj.auc.expression$seurat_clusters)
ari[6] <-adj.rand.index(obj$seurat_clusters, obj.ulm.lra$seurat_clusters)
ari[7] <-adj.rand.index(obj$seurat_clusters, obj.auc.lra$seurat_clusters)
ari[8] <-adj.rand.index(obj$seurat_clusters, obj.expression.tfs$seurat_clusters)
ari[9] <-adj.rand.index(obj$seurat_clusters, obj.lra.tfs$seurat_clusters)
ari[10] <-adj.rand.index(obj$seurat_clusters, obj.degree.tfs$seurat_clusters)
ari[11] <-adj.rand.index(obj$seurat_clusters, obj.ulm.combined$seurat_clusters)
ari[12] <-adj.rand.index(obj$seurat_clusters, obj.auc.combined$seurat_clusters)

## Calculate Silhouette index
sil <- rep(0,12)
names(sil) <- c("2000 VF expression", "NNet-decoupleR", "NNet-AUCell", "Expression-decoupleR",
                "Expression-AUCell", "LRA-decoupleR", "LRA-AUCell","288 TF expression",
                "288 TF LRA", "288 TF NNet degree", "Combined-decoupleR","Combined-AUCell")
sil[1] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[2] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.ulm@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[3] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.auc@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[4] <- cluster::silhouette(as.numeric(obj.ulm.expression$seurat_clusters),
                              dist(obj.ulm.expression@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[5] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.auc.expression@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[6] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.ulm.lra@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[7] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.auc.lra@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[8] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.expression.tfs@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[9] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                              dist(obj.lra.tfs@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[10] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                               dist(obj.degree.tfs@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[11] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                               dist(obj.ulm.combined@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median
sil[12] <- cluster::silhouette(as.numeric(obj$seurat_clusters),
                               dist(obj.auc.combined@reductions$pca@cell.embeddings[,1:10]))[,3] %>% median


pbmc3k.cluster.eval <- data.frame(ari = ari, sil = sil)
rownames(pbmc3k.cluster.eval) <- names(ari)

# Write quantitative analysis result.
# write.table(pbmc3k.cluster.eval, "result/pbmc3k.cluster.eval.txt")

# Read quantitative analysis result.
# pbmc3k.cluster.eval <- read.table("result/pbmc3k.cluster.eval.txt")
# ari <- pbmc3k.cluster.eval$ari
# sil <- pbmc3k.cluster.eval$sil

# Rename for better intepretation
rownames(pbmc3k.cluster.eval) <- c("VF expression", "CoETAS (decoupleR)", "CoETAS (AUCell)",
                                   "ETAS (decoupleR)", "ETAS (AUCell)",  "LTAS (decoupleR)", "LTAS (AUCell)",
                                   "TF expression", "TF LRA", "TF degree", "Combined (decoupleR)","Combined (AUCell)")

# Supplementary Figure B
ggplot() + geom_point(aes(ari, sil, col = rownames(pbmc3k.cluster.eval)), size = 3) +
  geom_text_repel(aes(ari, sil, label = rownames(pbmc3k.cluster.eval),
                      col = rownames(pbmc3k.cluster.eval)), size = 4) +
  theme_classic() +
  theme(text = element_text(size = 13, face = "bold"),
        legend.position = "blank")+
  scale_color_manual(values = c(ggsci::pal_uchicago()(9), "black","red","blue")) +
  xlab("Adjusted Rand Index") +
  ylab("Median Silhouette Index")+
  ggtitle("Clustering Evaluation")


# Figure C
pbmc3k.cluster.eval <- pbmc3k.cluster.eval[-c(6,7,9,11,12),]
ari <- pbmc3k.cluster.eval$ari
sil <- pbmc3k.cluster.eval$sil
ggplot() + geom_point(aes(ari, sil, col = rownames(pbmc3k.cluster.eval)), size = 3) +
  geom_text_repel(aes(ari, sil, label = rownames(pbmc3k.cluster.eval),
                      col = rownames(pbmc3k.cluster.eval)), size = 4) +
  theme_classic() + theme(text = element_text(size = 13, face = "bold"),
                          legend.position = "blank")+
  scale_color_uchicago() +
  xlab("Adjusted Rand Index") +
  ylab("Median Silhouette Index")+
  ggtitle("Clustering Evaluation")


################################# Figure D,E, DAT analysis #################################
# Difference in TF activity ranking
Idents(obj.ulm) <- "celltype"
Idents(obj.ulm.expression) <- "celltype"
tfs <- rownames(obj.ulm)
acts.ulm <- obj.ulm@assays$RNA$counts[tfs,]
acts.ulm.expression <- obj.ulm.expression@assays$RNA$counts[tfs,]

# TF activity t-test B cell vs others
group <- 1-as.numeric(obj.ulm$celltype=="B")
t.ulm <- apply(acts.ulm, 1, function(x) tryCatch(t.test(x~group)$statistic, error = function(cond) 0))
p.val.ulm  <- apply(acts.ulm, 1, function(x) tryCatch(t.test(x~group)$p.val, error = function(cond) 0)) %>% p.adjust(., method = "BH")
group <- 1-as.numeric(obj.ulm.expression$celltype=="B")
t.ulm.expression <- apply(acts.ulm.expression, 1, function(x) t.test(x~group)$statistic)
p.val.ulm.expression  <- apply(acts.ulm.expression, 1, function(x) tryCatch(t.test(x~group)$p.val, error = function(cond) 0)) %>% p.adjust(., method = "BH")

# Write
pbmc.dat.t.val <- data.frame(t.ulm = t.ulm, t.ulm.expression = t.ulm.expression, p.val.ulm = p.val.ulm, p.val.ulm.expression = p.val.ulm.expression)
# write.table(pbmc.dat.t.val, "result/pbmc.dat.t.val.txt")

# Read
# pbmc.dat.t.val <- read.table("result/pbmc.dat.t.val.txt")
# t.ulm <- pbmc.dat.t.val$t.ulm
# t.ulm.expression <- pbmc.dat.t.val$t.ulm.expression
# p.val.ulm <-  pbmc.dat.t.val$p.val.ulm
# p.val.ulm.expression <-  pbmc.dat.t.val$p.val.ulm.expression
names(t.ulm) <- names(t.ulm.expression) <- names(p.val.ulm) <- names(p.val.ulm.expression) <- rownames(pbmc.dat.t.val)

################################# Figure deprecated Venn diagram #################################
dat.ulm <- p.val.ulm[p.val.ulm<0.05 & t.ulm > 0] %>% names
dat.ulm.expression <- p.val.ulm.expression[p.val.ulm.expression<0.05 & t.ulm.expression > 0] %>% names
ggvenn(list("NNet-decoupleR" = dat.ulm, "Expression-decoupleR" = dat.ulm.expression), fill_color = c("red","darkblue"), text_size = 5)

################################# Figure E Violin plot #################################
p1 <- VlnPlot(obj.ulm , "TBX21", group.by = "celltype", alpha = 0.05) + scale_fill_bmj() +
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        axis.text.x = element_blank(),
        legend.position = "blank") +xlab("") + ylab("") + ggtitle("CoETAS")
p2 <-VlnPlot(obj.ulm.expression , "TBX21", group.by = "celltype", alpha = 0.05) + scale_fill_bmj()+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        axis.text.x = element_blank(),
        axis.title.y = ggtext::element_markdown(),
        legend.position = "blank")+xlab("")+ ylab("*TBX21*")+ ggtitle("ETAS")
p3 <-VlnPlot(obj, "TBX21", alpha = 0.05)+ scale_fill_bmj()+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        axis.text.x = element_text(angle = 30, hjust = 1),
        legend.position = "blank") +xlab("")+ ylab("")+ ggtitle("Expression") + xlab("Celltype")
p1/p2/p3

ggplot() + geom_density(aes(x = obj.ulm@assays$RNA$counts["TBX21",], fill = obj.ulm$celltype))+ scale_fill_bmj()+
  theme_classic()+
  theme(text = element_text(size = 13, face = "bold"),
        axis.title.x = ggtext::element_markdown()) +
  guides(fill = guide_legend(title = "Celltype"))+
  xlab("*TBX21* activity")+ ylab("Density")+ ggtitle("")

################################# Figure deprecated DotPlot of DAT expression #################################
# Get Top 10 DAT based on expression and co-expression
top.10.ulm <- sort(t.ulm) %>% tail(20) %>% names
top.10.ulm.expression <- sort(t.ulm.expression) %>% tail(20) %>% names
common.dat <- intersect(top.10.ulm, top.10.ulm.expression)
dot.features <- unique(c(top.10.ulm,top.10.ulm.expression))
names(dot.features)[dot.features%in%top.10.ulm] <- "DAT NNet-decoupleR"
names(dot.features)[dot.features%in%top.10.ulm.expression] <- "DAT expression-decoupleR"
names(dot.features)[dot.features%in%common.dat] <- "Common DAT"

# Figure D (discarded)
DotPlot(obj, features = dot.features)  + theme(text = element_text(size = 15, face = "bold"),
                                               axis.text.x = element_text(angle = 90))+
  scale_color_gradientn(colours = RColorBrewer::brewer.pal(11, "Reds"))

################################# Figure D DAT log fold change and sparsity #################################
top.10.ulm <- sort(t.ulm) %>% tail(100) %>% names
top.10.ulm.expression <- sort(t.ulm.expression) %>% tail(100) %>% names
dot.features1 <- top.10.ulm[!top.10.ulm %in% top.10.ulm.expression] %>% tail(15)
dot.features2 <- top.10.ulm.expression[!top.10.ulm.expression %in% top.10.ulm] %>% tail(15)
dot.features <- c(dot.features1, dot.features2)
names(dot.features)[dot.features%in%dot.features1] <- "DAT NNet-decoupleR"
names(dot.features)[dot.features%in%dot.features2] <- "DAT expression-decoupleR"

fc <- FoldChange(obj, ident.1 = "B", features = tfs)
fc$avg_log2FC[fc$pct.1==0] <- 0
ggplot() + geom_point(aes(t.ulm.expression[order(fc$avg_log2FC)], t.ulm[order(fc$avg_log2FC)],
                          col = fc$avg_log2FC[order(fc$avg_log2FC)], size = round(fc$pct.1*100)[order(fc$avg_log2FC)]))+
  geom_text_repel(aes(t.ulm.expression[dot.features], t.ulm[dot.features], label = dot.features),
                  col = "grey23", size = 4, fontface=4)+
  theme_classic() + theme(text = element_text(size = 13, face = "bold"),
                          plot.title = ggtext::element_markdown(),
                          axis.title.x = ggtext::element_markdown(),
                          axis.title.y = ggtext::element_markdown())+
  geom_hline(yintercept = 0, linetype = "dashed")+
  geom_vline(xintercept = 0, linetype = "dashed")+
  scale_color_gradientn("LFC",
                        colours = RColorBrewer::brewer.pal(n = 11,"RdBu") %>% rev(), limits = c(-7.1,7.1)) +
  scale_size_continuous("% Expressed")+
  ggtitle("DAT *t*-statistics: B vs others")+
  xlab("*t*-stat by ETAS") +
  ylab("*t*-stat by CoETAS")
