#' Extract model components from formulas
#'
#' @description
#' This function extracts model matrices and response vectors used by the
#' sample-selection model fitting functions. It standardizes the selection
#' response to a numeric indicator, where 1 denotes selected/observed and
#' 0 denotes not selected/not observed.
#'
#' @param selection A formula representing the selection equation.
#' @param outcome A formula representing the outcome equation.
#' @param data A data frame containing all variables used in the formulas.
#' @param outcomeS Optional matrix, vector, one-sided formula, or 1 for the
#'   log-scale model in the generalized Heckman model.
#' @param outcomeC Optional matrix, vector, one-sided formula, or 1 for the
#'   correlation model in the generalized Heckman model.
#' @param drop.levels Logical; if TRUE, unused factor levels are dropped.
#'
#' @return A list with selection, outcome, dispersion and correlation matrices.
#'
#' @importFrom stats model.frame model.matrix model.response terms na.pass
#' @export
extract_model_components <- function(selection, outcome, data,
                                     outcomeS = NULL, outcomeC = NULL,
                                     drop.levels = TRUE) {

  make_aux_matrix <- function(aux, data, n, prefix) {
    if (is.null(aux)) {
      mat <- matrix(1, nrow = n, ncol = 1)
      colnames(mat) <- "(Intercept)"
      return(mat)
    }

    if (length(aux) == 1L && is.numeric(aux)) {
      mat <- matrix(as.numeric(aux), nrow = n, ncol = 1)
      colnames(mat) <- "(Intercept)"
      return(mat)
    }

    if (inherits(aux, "formula")) {
      mf <- model.frame(formula = aux, data = data,
                        drop.unused.levels = drop.levels,
                        na.action = na.pass)
      mat <- model.matrix(terms(mf), mf)
      return(mat)
    }

    if (is.data.frame(aux) || is.matrix(aux)) {
      mat <- as.matrix(aux)
    } else {
      mat <- matrix(aux, ncol = 1)
      colnames(mat) <- deparse(substitute(aux))
    }

    if (nrow(mat) != n) {
      stop(prefix, " must have the same number of rows as data.")
    }

    if (is.null(colnames(mat))) {
      colnames(mat) <- paste0(prefix, seq_len(ncol(mat)))
    }

    first_is_intercept <- ncol(mat) >= 1L &&
      all(is.finite(mat[, 1])) && all(abs(mat[, 1] - 1) < 1e-12)

    if (!first_is_intercept) {
      mat <- cbind("(Intercept)" = 1, mat)
    }

    mat
  }

  mfS <- model.frame(
    formula = selection,
    data = data,
    drop.unused.levels = drop.levels,
    na.action = na.pass
  )
  mtS <- terms(mfS)
  XS <- model.matrix(mtS, mfS)
  YS_raw <- model.response(mfS)

  YS_factor <- as.factor(YS_raw)
  YSLevels <- levels(YS_factor)
  if (length(YSLevels) != 2L) {
    stop("The selection equation must have exactly two observed levels.")
  }

  if (is.numeric(YS_raw) && all(stats::na.omit(unique(YS_raw)) %in% c(0, 1))) {
    YS <- as.integer(YS_raw)
  } else {
    YS <- as.integer(YS_factor == YSLevels[2L])
  }

  mfO <- model.frame(
    formula = outcome,
    data = data,
    drop.unused.levels = drop.levels,
    na.action = na.pass
  )
  mtO <- terms(mfO)
  XO <- model.matrix(mtO, mfO)
  YO <- model.response(mfO)

  Msigma <- make_aux_matrix(outcomeS, data, nrow(XS), "outcomeS")
  Mrho   <- make_aux_matrix(outcomeC, data, nrow(XS), "outcomeC")

  if (nrow(XO) != nrow(XS) || length(YO) != length(YS)) {
    stop("Selection and outcome equations must refer to the same number of rows.")
  }

  keep <- !is.na(YS) & stats::complete.cases(XS)
  keep <- keep & (YS == 0 | (!is.na(YO) & stats::complete.cases(XO)))
  keep <- keep & (YS == 0 | (stats::complete.cases(Msigma) & stats::complete.cases(Mrho)))

  XS <- XS[keep, , drop = FALSE]
  YS <- YS[keep]
  XO <- XO[keep, , drop = FALSE]
  YO <- YO[keep]
  Msigma <- Msigma[keep, , drop = FALSE]
  Mrho <- Mrho[keep, , drop = FALSE]

  storage.mode(YS) <- "integer"

  list(
    XS = XS,
    YS = YS,
    NXS = ncol(XS),
    XO = XO,
    YO = YO,
    NXO = ncol(XO),
    Msigma = Msigma,
    NE = ncol(Msigma),
    Mrho = Mrho,
    NV = ncol(Mrho),
    YSLevels = YSLevels,
    kept = keep
  )
}
