---
title: "AIS data in adlaplaceExample"
author: "Patrick Brown"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{AIS data in adlaplaceExample}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

This vignette fits a skew-normal location model to body-mass index (BMI) from
the Australian Institute of Sport (`ais`) dataset. We use the
`adlaplaceExample` backend for the observation density and `adlaplace` for
Laplace approximation, inner optimization, and outer profiling.

## Data

The `sn` package provides the original `ais` data; if it is not installed we
build a small stand-in so the vignette still knits.

```{r}
if (requireNamespace("sn", quietly = TRUE)) {
  data("ais", package = "sn")
  hist(ais$BMI)
} else {
  ais <- data.frame(
    sex = rep(c("male", "female"), 5),
    sport = rep(1:5, each = 2),
    BMI = rnorm(10)
  )
}
```

## Model and AD setup

We specify a skew-normal response for `BMI` with an intercept, fixed effect of
`sex`, and an independent random effect for `sport`. Observation shards are
split across 100 groups for parallel derivative evaluation.

```{r}
model_stuff <- adlaplace::model_data(
  adlaplaceExample::skewnormal(BMI) ~
    adlaplace::intercept(init = 10, parscale = 10) +
    sex + adlaplace::iid(sport),
  data = ais,
  verbose = TRUE
)

config <- list(
  transform_theta = TRUE,
  shards = adlaplace::ad_shards(model_stuff$data$A, num_shards = 100),
  num_threads = 2L,
  verbose = FALSE
)
```

## Laplace log-likelihood at starting values

`ad_fun()` composes the observation, parameter, and random-effect shards into
a single callable object. `log_lik_laplace()` solves the inner problem for
random effects and returns the profile log-likelihood, gradient, and fitted
parameters at the initial outer vector.

```{r}
ad_fun <- adlaplace::ad_fun(
  model_stuff,
  config = config,
  num_threads = config$num_threads
)
res <- adlaplace::log_lik_laplace(
  x = model_stuff$data$info$parameters$init,
  ad_fun = ad_fun,
  config = config,
  deriv = TRUE
)
res$parameters
res$neg_log_lik
res$grad
```

## Outer optimization

The outer objective profiles out random effects via a cache. We evaluate the
objective and gradient once at the start, then run a short L-BFGS-B fit.

```{r trustOptimOuterWrappers}
cache <- new.env(parent = emptyenv())
cache$gamma <- res$opt$solution

x0 <- model_stuff$data$info$parameters$init

adlaplace::outer_fn(x = x0, cache = cache, config = config, ad_fun = ad_fun)
adlaplace::outer_gr(x = x0, cache = cache, config = config, ad_fun = ad_fun)

outer_fit <- stats::optim(
  par = x0,
  fn = adlaplace::outer_fn,
  gr = adlaplace::outer_gr,
  lower = model_stuff$data$info$parameters$lower,
  upper = model_stuff$data$info$parameters$upper,
  method = "L-BFGS-B",
  control = list(
    maxit = 10,
    parscale = model_stuff$data$info$parameters$parscale,
    trace = 3,
    REPORT = 1
  ),
  config = config,
  ad_fun = ad_fun,
  cache = cache
)
outer_fit$par
outer_fit$value

adlaplace::format_parameters(
  parameters = outer_fit$par,gamma = cache$gamma,
  info=ad_fun@info)$parameters
```

## Fitted values and density overlay

After refitting at the outer MLE, we store updated fixed effects, variance
components, and random-effect modes, then overlay the fitted skew-normal
density on the BMI histogram (when `sn` is available).

```{r}
res <- adlaplace::log_lik_laplace(
  x = outer_fit$par,
  ad_fun = ad_fun,
  config = modifyList(config, list(verbose = FALSE)),
  deriv = TRUE
)
res$grad
model_stuff$data$info$gamma$mode <- res$opt$solution
model_stuff$data$info$beta$mle <- res$parameters[seq(1, nrow(model_stuff$data$info$beta))]

model_stuff$data$info$theta$mle <- res$parameters[-(1:nrow(model_stuff$data$info$beta))]
model_stuff$data$info$theta[
  model_stuff$data$info$theta$transform, "mle"
] <-
  exp(model_stuff$data$info$theta[model_stuff$data$info$theta$transform, "mle"])
```

```{r}
hist(ais$BMI, prob = TRUE, breaks = 31)
x_seq <- seq(par("usr")[1], par("usr")[2], len = 1001)
if (requireNamespace("sn", quietly = TRUE)) {
  lines(x_seq,
    sn::dsn(x_seq,
      xi = model_stuff$data$info$beta$mle[1],
      omega = model_stuff$data$info$theta$mle[1],
      alpha = model_stuff$data$info$theta$mle[2]
    ),
    col = "red"
  )
}
```

## Derivative check along one parameter

Finally we scan the fourth outer parameter and compare the analytic gradient
to finite differences, along with inner-mode and Hessian-determinant
sensitivities. This requires the suggested `abind` package.

```{r testLogLgrad}
if (!requireNamespace("abind", quietly = TRUE)) {
  stop("Install suggested package 'abind' to run derivative checks.")
}

x <- cache$last_par_fn # outer_fit$par
Dpar <- 4
Ngrid <- 13L

par_grid <- matrix(x, nrow = Ngrid, ncol = length(x), byrow = TRUE)
Sx <- 5 * seq(-1, 1, len = Ngrid) + x[Dpar]
SxD <- Sx[-1] - diff(Sx) / 2
par_grid[, Dpar] <- Sx

res <- lapply(
  split(par_grid, row(par_grid)),
  adlaplace::log_lik_laplace,
  ad_fun = ad_fun,
  config = config,
  deriv = TRUE,
  gamma = cache$gamma
)

SnegLik <- vapply(res, `[[`, numeric(1), "neg_log_lik")
Sdet <- vapply(res, function(r) r$extra$hessian$half_log_det, numeric(1))
grad_mat <- do.call(rbind, lapply(res, `[[`, "grad"))
extra_df <- do.call(abind::abind, c(lapply(res, `[[`, "deriv"), along = 3))
dU <- do.call(
  abind::abind,
  c(lapply(res, function(r) as.matrix(r$extra$dU)), along = 3)
)
u_hat <- do.call(rbind, lapply(res, function(r) r$opt$solution))

par(mfrow = c(3, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0))
plot(Sx, SnegLik)

plot(Sx, grad_mat[, Dpar], type = "l")
points(SxD, diff(SnegLik) / diff(Sx))

Du <- 1L
plot(Sx, u_hat[, Du])

plot(Sx, dU[Du, Dpar, ])
lines(SxD, diff(u_hat[, Du]) / diff(Sx))

plot(Sx, Sdet)

plot(Sx, extra_df[Dpar, "d_det", ])
lines(SxD, diff(Sdet) / diff(Sx))
```