#' Heckman's two-step method
#'
#' @description
#' Estimate classical Heckman model starting values by Heckman's two-step method.
#'
#' @param YS Selection vector coded as 0/1.
#' @param XS Selection matrix.
#' @param YO Outcome vector.
#' @param XO Covariate matrix for the outcome equation.
#'
#' @return A numeric vector with selection coefficients, outcome coefficients,
#'   sigma and rho.
#'
#' @importFrom stats binomial glm lm.fit dnorm pnorm na.omit
#' @export
step2 <- function(YS, XS, YO, XO) {
  YS <- as.integer(YS)
  if (!all(stats::na.omit(unique(YS)) %in% c(0L, 1L))) {
    stop("YS must be coded as 0/1.")
  }

  XS <- as.matrix(XS)
  XO <- as.matrix(XO)
  selected <- YS == 1L

  if (sum(selected) <= ncol(XO)) {
    stop("There are not enough selected observations to estimate the outcome equation.")
  }

  fit1 <- stats::glm.fit(x = XS, y = YS, family = stats::binomial(link = "probit"))
  eta <- drop(XS %*% stats::coef(fit1))
  imr <- exp(stats::dnorm(eta, log = TRUE) - stats::pnorm(eta, log.p = TRUE))

  XO1 <- XO[selected, , drop = FALSE]
  YO1 <- YO[selected]
  imr1 <- imr[selected]
  eta1 <- eta[selected]

  fit2 <- stats::lm.fit(x = cbind(XO1, IMR = imr1), y = YO1)
  coef2 <- stats::coef(fit2)
  beta <- coef2[seq_len(ncol(XO))]
  imr_coef <- coef2[ncol(XO) + 1L]

  delta <- imr1 * (imr1 + eta1)
  variance <- (sum(fit2$residuals^2) + imr_coef^2 * sum(delta)) / sum(selected)
  if (!is.finite(variance) || variance <= 0) {
    variance <- stats::var(YO1, na.rm = TRUE)
  }
  sigma <- sqrt(variance)
  rho <- imr_coef / sigma
  rho <- .ssmodels_clip_rho(rho)

  out <- c(stats::coef(fit1), beta, sigma = sigma, rho = rho)
  names(out) <- c(colnames(XS), colnames(XO), "sigma", "rho")
  out
}
