# ============================================================
# Funções independentes para validação do modelo Heckman-BS
# Pacote: ssmodels
# Objetivo: comparar gradiente analítico contra numDeriv::grad()
# ============================================================

.bs_inv_logit_2 <- function(x) {
  2 / (1 + exp(-x)) - 1
}

.bs_rho_star_from_rho <- function(rho) {
  rho <- pmin(pmax(rho, -0.999999), 0.999999)
  log((1 + rho) / (1 - rho))
}

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

.bs_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 variável de seleção deve possuir dois níveis.")
    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 variável de seleção deve ser binária, codificada como 0/1, ou fator com dois níveis.")
    }
    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
  )
}

heckmanBS_prepare <- function(selection, outcome, data) {
  comp <- .bs_extract_components_independente(selection, outcome, data)
  if (any(comp$YS == 1L & (!is.finite(comp$YO) | comp$YO <= 0))) {
    stop("HeckmanBS requer valores estritamente positivos da variável resposta nas observações selecionadas.")
  }
  comp
}

heckmanBS_loglik_independente <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanBS_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)
  iphi1  <- max(ibeta) + 1L
  irho   <- iphi1 + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  phi1  <- par[iphi1]
  rho_star <- par[irho]

  if (!is.finite(phi1) || phi1 < 1e-3 || phi1 > 100) return(NA_real_)
  if (!is.finite(rho_star)) return(NA_real_)

  rho <- .bs_inv_logit_2(rho_star)
  phi2 <- 1
  r <- 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]
  YO1 <- YO[YS == 1L]
  XO1 <- XO[YS == 1L, , drop = FALSE]

  eta0 <- as.vector(XS0 %*% gamma)
  eta1 <- as.vector(XS1 %*% gamma)
  mu1  <- as.vector(exp(XO1 %*% beta))
  xi0  <- exp(eta0)
  xi1  <- exp(eta1)

  a1 <- sqrt(YO1 * (phi1 + 1) / (phi1 * mu1)) -
    sqrt((phi1 * mu1) / (YO1 * (phi1 + 1)))

  log_bs_density <-
    (-phi1 / 4) * a1^2 +
    log(
      sqrt((phi1 + 1) / (phi1 * mu1 * YO1)) +
        sqrt((phi1 * mu1) / ((phi1 + 1) * YO1^3))
    ) +
    log((1 / (2 * sqrt(2 * pi))) * sqrt(phi1 / 2))

  integral_arg <-
    sqrt((phi2 + 1) / (2 * xi1 * r)) * ((phi2 * xi1) / (phi2 + 1) - 1) +
    rho * sqrt(phi1 / (2 * r)) * a1

  log_selected <- log_bs_density + stats::pnorm(integral_arg, log.p = TRUE)

  cens_arg <- sqrt(phi2 / 2) *
    (sqrt((phi2 + 1) / (phi2 * xi0)) - sqrt((phi2 * xi0) / (phi2 + 1)))

  log_censored <- stats::pnorm(cens_arg, log.p = TRUE)

  ll <- sum(log_selected) + sum(log_censored)
  if (!is.finite(ll)) return(NA_real_)
  ll
}

heckmanBS_score_independente <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanBS_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)
  iphi1  <- max(ibeta) + 1L
  irho   <- iphi1 + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  phi1  <- par[iphi1]
  rho_star <- par[irho]

  if (!is.finite(phi1) || phi1 < 1e-3 || phi1 > 100 || !is.finite(rho_star)) {
    return(rep(NA_real_, length(par)))
  }

  rho <- .bs_inv_logit_2(rho_star)
  phi2 <- 1
  r <- 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]
  YO1 <- YO[YS == 1L]
  XO1 <- XO[YS == 1L, , drop = FALSE]

  xi0 <- as.vector(exp(XS0 %*% gamma))
  xi1 <- as.vector(exp(XS1 %*% gamma))
  mu1 <- as.vector(exp(XO1 %*% beta))

  a <- sqrt(YO1 * (phi1 + 1) / (phi1 * mu1)) -
    sqrt((phi1 * mu1) / (YO1 * (phi1 + 1)))

  b <- sqrt((phi1 + 1) / (phi1 * mu1 * YO1)) +
    sqrt((phi1 * mu1) / ((phi1 + 1) * YO1^3))

  c_arg <- sqrt((phi2 + 1) / (2 * xi1 * r)) * ((phi2 * xi1) / (phi2 + 1) - 1) +
    rho * sqrt(phi1 / (2 * r)) * a

  d_arg <- sqrt(phi2 / 2) *
    (sqrt((phi2 + 1) / (phi2 * xi0)) - sqrt((phi2 * xi0) / (phi2 + 1)))

  lambda_c <- .bs_safe_ratio(c_arg)
  lambda_d <- .bs_safe_ratio(d_arg)

  # Derivadas auxiliares em relação aos preditores lineares
  da_dbeta <- sweep(XO1, 1,
                    (-0.5) * (
                      sqrt(YO1 * (phi1 + 1) / (phi1 * mu1)) +
                        sqrt((phi1 * mu1) / (YO1 * (phi1 + 1)))
                    ), `*`)

  db_dbeta <- sweep(XO1, 1,
                    0.5 * (
                      sqrt((mu1 * phi1) / (YO1^3 * (phi1 + 1))) -
                        sqrt((phi1 + 1) / (YO1 * phi1 * mu1))
                    ), `*`)

  dd_dgamma <- sweep(XS0, 1,
                     (-0.5) * sqrt(phi2 / 2) * (
                       sqrt((phi2 + 1) / (phi2 * xi0)) +
                         sqrt((phi2 * xi0) / (phi2 + 1))
                     ), `*`)

  da_dphi <- (-0.5) * (
    sqrt(YO1 / (phi1^3 * mu1 * (phi1 + 1))) +
      sqrt(mu1 / (YO1 * (phi1 + 1)^3 * phi1))
  )

  db_dphi <- 0.5 * (
    sqrt(mu1 / (phi1 * (phi1 + 1)^3 * YO1^3)) -
      sqrt(1 / (phi1^3 * (phi1 + 1) * mu1 * YO1))
  )

  dlog_const_dphi <- 1 / (2 * phi1)

  dc_dphi <- rho / (2 * sqrt(2 * phi1 * r)) * a +
    rho * sqrt(phi1 / (2 * r)) * da_dphi

  dc_dbeta <- sweep(da_dbeta, 1, rho * sqrt(phi1 / (2 * r)), `*`)

  # c = A(xi, rho) + rho sqrt(phi/(2r)) a
  # dA/deta = (xi + 2)/(4 sqrt(xi*r)), with phi2=1
  dc_dgamma <- sweep(XS1, 1, (xi1 + 2) / (4 * sqrt(xi1 * r)), `*`)

  # Derivada em relação a rho_star
  A <- sqrt((phi2 + 1) / (2 * xi1 * r)) * ((phi2 * xi1) / (phi2 + 1) - 1)
  dA_drho <- A * rho / r
  dB_drho <- sqrt(phi1 / 2) * (r^(-0.5) + rho^2 * r^(-1.5)) * a
  dc_drho <- dA_drho + dB_drho
  drho_drhostar <- 2 * exp(-rho_star) / (1 + exp(-rho_star))^2
  dc_drhostar <- dc_drho * drho_drhostar

  grad <- numeric(length(par))

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

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

    grad[ibeta] <- grad[ibeta] +
      colSums((-phi1 / 2) * sweep(da_dbeta, 1, a, `*`) +
                sweep(db_dbeta, 1, 1 / b, `*`) +
                sweep(dc_dbeta, 1, lambda_c, `*`))

    grad[iphi1] <- grad[iphi1] +
      sum(-a^2 / 4 - (phi1 / 2) * a * da_dphi + db_dphi / b +
            dlog_const_dphi + lambda_c * dc_dphi)

    grad[irho] <- grad[irho] + sum(lambda_c * dc_drhostar)
  }

  grad
}

heckmanBS_density_components <- function(par, selection, outcome, data, components = NULL) {
  if (is.null(components)) components <- heckmanBS_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)
  iphi1  <- max(ibeta) + 1L
  irho   <- iphi1 + 1L

  gamma <- par[igamma]
  beta  <- par[ibeta]
  phi1  <- par[iphi1]
  rho <- .bs_inv_logit_2(par[irho])
  phi2 <- 1
  r <- 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)
    xi1 <- as.vector(exp(XS[idx1, , drop = FALSE] %*% gamma))
    mu1 <- as.vector(exp(XO[idx1, , drop = FALSE] %*% beta))
    y1 <- YO[idx1]
    a <- sqrt(y1 * (phi1 + 1) / (phi1 * mu1)) - sqrt((phi1 * mu1) / (y1 * (phi1 + 1)))
    log_bs_density <- (-phi1 / 4) * a^2 +
      log(sqrt((phi1 + 1) / (phi1 * mu1 * y1)) + sqrt((phi1 * mu1) / ((phi1 + 1) * y1^3))) +
      log((1 / (2 * sqrt(2 * pi))) * sqrt(phi1 / 2))
    c_arg <- sqrt((phi2 + 1) / (2 * xi1 * r)) * ((phi2 * xi1) / (phi2 + 1) - 1) +
      rho * sqrt(phi1 / (2 * r)) * a
    logc <- log_bs_density + stats::pnorm(c_arg, log.p = TRUE)
    out$value[idx1] <- exp(logc)
    out$log_contribution[idx1] <- logc
  }

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

  out
}
