# ============================================================
# Validação do gradiente do modelo Heckman-BS via numDeriv
# Banco: MEPS2001
# Pacote: ssmodels
# ============================================================

rm(list = ls())

pkgs <- c("pkgload", "numDeriv", "dplyr", "tibble", "readr")
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("Não foi possível localizar o banco MEPS2001.")

dados <- as.data.frame(MEPS2001)

# ------------------------------------------------------------
# 3. Fórmulas usadas no pacote
# ------------------------------------------------------------
selectEq <- dambexp ~ age + female + educ + blhisp + totchr + ins + income
outcomeBS <- ambexp ~ age + female + educ + blhisp + totchr + ins

# ------------------------------------------------------------
# 4. Carregar funções independentes
# ------------------------------------------------------------
local_candidates <- c(
  if (!is.na(pkg_root)) file.path(pkg_root, "inst", "validation", "HeckmanBS_loglik_score_independente.R"),
  "HeckmanBS_loglik_score_independente.R",
  "inst/validation/HeckmanBS_loglik_score_independente.R",
  "../inst/validation/HeckmanBS_loglik_score_independente.R",
  "../HeckmanBS_loglik_score_independente.R"
)
source_file <- local_candidates[file.exists(local_candidates)][1]
if (is.na(source_file)) stop("Arquivo HeckmanBS_loglik_score_independente.R não encontrado.")
source(source_file)

components_bs <- heckmanBS_prepare(selectEq, outcomeBS, dados)

# ------------------------------------------------------------
# 5. Ajustar HeckmanBS()
# ------------------------------------------------------------
fit_bs <- ssmodels::HeckmanBS(
  selection = selectEq,
  outcome   = outcomeBS,
  data      = dados
)

# A função retorna rho na escala natural. Quando disponível, usamos
# coefficients_link, que mantém rho_star na escala da otimização.
theta_hat <- fit_bs$coefficients
theta_raw <- if (!is.null(fit_bs$coefficients_link)) {
  fit_bs$coefficients_link
} else {
  theta_tmp <- theta_hat
  theta_tmp[length(theta_tmp)] <- .bs_rho_star_from_rho(theta_hat[length(theta_hat)])
  theta_tmp
}

param_names <- c(
  paste0("gamma:", colnames(components_bs$XS)),
  paste0("beta:", colnames(components_bs$XO)),
  "phi1",
  "rho_star"
)
names(theta_raw) <- param_names

ll_fit <- fit_bs$loglik
ll_ind <- heckmanBS_loglik_independente(theta_raw, selectEq, outcomeBS, dados, components_bs)

cat("\nLogLik HeckmanBS():          ", ll_fit, "\n")
cat("LogLik função independente: ", ll_ind, "\n\n")

# ------------------------------------------------------------
# 6. Funções para numDeriv
# ------------------------------------------------------------
fn_bs <- function(par) {
  heckmanBS_loglik_independente(par, selectEq, outcomeBS, dados, components_bs)
}

gr_bs <- function(par) {
  heckmanBS_score_independente(par, selectEq, outcomeBS, dados, components_bs)
}

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

  g_ana <- gr_bs(theta)
  g_num <- numDeriv::grad(fn_bs, 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_raw, "theta_hat")
cat("Resumo da validação em theta_hat:\n")
print(res_theta_hat, n = Inf)
cat("\nMáximo erro absoluto:", max(res_theta_hat$erro_abs, na.rm = TRUE), "\n")
cat("Máximo erro relativo :", max(res_theta_hat$erro_rel, na.rm = TRUE), "\n")

if (!is.null(fit_bs$gradient_link)) {
  diff_pkg <- as.numeric(fit_bs$gradient_link) - as.numeric(gr_bs(theta_raw))
  cat("Máxima diferença entre gradient_link e função independente:",
      max(abs(diff_pkg), na.rm = TRUE), "\n")
}

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

  # manter phi1 no domínio usado pela função do pacote
  cand[length(cand) - 1L] <- min(max(cand[length(cand) - 1L], 1e-2), 50)

  # manter rho longe da fronteira
  rho_cand <- .bs_inv_logit_2(cand[length(cand)])
  if (!is.finite(rho_cand) || abs(rho_cand) > 0.95) next

  ll_cand <- fn_bs(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)

# ------------------------------------------------------------
# 8. Validação da densidade/probabilidade individual
# ------------------------------------------------------------
dens_bs <- heckmanBS_density_components(theta_raw, selectEq, outcomeBS, dados, components_bs)
dens_bs_50 <- head(dens_bs, 50)

# ------------------------------------------------------------
# 9. Salvar resultados
# ------------------------------------------------------------
readr::write_csv(res_detalhado, "validacao_gradiente_HeckmanBS_detalhada.csv")
readr::write_csv(res_resumo, "validacao_gradiente_HeckmanBS_resumo.csv")
readr::write_csv(dens_bs_50, "validacao_densidade_HeckmanBS_primeiras_50.csv")
saveRDS(
  list(
    fit = fit_bs,
    theta_raw = theta_raw,
    logLik_fit = ll_fit,
    logLik_independente = ll_ind,
    validacao_detalhada = res_detalhado,
    validacao_resumo = res_resumo,
    densidade_primeiras_50 = dens_bs_50
  ),
  "validacao_gradiente_HeckmanBS.rds"
)

cat("\nArquivos gerados:\n")
cat("\n- validacao_gradiente_HeckmanBS_detalhada.csv")
cat("\n- validacao_gradiente_HeckmanBS_resumo.csv")
cat("\n- validacao_densidade_HeckmanBS_primeiras_50.csv")
cat("\n- validacao_gradiente_HeckmanBS.rds\n")
