# ============================================================
# Validacao da log-verossimilhanca e gradiente do Heckman-t
# Modelo de Marchenko & Genton (2012), usando MEPS2001.
# ============================================================

rm(list = ls())

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

library(dplyr)
library(tibble)

# ------------------------------------------------------------
# Localizar raiz do pacote
# ------------------------------------------------------------
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
  }
  stop("Nao encontrei a raiz do pacote. Execute dentro da pasta ssmodels ou em uma subpasta dela.")
}

pkg_root <- find_pkg_root()
message("Raiz do pacote: ", pkg_root)

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

# Carregar funcoes independentes
validation_file <- file.path(pkg_root, "inst", "validation", "HeckmantS_loglik_score_independente.R")
if (!file.exists(validation_file)) {
  validation_file <- file.path(getwd(), "HeckmantS_loglik_score_independente.R")
}
if (!file.exists(validation_file)) {
  stop("Nao encontrei HeckmantS_loglik_score_independente.R.")
}
source(validation_file)

# ------------------------------------------------------------
# Dados MEPS2001
# ------------------------------------------------------------
if (file.exists(file.path(pkg_root, "data", "MEPS2001.rda"))) {
  load(file.path(pkg_root, "data", "MEPS2001.rda"))
} else {
  data("MEPS2001", package = "ssmodels")
}

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

# ------------------------------------------------------------
# Ajuste do modelo do pacote
# ------------------------------------------------------------
fit_t <- ssmodels::HeckmantS(
  selection = selectEq,
  outcome   = outcomeEq,
  data      = dados,
  df        = 12
)

print(summary(fit_t))

comp <- .ss_get_tS_components(selectEq, outcomeEq, dados)

# Obter theta_hat na escala livre
if (!is.null(fit_t$coefficients_link)) {
  theta_hat_link <- fit_t$coefficients_link
} else {
  theta_hat_link <- heckmant_natural_to_link(fit_t$coefficients, comp)
}

theta_hat_nat <- heckmant_link_to_natural(theta_hat_link, comp)

ll_pkg <- fit_t$loglik
ll_ind <- heckmant_loglik_link(theta_hat_link, comp)

cat("\nLogLik HeckmantS():          ", ll_pkg, "\n")
cat("LogLik funcao independente: ", ll_ind, "\n")
cat("Diferenca:                  ", ll_ind - ll_pkg, "\n")

# ------------------------------------------------------------
# Validacao do gradiente em theta_hat
# ------------------------------------------------------------
loglik_fun <- function(th) heckmant_loglik_link(th, comp)
score_fun  <- function(th) heckmant_score_link(th, comp)

safe_num_grad <- function(th) {
  numDeriv::grad(loglik_fun, th, method = "Richardson")
}

compare_grad <- function(th, label) {
  g_ana <- score_fun(th)
  g_num <- safe_num_grad(th)
  tibble(
    ponto = label,
    parametro = names(th),
    analitico = as.numeric(g_ana),
    numerico = as.numeric(g_num),
    diferenca = analitico - numerico,
    erro_abs = abs(diferenca),
    erro_rel = abs(diferenca) / pmax(1, abs(numerico))
  )
}

res_hat <- compare_grad(theta_hat_link, "theta_hat")

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

# ------------------------------------------------------------
# Perturbacoes seguras
# ------------------------------------------------------------
set.seed(20260530)
perturbacoes <- list()
tries <- 0
while (length(perturbacoes) < 5 && tries < 200) {
  tries <- tries + 1
  th <- theta_hat_link + stats::rnorm(length(theta_hat_link), sd = 0.02)
  th_nat <- heckmant_link_to_natural(th, comp)
  ll <- loglik_fun(th)

  if (is.finite(ll) && th_nat["sigma"] > 0.05 && abs(th_nat["rho"]) < 0.95 &&
      th_nat["df"] > 2.05 && th_nat["df"] < 80) {
    g1 <- try(score_fun(th), silent = TRUE)
    g2 <- try(safe_num_grad(th), silent = TRUE)
    if (!inherits(g1, "try-error") && !inherits(g2, "try-error") &&
        all(is.finite(g1)) && all(is.finite(g2))) {
      perturbacoes[[length(perturbacoes) + 1L]] <- th
    }
  }
}

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

res_all <- dplyr::bind_rows(res_hat, res_pert)

resumo <- 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(resumo)

# ------------------------------------------------------------
# Checagem das contribuicoes de densidade/log-verossimilhanca
# ------------------------------------------------------------
dens <- heckmant_density_components(theta_hat_nat, comp)
cat("\nChecagem da soma das contribuicoes de loglik observado:\n")
cat("Soma por density_components:", sum(dens$loglik_observed), "\n")
cat("Soma por loglik independente:", ll_ind, "\n")

# ------------------------------------------------------------
# Exportar resultados
# ------------------------------------------------------------
readr::write_csv(res_all, "validacao_gradiente_HeckmantS_detalhada.csv")
readr::write_csv(resumo, "validacao_gradiente_HeckmantS_resumo.csv")
readr::write_csv(head(dens, 50), "validacao_densidade_HeckmantS_primeiras_50.csv")

saveRDS(
  list(
    fit = fit_t,
    theta_hat_link = theta_hat_link,
    theta_hat_natural = theta_hat_nat,
    logLik_package = ll_pkg,
    logLik_independent = ll_ind,
    gradiente_detalhado = res_all,
    resumo = resumo,
    densidade_primeiras_50 = head(dens, 50)
  ),
  "validacao_gradiente_HeckmantS.rds"
)

cat("\nArquivos gerados:\n")
cat("\n- validacao_gradiente_HeckmantS_detalhada.csv")
cat("\n- validacao_gradiente_HeckmantS_resumo.csv")
cat("\n- validacao_densidade_HeckmantS_primeiras_50.csv")
cat("\n- validacao_gradiente_HeckmantS.rds\n")
