library("MASS")
library("survival")
library("INLA")
# Section 2
# Example 1
inla.setOption(short.summary = TRUE)
data("aids", package = "JMbayes")
par(mfrow = c(1, 2))
interaction.plot(aids$obstime[1:100], aids$patient[1:100], aids$CD4[1:100], 
                 xlab = "Time(years)", ylab = "CD4 count", 
                 legend = FALSE, col = 1:100)
hist(aids$CD4, main = "", xlab = "CD4 count")
data1 <- aids
mtime <- max(data1$Time, data1$obstime)
mtime
data1$Time <- data1$Time / mtime
data1$obstime <- data1$obstime / mtime
datas <- data1[data1$obstime == 0, ]
datal <- data1[, c(1, 4:12)]
ns <- nrow(datas)
nl <- nrow(datal)
N <- length(unique(data1$patient))

fixed.covariate <- data.frame(
  beta0 = as.factor(c(rep(1, nl), rep(2, ns))), 
  l.drug = as.factor(c(as.factor(datal$drug), rep(NA, ns))), 
  l.gender = as.factor(c(as.factor(datal$gender), rep(NA, ns))), 
  l.prevOI = as.factor(c(as.factor(datal$prevOI), rep(NA, ns))), 
  l.AZT = as.factor(c(as.factor(datal$AZT), rep(NA, ns))), 
  s.drug = as.factor(c(rep(NA, nl), as.factor(datas$drug))), 
  s.gender = as.factor(c(rep(NA, nl), as.factor(datas$gender))), 
  s.prevOI = as.factor(c(rep(NA, nl), as.factor(datas$prevOI))), 
  s.AZT = as.factor(c(rep(NA, nl), as.factor(datas$AZT))), 
  l.time = c(datal$obstime, rep(NA, ns)), 
  s.time = c(rep(NA, nl), datas$Time))

random.covariate <- list(U11 = c(datal$patient, rep(NA, ns)), 
                         U21 = c(datal$patient, rep(NA, ns)), 
                         U12 = c(rep(NA, nl), datas$patient), 
                         U22 = c(rep(NA, nl), datas$patient))

joint.dataCD4 <- c(fixed.covariate, random.covariate)

y.long <- c(round(datal$CD4), rep(NA, ns))
y.surv <- inla.surv(time = c(rep(NA, nl), datas$Time), 
                    event = c(rep(NA, nl), rep(1, ns) - datas$death))
Yjoint <- list(y.long, y.surv)
joint.dataCD4$Y <- Yjoint

JointmodelCD4 <- inla(Y ~ -1 + beta0 + l.gender + l.drug + l.prevOI +
                          l.AZT + s.gender + s.drug + s.prevOI + s.AZT + 
                          f(U11, model = "iid2d", n = 2 * length(joint.dataCD4$beta0)) +
                          f(U21, l.time, copy = "U11", fixed = TRUE) + 
                          f(U12, copy = "U11", fixed = FALSE) +
                          f(U22, s.time, copy = "U11", fixed = FALSE), 
                      family = c("poisson", "weibullsurv"), 
                      data = joint.dataCD4, verbose = FALSE, 
                      control.compute = list(dic = TRUE), 
                      control.family = list(list(), 
                                            list(hyper = list(
                                                     theta = list(
                                                         prior = "pc.alphaw", 
                                                         param = 1)),
                                                 variant = 1)))

summary(JointmodelCD4)

datas <- data1[data1$obstime == 0, ] # Observed event times
datal <- data1[, c(1, 4:12)] # Observed longitudinal responses and covariates

# Make a dataframe for the estimated survival functions
dataH <- data.frame(datas, 
                    lambda1 = JointmodelCD4$summary.fitted.values$mean[(nl+1):(nl+ns)])
# Make a dataframe for the estimated longitudinal trajectories
dataL1 <- data.frame(datal, 
                     fitted_l = JointmodelCD4$summary.fitted.values$mean[1:nl], 
                     random_l = JointmodelCD4$summary.random$U11$mean[1:nl], 
                     randoms_l = JointmodelCD4$summary.random$U21$mean[1:nl])

patients <- c(4, 35)
par(mfrow = c(2, 2), mar = c(4, 4, 4, 4))
alpha <- JointmodelCD4$summary.hyperpar$mean[1]
j_est <- JointmodelCD4$summary.hyperpar$mean[2:6]
f_est <- JointmodelCD4$summary.fixed$mean

for (patientnr in patients) {
  dataHi <- dataH[dataH$patient == patientnr, ]
  # Observed trajectory
  with(datal[datal$patient == patientnr, ], plot(obstime * mtime, CD4,
       ylab = "CD4 count", xlab = "Time (months)", type = "l", xlim = c(0, 21.4), 
       ylim = c(0, 20), main = paste("CD4 trajectory - patient", patientnr)))
  # Model-based trajectory
  with(dataL1[dataL1$patient == patientnr, ], lines(obstime * mtime, 
    fitted_l + random_l + randoms_l * obstime, col = "blue", lty = 2))
  
  # Model-based survival curve
  pred_time <- seq(0, 1, by = 0.01)
  lambda <- dataHi$lambda1 
  # Estimated survival function
  plot((pred_time * mtime),
       exp(-(pred_time * lambda)^alpha), 
       type = "l", ylab = "Survival probability", xlab = "Time (months)", 
       main = paste("Survival curve - patient", patientnr))
  abline(h = 0.5, col = "red")
}

# Example 2
library("JointModel")
inla.setOption(short.summary = TRUE)
data1 <- prostate
data2 <- dropout
ng <- nrow(data1)
ns <- nrow(data2)
data1 <- data1[order(data1$VisitTime), ] 

joint.dataPSA <- readRDS(system.file("exampledata/psa/jointdataPSA.rds", 
                                     package = "INLA"))
joint.dataPSA$beta0 <- joint.dataPSA$mu

JointmodelPSA <-
    inla(Y ~ -1 + beta0 +
             f(inla.group(V1, n = 50), model = "rw2", scale.model = TRUE,
               hyper = list(prec = list(prior = "pc.prec", param = c(1, 0.01)))) +
             b13.PSAbase +
             f(u, w, model = "iid", hyper = list(prec = list(initial = -6, fixed = TRUE))) +
             f(b.eta, copy = "u", hyper = list(beta = list(fixed = FALSE))),
         family = c("gaussian", "gaussian", "weibullsurv"),
         data = joint.dataPSA, verbose = FALSE,
         control.compute = list(dic = TRUE, config = TRUE),
         control.family = list(list(),
                               list(hyper = list(
                                        prec = list(initial = 10, 
                                                    fixed = TRUE))),
                               list()))
summary(JointmodelPSA)
par(mfrow = c(1, 1))
plot(data1$VisitTime, data1$logPSA.postRT, xlab = "Time (years)", 
     ylab = "log(PSA+0.1)", col = "lightgrey") 
points(data1$VisitTime, JointmodelPSA$summary.fitted.values[1:ng, 1], 
       col = "blue", lwd = 1)
lines(JointmodelPSA$summary.random$`inla.group(V1, n = 50)`[1:50, 1], 
      JointmodelPSA$summary.random$`inla.group(V1, n = 50)`[1:50, 2], 
      col = "red", lwd = 3)

# Section 3
library("fields")
library("viridisLite")
set.seed(2019)
t.max <- 8
mesh.time <- inla.mesh.1d(1:t.max)
fake.locations <- matrix(c(0, 0, 10, 10, 0, 10, 10, 0), nrow = 4, byrow = TRUE)
mesh.space <- inla.mesh.2d(loc = fake.locations, max.edge = c(1.5, 2))
# Hyper-parameters of our random effects
# For the separable model
range.time <- 20 
range.space <- 6
sigma.u <- 1
# For the non-separable model
gt <- 2.23
gs2 <- 0.2222
ge2 <- 0.0805
# Pre-processing: Spatial and temporal FEM matrices
sfe <- inla.mesh.fem(mesh.space, order = 4)
tfe <- inla.mesh.fem(mesh.time, order = 2)
M0 <- tfe$c0
N.t <- nrow(M0)
M1 <- sparseMatrix(i = c(1, N.t), j = c(1, N.t), x = 0.5)
M2 <- tfe$g1
kappa <- 2/range.time
Q.M <- kappa^2*M0 + 2*kappa*M1 + M2
Q.M <- Q.M/2/kappa
Q.space.alpha2 <- gs2^2*sfe$c0 + 2*gs2*sfe$g1 + sfe$g2
Q.space.alpha2 <- Q.space.alpha2/(4*pi*gs2)
Q.separ <- kronecker(Q.M, Q.space.alpha2)
Q.nonsep <- (kronecker(gt^2*M2, gs2*sfe$c0 + sfe$g1) + 
             kronecker(M0, gs2^3*sfe$c0 + gs2^2*sfe$g1 +
                           gs2*sfe$g2 + sfe$g3) +
             kronecker(2*gt*M1, gs2^2*sfe$c0 +
                                2*gs2*sfe$g1 + sfe$g2 )) * ge2
summary(diag(inla.qinv(Q.separ)))
summary(diag(inla.qinv(Q.nonsep)))
# Prior simulation
u <- inla.qsample(n = 1, Q.nonsep, seed = 2019, num.threads = 1)[, 1]
N.st <- length(u)
# == nrow(Q.separ)
# == nrow(Q.nonsep)
sig.eps <- 0.01
noise <- rnorm(N.st, 0, 1) * sig.eps
df <- data.frame(y = u*sigma.u + noise, st = 1:N.st)
df$y[-(1:mesh.space$n)] <- NA
# Precision matrix for observation noise
Qeps <- Diagonal(n = mesh.space$n)
# Projection matrix for the observed latent field
A.observe <- sparseMatrix(i = 1:mesh.space$n, j = 1:mesh.space$n, 
                          dims = c(mesh.space$n, N.st))
# Posterior/conditional precision matrix
post.Q <- function(sig.eps = 0.01, Q.model) {
  Q <- sig.eps^-2*t(A.observe) %*% Qeps %*% A.observe + Q.model
  return(Q)
}
# Point estimate (posterior mean)
post.mu <- function (sig.eps = 0.01, Q.model) {
  a <- df$y[1:mesh.space$n]
  b <- sig.eps^-2 * t(A.observe) %*% Qeps %*% a
  res <- inla.qsolve(post.Q(sig.eps, Q.model = Q.model), b)
  return(res)
} 
mu.post.separ <- post.mu(Q.model = Q.separ)
mu.post.nonsep <- post.mu(Q.model = Q.nonsep)
local.plot.field <- function(field, mesh, time = 1, ...) {
  # Only use the relevant part of the incoming vector
  field <- field[1:mesh$n + (time - 1) * mesh$n]
  # Project the mesh onto a 200x200 grid
  proj <- inla.mesh.projector(mesh, dims = c(200, 200))
  field.proj <- inla.mesh.project(proj, field)
  image.plot(list(x = proj$x, y = proj$y, z = field.proj), 
             col = plasma(64), ...)  
}
par(mfrow = c(3, 2))
zlim2 <- range(c(mu.post.separ, mu.post.nonsep))
for (tp in 1:3) {
  local.plot.field(mu.post.separ, mesh.space, time = tp, 
                   main = paste0("Separable mean, t = ", tp), 
                   xlim = c(0, 10), ylim = c(0, 10), zlim = zlim2)
  local.plot.field(mu.post.nonsep, mesh.space, time = tp, 
                   main = paste0("Non-separable mean, t = ", tp), 
                   xlim = c(0, 10), ylim = c(0, 10), zlim = zlim2)
}
# Posterior simulations
post.sim.separ <- inla.qsample(1, Q = post.Q(Q.model = Q.separ), 
                               reordering = "identity", seed = 1, 
                               num.threads = 1)
post.sim.separ <- drop(post.sim.separ + mu.post.separ)
post.sim.nonsep <- inla.qsample(1, Q = post.Q(Q.model = Q.nonsep), 
                                reordering = "identity", seed = 1, 
                                num.threads = 1)
post.sim.nonsep <- drop(post.sim.nonsep + mu.post.nonsep)
# Plot the simulations
zlim1 <- range(c(post.sim.separ, post.sim.nonsep))
par(mfrow = c(3, 2))
for (tp in c(1, 2, 3)) {
  local.plot.field(post.sim.separ, mesh.space, time = tp, 
                   main = paste0("Separable sim, t = ", tp), 
                   xlim = c(0, 10), ylim = c(0, 10), zlim = zlim1)
  local.plot.field(post.sim.nonsep, mesh.space, time = tp, 
                   main = paste0("Non-separable sim, t = ", tp), 
                   xlim = c(0, 10), ylim = c(0, 10), zlim = zlim1)
}

sessionInfo()