library(mstate)

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


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

# Data from Grein et al.

# State 0: Censored
# State 1: ECMO 
# State 2: mechanical ventilation
# State 3: Non-invasive oxygen support
# State 4: Ambient Air
# State 5: Dead
# State 6: Discharged

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

# Set transition matrix for mstate
tra <- transMat(x = list(c(2,3), c(1,3,4,5,6), c(2,4,5,6), c(3,6),c(),c()),
  names = c("ECMO", "MV", "Non-Inv Oxy","Ambient", "Dead","Discharged"))

# 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
cG <- coxph(Surv(entry, exit, status) ~ strata(trans), data= my.data_ext, method = "breslow")

# msfit calculates baseline hazards
msfG <- msfit(cG, newdata = newd, trans = tra)

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

# ELOS gives expected length of stay in the states for patients in the states at day 0
LOS_matG <- ELOS(ptG, 28)
rownames(LOS_matG) <- c("from ECMO", "from MV", "from NI Oxy", "from Ambient",  "from Death", "from Discharge")
colnames(LOS_matG) <- c("to ECMO", "to MV", "to NI Oxy", "to Ambient",  "to Death", "to Discharge")
print(LOS_matG)

# Distribution of the patients in the 6 states at day 0
init_dis <- c(0.0754717, 0.5660377, 0.3207547, 0.03773585, 0, 0)

# Multiply LOS_matG with initial distribution to get weighted average of expected lengths of stay for entire cohort
LOS_cohort_G <- (init_dis %*% LOS_matG)
print(LOS_cohort_G)

## Create weighted average of transition probabilities for progress of entire cohort
pt.acc <- ptG
acc <- ((pt.acc[[1]] * 0.0754717) + (pt.acc[[2]] * 0.5660377) + (pt.acc[[3]] * 0.3207547) + (pt.acc[[4]] * 0.03773585))
pt.acc[[1]] <- acc

## Plot of transition probabilities
plot(pt.acc, from = 1,  ord = c(5,1,2,3,4,6),type= "filled",
  cols = c("indianred1","darkgoldenrod1","khaki2", "bisque2", "gray","cornflowerblue"),
  lwd= 2, xlab = "Days Since Treatment Initiation", ylab = "Predicted Probabilities", 
  cex.lab = 1.25,main= "Predicted Proportions Over Time (Grein et al.)", legend = c("", "", "", "", "", ""))
text(5, 0.04, "ECMO", cex=1)
text(5, 0.70, "Noninvasive Ventilation", cex=1)
text(5, 0.37, "Invasive Ventilation", cex=1)
text(11, 0.78, "Ambient Air", cex=1)
text(25, 0.8, "Discharge", cex=1)
text(25, 0.08, "Death", cex=1)