# Tests for utility, sampling, and distribution functions

# ---------------------------------------------------------------------------
# rdirichlet
# ---------------------------------------------------------------------------

test_that("rdirichlet returns correct structure", {
  set.seed(1)
  result <- rdirichlet(n = 10L, alpha = c(1, 2, 3))
  expect_true(is.matrix(result))
  expect_equal(dim(result), c(10L, 3L))
  expect_true(all(result >= 0))
  expect_true(all(abs(rowSums(result) - 1) < 1e-10))
})

test_that("rdirichlet row sums are 1 for various alpha", {
  set.seed(2)
  result <- rdirichlet(n = 100L, alpha = c(0.5, 0.5, 0.5, 0.5))
  expect_true(all(abs(rowSums(result) - 1) < 1e-10))
})

test_that("rdirichlet input validation works", {
  expect_error(rdirichlet(n = 0L,  alpha = c(1, 1, 1)))
  expect_error(rdirichlet(n = 10L, alpha = c(-1, 1, 1)))
  expect_error(rdirichlet(n = 10L, alpha = c(0, 1, 1)))
})

# ---------------------------------------------------------------------------
# getjointbin
# Returns a named numeric vector of length 4: (p00, p01, p10, p11)
# ---------------------------------------------------------------------------

test_that("getjointbin returns a named vector of length 4", {
  result <- getjointbin(pi1 = 0.3, pi2 = 0.4, rho = 0.0)
  expect_type(result, "double")
  expect_length(result, 4L)
  expect_named(result, c("p00", "p01", "p10", "p11"))
})

test_that("getjointbin elements are non-negative and sum to 1", {
  result <- getjointbin(pi1 = 0.3, pi2 = 0.4, rho = 0.0)
  expect_true(all(result >= 0))
  expect_equal(sum(result), 1, tolerance = 1e-10)
})

test_that("getjointbin with rho = 0: p11 equals product of marginals", {
  pi1 <- 0.3; pi2 <- 0.4
  result <- getjointbin(pi1 = pi1, pi2 = pi2, rho = 0.0)
  expect_equal(unname(result["p11"]), pi1 * pi2, tolerance = 1e-10)
})

test_that("getjointbin marginals are consistent with pi1 and pi2", {
  pi1 <- 0.3; pi2 <- 0.4
  result <- getjointbin(pi1 = pi1, pi2 = pi2, rho = 0.0)
  # P(Y1 = 1) = p10 + p11
  expect_equal(unname(result["p10"] + result["p11"]), pi1, tolerance = 1e-10)
  # P(Y2 = 1) = p01 + p11
  expect_equal(unname(result["p01"] + result["p11"]), pi2, tolerance = 1e-10)
})

test_that("getjointbin input validation works", {
  expect_error(getjointbin(pi1 = -0.1, pi2 = 0.3, rho = 0))
  expect_error(getjointbin(pi1 = 0.3,  pi2 = 1.1, rho = 0))
  expect_error(getjointbin(pi1 = 0.5,  pi2 = 0.5, rho = 2))
})

# ---------------------------------------------------------------------------
# allmultinom
# Signature: allmultinom(n) — always returns a 4-column matrix
# ---------------------------------------------------------------------------

test_that("allmultinom returns correct structure", {
  result <- allmultinom(n = 5L)
  expect_true(is.matrix(result))
  expect_equal(ncol(result), 4L)
  expect_named(result[1L, ], c("x00", "x01", "x10", "x11"))
  expect_true(all(rowSums(result) == 5L))
  expect_true(all(result >= 0L))
})

test_that("allmultinom row count matches stars-and-bars formula", {
  # C(n+3, 3) = (n+1)(n+2)(n+3)/6
  n <- 5L
  result <- allmultinom(n)
  expect_equal(nrow(result), choose(n + 3L, 3L))
})

test_that("allmultinom n = 0 returns 1 row of zeros", {
  result <- allmultinom(0L)
  expect_equal(nrow(result), 1L)
  expect_true(all(result == 0L))
})

test_that("allmultinom input validation works", {
  expect_error(allmultinom(-1L))
})

# ---------------------------------------------------------------------------
# pbetadiff
# ---------------------------------------------------------------------------

test_that("pbetadiff returns a value in [0, 1]", {
  result <- pbetadiff(q = 0.1, alpha_t = 6, beta_t = 5, alpha_c = 3, beta_c = 8)
  expect_type(result, "double")
  expect_length(result, 1L)
  expect_true(result >= 0 && result <= 1)
})

test_that("pbetadiff lower.tail = TRUE and FALSE sum to 1", {
  q <- 0.1; a1 <- 6; b1 <- 5; a2 <- 3; b2 <- 8
  p_u <- pbetadiff(q, a1, b1, a2, b2, lower.tail = FALSE)
  p_l <- pbetadiff(q, a1, b1, a2, b2, lower.tail = TRUE)
  expect_equal(p_u + p_l, 1, tolerance = 1e-6)
})

test_that("pbetadiff at q = 0 with equal beta params gives ~0.5", {
  result <- pbetadiff(q = 0, alpha_t = 5, beta_t = 5, alpha_c = 5, beta_c = 5,
                      lower.tail = FALSE)
  expect_equal(result, 0.5, tolerance = 0.01)
})

test_that("pbetadiff input validation works", {
  expect_error(pbetadiff(q = 0.1, alpha_t = -1, beta_t = 5,  alpha_c = 3, beta_c = 8))
  expect_error(pbetadiff(q = 0.1, alpha_t = 6,  beta_t = 0,  alpha_c = 3, beta_c = 8))
  expect_error(pbetadiff(q = 0.1, alpha_t = 6,  beta_t = 5,  alpha_c = 3, beta_c = -1))
})

# ---------------------------------------------------------------------------
# pbetabinomdiff
# Signature: pbetabinomdiff(q, m_t, m_c, alpha_t, alpha_c, beta_t, beta_c, lower.tail)
# Posterior predictive: Y_k ~ BetaBinomial(m_k, alpha_k, beta_k)
# ---------------------------------------------------------------------------

test_that("pbetabinomdiff returns a value in [0, 1]", {
  result <- pbetabinomdiff(q = 0.1, m_t = 30, m_c = 30,
                           alpha_t = 10.5, alpha_c = 6.5,
                           beta_t = 10.5, beta_c = 14.5)
  expect_type(result, "double")
  expect_length(result, 1L)
  expect_true(result >= 0 && result <= 1)
})

test_that("pbetabinomdiff lower.tail = TRUE and FALSE sum to 1", {
  p_u <- pbetabinomdiff(q = 0.1, m_t = 30, m_c = 30,
                        alpha_t = 10.5, alpha_c = 6.5,
                        beta_t = 10.5, beta_c = 14.5, lower.tail = FALSE)
  p_l <- pbetabinomdiff(q = 0.1, m_t = 30, m_c = 30,
                        alpha_t = 10.5, alpha_c = 6.5,
                        beta_t = 10.5, beta_c = 14.5, lower.tail = TRUE)
  expect_equal(p_u + p_l, 1, tolerance = 1e-10)
})

test_that("pbetabinomdiff input validation works", {
  expect_error(pbetabinomdiff(q = 0.1, m_t = 0,  m_c = 30,
                              alpha_t = 1, alpha_c = 1, beta_t = 1, beta_c = 1))
  expect_error(pbetabinomdiff(q = 0.1, m_t = 30, m_c = 30,
                              alpha_t = -1, alpha_c = 1, beta_t = 1, beta_c = 1))
})

# ---------------------------------------------------------------------------
# ptdiff_NI
# ---------------------------------------------------------------------------

test_that("ptdiff_NI returns a value in [0, 1]", {
  result <- ptdiff_NI(q = 1, mu_t = 3, mu_c = 1, sd_t = 1, sd_c = 1,
                      nu_t = 15, nu_c = 15, lower.tail = FALSE)
  expect_type(result, "double")
  expect_length(result, 1L)
  expect_true(result >= 0 && result <= 1)
})

test_that("ptdiff_NI symmetric case returns ~0.5", {
  result <- ptdiff_NI(q = 0, mu_t = 1, mu_c = 1, sd_t = 1, sd_c = 1,
                      nu_t = 20, nu_c = 20, lower.tail = FALSE)
  expect_equal(result, 0.5, tolerance = 0.01)
})

test_that("ptdiff_NI lower.tail = TRUE and FALSE sum to 1", {
  p_u <- ptdiff_NI(q = 2, mu_t = 3, mu_c = 1, sd_t = 1.5, sd_c = 1.2,
                   nu_t = 15, nu_c = 18, lower.tail = FALSE)
  p_l <- ptdiff_NI(q = 2, mu_t = 3, mu_c = 1, sd_t = 1.5, sd_c = 1.2,
                   nu_t = 15, nu_c = 18, lower.tail = TRUE)
  expect_equal(p_u + p_l, 1, tolerance = 1e-6)
})

test_that("ptdiff_NI input validation works", {
  expect_error(ptdiff_NI(q = 1, mu_t = 3, mu_c = 1, sd_t = -1, sd_c = 1,
                         nu_t = 15, nu_c = 15))
  expect_error(ptdiff_NI(q = 1, mu_t = 3, mu_c = 1, sd_t = 1,  sd_c = 1,
                         nu_t = 1,  nu_c = 15))
})

# ---------------------------------------------------------------------------
# ptdiff_MC
# ---------------------------------------------------------------------------

test_that("ptdiff_MC returns a value in [0, 1]", {
  set.seed(1)
  result <- ptdiff_MC(nMC = 1000L, q = 1, mu_t = 3, mu_c = 1,
                      sd_t = 1, sd_c = 1, nu_t = 15, nu_c = 15,
                      lower.tail = FALSE)
  expect_type(result, "double")
  expect_length(result, 1L)
  expect_true(result >= 0 && result <= 1)
})

test_that("ptdiff_MC symmetric case returns ~0.5", {
  set.seed(2)
  result <- ptdiff_MC(nMC = 5000L, q = 0, mu_t = 1, mu_c = 1,
                      sd_t = 1, sd_c = 1, nu_t = 20, nu_c = 20,
                      lower.tail = FALSE)
  expect_equal(result, 0.5, tolerance = 0.05)
})

test_that("ptdiff_MC and ptdiff_NI agree closely", {
  set.seed(3)
  p_mc <- ptdiff_MC(nMC = 10000L, q = 2, mu_t = 3, mu_c = 1,
                    sd_t = 1.5, sd_c = 1.2, nu_t = 15, nu_c = 18,
                    lower.tail = FALSE)
  p_ni <- ptdiff_NI(q = 2, mu_t = 3, mu_c = 1,
                    sd_t = 1.5, sd_c = 1.2, nu_t = 15, nu_c = 18,
                    lower.tail = FALSE)
  expect_equal(p_mc, p_ni, tolerance = 0.03)
})

test_that("ptdiff_MC input validation works", {
  expect_error(ptdiff_MC(nMC = 0L,    q = 1, mu_t = 1, mu_c = 0,
                         sd_t = 1, sd_c = 1, nu_t = 10, nu_c = 10))
  expect_error(ptdiff_MC(nMC = 1000L, q = 1, mu_t = 1, mu_c = 0,
                         sd_t = 1, sd_c = 1, nu_t = 1,  nu_c = 10))
})

# ---------------------------------------------------------------------------
# ptdiff_MM
# ---------------------------------------------------------------------------

test_that("ptdiff_MM returns a value in [0, 1]", {
  result <- ptdiff_MM(q = 1, mu_t = 3, mu_c = 1, sd_t = 1, sd_c = 1,
                      nu_t = 15, nu_c = 15, lower.tail = FALSE)
  expect_type(result, "double")
  expect_length(result, 1L)
  expect_true(result >= 0 && result <= 1)
})

test_that("ptdiff_MM symmetric case returns ~0.5", {
  result <- ptdiff_MM(q = 0, mu_t = 1, mu_c = 1, sd_t = 1, sd_c = 1,
                      nu_t = 20, nu_c = 20, lower.tail = FALSE)
  expect_equal(result, 0.5, tolerance = 0.01)
})

test_that("ptdiff_MM and ptdiff_NI agree closely", {
  p_mm <- ptdiff_MM(q = 2, mu_t = 3, mu_c = 1,
                    sd_t = 1.5, sd_c = 1.2, nu_t = 15, nu_c = 18,
                    lower.tail = FALSE)
  p_ni <- ptdiff_NI(q = 2, mu_t = 3, mu_c = 1,
                    sd_t = 1.5, sd_c = 1.2, nu_t = 15, nu_c = 18,
                    lower.tail = FALSE)
  expect_equal(p_mm, p_ni, tolerance = 0.02)
})

test_that("ptdiff_MM vectorised input returns correct length", {
  result <- ptdiff_MM(q = 1, mu_t = c(2, 3, 4), mu_c = c(0, 1, 2),
                      sd_t = c(1, 1.2, 1.5), sd_c = c(1, 1.1, 1.3),
                      nu_t = 10, nu_c = 10, lower.tail = FALSE)
  expect_length(result, 3L)
  expect_true(all(result >= 0 & result <= 1))
})

test_that("ptdiff_MM input validation works", {
  expect_error(ptdiff_MM(q = 1, mu_t = 1, mu_c = 0, sd_t = 1,  sd_c = 1,
                         nu_t = 4, nu_c = 10))
  expect_error(ptdiff_MM(q = 1, mu_t = 1, mu_c = 0, sd_t = -1, sd_c = 1,
                         nu_t = 10, nu_c = 10))
})