---
title: "Case-crossover models with adlaplace"
author: "Patrick Brown"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{Case-crossover models with adlaplace}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

## Introduction

This vignette fits a time-stratified case-crossover model to daily mortality and
PM10 in London (2002–2006), using the example dataset from
[Tobias et al. (2024)](https://github.com/aureliotobias/casecrossover).

Strata are defined by day-of-week within month and year via
`dirichlet_multinom(..., by = "year_month_dow")`. PM10 enters through an
integrated Wiener process smooth (`iwp(pm10, ...)`), with relative humidity as
a linear fixed effect. Fitting is done with `adlaplace()`.

```{r setup, include=FALSE}
library(adlaplace)
```

## Data

```{r data}
data("london", package = "adlaplace")
head(london)
quantile(london$pm10, na.rm = TRUE)
quantile(london$tmean, na.rm = TRUE)
```

## Model specification

Case-crossover strata absorb the intercept, so the formula uses `0 +` on the
right-hand side.

```{r model-spec}
london$temp_div10 <- london$tmean / 10
cc_formula <- dirichlet_multinom(
    all,
    by = c("year", "month", "dow"),
    init = 0.2,
    parscale = 5
) ~ 0 + rh + iwp(
    pm10,
    p = 2,
    knots = seq(0, 100, by = 10),
    ref_value = 40,
    init = 0.001, parscale = 100
) + iwp(
    tmean,
    p = 3, knots = seq(-2, 30, by = 2),
    init = 0.01, parscale = 100,
    ref_value = 20
)
```


## Fit

```{r fit}
fit <- adlaplace(
    formula = cc_formula,
    data = london,
    num_threads = 2L,
    num_groups = 50L,
    config = list(transform_theta = TRUE),
    control = list(
        maxit = 200, trace = 0, REPORT = 1,
        factr = 1e3,
        pgtol = 1e-8
    ),
    verbose = FALSE
)
```

```{r summary}
fit
summary(fit)
```


```{r plot_results}
num_sim <- 200L
num_grid <- 101L
gamma_sims <- rmvnldl(n = num_sim, fit = fit)

smooth_vars <- unique(fit$model_data$term_data$info$gamma$term)


par(mfrow = c(1, 2), mar = c(2.25, 2.25, 0, 0), bty = "l", mgp = c(1, 0.5, 0))
for (D in smooth_vars) {
    sim_here <- sim_random(
        D,
        fit = fit, gamma_sims = gamma_sims, num_grid = num_grid
    )
    rr <- exp(sim_here$sim)
    if (requireNamespace("GET", quietly = TRUE)) {
        q <- GET::central_region(
            GET::create_curve_set(list(obs = rr)),
            coverage = 0.8
        )[, c("lo", "central", "hi")]
    } else {
        q <- t(apply(rr, 1, quantile, probs = c(0.1, 0.5, 0.9)))
    }


    matplot(
        sim_here$x, rr,
        type = "l", lty = 1, col = "#00000010",
        xlab = D, ylab = "RR", xaxs = "i", yaxs = "i",
        ylim = quantile(rr, prob = c(0.005, 0.95))
    )

    matlines(sim_here$x, q, lty = c(2, 1, 2), col = "red")
}
```


## AD development check

Build `model_data`, shard map, and `ad_pack` without running outer optimization,
then maximize the Laplace profile likelihood with `optim` directly.

```{r for_dev}
md <- model_data(cc_formula, data = london, verbose = FALSE)
config_dev <- list(
    transform_theta = TRUE,
    num_threads = 2L,
    verbose = FALSE
)
config_dev$obs_groups <- obs_groups(
    A = md$term_data$A,
    elgm_matrix = md$term_data$elgm_matrix,
    num_groups = 50L,
    min_groups = min(50L, 2L * 4L)
)
par_info <- md$term_data$info$parameters
cache <- new.env(parent = emptyenv())
cache$gamma <- rep(0, nrow(md$term_data$info$gamma))
control_dev <- list(
    maxit = 200, trace = 0, REPORT = 1,
    factr = 1e3,
    pgtol = 1e-8,
    parscale = par_info$parscale
)
control_inner <- list(report.level = 0)
af <- ad_pack(md, config_dev, num_threads = 2L)
slotNames(af)
dim(config_dev$obs_groups)
par_info
bob <- stats::optim(
    par = par_info$init,
    fn = outer_fn,
    gr = outer_gr,
    method = "L-BFGS-B",
    control = control_dev,
    lower = par_info$lower,
    upper = par_info$upper,
    config = config_dev,
    ad_pack = af,
    cache = cache,
    control_inner = control_inner,
    hessian = TRUE
)
```

## Joint log density

At starting values, evaluate the full log-density and its gradient over all
parameters (`beta`, inner `gamma`, and `theta`).

```{r joint-log-dens}
sz <- af@sizes
gamma0 <- rep(0, nrow(md$term_data$info$gamma))
x_full <- c(
    par_info$init[seq_len(sz["beta"])],
    gamma0,
    par_info$init[seq(sz["beta"] + 1L, length.out = sz["theta"])]
)

# Serial dens/grad/hess on a dens-safe clone (multi-thread affinity cleared).
dens_ptr <- clone_ad_pack_ptr(af@ptr)
shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(dens_ptr))
by_shard <- vapply(
    shards,
    function(s) joint_log_dens(dens_ptr, x_full, ad_shards = s),
    numeric(1)
)
by_shard[1:5]
by_shard[seq(to = length(by_shard), length.out = 5)]

sum(by_shard)
joint_log_dens(dens_ptr, x_full)
grad(dens_ptr, x_full)
bob <- hessian(dens_ptr, x_full)
```

## Likelihood

With outer parameters fixed at their starting values, `log_lik_laplace()`
optimizes random effects and returns the log-likelihood.

```{r likelihood}
x_outer <- par_info$init

lik <- log_lik_laplace(
    x = x_outer,
    ad_pack = af,
    config = modifyList(config_dev, list(verbose = TRUE)),
    gamma = gamma0,
    control = list(report.level = 2, report.freq = 1),
    deriv = TRUE
)

lik$neg_log_lik
lik$deriv$d_neg_log_lik
head(lik$inner_opt$solution)
```

`neg_log_lik` is the negative profile log-likelihood after inner optimization
over random effects.