R/summons.R

In asshah4/forks: Formula Spells

# Formula class ----------------------------------------------------------------

#' Summoning Formulas
#'
#' @param order Requested order that formulas are decomposed
#'   into. The default is to return ALL formula decompositions. Options include
#'   any integer inclusive from _1_ to _4_. The default is `2L:4L`, which
#'   includes every functional formula and its parent.
#' @name formulas
#' @export
fmls <- function(x = unspecified(),
                 role = list(),
                 tier = list(),
                 label = list(),
                 strata = character(),
                 pattern = character(),
                 order = 2L:4L,
                 ...) {

  # Early break and validation
  if (length(x) == 0) {
    return(new_formulas())
  }

  if (!any(
    c(
      "spell",
      "rune",
      "fmls",
      "character",
      "formula"
    ) %in% class(x)
  )) {
    stop("`fmls()` is not defined for a `",
      class(x)[1],
      "` object.",
      call. = FALSE
    )
  }

  if (is.numeric(order) & !all(order %in% 1:4)) {
    stop("The order should be an integer range between 1 to 4.")
  }

  if ("spell" %in% class(x) & length(x) > 1) {
    stop("`fmls()` can only accept 1 `spell` object at a time.")
  }

  # Extract and/or generate runes
  if ("character" %in% class(x)) {
    t <- distill_rune(stats::formula(x))
  } else if ("formula" %in% class(x)) {
    t <- distill_rune(x)
  } else if ("spell" %in% class(x)) {
    t <- distill_rune(x)
    pattern <- field(x, "pattern")
  } else if ("rune" %in% class(x)) {
    t <- x
  }

  # Add strata back in (if not already available)
  if (length(strata) > 0) {
    t <- add_strata(t, strata)
  }

  # Update runes
  t <-
    t |>
    set_roles(roles = formula_to_named_list(role)) |>
    set_tiers(tiers = formula_to_named_list(tier)) |>
    set_labels(labels = formula_to_named_list(label))

  # Generate or re-generate spell, should be only a single spell to start
  # Filter out just the second order formulas for now
  ancestor <- suppressMessages(deparse1(stats::formula(t)))
  n <- decipher(t)
  s <- suppressMessages(cast_spell(t, pattern = pattern))
  o <- min(order)
  while (n >= o) {
    s <- recompose_roles(s)
    n <- n - 1
  }

  # Turn each of these into formulas
  fa <- summon_formulas()
  for (i in seq_along(s)) {

    # From all of the spells, obtain the expansion patterns if possible
    if (field(s[i], "order") %in% c(1, 4)) {
      fl <- field(s[i], "formula")
    } else {
      fl <- decompose_patterns(s[i])
    }
    tl <- field(s[i], "runes")[[1]]

    for (j in seq_along(fl)) {
      fx <- stats::formula(fl[[j]])
      tms <-
        match_runes(tl, fx) |>
        add(get_runes(tl, "role", "strata"))

      f <- new_formulas(
        formula = fl[[j]],
        number = j,
        n = length(tms),
        order = decipher(tms),
        left = list(lhs(fx)),
        right = list(rhs(fx)),
        outcome = list(get_runes(tms, "role", "outcome")),
        predictor = list(get_runes(tms, "role", "predictor")),
        exposure = list(get_runes(tms, "role", "exposure")),
        confounder = list(get_runes(tms, "role", "confounder")),
        mediator = list(get_runes(tms, "role", "mediator")),
        interaction = list(get_runes(tms, "role", "interaction")),
        unknown = list(get_runes(tms, "role", "unknown")),
        # Strata may get lost unless brought in from above
        strata = list(get_runes(tl, "role", "strata")),
        pattern = field(s[i], "pattern"),
        ancestor = ancestor,
        source = class(x)[1]
      )

      fa <- append(fa, f)
    }
  }

  # Return formulas
  fa |>
    vec_data() |>
    {
      \(.x) {.x[.x$order %in% order, ]}
    }() |>
    vec_restore(to = summon_formulas()) |>
    unique()
}

#' @rdname formulas
#' @export
summon_formulas <- fmls

# Record definition ------------------------------------------------------------

#' Record of formula archetypes
#' @keywords internal
#' @noRd
new_formulas <- function(formula_list = character(),
                         number = integer(),
                         n = integer(),
                         order = integer(),
                         left = list(),
                         right = list(),
                         outcome = list(),
                         predictor = list(),
                         exposure = list(),
                         confounder = list(),
                         mediator = list(),
                         interaction = list(),
                         unknown = list(),
                         strata = list(),
                         pattern = character(),
                         ancestor = character(),
                         source = character()) {

  # Validation
  vec_assert(formula_list, ptype = character())
  vec_assert(number, ptype = integer())
  vec_assert(n, ptype = integer())
  vec_assert(order, ptype = integer())

  # Character vectors based on sides
  vec_assert(left, ptype = list())
  vec_assert(right, ptype = list())

  # Roles (in rune() format)
  vec_assert(outcome, ptype = list())
  vec_assert(predictor, ptype = list())
  vec_assert(exposure, ptype = list())
  vec_assert(confounder, ptype = list())
  vec_assert(mediator, ptype = list())
  vec_assert(interaction, ptype = list())
  vec_assert(unknown, ptype = list())
  vec_assert(strata, ptype = list())

  # Specification information
  vec_assert(pattern, ptype = character())
  vec_assert(ancestor, ptype = character())
  vec_assert(source, ptype = character())

  new_rcrd(
    fields = list(
      "formulas" = formula_list,
      "number" = number,
      "n" = n,
      "order" = order,
      "left" = left,
      "right" = right,
      "outcome" = outcome,
      "predictor" = predictor,
      "exposure" = exposure,
      "confounder" = confounder,
      "mediator" = mediator,
      "interaction" = interaction,
      "unknown" = unknown,
      "strata" = strata,
      "pattern" = pattern,
      "ancestor" = ancestor,
      "source" = source
    ),
    class = "fmls"
  )
}

#' @keywords internal
#' @noRd
methods::setOldClass(c("fmls", "rcrds_rcrd"))

# Output -----------------------------------------------------------------------

#' @export
format.fmls <- function(x, ...) {
  field(x, "formulas")
}


#' @export
obj_print_data.fmls <- function(x, ...) {

  # Colorful printing
  if (vec_size(x) == 0) {
    fmt <- new_spell()
  } else {

    # Depending on length
    if (length(x) > 1) {
      cat(format(x), sep = "\n")
    } else {
      cat(format(x))
    }
  }
}

#' @export
vec_ptype_full.fmls <- function(x, ...) {
  "fmls"
}

#' @export
vec_ptype_abbr.fmls <- function(x, ...) {
  "fmls"
}

# Casting and coercion ---------------------------------------------------------

#' @export
vec_ptype2.fmls.fmls <- function(x, y, ...) {
  x
}

#' @export
vec_cast.fmls.fmls <- function(x, to, ...) {
  x
}


#' @export
vec_ptype2.fmls.character <- function(x, y, ...) {
  y
}

#' @export
vec_ptype2.character.fmls <- function(x, y, ...) {
  x
}

#' @export
vec_cast.character.fmls <- function(x, to, ...) {
  format(x) # Returns a character class by default
}

#' @export
vec_ptype2.fmls.rune <- function(x, y, ...) {
  y
}

#' @export
vec_ptype2.rune.fmls <- function(x, y, ...) {
  x
}

#' @export
vec_cast.rune.fmls <- function(x, to, ...) {
  distill_rune.fmls(x)
}

#' @export
vec_ptype2.fmls.spell <- function(x, y, ...) {
  x
}

#' @export
vec_ptype2.spell.fmls <- function(x, y, ...) {
  y
}

#' @export
vec_cast.fmls.spell <- function(x, to, ...) {
  format(x) |>
    stats::as.formula() |>
    summon_formulas()
}

#' @export
#' @importFrom stats formula
formula.fmls <- function(x, ...) {
  lapply(x, FUN = function(.x) {
    format(.x) |>
      stats::as.formula()
  })
}