set.seed(4321)

# Corner cases for allocations work as expected.
local({
  expect_identical(
    prop_allocation(numeric(0), 0, factor(integer(0), levels = 1)),
    c("1" = 0L)
  )
  expect_identical(
    prop_allocation(0, 0, 1),
    c("1" = 0L)
  )
  expect_identical(
    prop_allocation(1:3, 0, 1:3),
    c("1" = 0L, "2" = 0L, "3" = 0L)
  )
  expect_identical(
    prop_allocation(c(1, 1, 0, 0), 2, gl(2, 2), divisor = identity),
    c("1" = 2L, "2" = 0L)
  )
  expect_identical(
    prop_allocation(1:5, 5, gl(1, 5)),
    c("1" = 5L)
  )
  expect_identical(
    prop_allocation(
      rep(1, 10),
      4,
      factor(rep(letters[1], 10), levels = c("a", "b"))
    ),
    c(a = 4L, b = 0L)
  )
  expect_identical(
    prop_allocation(1:4, 2, gl(1, 4), initial = 3),
    c("1" = 2L)
  )
})

# Argument checking works for allocations.
local({
  expect_error(prop_allocation(c(-1, 1, 0), 1, c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, NA), 1, c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), -1, c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), integer(0), c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), NA, c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), 3, c(1, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(NA, 1, 2)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1)))
  expect_error(prop_allocation(integer(0), 0, factor(integer(0))))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), initial = integer(0)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), initial = c(1, 1)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), initial = c(1, -1)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), initial = c(2, 0)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), initial = c(2, 0, 1)))
  expect_error(prop_allocation(c(1, 1, 0), 1, c(1, 1, 2), divisor = "a"))
  expect_error(prop_allocation(c(1, 0, 1), 2, c(1, 1, 2), initial = c(2, 0)))
})

# Simple allocations are correct.
local({
  expect_identical(
    prop_allocation(rep(1, 10), 4, rep(letters[1:2], 5)),
    c(a = 2L, b = 2L)
  )
  expect_identical(
    prop_allocation(rep(c(1, 10), 5), 4, rep(letters[1:2], 5)),
    c(a = 0L, b = 4L)
  )
  expect_identical(
    prop_allocation(rep(c(1, 10), 5), 4, rep(letters[1:2], 5), initial = 1),
    c(a = 1L, b = 3L)
  )
  expect_identical(
    prop_allocation(rep(c(1, 10), 5), 4, rep(letters[1:2], 5), initial = 2),
    c(a = 2L, b = 2L)
  )
  expect_identical(
    prop_allocation(
      c(rep(10, 8), 1, 1),
      5,
      c(rep("a", 8), "b", "b"),
      initial = 3
    ),
    c(a = 3L, b = 2L)
  )
  expect_identical(
    prop_allocation(
      rep(c(1, 10), 5),
      4,
      rep(letters[1:2], 5),
      initial = c(3, 1)
    ),
    c(a = 3L, b = 1L)
  )
  expect_identical(
    prop_allocation(
      rep(c(1, 10), 5),
      4,
      factor(rep(letters[1:2], 5), levels = letters[1:3]),
      initial = c(2, 1, 0)
    ),
    c(a = 2L, b = 2L, c = 0L)
  )
  expect_identical(
    prop_allocation(c(0, rep(1, 100), 0), 10, rep(1:4, c(11, 20, 30, 41))),
    c("1" = 1L, "2" = 2L, "3" = 3L, "4" = 4L)
  )
  expect_identical(
    prop_allocation(c(0, 0, 100, 1, 1, 1), 3, gl(2, 3)),
    c("1" = 1L, "2" = 2L)
  )
  expect_identical(
    prop_allocation(
      c(1, 1, 1, 1, 1, 1, 1, 10, 10, 10),
      5,
      rep(letters[1:2], c(7, 3))
    ),
    c(a = 2L, b = 3L)
  )
  expect_identical(
    prop_allocation(
      c(1, 1, 1, 1, 1, 1, 9, 9, 100, 100),
      5,
      rep(letters[1:3], c(6, 2, 2))
    ),
    c(a = 1L, b = 2L, c = 2L)
  )
  expect_identical(
    prop_allocation(
      c(1, 1, 1, 1, 1, 1, 9, 9, 100, 100),
      5,
      rep(letters[1:3], c(6, 2, 2)),
      divisor = identity
    ),
    c(a = 1L, b = 2L, c = 2L)
  )
})

# Ties are broken correctly.
local({
  expect_identical(
    prop_allocation(rep(1, 9), 8, rep(1:2, c(4, 5))),
    c("1" = 3L, "2" = 5L)
  )
  expect_identical(
    prop_allocation(rep(1, 9), 8, rep(1:2, c(4, 5)), ties = "first"),
    c("1" = 4L, "2" = 4L)
  )
  expect_identical(
    prop_allocation(rep(1, 5), 2, 1:5),
    c("1" = 1L, "2" = 1L, "3" = 0L, "4" = 0L, "5" = 0L)
  )
  expect_identical(
    prop_allocation(rep(1, 5), 2, 1:5, ties = "first"),
    c("1" = 1L, "2" = 1L, "3" = 0L, "4" = 0L, "5" = 0L)
  )
})

# Alabama paradox doesn't happen.
local({
  expect_identical(
    prop_allocation(rep(1, 14), 10, rep(1:3, c(6, 6, 2))),
    c("1" = 5L, "2" = 4L, "3" = 1L)
  )
  expect_identical(
    prop_allocation(rep(1, 14), 11, rep(1:3, c(6, 6, 2))),
    c("1" = 5L, "2" = 5L, "3" = 1L)
  )
})

# Allocations are correct for voting examples.
local({
  # example from https://en.wikipedia.org/wiki/Highest_averages_method
  x <- rep(1, 1e5)
  s <- factor(
    rep(1:7, c(47000, 16000, 15900, 12000, 6000, 0, 3100)),
    levels = 1:7
  )

  expect_identical(
    prop_allocation(x, 10, s),
    c("1" = 5L, "2" = 2L, "3" = 2L, "4" = 1L, "5" = 0L, "6" = 0L, "7" = 0L)
  )
  expect_identical(
    prop_allocation(x, 10, s, divisor = \(a) a + 0.5),
    c("1" = 4L, "2" = 2L, "3" = 2L, "4" = 1L, "5" = 1L, "6" = 0L, "7" = 0L)
  )
  expect_identical(
    prop_allocation(x, 10, s, divisor = \(a) sqrt(a * (a + 1))),
    c("1" = 4L, "2" = 2L, "3" = 1L, "4" = 1L, "5" = 1L, "6" = 0L, "7" = 1L)
  )
  expect_identical(
    prop_allocation(x, 10, s, divisor = \(a) a),
    c("1" = 3L, "2" = 2L, "3" = 2L, "4" = 1L, "5" = 1L, "6" = 0L, "7" = 1L)
  )
})

# Expected coverage works.
local({
  expect_equal(expected_coverage(0, 0, 1), 0)
  expect_equal(expected_coverage(1:6, 6, gl(1, 6)), 1)
  expect_equal(expected_coverage(1:6, 0, gl(1, 6)), 0)
  expect_equal(expected_coverage(1:6, 3, 1:6), 3)
  expect_equal(
    expected_coverage(1:10, 4, gl(2, 5)),
    expected_coverage(1:10, 4, gl(2, 5, labels = 1:3))
  )
  # bernoulli sampling
  expect_equal(
    expected_coverage(rep(1, 10), 4, c(rep(1, 4), rep(2, 6))),
    2 - (1 - 0.4)^4 - (1 - 0.4)^6
  )
  # simulation
  x <- c(
    0,
    20,
    16,
    32,
    14,
    35,
    9,
    6,
    2,
    33,
    29,
    40,
    27,
    38,
    47,
    26,
    46,
    12,
    11,
    39,
    24,
    100,
    0,
    1,
    6,
    6,
    9,
    20,
    15,
    25,
    14,
    0,
    100
  )
  s <- c(
    4,
    4,
    2,
    5,
    5,
    5,
    5,
    2,
    3,
    2,
    5,
    2,
    2,
    5,
    5,
    2,
    3,
    4,
    5,
    5,
    3,
    4,
    2,
    2,
    2,
    3,
    5,
    2,
    1,
    1,
    2,
    3,
    3
  )

  expect_equal(expected_coverage(x, 10, s), 4.41282552802)

  # simulate mean(replicate(1e3, length(unique(s[ps(x, 10)]))))
})

# Argument checking for expected coverage works.
local({
  expect_error(expected_coverage(numeric(0), 0, integer(0)))
  expect_error(expected_coverage(-1:4, 3, gl(2, 3)))
  expect_error(expected_coverage(c(1:5, NA), 3, gl(2, 3)))
  expect_error(expected_coverage(numeric(0), 3, gl(2, 3)))
  expect_error(expected_coverage(1:6, -3, gl(2, 3)))
  expect_error(expected_coverage(1:6, integer(0), gl(2, 3)))
  expect_error(expected_coverage(1:6, NA, gl(2, 3)))
  expect_error(expected_coverage(1:6, 3, gl(2, 2)))
  expect_error(expected_coverage(1:6, 3, gl(2, 3)[c(1:5, 7)]))
  expect_error(expected_coverage(1:6, 3, gl(2, 3), alpha = -0.5))
  expect_error(expected_coverage(1:6, 3, gl(2, 3), alpha = c(0, 1)))
  expect_error(expected_coverage(1:6, 3, gl(2, 3), alpha = numeric(0)))
  expect_error(expected_coverage(1:6, 3, gl(2, 3), alpha = NA))
})

# Divisors are correct.
local({
  expect_equal(divisor_method("Webster/Sainte-Lague")(1:3), 1:3 + 1 / 2)
  expect_equal(divisor_method("Imperiali")(1:3), 1:3 + 2)
  expect_equal(divisor_method("Huntington-Hill")(1:3), sqrt(1:3 * 2:4))
  expect_equal(divisor_method("Danish")(1:3), 1:3 + 1 / 3)
  expect_equal(divisor_method("Adams")(1:3), 1:3)
  expect_equal(divisor_method("Dean")(1:3), 1:3 * 2:4 / (1:3 + 1 / 2))
  expect_error(divisor_method("Adams")(0:3))
})
