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 Grein et al. (2020)

# State 0: Censored
# State 1: Noninvasive Positive Pressure ventilation, Nasal High-Flow Oxygen therapy, Low-Flow Oxygen, Ambient Air 
# State 2: Extracorporeal Membrane Oxygenation, Mechanical Ventilation
# State 3: Discharge
# State 4: Death

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

# Set transition matrix for mstate
tra <- transMat(x = list(c(2,3,4), c(1,3,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_2 <- coxph(Surv(entry, exit, status) ~ strata(trans), data= my.data_ext, method = "breslow")

# msfit calculates baseline hazards
msf_2 <- msfit(c_2, trans = tra)

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

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

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

# Patients starting MV
print(tail(pt_2[[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 3
LOS_mat_2 <- ELOS(pt_2, 28)
rownames(LOS_mat_2) <- c("from Non-MV", "from MV", "from Discharge", "from Death")
colnames(LOS_mat_2) <- c("to Non-MV", "to MV","to Discharge", "to Death")
print(LOS_mat_2)

# Inital Distribution of the patients in the 4 states at day 0
init_dis_2 <- c(0.3584906, 0.6415094, 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 3
LOS_cohort_2 <- (init_dis_2 %*% LOS_mat_2)
print(LOS_cohort_2)

# Create weighted average of transition probabilities for progress of entire cohort
# Used to produce full cohort results in Table 3 
pt_fc_2 <- pt_2
fc_2 <- pt_fc_2[[1]] * 0.3584906 + pt_fc_2[[2]] * 0.6415094
pt_fc_2[[1]] <- fc_2

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

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

# Figure 3, plot of transition probabilities
plot(pt_fc_2, from = 1,  ord = c(4,2,1,3),type= "filled",cols = c("khaki1","indianred1","cornflowerblue","gray"),
     lwd= 2, xlab = "Days Since Treatment Initiation", ylab = "Predicted Probabilities", 
     cex.lab = 1.25, legend = c("", "", "", ""),
     main= "Predicted Proportions Over Time (Grein et al.)")
text(3, 0.8, "Non-MV", cex = 1)
text(3, 0.3, "MV", cex = 1)
text(25, 0.9, "Discharge", cex = 1)
text(25, 0.10, "Death", cex = 1)  

# Figure 4, plot of transition probabilities for patients starting Non-MV
plot(pt_2, from = 1,ord = c(4,2,1,3),type= "filled", cols = c("khaki1","indianred1", "cornflowerblue", "gray"),
     lwd= 2, xlab = "Days Since Treatment Initiation", ylab = "Predicted Probabilities", cex.lab = 1.25,  
     legend = c("", "", "", ""), main= "Patients starting Non-MV (Grein et al.)")
text(5, 0.5, "Non-MV", cex = 1)
text(15.5, 0.085, "MV", cex = 1)
text(23, 0.8, "Discharge", cex = 1)
text(25, 0.04, "Death", cex = 1)

# Figure 4, plot of transition probabilities for patients starting MV
plot(pt_2, from = 2,ord = c(4,2,1,3),type= "filled", cols = c("khaki1","indianred1", "cornflowerblue", "gray"),
  lwd= 2, xlab = "Days Since Treatment Initiation",ylab = "Predicted Probabilities", cex.lab = 1.25,  
  legend = c("", "", "", ""), main= "Patients starting MV (Grein et al.)")
text(5, 0.5, "MV", cex = 1)
text(20, 0.60, "Non-MV", cex = 1 )
text(23, 0.8, "Discharge", cex = 1 )
text(25, 0.10, "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)
Example2_boot <- msboot(theta = LOS_boot, data=my.data_ext, id="id",B = 1000)

# Confidence Intervals shown in Table 3

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

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

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

# From MV, to MV
quantile(Example2_boot[4, ], probs=c(0.025, 0.975))

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

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