#' Generalized Heckman Model Estimation
#'
#' @description
#' Estimates the parameters of a generalized Heckman selection model,
#' allowing covariates in the scale and correlation components.
#'
#' @param selection Formula for the selection equation.
#' @param outcome Formula for the outcome equation.
#' @param outcomeS Covariates for the log-scale model. It may be a matrix,
#'   vector, one-sided formula, or 1 for an intercept-only scale model.
#' @param outcomeC Covariates for the correlation model. It may be a matrix,
#'   vector, one-sided formula, or 1 for an intercept-only correlation model.
#' @param data Data frame containing all variables.
#' @param start Optional starting values. If supplied on the natural classical
#'   scale, sigma and rho are internally mapped to log(sigma) and atanh(rho).
#'
#' @return A list with the fitted model parameters and diagnostics.
#' @export
HeckmanGe <- function(selection, outcome, outcomeS = 1, outcomeC = 1,
                      data = sys.frame(sys.parent()), start = NULL) {

  components <- extract_model_components(selection = selection,
                                         outcome = outcome,
                                         outcomeS = outcomeS,
                                         outcomeC = outcomeC,
                                         data = data)
  XS <- components$XS
  YS <- components$YS
  NXS <- components$NXS
  XO <- components$XO
  YO <- components$YO
  NXO <- components$NXO
  Msigma <- components$Msigma
  NE <- components$NE
  Mrho <- components$Mrho
  NV <- components$NV
  YSLevels <- components$YSLevels

  YO_clean <- ifelse(is.na(YO), 0, YO)

  loglik_gen <- function(par) {
    if (any(!is.finite(par))) return(NA_real_)

    istartS <- seq_len(NXS)
    istartO <- seq.int(NXS + 1L, NXS + NXO)
    ilambda <- seq.int(NXS + NXO + 1L, NXS + NXO + NE)
    ikappa  <- seq.int(NXS + NXO + NE + 1L, NXS + NXO + NE + NV)

    gamma <- par[istartS]
    beta <- par[istartO]
    lambda <- par[ilambda]
    kappa <- par[ikappa]

    mu2 <- drop(XS %*% gamma)
    mu1 <- drop(XO %*% beta)
    sigma <- as.numeric(exp(Msigma %*% lambda))
    rho <- as.numeric(tanh(Mrho %*% kappa))

    r <- sqrt(pmax(1 - rho^2, .Machine$double.eps))
    z <- (YO_clean - mu1) / sigma
    zeta <- (mu2 + rho * z) / r

    ll0 <- stats::pnorm(-mu2, log.p = TRUE)
    ll1 <- stats::dnorm(z, log = TRUE) - log(sigma) +
      stats::pnorm(zeta, log.p = TRUE)

    sum(ifelse(YS == 0L, ll0, ll1))
  }

  gradlik_gen <- function(par) {
    istartS <- seq_len(NXS)
    istartO <- seq.int(NXS + 1L, NXS + NXO)
    ilambda <- seq.int(NXS + NXO + 1L, NXS + NXO + NE)
    ikappa  <- seq.int(NXS + NXO + NE + 1L, NXS + NXO + NE + NV)
    nParam <- NXS + NXO + NE + NV

    if (any(!is.finite(par))) return(rep(NA_real_, nParam))

    gamma <- par[istartS]
    beta <- par[istartO]
    lambda <- par[ilambda]
    kappa <- par[ikappa]

    XS0 <- XS[YS == 0L, , drop = FALSE]
    XS1 <- XS[YS == 1L, , drop = FALSE]
    XO1 <- XO[YS == 1L, , drop = FALSE]
    YO1 <- YO_clean[YS == 1L]
    MS1 <- Msigma[YS == 1L, , drop = FALSE]
    MR1 <- Mrho[YS == 1L, , drop = FALSE]

    mu20 <- drop(XS0 %*% gamma)
    mu21 <- drop(XS1 %*% gamma)
    mu11 <- drop(XO1 %*% beta)

    sigma1 <- as.numeric(exp(MS1 %*% lambda))
    rho1 <- as.numeric(tanh(MR1 %*% kappa))
    r1 <- sqrt(pmax(1 - rho1^2, .Machine$double.eps))
    z1 <- (YO1 - mu11) / sigma1
    zeta1 <- (mu21 + rho1 * z1) / r1

    lambda0 <- exp(stats::dnorm(-mu20, log = TRUE) -
                     stats::pnorm(-mu20, log.p = TRUE))
    lambda1 <- exp(stats::dnorm(zeta1, log = TRUE) -
                     stats::pnorm(zeta1, log.p = TRUE))

    gradient <- matrix(0, nrow = length(YS), ncol = nParam)
    gradient[YS == 0L, istartS] <- -XS0 * lambda0
    gradient[YS == 1L, istartS] <- XS1 * (lambda1 / r1)
    gradient[YS == 1L, istartO] <- XO1 *
      (z1 / sigma1 - lambda1 * rho1 / (sigma1 * r1))
    gradient[YS == 1L, ilambda] <- MS1 *
      (z1^2 - 1 - lambda1 * rho1 * z1 / r1)
    gradient[YS == 1L, ikappa] <- MR1 *
      (lambda1 * (z1 + rho1 * mu21) / r1)

    colSums(gradient)
  }

  if (is.null(start)) {
    start_classic <- tryCatch(step2(YS, XS, YO, XO), error = function(e) NULL)
    if (is.null(start_classic) || any(!is.finite(start_classic))) {
      message("Two-step starting values failed; using neutral default start values.")
      start_classic <- c(rep(0, NXS + NXO), sigma = 1, rho = 0)
    }
  } else {
    start_classic <- start
  }

  sigma0 <- max(start_classic[NXS + NXO + 1L], .Machine$double.eps)
  rho0 <- .ssmodels_clip_rho(start_classic[NXS + NXO + 2L])

  lambda_start <- c(log(sigma0), rep(0, max(NE - 1L, 0L)))
  kappa_start <- c(atanh(rho0), rep(0, max(NV - 1L, 0L)))

  start_opt <- c(start_classic[seq_len(NXS + NXO)], lambda_start, kappa_start)

  sigma_names <- if (NE == 1L) "sigma" else paste0("log.sigma.", colnames(Msigma))
  rho_names <- if (NV == 1L) "rho" else paste0("atanh.rho.", colnames(Mrho))
  names(start_opt) <- c(colnames(XS), colnames(XO), sigma_names, rho_names)

  theta_HG <- stats::optim(start_opt, loglik_gen, gradlik_gen,
                           method = "BFGS", hessian = TRUE,
                           control = list(fnscale = -1, maxit = 1000))

  vcov_link <- .ssmodels_safe_vcov(theta_HG$hessian)
  se_link <- .ssmodels_safe_se(vcov_link)
  coef_out <- theta_HG$par
  se_out <- se_link

  ilambda <- seq.int(NXS + NXO + 1L, NXS + NXO + NE)
  ikappa  <- seq.int(NXS + NXO + NE + 1L, NXS + NXO + NE + NV)

  if (NE == 1L) {
    coef_out[ilambda] <- exp(theta_HG$par[ilambda])
    se_out[ilambda] <- abs(coef_out[ilambda]) * se_link[ilambda]
  }
  if (NV == 1L) {
    coef_out[ikappa] <- tanh(theta_HG$par[ikappa])
    se_out[ikappa] <- (1 - coef_out[ikappa]^2) * se_link[ikappa]
  }
  names(coef_out) <- names(start_opt)
  names(se_out) <- names(start_opt)

  nObs <- length(YS)
  nParam <- length(start_opt)
  result <- list(
    coefficients   = coef_out,
    coefficients_link = theta_HG$par,
    value          = theta_HG$value,
    loglik         = theta_HG$value,
    counts         = theta_HG$counts[2],
    hessian        = theta_HG$hessian,
    fisher_infoHG  = vcov_link,
    prop_sigmaHG   = se_out,
    level          = YSLevels,
    nObs           = nObs,
    nParam         = nParam,
    N0             = sum(YS == 0L),
    N1             = sum(YS == 1L),
    NXS            = NXS,
    NXO            = NXO,
    df             = nObs - nParam,
    aic            = -2 * theta_HG$value + 2 * nParam,
    bic            = -2 * theta_HG$value + nParam * log(nObs),
    initial.value  = start_opt,
    NE             = NE,
    NV             = NV,
    Msigma         = Msigma,
    Mrho           = Mrho
  )

  class(result) <- c("HeckmanGe", class(theta_HG))
  result
}
