# ============================================================
# Validacao do gradiente do modelo Heckman classico via numDeriv
# Banco: MEPS2001
# Pacote: ssmodels
# ============================================================

rm(list = ls())

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

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

# ------------------------------------------------------------
# 1. Carregar pacote local ou instalado
# ------------------------------------------------------------
find_pkg_root <- function(path = getwd()) {
  cur <- normalizePath(path, winslash = "/", mustWork = TRUE)
  for (i in seq_len(8)) {
    if (file.exists(file.path(cur, "DESCRIPTION")) && dir.exists(file.path(cur, "R"))) {
      return(cur)
    }
    parent <- dirname(cur)
    if (identical(parent, cur)) break
    cur <- parent
  }
  NA_character_
}

pkg_root <- find_pkg_root()
if (!is.na(pkg_root)) {
  pkgload::load_all(pkg_root, reset = TRUE, quiet = TRUE)
} else {
  library(ssmodels)
}

# ------------------------------------------------------------
# 2. Carregar banco MEPS2001
# ------------------------------------------------------------
ok_data <- tryCatch({
  data("MEPS2001", package = "ssmodels", envir = environment())
  exists("MEPS2001", inherits = FALSE)
}, error = function(e) FALSE)

if (!ok_data && !is.na(pkg_root)) {
  candidate_data <- file.path(pkg_root, "data", "MEPS2001.rda")
  if (file.exists(candidate_data)) load(candidate_data)
}
if (!exists("MEPS2001")) stop("Nao foi possivel localizar o banco MEPS2001.")

dados <- as.data.frame(MEPS2001)

# ------------------------------------------------------------
# 3. Formulas de validacao
# ------------------------------------------------------------
selectEq <- dambexp ~ age + female + educ + blhisp + totchr + ins + income
outcomeEq <- lnambx ~ age + female + educ + blhisp + totchr + ins

# ------------------------------------------------------------
# 4. Carregar funcoes independentes
# ------------------------------------------------------------
local_candidates <- c(
  if (!is.na(pkg_root)) file.path(pkg_root, "inst", "validation", "HeckmanCL_loglik_score_independente.R"),
  "HeckmanCL_loglik_score_independente.R",
  "inst/validation/HeckmanCL_loglik_score_independente.R",
  "../inst/validation/HeckmanCL_loglik_score_independente.R",
  "../HeckmanCL_loglik_score_independente.R"
)
source_file <- local_candidates[file.exists(local_candidates)][1]
if (is.na(source_file)) stop("Arquivo HeckmanCL_loglik_score_independente.R nao encontrado.")
source(source_file)

components_cl <- heckmanCL_prepare(selectEq, outcomeEq, dados)

# ------------------------------------------------------------
# 5. Ajustar HeckmanCL()
# ------------------------------------------------------------
fit_cl <- ssmodels::HeckmanCL(
  selection = selectEq,
  outcome   = outcomeEq,
  data      = dados
)

theta_hat <- fit_cl$coefficients
param_names <- c(
  paste0("gamma:", colnames(components_cl$XS)),
  paste0("beta:", colnames(components_cl$XO)),
  "sigma",
  "rho"
)
names(theta_hat) <- param_names

ll_fit <- fit_cl$loglik
ll_ind <- heckmanCL_loglik_independente(theta_hat, selectEq, outcomeEq, dados, components_cl)

cat("\nLogLik HeckmanCL():          ", ll_fit, "\n")
cat("LogLik funcao independente: ", ll_ind, "\n\n")

# ------------------------------------------------------------
# 6. Comparacao opcional com sampleSelection
# ------------------------------------------------------------
fit_sample_ml <- tryCatch(
  sampleSelection::selection(
    selection = selectEq,
    outcome   = outcomeEq,
    data      = dados,
    method    = "ml"
  ),
  error = function(e) NULL
)

if (!is.null(fit_sample_ml)) {
  ll_sample <- as.numeric(stats::logLik(fit_sample_ml))
  cat("LogLik sampleSelection::selection(method='ml'): ", ll_sample, "\n")
  cat("Diferenca ssmodels - sampleSelection:          ", ll_fit - ll_sample, "\n\n")
}

# ------------------------------------------------------------
# 7. Funcoes para numDeriv
# ------------------------------------------------------------
fn_cl <- function(par) {
  heckmanCL_loglik_independente(par, selectEq, outcomeEq, dados, components_cl)
}

gr_cl <- function(par) {
  heckmanCL_score_independente(par, selectEq, outcomeEq, dados, components_cl)
}

validar_ponto <- function(theta, nome_ponto) {
  ll <- fn_cl(theta)
  if (!is.finite(ll)) return(NULL)

  g_ana <- gr_cl(theta)
  g_num <- numDeriv::grad(fn_cl, theta, method = "Richardson")

  tibble(
    ponto = nome_ponto,
    parametro = param_names,
    analitico = as.numeric(g_ana),
    numerico = as.numeric(g_num),
    diferenca = analitico - numerico,
    erro_abs = abs(diferenca),
    erro_rel = erro_abs / pmax(1, abs(numerico))
  )
}

res_theta_hat <- validar_ponto(theta_hat, "theta_hat")
cat("Resumo da validacao em theta_hat:\n")
print(res_theta_hat, n = Inf)
cat("\nMaximo erro absoluto:", max(res_theta_hat$erro_abs, na.rm = TRUE), "\n")
cat("Maximo erro relativo :", max(res_theta_hat$erro_rel, na.rm = TRUE), "\n")

# ------------------------------------------------------------
# 8. Perturbacoes seguras
# ------------------------------------------------------------
set.seed(20260530)
perturbacoes <- list()
tentativas <- 0L
while (length(perturbacoes) < 5L && tentativas < 200L) {
  tentativas <- tentativas + 1L
  step <- stats::rnorm(length(theta_hat), mean = 0, sd = 0.01 * pmax(1, abs(theta_hat)))
  cand <- theta_hat + step

  idx_sigma <- length(cand) - 1L
  idx_rho <- length(cand)

  cand[idx_sigma] <- max(cand[idx_sigma], 0.05)
  cand[idx_rho] <- min(max(cand[idx_rho], -0.95), 0.95)

  ll_cand <- fn_cl(cand)
  if (!is.finite(ll_cand)) next

  perturbacoes[[length(perturbacoes) + 1L]] <- cand
}

res_pert <- dplyr::bind_rows(lapply(seq_along(perturbacoes), function(j) {
  validar_ponto(perturbacoes[[j]], paste0("perturbado_", j))
}))

res_detalhado <- dplyr::bind_rows(res_theta_hat, res_pert)
res_resumo <- res_detalhado |>
  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_resumo)

# ------------------------------------------------------------
# 9. Validacao da densidade/probabilidade individual
# ------------------------------------------------------------
dens_cl <- heckmanCL_density_components(theta_hat, selectEq, outcomeEq, dados, components_cl)
dens_cl_50 <- head(dens_cl, 50)

# ------------------------------------------------------------
# 10. Salvar resultados
# ------------------------------------------------------------
readr::write_csv(res_detalhado, "validacao_gradiente_HeckmanCL_detalhada.csv")
readr::write_csv(res_resumo, "validacao_gradiente_HeckmanCL_resumo.csv")
readr::write_csv(dens_cl_50, "validacao_densidade_HeckmanCL_primeiras_50.csv")
saveRDS(
  list(
    fit = fit_cl,
    fit_sampleSelection = fit_sample_ml,
    theta_hat = theta_hat,
    logLik_fit = ll_fit,
    logLik_independente = ll_ind,
    validacao_detalhada = res_detalhado,
    validacao_resumo = res_resumo,
    densidade_primeiras_50 = dens_cl_50
  ),
  "validacao_gradiente_HeckmanCL.rds"
)

cat("\nArquivos gerados:\n")
cat("\n- validacao_gradiente_HeckmanCL_detalhada.csv")
cat("\n- validacao_gradiente_HeckmanCL_resumo.csv")
cat("\n- validacao_densidade_HeckmanCL_primeiras_50.csv")
cat("\n- validacao_gradiente_HeckmanCL.rds\n")
