## ===========================================================================
## Median-of-Means gradient-descent estimation of the modality
## information gain H_m, for heavy-tailed responses.
## ===========================================================================

## ---- Per-observation log-likelihood at beta (family-dispatched) ------------
##
## Returns (1/n) * logL(beta), up to an additive constant that does not depend
## on beta and therefore cancels in the H_m difference ll_full - ll_reduced.
## An intercept column is prepended internally, so beta = c(intercept, coefs)
## has length ncol(X) + 1.
##   gaussian          : profile (concentrated-variance) log-likelihood,
##                       (1/n) logL = -1/2 * log(RSS / n)  (+ const)
##   negative.binomial : NB log-likelihood at fixed dispersion theta.
## A design matrix with zero columns encodes the intercept-only (null) model.
## (median, quantile, var and sd come from 'stats', imported package-wide.)
mom_loglik <- function(X, y, beta, family, theta = NULL) {
  Xa  <- cbind(1, as.matrix(X))                     # prepend intercept column
  eta <- as.vector(Xa %*% beta)                     # beta = c(intercept, coefs)
  if (family == "gaussian") {
    rss <- max(sum((y - eta)^2), 1e-12)
    return(-0.5 * log(rss / length(y)))
  }
  if (family == "negative.binomial") {
    eta <- pmax(pmin(eta, 10), -10)                 # clamp for mu overflow
    mu  <- exp(eta)
    ll  <- sum(lgamma(y + theta) - lgamma(theta) - lgamma(y + 1) +
                 theta * log(theta / (theta + mu)) +
                 y     * log(mu    / (theta + mu)))
    return(ll / length(y))
  }
  stop("mom_loglik: 'family' must be 'gaussian' or 'negative.binomial'.")
}

## ---- Score of (1/n_blk) logL w.r.t. beta on one block ---
##
## An intercept column is prepended internally; the returned gradient has length
## ncol(X) + 1 and matches the beta = c(intercept, coefs) parameterisation used
## by mom_loglik.
## Gaussian          : g = Xa' (y - Xa beta) / n_blk
## Negative-binomial : g = Xa' r / n_blk,  r_i = theta (y_i - mu_i)/(mu_i + theta)
mom_score <- function(X, y, beta, family, theta = NULL) {
  Xa    <- cbind(1, as.matrix(X))                   # prepend intercept column
  n_blk <- nrow(Xa)
  eta   <- as.vector(Xa %*% beta)
  if (family == "gaussian") {
    r <- y - eta
    return(as.vector(crossprod(Xa, r)) / n_blk)
  }
  if (family == "negative.binomial") {
    eta <- pmax(pmin(eta, 10), -10)
    mu  <- exp(eta)
    r   <- theta * (y - mu) / (mu + theta)
    return(as.vector(crossprod(Xa, r)) / n_blk)
  }
  stop("mom_score: 'family' must be 'gaussian' or 'negative.binomial'.")
}

## ---- MoM gradient descent: median-log-likelihood block selection -----------
mom_gd <- function(X, y, family, theta = NULL, nblock = 5,
                   lr = NULL, maxit = NULL, tol = 1e-5) {
  X <- as.matrix(X)
  q <- ncol(X)
  n <- nrow(X)

  ## family-specific defaults (match the reference simulation settings)
  if (is.null(lr))    lr    <- switch(family, gaussian = 0.02,
                                      negative.binomial = 1e-3, 0.01)
  if (is.null(maxit)) maxit <- switch(family, gaussian = 1000L,
                                      negative.binomial = 500L, 1000L)

  K      <- max(1L, min(as.integer(nblock), n))
  fold   <- sort(rep_len(seq_len(K), n))            # contiguous near-equal blocks
  blocks <- split(seq_len(n), fold)
  beta   <- rep(0, q + 1L)                          # cold start (intercept + coefs)

  for (it in seq_len(maxit)) {
    old <- beta
    ll  <- vapply(blocks, function(idx)
      mom_loglik(X[idx, , drop = FALSE], y[idx], beta, family, theta),
      numeric(1))
    m_idx <- which.min(abs(ll - median(ll)))        # median-log-lik block
    idx   <- blocks[[m_idx]]
    g     <- mom_score(X[idx, , drop = FALSE], y[idx], beta, family, theta)
    beta  <- beta + lr * g
    if (max(abs(beta - old)) < tol) break
  }
  beta
}

## ---- MoM-GD point estimate of H_m ------------------------------------------
##
## Fits the full (all selected features) and reduced (selected features in the
## conditioning modalities only) models by MoM and returns the per-
## observation log-likelihood gain ll_full - ll_reduced.
mom_Th <- function(X, y, var_id, Mbeta, family, theta = NULL, nblock = 5) {
  sel_full <- var_id
  sel_red  <- var_id[var_id %in% Mbeta]
  Xf <- X[, sel_full, drop = FALSE]
  Xr <- X[, sel_red,  drop = FALSE]
  bf <- mom_gd(Xf, y, family, theta, nblock)
  br <- mom_gd(Xr, y, family, theta, nblock)        # intercept-only if reduced is null
  mom_loglik(Xf, y, bf, family, theta) -
    mom_loglik(Xr, y, br, family, theta)
}

## ---- Basic (pivotal) bootstrap CI for the MoM estimate of H_m -----------
##
## Resamples observations with replacement 'nboot' times, recomputing the full
## MoM pipeline on each resample, and inverts the bootstrap distribution
## with the basic bootstrap:
##   two.sided : [ 2*Th - q_{1-a/2},  2*Th - q_{a/2} ]
##   one.sided : [ 2*Th - q_{1-a},    Inf )
## The lower bound is floored at 0 because H_m >= 0 (keeps pseudo-R^2 valid).
mom_boot_ci <- function(X, y, var_id, Mbeta, family, theta = NULL, nblock = 5,
                        Th, nboot = 200, conf.level = 0.95,
                        CI.type = c("two.sided", "one.sided")) {
  CI.type <- match.arg(CI.type)
  n     <- nrow(X)
  alpha <- 1 - conf.level
  Tb    <- rep(NA_real_, nboot)

  for (b in seq_len(nboot)) {
    id <- sample.int(n, n, replace = TRUE)
    Tb[b] <- tryCatch(
      mom_Th(X[id, , drop = FALSE], y[id], var_id, Mbeta, family, theta, nblock),
      error = function(e) NA_real_)
  }
  Tb <- Tb[is.finite(Tb)]
  if (length(Tb) < 2L) return(c(NA_real_, NA_real_))

  if (CI.type == "two.sided") {
    q  <- quantile(Tb, c(alpha / 2, 1 - alpha / 2), na.rm = TRUE)
    ci <- c(2 * Th - q[2], 2 * Th - q[1])
  } else {
    q  <- quantile(Tb, conf.level, na.rm = TRUE)
    ci <- c(2 * Th - q, Inf)
  }
  ci[1] <- max(ci[1], 0)                             # H_m >= 0
  names(ci) <- NULL
  ci
}

## ---- MoM estimation wrapper --------
##
## Returns the same list shape as mod_est: num.nonzeros, Th, CI.  The bootstrap
## CI is computed only when CI = TRUE (default FALSE for the robust path).
mom_est <- function(X, y, var_id, Mbeta, family, theta = NULL, nblock = 5,
                    CI = FALSE, nboot = 200, conf.level = 0.95,
                    CI.type = c("two.sided", "one.sided")) {
  CI.type <- match.arg(CI.type)
  if (length(var_id) == 0L)
    return(list(num.nonzeros = 0, Th = 0, CI = c(NA, NA)))
  num.nonzeros <- length(var_id) - length(var_id[var_id %in% Mbeta])
  if (num.nonzeros == 0)
    return(list(num.nonzeros = 0, Th = 0, CI = c(NA, NA)))

  Th <- mom_Th(X, y, var_id, Mbeta, family, theta, nblock)

  ci <- c(NA, NA)
  if (isTRUE(CI))
    ci <- mom_boot_ci(X, y, var_id, Mbeta, family, theta, nblock,
                      Th = Th, nboot = nboot, conf.level = conf.level,
                      CI.type = CI.type)

  list(num.nonzeros = num.nonzeros, Th = Th, CI = ci)
}

## ---- Fixed-dispersion estimate for the Negative-Binomial family ------------
##
## When the user does not supply theta, estimate it once (on the screened
## design) via MASS::glm.nb, falling back to a method-of-moments estimate.
#' @importFrom MASS glm.nb
est_theta_nb <- function(y, Xsub) {
  th <- tryCatch(
    as.numeric(suppressWarnings(glm.nb(y ~ as.matrix(Xsub)))$theta),
    error = function(e) NA_real_)
  if (is.na(th) || !is.finite(th) || th <= 0) {
    mu <- max(mean(y), 0.01)
    v  <- max(var(y), mu + 0.01)
    th <- min(max(mu^2 / (v - mu), 0.01), 100)      # MoM fallback
  }
  th
}
