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

wdir <- "./WPRE_mapping"

t <- readRDS(paste(wdir, "Full_dataset_counts.rds", sep = "/"))

ss <- stringr::str_split_fixed(colnames(t), pattern = "-", n = 2)[,1] %>% unique()
full_t <- list()
for (sample in ss) {
  print(sample)
  sample_cells <- grep(paste0(sample, "-"), colnames(t), value = T)
  t_sample <- t(t[,sample_cells])
  # keep genes expressed in at least 10 cells
  keep_cols <- colSums(t_sample > 0) > 10
  t_sample <- t_sample[,keep_cols]
  t_sample <- library.size.normalize(t_sample)
  t_sample <- sqrt(t_sample)
  t_sample_MAGIC <- magic(t_sample, genes = c("Cd8a","Cd4","Cd3e","Cd14","Mki67","Top2a","TGeneCAR-WPRE"))
  magic_res <- as.data.frame(t_sample_MAGIC$result)
  full_t[[sample]] <- magic_res
}
saveRDS(object = full_t, file = paste(wdir, "MAGICimputation.rds", sep = "/"))

# Add MAGIC imputation to Seurat meta-data
df <- data.frame()
for (sample in names(full_t)) {
  df_sample <- full_t[[sample]]
  df_sample$Sample <- sample
  df <- rbind(df, df_sample)
}
full_obj <- readRDS(paste(wdir, "Full_dataset.rds", sep = "/"))
md <- df[, "TGeneCAR-WPRE", drop = F]
colnames(md) <- "TGeneCAR_WPRE_MAGIC"
head(md)
full_obj <- AddMetaData(object = full_obj, metadata = md)

# Add WPRE counts to Seurat meta-data
wpre_expr <- data.frame("TGeneCAR_WPRE_count" = full_obj@assays$RNA@counts["TGeneCAR-WPRE",])
full_obj <- AddMetaData(object = full_obj, metadata = wpre_expr)