library(TransModel)
set.seed(100)

##### simulation #####
He.t <- function(t)  log(1 + t) + t^(3/2)
gen.t <- function(tt, lamb){
  dif2 <- (He.t(tt) - lamb)^2
  return(dif2)
}
gendata <- function(n, beta0 = c(1, -1), r, c){
  ctime <- c*runif(n)
  z1 <- rnorm(n, 0, 1)
  z2 <- 2*runif(n) - 1
  z <- cbind(z1, z2)
  exb <- exp(z%*%beta0)
  U <- runif(n)
  if(r == 0) Het <- -log(U)/exb
  if(r > 0)  Het <- (U^(-r) - 1)/r/exb
  dtime <- numeric()
  for(i in 1:n) dtime[i] <- optimize(gen.t, interval = c(0, 100), 
                                     lamb = Het[i])$minimum
  
  obs.t <- apply(cbind(ctime, dtime), 1, min)
  delta <- as.numeric(dtime <= ctime)
  
  outdt <- data.frame(time = obs.t, status = delta, z1 = z[, 1], z2 = z[, 2])
  return(outdt)
}

#settings
# PO(r = 1): 15%, c = 12.5;  40%, c = 2.6
# PH(r = 0): 15%, c = 5.7;   40%, c = 1.5
# r = 0.5  : 15%, c = 8.2;   40%, c = 2
# r = 2    : 15%, c = 30.5;  40%, c = 4.3

r <- 0; c <- 5.7
# rcens <- numeric()
# for(i in 1:100) {
#  dat <- gendata(n = 500, beta0 = c(1, -1), r = r, c = c)
#  rcens[i] <- 1 - mean(dat$status)
# }
# mean(rcens)

num.sim <- 1000
beta.tm <- matrix(0, nrow = num.sim, ncol = 2)
sd.tm <- matrix(0, nrow = num.sim, ncol = 2)
set.seed(350284)
for(i in 1:num.sim){
  mdata <- gendata(n = 200, beta0 = c(1, -1), r = r, c = c)
  fit <- TransModel(Surv(time, status) ~ z1 + z2, mdata, r = r)
  beta.tm[i,] <- coefficients(fit)
  sd.tm[i,] <- sqrt(diag(fit$vcov))
  cat("simu =", i, "\n")
}

# compare results 
# TransModel
est <- apply(beta.tm, 2, mean)
stderr <- apply(sd.tm, 2, mean)
stdev <- apply(beta.tm, 2, sd)
truth <- matrix(c(1, -1), nrow = num.sim, ncol = 2, byrow = TRUE)
cp <- apply(beta.tm - 1.96*sd.tm <= truth & beta.tm + 1.96*sd.tm >= truth, 2, mean)

round(cbind(est, stderr, stdev, cp), 3)


##### veteran data analysis #####
set.seed(100)
data(veteran)
veteran$celltype <- relevel(veteran$celltype, ref = "squamous")
fit <- TransModel(Surv(time, status) ~ karno + factor(celltype), data = veteran,
                  r = 1, CICB.st = TRUE, subset = (prior == 0))
fit$coefficient
fit$vcov
summary(fit)

## r = 0
fit_0 <- TransModel(Surv(time, status) ~ karno + factor(celltype), data = veteran, 
                    r = 0, CICB.st = TRUE, subset = (prior == 0))
summary(fit_0)

## r = 0.5
fit_05 <- TransModel(Surv(time, status) ~ karno + factor(celltype), data = veteran, 
                     r = 0.5, CICB.st = TRUE, subset = (prior == 0))
summary(fit_05)

## r = 2
fit_2 <- TransModel(Surv(time, status) ~ karno + factor(celltype), data = veteran, 
                     r = 2, CICB.st = TRUE, subset = (prior == 0))
summary(fit_2)

# Predict survival curves for each tumor type 
pred1 <- predict(fit, newdata = c(60, 0, 0, 0)) 
pred2 <- predict(fit, newdata = c(60, 1, 0, 0)) 
pred3 <- predict(fit, newdata = c(60, 0, 1, 0)) 
pred4 <- predict(fit, newdata = c(60, 0, 0, 1))

# Figure 1
pdf("veteran.pdf")
plot(pred1, lty = 1)
lines(pred2$time, pred2$survival, type = "s", lty = 2)
lines(pred3$time, pred3$survival, type = "s", lty = 3)
lines(pred4$time, pred4$survival, type = "s", lty = 4)
legend("topright", c("squamous", "adeno", "large", "smallcell"),
       title = "Tumor Type", lty = 1:4, bty = "n")
dev.off()

# Figure 2
pdf("cicb.pdf")
plot(pred1) 
lines(pred1$time, pred1$low.ci, lty = 2)
lines(pred1$time, pred1$up.ci, lty = 2)
lines(pred1$time, pred1$low.cb, lty = 3)
lines(pred1$time, pred1$up.cb, lty = 3)
legend("topright", c("Survival estimates", "95% CI", "95% CB"), 
       lty = 1:3, lwd = 1, bty = "n")
dev.off()

sessionInfo()