test_that("balEq computes balance for a valid blockMatch-like object", {
  m <- data.frame(
    mset = c(1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 0),
    age = c(40, 41, 50, 52)
  )
  all <- data.frame(
    z = c(1, 0, 1, 0, 1, 0),
    age = c(40, 41, 50, 52, 45, 48),
    matched = rep(c(TRUE, FALSE), each = 3L)
  )
  o <- list(m = m, all = all)
  res <- balEq("age", o)
  expect_true(is.matrix(res))
  expect_equal(nrow(res), 1L)
  expect_equal(ncol(res), 11L)
  expect_equal(colnames(res), c(
    "T-before", "C-before", "T-after", "C-after",
    "dif.before", "dif.after", "sdif.before", "sdif.after",
    "med", "q9", "pct0"
  ))
})

test_that("balEq standardized differences match hand-computed values", {
  # Two blocks of size 2; same four people in m and all -> before/after means align
  m <- data.frame(
    mset = c(1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 0),
    age = c(40, 50, 45, 55)
  )
  all <- data.frame(
    z = c(1, 0, 1, 0),
    age = c(40, 50, 45, 55),
    matched = TRUE
  )
  o <- list(m = m, all = all)
  res <- balEq("age", o)
  av <- all$age
  bt <- mean(av[all$z == 1])
  bc <- mean(av[all$z == 0])
  sbefore <- sqrt((var(av[all$z == 1]) + var(av[all$z == 0])) / 2)
  d_block <- c(40 - 50, 45 - 55)
  expect_equal(unname(res[1, "T-before"]), bt)
  expect_equal(unname(res[1, "C-before"]), bc)
  expect_equal(unname(res[1, "T-after"]), bt)
  expect_equal(unname(res[1, "C-after"]), bc)
  expect_equal(unname(res[1, "dif.before"]), bt - bc)
  expect_equal(unname(res[1, "dif.after"]), mean(d_block))
  expect_equal(unname(res[1, "sdif.before"]), (bt - bc) / sbefore)
  expect_equal(unname(res[1, "sdif.after"]), mean(d_block) / sbefore)
  expect_equal(unname(res[1, "med"]), 10)
  expect_equal(unname(res[1, "q9"]), 10)
  expect_equal(unname(res[1, "pct0"]), 0)

  det <- balEq("age", o, detail = TRUE)
  expect_equal(as.vector(det$d), d_block)
  expect_equal(det$z, rbind(c(1, 0), c(1, 0)))
  expect_equal(det$y, rbind(c(40, 50), c(45, 55)))
})

test_that("balEq detail = TRUE returns list with y, z, d", {
  m <- data.frame(
    mset = c(1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 0),
    score = c(1, 2, 3, 4)
  )
  all <- data.frame(
    z = c(1, 0, 1, 0),
    score = c(1, 2, 3, 4),
    matched = TRUE
  )
  o <- list(m = m, all = all)
  res <- balEq("score", o, detail = TRUE)
  expect_named(res, c("balance", "y", "z", "d"))
  expect_true(is.matrix(res$balance))
  expect_equal(dim(res$y), c(2L, 2L))
})

test_that("balEq rejects bad inputs", {
  expect_error(
    balEq("x", list(m = data.frame(), all = data.frame(z = 1, matched = TRUE))),
    "empty"
  )
  expect_error(balEq("missing", list(
    m = data.frame(mset = 1, z = 1, y = 1),
    all = data.frame(z = 1, y = 1)
  )), "must be a column of o\\$m")
})

test_that("balEq allows extra list elements on o", {
  m <- data.frame(
    mset = c(1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 0),
    x = c(1, 2, 3, 4)
  )
  all <- data.frame(
    z = c(1, 0, 1, 0),
    x = c(1, 2, 3, 4)
  )
  o <- list(m = m, all = all, extra = 99L, meta = list(a = 1L))
  expect_equal(dim(balEq("x", o)), c(1L, 11L))
})

test_that("balEq agrees with sorted rows when o$m is interleaved by mset", {
  m_sorted <- data.frame(
    mset = c(1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 0),
    x = c(10, 20, 30, 40)
  )
  # Same two blocks, rows interleaved; after sort by mset each block still has one T and one C
  m_interleaved <- data.frame(
    mset = c(1L, 2L, 1L, 2L),
    z = c(1, 0, 0, 1),
    x = c(10, 40, 20, 30)
  )
  all <- data.frame(z = c(1, 0, 1, 0), x = 1:4, matched = TRUE)
  o1 <- list(m = m_sorted, all = all)
  o2 <- list(m = m_interleaved, all = all)
  expect_equal(balEq("x", o1), balEq("x", o2))
  d1 <- balEq("x", o1, detail = TRUE)
  d2 <- balEq("x", o2, detail = TRUE)
  # Within-block row order in y, z can differ; block-level quantities match
  expect_equal(d1$d, d2$d)
})

test_that("balEq rejects unequal block sizes", {
  m <- data.frame(
    mset = c(1L, 1L, 1L, 2L, 2L),
    z = c(1, 0, 1, 1, 0),
    x = 1:5
  )
  all <- data.frame(z = c(1, 0), x = 1:2, matched = TRUE)
  o <- list(m = m, all = all)
  expect_error(balEq("x", o), "same number of rows")
})

test_that("balEq works on output from blockMatch", {
  set.seed(42)
  n <- 40L
  dat <- data.frame(x = rnorm(n), w = rnorm(n), z = rep(0:1, each = n / 2L))
  pr <- runif(n)
  prG <- 2L
  dat$prc <- as.integer(cut(pr, stats::quantile(pr, (0:prG) / prG), include.lowest = TRUE))
  bd <- basicDistance(dat)
  out <- suppressWarnings(
    blockMatch(dat, cost = bd$cost, J = 2L, ratio = 1L, rseed = 42L)
  )
  b <- balEq("x", out)
  expect_equal(dim(b), c(1L, 11L))
  # Cross-check vs manual formulas on the same matched sample (same row order as balEq)
  m <- out$m[order(out$m$mset, seq_len(nrow(out$m))), , drop = FALSE]
  a <- out$all
  v <- m$x
  av <- a$x
  I <- length(unique(m$mset))
  J <- length(m$mset) / I
  z <- t(matrix(m$z, J, I))
  y <- t(matrix(v, J, I))
  zsum <- apply(z, 1, sum)
  d_manual <- (apply(z * y, 1, sum) / zsum) -
    (apply((1 - z) * y, 1, sum) / (J - zsum))
  sbefore <- sqrt((var(av[a$z == 1]) + var(av[a$z == 0])) / 2)
  expect_equal(unname(b[1, "dif.after"]), mean(d_manual), tolerance = 1e-10)
  expect_equal(unname(b[1, "sdif.after"]), mean(d_manual) / sbefore, tolerance = 1e-10)

  bd <- balEq("x", out, detail = TRUE)
  expect_named(bd, c("balance", "y", "z", "d"))
  expect_equal(as.vector(bd$d), d_manual, tolerance = 1e-10)
})