#Precision Recall Analysis
library(tidyverse)
library(ggpubr)
library("viridis")
library(PRROC)

class1_scores=read.csv("Supplementary Table 8",stringsAsFactors = F)

class1_viper_wScores_MHCB=class1_scores %>% filter(Assay.Type.Category=="MHC Binding",!grepl("x-ray",Method))
class1_viper_wScores_MHCB$pmhc=paste(class1_viper_wScores_MHCB$Epitope.Sequence,class1_viper_wScores_MHCB$MHC.Allele.Name,sep = "_")
dups=unique(class1_viper_wScores_MHCB$pmhc)

### use strongest response in duplicated assays
reclass=NULL
for (i in dups){
  res_i=NULL
  dup_i=class1_viper_wScores_MHCB %>% filter(pmhc==i) %>% select(Assay.Result)
  tab_i=class1_viper_wScores_MHCB %>% filter(pmhc==i)
  if ("Positive-High" %in% dup_i$Assay.Result) {res_i="Positive-High"
  } else if ("Positive-Intermediate" %in% dup_i$Assay.Result) {res_i="Positive-Intermediate"
  } else if ("Positive-Low" %in% dup_i$Assay.Result) {res_i="Positive-Low"
  } else if ("Positive" %in% dup_i$Assay.Result) {res_i="Positive"
  } else if ("Negative" %in% dup_i$Assay.Result) {res_i="Negative"
  }
  reclass=rbind(reclass,data.frame(tab_i,correct=rep(res_i,nrow(dup_i))))
}

class1_viper_wScores_MHCB=reclass

class1_viper_wScores_MHCB_unique=class1_viper_wScores_MHCB %>% select(Epitope.Sequence,MHC.Allele.Name,correct,Neon.Percent.Rank)
class1_viper_wScores_MHCB_unique=unique(class1_viper_wScores_MHCB_unique)

class1_viper_wScores_MHCB_unique$correct=factor(class1_viper_wScores_MHCB_unique$correct,levels=c("Negative","Positive","Positive-Low","Positive-Intermediate","Positive-High"))

comparisons_to_neg=list(c("Negative","Positive"),c("Negative","Positive-Low"),c("Negative","Positive-Intermediate"),c("Negative","Positive-High"))


df = data.frame(predictions=class1_viper_wScores_MHCB_unique$Neon.Percent.Rank,labels=as.character(class1_viper_wScores_MHCB_unique$correct),stringsAsFactors = F)
df$labels[df$labels!="Negative"]=1
df$labels[df$labels=="Negative"]=0

prRes=pr.curve(100-df$predictions,weights.class0 = as.numeric(df$labels),curve = T)
forPlot=data.frame(prRes$curve)
forPlot$X3=log10(100-forPlot$X3)


ggplot(forPlot,aes(x=X1,y=X3,color=X2))+geom_point(size=1.5)+scale_x_continuous(limits = c(0,1))+scale_color_viridis()+theme_bw()+theme_classic()+ylab("%Rank Threshold")+xlab("Recall")+geom_text(x=0.15, y=0.5, label=paste0("No Skill AUC:",round(sum(df$labels=="1")/length(df$labels),digits = 2)),color="black")+geom_text(x=0.15, y=0.17, label=paste0("Curve AUC:",round(prRes$auc.integral,digits = 2)),color="black")+labs(color="Precision")+geom_hline(yintercept = -1,linetype="dashed",color="red")


###class2
class2=read.csv("C:/Users/aporan/Desktop/Projects/EpiSelection/nCoV19/Matt_Dewi/revision/updated supp tables/SuppTable_9_ViPR_classii_percent-rank.csv",stringsAsFactors = F)
class2=as_tibble(read.table("vipr_scores.tab",header=T,sep="\t",stringsAsFactors = F))
class2$Assay.Result=factor(class2$Assay.Result,levels=c("Negative","Positive","Positive-Low","Positive-Intermediate","Positive-High"))
class2$MHC.Allele.Name=gsub("\\*","",class2$MHC.Allele.Name)

class2_MHCB=class2 %>% filter(Assay.Type.Category=="MHC Binding",!grepl("x-ray",Method))
class2_MHCB$pmhc=paste(class2_MHCB$Epitope.Sequence,class2_MHCB$MHC.Allele.Name,sep = "_")
class2_MHCB_unique=class2_MHCB %>% select(Epitope.Sequence,MHC.Allele.Name,Assay.Result,Neon.Percent.Rank)
class2_MHCB_unique=unique(class2_MHCB_unique)

class2_MHCB_unique$Assay.Result2=class2_MHCB_unique$Assay.Result %>% as.character()
class2_MHCB_unique$Assay.Result2[class2_MHCB_unique$Assay.Result2 %in% c("Negative","Positive-Low")]="Negative-Positive-Low"
class2_MHCB_unique$Assay.Result2=factor(class2_MHCB_unique$Assay.Result2,levels=c("Negative-Positive-Low","Positive-Intermediate","Positive-High"))

comparisons_to_neg=list(c("Negative-Positive-Low","Positive-Intermediate"),c("Negative-Positive-Low","Positive-High"))


df = data.frame(predictions=class2_MHCB_unique$Neon.Percent.Rank,labels=as.character(class2_MHCB_unique$Assay.Result2),stringsAsFactors = F)
df$labels[df$labels!="Negative-Positive-Low"]=1
df$labels[df$labels=="Negative-Positive-Low"]=0


prRes=pr.curve(df$predictions,weights.class0 = as.numeric(df$labels),curve = T)
forPlot=data.frame(prRes$curve)
forPlot$X3=log10(100-forPlot$X3)

ggplot(forPlot,aes(x=X1,y=X3,color=X2))+geom_point(size=1.5)+scale_x_continuous(limits = c(0,1))+scale_color_viridis()+theme_bw()+theme_classic()+ylab("%Rank Threshold")+xlab("Recall")+geom_text(x=0.5, y=0.5, label=paste0("No Skill AUC:",round(sum(df$labels=="1")/length(df$labels),digits = 2)),color="black")+geom_text(x=0.5, y=0.17, label=paste0("Curve AUC:",round(prRes$auc.integral,digits = 2)),color="black")+labs(color="Precision")+geom_hline(yintercept = 0.75,linetype="dashed",color="red")



##### Code for Figure 5
data = read.csv("Supplementary Table 10")

data$Protein = fct_reorder(.f=data$Protein,.x=data$Abundance,.fun=median)

q = ggplot(data,aes(x=Protein,y=Abundance,fill=Study,group=Study)) + 
  geom_bar(stat='identity',position='dodge',color='black') +
  geom_hline(yintercept=0) +
  theme_bw() + theme_classic()+
  facet_grid(. ~ Protein,scales='free',space='free') +
  ylim(0,1) + scale_fill_manual(values = viridis(3)) +
  theme(text = element_text(size=20)) +
  ylab('Relative Abundance')

plot(q)

### Table 1 data
##From the ViPR class 1 table
scores_SARS2=read.csv("Supplementary Table 4.csv",header=T,stringsAsFactors = F)
scores_SARS2= scores_SARS2 %>% select(peptide, protein,alleles)

alleles_list=strsplit(scores_SARS2$alleles,split = ";")
proteins_list=strsplit(scores_SARS2$protein,split = ";")
proteins_list=sapply(proteins_list,tail,1)

scores_SARS2=data.frame(peptide=rep(scores_SARS2$peptide,sapply(alleles_list,length)),
                             allele=unlist(alleles_list),
                             protein=(rep(proteins_list,sapply(alleles_list,length))))

scores_SARS2$pmhc=paste(scores_SARS2$peptide,scores_SARS2$allele,sep="_")

hits_1per_counts=data.frame(table(scores_SARS2$protein))
names(hits_1per_counts)=c("protein","count")
hits_1per_counts$count[hits_1per_counts$protein=="orf1a polyprotein"]=length(setdiff(scores_SARS2$pmhc[scores_SARS2$protein=="orf1a polyprotein"],scores_SARS2$pmhc[scores_SARS2$protein=="orf1ab polyprotein"]))

hits_1per_counts$percent=round(100*hits_1per_counts$count/sum(hits_1per_counts$count),2)

class1=as_tibble(read.table("Supplementary Table 8.csv",header=T,sep=",",stringsAsFactors = F))
class1$Assay.Result=factor(class1$Assay.Result,levels=c("Negative","Positive","Positive-Low","Positive-Intermediate","Positive-High"))
class1$MHC.Allele.Name=gsub("\\*","",class1$MHC.Allele.Name)
class1_viper_wScores=merge(class1,scores_SARS2,by.x=c("Epitope.Sequence","MHC.Allele.Name"),by.y=c("peptide","allele"))

dups=unique(class1_viper_wScores$pmhc)

reclass=NULL
for (i in dups){
    res_i=NULL
    dup_i=class1_viper_wScores %>% filter(pmhc==i) %>% select(Assay.Result)
    tab_i=class1_viper_wScores %>% filter(pmhc==i)
    if ("Positive-High" %in% dup_i$Assay.Result) {res_i="Positive-High"
    } else if ("Positive-Intermediate" %in% dup_i$Assay.Result) {res_i="Positive-Intermediate"
    } else if ("Positive-Low" %in% dup_i$Assay.Result) {res_i="Positive-Low"
    } else if ("Positive" %in% dup_i$Assay.Result) {res_i="Positive"
    } else if ("Negative" %in% dup_i$Assay.Result) {res_i="Negative"
    }
    reclass=rbind(reclass,data.frame(tab_i,correct=rep(res_i,nrow(dup_i))))
}

class1_viper_wScores=reclass

class1_viper_wScores_unique=class1_viper_wScores %>% filter(Assay.Type.Category=="MHC Binding") %>% select(Epitope.Sequence,MHC.Allele.Name,correct,Neon.Percent.Rank,protein) 
class1_viper_wScores_unique=unique(class1_viper_wScores_unique)

class1_viper_wScores_unique$correct=factor(class1_viper_wScores_unique$correct,levels=c("Negative","Positive","Positive-Low","Positive-Intermediate","Positive-High"))
class1_viper_wScores_unique$pmhc=paste(class1_viper_wScores_unique$Epitope.Sequence,class1_viper_wScores_unique$MHC.Allele.Name,sep="_")
unique_to_1a=setdiff(scores_SARS2$pmhc[scores_SARS2$protein=="orf1a polyprotein"],scores_SARS2$pmhc[scores_SARS2$protein=="orf1ab polyprotein"])
unique_to_1a=nrow(class1_viper_wScores_unique %>% filter(protein=="orf1a polyprotein" & pmhc %in% unique_to_1a))

hits_1per_counts_SARS2=data.frame(table(class1_viper_wScores_unique$protein))
names(hits_1per_counts_SARS2)=c("protein","count")
hits_1per_counts_SARS2$count[hits_1per_counts_SARS2$protein=="orf1a polyprotein"]=unique_to_1a

hits_1per_counts_SARS2=cbind(hits_1per_counts_SARS2,matrix(table(class1_viper_wScores_unique$protein,class1_viper_wScores_unique$correct!="Negative"),ncol = 2))
names(hits_1per_counts_SARS2)=c("protein","count","Negative","Positive")
hits_1per_counts_SARS2$Negative[hits_1per_counts_SARS2$protein=="orf1a polyprotein"]=unique_to_1a
hits_1per_counts_SARS2$Positive[hits_1per_counts_SARS2$protein=="orf1a polyprotein"]=unique_to_1a

hits_1per_counts_SARS2$percent=round(100*hits_1per_counts_SARS2$Positive/hits_1per_counts_SARS2$count,2)

hits_1per_counts_full=merge(hits_1per_counts,hits_1per_counts_SARS2,by="protein",all=T)
hits_1per_counts_full[is.na(hits_1per_counts_full)]=0
names(hits_1per_counts_full)=c("protein","percent-rank 1 Epitope Count","fraction of all epitopes","Found in ViPR","Assay:Negative","Assay:Positive","percent positive")