# ==============================================================
# R code for B.1.1  Intensities featuring smooth time trends
#
# December 11, 2018
# ==============================================================


# --------------------------------------------------------------
# Generate data based on Figure B.1: an illness-death model
# --------------------------------------------------------------

library(truncdist)

generatedata.f <- function(N, CC, lam12, lam13, lam23, alp12, alp13, alp23) {
  outdata <- NULL
  for (i in 1:N) {
    t12 <- rweibull(1, shape=alp12, scale=(1/lam12))
    t13 <- rweibull(1, shape=alp13, scale=(1/lam13))
    t1  <- min(c(t12, t13))
    s1  <- ifelse(t12 < t13, 2, 3)
    if ((t12 < t13) && (t12 < CC)) {
      t23 <- rtrunc(1, spec="weibull", a=t1, b=Inf, shape=alp23, scale=(1/lam23))
      s23 <- 3

      estart  <- c(0, 0, t1)
      estop   <- c(t1, t1, min(t23,CC))
      estatus <- c(1, 0, ifelse(t23 < CC, 1, 0))
      from    <- c(1, 1, s1)
      to      <- c(s1, s23, s23)
      for.etm <- c(1,0,1)
    }
    else {
      estart  <- rep(0,2)
      estop   <- rep(min(t1,CC),2)
      estatus <- c(ifelse(t1 < CC, 1, 0), 0)
      from    <- rep(1,2)
      to      <- c(s1,2)
      for.etm <- c(1,0)
    }
    id <- rep(i, length(estart))
    outdata <- rbind(outdata, data.frame(id, estart, estop, estatus, from, to, for.etm))
  }
  outdata <- outdata[order(outdata$id, outdata$from, outdata$to),]
  return(outdata)
}


stats <- data.frame(CC=100, lam12=0.0163, lam13=0.0200, lam23=0.0224,
                            alp12=2, alp13=2, alp23=2)

set.seed(1000)
simdata <- generatedata.f(N=5000, CC=stats$CC, 
  lam12=stats$lam12, lam13=stats$lam13, lam23=stats$lam23,
  alp12=stats$alp12, alp13=stats$alp13, alp23=stats$alp23)
print(simdata[simdata$id %in% c(1,3,16,181),])


# --------------------------------------------------------------
# Nelson-Aalen Estimates
# --------------------------------------------------------------

library(survival)

fit12 <- coxph(Surv(estart, estop, estatus) ~ 1,
  data=simdata[(simdata$from == 1) & (simdata$to == 2),], method="breslow")
na12 <- survfit(fit12, type="aalen")

fit13 <- coxph(Surv(estart, estop, estatus) ~ 1,
  data=simdata[(simdata$from == 1) & (simdata$to == 3),], method="breslow")
na13 <- survfit(fit13, type="aalen")

fit23 <- coxph(Surv(estart, estop, estatus) ~ 1,
  data=simdata[(simdata$from == 2) & (simdata$to == 3),], method="breslow")
na23 <- survfit(fit23, type="aalen")


# --------------------------------------------------------------
# Aalen-Johansen Estimates
# --------------------------------------------------------------

library(etm)

trdata <- simdata[simdata$for.etm == 1,]
trdata$from <- as.character(trdata$from)
trdata$to <- ifelse((trdata$to == 3) & (trdata$estatus == 0),
                    "cens", as.character(trdata$to))
trdata <- trdata[,c("id","from","to","estart","estop")]
dimnames(trdata)[[2]] <- c("id","from","to","entry","exit")

tra <- matrix(ncol=3, nrow=3, FALSE)
tra[1,c(2,3)] <- TRUE
tra[2,3] <- TRUE

tr <- etm(trdata, c("1","2","3"), tra, "cens", 0)

aj <- data.frame(cbind(tr$time, t(tr$est[1,,])))
dimnames(aj)[[2]] <- c("tt","P11","P12","P13")


# --------------------------------------------------------------
# True values
# --------------------------------------------------------------

weib.f <- function(u, lam, alp) {
  Qt <- (lam*u)^alp
  qt <- alp*lam*(lam*u)^(alp-1)
  return( data.frame(qt, Qt) )
}


Pmat.weib.f <- function(tt, lam12, lam13, lam23, alp12, alp13, alp23) {
  mat12 <- weib.f(u=tt, lam=lam12, alp=alp12)
  mat13 <- weib.f(u=tt, lam=lam13, alp=alp13)
  mat23 <- weib.f(u=tt, lam=lam23, alp=alp23)

  P11 <- exp( (-1)*(mat12$Qt + mat13$Qt) )
  P22 <- exp( (-1)*mat23$Qt )

  func12.f <- function(u, tt, lam12, lam13, lam23, alp12, alp13, alp23) {
    mat12  <- weib.f(u=u,  lam=lam12, alp=alp12)
    mat13  <- weib.f(u=u,  lam=lam13, alp=alp13)
    mat23  <- weib.f(u=u,  lam=lam23, alp=alp23)
    mat23t <- weib.f(u=tt, lam=lam23, alp=alp23)

    pp <- mat12$qt*exp((-1)*(mat12$Qt + mat13$Qt))*exp((-1)*(mat23t$Qt - mat23$Qt))
    return(pp)
  }

  P12 <- integrate(func12.f, lower=0, upper=tt, tt=tt,
                   lam12=lam12, lam13=lam13, lam23=lam23,
                   alp12=alp12, alp13=alp13, alp23=alp23)$value

  Pmat <- matrix(0, nrow=3, ncol=3)
  Pmat[1,] <- c(P11, P12, 1 - P11 - P12)
  Pmat[2,] <- c(0, P22, 1 - P22)
  Pmat[3,] <- c(0, 0, 1)
  return(Pmat)
}

tau <- round(max(simdata$estop),1)

true <- data.frame(tt=0, P12=0, P13=0)
for (tt in seq(1,tau,by=1)) {
  PP <- Pmat.weib.f(tt=tt, lam12=stats$lam12, lam13=stats$lam13, lam23=stats$lam23,
                    alp12=stats$alp12, alp13=stats$alp13, alp23=stats$alp23)
  true <- rbind(true, data.frame(tt, P12=PP[1,2], P13=PP[1,3]))
}


# --------------------------------------------------------------
# Create Figure B.2
# --------------------------------------------------------------

pdf("figureB2.pdf", width=11.5, height=8.5)
par(mfrow=c(1,2), mai=c(0.8,0.8,0.4,0.4))

plot(0, 0, type="n", axes=F, xlim=c(0,100), ylim=c(0,5), xlab="", ylab="")
axis(side=1, at=seq(0,100,by=20), label=T)
axis(side=2, at=seq(0,5,by=1), label=T, las=2, adj=1)
mtext("TIME (t)", side=1, line=2.6)
mtext("CUMULATIVE  INTENSITY", side=2, line=2.8)
mtext("(a)", side=3, line=0.5, cex=1.2)

legend(0, 5, c("1 -> 2  CUMULATIVE  INTENSITY",
               "1 -> 3  CUMULATIVE  INTENSITY",
               "2 -> 3  CUMULATIVE  INTENSITY"),
       lty=c(1,2,3), bty="n")

tt  <- seq(0,stats$CC,by=1)
H12 <- (stats$lam12*tt)^2
H13 <- (stats$lam13*tt)^2
H23 <- (stats$lam23*tt)^2

lines(na12$time, na12$cumhaz, type="s", lty=1, lwd=4)
lines(tt, H12, lty=1, lwd=2.5, col="grey")

lines(na13$time, na13$cumhaz, type="s", lty=2, lwd=4)
lines(tt, H13, lty=2, lwd=2.5, col="grey")

lines(na23$time, na23$cumhaz, type="s", lty=3, lwd=4)
lines(tt, H23, lty=3, lwd=2.5, col="grey")
box()


plot(0, 0, type="n", axes=F, xlim=c(0,100), ylim=c(0,1.1), xlab="", ylab="")
axis(side=1, at=seq(0,100,by=20), label=T)
axis(side=2, at=seq(0,1,by=0.2), label=T, las=2, adj=1)
mtext("TIME (t)", side=1, line=2.6)
mtext("CUMULATIVE  PROBABILITY", side=2, line=2.8)
mtext("(b)", side=3, line=0.5, cex=1.2)

legend(0,1.1,c("PROBABILITY  OF  ALIVE  WITH  DISEASE",
               "CUMULATIVE  PROBABILITY  OF  DEATH"),
       lty=c(1,2), lwd=c(2,2), bty="n", cex=1)

lines(aj$tt, aj$P12, lty=1, lwd=4)
lines(true$tt, true$P12, lty=1, lwd=2.5, col="grey")

lines(aj$tt, aj$P13, lty=2, lwd=4)
lines(true$tt, true$P13, lty=2, lwd=2.5, col="grey")
box()

dev.off()