# ============================================================
# Funções independentes para validar a log-verossimilhança
# e o gradiente do modelo Heckman Generalizado.
#
# Parametrização:
#   sigma_i = exp(E_i %*% lambda)
#   rho_i   = tanh(V_i %*% kappa)
#
# A ordem dos parâmetros é:
#   theta = c(gamma, beta, lambda, kappa)
#
# Fórmula-base:
#   l_i(theta) = u_i{log Phi(zeta_i) + log phi(z_i) - log sigma_i}
#              + (1-u_i) log Phi(-mu2_i),
#
# em que
#   z_i     = (y_i - mu1_i)/sigma_i,
#   zeta_i  = (mu2_i + rho_i z_i)/sqrt(1-rho_i^2).
# ============================================================

.heckmanGE_mills <- function(x) {
  exp(stats::dnorm(x, log = TRUE) - stats::pnorm(x, log.p = TRUE))
}

.heckmanGE_sech <- function(x) 1 / cosh(x)

heckmanGE_make_components <- function(selection,
                                      outcome,
                                      dispersion,
                                      correlation,
                                      data,
                                      weights = NULL) {
  yS_name <- as.character(selection[[2L]])

  vars_selection  <- all.vars(selection)
  vars_outcome    <- all.vars(outcome)
  vars_dispersion <- all.vars(dispersion)
  vars_correlation <- all.vars(correlation)

  complete_selection <- stats::complete.cases(data[, vars_selection, drop = FALSE])
  complete_correlation <- stats::complete.cases(data[, vars_correlation, drop = FALSE])
  complete_outdisp <- stats::complete.cases(
    data[, unique(c(vars_outcome, vars_dispersion)), drop = FALSE]
  )

  ys_raw <- data[[yS_name]]
  if (is.factor(ys_raw)) {
    if (nlevels(ys_raw) != 2L) stop("A variável de seleção deve ter exatamente dois níveis.")
    ys_num <- as.integer(ys_raw == levels(ys_raw)[2L])
  } else {
    ys_num <- as.integer(ys_raw)
  }

  keep <- complete_selection & complete_correlation &
    ((ys_num == 0L) | ((ys_num == 1L) & complete_outdisp))

  data2 <- data[keep, , drop = FALSE]

  mfS <- stats::model.frame(selection, data = data2, na.action = stats::na.pass)
  YS_raw <- stats::model.response(mfS)
  if (is.factor(YS_raw)) {
    YS <- as.integer(YS_raw == levels(YS_raw)[2L])
  } else {
    YS <- as.integer(YS_raw)
  }
  XS <- stats::model.matrix(stats::terms(mfS), mfS)

  mfO <- stats::model.frame(outcome, data = data2, na.action = stats::na.pass)
  YO <- stats::model.response(mfO)
  XO <- stats::model.matrix(attr(mfO, "terms"), mfO)

  mfD <- stats::model.frame(dispersion, data = data2, na.action = stats::na.pass)
  Msigma <- stats::model.matrix(attr(mfD, "terms"), mfD)

  mfC <- stats::model.frame(correlation, data = data2, na.action = stats::na.pass)
  Mrho <- stats::model.matrix(attr(mfC, "terms"), mfC)

  if (is.null(weights)) {
    w <- rep(1, length(YS))
  } else {
    w_all <- eval(substitute(weights), data2, parent.frame())
    w <- as.numeric(w_all)
    if (length(w) != length(YS)) stop("O vetor de pesos possui comprimento incompatível.")
  }

  list(
    data = data2,
    YS = YS,
    YO = YO,
    XS = XS,
    XO = XO,
    Msigma = Msigma,
    Mrho = Mrho,
    w = w,
    index = heckmanGE_parameter_index(XS, XO, Msigma, Mrho)
  )
}

heckmanGE_parameter_index <- function(XS, XO, Msigma, Mrho) {
  NXS <- ncol(XS)
  NXO <- ncol(XO)
  NE  <- ncol(Msigma)
  NV  <- ncol(Mrho)

  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)

  list(
    NXS = NXS,
    NXO = NXO,
    NE = NE,
    NV = NV,
    istartS = istartS,
    istartO = istartO,
    ilambda = ilambda,
    ikappa = ikappa,
    npar = NXS + NXO + NE + NV,
    parameter_names = c(
      paste0("selection:", colnames(XS)),
      paste0("outcome:", colnames(XO)),
      paste0("dispersion:", colnames(Msigma)),
      paste0("correlation:", colnames(Mrho))
    )
  )
}

heckmanGE_loglik_indep <- function(theta,
                                   YS,
                                   XS,
                                   YO,
                                   XO,
                                   Msigma,
                                   Mrho,
                                   w = NULL) {
  idx <- heckmanGE_parameter_index(XS, XO, Msigma, Mrho)
  if (length(theta) != idx$npar) stop("Comprimento incorreto de theta.")
  if (is.null(w)) w <- rep(1, length(YS))

  gamma  <- theta[idx$istartS]
  beta   <- theta[idx$istartO]
  lambda <- theta[idx$ilambda]
  kappa  <- theta[idx$ikappa]

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

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

  if (any(!is.finite(sigma)) || any(sigma <= 0)) return(NA_real_)
  if (any(!is.finite(rho)) || any(abs(rho) >= 1)) return(NA_real_)

  r <- sqrt(1 - rho^2)
  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(w * ifelse(YS == 0L, ll0, ll1))
}

heckmanGE_score_i_indep <- function(theta,
                                    YS,
                                    XS,
                                    YO,
                                    XO,
                                    Msigma,
                                    Mrho,
                                    w = NULL) {
  idx <- heckmanGE_parameter_index(XS, XO, Msigma, Mrho)
  if (length(theta) != idx$npar) stop("Comprimento incorreto de theta.")
  if (is.null(w)) w <- rep(1, length(YS))

  gamma  <- theta[idx$istartS]
  beta   <- theta[idx$istartO]
  lambda <- theta[idx$ilambda]
  kappa  <- theta[idx$ikappa]

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

  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))
  eta_rho1 <- as.numeric(MR1 %*% kappa)
  rho1 <- tanh(eta_rho1)

  r1 <- sqrt(1 - rho1^2)
  A_rho <- 1 / r1
  A_rrho <- rho1 / r1

  z1 <- (YO1 - mu11) / sigma1
  zeta1 <- mu21 * A_rho + z1 * A_rrho

  Mmu2 <- .heckmanGE_mills(-mu20)
  Mzeta <- .heckmanGE_mills(zeta1)

  sech2 <- .heckmanGE_sech(eta_rho1)^2
  Q_rho <- mu21 * rho1 * A_rho^2 + z1 * (1 + A_rrho^2)

  gradient <- matrix(0, nrow = length(YS), ncol = idx$npar)
  colnames(gradient) <- idx$parameter_names

  gradient[YS == 0L, idx$istartS] <- -XS0 * Mmu2
  gradient[YS == 1L, idx$istartS] <- XS1 * (Mzeta * A_rho)
  gradient[YS == 1L, idx$istartO] <- XO1 *
    ((z1 - Mzeta * A_rrho) / sigma1)
  gradient[YS == 1L, idx$ilambda] <- MS1 *
    (z1^2 - 1 - Mzeta * z1 * A_rrho)
  gradient[YS == 1L, idx$ikappa] <- MR1 *
    (Mzeta * A_rho * sech2 * Q_rho)

  gradient * w
}

heckmanGE_score_indep <- function(theta,
                                  YS,
                                  XS,
                                  YO,
                                  XO,
                                  Msigma,
                                  Mrho,
                                  w = NULL) {
  colSums(
    heckmanGE_score_i_indep(
      theta = theta,
      YS = YS,
      XS = XS,
      YO = YO,
      XO = XO,
      Msigma = Msigma,
      Mrho = Mrho,
      w = w
    )
  )
}

heckmanGE_validate_gradient <- function(theta,
                                        YS,
                                        XS,
                                        YO,
                                        XO,
                                        Msigma,
                                        Mrho,
                                        w = NULL,
                                        ponto = "theta") {
  if (!requireNamespace("numDeriv", quietly = TRUE)) {
    stop("Instale o pacote numDeriv para usar esta função.")
  }

  idx <- heckmanGE_parameter_index(XS, XO, Msigma, Mrho)
  names(theta) <- idx$parameter_names

  loglik_fun <- function(par) {
    heckmanGE_loglik_indep(par, YS, XS, YO, XO, Msigma, Mrho, w)
  }

  g_ana <- heckmanGE_score_indep(theta, YS, XS, YO, XO, Msigma, Mrho, w)
  g_num <- numDeriv::grad(loglik_fun, theta, method = "Richardson")

  out <- data.frame(
    ponto = ponto,
    parametro = idx$parameter_names,
    analitico = as.numeric(g_ana),
    numerico = as.numeric(g_num),
    diferenca = as.numeric(g_ana - g_num),
    erro_abs = abs(as.numeric(g_ana - g_num)),
    stringsAsFactors = FALSE
  )

  out$erro_rel <- out$erro_abs / pmax(1, abs(out$numerico), abs(out$analitico))
  out
}
