

rm(list = ls())


library(rstudioapi)
library(latticeExtra) 
require(gridExtra)
library(RColorBrewer)


###########################
# Reading the data
###########################

setwd(dirname(getActiveDocumentContext()$path ))
data <- read.csv("Figure 4-source data.csv", header=TRUE,stringsAsFactors=FALSE)



IDs = unique(data[,"ID"])
ytypes = unique(data[,"ytype_def"])
pchID = c(21,22,23,24,25,
          21,22,23,24,25,
          21,22,23,24,25,21,22,23,24,25,21,22)

coloresID = c("steelblue4","slateblue3","navy","lightsteelblue4","dodgerblue2", 
              "salmon4","indianred3","brown2","coral1","chocolate2",
              "aquamarine3","darkolivegreen3","cadetblue4","chartreuse3","forestgreen","seagreen3",
              "palegreen3","springgreen4","darkseagreen4","green4","greenyellow","lightgreen")



######################################
# Reading Reff, fp and D from simulations
######################################

Results <- read.csv("Figure 6-source data.csv", header=TRUE,stringsAsFactors=FALSE)


animals=rev(c(6:9,11:16))#c(6,8:9,11:14,16)#


W = 5
pb = 1
VolB = 60*W*1e3

Rt_vals= c(0,0,0,0,0, 
           6.45e6,4.14e6,2.08e6,2.39e6,6.28e6,
           1.06e7,1.6e7,6.7e6,1.48e7,1.2e7,6.6e7,
           6.2e6,8.25e6,2.3e6,4.82e6,6.6e6,6.6e6)

Results <- read.csv("Figure 6-source data.csv", header=TRUE,stringsAsFactors=FALSE)


##########################
# Plotting Figure 6 C
##########################
Results_i = Results[which(Results$ID==IDs[8]),]

D = Results_i$D[sort(Results_i$Reff,index.return=TRUE)$ix]
R = Results_i$Pr[sort(Results_i$Reff,index.return=TRUE)$ix]

D_a = Rt_vals*5
P_a = 10^6.758
Te_a = D_a/P_a
fp_a = c(rep(0,10),unique(as.numeric(data[,11]))[c(2:7)])/100

R0  = sort(Results_i$Reff,index.return=TRUE)$x
fp  = (Results_i$fp)[sort(Results_i$Reff,index.return=TRUE)$ix]
Te  = D/R
RT = (Results_i$RT)[sort(Results_i$Reff,index.return=TRUE)$ix]
fp_DR = data.frame(cbind(fp,Te,R0))


fp_DR$Te=(log10(fp_DR$Te))

col.l <- rev(colorRampPalette(brewer.pal(10, "RdYlBu"))(299))
colors = c(seq(min(fp_DR$R0),0.95,length=130),
           seq(0.96,1.12,length=20),
           seq(1.13,1.25,length=20),
           seq(1.26,max(fp_DR$R0),length=130))

colorplot1 <- levelplot(
  R0~Te*fp,#data.matrix(t(fp_DR_hm)),
  data=fp_DR,
  ylab = list(expression("%"~italic(f[p])),cex=1.4),
  xlab = list(expression(paste(italic(D),":",italic(P[r]))),cex=1.4),
  col.regions=col.l,
  at=colors,
  colorkey = FALSE, region = TRUE,
  scales = list( x=list(at = seq(from=0, to=4, length.out=5),
                        labels=  expression(bold("1"),bold("10"),bold("10"^"2"),bold("10"^"3"),bold("10"^"4")),
                        cex=1.3),
                 y=list(at = seq(from=0, to=1, length.out = 5),
                        labels=as.character(seq(from=0, to=100, length.out =5)),
                        cex=1.3,font=2)
                 ),
)
print(colorplot1)

trellis.focus("panel", 1, 1, highlight=FALSE)

lpoints(log10(Te_a[animals]),fp_a[animals],
        col="black",
        margin=FALSE,
        fill =coloresID[animals],
        pch=pchID[animals],
        lwd=1,cex=1.4)

ltext(2.5,0.88,"Control",font=3,cex=1.2)#,col="gray50")
ltext(0.1,0.15,"No control",font=3,cex=1.2,col="white")

ltext(1,0.46,expression(paste("Actual ",italic(f[p])," and ",italic(D))),cex=0.9)
ltext(1,0.40,expression(paste("and assuming ",italic(P[r]),"=",10^7)),cex=0.9)

trellis.unfocus()



##################################
# Plotting heatmap for all animals
###################################

animals=c(6:9,11:16)
i=1

colorplot=list()


RT_vals= c()
Fp_vals= c()
DP_vals= c()

for(ID in IDs[animals]){#[c(6:9,11:16)]){
  
  ii = which(IDs %in% ID)
  fp_a = as.numeric(data[which(ID==data[,"ID"] ),11])[1]/100
  D_a = Rt_vals[ii]*5
  P_a = 10^6.758
  Te_a = D_a/P_a
  
  
  
  Results_i = Results[which(Results$ID==ID),]
  
  D = Results_i$D[sort(Results_i$Reff,index.return=TRUE)$ix]
  R = Results_i$Pr[sort(Results_i$Reff,index.return=TRUE)$ix]

  R0  = sort(Results_i$Reff,index.return=TRUE)$x
  fp  = (Results_i$fp)[sort(Results_i$Reff,index.return=TRUE)$ix]
  Te  = D/R
  RT = (Results_i$RT)[sort(Results_i$Reff,index.return=TRUE)$ix]
  fp_DR = data.frame(cbind(fp,Te,R0))
  
  index_control = which(fp_DR$R0<1)
  
  RT_vals= c(RT_vals,RT[1])
  Fp_vals= c(Fp_vals,min(fp_DR$fp[index_control]))
  DP_vals= c(DP_vals,min(fp_DR$Te[index_control]))

  
  fp_DR$Te=(log10(fp_DR$Te))

  col.l <- rev(colorRampPalette(brewer.pal(10, "RdYlBu"))(299))
  colors = c(seq(min(fp_DR$R0),0.95,length=130),
             seq(0.96,1.12,length=20),
             seq(1.13,1.25,length=20),
             seq(1.26,max(fp_DR$R0),length=130))
  if(i==1){
    colorplot1 <- levelplot(
      R0~Te*fp,
      data=fp_DR,
      ylab = list(expression("%"~italic(f[p])),cex=1.4),
      xlab = list(expression(paste(italic(D),":",italic(P[r]))),cex=1.4),
      col.regions=col.l,
      at=colors,
      colorkey = FALSE, region = TRUE,

      scales = list( x=list(at = seq(from=0, to=4, length.out=5),
                            labels=  expression(bold("1"),bold("10"),bold("10"^"2"),bold("10"^"3"),bold("10"^"4")),
                            cex=1.3),
                     y=list(at = seq(from=0, to=1, length.out = 5),
                            labels=as.character(seq(from=0, to=100, length.out =5)),
                            cex=1.3,font=2))
    )
    print(colorplot1)
    
    index_control = which(fp_DR$R0<1)
    
    trellis.focus("panel", 1, 1, highlight=FALSE)
    lpoints(log10(Te_a),fp_a,
            col="black",
            margin=FALSE,
           fill =coloresID[animals[i]],
           pch=pchID[animals[i]],
           lwd=1,cex=1.4)
    ltext(2.5,0.95,"Control",font=3,cex=1.2)#,col="gray50")
    ltext(0.1,0.15,"No control",font=3,cex=1.2,col="white")
    ltext(0.1,0.4,ID,font=2,cex=1.0,col="white")
    ltext(0.1,0.3,as.expression(bquote(R[T]==.(round(RT,digits=1)))),font=2,cex=1.0,col="white")
    
    trellis.unfocus()
   
    
   }
  else{
    colorplot1 <- levelplot(
      R0~Te*fp,
      data=fp_DR,
      ylab = list(expression("%"~italic(f[p])),cex=1.4),
      xlab = list(expression(paste(italic(D),":",italic(P[r]))),cex=1.4),
      col.regions=col.l,
      at=colors,
      colorkey = FALSE, region = TRUE,

      scales = list( x=list(at = seq(from=0, to=4, length.out=5),
                            labels=  expression(bold("1"),bold("10"),bold("10"^"2"),bold("10"^"3"),bold("10"^"4")),
                            cex=1.3),
                     y=list(at = seq(from=0, to=1, length.out = 5),
                            labels=as.character(seq(from=0, to=100, length.out =5)),
                            cex=1.3,font=2))
    )
    print(colorplot1)
    trellis.focus("panel", 1, 1, highlight=FALSE)
    lpoints(log10(Te_a),fp_a,
            col="black",
            fill =coloresID[animals[i]],
            pch=pchID[animals[i]],
            lwd=1,cex=1.4)
    ltext(2.5,0.95,"Control",font=3,cex=1.2)#,col="gray50")
    ltext(0.1,0.15,"No control",font=3,cex=1.2,col="white")
    ltext(0.1,0.4,ID,font=2,cex=1.0,col="white")
    ltext(0.1,0.3,as.expression(bquote(R[T]==.(round(RT,digits=1)))),font=2,cex=1.0,col="white")
    trellis.unfocus()
    
    
    print(c(D_a/(P_a),fp_a))
  }
  i=i+1
}


##################################
# Plotting Rt vs min fp for control
##################################
animals=c(6:9,11:16)#

k=1


for(ID in IDs[animals]){
  
  i = which(IDs %in% ID)
  ####
  
  if(k==1){
    plot(RT_vals[k],Fp_vals[k]*100,type ="p",log="x",
         ylab=expression("min  %"~italic(f[p])~"  for viral control"),
         xlab=expression(italic(R[T])),cex=1.5,
         cex.main=1.4,cex.axis=1.4,cex.lab=1.4,
         xlim=c(4,18),ylim =c(0.7,0.97)*100,
         xaxt = 'n',
         yaxt = 'n',
         lwd=1,bty="n",
         col =coloresID[i],
         pch=pchID[i]
    )
    
 
    points(RT_vals[k],Fp_vals[k]*100,
           bg =coloresID[i],#pch=19,
           pch=pchID[i],
           lwd=1,cex=1.6)
    

    axis(2,at=seq(70,95,by = 5),labels=seq(70,95,by = 5),cex.axis=1.4,lwd=2,font=2)
    axis(1,at=c(4,6,8,11,16),labels=c(4,6,8,11,16),cex.axis=1.4,lwd=2,font=2)

  }
  else{
    
   points(RT_vals[k],Fp_vals[k]*100,
           bg =coloresID[i],#pch=19,
           pch=pchID[i],
           lwd=1,cex=1.6)
    
  }

  k=k+1
}

# Correrlation test for min fp vs RT

R00test = cor.test(RT_vals,Fp_vals,method="pearson")
text(10,75.1,paste("p =",round(R00test$p.value,digits=4),"\nr =",round(R00test$estimate,digits=2)),cex=1.2,pos=4)

text(10,85,expression(paste("(",italic(D),":",italic(P[r]),")>",10^2)),cex=1.4,pos=4)

