# ============================================================
# Validação do gradiente analítico da função HeckmanSK()
# Modelo: Selection skew-normal model de Ogundimu e Hutton
# Estratégia: comparar gradiente analítico corrigido com numDeriv::grad()
# Banco: MEPS2001
# ============================================================

rm(list = ls())

# ------------------------------------------------------------
# 1. Pacotes necessários
# ------------------------------------------------------------
# pkgs <- c("pkgload", "numDeriv", "sn", "dplyr", "tibble")
#
# for (p in pkgs) {
#   if (!requireNamespace(p, quietly = TRUE)) {
#     install.packages(p)
#   }
# }

library(pkgload)
library(numDeriv)
library(sn)
library(dplyr)
library(tibble)

# ------------------------------------------------------------
# 2. Carregar pacote local
#    O script funciona tanto na raiz do pacote quanto em ssmodels/Teste.
# ------------------------------------------------------------
if (file.exists("DESCRIPTION")) {
  pkg_root <- "."
} else if (file.exists(file.path("..", "DESCRIPTION"))) {
  pkg_root <- ".."
} else if (file.exists(file.path("..", "..", "DESCRIPTION"))) {
  pkg_root <- file.path("..", "..")
} else {
  stop("Execute este script na raiz do pacote ssmodels, dentro de ssmodels/Teste ou em inst/validation.")
}

pkgload::load_all(pkg_root, reset = TRUE, quiet = TRUE)

# ------------------------------------------------------------
# 3. Carregar MEPS2001
# ------------------------------------------------------------
if (exists("MEPS2001")) rm(MEPS2001)

data("MEPS2001", package = "ssmodels", envir = environment())
dados <- as.data.frame(MEPS2001)

# Padronização para comparação e estimação.
dados <- dados |>
  mutate(
    dambexp_num = as.integer(dambexp == 1),
    lnambx_obs  = ifelse(dambexp_num == 1, lnambx, NA_real_)
  )

selectEq <- dambexp_num ~ age + female + educ + blhisp + totchr + ins + income
outcomeEq <- lnambx_obs ~ age + female + educ + blhisp + totchr + ins

# ------------------------------------------------------------
# 4. Funções independentes de log-verossimilhança e gradiente
# ------------------------------------------------------------
.sk_components <- function(selection, outcome, data) {
  comp <- ssmodels::extract_model_components(selection = selection,
                                             outcome = outcome,
                                             data = data)
  comp$YO[is.na(comp$YO)] <- 0
  comp
}

.sk_loglik <- function(theta, comp) {
  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO <- comp$YO
  NXS <- ncol(XS)
  NXO <- ncol(XO)

  istartS <- seq_len(NXS)
  istartO <- seq(from = NXS + 1L, length.out = NXO)
  isigma  <- NXS + NXO + 1L
  irho    <- NXS + NXO + 2L
  ilamb1  <- NXS + NXO + 3L

  gamma  <- theta[istartS]
  beta   <- theta[istartO]
  sigma  <- theta[isigma]
  rho    <- theta[irho]
  lambda <- theta[ilamb1]

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

  eta <- as.numeric(XS %*% gamma)
  mu  <- as.numeric(XO %*% beta)
  z   <- (YO - mu) / sigma
  r   <- 1 - rho^2
  u   <- 1 + lambda^2 - (lambda * rho)^2
  q   <- lambda * rho / sqrt(u)

  omega <- (eta + rho * z) / sqrt(r)

  ll_selected <-
    log(2) - log(sigma) +
    stats::dnorm(z, log = TRUE) +
    stats::pnorm(lambda * z, log.p = TRUE) +
    stats::pnorm(omega, log.p = TRUE)

  # P(S = 0 | w) = F_SN(-eta; q)
  F0 <- sn::psn(-eta, xi = 0, omega = 1, alpha = q)
  F0 <- pmax(F0, .Machine$double.xmin)
  ll_censored <- log(F0)

  sum(ifelse(YS == 1, ll_selected, ll_censored))
}

.sk_grad <- function(theta, comp) {
  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO <- comp$YO
  NXS <- ncol(XS)
  NXO <- ncol(XO)
  nObs <- length(YS)
  nParam <- NXS + NXO + 3L

  istartS <- seq_len(NXS)
  istartO <- seq(from = NXS + 1L, length.out = NXO)
  isigma  <- NXS + NXO + 1L
  irho    <- NXS + NXO + 2L
  ilamb1  <- NXS + NXO + 3L

  gamma  <- theta[istartS]
  beta   <- theta[istartO]
  sigma  <- theta[isigma]
  rho    <- theta[irho]
  lambda <- theta[ilamb1]

  if (!is.finite(sigma) || sigma <= 0) return(rep(NA_real_, nParam))
  if (!is.finite(rho) || abs(rho) >= 1) return(rep(NA_real_, nParam))
  if (!is.finite(lambda)) return(rep(NA_real_, nParam))

  sel0 <- YS == 0
  sel1 <- YS == 1

  XS0 <- XS[sel0, , drop = FALSE]
  XS1 <- XS[sel1, , drop = FALSE]
  XO1 <- XO[sel1, , drop = FALSE]
  YO1 <- YO[sel1]

  eta0 <- as.numeric(XS0 %*% gamma)
  eta1 <- as.numeric(XS1 %*% gamma)
  mu1  <- as.numeric(XO1 %*% beta)

  z <- (YO1 - mu1) / sigma
  r <- 1 - rho^2
  u <- 1 + lambda^2 - (lambda * rho)^2
  q <- lambda * rho / sqrt(u)

  omega <- (eta1 + rho * z) / sqrt(r)
  a     <- lambda * z

  K1 <- exp(stats::dnorm(omega, log = TRUE) -
              stats::pnorm(omega, log.p = TRUE))
  K3 <- exp(stats::dnorm(a, log = TRUE) -
              stats::pnorm(a, log.p = TRUE))

  F0 <- sn::psn(-eta0, xi = 0, omega = 1, alpha = q)
  F0 <- pmax(F0, .Machine$double.xmin)

  K2 <- exp(stats::dnorm(eta0, log = TRUE) +
              stats::pnorm(-q * eta0, log.p = TRUE) -
              log(F0))

  K4 <- exp(stats::dnorm(eta0 * sqrt(1 + lambda^2) / sqrt(u), log = TRUE) -
              log(F0))

  G <- matrix(0, nrow = nObs, ncol = nParam)

  if (nrow(XS1) > 0L) {
    G[sel1, istartS] <- XS1 * as.numeric(K1 / sqrt(r))
  }
  if (nrow(XS0) > 0L) {
    G[sel0, istartS] <- XS0 * as.numeric(-2 * K2)
  }

  if (nrow(XO1) > 0L) {
    G[sel1, istartO] <- XO1 * as.numeric(
      z / sigma -
      lambda * K3 / sigma -
      rho * K1 / (sigma * sqrt(r))
    )
  }

  if (sum(sel1) > 0L) {
    G[sel1, isigma] <-
      -1 / sigma +
      z^2 / sigma -
      lambda * z * K3 / sigma -
      rho * z * K1 / (sigma * sqrt(r))

    G[sel1, irho] <-
      K1 * (z + rho * eta1) / (r^(3 / 2))

    G[sel1, ilamb1] <-
      K3 * z
  }

  if (sum(sel0) > 0L) {
    G[sel0, irho] <-
      -2 * lambda * K4 / sqrt(2 * pi * u)

    G[sel0, ilamb1] <-
      -2 * rho * K4 / ((1 + lambda^2) * sqrt(2 * pi * u))
  }

  colSums(G)
}

.validar_um_ponto <- function(theta, comp, nome = "theta") {
  g_ana <- .sk_grad(theta, comp)
  g_num <- numDeriv::grad(func = .sk_loglik, x = theta, comp = comp,
                          method = "Richardson")

  tibble(
    ponto = nome,
    parametro = names(theta),
    analitico = as.numeric(g_ana),
    numerico = as.numeric(g_num),
    diferenca = analitico - numerico,
    erro_abs = abs(diferenca),
    erro_rel = abs(diferenca) / (1 + abs(numerico))
  )
}

# ------------------------------------------------------------
# 5. Ajuste do modelo e validação no máximo encontrado
# ------------------------------------------------------------
fit_sk <- ssmodels::HeckmanSK(selection = selectEq,
                              outcome   = outcomeEq,
                              data      = dados)

comp <- .sk_components(selectEq, outcomeEq, dados)
theta_hat <- fit_sk$coefficients

cat("\nLogLik HeckmanSK(): ", fit_sk$loglik, "\n")
cat("LogLik função independente: ", .sk_loglik(theta_hat, comp), "\n")

valid_hat <- .validar_um_ponto(theta_hat, comp, nome = "theta_hat")

cat("\nResumo da validação em theta_hat:\n")
print(valid_hat, n = Inf)
cat("\nMáximo erro absoluto:", max(valid_hat$erro_abs, na.rm = TRUE), "\n")
cat("Máximo erro relativo :", max(valid_hat$erro_rel, na.rm = TRUE), "\n")

# ------------------------------------------------------------
# 6. Validação em pontos perturbados, para evitar validação apenas no ótimo
# ------------------------------------------------------------
set.seed(123)
B <- 5L
NXS <- comp$NXS
NXO <- comp$NXO
isigma <- NXS + NXO + 1L
irho   <- NXS + NXO + 2L
ilamb  <- NXS + NXO + 3L

valid_list <- vector("list", B + 1L)
valid_list[[1L]] <- valid_hat

for (b in seq_len(B)) {
  theta_b <- theta_hat + stats::rnorm(length(theta_hat), mean = 0, sd = 0.03)
  theta_b[isigma] <- max(0.2, theta_b[isigma])
  theta_b[irho]   <- max(-0.85, min(0.85, theta_b[irho]))
  theta_b[ilamb]  <- max(-4, min(4, theta_b[ilamb]))
  names(theta_b)  <- names(theta_hat)

  valid_list[[b + 1L]] <- .validar_um_ponto(theta_b, comp, nome = paste0("perturbado_", b))
}

valid_all <- dplyr::bind_rows(valid_list)

resumo <- valid_all |>
  group_by(ponto) |>
  summarise(
    max_erro_abs = max(erro_abs, na.rm = TRUE),
    max_erro_rel = max(erro_rel, na.rm = TRUE),
    .groups = "drop"
  )

cat("\nResumo geral:\n")
print(resumo, n = Inf)

write.csv(valid_all, "validacao_gradiente_HeckmanSK_detalhada.csv", row.names = FALSE)
write.csv(resumo, "validacao_gradiente_HeckmanSK_resumo.csv", row.names = FALSE)

saveRDS(
  list(
    fit_sk = fit_sk,
    theta_hat = theta_hat,
    validacao_detalhada = valid_all,
    validacao_resumo = resumo
  ),
  "validacao_gradiente_HeckmanSK.rds"
)

cat("\nArquivos gerados:\n")
cat("\n- validacao_gradiente_HeckmanSK_detalhada.csv")
cat("\n- validacao_gradiente_HeckmanSK_resumo.csv")
cat("\n- validacao_gradiente_HeckmanSK.rds\n")
