# ===========================================================
# R code to illustrate data generation steps described in 
# 2.5 Simulation of Multistate Processes
#
# December 11, 2018
# ===========================================================
# Note: Based on a reversible illness-death process shown in
#       left panel of Figure 2.1
# ===========================================================


library(msm)


# -----------------------------------------------------------
# Setup parameters
# Note:
#   P(Z(1) < 3 | Z(0) = 1) = P13(0,1) = 0.8
#   1/(q12 + q13) = mean1
#   1/(q21 + q23) = mean2
#   q23 = 1.2*q13
# -----------------------------------------------------------

getstats.f <- function(mean1, mean2, P13tau, tau) {
  Qmat.f <- function(q12, q13, q21, q23) {
    qmat <- matrix(0, nrow=3, ncol=3)
    qmat[1,] <- c(-(q12 + q13), q12, q13)
    qmat[2,] <- c(q21, -(q21 + q23), q23)
    return(qmat)
  }

  funcP.f <- function(q13, mean1, mean2, P13tau, tau) {
    q12 <- (1/mean1) - q13
    q23 <- 1.2*q13
    q21 <- (1/mean2) - q23

    Qmat <- Qmat.f(q12=q12, q13=q13, q21=q21, q23=q23)
    Pt <- MatrixExp(Qmat, t=tau, n=50, k=3, method="series")
    eq <- Pt[1,ncol(Pt)] - P13tau
    return(eq)
  }

  q13 <- uniroot(funcP.f, interval=c(0,10), 
                 mean1=mean1, mean2=mean2, P13tau=P13tau, tau=tau)$root
  q12 <- (1/mean1) - q13
  q23 <- 1.2*q13
  q21 <- (1/mean2) - q23

  return( data.frame(q12, q13, q21, q23, tau) )
}


# -----------------------------------------------------------
# Generate data
# -----------------------------------------------------------

generatedata.f <- function(M, CC, q12, q13, q21, q23) {
  qmat <- trans <- matrix(0, nrow=2, ncol=2)
  qmat[1,] <- c(q12, q13); trans[1,] <- c(2, 3)
  qmat[2,] <- c(q21, q23); trans[2,] <- c(1, 3)

  outdata <- lapply(1:M, function(ith, CC, qmat, trans) {
    cur.tt <- 0; cur.state <- 1

    tt <- cur.tt; state <- cur.state; status <- 1
    while ( (cur.tt < CC) && (cur.state != 3) ) {
      pre.tt <- cur.tt; pre.state <- cur.state

      cur.tt    <- pre.tt + rexp(1, rate=sum(qmat[pre.state,]))
      cur.state <- sample(trans[pre.state,], size=1, 
                          prob=(qmat[pre.state,]/sum(qmat[pre.state,])), 
                          replace=FALSE)

      tt     <- c(tt, ifelse(cur.tt < CC, cur.tt, CC))
      state  <- c(state, cur.state)
      status <- c(status, ifelse(cur.tt < CC, 1, 0))
    }

    nlen <- length(tt)
    dataE <- data.frame(id=rep(ith,(nlen-1)), 
                        from=state[-nlen], 
                        to=state[-1], 
                        start=tt[-nlen], 
                        stop=tt[-1], 
                        status=status[-1], 
                        sojourn=c(1:(nlen-1)), 
                        enum=rep(1,(nlen-1)))

    dataC <- dataE
    dataC$to <- ifelse((dataE$from == 1) & (dataE$to == 2), 3, dataC$to)
    dataC$to <- ifelse((dataE$from == 1) & (dataE$to == 3), 2, dataC$to)
    dataC$to <- ifelse((dataE$from == 2) & (dataE$to == 1), 3, dataC$to)
    dataC$to <- ifelse((dataE$from == 2) & (dataE$to == 3), 1, dataC$to)
    dataC$status <- rep(0,(nlen-1))
    dataC$enum   <- rep(2,(nlen-1))

    outdata <- rbind(dataE, dataC)
    outdata <- outdata[order(outdata$sojourn, outdata$enum),]
    return(outdata)
  }, CC=CC, qmat=qmat, trans=trans)
  outdata <- do.call("rbind", outdata)
  outdata <- data.frame(outdata, row.names=c(1:nrow(outdata)))
  return(outdata)
}


## Setup parameters
stats <- getstats.f(mean1=0.25, mean2=0.1, P13tau=0.8, tau=1)

q12 <- stats$q12; q13 <- stats$q13
q21 <- stats$q21; q23 <- stats$q23
tau <- stats$tau


## Generate data for 3 individuals
set.seed(80)
simdata <- generatedata.f(M=3, CC=tau, q12=q12, q13=q13, q21=q21, q23=q23)
print(simdata)


# -----------------------------------------------------------
# Step-by-step to generate data for the first individual
# in simdata data set
# -----------------------------------------------------------

set.seed(80)
T1 <- rexp(1, rate=(q12 + q13))
Z1 <- rbinom(1, size=1, prob=(q12/(q12 + q13)))

T2 <- T1 + rexp(1, rate=(q21 + q23))
Z2 <- rbinom(1, size=1, prob=(q21/(q21 + q23)))

T3 <- T2 + rexp(1, rate=(q12 + q13))
Z3 <- rbinom(1, size=1, prob=(q12/(q12 + q13)))

#print( round(c(T1,Z1),6) )
#print( round(c(T2,Z2),6) )
#print( round(c(T3,Z3),6) )


to      <- c(ifelse(Z1 == 1, 2, 3), ifelse(Z2 == 1, 1, 3), ifelse(Z3 == 1, 2, 3))
from    <- c(1, to[-length(to)])
start   <- c(0,  T1, T2)
stop    <- c(T1, T2, T3)
status  <- c(1, 1, ifelse(T3 < 1, 1, 0))
sojourn <- 1:3
enum    <- rep(1,3)
id      <- rep(1,3)


dataE <- data.frame(id, from, to, start, stop, status, sojourn, enum)

dataC <- dataE
dataC$to <- ifelse((dataE$from == 1) & (dataE$to == 2), 3, dataC$to)
dataC$to <- ifelse((dataE$from == 1) & (dataE$to == 3), 2, dataC$to)
dataC$to <- ifelse((dataE$from == 2) & (dataE$to == 1), 3, dataC$to)
dataC$to <- ifelse((dataE$from == 2) & (dataE$to == 3), 1, dataC$to)
dataC$status <- rep(0,3)
dataC$enum   <- rep(2,3)


dataEC <- rbind(dataE, dataC)
dataEC <- dataEC[order(dataEC$sojourn, dataEC$enum),]
dataEC <- data.frame(dataEC, row.names=c(1:nrow(dataEC)))
print(dataEC)