# Load packages JM and lcmm
library("JM")
library("lcmm") # this can be installed using install.packages("lcmm")

# indicator for the composite event for the PBC dataset
pbc2$status2 <- as.numeric(pbc2$status != "alive")
pbc2.id$status2 <- as.numeric(pbc2.id$status != "alive")


#################
# Section 5.1.1 #
#################

lmeFit.pbc <- lme(log(serBilir) ~ drug * (year + I(year^2)),
    random = ~ year + I(year^2) | id, data = pbc2)

coxFit.pbc <- coxph(Surv(years, status2) ~ drug + hepatomegaly,
    data = pbc2.id, x = TRUE)

jointFit.pbc <- jointModel(lmeFit.pbc, coxFit.pbc, timeVar = "year",
    method = "piecewise-PH-aGH",
    interFact = list(value = ~ hepatomegaly, data = pbc2.id))

summary(jointFit.pbc)


#################
# Section 5.1.2 #
#################

prothro$t0 <- as.numeric(prothro$time == 0)
lmeFit.pro <- lme(pro ~ treat * (ns(time, 3) + t0),
    random = list(id = pdDiag(form = ~ ns(time, 3))),
    data = prothro)

coxFit.pro <- coxph(Surv(Time, death) ~ treat, data = prothros,
    x = TRUE)

# standard joint model fit m_i(t)
jointFit.pro <- jointModel(lmeFit.pro, coxFit.pro, timeVar = "time",
    method = "piecewise-PH-aGH")

# lagged effect m_i(t - c), with c = 2
jointFit2.pro <- update(jointFit.pro, lag = 2)

# confidence intervals for the coefficients in the survival submodel
confint(jointFit.pro, parm = "Event")
confint(jointFit2.pro, parm = "Event")

# compare the two models using AIC and BIC
anova(jointFit.pro, jointFit2.pro, test = FALSE)


#################
# Section 5.1.3 #
#################

# update the fit from Section 5.1.1 and exclude the
# interaction factors
jointFit2.pbc <- update(jointFit.pbc, interFact = NULL)

# list with formulas to compute the slope
dform <- list(fixed = ~ I(2*year) + drug + I(2*year):drug,
    indFixed = 3:6, random = ~ I(2*year), indRandom = 2:3)

# joint model with time-dependent slopes
jointFit3.pbc <- update(jointFit2.pbc, parameterization = "both",
    derivForm = dform)

summary(jointFit3.pbc)

# marginal Wald tests or the coefficients in the survival submodel
anova(jointFit3.pbc, process = "Event")


#################
# Section 5.1.4 #
#################

# list with formulas to compute the cumulative effect
iform <- list(fixed = ~ -1 + year + I(year * (drug == "D-penicil"))
    + I(year^2/2) + I(year^3/3) + I(year^2/2 * (drug == "D-penicil"))
    + I(year^3/3 * (drug == "D-penicil")),
    indFixed = 1:6,
    random = ~ -1 + year + I(year^2/2) + I(year^3/3),
    indRandom = 1:3)

# joint model with cumulative effect
jointFit4.pbc <- update(jointFit3.pbc, parameterization = "slope",
    derivForm = iform)

summary(jointFit4.pbc)


# function to compute weighted cumulative effects
g <- function (u, pow = 0) {
    f <- function (t)
        integrate(function (s) s^pow * dnorm(t - s), 0, t)$value
    sapply(u, f)
}

# list with formulas to compute the weighted cumulative effect
iformW <- list(fixed = ~ -1 + I(g(year)) +
    I(g(year) * (drug == "D-penicil")) +
    I(g(year, 1)) + I(g(year, 2)) +
    I(g(year, 1) * (drug == "D-penicil")) +
    I(g(year, 2) * (drug == "D-penicil")),
    indFixed = 1:6,
    random = ~ -1 + I(g(year)) + I(g(year, 1)) + I(g(year, 2)),
    indRandom = 1:3)

# joint model with weighted cumulative effect
jointFit5.pbc <- update(jointFit3.pbc, parameterization = "slope",
    derivForm = iformW)

summary(jointFit5.pbc)


#################
# Section 5.1.5 #
#################

# re-fit the mixed model with only random intercepts and
# random slopes
lmeFit2.pbc <- lme(log(serBilir) ~ year * drug,
    random = ~ year | id, data = pbc2)

# list with formulas to compute the slope
dform2 <- list(fixed = ~ 1, indFixed = 3,
    random = ~ 1, indRandom = 2)

# joint model with the slope term
jointFit6.pbc <- update(jointFit3.pbc, lmeObject = lmeFit2.pbc,
    parameterization = "slope", derivForm = dform2)

summary(jointFit6.pbc)

# list with formulas to compute the re-scaled slope
dform3 <- list(fixed = ~ -1 + I(rep(1/0.18, length(drug))),
    random = ~ -1 + I(rep(1/0.18, length(drug))),
    indFixed = 3, indRandom = 2)

# joint model with the re-scaled slope term
jointFit7.pbc <- update(jointFit6.pbc, derivForm = dform3)

summary(jointFit7.pbc)


###############
# Section 5.2 #
###############

# construct the 'start', 'stop' and 'event' variables
pbc <- pbc2[c("id", "serBilir", "drug", "year", "years",
    "status2", "spiders")]

pbc$start <- pbc$year

splitID <- split(pbc[c("start", "years")], pbc$id)
pbc$stop <- unlist(lapply(splitID,
    function (d) c(d$start[-1], d$years[1]) ))

pbc$event <- with(pbc, ave(status2, id,
    FUN = function (x) c(rep(0, length(x)-1), x[1]) ))


# time-dependent Cox model fit to be included in jointModel()
tdCox.pbc <- coxph(Surv(start, stop, event) ~ drug * spiders + cluster(id),
        data = pbc, x = TRUE, model = TRUE)

# joint model fit with exogenous covariate
jointFit8.pbc <- jointModel(lmeFit.pbc, tdCox.pbc, timeVar = "year",
    method = "spline-PH-aGH")

summary(jointFit8.pbc)

# time-dependent Cox model fit, including log serum bilirubin
tdCox2.pbc <- coxph(Surv(start, stop, event) ~ drug * spiders +
    log(serBilir), data = pbc)

tdCox2.pbc

# confidence intervals for hazard ratios under the joint and
# Cox models
exp(confint(jointFit8.pbc, parm = "Event"))
exp(confint(tdCox2.pbc))


###############
# Section 5.3 #
###############

lmeFit.pbc <- lme(log(serBilir) ~ drug * (year + I(year^2)),
    random = ~ year + I(year^2) | id, data = pbc2)

coxFit2.pbc <- coxph(Surv(years, status2) ~ drug + strata(hepatomegaly),
    data = pbc2.id, x = TRUE)

jointFit9.pbc <- jointModel(lmeFit.pbc, coxFit2.pbc, timeVar = "year",
    method = "spline-PH-aGH")

summary(jointFit9.pbc)

# Wald test for stratification factors
wald.strata(jointFit9.pbc)


# include interactions with strata
coxFit3.pbc <- coxph(Surv(years, status2) ~ drug * hepatomegaly +
    strata(hepatomegaly), data = pbc2.id, x = TRUE)

jointFit10.pbc <- update(jointFit9.pbc, survObject = coxFit3.pbc,
    interFact = list(value = ~ hepatomegaly, data = pbc2.id))

summary(jointFit10.pbc)

# plot of marginal survival functions for the two strata
plot(jointFit9.pbc, 3, lwd = 3)


###############
# Section 5.4 #
###############

# latent class joint models are fitted with the functions
# from the 'lcmm' package

# latent class joint model with 3 classes
# Note: version 1.4-3 of lcmm used in the book does not automatically exclude missing
# data, and does not create dummy variables for factors; the user is required to do
# these operations manually
aidsLC <- aids[c("patient", "CD4", "obstime", "drug", "Time", "death")]
aidsLC$drug <- c(aidsLC$drug) - 1
aidsLC <- aidsLC[complete.cases(aidsLC), ]

lcjmFit.aids <- Jointlcmm(fixed = CD4 ~ obstime + drug,
    mixture = ~ obstime + drug, random = ~ obstime,
    classmb = ~ drug, subject = "patient", ng = 3, data = aidsLC,
    survival = Surv(Time, death) ~ mixture(drug),
    hazard = "6-quant-piecewise", hazardtype = "Specific")

summary(lcjmFit.aids)

# posterior probabilities
postprob(lcjmFit.aids)


#################
# Section 5.5.1 #
#################

# convert data to the competing risks long format
pbc2.idCR <- crLong(pbc2.id, statusVar = "status",
    censLevel = "alive", nameStrata = "CR")


lmeFit.pbc <- lme(log(serBilir) ~ drug * (year + I(year^2)),
    random = ~ year + I(year^2) | id, data = pbc2)

coxFit4.pbc <- coxph(Surv(years, status2) ~ (drug + age) * CR + strata(CR),
        data = pbc2.idCR, x = TRUE)

# joint model for competing risks data
jointFit11.pbc <- jointModel(lmeFit.pbc, coxFit4.pbc,
    timeVar = "year", method = "spline-PH-aGH", CompRisk = TRUE,
    interFact = list(value = ~ CR, data = pbc2.idCR))

summary(jointFit11.pbc)

# include time-dependent slopes and interaction terms
dform <- list(fixed = ~ I(2*year) * drug, indFixed = 3:6,
    random = ~ I(2*year), indRandom = 2:3)

jointFit12.pbc <- update(jointFit11.pbc,
    parameterization = "both", derivForm = dform,
    interFact = list(value = ~ CR, slope = ~ CR, data = pbc2.idCR))

summary(jointFit12.pbc)

# likelihood ratio test
anova(jointFit11.pbc, jointFit12.pbc)


###############
# Section 5.6 #
###############

lmeFit.aids <- lme(CD4 ~ obstime + obstime:drug,
    random = ~ obstime | patient, data = aids)

coxFit.aids <- coxph(Surv(Time, death) ~ drug + gender + AZT,
    data = aids.id, x = TRUE)

# joint model with Weibull baseline hazard under PH formulation
jointFit3.aids <- jointModel(lmeFit.aids, coxFit.aids,
    timeVar = "obstime")

# joint model with Weibull baseline hazard under AFT formulation
jointFit4.aids <- update(jointFit3.aids, method = "weibull-AFT-aGH")

# compare the two models using AIC and BIC
anova(jointFit3.aids, jointFit4.aids, test = FALSE)