## Code for Lindgren, Rue, Lindstrom (2011), partially updated to newer R-INLA
## Note: Will not reproduce exactly the same results as in the paper.

require(rgl)
require(INLA)
require(lattice)
require(fmesher)
source("utils.R")

options(lifecycle_verbosity = "warning")

## Read data:
map <- read.table("Leuk.map")
data(Leuk)
Leuk$id <- 1:dim(Leuk)[1]

## Build triangular mesh:
loc <- cbind(Leuk$xcoord, Leuk$ycoord)
bnd1 <- fm_nonconvex_hull(loc, convex = 0.05)
bnd2 <- fm_nonconvex_hull(loc, convex = 0.25)
mesh <- (fm_mesh_2d(
    ## Data locations to use as location seeds:
    loc = cbind(Leuk$xcoord, Leuk$ycoord),
    ## Encapsulate data region:
    boundary = list(bnd1, bnd2),
    ## Refined triangulation,
    ## minimal angles >=26 degrees,
    ## interior maximal edge lengths 0.05,
    ## exterior maximal edge lengths 0.2,
    ## don't add input points closer than 0.05:
    min.angle = 24,
    max.edge = c(0.05, 0.2),
    cutoff = 0.005,
    ## Set to >=0 for visual (no effect Windows):
    plot.delay = 0.5
))
## Store the data-->vertex mapping:
## ( mesh$idx$loc[k] is the mesh vertex for data location nr. k )
Leuk$spatial <- mesh$idx$loc

## Create the SPDE/GMRF model, (kappa^2-Delta)(tau x) = W:
spde <- inla.spde2.matern(mesh, alpha = 2)

## Scale the data
Leuk$time <- Leuk$time / max(Leuk$time)
## Build the GLM model:
# Plain INLA:
formula <- inla.surv(Leuk$time, Leuk$cens) ~ 1 + sex + age + wbc + tpi +
    ## Add the spatial effect model:
    f(spatial, model = spde)

## Run INLA:
r <- (inla(formula,
    family = "weibullsurv",
    data = Leuk,
    ## We don't need the marginals:
    control.compute = list(return.marginals = FALSE),
    ## Verbose output:
    verbose = TRUE
))

if (FALSE) {
    ## In inlabru:
    formula_bru <- inla.surv(Leuk$time, Leuk$cens) ~ 1 + sex + age + wbc + tpi +
        ## Add the spatial effect model:
        spatial(cbind(Leuk$xcoord, Leuk$ycoord), model = spde)
    ## Run inlabru:
    r_bru <- (inlabru::bru(formula_bru,
        family = "weibullsurv",
        data = Leuk,
        options = list(
            ## We don't need the marginals:
            control.compute = list(return.marginals = FALSE),
            ## Verbose output:
            verbose = TRUE
        )
    ))
}

## Extract the SPDE parameters:
tau <- exp(r$summary.hyperpar[2, "mean"])
kappa <- exp(r$summary.hyperpar[3, "mean"])

## Get the precision matrix:
Q <- inla.spde2.precision(spde, theta = log(c(tau, kappa)))

## Get a random sample (not used here),
## and the index reordering,
## using an undocumented ad hoc function:
reo <- inla.qreordering(Q, reordering = "metis")
## Need to invert the indexing:
neworder <- reo$reordering
neworder[neworder] <- 1:length(neworder)
## Reorder the matrix:
Q.reordered <- Q[neworder, neworder]


## Reference point for covariance/correlation comparisons:
ref.s <- (which.min((mesh$loc[, 1] - mean(range(mesh$loc[, 1])))^2 +
    (mesh$loc[, 2] - mean(range(mesh$loc[, 2])))^2))

## Calculate covariances (S) and correlations (SS):
S <- solve(Q)
SS <- diag(1 / sqrt(diag(S))) %*% S %*% diag(1 / sqrt(diag(S)))
D <- as.matrix(dist(mesh$loc))

## Theoretical Matern correlations and covariances:
dd <- (0:1000) / 1000
SS.theory <- (dd * kappa) * besselK(dd * kappa, 1)
SS.theory[1] <- 1
S.theory <- SS.theory / (4 * pi * kappa^2) / tau^2




###########################
## Prepare for plotting:

## Calculate mapping between triangulation vertices and grid points:
## Resolution for gridded output was dims=c(100,100) in the paper.
proj <- fm_evaluator(mesh, dims = c(200, 200))

## Construct greyscale palette function:
my.grey.palette <- function(n, ...) {
    return(grey.colors(n, 0.05, 0.95, ...))
}
## Use it:
my.palette <- my.grey.palette

## Construct map data appropriate for easy plotting:
mm <- calc.map(map)


#####################
## Plot results:

## Compare correlations:
dev.new()
plot(D[ref.s, ], SS[ref.s, ],
    type = "p", pch = 20,
    xlab = "Distance", ylab = "Correlation",
    xlim = c(0.0, 1.0), ylim = c(-0.005, 1.005)
)
lines(dd, SS.theory,
    type = "l",
    col = rgb(0.5, 0.5, 0.5), lwd = 2
)


## Don't plot the full precision pattern; reduce the dimension first:
Q.coarse <- sparse.pattern.coarsen(Q.reordered, 4)
## Plot the reordered precision pattern:
dev.new()
image(as.matrix(Q.coarse) > 0,
    col = grey.colors(2, start = 1, end = 0),
    ylim = c(1, 0), axes = FALSE, xlim = c(0, 1)
)
box()



## Map resulting posterior mean field to a grid:
plotdata <- fm_evaluate(proj, r$summary.random$spatial[, "mean"])
## Plot PM contours:
dev.new()
bbb <- (levelplot(
    row.values = proj$x, column.values = proj$y, x = plotdata,
    mm = mm, panel = levelplotmap,
    col.regions = my.palette,
    xlim = range(proj$x), ylim = range(proj$y), aspect = "iso",
    contour = TRUE, cuts = 11, labels = FALSE, pretty = TRUE,
    xlab = "Easting", ylab = "Northing"
))
print(bbb)


## Map resulting posterior standard deviation field to a grid:
plotdata <- fm_evaluate(proj, r$summary.random$spatial[, "sd"])
## Plot std.dev. contours:
dev.new()
bbb <- (levelplot(
    row.values = proj$x, column.values = proj$y, x = plotdata,
    mm = mm, panel = levelplotmap,
    col.regions = my.palette,
    xlim = range(proj$x), ylim = range(proj$y), aspect = "iso",
    contour = TRUE, cuts = 11, labels = FALSE, pretty = TRUE,
    xlab = "Easting", ylab = "Northing"
))
print(bbb)


## Plot data points+map:
rgl.win1 <- open3d()
plot.rgl.map(map, color = "black", lwd = 2, zoffset = 0.1)
plot.fake.points(cbind(Leuk$xcoord, Leuk$ycoord, 0.1),
    radius = 0.005,
    color = "black", specular = "black"
)
plot_rgl(mesh,
    add = TRUE,
    lwd = 1, draw.vertices = FALSE, draw.edges = FALSE,
    col = "white"
)
view3d(0, 0, fov = 0, zoom = 0.8)

## Plot triangulation+map
rgl.win2 <- open3d()
plot.rgl.map(map, color = "black", lwd = 1, zoffset = 0.1)
plot_rgl(mesh,
    add = TRUE,
    draw.vertices = FALSE, draw.edges = TRUE,
    edge.color = rgb(0.6, 0.6, 0.6), lwd = 1, col = "white"
)
view3d(0, 0, fov = 0, zoom = 0.8)