## Tests for the compiled routines in ./src, exercised through .C().
##
## Memory layout used by the C code (all buffers are passed row-major, i.e.
## as.vector(t(.))):
##   lnk  Np x H         log kernel value at the draw from component j
##   lnd  Np x H^2       column (j-1)*H + k holds the log density of component
##                       k evaluated at the draw coming from component j

## Independent R implementation of the objective minimised by 'fn.optp':
## log(f1) - 2 log(f2), with f1 = E[p_j w^2] and f2 = E[p_j w].
ref_objective <- function(p, lnk, lnd_arr)
{
  Np <- dim(lnd_arr)[1]
  H  <- dim(lnd_arr)[2]
  f1 <- f2 <- 0
  for (i in seq_len(Np))
    for (j in seq_len(H))
      {
        d <- sum(p * exp(lnd_arr[i, j, ])) ## mixture density at that draw
        w <- exp(lnk[i, j]) / d
        f1 <- f1 + p[j] * w^2
        f2 <- f2 + p[j] * w
      }
  log(f1 / Np) - 2 * log(f2 / Np)
}

call_fnlnf <- function(p, lnk, lnd_arr)
{
  Np <- dim(lnd_arr)[1]
  H  <- dim(lnd_arr)[2]
  lnd_mat <- matrix(NA_real_, Np, H * H)
  for (j in seq_len(H))
    for (k in seq_len(H))
      lnd_mat[, (j - 1) * H + k] <- lnd_arr[, j, k]
  .C("fnlnf_C",
     lnp  = as.double(log(p)),
     lnk  = as.double(as.vector(t(lnk))),
     lnd  = as.double(as.vector(t(lnd_mat))),
     Np   = as.integer(Np),
     H    = as.integer(H),
     f    = as.double(0),
     grad = vector("double", H),
     PACKAGE = "AdMit",
     NAOK = TRUE)
}

test_that("fnlnf_C reproduces the objective function", {
  set.seed(101)
  Np <- 7L; H <- 3L
  p       <- c(0.5, 0.3, 0.2)
  lnk     <- matrix(rnorm(Np * H), Np, H)
  lnd_arr <- array(rnorm(Np * H * H, -1, 0.5), c(Np, H, H))

  expect_equal(call_fnlnf(p, lnk, lnd_arr)$f,
               ref_objective(p, lnk, lnd_arr))
})

test_that("fnlnf_C reproduces the gradient of the objective function", {
  set.seed(102)
  Np <- 5L; H <- 3L
  p       <- c(0.45, 0.35, 0.20)
  lnk     <- matrix(rnorm(Np * H), Np, H)
  lnd_arr <- array(rnorm(Np * H * H, -1, 0.5), c(Np, H, H))

  ## central differences of the R reference with respect to p
  eps  <- 1e-6
  num <- vapply(seq_len(H), function(j) {
    ph <- pl <- p
    ph[j] <- p[j] + eps
    pl[j] <- p[j] - eps
    (ref_objective(ph, lnk, lnd_arr) - ref_objective(pl, lnk, lnd_arr)) / (2 * eps)
  }, numeric(1))

  expect_equal(call_fnlnf(p, lnk, lnd_arr)$grad, num, tolerance = 1e-5)
})

test_that("fnlnf_C survives a large Np (regression: C stack overflow)", {
  ## 'AdMit' allows any Np in [100, Ns], and Ns defaults to 1e5. With the
  ## scratch buffers on the C stack, Np = 1e5 and H = 3 asked for ~11 MB and
  ## crashed the R session with "segfault from C stack overflow".
  set.seed(103)
  Np <- 100000L; H <- 3L
  p       <- c(0.5, 0.3, 0.2)
  lnk     <- matrix(rnorm(Np * H), Np, H)
  lnd_arr <- array(rnorm(Np * H * H, -1, 0.5), c(Np, H, H))

  r <- call_fnlnf(p, lnk, lnd_arr)
  expect_true(is.finite(r$f))
  expect_true(all(is.finite(r$grad)))
})

test_that("fnKernelMixtureArch_C is registered and computes the ARCH kernel", {
  ## Used by demo/AdMit.R. It was compiled but absent from the registration
  ## table, and R_useDynamicSymbols(dll, FALSE) then made it unreachable.
  set.seed(104)
  y     <- rnorm(40, sd = 0.5)
  theta <- rbind(c(0.10, 0.50, 0.10, 0.50),
                 c(0.05, 0.20, 0.30, 0.25),
                 c(0.10, 0.50, 0.10, 0.50))
  prior <- rbind(c(1, -1.5), c(1, -2.0), c(0, -Inf)) ## third draw violates the prior

  r <- .C("fnKernelMixtureArch_C",
          theta = as.double(as.vector(t(theta))),
          N     = as.integer(nrow(theta)),
          y     = as.double(y),
          n     = as.integer(length(y)),
          prior = as.double(as.vector(t(prior))),
          d     = vector("double", nrow(theta)),
          PACKAGE = "AdMit",
          NAOK = TRUE)$d

  ## R reference: mixture-of-two-ARCH(1) log-likelihood (Gaussian kernel,
  ## normalising constant omitted) plus the log prior
  ref <- vapply(seq_len(nrow(theta)), function(j) {
    if (prior[j, 1] != 1)
      return(prior[j, 2])
    s1 <- sqrt(theta[j, 1] + theta[j, 3] * head(y, -1)^2)
    s2 <- sqrt(theta[j, 2] + theta[j, 3] * head(y, -1)^2)
    t1 <- -0.5 * (tail(y, -1) / s1)^2 - log(s1)
    t2 <- -0.5 * (tail(y, -1) / s2)^2 - log(s2)
    sum(log(theta[j, 4] * exp(t1) + (1 - theta[j, 4]) * exp(t2))) + prior[j, 2]
  }, numeric(1))

  expect_equal(r, ref)
  expect_identical(r[3], -Inf) ## prior not satisfied
})

test_that("fnMH_C produces a valid independence chain", {
  set.seed(105)
  N <- 200L; k <- 2L
  theta <- matrix(rnorm(N * k), N, k)
  lnw   <- rnorm(N)
  u     <- runif(N)

  r <- .C("fnMH_C",
          theta  = as.double(as.vector(t(theta))),
          N      = as.integer(N),
          k      = as.integer(k),
          lnw    = as.double(lnw),
          u      = as.double(u),
          draws  = vector("double", N * k),
          ns     = as.integer(0),
          PACKAGE = "AdMit",
          NAOK = TRUE)
  draws <- matrix(r$draws, N, k, byrow = TRUE)

  ## R reference of the same independence chain
  ref <- matrix(NA_real_, N, k)
  ref[1, ] <- theta[1, ]
  s <- 1L
  ns <- 0L
  for (i in 2:N)
    {
      if (u[i] <= min(1, exp(lnw[i] - lnw[s])))
        {
          s  <- i
          ns <- ns + 1L
        }
      ref[i, ] <- theta[s, ]
    }

  expect_equal(draws, ref)
  expect_identical(r$ns, ns)
  ## every retained draw must be one of the candidates
  expect_true(all(draws[, 1] %in% theta[, 1]))
})