R/strata-utils.R

In R4VN: Health Data Analysis and Publication-Ready Reporting

# ============================================================================
# R4VN shared variable-list and hierarchical-by helpers
# ============================================================================
# These helpers implement one consistent R4VN convention across statistical
# commands and graphs:
#
#   by = vars(region, sex, outcome)
#
# means that `region` and `sex` are nested stratification (superby) variables
# and `outcome` is the innermost grouping/outcome variable used by the analysis.
# A one-variable specification, by = vars(outcome), is equivalent to
# by = outcome.  This convention lets commands share the same mental model
# without forcing each function to invent its own superby argument.

.r4vn_is_vars_call <- function(expr) {
  is.call(expr) && length(expr) >= 1L && identical(as.character(expr[[1L]]), "vars")
}

# TRUE for either a literal vars(...) call or a saved r4vn_vars selector.
# A literal data-column name wins over an object with the same name so legacy
# calls such as by = group retain their established meaning.
.r4vn_is_vars_spec_expr <- function(expr, data = NULL, env = parent.frame()) {
  if (.r4vn_is_vars_call(expr)) return(TRUE)
  if (is.symbol(expr) && !is.null(data) && as.character(expr) %in% names(data)) return(FALSE)
  value <- tryCatch(eval(expr, envir = env), error = function(e) NULL)
  inherits(value, "r4vn_vars")
}

.r4vn_expr_is_null <- function(expr) {
  is.null(expr) || identical(expr, quote(NULL)) || identical(expr, quote(expr = ))
}

.r4vn_resolve_name_spec <- function(expr, data, env = parent.frame(), arg = "variable",
                                    allow_null = FALSE, multiple = TRUE,
                                    default_type = "auto") {
  if (.r4vn_expr_is_null(expr)) {
    if (isTRUE(allow_null)) return(character())
    stop(sprintf("`%s` is required.", arg), call. = FALSE)
  }

  # vars(...) is intentionally evaluated only after `data` is known so that
  # deferred selectors (., wildcards, exclusions) keep working.
  if (.r4vn_is_vars_call(expr)) {
    spec <- eval(expr, envir = env)
    resolved <- .r4vn_resolve_vars(
      spec,
      data = data,
      default_type = default_type,
      strict = TRUE
    )
    out <- as.character(resolved$variable)
  } else if (inherits(expr, "r4vn_vars")) {
    resolved <- .r4vn_resolve_vars(
      expr,
      data = data,
      default_type = default_type,
      strict = TRUE
    )
    out <- as.character(resolved$variable)
  } else if (is.symbol(expr)) {
    nm <- as.character(expr)
    # First prefer a literal data-column name (including R4VN declaration
    # prefixes such as c.age or b2.group). If the symbol is not a column, try
    # evaluating it so saved selectors such as `g <- vars(region, sex)` and
    # character-name objects remain usable.
    nm0 <- sub("^(?:b[1-9][0-9]*\\.|[cqf]\\.)", "", nm, perl = TRUE)
    if (nm0 %in% names(data)) {
      out <- nm0
    } else {
      value <- tryCatch(eval(expr, envir = env), error = function(e) NULL)
      if (inherits(value, "r4vn_vars")) {
        resolved <- .r4vn_resolve_vars(value, data = data, default_type = default_type, strict = TRUE)
        out <- as.character(resolved$variable)
      } else if (is.character(value) && length(value) && all(value %in% names(data))) {
        out <- value
      } else {
        stop(sprintf("Variable `%s` from `%s` was not found in `data`.", nm0, arg), call. = FALSE)
      }
    }
  } else if (is.character(expr)) {
    out <- as.character(expr)
    missing_names <- setdiff(out, names(data))
    if (length(missing_names)) {
      stop(sprintf(
        "Variable(s) from `%s` were not found in `data`: %s.",
        arg, paste(missing_names, collapse = ", ")
      ), call. = FALSE)
    }
  } else {
    value <- tryCatch(eval(expr, envir = env), error = function(e) NULL)
    if (inherits(value, "r4vn_vars")) {
      resolved <- .r4vn_resolve_vars(value, data = data, default_type = default_type, strict = TRUE)
      out <- as.character(resolved$variable)
    } else if (is.character(value) && length(value) && all(value %in% names(data))) {
      out <- value
    } else {
      stop(sprintf(
        "`%s` must be a variable name, a character variable name, or `vars(...)`.",
        arg
      ), call. = FALSE)
    }
  }

  out <- unique(out[nzchar(out)])
  if (!length(out) && !isTRUE(allow_null)) stop(sprintf("`%s` selected no variables.", arg), call. = FALSE)
  if (!isTRUE(multiple) && length(out) != 1L) {
    stop(sprintf("`%s` must select exactly one variable.", arg), call. = FALSE)
  }
  out
}

.r4vn_by_spec <- function(expr, data, env = parent.frame(), arg = "by", allow_null = TRUE) {
  names <- .r4vn_resolve_name_spec(
    expr,
    data = data,
    env = env,
    arg = arg,
    allow_null = allow_null,
    multiple = TRUE,
    default_type = "auto"
  )
  if (!length(names)) {
    return(list(
      all = character(), strata = character(), by = NULL,
      all_labels = character(), strata_labels = character(), by_label = NULL
    ))
  }
  out <- list(
    all = names,
    strata = if (length(names) > 1L) head(names, -1L) else character(),
    by = tail(names, 1L)
  )
  out$all_labels <- vapply(out$all, function(nm) .r4vn_variable_label(data, nm), character(1))
  out$strata_labels <- vapply(out$strata, function(nm) .r4vn_variable_label(data, nm), character(1))
  out$by_label <- .r4vn_variable_label(data, out$by)
  out
}

.r4vn_vars_spec_names <- function(expr, data, env = parent.frame(), arg = "vars",
                                  allow_null = TRUE, numeric_only = FALSE) {
  out <- .r4vn_resolve_name_spec(
    expr,
    data = data,
    env = env,
    arg = arg,
    allow_null = allow_null,
    multiple = TRUE,
    default_type = "auto"
  )
  if (isTRUE(numeric_only) && length(out)) {
    bad <- out[!vapply(data[out], is.numeric, logical(1))]
    if (length(bad)) {
      stop(sprintf("`%s` must select numeric variables; nonnumeric: %s.", arg, paste(bad, collapse = ", ")), call. = FALSE)
    }
  }
  out
}

.r4vn_strata_indices <- function(data, strata, drop_missing = TRUE) {
  if (!length(strata)) {
    return(list(Overall = seq_len(nrow(data))))
  }
  if (!all(strata %in% names(data))) stop("Internal error: unknown stratification variable.", call. = FALSE)

  d <- data[strata]
  complete <- stats::complete.cases(d)
  idx <- seq_len(nrow(data))
  if (isTRUE(drop_missing)) {
    d <- d[complete, , drop = FALSE]
    idx <- idx[complete]
  } else {
    for (nm in names(d)) {
      z <- as.character(d[[nm]])
      z[is.na(z)] <- "Missing"
      d[[nm]] <- z
    }
  }
  if (!nrow(d)) return(list())

  # Preserve observed/factor order rather than alphabetical interaction order.
  parts <- lapply(d, function(z) {
    if (is.factor(z)) factor(z, levels = levels(z))
    else factor(z, levels = unique(z[!is.na(z)]))
  })
  key <- do.call(interaction, c(parts, list(drop = TRUE, sep = " | ", lex.order = TRUE)))
  split(idx, key, drop = TRUE)
}

.r4vn_variable_label <- function(data, name) {
  if (!name %in% names(data)) return(as.character(name)[1L])
  label <- attr(data[[name]], "label", exact = TRUE)
  if (is.null(label) || !length(label) || is.na(label[1L]) ||
      !nzchar(trimws(as.character(label[1L])))) {
    return(as.character(name)[1L])
  }
  trimws(as.character(label[1L]))
}

.r4vn_value_label <- function(x, index) {
  value <- x[index][1L]
  if (is.na(value)) return("Missing")
  if (is.factor(x)) return(as.character(value))

  value_labels <- attr(x, "labels", exact = TRUE)
  if (!is.null(value_labels) && length(value_labels) &&
      !is.null(names(value_labels))) {
    hit <- which(!is.na(value_labels) &
                   as.character(value_labels) == as.character(value))
    if (length(hit) && nzchar(names(value_labels)[hit[1L]])) {
      return(names(value_labels)[hit[1L]])
    }
  }
  as.character(value)
}

.r4vn_stratum_component <- function(data, name, index) {
  paste0(.r4vn_variable_label(data, name), " (",
         .r4vn_value_label(data[[name]], index), ")")
}

.r4vn_stratum_label <- function(data, strata, index) {
  if (!length(strata)) return("Overall")
  if (!length(index)) return("")
  i <- index[1L]
  paste(vapply(
    strata,
    function(nm) .r4vn_stratum_component(data, nm, i),
    character(1)
  ), collapse = " > ")
}

.r4vn_hierarchy_note <- function(data, spec, innermost = "innermost group") {
  strata_labels <- vapply(
    spec$strata, function(nm) .r4vn_variable_label(data, nm), character(1)
  )
  by_label <- if (is.null(spec$by)) NULL else .r4vn_variable_label(data, spec$by)
  paste0(
    "Hierarchical by: ", paste(strata_labels, collapse = " > "),
    if (length(by_label)) paste0("; ", innermost, ": ", by_label, ".") else "."
  )
}

.r4vn_prepend_context <- function(tab, variable = NULL, stratum = NULL) {
  if (!is.data.frame(tab)) return(tab)
  # split()/vapply() can leave names on scalar labels. Passing a named scalar
  # directly to data.frame()/cbind() makes base R interpret that name as a row
  # name and emits "row names were found from a short variable". Strip names
  # and recycle explicitly so hierarchical results remain warning-free.
  if (!is.null(variable)) {
    variable <- unname(as.character(variable))[1L]
    tab <- cbind(Variable = rep(variable, nrow(tab)), tab, stringsAsFactors = FALSE)
  }
  if (!is.null(stratum)) {
    stratum <- unname(as.character(stratum))[1L]
    tab <- cbind(Stratum = rep(stratum, nrow(tab)), tab, stringsAsFactors = FALSE)
  }
  rownames(tab) <- NULL
  tab
}

.r4vn_bind_section_tables <- function(results, labels = NULL, variables = NULL) {
  if (!length(results)) return(list())
  section_names <- unique(unlist(lapply(results, function(z) names(z$sections)), use.names = FALSE))
  out <- setNames(vector("list", length(section_names)), section_names)

  for (nm in section_names) {
    pieces <- vector("list", length(results))
    keep <- logical(length(results))
    for (i in seq_along(results)) {
      z <- results[[i]]$sections[[nm]]
      if (is.null(z)) next
      if (is.matrix(z)) z <- as.data.frame(z, stringsAsFactors = FALSE, check.names = FALSE)
      if (!is.data.frame(z)) next
      if (!is.null(variables)) z <- .r4vn_prepend_context(z, variable = variables[i])
      if (!is.null(labels)) z <- .r4vn_prepend_context(z, stratum = labels[i])
      pieces[[i]] <- z
      keep[i] <- TRUE
    }
    pieces <- pieces[keep]
    if (!length(pieces)) next
    # Union columns so sections can still combine when a subgroup omits a
    # diagnostic column because it was not estimable.
    cols <- unique(unlist(lapply(pieces, names), use.names = FALSE))
    pieces <- lapply(pieces, function(z) {
      miss <- setdiff(cols, names(z))
      for (m in miss) z[[m]] <- ""
      z[cols]
    })
    out[[nm]] <- do.call(rbind, pieces)
    rownames(out[[nm]]) <- NULL
  }
  out[!vapply(out, is.null, logical(1))]
}

.r4vn_stat_collection <- function(title, results, labels = NULL, variables = NULL,
                                  notes = NULL, call = NULL) {
  sections <- .r4vn_bind_section_tables(results, labels = labels, variables = variables)
  .r4vn_result(
    title,
    sections = sections,
    notes = notes,
    raw = list(results = results, labels = labels, variables = variables),
    call = call
  )
}

# Graph result collections ----------------------------------------------------
.r4vn_graph_set <- function(graphs, labels = NULL, call = NULL,
                            combined = FALSE, ncol = NULL, file = NULL) {
  structure(
    list(
      graphs = graphs,
      labels = labels %||% names(graphs),
      call = call,
      combined = isTRUE(combined),
      ncol = ncol,
      file = file
    ),
    class = "r4vn_graph_set"
  )
}

#' Print a collection of R4VN graphs
#'
#' Returns the collection invisibly without writing graph status or row counts
#' to the console. Use `plot()` to redraw the collection.
#'
#' @param x An object returned when an R4VN graph command creates more than one
#'   graph because `vars()` and/or hierarchical `by = vars(...)` were used.
#' @param ... Not used.
#' @return `x`, invisibly.
#' @export
print.r4vn_graph_set <- function(x, ...) {
  invisible(x)
}