# ============================================================
# Funcoes independentes para validacao do modelo Heckman classico
# Pacote: ssmodels
# Objetivo: comparar gradiente analitico contra numDeriv::grad()
# ============================================================

.cl_safe_ratio <- function(x) {
  exp(stats::dnorm(x, log = TRUE) - stats::pnorm(x, log.p = TRUE))
}

.cl_extract_components_independente <- function(selection, outcome, data) {
  if (exists("extract_model_components", mode = "function")) {
    return(extract_model_components(selection = selection, outcome = outcome, data = data))
  }

  mf_s <- stats::model.frame(selection, data = data, na.action = stats::na.pass)
  YS_raw <- stats::model.response(mf_s)
  XS <- stats::model.matrix(selection, data = mf_s)

  if (is.factor(YS_raw)) {
    lev <- levels(YS_raw)
    if (length(lev) != 2L) stop("A variavel de selecao deve possuir dois niveis.")
    YS <- as.integer(YS_raw == lev[2L])
  } else {
    vals <- sort(unique(stats::na.omit(YS_raw)))
    if (!all(vals %in% c(0, 1))) {
      stop("A variavel de selecao deve ser binaria, codificada como 0/1, ou fator com dois niveis.")
    }
    lev <- vals
    YS <- as.integer(YS_raw)
  }

  mf_o <- stats::model.frame(outcome, data = data, na.action = stats::na.pass)
  YO <- stats::model.response(mf_o)
  XO <- stats::model.matrix(outcome, data = mf_o)

  keep <- !is.na(YS) &
    stats::complete.cases(XS) &
    stats::complete.cases(XO) &
    (YS == 0L | !is.na(YO))

  XS <- XS[keep, , drop = FALSE]
  XO <- XO[keep, , drop = FALSE]
  YS <- YS[keep]
  YO <- YO[keep]

  list(
    XS = XS,
    YS = YS,
    NXS = ncol(XS),
    XO = XO,
    YO = YO,
    NXO = ncol(XO),
    YSLevels = lev
  )
}

heckmanCL_prepare <- function(selection, outcome, data) {
  .cl_extract_components_independente(selection, outcome, data)
}

heckmanCL_loglik_independente <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanCL_prepare(selection, outcome, data)

  XS <- components$XS
  YS <- components$YS
  NXS <- components$NXS
  XO <- components$XO
  YO <- components$YO
  NXO <- components$NXO

  if (any(!is.finite(par))) return(NA_real_)

  igamma <- seq_len(NXS)
  ibeta  <- seq.int(max(igamma) + 1L, length.out = NXO)
  isigma <- max(ibeta) + 1L
  irho   <- isigma + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  sigma <- par[isigma]
  rho   <- par[irho]

  if (!is.finite(sigma) || sigma <= 0) return(NA_real_)
  if (!is.finite(rho) || abs(rho) >= 1) return(NA_real_)

  r <- sqrt(1 - rho^2)
  if (!is.finite(r) || r <= 1e-8) return(NA_real_)

  XS0 <- XS[YS == 0L, , drop = FALSE]
  XS1 <- XS[YS == 1L, , drop = FALSE]
  XO1 <- XO[YS == 1L, , drop = FALSE]
  YO1 <- YO[YS == 1L]

  eta0 <- as.vector(XS0 %*% gamma)
  eta1 <- as.vector(XS1 %*% gamma)
  mu1  <- as.vector(XO1 %*% beta)
  z1   <- (YO1 - mu1) / sigma

  B1 <- (eta1 + rho * z1) / r

  ll0 <- stats::pnorm(-eta0, log.p = TRUE)
  ll1 <- stats::dnorm(z1, log = TRUE) - log(sigma) + stats::pnorm(B1, log.p = TRUE)

  ll <- sum(ll0) + sum(ll1)
  if (!is.finite(ll)) return(NA_real_)
  ll
}

heckmanCL_score_independente <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanCL_prepare(selection, outcome, data)

  XS <- components$XS
  YS <- components$YS
  NXS <- components$NXS
  XO <- components$XO
  YO <- components$YO
  NXO <- components$NXO

  if (any(!is.finite(par))) return(rep(NA_real_, length(par)))

  igamma <- seq_len(NXS)
  ibeta  <- seq.int(max(igamma) + 1L, length.out = NXO)
  isigma <- max(ibeta) + 1L
  irho   <- isigma + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  sigma <- par[isigma]
  rho   <- par[irho]

  if (!is.finite(sigma) || sigma <= 0) return(rep(NA_real_, length(par)))
  if (!is.finite(rho) || abs(rho) >= 1) return(rep(NA_real_, length(par)))

  r <- sqrt(1 - rho^2)
  if (!is.finite(r) || r <= 1e-8) return(rep(NA_real_, length(par)))

  XS0 <- XS[YS == 0L, , drop = FALSE]
  XS1 <- XS[YS == 1L, , drop = FALSE]
  XO1 <- XO[YS == 1L, , drop = FALSE]
  YO1 <- YO[YS == 1L]

  eta0 <- as.vector(XS0 %*% gamma)
  eta1 <- as.vector(XS1 %*% gamma)
  mu1  <- as.vector(XO1 %*% beta)
  u1   <- YO1 - mu1
  z1   <- u1 / sigma
  B1   <- (eta1 + rho * z1) / r

  lambda0 <- .cl_safe_ratio(-eta0)
  lambdaB <- .cl_safe_ratio(B1)

  grad <- numeric(length(par))

  if (nrow(XS0) > 0L) {
    grad[igamma] <- grad[igamma] - colSums(sweep(XS0, 1, lambda0, `*`))
  }

  if (nrow(XS1) > 0L) {
    grad[igamma] <- grad[igamma] + colSums(sweep(XS1, 1, lambdaB / r, `*`))

    grad[ibeta] <- grad[ibeta] +
      colSums(sweep(XO1, 1, z1 / sigma - rho * lambdaB / (sigma * r), `*`))

    grad[isigma] <- grad[isigma] +
      sum(z1^2 / sigma - 1 / sigma - rho * z1 * lambdaB / (sigma * r))

    grad[irho] <- grad[irho] +
      sum(lambdaB * (z1 + rho * eta1) / r^3)
  }

  names(grad) <- names(par)
  grad
}

heckmanCL_density_components <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanCL_prepare(selection, outcome, data)

  XS <- components$XS
  YS <- components$YS
  NXS <- components$NXS
  XO <- components$XO
  YO <- components$YO
  NXO <- components$NXO

  igamma <- seq_len(NXS)
  ibeta  <- seq.int(max(igamma) + 1L, length.out = NXO)
  isigma <- max(ibeta) + 1L
  irho   <- isigma + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  sigma <- par[isigma]
  rho   <- par[irho]
  r <- sqrt(1 - rho^2)

  out <- data.frame(
    row = seq_along(YS),
    selected = YS,
    value = NA_real_,
    log_contribution = NA_real_,
    type = ifelse(YS == 1L, "selected_density", "censored_probability")
  )

  if (any(YS == 1L)) {
    idx1 <- which(YS == 1L)
    eta1 <- as.vector(XS[idx1, , drop = FALSE] %*% gamma)
    mu1  <- as.vector(XO[idx1, , drop = FALSE] %*% beta)
    z1   <- (YO[idx1] - mu1) / sigma
    B1   <- (eta1 + rho * z1) / r
    logc <- stats::dnorm(z1, log = TRUE) - log(sigma) + stats::pnorm(B1, log.p = TRUE)
    out$value[idx1] <- exp(logc)
    out$log_contribution[idx1] <- logc
  }

  if (any(YS == 0L)) {
    idx0 <- which(YS == 0L)
    eta0 <- as.vector(XS[idx0, , drop = FALSE] %*% gamma)
    logd <- stats::pnorm(-eta0, log.p = TRUE)
    out$value[idx0] <- exp(logd)
    out$log_contribution[idx0] <- logd
  }

  out
}
