# load packages
library(bshazard)
library(etm)

# load R-function to calculate parametric transition probabilities 
source('/home/Additional file 2.R')

# load dataset
data("los.data")

###################################################################################
# Prepare data for analysis (data is described by an extended illness-death model):
###################################################################################
# We create the columns entry (entry time into a state), 
# exit (exit time from a state), from (starting state of transition), 
# to (target state of transition)
# State 0: Admission to the ICU, state 1: Acquistion of infection,
# state 2: discharge without infection, state 3: death without infection,
# state 4: discharge with infection, state 5: death with infection
# Patients not acquiring an infection are represented by one column.
# Patients acquiring an infection by two.
my.data <- prepare.los.data(los.data)

my.data$entry<-0
my.data$exit<-my.data$time
my.data<-my.data[order(my.data$id),]
for(i in 1:(nrow(my.data)-1)){
  if(my.data$id[i]==my.data$id[i+1]){
    my.data$entry[i+1]<-my.data$exit[i]
  }
}

my.data<-my.data[,c(1:3,6:7)]
levels(my.data$to)<-c(levels(my.data$to), "4", "5")
my.data[my.data$from==1,]$to<-ifelse(my.data[my.data$from==1,]$to==2,4,5)

# Transition matrix (describes transitions in an extended illness-death model)

traEx<-matrix(FALSE, 6, 6)
traEx[1,2]<-TRUE
traEx[1,3]<-TRUE
traEx[1,4]<-TRUE
traEx[2,5]<-TRUE
traEx[2,6]<-TRUE

#### calculation of the constant hazards (exposure/occurrence rate):
#  (total number of i->j)/(total number of patient days in i)
alpha02<-sum(my.data[my.data$from==0,]$to==2)/
  sum(my.data[my.data$from==0,]$exit)
alpha03<-sum(my.data[my.data$from==0,]$to==3)/
  sum(my.data[my.data$from==0,]$exit)
alpha01<-nrow(my.data[my.data$to==1,])/
  sum(my.data[my.data$from==0,]$exit)
alpha14<-sum(my.data[my.data$from==1,]$to==4)/
  sum(my.data[my.data$from==1,]$exit-my.data[my.data$from==1,]$entry)
alpha15<-sum(my.data[my.data$from==1,]$to==5)/
  sum(my.data[my.data$from==1,]$exit-my.data[my.data$from==1,]$entry)

# total hazards:
alpha0<-(alpha01+alpha02+alpha03)
alpha1<-(alpha14+alpha15)

# estimation of transition probabilities with etm at starting
# time-points 0, 4, 10
TransProbMyData<-etm(my.data, as.character(0:5), traEx, "cens", s=0)
TransProbMyData4<-etm(my.data, as.character(0:5), traEx, "cens", s=4)
TransProbMyData10<-etm(my.data, as.character(0:5), traEx, "cens", s=10)


t<-0:100
P<-calc_transProbs_constHaz(alpha01, alpha02, alpha03, alpha14, alpha15,t)


################# new plots with LM = 4, 10 ##################

par(mfrow=c(2,3))
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 0")), type="l",
     lwd=3, ylab="Trans. Prob. 0 0", main="Stay in ICU without HAI",
     xlab="Time since ICU admission", xlim=c(0,40), ylim=c(0,1))
lines(t, P[,"P00"], lwd=3, lty=2)
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 2")), type="l",
     lwd=3, ylab="Trans. Prob. 0 2", main="Discharge without HAI",
     xlab="Time since ICU admission", xlim=c(0,40),ylim=c(0,1))
lines(t, P[,"P02"], lwd=3,lty=2)
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 3")), type="l",
     lwd=3, ylab="Trans. Prob. 0 3", main="Death without HAI",
     xlab="Time since ICU admission", xlim=c(0,40), ylim=c(0,1))
lines(t, P[,"P03"], lwd=3,lty=2)
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 1")), type="l",
     lwd=3, ylab="Trans. Prob. 0 1", main="Acquisition of HAI",
     xlab="Time since ICU admission", xlim=c(0,40), ylim=c(0,0.15))
lines(t, P[,"P01"], lwd=3, lty=2)
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 4")), type="l",
     lwd=3, ylab="Trans. Prob. 0 4", main="Discharge with HAI",
     xlab="Time since ICU admission", xlim=c(0,40), ylim=c(0,0.15))
lines(t, P[,"P04"], lwd=3, lty=2)
plot(TransProbMyData$time, trprob(TransProbMyData, tr.choice=c("0 5")), type="l",
     lwd=3, ylab="Trans. Prob. 0 5", main="Death with HAI",
     xlab="Time since ICU admission", xlim=c(0,40), ylim=c(0,0.15))
lines(t, P[,"P05"], lwd=3, lty=2)
legend("topright", c("non-parametric", "parametric"), col=c(1,1), lwd=c(3,3),
       lty=c(1,2), bty="n")

t_l<-length(t)
par(mfrow=c(2,3))
plot(TransProbMyData4$time, trprob(TransProbMyData4, tr.choice=c("1 1")), type="l",
     lwd=3, ylab="Trans. Prob. 1 1", main="Stay in ICU, given HAI at day 4",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(4,40))
lines(t[-(1:3)], P[1:(t_l-3), "P11"], lwd=3, lty=2)
plot(TransProbMyData4$time, trprob(TransProbMyData4, tr.choice=c("1 4")), type="l",
     lwd=3, ylab="Trans. Prob. 1 4", main="Discharge, given HAI at day 4",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(4,40))
lines(t[-(1:3)], P[1:(t_l-3), "P14"], lwd=3, lty=2)

plot(TransProbMyData4$time, trprob(TransProbMyData4, tr.choice=c("1 5")), type="l",
     lwd=3, ylab="Trans. Prob. 1 5", main="Death, given HAI at day 4",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(4,40))
lines(t[-(1:3)], P[1:(t_l-3), "P15"], lwd=3, lty=2)

plot(TransProbMyData10$time, trprob(TransProbMyData10, tr.choice=c("1 1")), type="l",
     lwd=3, ylab="Trans. Prob. 1 1", main="Stay in ICU, given HAI at day 10",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(10,40))
lines(t[-(1:9)], P[1:(t_l-9), "P11"], lwd=3, lty=2)


plot(TransProbMyData10$time, trprob(TransProbMyData10, tr.choice=c("1 4")), type="l",
     lwd=3, ylab="Trans. Prob. 1 4", main="Discharge, given HAI at day 10",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(10,40))
lines(t[-(1:9)], P[1:(t_l-9), "P14"], lwd=3, lty=2)
plot(TransProbMyData10$time, trprob(TransProbMyData10, tr.choice=c("1 5")), type="l",
     lwd=3, ylab="Trans. Prob. 1 5", main="Death, given HAI at day 10",
     xlab="Time since ICU admission", ylim=c(0,1), xlim=c(10,40))
lines(t[-(1:9)], P[1:(t_l-9), "P15"], lwd=3, lty=2)
legend("topright", c("non-parametric", "parametric"), col=c(1,1), lwd=c(3,3),
       lty=c(1,2), bty="n")




########## the hazard rates ##############
my.data2<-my.data
#my.data2[my.data2$exit>50,]$to<-"cens"
#my.data2[my.data2$exit>50,]$exit<-50
#death
hazard02<-bshazard(Surv(entry, exit, to==2)~1, data=my.data2[my.data2$from==0,])
# discharge
hazard03<-bshazard(Surv(entry, exit, to==3)~1, data=my.data2[my.data2$from==0,])
#inf
hazard01<-bshazard(Surv(entry, exit, to==1)~1, data=my.data2[my.data2$from==0,])
hazard14<-bshazard(Surv(entry, exit, to==4)~1, data=my.data2[my.data2$from==1,])
hazard15<-bshazard(Surv(entry, exit, to==5)~1, data=my.data2[my.data2$from==1,])

par(mfrow=c(2,3))
plot(hazard01$time, hazard01$hazard, type="l", lwd=3, main="Infection hazard", xlim=c(0, 50),
     xlab="Time since ICU admission", ylab="Hazard rate", ylim=c(0,0.15))
abline(h=alpha01, lty=2, lwd=3)
plot(hazard02$time, hazard02$hazard, type="l", lwd=3, main="Discharge without HAI", xlim=c(0, 50),
     xlab="Time since ICU admission", ylab="Hazard rate", ylim=c(0,0.15))
abline(h=alpha02, lty=2, lwd=3)
plot(hazard03$time, hazard03$hazard, type="l", lwd=3, main="Death without HAI", xlim=c(0, 50),
     xlab="Time since ICU admission", ylab="Hazard rate", ylim=c(0,0.15))
abline(h=alpha03, lty=2, lwd=3)
plot.new()
legend("center", c("non-parametric", "parametric"), col=c(1,1), lwd=c(3,3),
       lty=c(1,2), bty="n")
plot(hazard14$time, hazard14$hazard, type="l", lwd=3, main="Discharge with HAI", xlim=c(0, 50),
     xlab="Time since ICU admission", ylab="Hazard rate", ylim=c(0,0.15))
abline(h=alpha14, lty=2, lwd=3)
plot(hazard15$time, hazard15$hazard, type="l", lwd=3, main="Death with HAI", xlim=c(0, 50),
     xlab="Time since ICU admission", ylab="Hazard rate", ylim=c(0,0.15))
abline(h=alpha15, lty=2, lwd=3)


################ AM and PAF #############

D1los<-trprob(TransProbMyData, tr.choice=c("0 3"))+trprob(TransProbMyData, tr.choice=c("0 5"))
D1E0los<-trprob(TransProbMyData, tr.choice=c("0 3"))/(trprob(TransProbMyData, tr.choice=c("0 3"))+
                                                        trprob(TransProbMyData, tr.choice=c("0 2"))+trprob(TransProbMyData, tr.choice=c("0 0")))  

D1E1los<-trprob(TransProbMyData, tr.choice=c("0 5"))/(trprob(TransProbMyData, tr.choice=c("0 1"))+
                                                        trprob(TransProbMyData, tr.choice=c("0 4"))+trprob(TransProbMyData, tr.choice=c("0 5"))) 


par(mfrow=c(1,1))
plot(TransProbMyData$time, D1E0los, lwd=3, type="l", ylab="Mortality risk", 
     xlab = "Time since hospital admission", col=3, ylim=c(0,0.3))
lines(t, P[,"P(D(t)=1|E(t)=0)"], lwd=3, lty=2, col=3)
#abline(h=0.28)
lines(TransProbMyData$time, D1los, lwd=3)
lines(t, P[,"P(D(t)=1)"], lwd=3, lty=2)
lines(TransProbMyData$time, D1E1los, lwd=3, col=2)
lines(t, P[,"P(D(t)=1|E(t)=1)"], lwd=3, col=2, lty=2)
legend("bottomright", c("P(D(t)=1), non-parametric", "P(D(t)=1|E(t)=0), non-parametric", 
                        "P(D(t)=1|E(t)=1), non-parametric", "P(D(t)=1), parametric", 
                        "P(D(t)=1|E(t)=0), parametric", 
                        "P(D(t)=1|E(t)=1), parametric"), 
       col=c(1,3,2,1,3,2), lty=c(1,1,1,2,2,2),
       bty="n", lwd=c(3,3,3,3,3,3))


####### each in one plot #########

par(mfrow=c(2,2))
plot(TransProbMyData$time, D1E0los, lwd=3, type="l", ylab="Mortality risk", 
     xlab = "Time since hospital admission", col=1, ylim=c(0,0.3), 
     main="Mortality risk, given no HAI at t")
lines(t, P[,"P(D(t)=1|E(t)=0)"], lwd=3, lty=2, col=1)
#abline(h=0.28)
plot(TransProbMyData$time, D1los, lwd=3,type="l", ylab="Mortality risk", 
     xlab = "Time since hospital admission", col=1, ylim=c(0,0.3), main="Overall mortality risk")
lines(t, P[,"P(D(t)=1)"], lwd=3, lty=2)
plot(TransProbMyData$time, D1E1los, lwd=3, type="l", ylab="Mortality risk", 
     xlab = "Time since hospital admission", col=1, ylim=c(0,0.3), 
     main="Mortality risk, given HAI at t")
lines(t, P[,"P(D(t)=1|E(t)=1)"], lwd=3, lty=2)
plot.new()
legend("center", c("non-parametric", "parametric"), col=c(1,1), lwd=c(3,3),
       lty=c(1,2), bty="n")



PAFlos<-(D1los-D1E0los)/D1los
AMlos<-D1E1los-D1E0los


par(mfrow=c(1,2))
plot(TransProbMyData$time, AMlos, lwd=3, type="l", xlab="Time since hospital admission",
     ylab="Attributable mortality of HAI")
lines(t, P[,"AM"], lwd=3, lty=2)
abline(h=0, lty=3)
#abline(h=0.025, lty=3)
plot(TransProbMyData$time, PAFlos, lwd=3, type="l", xlab="Time since hospital admission",
     ylab="Population attributable fraction of HAI")
abline(h=0, lty=3)
lines(t, P[,"PAF"], lwd=3, lty=2)
legend("bottomright", c("non-parametric","parametric"), col=c(1,1), lty=c(1,2), lwd=c(3,3),
       bty="n")