#####Fluvial
library(Seurat)
library(openxlsx)
library(readxl)
library(tidyr)
library("RColorBrewer")
library(dplyr)
library(reshape2)
library(fgsea)
library(ggalluvial)
library("gridExtra")

cluster_dir <- "/beegfs/scratch/ric.squadrito/ric.squadrito/90-935462466_scRNAseq_NotaroTAD_A018/Analysis_MM/results3"
wdir<- paste0(cluster_dir, "/Result3_MNOT")
ddir<- paste0(cluster_dir, "/Liver_Tumor_TnK/")
pdir<- paste0(cluster_dir, "/Result3_MNOT/Plots")

# Load the RDS
obs1 <- readRDS(paste0(ddir, "Liver_Tumor_TnK_final.rds"))

obs1@meta.data$Tissue <- obs1@meta.data$PROP.Source
obs1@meta.data$RNA_Group <- obs1@meta.data$PROP.Condition

obs1@meta.data$RNA_Group[obs1@meta.data$RNA_Group == "liOVA"] <- "OVA"
obs1@meta.data$RNA_Group[obs1@meta.data$RNA_Group == "liOVA_IL12"] <- "OVA.IL12"
obs1@meta.data$RNA_Group[obs1@meta.data$RNA_Group == "liOVA_IFNa"] <- "OVA.IFNa"
obs1@meta.data$RNA_Group[obs1@meta.data$RNA_Group == "Combo3x"] <- "Combo"
obs1@meta.data$RNA_Group <- factor(x = obs1@meta.data$RNA_Group, levels = c("OVA","OVA.IFNa","OVA.IL12","Combo"))


# Sample origin
obs1$Tube <- paste0(obs1$RNA_Group, obs1$hash.ID)
obs1@meta.data$Orig<-paste0(obs1@meta.data$Tissue, ".", obs1@meta.data$RNA_Group)

obs1@meta.data$RNA_Group <- factor(x = obs1@meta.data$RNA_Group, levels = c("OVA","OVA.IFNa","OVA.IL12","Combo"))



#################Create labels for all CD8 and CD4 T cells#############
obs1@meta.data$NewlablesG <- obs1@meta.data$NewlablesM
obs1@meta.data$NewlablesG[grep("CD4",obs1@meta.data$NewlablesM) ] <- "CD4 T cells"
obs1@meta.data$NewlablesG[grep("CD8",obs1@meta.data$NewlablesM)] <- "CD8 T cells"


###############create labels CD4 and CD8 clusters without naive####################
obs1@meta.data$NewlablesS <- obs1@meta.data$NewlablesM
indCd4<- grepl("CD4",obs1@meta.data$NewlablesM) & !grepl("Naive",obs1@meta.data$NewlablesM)
indCd8<- grepl("CD8",obs1@meta.data$NewlablesM) & !grepl("Naive",obs1@meta.data$NewlablesM)
obs1@meta.data$NewlablesS[indCd4] <- "CD4 T cells"
obs1@meta.data$NewlablesS[indCd8] <- "CD8 T cells"
obs1@meta.data$NewlablesST <- paste0(obs1@meta.data$Tissue, ".", obs1@meta.data$NewlablesS)


################################################
###Calculating hyperexpanded T cells#############
################################################
obs1$aa3.Tube <- paste0(obs1$cdr3_aa2,".",obs1$Tube)
obs1$aa3.Tube[is.na(obs1$cdr3_aa2)] <- NA
obs1@meta.data$ClonoFreq <- obs1@meta.data  %>% group_by(aa3.Tube)  %>% mutate(count = n()) %>% ungroup() %>%
  as.data.frame() %>% select (count)
obs1@meta.data$ClonoFreq[is.na(obs1$cdr3_aa2),] <- NA

obs1@meta.data$Clonetype  <- NA
obs1@meta.data$Clonetype[obs1@meta.data$ClonoFreq==1]  <- "unique"
obs1@meta.data$Clonetype[obs1@meta.data$ClonoFreq > 1 & obs1@meta.data$ClonoFreq <= 5]  <- "small"
obs1@meta.data$Clonetype[obs1@meta.data$ClonoFreq > 5 & obs1@meta.data$ClonoFreq <= 30]  <- "large"
obs1@meta.data$Clonetype[obs1@meta.data$ClonoFreq > 30]  <- "hyperexpanded"

obs1$Clonetype <- factor(x = obs1$Clonetype, levels = c("unique", "small", "large", "hyperexpanded")) #order clonotype variable



############################################
######FINDING Tetramer reactive TCRs##########
############################################
ind1 <- obs1$ADT.Tetramer
ind1 <- ind1[order(ind1)]
ind1 <- ind1[order(ind1)]

TetramerCount <- obs1@meta.data$ADT.Tetramer[!is.na(obs1@meta.data$ADT.Tetramer)]
TetramerCount <- TetramerCount[TetramerCount > 0]

obs1@meta.data$TetramerT <- "Other"
obs1@meta.data$TetramerT[obs1@meta.data$ADT.Tetramer>2] <- "OVA.TCR"

df1 <- data.frame(obs1$cdr3_aa2, obs1$TetramerT,obs1$ClonoFreq,obs1$Tissue)
df1$aa3.OVA <-  paste0(obs1$cdr3_aa2,".",obs1$TetramerT)
df1 <- na.omit(df1)
df1l <- df1[df1$obs1.Tissue=="Liver",]
df2 <- df1l %>% group_by(aa3.OVA) %>% mutate(count = n()) %>% ungroup() %>% as.data.frame()
df3 <- df2 %>% select(obs1.cdr3_aa2,obs1.TetramerT,obs1.ClonoFreq,count) %>% unique() %>%
  pivot_wider(names_from = obs1.TetramerT, values_from = count) %>% as.data.frame()
df3[is.na(df3)] <- 0
df3$Tumor <- df3[,2]-df3[,3]-df3[,4]
OVA.TCRs <- df3$obs1.cdr3_aa2[df3[,4]/df3[,3]>1]

df1 <- data.frame(obs1$cdr3_aa2, obs1$TetramerT,obs1$ClonoFreq,obs1$Tissue,obs1$Tube)
df1$aa3.Tube <-  paste0(obs1$cdr3_aa2,".",obs1$TetramerT,".",obs1$Tube)
df1 <- na.omit(df1)
df1l <- df1[df1$obs1.Tissue=="Liver",]
df2 <- df1l %>% group_by(aa3.Tube) %>% mutate(count = n()) %>% ungroup() %>% as.data.frame()
df3T <- df2 %>% select(obs1.cdr3_aa2,obs1.TetramerT,obs1.Tube,obs1.ClonoFreq,count) %>% unique() %>%
  pivot_wider(names_from = obs1.TetramerT, values_from = count) %>% as.data.frame()
df3T[is.na(df3T)] <- 0


######Indicating Tetramer OVA in obs1######
obs1$OVA.reactive <- ifelse(obs1$cdr3_aa2 %in% OVA.TCRs, "OVA", "Other")
obs1@meta.data$TetramerTS <- paste0(obs1@meta.data$OVA.reactive,".",obs1@meta.data$NewlablesS)


######TISSUE SHARED CLONOTYPES########
#Identify cells with the same Cdr3aa sequence in both tissues
df1<-data.frame(obs1@meta.data$cdr3_aa2,obs1@meta.data$hash.ID,obs1@meta.data$RNA_Group, obs1@meta.data$Tissue)
df2<-na.omit(df1)
df2$cdr3_mouse<-paste0(df2$obs1.meta.data.cdr3_aa2,".",df2$obs1.meta.data.hash.ID,".",df2$obs1.meta.data.RNA_Group)
shared_cdr3aa <- intersect(df2$cdr3_mouse[df2$obs1.meta.data.Tissue == "Liver"], df2$cdr3_mouse[df2$obs1.meta.data.Tissue == "Tumor"])
obs1$cdr3_mouse<-paste0(obs1$cdr3_aa2,".",obs1$hash.ID,".",obs1$RNA_Group)
obs1$Shared <- ifelse(obs1$cdr3_mouse %in% shared_cdr3aa, "Shared", "Unique")

# Create variables share tissue and shared OVA
obs1@meta.data$Sharedtissue<-paste0(obs1@meta.data$Tissue,".",obs1@meta.data$Shared)
obs1@meta.data$SharedOVA<-paste0(obs1@meta.data$Shared,".",obs1@meta.data$OVA.reactive)

####Alluvial plot#####
######Figure S3B-allOVA the same color########
df1 <- obs1@meta.data
df1 <- df1[df1$NewlablesG=="CD8 T cells",]
#df2<- na.omit(df1[,c(13,38,39,49)])
df2<- na.omit(df1[,c("cdr3_aa2","RNA_Group","Tissue","OVA.reactive")])
for (i in unique(df1$RNA_Group)){
  # i <- unique(df1$RNA_Group)[1]
  df3 <- df2[df2$RNA_Group==i&df2$Tissue=="Liver",]
  df4 <- df2[df2$RNA_Group==i&df2$Tissue=="Tumor",]
  df3 <- df3 %>% group_by(cdr3_aa2) %>% mutate(ClonoFreq = n()) %>% unique() %>% ungroup() %>% as.data.frame() 
  df4 <- df4 %>% group_by(cdr3_aa2) %>% mutate(ClonoFreq = n()) %>% unique() %>% ungroup() %>% as.data.frame()
  df3 <- df3[df3$ClonoFreq>2,]
  df4 <- df4[df4$ClonoFreq>2,]
  df3$ClonoFreq.norm <- sapply(df3$ClonoFreq,function(x){100*x/sum(df3$ClonoFreq)})
  df4$ClonoFreq.norm <- sapply(df4$ClonoFreq,function(x){100*x/sum(df4$ClonoFreq)})
  df6 <- rbind(df3,df4)
  df6 <- df6[order(df6$cdr3_aa2),]
  df6$col1[!duplicated(df6$cdr3_aa2)] <- "#CCCCCC"
  Dtests <- duplicated(df6$cdr3_aa2)|duplicated(df6$cdr3_aa2, fromLast=T)
  df6$col1 [df6$OVA.reactive=="OVA"] <- "#3399FF" #"#3399FF"
  df6$col1 [df6$OVA.reactive=="Other"&Dtests] <- "#CC6600"
  plotA <- ggplot(data = df6,
                  aes(x=Tissue, y=ClonoFreq.norm, alluvium=cdr3_aa2, stratum=cdr3_aa2)) + 
    geom_flow(aes(fill=col1), decreasing=F)+
    geom_stratum(aes(fill=col1, size=0.2), decreasing=F) +
    scale_fill_manual(values=setNames(unique(df6$col), unique(df6$col)))+
    theme_bw() + theme(legend.position="none")+
    scale_size_identity()+
    ggtitle(i)
  nameR <- paste0(i, ".plot") 
  assign(nameR, plotA)}

pdf(paste0(pdir, "/Fluvial.Plot.18.CD8.norm.2clones.OVAblue_FigS3B.pdf"), width=4, height=4)
grid.arrange(Combo.plot,OVA.IFNa.plot,OVA.IL12.plot,OVA.plot, ncol=2, nrow=2)
dev.off()

######CD4 T cells########
####Figure 3H - normalized test CD4#######
df1 <- obs1@meta.data
df1 <- df1[df1$NewlablesG=="CD4 T cells",]
#df2<- na.omit(df1[,c(13,38,39,49)])
df2<- na.omit(df1[,c("cdr3_aa2","RNA_Group","Tissue","OVA.reactive")])
for (i in unique(df1$RNA_Group)){
 # i <- unique(df1$RNA_Group)[1]
  df3 <- df2[df2$RNA_Group==i&df2$Tissue=="Liver",]
  df4 <- df2[df2$RNA_Group==i&df2$Tissue=="Tumor",]
  df3 <- df3 %>%  group_by(cdr3_aa2) %>% mutate(ClonoFreq = n()) %>% unique() %>% ungroup() %>% as.data.frame() 
  df4 <- df4 %>%  group_by(cdr3_aa2) %>% mutate(ClonoFreq = n()) %>% unique() %>% ungroup() %>% as.data.frame()
  df3 <- df3[df3$ClonoFreq>1,]
  df4 <- df4[df4$ClonoFreq>1,]
  df3$ClonoFreq.norm <- sapply(df3$ClonoFreq,function(x){100*x/sum(df3$ClonoFreq)})
  df4$ClonoFreq.norm <- sapply(df4$ClonoFreq,function(x){100*x/sum(df4$ClonoFreq)})
  df6 <- rbind(df3,df4)
  df6 <- df6[order(df6$cdr3_aa2),]
  df6$col1[!duplicated(df6$cdr3_aa2)] <- "#CCCCCC"
  Dtests <- duplicated(df6$cdr3_aa2)|duplicated(df6$cdr3_aa2, fromLast=T)
  df6$col1 [df6$OVA.reactive=="OVA"&Dtests] <- "#CC6600" #"#3399FF"
  df6$col1 [df6$OVA.reactive=="Other"&Dtests] <- "#CC6600"
  plotA <- ggplot(data = df6,
                  aes(x = Tissue, y=ClonoFreq.norm, alluvium = cdr3_aa2,stratum=cdr3_aa2)) + 
    geom_flow(aes(fill = col1),decreasing = F)+
    geom_stratum(aes(fill = col1,size=0.2),decreasing = F) +
    scale_fill_manual(values = setNames(unique(df6$col), unique(df6$col)))+
    theme_bw() + theme(legend.position = "none")+
    scale_size_identity()+
    ggtitle(i)
  nameR <- paste0(i,".plot") 
  assign(nameR,plotA)}


pdf(paste0(pdir, "/Fluvial.Plot.18.CD4.norm.2clones_Fig3H.pdf"), width=4, height=4)
grid.arrange(Combo.plot,OVA.IFNa.plot,OVA.IL12.plot,OVA.plot, ncol=2, nrow=2)
dev.off()