archetypes: Formula Spells

test_that("new distill_runes can be made from character/atomic components", {
  ty <- distill_rune(
    x = "Y",
    side = "left",
    role = "outcome",
    label = "Dependent Variable",
    description = "Artificially created",
    distribution = "normal",
    type = "numeric",
    subtype = "continuous"
  )

  tx <- distill_rune(
    "X",
    side = "right",
    role = "exposure",
    label = "Independent Variable",
    description = "Artificially created",
    distribution = "normal",
    type = "numeric",
    subtype = "dichotomous"
  )

  tm <- distill_rune(
    "M",
    side = "right",
    role = "mediator",
    label = "Independent Variable",
    description = "Artificially created",
    distribution = "normal",
    type = "integer",
    subtype = "continuous"
  )

  tc <- distill_rune(
    "C",
    side = "right",
    role = "covariate",
    label = "Independent Variable",
    description = "Artificially created",
    distribution = "normal",
    type = "character",
    subtype = "categorical"
  )

  ts <- distill_rune(
    "S",
    side = "meta",
    role = "strata",
    label = "Stratification Variable",
    description = "Levels for data set",
    distribution = "binary",
    type = "character",
    subtype = "dichotomous"
  )

  ti <- distill_rune(
    "I",
    side = "right",
    role = "interaction",
    label = "Interaction Variable",
    description = "Interaction for the exposure variable",
    distribution = "binary",
    type = "character",
    subtype = "dichotomous"
  )

  t <- c(ty, tx, tm, tc, ts, ti)

  expect_length(t, 6)
  expect_true(is_rune(t))

  # Inappropriate variables should lead to stop
  x <- distill_rune("M", side = "right")
  y <- distill_rune("Y", side = "left")
  role = list(M ~ "mediator", Y ~ "covariate")
  expect_error(set_roles(c(x, y), roles = formula_to_named_list(role)))

})

test_that("distill_rune() makes distill_rune object or errors", {

  # Messages for zero length objects
  expect_message(distill_rune(formula()))
  expect_message(distill_rune(character()))
  #expect_message(distill_rune(data.frame()))

  t1 <- distill_rune("y", side = "left", role = "outcome", label = "Dependent Variable")
  t2 <- distill_rune("x", side = "right", role = "exposure", label = "Independent Variable")
  expect_s3_class(t1, "rune")
  expect_true(is_rune(t1))
  expect_error(new_rune("x"))
  expect_length(t1, 1)
  expect_length(t2, 1)
  expect_length(suppressMessages(distill_rune(formula())), 0)

  # Field size should be the same
  expect_error(distill_rune(c("x", "y")))

  # Expected class of input matters
  expect_error(distill_rune(as.name("x")))
})

test_that("formatting is correct", {
  t1 <- distill_rune("y", side = "left", role = "outcome", label = "Dependent Variable")
  t2 <- distill_rune("x", side = "right", role = "exposure", label = "Independent Variable")
  vt <- c(t1, t2)
  expect_output(print(t1), "y")
  expect_output(print(new_rune()), "[0]")

  if (isTRUE(requireNamespace("tibble", quietly = TRUE))) {
    tibble::tibble(vt) |>
      print() |>
      expect_output("<rx>")
  }
})

test_that("casting and coercion for different dispatches work", {

  # Basic cast into character
  x1 <- distill_rune("x1", side = "right", role = "exposure", label = "Independent Variable")
  x2 <- distill_rune("x2", side = "right", role = "confounder", label = "Independent Variable")
  y <- "y"
  expect_type(c(x1, y), "character")
  expect_s3_class(c(x1, x2), "rune")
  expect_type(vec_c(x1, y), "character")

  # Formula archetypes
  s <- cast_spell(mpg ~ X(wt) + M(cyl) + hp)
  f <- summon_formulas(s)
  t <- distill_rune(f)
  expect_length(t, 4)
  expect_equal(decipher(t), 3)
  expect_equal(decipher(t), field(s, "order"))
  expect_equal(length(t), length(field(s, "runes")[[1]]))

})


test_that("runes can be generated from formulas", {

  # Simple formula for distill_runes to be broken down
  ts <- distill_rune(
    x = mpg + wt ~ hp + cyl + gear,
    tier = list(cyl ~ "engine", gear ~ "engine"),
    label = list(mpg ~ "Mileage")
  )
  expect_length(ts, 5)

  # Complex formula with distill_rune and data operations
  f <- mpg + wt ~ X(hp) + M(cyl) + gear + drat + log(qsec)
  t <- distill_rune(
    x = f,
    tier = list(drat + qsec ~ "spec"),
    label = list(mpg ~ "Mileage", wt ~ "Weight")
  )
  expect_length(t, 7)

  # Reversing a rune object into a formula
  expect_s3_class(stats::formula(t), "formula")

  t1 <- distill_rune(f)
  t2 <- distill_rune(f, label = list(mpg ~ "Mileage"), tier = list(qsec + drat ~ "speed"))
  expect_equal(vec_size(t1), 7)
  expect_equal(vec_size(t1), length(t1))
  expect_length(tiers(t2), 2)

  # Adding roles and labels works, including strata
  f <- mpg ~ X(hp) + M(gear) + drat + S(cyl)
  x <-
    distill_rune(f, label = list(gear ~ "Gears")) |>
    vec_data()
  expect_equal(x$role[x$runes == "gear"], "mediator")
  expect_equal(x$label[x$runes == "gear"], "Gears")
  expect_equal(x$side[x$runes == "cyl"], "meta")
  expect_equal(x$role[x$runes == "cyl"], "strata")

})

test_that("interaction terms are appropriately included", {

  f <- mpg ~ X(hp) + gear + M(cyl) + In(am)
  x <- distill_rune(f)
  expect_length(x, 5)
  expect_match(rhs(f, tidy = TRUE)[4], "am")

  f <- mpg ~ X(hp) + am + hp:am
  x <- distill_rune(f)
  expect_length(x, 4)
  expect_match(rhs(f, tidy = TRUE)[3], "hp:am")

  # Expect warning
  f <- mpg ~ hp
  x <- distill_rune(f)
  y <- distill_rune(
    "I",
    side = "right",
    role = "interaction",
    label = "Interaction Variable",
    description = "Interaction for the exposure variable",
    distribution = "binary",
    type = "character",
    subtype = "dichotomous"
  )

  expect_warning(distill_rune(mpg ~ hp + In(am)))

})

test_that("runes can be made from a fitted model", {

  # lm models
  m_lm <- lm(mpg ~ wt + hp + cyl, mtcars)
  t_lm <- distill_rune(m_lm)
  expect_length(t_lm, 4)

  # glm
  m_glm <- glm(am ~ wt + hp, mtcars, family = "binomial")
  t_glm <- distill_rune(m_glm, label = list(am ~ "Automatic Transmission"))
  expect_length(t_glm, 3)
  expect_equal(labels(t_glm)$am, "Automatic Transmission")

  # Model spec of parsip models
  if (isTRUE(requireNamespace("parsnip", quietly = TRUE))) {
    m_parsnip <-
      parsnip::linear_reg() |>
      parsnip::set_engine("lm") |>
      parsnip::fit(mpg ~ ., data = mtcars)
    t_parsnip <- distill_rune(m_parsnip)
    expect_length(t_parsnip, 11)
  }
})