# ============================================================
# Validação do gradiente do HeckmanGe/HeckmanGE com numDeriv
# Banco: MEPS2001
# ============================================================

rm(list = ls())

find_pkg_root <- function() {
  candidates <- c(".", "..", "../..", "../../..")
  for (cc in candidates) {
    if (file.exists(file.path(cc, "DESCRIPTION")) && dir.exists(file.path(cc, "R"))) {
      return(normalizePath(cc, winslash = "/"))
    }
  }
  stop("Não encontrei a raiz do pacote. Execute o script na raiz do pacote, em Teste, ou em inst/validation.")
}

pkg_root <- find_pkg_root()

pkgs <- c("pkgload", "numDeriv", "dplyr", "tibble")
for (p in pkgs) {
  if (!requireNamespace(p, quietly = TRUE)) {
    install.packages(p)
  }
}

suppressPackageStartupMessages({
  library(dplyr)
  library(tibble)
})

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

source_file <- file.path(pkg_root, "inst", "validation", "HeckmanGE_loglik_score_independente.R")
if (!file.exists(source_file)) {
  source_file <- "HeckmanGE_loglik_score_independente.R"
}
source(source_file)

# ------------------------------------------------------------
# Carrega MEPS2001
# ------------------------------------------------------------
ok_data <- tryCatch({
  data("MEPS2001", package = "ssmodels", envir = environment())
  exists("MEPS2001")
}, error = function(e) FALSE)

if (!ok_data) {
  candidate_data <- c(
    file.path(pkg_root, "data", "MEPS2001.rda"),
    file.path(pkg_root, "data", "MEPS2001.RData"),
    file.path("data", "MEPS2001.rda"),
    file.path("..", "data", "MEPS2001.rda")
  )
  candidate_data <- candidate_data[file.exists(candidate_data)]
  if (length(candidate_data) == 0L) stop("Não encontrei o banco MEPS2001.")
  load(candidate_data[1])
}

dados <- as.data.frame(MEPS2001)

dados <- dados |>
  mutate(
    dambexp_num = as.integer(dambexp == 1),
    lnambx_obs = ifelse(dambexp_num == 1, lnambx, NA_real_)
  )

# ------------------------------------------------------------
# Fórmulas: exemplo recomendado pelo pacote heckmanGE
# ------------------------------------------------------------
selectEq <- dambexp_num ~ age + female + educ + blhisp + totchr + ins + income
outcomeEq <- lnambx_obs ~ age + female + educ + blhisp + totchr + ins
dispersionEq <- ~ age + female + totchr + ins
correlationEq <- ~ age

# ------------------------------------------------------------
# Ajuste pelo ssmodels::HeckmanGe
# ------------------------------------------------------------
fit_ge <- ssmodels::HeckmanGe(
  selection = selectEq,
  outcome = outcomeEq,
  outcomeS = dispersionEq,
  outcomeC = correlationEq,
  data = dados
)

summary(fit_ge)

# ------------------------------------------------------------
# Componentes independentes
# ------------------------------------------------------------
comp <- heckmanGE_make_components(
  selection = selectEq,
  outcome = outcomeEq,
  dispersion = dispersionEq,
  correlation = correlationEq,
  data = dados
)

idx <- heckmanGE_parameter_index(comp$XS, comp$XO, comp$Msigma, comp$Mrho)

theta_hat <- if (!is.null(fit_ge$coefficients_link)) {
  as.numeric(fit_ge$coefficients_link)
} else {
  as.numeric(fit_ge$coefficients)
}

names(theta_hat) <- idx$parameter_names

ll_fit <- fit_ge$loglik
ll_ind <- heckmanGE_loglik_indep(
  theta = theta_hat,
  YS = comp$YS,
  XS = comp$XS,
  YO = comp$YO,
  XO = comp$XO,
  Msigma = comp$Msigma,
  Mrho = comp$Mrho,
  w = comp$w
)

cat("\nLogLik HeckmanGe(): ", ll_fit, "\n")
cat("LogLik função independente: ", ll_ind, "\n")
cat("Diferença: ", ll_ind - ll_fit, "\n")

# ------------------------------------------------------------
# Validação no estimador
# ------------------------------------------------------------
res_hat <- heckmanGE_validate_gradient(
  theta = theta_hat,
  YS = comp$YS,
  XS = comp$XS,
  YO = comp$YO,
  XO = comp$XO,
  Msigma = comp$Msigma,
  Mrho = comp$Mrho,
  w = comp$w,
  ponto = "theta_hat"
)

cat("\nResumo da validação em theta_hat:\n")
print(tibble::as_tibble(res_hat), n = Inf)

cat("\nMáximo erro absoluto em theta_hat: ", max(res_hat$erro_abs, na.rm = TRUE), "\n")
cat("Máximo erro relativo em theta_hat : ", max(res_hat$erro_rel, na.rm = TRUE), "\n")

# ------------------------------------------------------------
# Perturbações seguras
# ------------------------------------------------------------
set.seed(123)
res_list <- list(theta_hat = res_hat)
accepted <- 0L
attempt <- 0L

loglik0 <- function(par) {
  heckmanGE_loglik_indep(
    theta = par,
    YS = comp$YS,
    XS = comp$XS,
    YO = comp$YO,
    XO = comp$XO,
    Msigma = comp$Msigma,
    Mrho = comp$Mrho,
    w = comp$w
  )
}

while (accepted < 5L && attempt < 100L) {
  attempt <- attempt + 1L
  theta_pert <- theta_hat + stats::rnorm(length(theta_hat), sd = 0.01)
  if (!is.finite(loglik0(theta_pert))) next

  accepted <- accepted + 1L
  res_list[[paste0("perturbado_", accepted)]] <- heckmanGE_validate_gradient(
    theta = theta_pert,
    YS = comp$YS,
    XS = comp$XS,
    YO = comp$YO,
    XO = comp$XO,
    Msigma = comp$Msigma,
    Mrho = comp$Mrho,
    w = comp$w,
    ponto = paste0("perturbado_", accepted)
  )
}

res_all <- dplyr::bind_rows(res_list)
res_summary <- res_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(res_summary, n = Inf)

# ------------------------------------------------------------
# Comparação opcional com heckmanGE::estfun()
# ------------------------------------------------------------
comparacao_heckmanGE <- NULL
if (requireNamespace("heckmanGE", quietly = TRUE)) {
  fit_ref <- tryCatch(
    heckmanGE::heckmanGE(
      selection = selectEq,
      outcome = outcomeEq,
      dispersion = dispersionEq,
      correlation = correlationEq,
      data = dados
    ),
    error = function(e) e
  )

  if (!inherits(fit_ref, "error")) {
    g_ref <- colSums(heckmanGE::estfun(fit_ref))
    comparacao_heckmanGE <- tibble::tibble(
      parametro = names(g_ref),
      gradiente_heckmanGE_estfun = as.numeric(g_ref)
    )
    cat("\nGradiente agregado de heckmanGE::estfun() calculado com sucesso.\n")
  } else {
    cat("\nNão foi possível ajustar heckmanGE::heckmanGE(): ", conditionMessage(fit_ref), "\n")
  }
} else {
  cat("\nPacote heckmanGE não instalado; comparação opcional com estfun() foi ignorada.\n")
}

# ------------------------------------------------------------
# Salva arquivos
# ------------------------------------------------------------
utils::write.csv(res_all, "validacao_gradiente_HeckmanGE_detalhada.csv", row.names = FALSE)
utils::write.csv(res_summary, "validacao_gradiente_HeckmanGE_resumo.csv", row.names = FALSE)

if (!is.null(comparacao_heckmanGE)) {
  utils::write.csv(comparacao_heckmanGE, "comparacao_gradiente_heckmanGE_estfun.csv", row.names = FALSE)
}

saveRDS(
  list(
    fit_ge = fit_ge,
    theta_hat = theta_hat,
    logLik_fit = ll_fit,
    logLik_indep = ll_ind,
    componentes = comp,
    validacao_detalhada = res_all,
    validacao_resumo = res_summary,
    comparacao_heckmanGE = comparacao_heckmanGE
  ),
  "validacao_gradiente_HeckmanGE.rds"
)

cat("\nArquivos gerados:\n")
cat("\n- validacao_gradiente_HeckmanGE_detalhada.csv")
cat("\n- validacao_gradiente_HeckmanGE_resumo.csv")
cat("\n- validacao_gradiente_HeckmanGE.rds")
if (!is.null(comparacao_heckmanGE)) cat("\n- comparacao_gradiente_heckmanGE_estfun.csv")
cat("\n")
