## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4.5
)
set.seed(20260729)

## ----design-------------------------------------------------------------------
library(Matrix)
library(mvtnorm)
library(ZINB.GP)

make_design <- function(n_space, n_time, mean_replicates) {
  cell_n <- matrix(
    rpois(n_space * n_time, mean_replicates),
    nrow = n_space
  )
  cells <- expand.grid(
    spatial = seq_len(n_space),
    temporal = seq_len(n_time)
  )

  spatial_id <- rep(cells$spatial, times = as.vector(cell_n))
  temporal_id <- rep(cells$temporal, times = as.vector(cell_n))
  n <- length(spatial_id)

  Vs_full <- as.matrix(sparseMatrix(
    i = seq_len(n), j = spatial_id, x = 1,
    dims = c(n, n_space)
  ))
  Vt_full <- as.matrix(sparseMatrix(
    i = seq_len(n), j = temporal_id, x = 1,
    dims = c(n, n_time)
  ))

  list(Vs_full = Vs_full, Vt_full = Vt_full, cell_n = cell_n)
}

n_space <- 12
n_time <- 8
design <- make_design(n_space, n_time, mean_replicates = 3)
N <- nrow(design$Vs_full)
N

## ----distances----------------------------------------------------------------
coords <- cbind(runif(n_space), runif(n_space)) * 1000
time_coord <- matrix(0:(n_time - 1) * 50, ncol = 1)

Ds <- as.matrix(dist(coords))
Dt <- as.matrix(dist(time_coord))

Vs <- design$Vs_full[, -1, drop = FALSE]
Vt <- design$Vt_full[, -1, drop = FALSE]

stopifnot(
  ncol(Vs) + 1 == nrow(Ds),
  ncol(Vt) + 1 == nrow(Dt),
  nrow(Vs) == nrow(Vt)
)

## ----simulate-----------------------------------------------------------------
noisy_covariance <- function(distance, length_scale, sigma, kappa) {
  correlation <- exp(-(distance^2) / length_scale^2)
  sigma^2 * (kappa * correlation + (1 - kappa) * diag(nrow(distance)))
}

spatial_distance <- Ds[-1, -1, drop = FALSE]
temporal_distance <- Dt[-1, -1, drop = FALSE]

a <- drop(rmvnorm(1, sigma = noisy_covariance(
  spatial_distance, length_scale = 350, sigma = 1, kappa = 0.5
)))
c <- drop(rmvnorm(1, sigma = noisy_covariance(
  spatial_distance, length_scale = 250, sigma = 1, kappa = 0.5
)))
b <- drop(rmvnorm(1, sigma = noisy_covariance(
  temporal_distance, length_scale = 100, sigma = 0.5, kappa = 0.2
)))
d <- drop(rmvnorm(1, sigma = noisy_covariance(
  temporal_distance, length_scale = 150, sigma = 0.5, kappa = 0.2
)))

x <- rnorm(N)
X <- cbind("(Intercept)" = 1, x = x)
alpha <- c(-0.25, 0.25)
beta <- c(0.50, -0.25)
r <- 1

eta_at_risk <- drop(X %*% alpha + Vs %*% a + Vt %*% b)
p_at_risk <- plogis(eta_at_risk)
at_risk <- rbinom(N, size = 1, prob = p_at_risk)

eta_count <- drop(X %*% beta + Vs %*% c + Vt %*% d)
mu_count <- r * exp(eta_count)
y <- integer(N)
y[at_risk == 1] <- rnbinom(
  sum(at_risk == 1),
  size = r,
  mu = mu_count[at_risk == 1]
)

c(observations = N, zeros = sum(y == 0), positive = sum(y > 0))

## ----simulated-heatmap--------------------------------------------------------
cell_id <- max.col(design$Vs_full) +
  n_space * (max.col(design$Vt_full) - 1)
cell_sum <- tapply(y, cell_id, sum)
cell_total <- numeric(n_space * n_time)
cell_total[as.integer(names(cell_sum))] <- cell_sum
cell_total <- matrix(cell_total, nrow = n_space)

image(
  x = seq_len(n_time),
  y = seq_len(n_space),
  z = t(log1p(cell_total)),
  xlab = "Time",
  ylab = "Spatial location",
  main = "Simulated cell totals: log(1 + count)",
  col = hcl.colors(20, "YlOrRd", rev = TRUE)
)

## ----fit, eval=FALSE----------------------------------------------------------
#  fit <- ZINB_GP(
#    X = X,
#    y = y,
#    coords = coords,
#    Vs = Vs,
#    Vt = Vt,
#    Ds = Ds,
#    Dt = Dt,
#    nsim = 20000,
#    burn = 5000,
#    thin = 5,
#    save_ypred = TRUE,
#    print_progress = TRUE,
#    use_count_gp = TRUE,
#    use_inflation_gp = TRUE
#  )

## ----count-only, eval=FALSE---------------------------------------------------
#  fit_count_gp <- ZINB_GP(
#    X = X, y = y, coords = coords,
#    Vs = Vs, Vt = Vt, Ds = Ds, Dt = Dt,
#    nsim = 20000, burn = 5000, thin = 5,
#    use_count_gp = TRUE,
#    use_inflation_gp = FALSE
#  )

## ----prediction-inputs, eval=FALSE--------------------------------------------
#  coords_future <- rbind(
#    c(250, 250),
#    c(500, 500),
#    c(750, 750)
#  )
#  times_future <- matrix(c(400, 450), ncol = 1)
#  prediction_grid <- expand.grid(
#    spatial = seq_len(nrow(coords_future)),
#    temporal = seq_len(nrow(times_future))
#  )
#  
#  coords_new <- coords_future[prediction_grid$spatial, , drop = FALSE]
#  time_coords_new <- times_future[
#    prediction_grid$temporal, ,
#    drop = FALSE
#  ]
#  X_new <- cbind(
#    "(Intercept)" = 1,
#    x = rep(0, nrow(prediction_grid))
#  )
#  
#  prediction_inputs <- make_prediction_inputs(
#    coords = coords,
#    time_coords = time_coord,
#    coords_new = coords_new,
#    time_coords_new = time_coords_new
#  )

## ----predict-new, eval=FALSE--------------------------------------------------
#  predicted <- do.call(
#    predict,
#    c(
#      list(object = fit, X = X_new),
#      prediction_inputs
#    )
#  )
#  
#  apply(predicted$Y_pred, 2, quantile, probs = c(0.025, 0.5, 0.975))

## ----intervals, eval=FALSE----------------------------------------------------
#  apply(fit$Alpha, 2, quantile, probs = c(0.025, 0.5, 0.975))
#  apply(fit$Beta, 2, quantile, probs = c(0.025, 0.5, 0.975))