# ===============================================================
# R Code for B.1.2 Processes with piecewise-constant intensities
#
# December 11, 2018
# ===============================================================


qmat <- rbind(c(-0.03, 0.01, 0.02), c(0, -0.02, 0.02), c(0, 0, 0))

btime <- c(0,30,60)
bmat  <- matrix(0, nrow=3, ncol=2)
bmat[2,1] <- 1
bmat[3,2] <- 1

beta <- matrix(0, nrow=2, ncol=3)
beta[,1] <- (c(2:3) - 1)*log(4)
beta[,2] <- (c(2:3) - 1)*log(2)
beta[,3] <- rep(log(1.2), 2)

CC <- 1000


library(msm)

# ---------------------------------------------------------------
# Generate data using sim.msm function
# ---------------------------------------------------------------

set.seed(1000)

simdata <- NULL
for (i in 1:1000) {
  sim <- sim.msm(qmatrix=qmat, maxtime=CC, covs=bmat, beta=beta,
                 obstimes=btime, start=1, mintime=0)
  times  <- sim$times
  states <- sim$states

  if (length(times) == 3) {
    estart  <- c(times[1],  times[1], times[2])
    estop   <- c(times[2],  times[2], times[3])
    estatus <- c(1, 0, ifelse(times[3] == CC, 0, 1))
    from    <- c(states[1], states[1], states[2])
    to      <- c(states[2], 3, 3)
    for.etm <- c(1, 0, 1) 
  }
  else {
    estart  <- rep(times[1],2)
    estop   <- rep(times[2],2)
    estatus <- c(ifelse(times[2] == CC, 0, 1), 0)
    from    <- rep(states[1],2)
    to      <- c(states[2],2)
    for.etm <- c(1,0)
  }

  id <- rep(i, length(estart))
  simdata <- rbind(simdata, data.frame(id, estart, estop, estatus, from, to, for.etm))
}

print( simdata[simdata$id %in% c(1,3,8),] )


# ---------------------------------------------------------------
# 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")


# ---------------------------------------------------------------
# True cumulative transition intensities
# ---------------------------------------------------------------

getHt.f <- function() {
  tt <- seq(0,30,by=1)
  H12 <- 0.01*tt
  H13 <- 0.02*tt
  H23 <- 0.02*tt
  H12p <- H12[length(H12)]; H13p <- H13[length(H13)]; H23p <- H23[length(H23)]

  tt <- seq(31,60,by=1)
  H12 <- c(H12, H12p + 0.01*exp(1*log(4))*(tt-30))
  H13 <- c(H13, H13p + 0.02*exp(1*log(2))*(tt-30))
  H23 <- c(H23, H23p + 0.02*exp(log(1.2))*(tt-30))
  H12p <- H12[length(H12)]; H13p <- H13[length(H13)]; H23p <- H23[length(H23)]

  tt <- seq(61,100,by=1)
  H12 <- c(H12, H12p + 0.01*exp(2*log(4))*(tt-60))
  H13 <- c(H13, H13p + 0.02*exp(2*log(2))*(tt-60))
  H23 <- c(H23, H23p + 0.02*exp(log(1.2))*(tt-60))

  tt <- seq(0,100,by=1)
  return( data.frame(tt, H12, H13, H23) )
}

Ht <- getHt.f()


# ---------------------------------------------------------------
# Create Figure B.3
# ---------------------------------------------------------------

pdf("figureB3.pdf", width=11.5, height=8.5)
par(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=10), label=T)
axis(side=2, at=seq(0,5,by=1), label=T, las=2, adj=1)
mtext("TIME (t)", side=1, line=2.4)
mtext("CUMULATIVE  INTENSITY", side=2, line=2.8)

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

lines(na12$time, na12$cumhaz, type="s", lty=1, lwd=3)
lines(na13$time, na13$cumhaz, type="s", lty=2, lwd=3)
lines(na23$time, na23$cumhaz, type="s", lty=3, lwd=3)

lines(Ht$tt, Ht$H12, lty=1, lwd=3, col="grey")
lines(Ht$tt, Ht$H13, lty=2, lwd=3, col="grey")
lines(Ht$tt, Ht$H23, lty=3, lwd=3, col="grey")

dev.off()