# Pure-R tests for calculate_loo(), calculate_waic(), and mixed-criterion
# compare_models(). No Stan compilation — we construct mock mlumr_fit
# objects whose $draws carries the vector-indexed `log_lik_ipd[i]` and
# `log_lik_agd[i]` columns that all six shipping Stan models emit.

make_ll_fit <- function(model = "spfa", n_draws = 400,
                        n_ipd = 12L, n_agd = 4L, seed = 1L) {
  set.seed(seed)
  ll_ipd <- matrix(rnorm(n_draws * n_ipd, -5, 0.5), nrow = n_draws)
  ll_agd <- matrix(rnorm(n_draws * n_agd, -6, 0.5), nrow = n_draws)
  colnames(ll_ipd) <- sprintf("log_lik_ipd[%d]", seq_len(n_ipd))
  colnames(ll_agd) <- sprintf("log_lik_agd[%d]", seq_len(n_agd))

  draws <- data.frame(
    ll_ipd, ll_agd,
    lor_index = rnorm(n_draws, 0.5, 0.2),
    lor_comparator = rnorm(n_draws, 0.4, 0.2),
    check.names = FALSE
  )

  out <- list(
    draws = draws,
    summary = data.frame(
      variable = c("mu_index", "lor_index"),
      mean = c(-0.5, 0.5), sd = c(0.1, 0.2),
      Rhat = c(1.001, 1.002), n_eff = c(800, 700),
      stringsAsFactors = FALSE
    ),
    model = model,
    diagnostics = list(n_divergent = 0, n_max_treedepth = 0),
    sampling_args = list(adapt_delta = 0.95, max_treedepth = 15, chains = 4)
  )
  class(out) <- c("mlumr_fit", "list")
  out
}


# extract_log_lik is internal but we can probe it via mlumr:::
test_that("extract_log_lik orders columns and returns a numeric matrix", {
  fit <- make_ll_fit(n_ipd = 5L, n_agd = 2L)
  ll <- mlumr:::extract_log_lik(fit)
  expect_true(is.matrix(ll))
  expect_equal(dim(ll), c(400L, 7L))
  # The last IPD column (index 5) should precede the first AgD column
  expect_equal(colnames(ll)[5], "log_lik_ipd[5]")
  expect_equal(colnames(ll)[6], "log_lik_agd[1]")
})


test_that("extract_log_lik rejects fits with no pointwise columns", {
  fit <- make_ll_fit()
  fit$draws <- fit$draws[, c("lor_index", "lor_comparator")]
  expect_error(mlumr:::extract_log_lik(fit),
               "Pointwise log-likelihood columns not found")
})


test_that("calculate_loo returns a psis_loo object", {
  skip_if_not_installed("loo")
  fit <- make_ll_fit()
  result <- suppressWarnings(calculate_loo(fit))
  expect_s3_class(result, "psis_loo")
  expect_s3_class(result, "loo")
  # The `loo` object must expose `elpd_loo`, p_loo, looic
  expect_true(all(c("elpd_loo", "p_loo", "looic") %in%
                    rownames(result$estimates)))
})


test_that("calculate_loo uses stable relative efficiency for tiny likelihoods", {
  skip_if_not_installed("loo")
  fit <- make_ll_fit()
  ll_cols <- grep("^log_lik_", names(fit$draws), value = TRUE)
  fit$draws[, ll_cols] <- fit$draws[, ll_cols] - 10000

  result <- suppressWarnings(calculate_loo(fit))

  expect_s3_class(result, "psis_loo")
})


test_that("calculate_waic returns a waic object", {
  skip_if_not_installed("loo")
  fit <- make_ll_fit()
  result <- suppressWarnings(calculate_waic(fit))
  expect_s3_class(result, "waic")
  expect_s3_class(result, "loo")
  expect_true(all(c("elpd_waic", "p_waic", "waic") %in%
                    rownames(result$estimates)))
})


test_that("calculate_loo errors informatively without pointwise log_lik", {
  fit <- make_ll_fit()
  fit$draws <- fit$draws[, c("lor_index", "lor_comparator")]
  expect_error(calculate_loo(fit),
               "Pointwise log-likelihood columns not found")
})


test_that("compare_models with criterion = 'loo' returns a compare.loo", {
  skip_if_not_installed("loo")
  fit1 <- make_ll_fit("spfa", seed = 1L)
  fit2 <- make_ll_fit("relaxed", seed = 2L)
  out <- suppressWarnings(
    capture.output(cmp <- compare_models(fit1, fit2, criterion = "loo"))
  )
  expect_s3_class(cmp, "compare.loo")
  # Two rows, one per model
  expect_equal(nrow(cmp), 2L)
  expect_true(any(grepl("Model Comparison \\(LOO\\)", out)))
})


test_that("compare_models with criterion = 'waic' returns a compare.loo", {
  skip_if_not_installed("loo")
  fit1 <- make_ll_fit("spfa", seed = 1L)
  fit2 <- make_ll_fit("relaxed", seed = 2L)
  out <- suppressWarnings(
    capture.output(cmp <- compare_models(fit1, fit2, criterion = "waic"))
  )
  expect_s3_class(cmp, "compare.loo")
  expect_equal(nrow(cmp), 2L)
})


test_that("compare_models rejects mixed mlumr_fit + mlumr_dic for LOO/WAIC", {
  skip_if_not_installed("loo")
  fit <- make_ll_fit("spfa")
  dic <- calculate_dic(fit)
  expect_error(
    compare_models(fit, dic, criterion = "loo"),
    "all arguments must be mlumr_fit"
  )
})


test_that(".chain_id recovers chain ids in chain-major order", {
  fit <- make_ll_fit(n_draws = 400)  # chains = 4 -> 100 per chain
  ids <- mlumr:::.chain_id(fit)
  expect_equal(length(ids), 400L)
  expect_equal(sort(unique(ids)), 1:4)
  # First and last 100 draws belong to different chains
  expect_equal(ids[1], 1L)
  expect_equal(ids[400], 4L)
})
