# ============================================================
# Validação SEGURA 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
#
# Este script foi escrito para evitar falsos alarmes causados por pontos
# perturbados numericamente instáveis. A validação principal é feita em
# theta_hat; as perturbações são pequenas e aceitas apenas quando ficam
# em região numericamente regular da log-verossimilhança.
# ============================================================

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 na raiz do pacote, em ssmodels/Teste ou
#    quando copiado para outro diretório indicando pkg_root manualmente.
# ------------------------------------------------------------
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 dentro de 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)

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)

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

  out <- sum(ifelse(YS == 1, ll_selected, ll_censored))
  if (!is.finite(out)) return(NA_real_)
  out
}

.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)
}

# Diagnóstico simples para rejeitar pontos com alta chance de instabilidade
# numérica antes da chamada a numDeriv::grad().
.sk_is_safe_theta <- function(theta, comp) {
  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO <- comp$YO
  NXS <- ncol(XS)
  NXO <- ncol(XO)

  sigma <- theta[NXS + NXO + 1L]
  rho   <- theta[NXS + NXO + 2L]
  lambda <- theta[NXS + NXO + 3L]

  if (!is.finite(sigma) || sigma <= 0.10) return(FALSE)
  if (!is.finite(rho) || abs(rho) >= 0.90) return(FALSE)
  if (!is.finite(lambda) || abs(lambda) > 5) return(FALSE)

  ll <- .sk_loglik(theta, comp)
  if (!is.finite(ll)) return(FALSE)

  gamma <- theta[seq_len(NXS)]
  beta  <- theta[seq(from = NXS + 1L, length.out = NXO)]
  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)

  # Evita pontos nos quais pequenas perturbações podem levar a probabilidade
  # numericamente nula em alguma das componentes da log-verossimilhança.
  log_p1 <- stats::pnorm(lambda * z[YS == 1], log.p = TRUE)
  log_p2 <- stats::pnorm(omega[YS == 1], log.p = TRUE)
  F0 <- sn::psn(-eta[YS == 0], xi = 0, omega = 1, alpha = q)
  log_p0 <- log(pmax(F0, .Machine$double.xmin))

  min_log_prob <- min(c(log_p1, log_p2, log_p0), na.rm = TRUE)
  is.finite(min_log_prob) && min_log_prob > -650
}

.validar_um_ponto <- function(theta, comp, nome = "theta") {
  if (!.sk_is_safe_theta(theta, comp)) {
    return(tibble(
      ponto = nome,
      parametro = names(theta),
      analitico = NA_real_,
      numerico = NA_real_,
      diferenca = NA_real_,
      erro_abs = NA_real_,
      erro_rel = NA_real_,
      status = "descartado_regiao_instavel"
    ))
  }

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

  status <- if (all(is.finite(g_ana)) && all(is.finite(g_num))) {
    "avaliado"
  } else {
    "gradiente_nao_finito"
  }

  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)),
    status = status
  )
}

# ------------------------------------------------------------
# 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 em theta_hat:", max(valid_hat$erro_abs, na.rm = TRUE), "\n")
cat("Máximo erro relativo em theta_hat :", max(valid_hat$erro_rel, na.rm = TRUE), "\n")

# ------------------------------------------------------------
# 6. Validação em pontos perturbados seguros
# ------------------------------------------------------------
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 <- list(theta_hat = valid_hat)
accepted <- 0L
attempts <- 0L
max_attempts <- 200L
sd_perturb <- 0.005

while (accepted < B && attempts < max_attempts) {
  attempts <- attempts + 1L

  theta_b <- theta_hat + stats::rnorm(length(theta_hat), mean = 0, sd = sd_perturb)
  theta_b[isigma] <- max(0.2, theta_b[isigma])
  theta_b[irho]   <- max(-0.80, min(0.80, theta_b[irho]))
  theta_b[ilamb]  <- max(-4, min(4, theta_b[ilamb]))
  names(theta_b)  <- names(theta_hat)

  if (!.sk_is_safe_theta(theta_b, comp)) next

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

if (accepted < B) {
  warning("Foram aceitos apenas ", accepted, " pontos perturbados seguros em ", attempts, " tentativas.")
}

valid_all <- dplyr::bind_rows(valid_list)

resumo <- valid_all |>
  group_by(ponto, status) |>
  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)
cat("\nPontos perturbados seguros aceitos:", accepted, "de", attempts, "tentativas.\n")

# 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,
#     pontos_perturbados_aceitos = accepted,
#     tentativas = attempts
#   ),
#   "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")
