# Dean, Halloran, Longini
# 2016
# Code accompanying "Design of Vaccine Trials during Outbreaks with and without a Delayed Vaccination Comparator"


######################################
library(ggplot2)
library(scales)
library(gridExtra)


######################################
# FUNCTIONS

# product function (put in infection time + incubation, returns product of terms)
product_u <- function(ui,ti,R,sj,VE0,shape,scale,hz_shape,hz_value1,hz_value2){
  
  # density for that incubation time 
  term1 <- dgamma(ui,shape=shape,scale=scale)
  
  # put in infection time wi=ti-ui, returns hazard at that time
  wi <- ti-ui
  if(hz_shape=="constant"){term2 <- hz_value1}
  # code as linear function decreasing from max [1] to 0 over period [2]
  else if(hz_shape=="decreasing"){term2 <- max(hz_value1-hz_value1/hz_value2*wi,0)}
  # code as linear function increasing from 0 to max [1] over period [2]
  else if(hz_shape=="increasing"){term2 <- max(hz_value1/hz_value2*wi,0)}
  
  # beta function (put in infection time, returns vaccine effect at that time)
  if(wi < sj){ term3 <- 1 }
  else if(wi >= sj & wi < sj + R){ term3 <- 1-(wi-sj)/R*(VE0) } # can change
  else if(wi >= sj + R){ term3 <- 1-VE0 }

  # slow growth, first 10% efficacy in R/2 time, last 90% efficacy in R/2 time
  #if(wi < sj){ term3 <- 1 }
  #else if(wi >= sj & wi < sj + R/2){ term3 <- 1-(wi-sj)*(2*0.10/R)*VE0 }
  #else if(wi >= sj + R/2 & wi < sj + R){ term3 <- VE0*.9 -(wi-sj-R/2)*(2*0.90/R)*VE0 }
  #else if(wi >= sj + R){ term3 <- 1-VE0 }
  
  # fast growth, first 90% efficacy in R/2 time, last 90% efficacy in R/2 time
  #if(wi < sj){ term3 <- 1 }
  #else if(wi >= sj & wi < sj + R/2){ term3 <- 1-(wi-sj)*(2*0.90/R)*VE0 }
  #else if(wi >= sj + R/2 & wi < sj + R){ term3 <- VE0*.1 -(wi-sj-R/2)*(2*0.10/R)*VE0 }
  #else if(wi >= sj + R){ term3 <- 1-VE0 }
  
  return(term1*term2*term3)
}
product_u <- Vectorize(product_u)


# lambda_T function (put in symptom time, returns group hazard at that time)
lambda_T <- function(ti,R,sj,VE0,shape,scale,hz_shape,hz_value1,hz_value2){
  # integrate over all possible incubation periods
  return(integrate(f=product_u,lower=0,upper=qgamma(0.99999,shape=shape,scale=scale),
              ti=ti,R=R,sj=sj,VE0=VE0,
              shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,
              rel.tol=1.e-10)$value)
}
lambda_T <- Vectorize(lambda_T)
#lambda_T(ti=ti,R=R,sj=s1,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)
  

# apparent VE function (put in symptom time, returns 1-hazard ratio at that time)
VE_T <- function(ti,R,s1,s0,VE0,shape,scale,hz_shape,hz_value1,hz_value2){
  # immediate arm
  lambda_1t <- lambda_T(ti=ti,R=R,sj=s1,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)
  # delayed (or control) arm
  lambda_0t <- lambda_T(ti=ti,R=R,sj=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)
  return(1-lambda_1t/lambda_0t)
}
VE_T <- Vectorize(VE_T)
#VE_T(ti=ti,R=R,s1=s1,s0=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)


# VE_D function (analysis delay d and width c and returns approximated apparent VE for that design)
VE_D <- function(d,c,R,s1,s0,VE0,shape,scale,hz_shape,hz_value1,hz_value2){
  num <- integrate(f=lambda_T,lower=d,upper=d+c,R=R,sj=s1,VE0=VE0,
                   shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,
                   rel.tol=1.e-7)$value/c
  den <- integrate(f=lambda_T,lower=d,upper=d+c,R=R,sj=s0,VE0=VE0,
                   shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,
                   rel.tol=1.e-7)$value/c
  return(1-num/den)
}
VE_D <- Vectorize(VE_D)
#VE_D(d=d,c=21,R=R,s1=s1,s0=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)


# power function given sample size n (per arm)
# requires VE_D as estimated via integration
power_D <- function(d,c,R,s1,s0,VE0,shape,scale,hz_shape,hz_value1,hz_value2,VE_Dobs,n){
  # number of events in immediate arm
  m1 <- n*(exp(-integrate(f=lambda_T,lower=0,upper=d,R=R,sj=s1,VE0=VE0,
                            shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,rel.tol=1.e-7)$value)
          )*(1-exp(-integrate(f=lambda_T,lower=d,upper=d+c,R=R,sj=s1,VE0=VE0,
                              shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,
                              rel.tol=1.e-7)$value))
  # number of events in delayed arm
  m0 <- n*(exp(-integrate(f=lambda_T,lower=0,upper=d,R=R,sj=s0,VE0=VE0,
                            shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,rel.tol=1.e-7)$value)
          )*(1-exp(-integrate(f=lambda_T,lower=d,upper=d+c,R=R,sj=s0,VE0=VE0,
                              shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,
                              rel.tol=1.e-7)$value))
  # power
  B <- pnorm(sqrt(m0+m1)*abs(VE_Dobs)/(2-VE_Dobs)-1.96)
  return(B)
}
#power_D(d=d,c=c,R=R,s1=s1,s0=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,VE_Dobs=VE_Dobs,n=n)


# simplified function for plotting that uses existing lambda_obs vectors (faster computations)
VE_D_simple <- function(d,c,ti,lambda_1T_obs,lambda_0T_obs){
  num <- mean(lambda_1T_obs[ti>=d & ti<(d+c)])
  den <- mean(lambda_0T_obs[ti>=d & ti<(d+c)])
  return(1-num/den)
}
#VE_D_simple(d=d,c=c,ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs)


# simplifed function for plotting power for sample size n per arm (faster computations)
power_D_simple <- function(d,c,ti,VE_Tobs,lambda_1T_obs,lambda_0T_obs,n){
  # predicted VE
  VE_D <- VE_D_simple(d=d,c=c,ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs)
  # account for cases necessarily excluded that occur before d
  if(d>0){
    # number of events in immediate arm
    m1 <- n*(exp(-d*mean(lambda_1T_obs[ti>=0 & ti<(d)])))*(1-exp(-c*mean(lambda_1T_obs[ti>=d & ti<(d+c)])))
    # number of events in delayed arm
    m0 <- n*(exp(-d*mean(lambda_0T_obs[ti>=0 & ti<(d)])))*(1-exp(-c*mean(lambda_0T_obs[ti>=d & ti<(d+c)])))
  }else{
    # number of events in immediate arm
    m1 <- n*(1-exp(-c*mean(lambda_1T_obs[ti>=d & ti<(d+c)])))
    # number of events in delayed arm
    m0 <- n*(1-exp(-c*mean(lambda_0T_obs[ti>=d & ti<(d+c)])))
  }
  # power
  B <- pnorm(sqrt(m0+m1)*abs(VE_D)/(2-VE_D)-1.96)
  return(B)
}
#power_D_simple(d=d,c=c,ti=ti,VE_Tobs=VE_Tobs,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs,n=n)


######################################
# SPECIFY MODEL ASSUMPTIONS

# notation as in paper
R=4 # ramp-up period
VE0=0.6 # vaccine efficacy

# gamma distributed incubation period
inc=6 # mean of incubation period
scale=0.5
shape=inc/scale

s1=0 # date of immediate vaccination
s0=21 # date of delayed vaccination (set to large value for placebo vacc)

# hazard shape (constant, decreasing, or increasing)
hz_shape="constant"
hz_value1=0.001
hz_value2=NA

#hz_shape="decreasing"
#hz_value1=0.0015 # specifies starting value for the background infection hazard rate
#hz_value2=60 # specifies days after which the hazard goes to 0

#hz_shape="increasing" 
#hz_value1=0.002
#hz_value2=30

n <- 1000 # sample size per arm


#############################
# ESTIMATION

# These should be used to produce approximated bias and power for a given trial design
# As they can be slow to integrate, simpler versions (below) are used for plotting

d <- 12 # analysis period start (days)
c <- 21 # analysis period width (days)

# bias (from full integral)
VE_Dobs_int <- VE_D(d=d,c=c,R=R,s1=s1,s0=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)
round(VE_Dobs_int,3)

# power (from full integral)
power_Dobs_int <- power_D(d=d,c=c,R=R,s1=s1,s0=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2,VE_Dobs=VE_Dobs_int,n=n)
round(power_Dobs_int,3)


#############################
# PLOTTING

# ILLNESS ONSETS
# produce plots of illness onset in each arm and apparent VE

# grid of illness onset times considered for plotting, can revise to be a finer grid
# be sure to select an upper bound larger than highest possible analysis period (d+c)
ti <- seq(0,100,by=0.1)

# calculate illness onset hazard rate in immediate arm for each time
lambda_1T_obs <- lambda_T(ti=ti,R=R,sj=s1,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2) 

# calculate illness onset hazard rate in delayed (or control) arm
lambda_0T_obs <- lambda_T(ti=ti,R=R,sj=s0,VE0=VE0,shape=shape,scale=scale,hz_shape=hz_shape,hz_value1=hz_value1,hz_value2=hz_value2)

# calculate apparent VE (1 - hazard ratio) at each illness onset time
VE_Tobs <- 1-lambda_1T_obs/lambda_0T_obs


# save data for plotting
data_T <- data.frame(ti=ti,lambda_1T_obs,lambda_0T_obs,VE_Tobs=VE_Tobs)

# prepare to plot
inc50 <- qgamma(0.5,shape=shape,scale=scale)
inc90 <- qgamma(0.9,shape=shape,scale=scale)
inc95 <- qgamma(0.95,shape=shape,scale=scale)
incmax <- qgamma(0.999,shape=shape,scale=scale)

# specify plot upperbound
tmax <- 50

# plot of lambda1
p_lambda_1T <- ggplot() +
  theme_bw() +
  labs(title="A") +
  ylab(expression(paste("Illness onset hazard rate in Arm 1, ",lambda[1][T](t),sep="")))+
  xlab("Illness onset time, t")+
  scale_x_continuous(expand = c(0,0) , limits = c(0,tmax))+
  scale_y_continuous(expand=c(0,0),limits=c(-hz_value1*.01,hz_value1*1.2),breaks=seq(0,hz_value1*1.2,by=hz_value1/4))+
  annotate("rect", xmin=s1, xmax=min(s1+R,tmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="forestgreen") +
  annotate("rect", xmin=s1+R, xmax=min(s1+R+inc50,tmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc50, xmax=min(s1+R+inc90,tmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc90, xmax=min(s1+R+incmax,tmax), ymin=-Inf, ymax=Inf,alpha=0.2,fill="forestgreen")+
  geom_line(data=data_T,aes(x=ti,y=lambda_1T_obs))
#p_lambda_1T

p_lambda_0T <- ggplot() +
  theme_bw() +
  labs(title="B") +
  ylab(expression(paste("Illness onset hazard rate in Arm 0, ",lambda[0][T](t),sep="")))+
  xlab("Illness onset time, t")+
  scale_x_continuous(expand = c(0,0) , limits = c(0,tmax))+
  scale_y_continuous(expand = c(0,0), limits = c(-hz_value1*.01,hz_value1*1.2),breaks = seq(0,hz_value1*1.2,by=hz_value1/4))+
  annotate("rect", xmin=s0, xmax=min(s0+R,tmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="royalblue") +
  annotate("rect", xmin=s0+R, xmax=min(s0+R+inc50,tmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc50, xmax=min(s0+R+inc90,tmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc90, xmax=min(s0+R+incmax,tmax), ymin=-Inf, ymax=Inf, alpha=0.2, fill="royalblue")+
  geom_line(data=data_T,aes(x=ti,y=lambda_0T_obs)) 
#p_lambda_0T

p_VE_T <- ggplot() +
  theme_bw() +
  labs(title="C") +
  ylab(expression(paste("Apparent vaccine efficacy,  ",VE[T](t),sep="")))+
  xlab("Illness onset time, t")+
  scale_x_continuous(expand = c(0,0) , limits = c(0,tmax))+
  scale_y_continuous(expand = c(0,0) , limits = c(-.01,1.01),breaks = seq(0,1,by=0.25),labels=percent)+
  annotate("rect", xmin=s1, xmax=min(s1+R,tmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="forestgreen") +
  annotate("rect", xmin=s1+R, xmax=min(s1+R+inc50,tmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc50, xmax=min(s1+R+inc90,tmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc90, xmax=min(s1+R+incmax,tmax), ymin=-Inf, ymax=Inf,alpha=0.2,fill="forestgreen")+
  annotate("rect", xmin=s0, xmax=min(s0+R,tmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="royalblue") +
  annotate("rect", xmin=s0+R, xmax=min(s0+R+inc50,tmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc50, xmax=min(s0+R+inc90,tmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc90, xmax=min(s0+R+incmax,tmax), ymin=-Inf, ymax=Inf, alpha=0.2, fill="royalblue")+
  geom_hline(yintercept=VE0,lty=2,color="gray20") +
  geom_line(data=data_T,aes(x=ti,y=VE_Tobs))
#p_VE_T

# plot
grid.arrange(p_lambda_1T, p_lambda_0T, p_VE_T, ncol=3)

# save
plot_title <- paste("lambda3panel_rate",hz_value1,hz_shape,"_R",R,"_VE",VE0,"_b",s0,"_inc",inc,".pdf",sep="")
pdf(plot_title,width=12,height=4)
grid.arrange(p_lambda1_1T, p_lambda_0T, p_VE_T, ncol=3)
dev.off()


# TRIAL ANALYSIS
# produce bias and power plots

###
# if you choose a fixed width analysis period (fixed c), run through #*#
d_vec <- 0:29 # range of analysis period starts considered
c <- 21 # analysis period width

# bias
VE_Dobs <- sapply(FUN=VE_D_simple,X=d_vec,c=c,ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs)

# power
power_Dobs <- sapply(FUN=power_D_simple,X=d_vec,c=c,ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs,n=n)

# save data
data_D <- data.frame(d=d_vec,VE_Dobs=VE_Dobs,VE_Dobs=VE_Dobs,power_Dobs=power_Dobs)

#*#

###
# if you choose a fixed analysis period end (fixed d+c), run through #&#
d_vec <- 0:20 # range of analysis period starts considered
dc <- 21 # analysis period end

# bias and power
VE_Dobs <- rep(NA,length(d_vec))
power_Dobs <- rep(NA,length(d_vec))
for(i in 1:length(d_vec)){
  VE_Dobs[i] <- VE_D_simple(d=d_vec[i],c=dc-d_vec[i],ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs)
  power_Dobs[i] <- power_D_simple(d=d_vec[i],c=dc-d_vec[i],ti=ti,lambda_1T_obs=lambda_1T_obs,lambda_0T_obs=lambda_0T_obs,n=n)
}
  
# save data
data_D <- data.frame(d=d_vec,VE_Dobs=VE_Dobs,power_Dobs=power_Dobs)

#&#

# now that you have your data_D with apparent bias and power for each design
# prepare for plotting

inc50 <- qgamma(0.5,shape=shape,scale=scale)
inc90 <- qgamma(0.9,shape=shape,scale=scale)
inc95 <- qgamma(0.95,shape=shape,scale=scale)
incmax <- qgamma(0.999,shape=shape,scale=scale)
dmax <- max(data_D$d)

p_biasD <- ggplot() +
  theme_bw() +
  labs(title="A") +
  ylab(expression(paste("Apparent vaccine efficacy,  ",VE[d](d,d+c),sep=""))) +
  xlab("Analysis period start, d")+ 
  scale_x_continuous(expand = c(0,0) , limits = c(0,dmax))+
  scale_y_continuous(expand = c(0,0) , limits = c(-.01,1.01),breaks = seq(0,1,by=0.25),labels=percent)+
  annotate("rect", xmin=s1, xmax=min(s1+R,dmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="forestgreen") +
  annotate("rect", xmin=s1+R, xmax=min(s1+R+inc50,dmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc50, xmax=min(s1+R+inc90,dmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc90, xmax=min(s1+R+incmax,dmax), ymin=-Inf, ymax=Inf,alpha=0.2,fill="forestgreen")+
  annotate("rect", xmin=s0, xmax=min(s0+R,dmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="royalblue") +
  annotate("rect", xmin=s0+R, xmax=min(s0+R+inc50,dmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc50, xmax=min(s0+R+inc90,dmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc90, xmax=min(s0+R+incmax,dmax), ymin=-Inf, ymax=Inf, alpha=0.2, fill="royalblue")+
  geom_hline(yintercept=VE0,lty=2,color="gray20") + 
  geom_line(data=data_D,aes(x=d,y=VE_Dobs))

#p_biasD

p_powerD <- ggplot() +
  theme_bw() +
  labs(title="B") +
  ylab(expression(paste("Apparent power"))) +
  xlab("Analysis period start, d")+ 
  scale_x_continuous(expand = c(0,0) , limits = c(0,dmax))+
  scale_y_continuous(expand = c(0,0) , limits = c(-.01,1.01),breaks = seq(0,1,by=0.25),labels=percent)+
  annotate("rect", xmin=s1, xmax=min(s1+R,dmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="forestgreen") +
  annotate("rect", xmin=s1+R, xmax=min(s1+R+inc50,dmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc50, xmax=min(s1+R+inc90,dmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="forestgreen") +
  annotate("rect", xmin=s1+R+inc90, xmax=min(s1+R+incmax,dmax), ymin=-Inf, ymax=Inf,alpha=0.2,fill="forestgreen")+
  annotate("rect", xmin=s0, xmax=min(s0+R,dmax), ymin=-Inf, ymax=Inf, alpha=0.7, fill="royalblue") +
  annotate("rect", xmin=s0+R, xmax=min(s0+R+inc50,dmax), ymin=-Inf, ymax=Inf, alpha=0.5, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc50, xmax=min(s0+R+inc90,dmax), ymin=-Inf, ymax=Inf, alpha=0.3, fill="royalblue") +
  annotate("rect", xmin=s0+R+inc90, xmax=min(s0+R+incmax,dmax), ymin=-Inf, ymax=Inf, alpha=0.2, fill="royalblue")+
  geom_line(data=data_D,aes(x=d,y=power_Dobs))

#p_powerD

# plot
grid.arrange(p_biasD, p_powerD, ncol=2)

# title if fixed width c used
plot_title <- paste("trial2panel_rate",hz_value1,hz_shape,"_R",R,"_VE",VE0,"_b",s0,"_c",c,"_inc",inc,"_n",n,".pdf",sep="")

# title if fixed cutoff d+c (dc) used
plot_title <- paste("trial2panel_rate",hz_value1,hz_shape,"_R",R,"_VE",VE0,"_b",s0,"_dc",dc,"_inc",inc,"_n",n,".pdf",sep="")

# save plot
pdf(plot_title,width=8,height=4)
grid.arrange(p_biasD, p_powerD, ncol=2)
dev.off()

###########################