library(mstate)

# function for preparing data set for mstate package
source("ext_mstate.R")

# function for bootstrap confidence intervals
source("LOS_boot.R")

# Data preparation-------------------------------------------------------------

# Data from Bhatraju et al. (2020)

# State 0: Censored
# State 1: ICU without Mechanical Ventilation and Acute Care 
# State 2: Mechanical ventilation
# State 3: Discharge
# State 4: Death


# read data set
my.data <- read.csv(file = 'Example1_Data.csv')

# Set transition matrix for mstate
tra <- transMat(x = list(c(2,3,4), c(1,4), c(),c()),
  names = c("Non-MV", "MV", "Discharge","Death"))

# Add transition vectors
my.data$trans <- NA
for (i in 1:nrow(tra))
{
  for (j in 1:ncol(tra))
  {
    my.data$trans[which(my.data$from == i & my.data$to ==j)] <- tra[i, j]  
  }
}

# Add status vector, indicates observed transition
my.data$status <- 1

# Status vector for censored observations set to '0'
my.data$status[my.data$to == 0] <- 0

# Rename 'to' == 0  to 'cens' for use in 'ext_mstate' function
my.data$to[my.data$to == 0] <- 'cens'

# Create data frame with all possible transitions for when a patient is at risk
my.data_ext <- ext_mstate(my.data, tra)

# Analysis ------------------------------------------------------------------------

## Cox model stratified by transition
c_1 <- coxph(Surv(entry, exit, status) ~ strata(trans), data= my.data_ext, method = "breslow")

# msfit calculates baseline hazards
msf_1 <- msfit(c_1, trans = tra)

# probtrans calculates transition probabilities, prediction from day 0
pt_1 <- probtrans(msf_1, predt = 0)

# Retrieve predicted mortality at day 28
# Results in Table 2

# Patients starting Non-MV
print(tail(pt_1[[1]]))

# Patients starting MV
print(tail(pt_1[[2]]))

# ELOS gives expected length of stay for 28 days in the states for patients in the states at day 0
# Results shown in Table 2
LOS_mat_1 <- ELOS(pt_1, 28)
rownames(LOS_mat_1) <- c("from Non-MV", "from MV", "from Discharge", "from Death")
colnames(LOS_mat_1) <- c("to Non-MV", "to MV","to Discharge", "to Death")
print(LOS_mat_1)

# Inital Distribution of the patients in the 4 states at day 0
init_dis_1 <- c(0.5416667, 0.4583333, 0, 0)

# Multiply LOS_mat_1 with initial distribution to get weighted average of expected lengths of stay for entire cohort
# Results shown in Table 2
LOS_cohort_1 <- (init_dis_1 %*% LOS_mat_1)
print(LOS_cohort_1)

# Create weighted average of transition probabilities for progress of entire cohort
# Used to produce full cohort results in Table 2, Figure 2, 
pt_fc_1 <- pt_1
fc_1 <- pt_fc_1[[1]] * 0.5416667 + pt_fc_1[[2]] * 0.4583333
pt_fc_1[[1]] <- fc_1

# Retrieve predicted mortality at day 28
# Results in Table 2

# All patients
print(tail(pt_fc_1[[1]]))

# Figure 2, plot of transition probabilities
plot(pt_fc_1, from = 1,  ord = c(4,2,1,3),type= "filled",cols = c("khaki1","indianred1","cornflowerblue","gray"),
  lwd= 2, xlab = "Days Since ICU Admission", ylab = "Predicted Probabilities", 
  cex.lab = 1.25, legend = c("", "", "", ""),
  main= "Predicted Proportions Over Time (Bhatraju et al.)")
text(25, 0.7, "Non-MV", cex = 1)
text(25, 0.55, "MV", cex = 1)
text(25, 0.9, "Discharge", cex = 1)
text(25, 0.25, "Death", cex = 1)  

# Bootstrapped confidence intervals

# change class of data frame
class(my.data_ext) <- c("msdata"  , "data.frame")

# give transition matrix attribute
tmat <- attr(my.data_ext,"trans")

# bootstrap performed with msboot function in mstate, 1000 samples
set.seed(1905)
Example1_boot <- msboot(theta = LOS_boot, data=my.data_ext, id="id",B=1000)

# Confidence Intervals shown in Table 2

# From non-MV, to non-MV
quantile(Example1_boot[1, ], probs=c(0.025, 0.975))

# From non-MV, to MV
quantile(Example1_boot[2, ], probs=c(0.025,0.975))

# From MV, to non-MV
quantile(Example1_boot[3 ,], probs=c(0.025, 0.975))
 
# From MV, to MV
quantile(Example1_boot[4, ], probs=c(0.025, 0.975))

# Length of stay entire cohort: Non-MV
quantile(Example1_boot[5, ], probs=c(0.025, 0.975))

# Length of stay entire cohort: MV
quantile(Example1_boot[6, ], probs=c(0.025, 0.975))