#' Heckman-t Model fit Function
#'
#' @description
#' Estimates the parameters of the Heckman-t model. The scale, correlation and
#' degrees of freedom are optimized on unconstrained scales to improve numerical
#' stability: sigma = exp(eta_sigma), rho = tanh(eta_rho), and df = 2 + exp(eta_df).
#'
#' @param selection Selection equation.
#' @param outcome Primary regression equation.
#' @param data Database.
#' @param df Initial value for the degrees of freedom. Values greater than 2 are
#'   recommended because the model uses a finite-variance Student-t error.
#' @param start Optional starting values on the natural scale
#'   (selection, outcome, sigma, rho, df).
#'
#' @return A list with estimated coefficients, log-likelihood and diagnostics.
#'
#' @examples
#' data(MEPS2001)
#' dados <- MEPS2001[seq_len(500), ]
#' selectEq <- dambexp ~ age + female + educ + blhisp + totchr + ins + income
#' outcomeEq <- lnambx ~ age + female + educ + blhisp + totchr + ins
#' HeckmantS(selectEq, outcomeEq, data = dados, df = 12)
#'
#' @importFrom Rdpack reprompt
#' @references Marchenko, Y. V. and Genton, M. G. (2012). A Heckman selection-t
#' model. \emph{Journal of the American Statistical Association}, 107(497),
#' 304-317.
#' @export HeckmantS
#' @export
HeckmantS <- function(selection, outcome, data = sys.frame(sys.parent()),
                      df, start = NULL) {

  components <- extract_model_components(selection = selection, outcome = outcome, data = data)
  XS <- components$XS
  YS <- components$YS
  NXS <- components$NXS
  XO <- components$XO
  YO <- components$YO
  NXO <- components$NXO
  YSLevels <- components$YSLevels
  YO_clean <- ifelse(is.na(YO), 0, YO)

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

    istartS <- seq_len(NXS)
    istartO <- seq.int(NXS + 1L, NXS + NXO)
    isigma <- NXS + NXO + 1L
    irho <- isigma + 1L
    iv <- irho + 1L

    gamma <- par[istartS]
    beta <- par[istartO]
    sigma <- exp(par[isigma])
    rho <- tanh(par[irho])
    nu <- 2 + exp(par[iv])

    mu2 <- drop(XS %*% gamma)
    mu1 <- drop(XO %*% beta)
    z <- (YO_clean - mu1) / sigma
    r <- sqrt(pmax(1 - rho^2, .Machine$double.eps))
    q <- sqrt((nu + 1) / (nu + z^2))
    eta <- q * ((rho * z + mu2) / r)

    ll0 <- stats::pt(-mu2, df = nu, log.p = TRUE)
    ll1 <- stats::dt(z, df = nu, log = TRUE) - log(sigma) +
      stats::pt(eta, df = nu + 1, log.p = TRUE)

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

  if (missing(df) || is.null(df) || !is.finite(df) || df <= 2) {
    df <- 12
  }

  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)
    }
    start <- c(start_classic[seq_len(NXS + NXO)],
               sigma = start_classic[NXS + NXO + 1L],
               rho = start_classic[NXS + NXO + 2L],
               df = df)
  }

  if (length(start) != NXS + NXO + 3L) {
    stop("start must have length NXS + NXO + 3: selection, outcome, sigma, rho and df.")
  }

  start_opt <- start
  start_opt[NXS + NXO + 1L] <- log(max(start[NXS + NXO + 1L], .Machine$double.eps))
  start_opt[NXS + NXO + 2L] <- atanh(.ssmodels_clip_rho(start[NXS + NXO + 2L]))
  start_opt[NXS + NXO + 3L] <- log(max(start[NXS + NXO + 3L] - 2, .Machine$double.eps))
  names(start_opt) <- c(colnames(XS), colnames(XO), "log.sigma", "atanh.rho", "log.df.minus.2")

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

  vcov_link <- .ssmodels_safe_vcov(theta_tS$hessian)
  se_link <- .ssmodels_safe_se(vcov_link)

  coef_out <- theta_tS$par
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  iv <- irho + 1L
  coef_out[isigma] <- exp(theta_tS$par[isigma])
  coef_out[irho] <- tanh(theta_tS$par[irho])
  coef_out[iv] <- 2 + exp(theta_tS$par[iv])
  names(coef_out) <- c(colnames(XS), colnames(XO), "sigma", "rho", "df")

  jac <- rep(1, length(theta_tS$par))
  jac[isigma] <- coef_out[isigma]
  jac[irho] <- 1 - coef_out[irho]^2
  jac[iv] <- coef_out[iv] - 2
  vcov_nat <- diag(jac, nrow = length(jac)) %*% vcov_link %*% diag(jac, nrow = length(jac))
  se_nat <- .ssmodels_safe_se(vcov_nat)
  names(se_nat) <- names(coef_out)

  nObs <- length(YS)
  nParam <- length(start_opt)
  result <- list(
    coefficients    = coef_out,
    coefficients_link = theta_tS$par,
    value           = theta_tS$value,
    loglik          = theta_tS$value,
    counts          = theta_tS$counts[1],
    hessian         = theta_tS$hessian,
    fisher_infotS   = vcov_nat,
    prop_sigmatS    = se_nat,
    level           = YSLevels,
    nObs            = nObs,
    nParam          = nParam,
    N0              = sum(YS == 0L),
    N1              = sum(YS == 1L),
    NXS             = NXS,
    NXO             = NXO,
    df              = nObs - nParam,
    aic             = -2 * theta_tS$value + 2 * nParam,
    bic             = -2 * theta_tS$value + nParam * log(nObs),
    initial.value   = start_opt
  )
  class(result) <- c("HeckmantS", class(theta_tS))
  result
}
