# Source code for producing figure 2H

library(ggplot2)
library(dplyr)
library(RColorBrewer)
library(UCell)
library(scGate)
library(viridis)
library(stringr)   
library(patchwork)
library(dittoSeq)
library(gridExtra)
library(Seurat)
library(openxlsx)
library(ggalluvial)
library(gridExtra)

obj <- readRDS(file = "Fig2H_dataset.rds")

colpalette <- c("#A6CEE3","#B2DF8A","#FDBF6F","#E78AC3","#FFFF99","#E31A1C","darkgrey","#C4A484")

umap_annot <- dittoSeq::dittoDimPlot(object = obj, reduction.use = "umap.harmony.pt.tp",
                                     var = "RefinedAnnotation", opacity = 0.8, do.label = T,labels.size = 4,
                                     labels.highlight = F, labels.repel = T, size = 1,order = "decreasing") + 
  scale_color_manual(values = colpalette) + 
  xlab("UMAP_1") +
  ylab(label = "UMAP_2") +
  theme(legend.position = "none",
        legend.title = element_blank(),
        legend.key.size = unit(0.001, "in"),
        legend.key.spacing = unit(0,"in"),
        plot.title = element_blank(),
        axis.text = element_text(size = 16)
  )

alluvial_list <- list()

for(pt in c("MPSI001","MPSI004")){
  
  meta_tmp <- subset.data.frame(obj@meta.data, subset = Patient == pt)
  g <- as.data.frame(prop.table(table(meta_tmp[["Timepoint"]], meta_tmp[["RefinedAnnotation"]]), margin = 1)*100)
  g$Var1 <- factor(g$Var1, levels = c("dp","day4","day7"))
  bplot_list[[pt]] <- ggplot(data = g, mapping = aes(Var1, Freq, fill = Var2, label = Var2)) + 
    scale_fill_manual(values = colpalette) + 
    ggtitle(label = "Classification distribution within Timepoint", subtitle = paste0("Patient ", pt)) + 
    geom_bar(stat = "identity", width = 0.5, color = "black") + 
    theme(legend.position = "right",
          plot.title = element_text(size = 16, hjust = 0.5, face = "plain"),
          plot.subtitle = element_text(size = 16, hjust = 0.5, face = "plain"),
          panel.background = element_rect(fill = "white"),
          axis.line = element_line(colour = "black", size = 0.2),
          axis.title = element_blank(), 
          strip.text.x = element_text(size = 12),
          plot.margin = margin(0.2, 0.2,0.2,0.2, unit = "cm"),
          axis.text = element_text(size = 16, colour = "black"),
          legend.title = element_blank(),
          legend.text = element_text(size = 8),
          legend.spacing.y = unit(0.1, 'cm'),
          legend.spacing.x = unit(0.1, 'cm'),
          legend.key.size = unit(0.5, 'cm'), 
          legend.key = element_rect(fill = "white"),
          legend.background = element_blank(),
          #legend.key.height = unit(0, 'cm'),
          legend.key.width = unit(0.5, 'cm'),
          legend.box.margin=margin(-10,0,0,0)
    ) 
  alluvial_list[[pt]] <- ggplot(data = g, aes(x = Var1, y = Freq, alluvium = Var2)) +
    geom_alluvium(aes(fill = Var2),
                  alpha = .75, decreasing = FALSE, width = 1/3) +
    geom_stratum(aes(stratum = Var2, fill = Var2), alpha = .50, decreasing = FALSE, width = 1/3) + 
    scale_fill_manual(values = alpha(colour = colpalette3, alpha = 0.50)) + 
    theme(legend.position = "right", plot.title = element_blank(),
          axis.line = element_line(color = "black"),
          axis.title = element_blank(),
          panel.background = element_rect(fill = "white"),
          axis.text = element_text(size = 16))
}

alluvial_mps1 <- alluvial_list$MPSI001 + NoLegend() + theme(plot.title = element_blank())
alluvial_mps4 <- alluvial_list$MPSI004 + NoLegend() + theme(plot.title = element_blank())

pdf(file = "Fig2H_panel.pdf", width = 9, height = 3)
grid.arrange(umap_annot, alluvial_mps1, alluvial_mps4, ncol = 3)
dev.off()