`step_linadjust` for linear adjustment of covariates

# 1. User-facing function
step_linadjust <- function(
  recipe,
  ...,
  role = NA,
  trained = FALSE,
  remove_vars = NULL, # The "Batch" to remove
  keep_vars = NULL, # The "Design" to preserve
  models = NULL,
  drop = c("both", "remove", "none"), # drop the nuisance variables from the output?
  skip = FALSE,
  id = rand_id("linadjust")
) {
  add_step(
    recipe,
    step_linadjust_new(
      terms = enquos(...),
      role = role,
      trained = trained,
      remove_vars = remove_vars,
      keep_vars = keep_vars,
      models = models,
      drop = drop,
      skip = skip,
      id = id
    )
  )
}

# 2. Step Constructor
step_linadjust_new <- function(
  terms,
  role,
  trained,
  remove_vars,
  keep_vars,
  models,
  drop,
  skip,
  id
) {
  step(
    subclass = "linadjust",
    terms = terms,
    role = role,
    trained = trained,
    remove_vars = remove_vars,
    keep_vars = keep_vars,
    models = models,
    drop = drop,
    skip = skip,
    id = id
  )
}

# 3. Prep Method (Fits the model Y ~ Remove + Keep)
prep.step_linadjust <- function(x, training, info = NULL) {
  # Identify target columns (Y)

  col_names <- recipes_eval_select(x$terms, training, info)

  # Identify Nuisance columns (Batch)
  if (is.null(x$remove_vars)) {
    cli::cli_abort(
      c(
        "The `remove_vars` argument must be specified.",
        "i" = "This is the variable(s) you want to remove the effect of."
      )
    )
  }

  remove_names <- recipes_eval_select(x$remove_vars, training, info)

  # Identify Preserved columns (Design)
  # Handle case where keep_vars is NULL
  if (!is.null(x$keep_vars)) {
    keep_names <- recipes_eval_select(x$keep_vars, training, info)

    if (any(keep_names %in% remove_names)) {
      cli::cli_abort(
        c(
          "The `keep_vars` and `remove_vars` selectors must be disjoint.",
          "x" = "The following variables are in both: {intersect(keep_names, remove_names)}"
        )
      )
    }
  } else {
    keep_names <- NULL
  }

  for (ic in c(remove_names, keep_names)) {
    if (is.factor(training[[ic]])) {
      training[[ic]] <- C(droplevels(training[[ic]]), contr.sum)
    } else if (is.numeric(training[[ic]])) {
      training[[ic]] <- scale(training[[ic]])
    } else {
      cli::cli_abort(
        c(
          "The `remove_vars` and `keep_vars` selectors must be either factors or numeric.",
          "x" = "The following variable is neither: {ic}"
        )
      )
    }
  }

  model_list <- list()

  for (col in col_names) {
    # Create formula: Target ~ Remove1 + Keep1 + ...
    # We combine both sets of variables for the fit
    ff <- reformulate(
      response = col,
      termlabels = c(remove_names, keep_names)
    )

    # Fit and store the model
    model_list[[col]] <- butcher::butcher(lm(ff, data = training))
  }

  drop <- match.arg(x$drop, choices = c("both", "remove", "none"))

  step_linadjust_new(
    terms = col_names,
    role = x$role,
    trained = TRUE,
    remove_vars = remove_names,
    keep_vars = keep_names,
    models = model_list,
    drop = drop,
    skip = x$skip,
    id = x$id
  )
}

# 4. Bake Method (Subtracts ONLY the Remove effect)
bake.step_linadjust <- function(object, new_data, ...) {
  # Get names of the variables we want to remove effects for
  # We need to re-evaluate the selector to get string names
  remove_names <- names(object$remove_vars)
  keep_names <- names(object$keep_vars)

  for (col in names(object$models)) {
    model <- object$models[[col]]

    # Crucial Step: use type = "terms"
    # This returns a matrix with one column per independent variable,
    # representing that variable's contribution to the prediction.
    # It handles factors (dummification) automatically.
    term_preds <- predict(model, newdata = new_data, type = "terms")

    # Identify which columns in the term matrix correspond to our `remove_vars`
    # Note: `predict` names columns by the variable name.
    cols_to_subtract <- intersect(colnames(term_preds), remove_names)

    if (length(cols_to_subtract) > 0) {
      # Sum the effects of the nuisance variables
      nuisance_effect <- rowSums(term_preds[, cols_to_subtract, drop = FALSE])

      # Subtract nuisance effect from original data
      # Result = (Signal + Batch + Noise) - (Batch) = Signal + Noise
      new_data[[col]] <- new_data[[col]] - nuisance_effect
    }
  }

  if (object$drop == "remove") {
    new_data <- new_data[,
      !(colnames(new_data) %in% remove_names),
      drop = FALSE
    ]
  } else if (object$drop == "both") {
    new_data <- new_data[,
      !(colnames(new_data) %in% c(remove_names, keep_names)),
      drop = FALSE
    ]
  }

  tibble::as_tibble(new_data)
}