#' @title Modality Assessment in Multimodal Generalized Linear Models
#'
#' @description Provides statistical inference for modality-specific information gain in
#' multimodal GLMs. Estimates ERE and pseudo-\eqn{R^2} with confidence intervals using
#' Sure Independence Screening for variable selection and penalized likelihood for inference.
#'
#'
#' @param X The finite numeric \eqn{n \times p} data matrix consisting of features from all modalities.
#' @param y The response. For \code{family = "gaussian"} or \code{"binomial"} this is the
#' length-\eqn{n} vector (or one-column matrix/data frame) of finite responses; binomial responses
#' must contain both 0 and 1.
#' Negative-binomial responses must be non-negative integer counts. For \code{family = "cox"} this is a
#' \code{\link[survival]{Surv}} object, or equivalently an \eqn{n \times 2} matrix / data frame
#' whose first column is the observed time and whose second column is the event indicator.
#' @param mod.idx A non-empty list of disjoint column-index vectors for all modalities in the
#' concatenated data matrix \code{X}. Every column of \code{X} must be assigned to exactly one modality.
#' @param family A description of the error distribution and link function to be used in the model.
#' With \code{robust = "none"} the Gaussian (\code{"gaussian"}) and Binomial (\code{"binomial"})
#' families with canonical links and the Cox proportional hazards model (\code{"cox"}) are available.
#' With \code{robust = "mom"} the Gaussian (\code{"gaussian"}) and Negative-Binomial
#' (\code{"negative.binomial"}) families are available.
#' @param robust Estimator for the modality information gain. \code{"none"} (default) uses the
#' standard maximum-likelihood / partial-likelihood estimator with closed-form non-central
#' chi-square confidence intervals. \code{"mom"} uses the Median-of-Means gradient-descent
#' (MoM-GD) estimator, which is robust to heavy-tailed responses, with basic bootstrap
#' confidence intervals.
#' @param iter Specifies whether to perform iterative SIS. For the Cox family this toggles between
#' a single BIC-driven screening pass (\code{iter=FALSE}) and the full iterative (ISIS) procedure
#' (\code{iter=TRUE}). The default is \code{iter=TRUE}.
#' @param penalty Specifies the type of penalty to be used in the variable selection and
#' inference procedure.
#' Options include \code{'MCP'}, \code{'SCAD'}, and \code{'lasso'}. The default is \code{penalty='SCAD'}.
#' @param tune Specifies the method for selecting the optimal tuning parameters in (I)SIS and
#' penalized likelihood procedure. Options include \code{'bic'}, \code{'ebic'} and \code{'aic'}.
#' The default is \code{tune='bic'}.
#' @param lambda A user-specified finite, positive, strictly decreasing sequence of at least two
#' lambda values for penalized likelihood
#' procedure. By default, a sequence of values of length \code{nlambda} is automatically computed and
#' equally spaced on the log scale.
#' @param nlambda The number of lambda values, an integer of at least 2. The default is 100.
#' @param conf.level Confidence level, a finite number strictly between 0 and 1. The default is
#' \code{conf.level=0.95}.
#' @param CI.type A string specifying the type of the confidence interval. Options include
#' \code{'two.sided'} and \code{'one.sided'}. The default is \code{CI.type='two.sided'}.
#' @param CI Logical; whether to compute confidence intervals. When \code{robust = "none"} the
#' default is \code{TRUE} (closed-form non-central chi-square intervals). When
#' \code{robust = "mom"} the default is \code{FALSE}; set \code{CI = TRUE} to run the basic
#' bootstrap. If \code{CI = FALSE}, only point estimates are returned (CI bounds are \code{NA}).
#' @param nboot Number of bootstrap resamples (at least 2) used for the MoM basic bootstrap confidence interval
#' (used only when \code{robust = "mom"} and \code{CI = TRUE}). The default is 200.
#' @param nblock Number of Median-of-Means blocks used by the MoM-GD estimator (used only when
#' \code{robust = "mom"}). Must be an integer between 1 and \code{nrow(X)}. The rows of \code{X}
#' are partitioned into \code{nblock} contiguous,
#' near-equal blocks. The default is 5.
#' @param theta Negative-Binomial dispersion parameter (used only when
#' \code{family = "negative.binomial"}), a finite positive number treated as fixed in the MoM fit. If \code{NULL}
#' (default), it is estimated once from the screened design.
#' @param trace Specifies whether to print out logs of iterations in SIS procedure. The default is
#' \code{trace=FALSE}.
#'
#' @return An object with S3 class "mglm.test" containing:
#' \item{sel.idx}{ List of indices of selected features by (I)SIS in each modality.}
#' \item{num.nonzeros}{Number of selected features by (I)SIS in each modality.}
#' \item{ERE}{Point estimation of ERE for each modality.}
#' \item{ERE.CI.L}{Lower bound of the confidence interval of ERE for each modality}
#' \item{ERE.CI.U}{Upper bound of the confidence interval of ERE for each modality}
#' \item{R2}{Point estimate of pseudo-\eqn{R^2} for each modality.}
#' \item{R2.CI.L}{Lower bound of the confidence interval of pseudo-\eqn{R^2} for each modality}
#' \item{R2.CI.U}{Upper bound of the confidence interval of pseudo-\eqn{R^2} for each modality}
#' \item{conf.level}{ Level of confidence intervals.}
#' @export
#'
#' @importFrom MASS mvrnorm glm.nb
#' @importFrom SIS SIS
#' @importFrom glmnet glmnet
#' @importFrom glmnet cv.glmnet
#' @importFrom ncvreg ncvreg
#' @importFrom ncvreg cv.ncvreg
#' @importFrom MBESS conf.limits.nc.chisq
#' @importFrom survival coxph Surv
#' @importFrom dplyr case_when
#' @importFrom utils capture.output

#' @examples
#' ## Example 1: Linear model
#' data(data_linear_model)
#' X <- data_linear_model$X
#' y <- data_linear_model$y
#' mod.idx <- data_linear_model$mod.idx
#' test <- mglm.test(X = X, y = y, mod.idx = mod.idx, family = "gaussian",
#'                iter = TRUE, penalty = "SCAD", tune = "bic",
#'                conf.level = 0.95, CI.type = "one.sided")
#' summary(test)
#'
#'
#' ## Example 2: Logistic regression
#' data(data_logistic_model)
#' X <- data_logistic_model$X
#' y <- data_logistic_model$y
#' mod.idx <- data_logistic_model$mod.idx
#' test <- mglm.test(X = X, y = y, mod.idx = mod.idx, family = "binomial",
#'                iter = TRUE, penalty = "SCAD", tune = "bic",
#'                conf.level = 0.95, CI.type = "two.sided")
#' sum.test <- summary(test)
#'
#'
#' ## Example 3: Cox proportional hazards model
#' \donttest{
#' data(data_cox_model)
#' X <- data_cox_model$X
#' y <- data_cox_model$y          # two-column matrix: (time, status)
#' mod.idx <- data_cox_model$mod.idx
#' test <- mglm.test(X = X, y = y, mod.idx = mod.idx, family = "cox",
#'                iter = TRUE, penalty = "SCAD", tune = "ebic",
#'                conf.level = 0.95, CI.type = "two.sided")
#' summary(test)
#' }
#'
#'
#' ## Example 4: Robust (MoM) estimation under heavy-tailed responses
#' data(data_nb_model)
#' X <- data_nb_model$X
#' y <- data_nb_model$y
#' mod.idx <- data_nb_model$mod.idx
#' ## Point estimates only (CI = FALSE is the default when robust = "mom"):
#' test <- mglm.test(X = X, y = y, mod.idx = mod.idx,
#'                family = "negative.binomial", robust = "mom",
#'                iter = FALSE, penalty = "SCAD", tune = "bic",
#'                nblock = 5, theta = data_nb_model$theta)
#' summary(test)
#'
#'
mglm.test <- function(X, y, mod.idx,
                 family = c("gaussian", "binomial", "cox", "negative.binomial"),
                 robust = c("none", "mom"), iter = TRUE,
                 penalty = c("SCAD", "MCP", "lasso"),
                 tune = c("bic", "ebic", "aic"),
                 lambda = NULL, nlambda = 100,
                 conf.level = 0.95, CI.type = c("two.sided", "one.sided"),
                 CI = NULL, nboot = 200, nblock = 5, theta = NULL,
                 trace = FALSE){

  ###########################################################################################
  # Input argument checks and modifications
  ###########################################################################################

  family = match.arg(family)
  robust = match.arg(robust)
  penalty = match.arg(penalty)
  tune = match.arg(tune)
  CI.type = match.arg(CI.type)

  if (is.null(X) || is.null(y))
    stop("X and y must both be supplied.")
  if (!is.matrix(X) || !is.numeric(X))
    stop("X must be a numeric matrix.")
  if (nrow(X) < 2L || ncol(X) < 1L)
    stop("X must contain at least two rows and one column.")
  if (anyNA(X) || any(!is.finite(X)))
    stop("X must not contain missing or non-finite values.")

  nobs = nrow(X)
  nvars = ncol(X)

  if (!is.list(mod.idx) || length(mod.idx) == 0L)
    stop("mod.idx must be a non-empty list of modality column indices.")
  valid_modality <- vapply(mod.idx, function(idx) {
    is.numeric(idx) && length(idx) > 0L && !anyNA(idx) &&
      all(is.finite(idx)) && all(idx == floor(idx)) &&
      all(idx >= 1L & idx <= nvars) && !anyDuplicated(idx)
  }, logical(1))
  if (!all(valid_modality))
    stop("Each element of mod.idx must contain unique integer indices between 1 and ncol(X).")
  all_modality_indices <- unlist(mod.idx, use.names = FALSE)
  if (anyDuplicated(all_modality_indices))
    stop("Modalities in mod.idx must be disjoint.")
  if (!identical(sort(as.integer(all_modality_indices)), seq_len(nvars)))
    stop("mod.idx must assign every column of X to exactly one modality.")

  is_single_logical <- function(z) is.logical(z) && length(z) == 1L && !is.na(z)
  is_integer_count <- function(z, minimum) {
    is.numeric(z) && length(z) == 1L && !is.na(z) && is.finite(z) &&
      z == floor(z) && z >= minimum
  }

  if (!is_single_logical(iter))
    stop("iter must be a single non-missing logical value.")
  if (!is_single_logical(trace))
    stop("trace must be a single non-missing logical value.")
  if (!is.numeric(conf.level) || length(conf.level) != 1L ||
      is.na(conf.level) || !is.finite(conf.level) ||
      conf.level <= 0 || conf.level >= 1)
    stop("conf.level must be a single finite number strictly between 0 and 1.")
  if (!is_integer_count(nlambda, 2L))
    stop("nlambda must be a single integer greater than or equal to 2.")
  if (!is_integer_count(nboot, 2L))
    stop("nboot must be a single integer greater than or equal to 2.")
  if (!is_integer_count(nblock, 1L) || nblock > nobs)
    stop("nblock must be a single integer between 1 and nrow(X).")
  if (!is.null(lambda)) {
    if (!is.numeric(lambda) || length(lambda) < 2L || anyNA(lambda) ||
        any(!is.finite(lambda)) || any(lambda <= 0) || any(diff(lambda) >= 0))
      stop("lambda must be a finite, positive, strictly decreasing numeric sequence of length at least 2.")
  }
  if (!is.null(theta) &&
      (!is.numeric(theta) || length(theta) != 1L || is.na(theta) ||
       !is.finite(theta) || theta <= 0))
    stop("theta must be NULL or a single finite positive number.")

  ## Default CI behaviour depends on the estimator: closed-form CIs are cheap
  ## for the non-robust path (default on), whereas the robust path needs a
  ## bootstrap (default off unless the user asks for it).
  if (is.null(CI)) CI = (robust == "none")
  if (!is_single_logical(CI))
    stop("CI must be a single logical value (TRUE or FALSE).")

  ## Family must be compatible with the chosen estimator.
  if (robust == "none") {
    if (!(family %in% c("gaussian", "binomial", "cox")))
      stop("With robust='none', 'family' must be 'gaussian', 'binomial', or 'cox'. ",
           "The Negative-Binomial family is only available with robust='mom'.")
  } else { ## robust == "mom"
    if (!(family %in% c("gaussian", "negative.binomial")))
      stop("With robust='mom', 'family' must be 'gaussian' or 'negative.binomial'. ",
           "Median-of-Means estimation is not implemented for family='",
           family, "'.")
  }

  ## For the Cox model the response must be a survival object. Accept either a
  ## 'Surv' object or a two-column (time, status) matrix / data frame.
  if (family == "cox") {
    if (inherits(y, "Surv")) {
      if (nrow(y) != nobs)
        stop("X and y must contain the same number of observations.")
      if (anyNA(y) || any(!is.finite(unclass(y))))
        stop("The Cox response must not contain missing or non-finite values.")
      ysurv = y
    } else {
      ymat = as.matrix(y)
      if (!is.numeric(ymat) || ncol(ymat) != 2)
        stop("For family='cox', 'y' must be a 'Surv' object or a two-column matrix/data.frame with columns (time, status).")
      if (nrow(ymat) != nobs)
        stop("X and y must contain the same number of observations.")
      if (anyNA(ymat) || any(!is.finite(ymat)))
        stop("The Cox response must not contain missing or non-finite values.")
      if (any(ymat[, 1] < 0))
        stop("Cox follow-up times must be non-negative.")
      if (!all(ymat[, 2] %in% c(0, 1)))
        stop("Cox event indicators must contain only 0 and 1.")
      ysurv = Surv(ymat[, 1], ymat[, 2])
    }
    yuse = ysurv
  } else {
    if (is.matrix(y) || is.data.frame(y)) {
      ymat = as.matrix(y)
      if (!is.numeric(ymat) || ncol(ymat) != 1L)
        stop("For non-Cox families, y must be a numeric or logical vector or a one-column matrix/data frame.")
      if (nrow(ymat) != nobs)
        stop("X and y must contain the same number of observations.")
      y = as.numeric(ymat[, 1])
    } else {
      if ((!is.numeric(y) && !is.logical(y)) || !is.null(dim(y)))
        stop("For non-Cox families, y must be a numeric or logical vector or a one-column matrix/data frame.")
      if (length(y) != nobs)
        stop("For non-Cox families, y must have nrow(X) observations.")
      y = as.numeric(y)
    }
    if (anyNA(y) || any(!is.finite(y)))
      stop("y must not contain missing or non-finite values.")
    if (family == "binomial") {
      if (!all(y %in% c(0, 1)) || length(unique(y)) != 2L)
        stop("For family='binomial', y must contain both 0 and 1 and no other values.")
    }
    if (family == "negative.binomial" &&
        (any(y < 0) || any(y != floor(y))))
      stop("For family='negative.binomial', y must contain non-negative integer counts.")
    yuse = y
  }

  nMod = length(mod.idx)
  sel.idx = list()
  ERE = c()
  num.nonzeros = c()
  ERE.CI.L = c()
  ERE.CI.U = c()


  ## Working family used for screening and penalised selection. The Negative-
  ## Binomial mean model shares the log link with the Poisson, so a Poisson
  ## working likelihood is used to screen/select features; MoM-GD then supplies
  ## the robust NB coefficient estimates on the selected set.
  work.family <- if (robust == "mom" && family == "negative.binomial") "poisson" else family

  ## Conduct (I)SIS for dimension reduction
  if(family == "cox"){
    ## (Iterative) sure independence screening for the Cox model
    var_id0 <- screen_cox(X, yuse, iter = iter, trace = trace)
  }else{
    if(trace){
      sis <- SIS(x=X, y=y, family=work.family, iter = iter, tune = tune)
    }else{
      invisible(capture.output(sis <- SIS(x=X, y=y, family=work.family, iter = iter, tune = tune)))
    }
    var_id0 <- sis$ix
  }
  if(length(var_id0)==0){
    warning("No variables selected after SIS")
  }

  ## Negative-Binomial dispersion (fixed for MoM-GD). Estimate once on the
  ## screened design if the user did not supply it.
  if (robust == "mom" && family == "negative.binomial" && is.null(theta)) {
    theta <- if (length(var_id0) != 0)
      est_theta_nb(y, X[, var_id0, drop = FALSE]) else est_theta_nb(y, matrix(0, length(y), 0))
  }

  for(i in seq_len(nMod)){
    Mbeta = setdiff(1:ncol(X),mod.idx[[i]])

    ## partially penalized regularization procedure
    ##
    ## Skipped when screening retained fewer than two features: a single
    ## screened feature is kept as-is for every penalty. The penalized step
    ## could not remove it anyway when it belongs to the conditioning set
    ## (penalty.factor = 0), a one-variable penalized fit is degenerate, and
    ## glmnet (the lasso backend) rejects one-column design matrices outright.
    pen.factor = as.numeric(var_id0 %in% Mbeta)
    if(length(var_id0) >= 2 && sum(pen.factor)!=0){
      ## coerce to a matrix so ncvsurv / coxph handle a single screened feature
      Xsel = as.matrix(X[, var_id0, drop = FALSE])
      tunefit = tune.fit.adj(Xsel, yuse, family=work.family , penalty = penalty, tune=tune,
                             lambda = lambda, nlambda = nlambda,
                             penalty.factor = pen.factor)
      var_id = var_id0[tunefit$ix]   # empty when the penalized fit selects nothing
    }else{
      var_id = var_id0
    }

    ## Estimation of the ERE / H_m
    if(robust == "mom"){
      fit = mom_est(X, y, var_id, Mbeta, family, theta = theta, nblock = nblock,
                    CI = CI, nboot = nboot, conf.level = conf.level, CI.type = CI.type)
    }else{
      fit = mod_est(X, yuse, var_id, Mbeta, family, conf.level = conf.level,
                    CI.type = CI.type, CI = CI)
    }
    sel.idx[[i]] <- var_id[!(var_id %in% Mbeta)]
    num.nonzeros <- c(num.nonzeros,fit$num.nonzeros)
    ERE <- c(ERE,fit$Th)
    ERE.CI.L <- c(ERE.CI.L,fit$CI[1])
    ERE.CI.U <- c(ERE.CI.U,fit$CI[2])
  }

  R2 = (1 - exp(-ERE))
  R2.CI.L = (1-exp(-ERE.CI.L))
  R2.CI.U = (1-exp(-ERE.CI.U))

  result = list(sel.idx=sel.idx, num.nonzeros=num.nonzeros,
                ERE=ERE, ERE.CI.L=ERE.CI.L, ERE.CI.U=ERE.CI.U,
                R2=R2, R2.CI.L=R2.CI.L, R2.CI.U=R2.CI.U,
                conf.level=conf.level)
  class(result) <- "mglm.test"
  return(result)
}


