#-----------------------------------------------------------------------------#
# Code for "CQCoVID19"
# Author(s): Liu, Wang, Ma, Wang
# 2020\4\3
#-----------------------------------------------------------------------------#

library(readxl)
library(splines)
library(ggplot2)
library(quantreg)
library(SparseM)



#-----------------------------------------------------------------------#
# load data 
#-----------------------------------------------------------------------#
data = read_excel("Additional file 1.xlsx", sheet=1, na="NA")

Age      = as.numeric(data$Age)
InDate   = as.numeric(data$InDate)
OutDate   = as.numeric(data$OutDate)
InPeriod = as.numeric(data$InPeriod)


#--------------------------------------------------------------------------------#
# Plot Figure 1 -- scatter plot of Incubation period vs. Age
#--------------------------------------------------------------------------------#
plot(Age, InPeriod, xlab = "Age", ylab = "Incubation period")


#--------------------------------------------------------------------------------#
# Plot Figure 2 -- histogram
#--------------------------------------------------------------------------------#

#the estimator -- normal distribution of ages
logFunC = function(paramters, Age){
  mean = paramters[1]
  sd = paramters[2]
  dens = dnorm(Age, mean, sd)
  dens[dens < 1e-12] = 1e-12
  loglk = -sum(log(dens))
  return(loglk)
}
OFunC = function(paramters){ # objective function
  Qval = logFunC(paramters, Age)
  return(Qval)
}
estimatorC = nlminb(c(1, 1), OFunC, lower = c(-Inf, 1e-20))$par
mean = estimatorC[1]
sd = estimatorC[2]
#plot histogram
hist(Age, probability = T, ylim = c(0,0.030), main = "" )
age = sort(unique(Age))
lines(dnorm(age, mean, sd), col =2, lwd = 3)


#--------------------------------------------------------------------------------#
# Plot Figure 3 -- scatter plot of Number of infected cases vs. Infected time
#--------------------------------------------------------------------------------#

years = seq(min(InDate),max(InDate))
InDate1 = rep(NA, 46)
for (ii in 1 :46) {
  InDate1[ii] = length( as.numeric(data$InDate[data$InDate >= ii-3 &data$InDate< ii-2]))
}  #counting number infected cases of day
plot(years, InDate1, col="red",xlab = "Infected time",ylab = "Number of infected cases")
c = cbind(years, InDate1)
sy.hat <- smooth.spline(c,df = 15)    #using spine 
lines(sy.hat, col = 'SlateGrey',lty = 1, lwd = 2)


#--------------------------------------------------------------------------------#
# Table 1 -- parameter estimator
#--------------------------------------------------------------------------------#

tstar = max(OutDate) + 1
# initial function
InitFun = function(Ages, InPeriod){
  # eta  = theta[1], beta = theta[-1]
  logF0 = function(theta, Ages, InPeriod){
    dens = dweibull(InPeriod, theta[1], Ages %*% theta[-1])
    dens[dens < 1e-14] = 1e-14
    return(-mean(log(dens)))
  } 
  OFun0 = function(theta){return(logF0(theta, Ages, InPeriod))}
  
  # initial value
  init0 = constrOptim(c(1, 1, rep(0, ncol(Ages) - 1)), OFun0, NULL, 
                      ui = rbind(c(1, rep(0, ncol(Ages))), cbind(0, Ages)),
                      ci= rep(1e-14, 1 + nrow(Ages)))$par
  return(init0)
}
# the negtive log-likelihood function
logFun = function(theta, tstar, tvec, X, Yvec, Uvec, thetak){
  print(theta)
  Delta = tstar - tvec
  eta =  theta[1]
  beta = theta[-1]
  lambda = X %*% beta
  dens = dweibull(Yvec, eta, lambda)
  dens[dens < 1e-14] = 1e-14
  GDelta = pweibull(Delta, eta, lambda)
  
  part1 = mean(GDelta * log(dens))
  
  lambda1 = lambda
  lambda1[lambda1 < 1e-14] = 1e-14
  part2 = mean((log(eta) - eta * log(lambda1)) * (1 - GDelta))
  
  DLamb = (Delta / lambda)^eta
  part3 = -mean((DLamb + 1) * exp(-DLamb))
  
  etak = thetak[1]
  betak = thetak[-1]
  lambdak = X %*% betak

  pp = pweibull(Delta, etak, lambdak)
  xk = qweibull(Uvec, etak, lambdak)
  xk[xk > 1e14] = 1e14
  logg = log(xk)
  
  n = length(Yvec)
  intlog = numeric(n)
  for(ii in 1:n){
    if(length(logg[Uvec >= pp[ii]])){
      intlog[ii] = (1 - pp[ii]) * mean(logg[Uvec >= pp[ii]])
    }else{
      intlog[ii] = 0
    }
  }
  part4 = (eta - 1) * mean(intlog)
  
  fval = -(part1 + part2 + part3 + part4)
  return(fval)
}

# the estimator -- EM Algorithm
EMfun = function(init0, tstar, InDate, X, InPeriod, m = 1000, maxit = 20){
  minD = min(tstar - InDate)
  thetak = init0
  etak = thetak[1]
  betak = thetak[-1]
  n = length(InDate)
  lamdbak = X %*% betak
  
  is.cont = TRUE
  kk = 1
  while(is.cont & (kk < maxit)){
    # E-step
    pmin = min(pweibull(minD, etak, lamdbak))
    Uvec = runif(m, pmin, 0.999999)
    OFunE = function(theta){
      loglkE = logFun(theta, tstar, InDate, X, InPeriod, Uvec, thetak)
      return(loglkE)
    }
    
    # M-step
    thetaup = constrOptim(thetak, OFunE, NULL, 
                          ui = rbind(c(1, rep(0, ncol(X))), cbind(0, X)),
                          ci= rep(1e-5, 1 + nrow(X)))$par
    
    # check whether or not continue
    if(sum(abs(thetaup - thetak)) < n^(-0.51)){is.cont = FALSE}
    thetak = thetaup
    kk = kk + 1
    print(thetaup)
  }
  return(thetaup)
}
X = cbind(1, Age, Age^2, Age^3)
set.seed(666)
init0 = InitFun(X, InPeriod)
theta = EMfun(init0, tstar, InDate, X, InPeriod, m = 1000, maxit = 20)
eta = theta[1]
beta = theta[-1]



#--------------------------------------------------------------------------------#
# Plot Figure 4 -- conditional quantiles
#--------------------------------------------------------------------------------#

# compute quantiles
qumat = function(quantiles = c(0.05, 0.25, 0.5, 0.75, 0.90, 0.95), x,eta, beta){
  r = length(quantiles)
  c = nrow(x)
  lambda = x %*% beta
  quantilematrix = matrix(NA, ncol = c, nrow = r)
  for (a in 1:c) {
    quantilematrix[,a] = as.matrix(qweibull(quantiles, eta, lambda[a]))
  }
  return(quantilematrix)
}
# plot conditional quantiles 
plotqs = function(quantilematrix, quantiles = c(0.05, 0.25, 0.5, 0.75, 0.90, 0.95), 
                  x, ym, tb, ty = 1, wd =2, ex = 0.3, points = T, text = T, add = F,
                  color = c("blue","orange","red", "orange", "darkgreen", "blue"))
{
  r = nrow(quantilematrix)
  if(!add){
    plot(age,quantilematrix[1,], ylim = c(0,ym), ylab = 'Incubation period',
         xlab = 'Age', type = 'l', col = "white")
  }
  for (ii in 1:r) {
    lines(age,quantilematrix[ii,], type = 'l', lty = ty, lwd = wd, col = color[ii])
    if(points){
      points(age, quantilematrix[ii,], col = color[ii], pch = 20, cex = ex)
    }
    if(text){
      sca = quantilematrix[, tb] - 1
      text(tb, sca[ii], quantiles[ii], col = color[ii])
    }
  }
}

# plot quantiles
age = sort(unique(Age))

x = cbind(1, age, age^2, age^3)

quantilemat = qumat(quantiles = c(0.05, 0.25, 0.5, 0.75, 0.90, 0.95),x = x, eta = eta, beta = beta)
range(quantilemat)

plot(Age, InPeriod, col = 'grey90', cex = 0.8, ylab = "Incubation period")

plotqs(quantilematrix = quantilemat, x = x, ym = max(quantilemat), 
       tb = 88, wd = 3, ex = 1, points = F, add = T)

abline(h = 15.05, col = "SlateGrey", lty = 2, lwd = 2)
text(33, 16, 15.05, col = "SlateGrey")


#--------------------------------------------------------------------------------#
# Plot Figure 5 -- left
#--------------------------------------------------------------------------------#

# choose seperate points and split data accordingly
dataA = data.frame(Age, InDate, OutDate, InPeriod)
data1 = subset(dataA, Age <= 25)
data2 = subset(dataA, Age > 25 & Age<=60)
data3 = subset(dataA, Age> 60)

# estimate parameters under no covariant
InitFunA = function(Y){
  logF0 = function(theta, Y){
    dens = dweibull(Y, theta[1], theta[2])
    dens[dens < 1e-12] = 1e-12
    return(- mean(log(dens)))
  } 
  OFun0 = function(lambda){return(logF0(lambda, Y))}
  return(nlminb(c(1, 1), OFun0, lower = c(1e-12, 1e-12))$par)
}
logFunA = function(theta, tstar, tvec, Yvec, Uvec, thetak){
  eta = theta[1]
  lambda = theta[2]
  
  Delta = tstar - tvec
  dens = dweibull(Yvec, eta, lambda)
  dens[dens < 1e-14] = 1e-14
  GDelta = pweibull(Delta, eta, lambda)
  
  part1 = mean(GDelta * log(dens))
  if(lambda < 1e-14){
    part2 = (log(eta) - eta * log(1e-14)) * 
      (1 - mean(GDelta))
  }
  else{
    part2 = (log(eta) - eta * log(lambda)) * 
      (1 - mean(GDelta))
  }
  
  DLamb = (Delta / lambda)^eta
  part3 = -mean((DLamb + 1) * exp(-DLamb))
  
  pp = pweibull(Delta, thetak[1], thetak[2])
  xk = qweibull(Uvec, thetak[1], thetak[2])
  xk[xk > 1e14] = 1e14
  logg = log(xk)
  
  n = length(Yvec)
  intlog = numeric(n)
  for(ii in 1:n){
    if(length(logg[Uvec >= pp[ii]])){
      intlog[ii] = (1 - pp[ii]) * mean(logg[Uvec >= pp[ii]])
    }else{
      intlog[ii] = 0
    }
  }
  part4 = (eta - 1) * mean(intlog)
  
  fval = -(part1 + part2 + part3 + part4)
  print(c(eta,lambda,fval))
  return(fval)
}

EMfunA = function(init0, tstar, InDate, InPeriod, m = 1000, maxit = 20){
  minD = min(tstar - InDate)
  thetak = init0
  n = length(InDate)
  is.cont = TRUE
  kk = 1
  while(is.cont & (kk < maxit)){
    # E-step
    pmin = pweibull(minD, thetak[1], thetak[2])
    Uvec = runif(m, pmin, 0.999999)
    OFunE = function(theta){
      loglkE = logFunA(theta, tstar, InDate, InPeriod, Uvec, thetak)
      return(loglkE)
    }
    
    # M-step
    thetaup = nlminb(thetak, OFunE, lower = c(1e-12, 1e-12))$par
    
    # check whether or not continue
    if(sum(abs(thetaup - thetak)) < n^(-0.51)){is.cont = FALSE}
    thetak = thetaup
    kk = kk + 1
    print(thetaup)
  }
  return(thetaup)
}

set.seed(666)
init01 = InitFunA(data1$InPeriod)
init02 = InitFunA(data2$InPeriod)
init03 = InitFunA(data3$InPeriod)

estimator1 = EMfunA(init01, max(data1$OutDate)+1, data1$InDate, data1$InPeriod)
estimator2 = EMfunA(init02, max(data2$OutDate)+1, data2$InDate, data2$InPeriod)
estimator3 = EMfunA(init03, max(data3$OutDate)+1, data3$InDate, data3$InPeriod)


#
x = seq(0, 25, length.out = 10000)
# compute density
y1 = dweibull(x, estimator1[1], estimator1[2])
y2 = dweibull(x, estimator2[1], estimator2[2])
y3 = dweibull(x, estimator3[1], estimator3[2])

h1 = hist(data1$InPeriod, 
          freq=FALSE, 
          breaks=12, 
          xlab="Incubation Period", ylim = c(0.0,0.12),
          main="")  
h2 = hist(data2$InPeriod, 
          freq=FALSE, 
          breaks=12, 
          xlab="Incubation Period", ylim = c(0.0,0.12),
          main="")
h3 = hist(data3$InPeriod, 
          freq=FALSE, 
          breaks=12, 
          xlab="Incubation Period", ylim = c(0.0,0.12),
          main="")
plot(h1, col=rgb(0.8,0.8,0.8), xlim=c(0,max(data$InPeriod)),ylim = c(0.0,0.12),freq=FALSE, xlab = "Incubation Period", main="")  # first histogram
plot(h2, col=rgb(0.9,0.9,0.9), xlim=c(0,max(data$InPeriod)),ylim = c(0.0,0.12),freq=FALSE, xlab = "Incubation Period", add=T, main="")  # second
plot(h3, col=rgb(1,1,1), xlim=c(0,max(data$InPeriod)),ylim = c(0.0,0.12),freq=FALSE, xlab = "Incubation Period", add=T, main="")


x1 = seq(0, max(data1$InPeriod), 0.01)
x2 = seq(0, max(data2$InPeriod), 0.01)
x3 = seq(0, max(data3$InPeriod), 0.01)
lines(x1,dweibull(x1,estimator1[1],estimator1[2]),type = "l",lwd=3,lty=4,col="blue", ylim=c(0,1), main="")
lines(x2,dweibull(x2,estimator2[1],estimator2[2]),type = "l",lwd=2,lty=1,col="darkgreen", ylim=c(0,1), main="")
lines(x3,dweibull(x3,estimator3[1],estimator3[2]),type = "l",lwd=3,lty=5,col="red", ylim=c(0,1), main="")
words=c("0-25","26-60","over 60")
colors=c("blue","darkgreen","red")
line=c(4,1,5)
width=c(3,2,3)
legend("topright",legend = words,text.width = 3.5, col = colors, lty = line, lwd = width, cex = 1)


#--------------------------------------------------------------------------------#
# Plot Figure 5 -- right
#--------------------------------------------------------------------------------#
a = 25; b = 60
x = seq(0, 25, length.out = 10000)
# compute probability
y11 = pweibull(x, estimator1[1], estimator1[2])
y22 = pweibull(x, estimator2[1], estimator2[2])
y33 = pweibull(x, estimator3[1], estimator3[2])

# compute 0.95quantiles
q11 = qweibull(0.95, estimator1[1], estimator1[2])
q22 = qweibull(0.95, estimator2[1], estimator2[2])
q33 = qweibull(0.95, estimator3[1], estimator3[2])


# plot distribution and 0.95 quantiles
plot(x, y11, type = 'l', lwd = 3, col = "blue", lty = 4, xlab = "Incubation Period", ylab = "Cumulative distribution")
lines(x, y22, type = 'l', lwd = 2, col = "darkgreen", lty = 1)
lines(x, y33, type = 'l', lwd = 3, col = "red", lty = 5)
abline(h = 0.95, lty = 2, col = "SlateGrey")
abline(v = q11, lty = 2, col = "blue")
abline(v = q22, lty = 2, col = "darkgreen")
abline(v = q33, lty = 2, col = "red")
axis(1, at = q22, label = round(q22,1), cex.axis = 0.7, col.ticks = 'darkgreen', col.axis = 'darkgreen')
axis(1, at = q33, label = round(q33,1), cex.axis = 0.7, col.ticks = 'red', col.axis = 'red')
legend("left",legend = c(paste("0 -", a), paste(a + 1,"-",b), paste("over", b)), lwd = c(3, 2, 3),
       lty = c(4, 1, 5), col = c("blue", "darkgreen", "red"), cex = 1, text.width = 3.5)



#--------------------------------------------------------------------------------#
# Plot Figure 6 --  quantile regression
#--------------------------------------------------------------------------------#
x1 = Age
x2 = Age^2
x3 = Age^3
xy = data.frame(x1 = Age, x2 = Age^2, x3 = Age^3, y = InPeriod)
fit1 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.05)
fit2 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.25)
fit3 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.50)
fit4 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.75)
fit5 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.95)
fit6 <- rq(InPeriod ~ x1 + x2 + x3, data = xy, tau=.90)
xx = seq(from = min(Age), to = max(Age), by = 0.001)
y1=cbind(xx, fit1$coefficients[1] + fit1$coefficients[2] * xx 
         + fit1$coefficients[3] * xx^2 + fit1$coefficients[4] * xx^3)
y2=cbind(xx, fit2$coefficients[1] + fit2$coefficients[2] * xx 
         + fit2$coefficients[3] * xx^2 + fit2$coefficients[4] * xx^3)
y3=cbind(xx, fit3$coefficients[1] + fit3$coefficients[2] * xx 
         + fit3$coefficients[3] * xx^2 + fit3$coefficients[4] * xx^3)
y4=cbind(xx, fit4$coefficients[1] + fit4$coefficients[2] * xx 
         + fit4$coefficients[3] * xx^2 + fit4$coefficients[4] * xx^3)
y5=cbind(xx, fit5$coefficients[1] + fit5$coefficients[2] * xx 
         + fit5$coefficients[3] * xx^2 + fit5$coefficients[4] * xx^3)
y6=cbind(xx, fit6$coefficients[1] + fit6$coefficients[2] * xx 
         + fit6$coefficients[3] * xx^2 + fit6$coefficients[4] * xx^3)
#quantile regression plot
plot(Age, InPeriod, col = 'grey90', cex = 0.8, ylab = "Incubation period")
lines(y1, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "blue")
lines(y2, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "orange")
lines(y3, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "red")
lines(y4, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "orange")
lines(y6, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "darkgreen")
lines(y5, ylim = c(0, 25), ylab = 'Incubation period',
      xlab = 'Age', type = "l",lwd = 2, col = "blue")
abline(h = 15.05, col = "SlateGrey", lty = 2, lwd = 2)
text(33, 16, 15.05, col = "SlateGrey")
text(88, 1.5, 0.05, col = "blue")
text(88, 5.5, 0.25, col = "orange")
text(88, 10,  0.50, col = "red")
text(88, 15, 0.75, col = "orange")
text(88, 19, 0.90, col = "darkgreen")
text(88, 22, 0.95, col = "blue")