ms_prices$rel <- price_relative(ms_prices, price ~ period + product)
ms_prices$w1 <- 1:40
ms_prices$w2 <- 40:1

epr1 <- elementary_index(ms_prices, rel ~ period + business, contrib = TRUE)

epr2 <- elementary_index(
  ms_prices,
  rel ~ period + business,
  order = -1,
  contrib = TRUE,
  na.rm = TRUE
)

# Test a Fisher calculation
fw <- function(x, w1, w2) {
  v1 <- transmute_weights(x, w1, 1, 0)
  v2 <- transmute_weights(x, w2, -1, 0)
  v1 + v2
}

ms_prices2 <- na.omit(ms_prices)
f <- interaction(ms_prices2$period, ms_prices2$business)
w <- unsplit(lapply(split(ms_prices2, f), \(df) fw(df$rel, df$w1, df$w2)), f)

epr3 <- elementary_index(
  ms_prices2,
  rel ~ period + business,
  weights = w,
  contrib = TRUE
)

# Agrees with alternate implementations.
local({
  epr11 <- aggregate(
    rel ~ as.character(business) + period,
    ms_prices,
    function(x) exp(weighted.mean(log(x))),
    na.action = na.pass,
    drop = FALSE
  )
  epr22 <- aggregate(
    rel ~ as.character(business) + period,
    ms_prices,
    function(x) 1 / weighted.mean(1 / x),
    na.action = na.omit,
    drop = FALSE
  )
  names(epr11) <- names(epr22) <- c("level", "period", "value")

  expect_equal(as.data.frame(epr1), epr11[c(2, 1, 3)])
  expect_equal(as.data.frame(epr2), epr22[c(2, 1, 3)])
})

# Fisher calculation agrees with manual calculation.
local({
  l <- elementary_index(
    ms_prices2,
    rel ~ period + business,
    weights = w1,
    order = 1
  )
  p <- elementary_index(
    ms_prices2,
    rel ~ period + business,
    weights = w2,
    order = -1
  )
  expect_equal(sqrt(as.matrix(l) * as.matrix(p)), as.matrix(epr3))

  # Should work for other kinds of superlative indexes
  fw <- function(x, w1, w2) {
    v1 <- transmute_weights(x, order = 1.5, to = 0)
    v2 <- transmute_weights(x, w2, order = -1.5, to = 0)
    v1 + v2
  }

  f <- interaction(ms_prices2$period, ms_prices2$business)
  w <- unsplit(lapply(split(ms_prices2, f), \(df) fw(df$rel, df$w1, df$w2)), f)

  sepr <- elementary_index(ms_prices2, rel ~ period + business, weights = w)

  l <- elementary_index(ms_prices2, rel ~ period + business, order = 1.5)
  p <- elementary_index(
    ms_prices2,
    rel ~ period + business,
    weights = w2,
    order = -1.5
  )
  expect_equal(sqrt(as.matrix(l) * as.matrix(p)), as.matrix(sepr))
})

# Contributions add up.
local({
  sum_contrib <- \(x, ...) sum(x, ...) / (sum(!is.na(x)) > 0)
  expect_equal(
    as.numeric(epr1$index),
    sapply(epr1$contrib, sum_contrib) + 1
  )

  expect_equal(
    as.numeric(epr2$index),
    sapply(epr2$contrib, sum_contrib, na.rm = TRUE) + 1
  )

  expect_equal(
    as.numeric(epr3$index),
    sapply(epr3$contrib, sum_contrib) + 1
  )
})

# Argument checking.
local({
  expect_error(elementary_index(1:3, period = 1:2, ea = 1:3))
  expect_error(elementary_index(1:3, period = 1:3, ea = 1:4))
  expect_error(elementary_index(1:3, period = 1:3, ea = 1:3, weights = 1:2))
  expect_error(elementary_index(1:3, period = factor(1:3, levels = numeric(0))))
  expect_error(elementary_index(1:3, ea = factor(1:3, levels = numeric(0))))
  expect_error(elementary_index(setNames(1:3, c("", 1, 2)), contrib = TRUE))
  expect_error(elementary_index(-1:1, period = 1:3, ea = 1:3, order = 1))
  expect_warning(elementary_index(
    setNames(1:3, rep(1, 3)),
    period = gl(1, 3),
    ea = gl(1, 3),
    contrib = TRUE
  ))
})

# Non-standard evaluation works.
local({
  expect_equal(
    elementary_index(ms_prices, rel ~ period:business, contrib = TRUE),
    epr1
  )
  expect_equal(
    elementary_index(ms_prices[c(1:2, 5)], rel ~ ., contrib = TRUE),
    epr1
  )
  expect_error(elementary_index(ms_prices, rel ~ .))
  expect_error(elementary_index(ms_prices, ~ period + business))
})
