# ============================================================
# Funcoes independentes para validar o modelo Heckman-t
# Marchenko & Genton (2012), Equacao (14) e Apendice A.
# ============================================================

.ss_t_ratio <- function(x, df) {
  exp(stats::dt(x, df = df, log = TRUE) - stats::pt(x, df = df, log.p = TRUE))
}

.ss_get_tS_components <- function(selection, outcome, data) {
  if (exists("extract_model_components", mode = "function")) {
    comp <- extract_model_components(selection = selection, outcome = outcome, data = data)
  } else if (requireNamespace("ssmodels", quietly = TRUE)) {
    comp <- ssmodels::extract_model_components(selection = selection, outcome = outcome, data = data)
  } else {
    stop("A funcao extract_model_components() nao foi encontrada. Carregue o pacote ssmodels.")
  }
  comp$YO_clean <- ifelse(is.na(comp$YO), 0, comp$YO)
  comp
}

# ------------------------------------------------------------
# Transforma parametro em escala livre para escala natural.
# Ordem na escala livre:
#   gamma, beta, log.sigma, atanh.rho, log.df.minus.2
# Ordem natural:
#   gamma, beta, sigma, rho, nu
# ------------------------------------------------------------
heckmant_link_to_natural <- function(theta_link, comp) {
  NXS <- comp$NXS
  NXO <- comp$NXO
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  inu <- irho + 1L

  theta_nat <- theta_link
  theta_nat[isigma] <- exp(theta_link[isigma])
  theta_nat[irho] <- tanh(theta_link[irho])
  theta_nat[inu] <- 2 + exp(theta_link[inu])
  names(theta_nat) <- c(colnames(comp$XS), colnames(comp$XO), "sigma", "rho", "df")
  theta_nat
}

heckmant_natural_to_link <- function(theta_nat, comp) {
  NXS <- comp$NXS
  NXO <- comp$NXO
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  inu <- irho + 1L

  theta_link <- theta_nat
  theta_link[isigma] <- log(theta_nat[isigma])
  theta_link[irho] <- atanh(max(min(theta_nat[irho], 0.999999), -0.999999))
  theta_link[inu] <- log(max(theta_nat[inu] - 2, .Machine$double.eps))
  names(theta_link) <- c(colnames(comp$XS), colnames(comp$XO),
                         "log.sigma", "atanh.rho", "log.df.minus.2")
  theta_link
}

# ------------------------------------------------------------
# Log-verossimilhanca em escala natural.
# Equacao (14) de Marchenko & Genton (2012):
# l_i = u_i[log t_nu(z_i) - log sigma + log T_{nu+1}(eta_i)]
#       + (1-u_i)log T_nu(-w_i'gama)
# ------------------------------------------------------------
heckmant_loglik_natural <- function(theta_nat, comp) {
  NXS <- comp$NXS
  NXO <- comp$NXO
  istartS <- seq_len(NXS)
  istartO <- seq.int(NXS + 1L, NXS + NXO)
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  inu <- irho + 1L

  gamma <- theta_nat[istartS]
  beta <- theta_nat[istartO]
  sigma <- theta_nat[isigma]
  rho <- theta_nat[irho]
  nu <- theta_nat[inu]

  if (any(!is.finite(theta_nat)) || !is.finite(sigma) || sigma <= 0 ||
      !is.finite(rho) || abs(rho) >= 1 || !is.finite(nu) || nu <= 2) {
    return(-Inf)
  }

  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO_clean <- comp$YO_clean

  eta_s <- drop(XS %*% gamma)
  mu_y <- drop(XO %*% beta)
  z <- (YO_clean - mu_y) / sigma
  r <- sqrt(1 - rho^2)
  Q <- sqrt((nu + 1) / (nu + z^2))
  eta <- Q * ((rho * z + eta_s) / r)

  ll0 <- stats::pt(-eta_s, df = nu, log.p = TRUE)
  ll1 <- stats::dt(z, df = nu, log = TRUE) - log(sigma) +
    stats::pt(eta, df = nu + 1, log.p = TRUE)

  val <- sum(ifelse(YS == 0L, ll0, ll1))
  if (!is.finite(val)) return(-Inf)
  val
}

heckmant_loglik_link <- function(theta_link, comp) {
  heckmant_loglik_natural(heckmant_link_to_natural(theta_link, comp), comp)
}

# ------------------------------------------------------------
# Gradiente analitico em escala natural.
# Os componentes beta, gamma, rho e sigma seguem o Apendice A.
# O componente nu e calculado por diferenciacao numerica escalar,
# como recomendado no proprio artigo para os termos envolvendo T_nu.
# ------------------------------------------------------------
heckmant_score_natural <- function(theta_nat, comp) {
  if (!requireNamespace("numDeriv", quietly = TRUE)) {
    stop("Instale o pacote numDeriv para calcular o componente em df.")
  }

  NXS <- comp$NXS
  NXO <- comp$NXO
  istartS <- seq_len(NXS)
  istartO <- seq.int(NXS + 1L, NXS + NXO)
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  inu <- irho + 1L

  gamma <- theta_nat[istartS]
  beta <- theta_nat[istartO]
  sigma <- theta_nat[isigma]
  rho <- theta_nat[irho]
  nu <- theta_nat[inu]

  if (any(!is.finite(theta_nat)) || sigma <= 0 || abs(rho) >= 1 || nu <= 2) {
    return(rep(NA_real_, length(theta_nat)))
  }

  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO <- comp$YO_clean

  idx1 <- YS == 1L
  idx0 <- YS == 0L

  XS1 <- XS[idx1, , drop = FALSE]
  XO1 <- XO[idx1, , drop = FALSE]
  Y1 <- YO[idx1]
  XS0 <- XS[idx0, , drop = FALSE]

  eta1 <- drop(XS1 %*% gamma)
  eta0 <- drop(XS0 %*% gamma)
  mu1 <- drop(XO1 %*% beta)

  z <- (Y1 - mu1) / sigma
  r <- sqrt(1 - rho^2)
  A <- 1 / r
  Arho <- rho / r
  Q <- sqrt((nu + 1) / (nu + z^2))
  zeta <- (rho * z + eta1) / r
  eta <- Q * zeta
  M <- .ss_t_ratio(eta, df = nu + 1)

  # gamma
  score_gamma <- rep(0, NXS)
  if (any(idx1)) {
    score_gamma <- score_gamma + colSums(XS1 * as.numeric(A * Q * M))
  }
  if (any(idx0)) {
    M0 <- .ss_t_ratio(-eta0, df = nu)
    score_gamma <- score_gamma - colSums(XS0 * as.numeric(M0))
  }

  # beta
  score_beta <- rep(0, NXO)
  if (any(idx1)) {
    term_beta <- (Q / sigma) *
      (Q * z + (zeta * z / (nu + z^2) - Arho) * M)
    score_beta <- colSums(XO1 * as.numeric(term_beta))
  }

  # sigma
  score_sigma <- 0
  if (any(idx1)) {
    term_sigma <- (-1 + Q^2 * z^2 +
                     Q * z * (zeta * z / (nu + z^2) - Arho) * M) / sigma
    score_sigma <- sum(term_sigma)
  }

  # rho
  score_rho <- 0
  if (any(idx1)) {
    score_rho <- sum(Q * A^3 * (z + rho * eta1) * M)
  }

  # nu: derivada numerica escalar da log-verossimilhanca em escala natural
  score_nu <- numDeriv::grad(
    func = function(nu_value) {
      th <- theta_nat
      th[inu] <- nu_value
      heckmant_loglik_natural(th, comp)
    },
    x = nu,
    method = "Richardson"
  )

  out <- c(score_gamma, score_beta, score_sigma, score_rho, score_nu)
  names(out) <- c(colnames(comp$XS), colnames(comp$XO), "sigma", "rho", "df")
  out
}

heckmant_score_link <- function(theta_link, comp) {
  theta_nat <- heckmant_link_to_natural(theta_link, comp)
  score_nat <- heckmant_score_natural(theta_nat, comp)

  NXS <- comp$NXS
  NXO <- comp$NXO
  isigma <- NXS + NXO + 1L
  irho <- isigma + 1L
  inu <- irho + 1L

  jac <- rep(1, length(theta_link))
  jac[isigma] <- theta_nat[isigma]       # d sigma / d log.sigma
  jac[irho] <- 1 - theta_nat[irho]^2     # d rho / d atanh.rho
  jac[inu] <- theta_nat[inu] - 2         # d nu / d log(nu - 2)

  out <- score_nat * jac
  names(out) <- names(theta_link)
  out
}

# ------------------------------------------------------------
# Contribuicoes de densidade/likelihood para checagem didatica.
# Retorna densidade condicional f(y | U=1) e contribuicao observada.
# ------------------------------------------------------------
heckmant_density_components <- function(theta_nat, comp) {
  NXS <- comp$NXS
  NXO <- comp$NXO
  gamma <- theta_nat[seq_len(NXS)]
  beta <- theta_nat[seq.int(NXS + 1L, NXS + NXO)]
  sigma <- theta_nat[NXS + NXO + 1L]
  rho <- theta_nat[NXS + NXO + 2L]
  nu <- theta_nat[NXS + NXO + 3L]

  XS <- comp$XS
  XO <- comp$XO
  YS <- comp$YS
  YO <- comp$YO_clean

  eta_s <- drop(XS %*% gamma)
  mu_y <- drop(XO %*% beta)
  z <- (YO - mu_y) / sigma
  r <- sqrt(1 - rho^2)
  Q <- sqrt((nu + 1) / (nu + z^2))
  eta <- Q * ((rho * z + eta_s) / r)

  f_y <- stats::dt(z, df = nu) / sigma
  p_sel_given_y <- stats::pt(eta, df = nu + 1)
  p_sel <- stats::pt(eta_s, df = nu)
  p_not_sel <- stats::pt(-eta_s, df = nu)

  f_cond <- f_y * p_sel_given_y / p_sel
  likelihood_observed <- ifelse(YS == 1L, f_y * p_sel_given_y, p_not_sel)

  data.frame(
    selected = YS,
    z = z,
    eta_selection = eta_s,
    eta_conditional = eta,
    f_y = f_y,
    p_sel_given_y = p_sel_given_y,
    p_sel = p_sel,
    f_cond_y_given_selected = f_cond,
    likelihood_observed = likelihood_observed,
    loglik_observed = log(likelihood_observed)
  )
}
