library(Seurat)
library(Rmagic)
library(ggplot2)
library(viridis)
library(dplyr)
library(ggpubr)
library(openxlsx)
library(gridExtra)

wdir <- "./WPRE_mapping"

# Fig5 A - CART Selection
df_t <- read.xlsx(xlsxFile = paste("./Clustering_Analysis", "Full_T_Myelo_metadata.xlsx", sep = "/"), sheet = "Tcell", rowNames = T)
ggplot(df_t, aes(x = UMAPh_1, y = UMAPh_2)) +
  theme_void(base_size = 12) +
  theme(legend.position = "none") +
  geom_point(data = subset(df_t, is.na(CART_MAGIC)), size = pt_size, color = "grey", alpha = .5) +
  geom_point(data = subset(df_t, CART_MAGIC == "CART_MAGIC"), size = pt_size, color = "darkred", alpha = .5)

# Supplementary Fig11 B - WPRE MAGIC imputation
df <- read.xlsx(xlsxFile = paste("./Clustering_Analysis", "Full_T_Myelo_metadata.xlsx", sep = "/"), sheet = "Full", rowNames = T)

pt_size <- 0.6
p1 <- ggplot(df %>% arrange(TGeneCAR_WPRE_count), aes(x = UMAPh_1, y = UMAPh_2, color = TGeneCAR_WPRE_count)) +
  theme_bw() +
  theme(legend.position = "top") +
  xlab("UMAP 1") +
  ylab("UMAP 2") +
  scale_color_gradient(low = "grey", high = "darkred", name = "Count") +
  geom_point(size = pt_size)
p2 <- ggplot(df %>% arrange(TGeneCAR_WPRE_MAGIC), aes(x = UMAPh_1, y = UMAPh_2, color = TGeneCAR_WPRE_MAGIC)) +
  theme_bw() +
  theme(legend.position = "top") +
  xlab("UMAP 1") +
  ylab("UMAP 2") +
  scale_color_gradient(low = "grey", high = "darkred", name = "Imputed") +
  geom_point(size = pt_size)
p3 <- ggplot(df, aes(x = TGeneCAR_WPRE_MAGIC)) +
  theme_bw() +
  geom_histogram(bins = 150) +
  geom_vline(xintercept = 0.05, color = "red", linetype = "dashed")

png(filename = paste(wdir, "MAGIC_imputation.png"), width = 1400, height = 400)
grid.arrange(p1, p2, p3, nrow = 1)
dev.off()