#' Frank-Wolfe homotopy loop for a single subset size `k`.
#'
#' This is an internal helper shared by all three GLM families.
#'
#' @param t_init Initial relaxed inclusion vector of length `p_sel`.
#' @param k Target subset size.
#' @param X Normalised design matrix (no intercept column).
#' @param y Response.
#' @param delta_min,delta_max Homotopy schedule endpoints.
#' @param Niter Number of homotopy iterations (loop runs `2 * Niter` steps).
#' @param alpha Frank-Wolfe step size.
#' @param lambda Ridge penalty.
#' @param mand_idx,sel_idx Mandatory and selectable column indices.
#' @param solve_inner Inner ridge solver function.
#' @return List with `selected` (1-indexed columns of `X`), `t`, `s`.
#' @keywords internal
#' @noRd
fw_loop <- function(t_init, k, X, y,
                    delta_min, delta_max,
                    Niter, alpha, lambda,
                    mand_idx, sel_idx,
                    solve_inner) {
  p_sel <- length(sel_idx)
  r <- (delta_max / delta_min)^(1 / Niter)
  t <- t_init
  s <- rep(0, p_sel)

  for (i in seq_len(2 * Niter)) {
    delta <- min(delta_min * r^i, delta_max)
    xi <- solve_inner(t, X, y, delta, lambda, mand_idx, sel_idx)
    g  <- grad_envelope(xi, t, delta, lambda, sel_idx)
    s <- rep(0, p_sel)
    s[order(g)[seq_len(k)]] <- 1
    t <- (1 - alpha) * t + alpha * s
    t <- pmin(pmax(t, 1e-4), 1 - 1e-4)
  }

  selected_sel  <- which(s == 1)
  selected_full <- sel_idx[selected_sel]
  if (length(mand_idx) > 0) {
    selected_full <- sort(c(mand_idx, selected_full))
  }
  list(selected = selected_full, t = t, s = s)
}

#' Refit on a subset of features and predict on test data.
#'
#' @keywords internal
#' @noRd
refit_predict <- function(X_train, y_train, X_test, idx, family) {
  if (length(idx) == 0) {
    if (family == "gaussian") {
      pred <- rep(mean(y_train), nrow(X_test))
    } else {
      tab <- table(y_train)
      majority <- names(which.max(tab))
      pred <- rep(if (family == "binomial") as.numeric(majority) else majority,
                  nrow(X_test))
    }
    return(pred)
  }

  X_sub_train <- X_train[, idx, drop = FALSE]
  X_sub_test  <- X_test[, idx, drop = FALSE]

  if (family == "gaussian") {
    if (ncol(X_sub_train) == 1) {
      df <- data.frame(y = y_train, x = as.numeric(X_sub_train))
      fit <- stats::lm(y ~ x, data = df)
      return(as.numeric(stats::predict(fit,
        newdata = data.frame(x = as.numeric(X_sub_test)))))
    }
    fit <- glmnet::glmnet(x = X_sub_train, y = y_train, family = "gaussian",
                          alpha = 0, lambda = 1e-8,
                          intercept = TRUE, standardize = FALSE)
    return(as.numeric(stats::predict(fit, newx = X_sub_test, s = 1e-8)))
  }

  if (family == "binomial") {
    if (ncol(X_sub_train) == 1) {
      df <- data.frame(y = y_train, x = as.numeric(X_sub_train))
      fit <- stats::glm(y ~ x, data = df, family = stats::binomial())
      prob <- stats::predict(fit,
        newdata = data.frame(x = as.numeric(X_sub_test)),
        type = "response")
    } else {
      fit <- glmnet::glmnet(x = X_sub_train, y = y_train, family = "binomial",
                            alpha = 0, lambda = 1e-8,
                            intercept = TRUE, standardize = FALSE)
      prob <- as.numeric(stats::predict(fit, newx = X_sub_test, s = 1e-8,
                                        type = "response"))
    }
    return(as.numeric(prob > 0.5))
  }

  # multinomial
  X_refit <- X_sub_train
  X_test_refit <- X_sub_test
  if (ncol(X_refit) == 1) {
    X_refit      <- cbind(X_refit, 0)
    X_test_refit <- cbind(X_test_refit, 0)
  }
  fit <- glmnet::glmnet(x = X_refit, y = y_train, family = "multinomial",
                        type.multinomial = "grouped", alpha = 0,
                        nlambda = 50, intercept = TRUE, standardize = TRUE,
                        thresh = 1e-5, maxit = 1e5)
  lam_refit <- min(fit$lambda)
  pred <- stats::predict(fit, newx = X_test_refit, s = lam_refit, type = "class")
  as.character(if (is.matrix(pred)) pred[, 1] else pred)
}
