library(ProjecTILs)
library(Seurat)
library(ggplot2)
library(patchwork)
library(SignatuR)
library(UCell)

# Load Referece
wdir <- "./ProjectTILs"
ref <- load.reference.map(paste(wdir, "ref_TILAtlas_mouse_v1.rds", sep = "/"))

####################
#### CART MAGIC ####
####################
obj <- readRDS(paste(wdir, "Full_CART.rds", sep = "/"))

# Project data by subject (two mice separately)
query.list <- SplitObject(obj, split.by = "orig.ident")
query.projected <- make.projection(query.list, ref = ref, ncores = 2)
query.projected.merged <- Reduce(f = merge.Seurat.embeddings, x = query.projected)
query.projected <- SplitObject(query.projected.merged, split.by = "PROP.Condition")

plots <- list()
lgt <- length(query.projected)
palette <- ref@misc$atlas.palette
for (i in seq_along(query.projected)) {
  sample <- names(query.projected)[i]
  plots[[i]] <- plot.projection(ref, query.projected[[i]], linesize = 0.5, pointsize = 0) +
    ggtitle(sample) +
    NoLegend() +
    ylim(-10, 10)
}

png(file = paste(wdir, "Full_WPRE_CART_TILs.png", sep = "/"), width = 1400, height = 400)
wrap_plots(plots, ncol = 4)
dev.off()

#### Barplot
tb.full <- data.frame()
palette <- ref@misc$atlas.palette
for (i in seq_along(query.projected)) {
  sample <- names(query.projected)[i]
  print(sample)
  query.projected[[i]] <- cellstate.predict(ref = ref, query = query.projected[[i]], reduction = "umap")
  metric <- "percent"
  labels.col <- "functional.cluster"
  labels <- ref[[labels.col]][,1]
  states_all <- levels(factor(labels))
  query = query.projected[[i]]
  tb <- table(factor(query[[labels.col]][,1], levels=states_all))
  if (metric=="percent") {  #normalize
    tb <- tb*100/sum(tb)
    tb.m <- reshape2::melt(tb)
    colnames(tb.m) <- c("Cell_state","Perc_cells")
  } else if (metric=="count") {
    tb.m <- reshape2::melt(tb)
    colnames(tb.m) <- c("Cell_state","Ncells")
  } else {
    stop("Unknown metric specified (Must be either 'count' or 'percent')")
  }
  tb.m$Group <- sample
  tb.full <- rbind(tb.full, tb.m)
}

tb.full$Group <- factor(tb.full$Group, levels = c("Empty", "Ortho", "IFN", "P2A"))
p <- ggplot(data = tb.full, mapping = aes(x = Group, y = Perc_cells, fill = Cell_state)) +
  theme_bw(base_size = 18) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank(),
        strip.background = element_rect(fill = "white", color = "white")) +
  geom_bar(stat = "identity", position = "stack", width = 1) +
  scale_fill_manual(values = palette, name = "Treatment") +
  scale_y_continuous(n.breaks = 10) +
  ggtitle("Cd8+ WPRE+") +
  ylab("% Cells") +
  xlab("") +
  facet_grid(.~Group, scales = "free_x")

ggsave(filename = paste(wdir, "Full_WPRE_CART_TILs_barplot.png", sep = "/"), width = 7, height = 7)


########################
#### CART MAGIC Neg ####
########################
obj <- readRDS(paste(wdir, "Full_CART_Neg_final.rds", sep = "/"))

# Project data by subject (two mice separately)
query.list <- SplitObject(obj, split.by = "orig.ident")
query.projected <- make.projection(query.list, ref = ref, ncores = 2)
query.projected.merged <- Reduce(f = merge.Seurat.embeddings, x = query.projected)
query.projected <- SplitObject(query.projected.merged, split.by = "PROP.Condition")

plots <- list()
lgt <- length(query.projected)
palette <- ref@misc$atlas.palette
for (i in seq_along(query.projected)) {
  sample <- names(query.projected)[i]
  plots[[i]] <- plot.projection(ref, query.projected[[i]], linesize = 0.5, pointsize = 0) +
    ggtitle(sample) +
    NoLegend() +
    ylim(-10, 10)
}

png(file = paste(wdir, "Full_WPRE_CART_Neg_TILs.png", sep = "/"), width = 1400, height = 400)
wrap_plots(plots, ncol = 4)
dev.off()

#### Barplot
tb.full <- data.frame()
palette <- ref@misc$atlas.palette
for (i in seq_along(query.projected)) {
  sample <- names(query.projected)[i]
  print(sample)
  query.projected[[i]] <- cellstate.predict(ref = ref, query = query.projected[[i]], reduction = "umap")
  metric <- "percent"
  labels.col <- "functional.cluster"
  labels <- ref[[labels.col]][,1]
  states_all <- levels(factor(labels))
  query = query.projected[[i]]
  tb <- table(factor(query[[labels.col]][,1], levels=states_all))
  if (metric=="percent") {  #normalize
    tb <- tb*100/sum(tb)
    tb.m <- reshape2::melt(tb)
    colnames(tb.m) <- c("Cell_state","Perc_cells")
  } else if (metric=="count") {
    tb.m <- reshape2::melt(tb)
    colnames(tb.m) <- c("Cell_state","Ncells")
  } else {
    stop("Unknown metric specified (Must be either 'count' or 'percent')")
  }
  tb.m$Group <- sample
  tb.full <- rbind(tb.full, tb.m)
}

tb.full$Group <- factor(tb.full$Group, levels = c("Empty", "Ortho", "IFN", "P2A"))
p <- ggplot(data = tb.full, mapping = aes(x = Group, y = Perc_cells, fill = Cell_state)) +
  theme_bw(base_size = 18) +
  theme(axis.text.x = element_blank(),
        axis.ticks.x = element_blank(),
        strip.background = element_rect(fill = "white", color = "white")) +
  geom_bar(stat = "identity", position = "stack", width = 1) +
  scale_fill_manual(values = palette, name = "Treatment") +
  scale_y_continuous(n.breaks = 10) +
  ggtitle("Cd8+ WPRE-") +
  ylab("% Cells") +
  xlab("") +
  facet_grid(.~Group, scales = "free_x")

ggsave(filename = paste(wdir, "Full_WPRE_CART_Neg_TILs_barplot.png", sep = "/"), width = 7, height = 7)