# ---------------------------------------------------------------------------
# adaptive_importance_sampling(): normalising constants of known targets
# ---------------------------------------------------------------------------

# log g(z) = log_c + log N(z; mu, S), so that integral g(z) dz = exp(log_c).
gaussian_log_g <- function(mu, S, log_c) {
  R <- chol(S)
  d <- length(mu)
  function(Z) {
    dev <- sweep(Z, 2, mu, "-")
    mahal <- colSums(backsolve(R, t(dev), transpose = TRUE)^2)
    list(log_g = log_c - 0.5 * mahal - sum(log(diag(R))) - (d / 2) * log(2 * pi))
  }
}

test_that("AIS recovers the normalising constant of a Gaussian (d = 1..4)", {
  # Error with 1024 draws: at most 0.008 over 20 seeds for d = 4 (where one
  # adaptation step can be needed)
  set.seed(11)
  log_c <- -123.456
  for (d in 1:4) {
    A <- matrix(rnorm(d * d), d, d)
    S <- crossprod(A) / d + diag(0.2, d)
    mu <- seq(0.8, -0.6, length.out = d)
    ais <- adaptive_importance_sampling(gaussian_log_g(mu, S, log_c),
                                               z_start = rep(0, d), n_draw = 1024,
                                               min_ess = 256)
    expect_lte(ais$n_iter, 2L)
    expect_lt(abs(ais$log_evidence - log_c), 0.02)
    expect_lt(max(abs(colSums(ais$w * ais$Z) - mu)), 0.05)
  }
})

test_that("AIS adapts and pools draws until min_ess is reached", {
  # Skewed target: the first step has ESS ~ 130 < min_ess, so a second
  # iteration is needed
  log_c <- -50
  skew_log_g <- function(Z) {
    list(log_g = log_c - (Z[, 1] + exp(-Z[, 1])) + stats::dlogis(Z[, 2], log = TRUE))
  }
  set.seed(1)
  ais <- adaptive_importance_sampling(skew_log_g, z_start = c(0, 0),
                                             n_draw = 256, min_ess = 0.9 * 256)
  expect_gt(ais$n_iter, 1L)
  expect_gte(ais$ess, 0.9 * 256)
  expect_equal(nrow(ais$Z), ais$n_iter * 256)
  expect_lt(abs(ais$log_evidence - log_c), 0.02)
})

# ---------------------------------------------------------------------------
# fit_one_model() with AIS: interface
# ---------------------------------------------------------------------------

test_that("fit_one_model() with AIS returns weighted particles, ESS and proposal", {
  set.seed(3)
  y <- ts(cumsum(rnorm(30)) + 10, frequency = 4)
  for (mc in list(c("A", "N", "N", "FALSE"), c("A", "A", "A", "TRUE"))) {
    ctrl <- resolve_bets_control(list())
    ctrl <- utils::modifyList(ctrl, prior_scales(y, ctrl))
    res <- fit_one_model(y, mc, ctrl, "ais")
    d <- if (mc[[2]] == "N") 1 else 4
    expect_named(res, c("particles", "ess", "n_iter", "prop_params", "log_evidence", "timing"))
    expect_true(is.finite(res$log_evidence))
    expect_gte(res$ess, ctrl$N_draw[d] / 4)
    expect_gte(res$n_iter, 1L)
    expect_named(res$prop_params, c("mus", "Sigma", "df"))
    expect_equal(ncol(res$particles$theta), d)
    expect_gte(nrow(res$particles$theta), ctrl$N_draw[d] / 4)
  }
})

test_that("fit_one_model() with AIS gives zero weight when min_ess is not reached", {
  set.seed(1)
  y <- ts(rnorm(20))
  ctrl <- resolve_bets_control(list(min_ess = 1e6, N_iter_max = 2))
  ctrl <- utils::modifyList(ctrl, prior_scales(y, ctrl))
  expect_warning(res <- fit_one_model(y, c("A", "N", "N", "FALSE"), ctrl, "ais"),
                 "zero weight")
  expect_equal(res$log_evidence, -Inf)
  expect_null(res$particles)
  expect_equal(res$n_iter, 2L)
})