R/boot_lucid.R

In LUCIDus: LUCID with Multiple Omics Data

#' @title Inference of LUCID model based on bootstrap resampling
#'
#' @description Generate \code{R} bootstrap replicates of LUCID parameters and
#' derive confidence interval (CI) based on bootstrap. Bootstrap replicates are
#' generated by nonparametric resampling, implemented with the \code{ordinary}
#' method of \code{boot::boot}. Supports \code{lucid_model = "early"},
#' \code{lucid_model = "parallel"}, and \code{lucid_model = "serial"}.
#'
#' @param G Exposures, a numeric vector, matrix, or data frame. Categorical variable
#' should be transformed into dummy variables. If a matrix or data frame, rows
#' represent observations and columns correspond to variables.
#' @param Z Omics data: for LUCID early integration, a numeric matrix/data frame; for LUCID in
#' parallel, a list of numeric matrices/data frames. Rows correspond to observations
#' and columns correspond to variables.
#' @param Y Outcome, a numeric vector. Categorical variable is not allowed. Binary
#' outcome should be coded as 0 and 1.
#' @param lucid_model Optional; "early", "parallel", or "serial". Auto-detected
#' from \code{class(model)} when omitted (the normal case), so this rarely
#' needs to be set explicitly -- it exists for backward compatibility with
#' scripts written before auto-detection. If supplied, it is cross-checked
#' against \code{model}'s actual class and an error is raised on a mismatch.
#' Bootstrap inference is implemented for all three model types.
#' @param CoG Optional, covariates to be adjusted for estimating the latent cluster.
#' A numeric vector, matrix or data frame. Categorical variable should be transformed
#' into dummy variables.
#' @param CoY Optional, covariates to be adjusted for estimating the association
#' between latent cluster and the outcome. A numeric vector, matrix or data frame.
#' Categorical variable should be transformed into dummy variables.
#' @param model A LUCID model fitted by \code{estimate_lucid}.
#' If the fitted model uses nonzero penalties, \code{boot_lucid} will
#' automatically refit a zero-penalty model as fallback because bootstrap
#' inference is only supported for \code{Rho_G = Rho_Z_Mu = Rho_Z_Cov = 0}.
#' @param conf A numeric scalar between 0 and 1 to specify confidence level(s)
#' of the required interval(s).
#' @param R An integer to specify number of bootstrap replicates for LUCID model.
#' If feasible, it is recommended to set R >= 1000.
#' @param verbose A flag indicates whether detailed information
#' is printed in console. Default is FALSE.
#' @param min_valid Minimum number of bootstrap replicates that must yield finite
#' estimates before confidence limits can be formed. The default, 2, is the
#' mathematical floor. Replicates that fail are counted and warned about, and a
#' small number of replicates raises a warning that the limits are unstable, but
#' neither suppresses the limits; only fewer than \code{min_valid} surviving
#' replicates yields NA limits.
#' 
#' @return A list containing:
#' \item{beta}{Bootstrap CI table(s) for G-to-X effects. For
#' \code{lucid_model = "parallel"}, this is a list by omics layer and includes
#' the multinomial intercept plus exposures in \code{G} (not \code{CoG}).}
#' \item{mu}{Bootstrap CI table(s) for cluster-specific means of omics features.
#' For \code{lucid_model = "parallel"}, this is a list by omics layer.}
#' \item{gamma}{Bootstrap CI table for X-to-Y parameters.}
#' \item{stage}{For \code{lucid_model = "serial"}, a list of stage-wise CI tables
#' (each stage contains \code{beta}, \code{mu}, and \code{gamma} for the final stage only).}
#' \item{bootstrap}{The \code{boot} object returned by \code{boot::boot}.}
#'
#' @export
#'
#' @import boot
#' @import progress
#'
#' @examples
#' \donttest{
#' # use simulated data (a small subset keeps the example quick)
#' G <- sim_data$G[1:150, , drop = FALSE]
#' Z <- sim_data$Z[1:150, , drop = FALSE]
#' Y_normal <- sim_data$Y_normal[1:150]
#'
#' # fit lucid model
#' fit1 <- estimate_lucid(G = G, Z = Z, Y = Y_normal, lucid_model = "early",
#' family = "normal", K = 2,
#' seed = 1008, max_itr = 20, max_tot.itr = 50)
#'
#' # conduct bootstrap resampling (lucid_model is auto-detected from fit1's class)
#' # a small R keeps the example quick; `conf` sets the CI level (default 0.95)
#' boot1 <- suppressWarnings(
#'   boot_lucid(G = G, Z = Z, Y = Y_normal, model = fit1, R = 3, conf = 0.9)
#' )
#' }
boot_lucid <- function(G,
                       Z,
                       Y,
                       lucid_model = NULL,
                       CoG = NULL,
                       CoY = NULL,
                       model,
                       conf = 0.95,
                       R = 100,
                       verbose = FALSE,
                       min_valid = 2L) {
  # `model`'s own class already says whether it's early/parallel/serial, so
  # lucid_model is auto-detected from it by default. A caller may still name
  # it explicitly (e.g. for backward compatibility with older scripts), in
  # which case it is cross-checked against model's actual class exactly as
  # before -- this branch is unchanged from prior behavior.
  if (is.null(lucid_model)) {
    lucid_model <- .detect_lucid_model(model)
  } else {
    lucid_model <- match.arg(lucid_model, c("early", "parallel", "serial"))
    expected_class <- switch(lucid_model,
                             early = "early_lucid",
                             parallel = "lucid_parallel",
                             serial = "lucid_serial")
    if (!inherits(model, expected_class)) {
      stop("'model' should be an object of class '", expected_class,
          "' to match lucid_model = '", lucid_model, "', but has class '",
          paste(class(model), collapse = "/"), "'.", call. = FALSE)
    }
  }
  check_complete_input(G, "G")
  check_complete_input(Y, "Y")
  check_complete_input(CoG, "CoG")
  check_complete_input(CoY, "CoY")
  model <- normalize_bootstrap_model(
    model = model,
    lucid_model = lucid_model,
    G = G,
    Z = Z,
    Y = Y,
    CoG = CoG,
    CoY = CoY
  )

  # prepare data for bootstrap (boot function require data in a matrix form,
  # list data structure doesn't work)
  if(!is.null(model$select) &&
     (!is.null(model$select$selectG) || !is.null(model$select$selectZ)) &&
     (has_unselected_feature(model$select$selectG) || has_unselected_feature(model$select$selectZ))) {
    stop("Refit LUCID model with selected feature first then conduct bootstrap inference")
  }
  if (lucid_model == "serial" && has_unselected_feature_serial(model)) {
    stop("Refit serial LUCID model with selected feature first then conduct bootstrap inference")
  }
  if (lucid_model == "early"){

    # ========================== Early Integration ==========================
    G <- as.matrix(G)
    Z <- as.matrix(Z)
    Y <- as.matrix(Y)
    dimG <- ncol(G)
    dimZ <- ncol(Z)
    dimCoG <- if (is.null(CoG)) 0 else ncol(as.matrix(CoG))
    dimCoY <- if (is.null(CoY)) 0 else ncol(as.matrix(CoY))
    K <- model$K
    alldata <- as.data.frame(cbind(G, Z, Y, CoG, CoY))

    # bootstrap
    if(verbose){
      cat(paste0("Use Bootstrap resampling to derive ", 100 * conf, "% CI for LUCID \n"))
    }
    #initialize progress bar object
    pb <- progress::progress_bar$new(total = R + 1)
    bootstrap <- boot::boot(data = alldata,
                            statistic = lucid_par_early,
                            R = R,
                            dimG = dimG,
                            dimZ = dimZ,
                            dimCoY = dimCoY,
                            dimCoG = dimCoG,
                            model = model,
                            prog = pb)

    # D8: boot() obtains t0 by calling the statistic on the original row order,
    # which refits the model under a fresh random seed.  That made bootstrap$t0
    # disagree with the model the user supplied, so summary(fit, boot.se = ...)
    # could print an `estimate` column inconsistent with summary(fit).  The
    # estimand is the supplied model, so take t0 from it directly.
    bootstrap$t0 <- lucid_early_par_vector(model, dimG)

    # bootstrap CIs
    ci <- gen_ci(bootstrap,
                conf = conf, min_valid = min_valid)

    # organize CIs
    # drop = FALSE: with a single exposure (or K = 2) these slices are a single
    # row and would otherwise collapse to a vector, breaking summary().
    # beta block = the whole G->X matrix (intercept + exposures + CoG) for each
    # non-reference cluster, matching lucid_early_par_vector().
    nBeta <- (K - 1) * ncol(model$res_Beta)
    beta <- ci[seq_len(nBeta), , drop = FALSE]
    mu <- ci[(nBeta + 1):(nBeta + K * dimZ), , drop = FALSE]
    gamma <- ci[-(seq_len(nBeta + K * dimZ)), , drop = FALSE]
    return(list(beta = beta,
                mu = mu,
                gamma = gamma,
                bootstrap = bootstrap))
  } else if (lucid_model == "parallel"){
      # ========================== Lucid in Parallel ==========================
      G <- as.matrix(G)
      Gnames_exposure <- colnames(G)
      if(is.null(Gnames_exposure) || length(Gnames_exposure) != ncol(G)) {
        Gnames_exposure <- paste0("G", seq_len(ncol(G)))
      }
      Y <- as.matrix(Y)
      if(!is.list(Z)) {
        stop("For lucid_model = 'parallel', input 'Z' should be a list of matrices/data frames")
      }
      Z <- lapply(Z, as.matrix)
      nOmics <- length(Z)
      if(nOmics == 0) {
        stop("For lucid_model = 'parallel', input 'Z' should contain at least one omics layer")
      }
      dimG <- ncol(G)
      dimZ <- sapply(Z, ncol)
      dimCoG <- if (is.null(CoG)) 0 else ncol(as.matrix(CoG))
      dimCoY <- if (is.null(CoY)) 0 else ncol(as.matrix(CoY))
      K <- model$K
      if(length(K) != nOmics) {
        stop("Length of model$K does not match number of omics layers in Z")
      }

      # The G->X beta block carried through the bootstrap is the whole coefficient
      # matrix -- exposures AND any CoG covariates -- so summary(fit, boot.se=)'s
      # (3) E table can show a CI for the covariate rows, matching summary(fit).
      CoGnames <- if (dimCoG > 0) {
        cn <- colnames(as.matrix(CoG))
        if (is.null(cn) || length(cn) != dimCoG) paste0("CoG", seq_len(dimCoG)) else cn
      } else character(0)
      Gnames_beta <- c(Gnames_exposure, CoGnames)
      dimG_beta <- dimG + dimCoG

      z_combined <- do.call(cbind, Z)
      alldata <- as.data.frame(cbind(G, z_combined, Y, CoG, CoY))

      # define template from fitted model to enforce fixed bootstrap statistic length
      template <- extract_parallel_boot_vector(
        model = model,
        dimG = dimG_beta,
        dimZ = dimZ,
        Gnames_exposure = Gnames_beta
      )
      template_names <- names(template)
      n_template <- length(template)

      if(verbose){
        cat(paste0("Use Bootstrap resampling to derive ", 100 * conf, "% CI for LUCID in parallel \n"))
      }
      pb <- progress::progress_bar$new(total = R + 1)
      bootstrap <- boot::boot(data = alldata,
                              statistic = lucid_par_parallel,
                              R = R,
                              dimG = dimG,
                              dimZ = dimZ,
                              dimCoY = dimCoY,
                              dimCoG = dimCoG,
                              model = model,
                              Gnames_exposure = Gnames_beta,
                              template_names = template_names,
                              n_template = n_template,
                              prog = pb)

      ci <- gen_ci(bootstrap, conf = conf, min_valid = min_valid)

      # split outputs by layer for easier consumption
      n_beta <- as.integer((K - 1) * (dimG_beta + 1))
      n_mu <- as.integer(K * dimZ)
      beta <- vector("list", nOmics)
      mu <- vector("list", nOmics)
      idx_start <- 1
      for(i in seq_len(nOmics)) {
        idx_end <- idx_start + n_beta[i] - 1
        beta[[i]] <- ci[idx_start:idx_end, , drop = FALSE]
        idx_start <- idx_end + 1
      }
      for(i in seq_len(nOmics)) {
        idx_end <- idx_start + n_mu[i] - 1
        mu[[i]] <- ci[idx_start:idx_end, , drop = FALSE]
        idx_start <- idx_end + 1
      }
      gamma <- ci[idx_start:nrow(ci), , drop = FALSE]

      names(beta) <- paste0("Layer", seq_len(nOmics))
      names(mu) <- paste0("Layer", seq_len(nOmics))

      return(list(beta = beta,
                  mu = mu,
                  gamma = gamma,
                  bootstrap = bootstrap))
  } else if (lucid_model == "serial"){
      # ========================== Lucid in Serial ==========================
      G <- as.matrix(G)
      Y <- as.matrix(Y)
      dimG <- ncol(G)
      dimCoG <- if (is.null(CoG)) 0 else ncol(as.matrix(CoG))
      dimCoY <- if (is.null(CoY)) 0 else ncol(as.matrix(CoY))
      z_flat <- flatten_serial_Z(Z)
      alldata <- as.data.frame(cbind(G, z_flat$flat, Y, CoG, CoY))

      template <- extract_serial_boot_template(model)
      template_names <- names(template$vector)
      n_template <- length(template_names)

      if(verbose){
        cat(paste0("Use Bootstrap resampling to derive ", 100 * conf, "% CI for LUCID in serial \n"))
      }
      pb <- progress::progress_bar$new(total = R + 1)
      bootstrap <- boot::boot(
        data = alldata,
        statistic = lucid_par_serial,
        R = R,
        dimG = dimG,
        dimCoY = dimCoY,
        dimCoG = dimCoG,
        z_meta = z_flat$meta,
        model = model,
        template_names = template_names,
        n_template = n_template,
        prog = pb
      )

      ci <- gen_ci(bootstrap, conf = conf, min_valid = min_valid)
      stage_ci <- split_serial_boot_ci(ci = ci, stage_layout = template$stage_layout)

      return(list(
        stage = stage_ci,
        bootstrap = bootstrap
      ))
  }
}



#' Extract the early-model bootstrap parameter vector from a fitted object
#'
#' Used for both the observed-data statistic (\code{t0}) and every
#' replicate, so the two can never disagree in layout or in value.
#'
#' @param fit A fitted \code{early_lucid} object.
#' @param dimG Number of true exposure columns (excluding covariates).
#' @param K Number of clusters; taken from \code{fit} if \code{NULL}. Pass
#'   the original model's \code{K} explicitly for a replicate fit, so the
#'   parameter vector keeps a fixed length across replicates.
#' @return A named numeric vector: exposure coefficients, then omics means,
#'   then outcome coefficients.
#' @noRd
lucid_early_par_vector <- function(fit, dimG, K = NULL) {
  # K is taken from the ORIGINAL model, not the replicate fit, so the parameter
  # vector keeps a fixed length across replicates.
  if (is.null(K)) K <- fit$K
  # The whole G->X coefficient matrix: intercept, then every exposure, then any
  # CoG covariate columns -- so summary(fit, boot.se=)'s (3) E table can show a
  # CI for the intercept and the covariates, matching summary(fit).
  beta_col_names <- colnames(fit$res_Beta)
  if (is.null(beta_col_names) || length(beta_col_names) != ncol(fit$res_Beta)) {
    beta_col_names <- c("intercept", paste0("G", seq_len(ncol(fit$res_Beta) - 1L)))
  }
  beta_col_names[1] <- "intercept"
  beta_block <- fit$res_Beta[-1, , drop = FALSE]
  out <- c(as.vector(t(beta_block)),
           as.vector(t(fit$res_Mu)),
           fit$res_Gamma$beta)
  G_names <- as.vector(sapply(2:K, function(x) {
    paste0(beta_col_names, ".cluster", x)
  }))
  Z_names <- as.vector(sapply(1:K, function(x) {
    paste0(colnames(fit$res_Mu), ".cluster", x)
  }))
  Y_names <- if (is.null(names(fit$res_Gamma$beta))) {
    paste0("cluster", 1:K)
  } else {
    names(fit$res_Gamma$beta)
  }
  names(out) <- c(G_names, Z_names, Y_names)
  out
}

#' Bootstrap replicate statistic for the early model
#'
#' \code{boot::boot()}'s \code{statistic} function for an early-model
#' bootstrap: refits on the resampled rows and extracts the parameter
#' vector. Refit failures are recorded as an \code{NA}-filled vector (of the
#' right length, so the replicate matrix stays rectangular) rather than
#' propagating the error and aborting the whole bootstrap run.
#'
#' @param data The combined data frame (G, Z, Y, CoG, CoY columns) passed to
#'   \code{boot::boot()}.
#' @param indices Row indices for this replicate.
#' @param model The original fitted model, supplying \code{K} and the
#'   fitting controls to reuse.
#' @param dimG,dimZ,dimCoY,dimCoG Column-block widths within \code{data}.
#' @param prog A \code{progress::progress_bar} to tick.
#' @return A named numeric vector (see
#'   \code{\link{lucid_early_par_vector}}), or all-\code{NA} if the refit
#'   failed.
#' @noRd
lucid_par_early <- function(data, indices, model, dimG, dimZ, dimCoY, dimCoG, prog) {
  #display progress with each run of the function
  prog$tick()

  # prepare data
  d <- data[indices, ]
  G <- as.matrix(d[, 1:dimG])
  Z <- as.matrix(d[, (dimG + 1):(dimG + dimZ)])
  Y <- as.matrix(d[, (dimG + dimZ + 1)])
  CoG <- CoY <- NULL
  K <- model$K
  if(dimCoG > 0){
    CoG <- as.matrix(d[, (dimG + dimZ + 2):(dimG + dimZ + dimCoG + 1)])
  }
  if(dimCoY > 0 && dimCoG > 0){
    CoY <- as.matrix(d[, (dimG + dimZ + dimCoG + 2):ncol(d)])
  }
  if(dimCoY > 0 && dimCoG == 0){
    CoY <- as.matrix(d[, (dimG + dimZ + 2):ncol(d)])
  }

  # fit lucid model
  seed <- sample(1:2000, 1)
  em_ctrl <- model$em_control
  rG <- if(!is.null(model$Rho$Rho_G)) model$Rho$Rho_G else 0
  rMu <- if(!is.null(model$Rho$Rho_Z_Mu)) model$Rho$Rho_Z_Mu else 0
  rCov <- if(!is.null(model$Rho$Rho_Z_Cov)) model$Rho$Rho_Z_Cov else 0
  tol_fit <- if(!is.null(em_ctrl$tol)) em_ctrl$tol else 0.001
  max_itr_fit <- if(!is.null(em_ctrl$max_itr)) em_ctrl$max_itr else 1000
  max_tot_fit <- if(!is.null(em_ctrl$max_tot.itr)) em_ctrl$max_tot.itr else 10000
  invisible(capture.output(try_lucid <- try(est_lucid(G = G,
                                                      Z = Z,
                                                      Y = Y,
                                                      CoY = CoY,
                                                      CoG = CoG,
                                                      lucid_model = "early",
                                                      family = model$family,
                                                      init_omic.data.model = model$init_omic.data.model,
                                                      K = K,
                                                      useY = model$useY,
                                                      tol = tol_fit,
                                                      max_itr = max_itr_fit,
                                                      max_tot.itr = max_tot_fit,
                                                      Rho_G = rG,
                                                      Rho_Z_Mu = rMu,
                                                      Rho_Z_Cov = rCov,
                                                      init_impute = model$init_impute,
                                                      init_par = model$init_par,
                                                      seed = seed), silent = TRUE)))
  if("try-error" %in% class(try_lucid)){
    n_par <- (K - 1) * ncol(model$res_Beta) + K * dimZ + length(model$res_Gamma$beta)
    par_lucid <- rep(NA_real_, n_par)
  } else{
    # Align this replicate's cluster labels to the reference fit before
    # extracting coefficients: an independent refit's cluster k need not be the
    # reference's cluster k, and stacking by position otherwise mixes the two.
    try_lucid <- tryCatch(align_replicate_early(try_lucid, model, indices), error = function(e) try_lucid)
    par_lucid <- lucid_early_par_vector(try_lucid, dimG, K = K)
    converge <- TRUE
  }
  return(par_lucid)
}

#' Reorder an early-model replicate fit's clusters to match the reference fit
#'
#' @param rep_fit The replicate \code{early_lucid} fit.
#' @param model The reference (point-estimate) fit.
#' @param indices The replicate's resampled row indices into the original data.
#' @return \code{rep_fit} with \code{res_Beta}, \code{res_Mu}, \code{res_Sigma},
#'   \code{res_Gamma} and \code{inclusion.p} reordered so cluster k lines up with
#'   the reference fit's cluster k. Returned unchanged when a match cannot be
#'   determined.
#' @noRd
align_replicate_early <- function(rep_fit, model, indices) {
  P_ref <- tryCatch(model$inclusion.p[indices, , drop = FALSE], error = function(e) NULL)
  perm <- match_boot_clusters(P_ref, rep_fit$inclusion.p)
  if (is.null(perm) || identical(as.integer(perm), seq_len(rep_fit$K))) {
    return(rep_fit)
  }
  rl <- relabel_early_parameters(rep_fit$res_Beta, rep_fit$res_Mu,
                                 rep_fit$res_Sigma, rep_fit$res_Gamma,
                                 rep_fit$K, index = perm)
  rep_fit$res_Beta <- rl$beta
  rep_fit$res_Mu <- rl$mu
  rep_fit$res_Sigma <- rl$sigma
  rep_fit$res_Gamma <- rl$gamma
  rep_fit$inclusion.p <- rep_fit$inclusion.p[, perm, drop = FALSE]
  rep_fit
}


#' Bootstrap replicate statistic for the parallel model
#'
#' \code{boot::boot()}'s \code{statistic} function for a parallel-model
#' bootstrap: refits on the resampled rows and extracts the parameter
#' vector, aligned to \code{template_names} so every replicate's vector has
#' the same layout regardless of which features that replicate selects.
#'
#' @param data The combined data frame passed to \code{boot::boot()}.
#' @param indices Row indices for this replicate.
#' @param model The original fitted model.
#' @param dimG,dimZ,dimCoY,dimCoG Column-block widths within \code{data}
#'   (\code{dimZ} one value per layer).
#' @param Gnames_exposure Exposure column names (excluding covariates).
#' @param template_names,n_template The observed-data statistic's parameter
#'   names/count, from \code{\link{extract_parallel_boot_vector}}.
#' @param prog A \code{progress::progress_bar} to tick.
#' @return A named numeric vector aligned to \code{template_names}, or
#'   all-\code{NA} if the refit failed.
#' @noRd
lucid_par_parallel <- function(data, indices, model, dimG, dimZ, dimCoY, dimCoG,
                               Gnames_exposure, template_names, n_template, prog) {
  prog$tick()

  d <- data[indices, , drop = FALSE]
  col_start <- 1

  G <- as.matrix(d[, col_start:(col_start + dimG - 1), drop = FALSE])
  col_start <- col_start + dimG

  nOmics <- length(dimZ)
  Z <- vector("list", nOmics)
  for(i in seq_len(nOmics)) {
    zdim <- dimZ[i]
    Z[[i]] <- as.matrix(d[, col_start:(col_start + zdim - 1), drop = FALSE])
    col_start <- col_start + zdim
  }

  Y <- as.matrix(d[, col_start, drop = FALSE])
  col_start <- col_start + 1

  CoG <- CoY <- NULL
  if(dimCoG > 0) {
    CoG <- as.matrix(d[, col_start:(col_start + dimCoG - 1), drop = FALSE])
    col_start <- col_start + dimCoG
  }
  if(dimCoY > 0) {
    CoY <- as.matrix(d[, col_start:(col_start + dimCoY - 1), drop = FALSE])
  }

  seed <- sample(1:2000, 1)
  rG <- if(!is.null(model$Rho$Rho_G)) model$Rho$Rho_G else 0
  rMu <- if(!is.null(model$Rho$Rho_Z_Mu)) model$Rho$Rho_Z_Mu else 0
  rCov <- if(!is.null(model$Rho$Rho_Z_Cov)) model$Rho$Rho_Z_Cov else 0
  em_ctrl <- model$em_control
  tol_fit <- if(!is.null(em_ctrl$tol)) em_ctrl$tol else 0.001
  max_itr_fit <- if(!is.null(em_ctrl$max_itr)) em_ctrl$max_itr else 1000
  max_tot_fit <- if(!is.null(em_ctrl$max_tot.itr)) em_ctrl$max_tot.itr else 10000

  invisible(capture.output(try_lucid <- try(est_lucid(
    G = G,
    Z = Z,
    Y = Y,
    CoY = CoY,
    CoG = CoG,
    lucid_model = "parallel",
    family = model$family,
    init_omic.data.model = model$init_omic.data.model,
    K = model$K,
    tol = tol_fit,
    max_itr = max_itr_fit,
    max_tot.itr = max_tot_fit,
    init_impute = model$init_impute,
    init_par = model$init_par,
    useY = model$useY,
    Rho_G = rG,
    Rho_Z_Mu = rMu,
    Rho_Z_Cov = rCov,
    seed = seed
  ), silent = TRUE)))

  if("try-error" %in% class(try_lucid)) {
    par_lucid <- rep(NA_real_, n_template)
    names(par_lucid) <- template_names
  } else {
    try_lucid <- tryCatch(align_replicate_parallel(try_lucid, model, indices), error = function(e) try_lucid)
    # `Gnames_exposure` here is the full beta-column name set (exposures + CoG);
    # `dimG` stays the true exposure count used to carve `d` above.
    par_raw <- extract_parallel_boot_vector(
      model = try_lucid,
      dimG = length(Gnames_exposure),
      dimZ = dimZ,
      Gnames_exposure = Gnames_exposure
    )
    par_lucid <- align_boot_vector(par_raw, template_names = template_names)
  }

  return(par_lucid)
}

#' Reorder a parallel-model replicate fit's per-layer clusters to match the
#' reference fit (see \code{\link{align_replicate_early}})
#' @noRd
align_replicate_parallel <- function(rep_fit, model, indices) {
  K <- as.integer(rep_fit$K)
  nOmics <- length(K)
  ref_pp <- model$inclusion.p
  rep_pp <- rep_fit$inclusion.p
  if (is.null(ref_pp) || is.null(rep_pp) || length(ref_pp) != nOmics) return(rep_fit)

  perms <- lapply(seq_len(nOmics), function(i) {
    P_ref <- tryCatch(as.matrix(ref_pp[[i]])[indices, , drop = FALSE],
                      error = function(e) NULL)
    p <- match_boot_clusters(P_ref, rep_pp[[i]])
    if (is.null(p)) seq_len(K[i]) else as.integer(p)
  })
  if (all(vapply(seq_len(nOmics),
                 function(i) identical(perms[[i]], seq_len(K[i])), logical(1)))) {
    return(rep_fit)
  }

  rl <- relabel_parallel_parameters(
    Beta = rep_fit$res_Beta$Beta,
    Mu = rep_fit$res_Mu,
    Sigma = rep_fit$res_Sigma,
    Delta = rep_fit$res_Gamma$Gamma,
    r = rep_fit$z,
    K = K,
    selectZ = rep_fit$select$selectZ,
    permutations = perms
  )
  rep_fit$res_Beta$Beta <- rl$Beta
  rep_fit$res_Mu <- rl$Mu
  rep_fit$res_Sigma <- rl$Sigma
  rep_fit$res_Gamma <- list(fit = rl$Delta$fit, Gamma = rl$Delta)
  rep_fit$z <- rl$r
  if (!is.null(rl$selectZ)) rep_fit$select$selectZ <- rl$selectZ
  rep_fit$inclusion.p <- lapply(seq_len(nOmics), function(i) {
    rep_pp[[i]][, perms[[i]], drop = FALSE]
  })
  rep_fit
}


#' Ensure a model passed to \code{boot_lucid()} has zero penalty
#'
#' Bootstrap CI is only supported for zero-penalty models: a penalized fit's
#' selection can differ across resamples, which the bootstrap machinery
#' doesn't reconcile. If \code{model} has any nonzero penalty, warns and
#' refits it at zero penalty (\code{\link{refit_bootstrap_zero_penalty}})
#' before returning.
#'
#' @param model A fitted model.
#' @param lucid_model "early", "parallel", or "serial".
#' @param G,Z,Y,CoG,CoY The data \code{model} was fitted on, needed for the
#'   zero-penalty refit if one is required.
#' @return \code{model} unchanged if already zero-penalty, otherwise the
#'   zero-penalty refit.
#' @noRd
normalize_bootstrap_model <- function(model, lucid_model, G, Z, Y, CoG = NULL, CoY = NULL) {
  rho_vals <- get_model_rho_values(model)
  if (all(abs(rho_vals) <= sqrt(.Machine$double.eps))) {
    return(model)
  }

  warning(
    paste0(
      "Bootstrap CI is only supported for zero-penalty models ",
      "(Rho_G = Rho_Z_Mu = Rho_Z_Cov = 0). ",
      "Detected nonzero penalty (Rho_G=", format(rho_vals[1], digits = 6),
      ", Rho_Z_Mu=", format(rho_vals[2], digits = 6),
      ", Rho_Z_Cov=", format(rho_vals[3], digits = 6),
      "). Falling back to a zero-penalty refit before bootstrap."
    ),
    call. = FALSE
  )

  refit_bootstrap_zero_penalty(
    model = model,
    lucid_model = lucid_model,
    G = G,
    Z = Z,
    Y = Y,
    CoG = CoG,
    CoY = CoY
  )
}


#' Extract a model's three penalty values as a numeric vector
#'
#' Robust to missing/non-numeric/non-finite \code{Rho} entries, returning 0
#' for any of those cases.
#'
#' @param model A fitted model.
#' @return A named numeric vector: \code{Rho_G}, \code{Rho_Z_Mu},
#'   \code{Rho_Z_Cov}.
#' @noRd
get_model_rho_values <- function(model) {
  get_scalar <- function(x) {
    if (is.null(x)) return(0)
    v <- suppressWarnings(as.numeric(x))
    v <- v[is.finite(v)]
    if (length(v) == 0) return(0)
    v[1]
  }
  rho <- model$Rho
  c(
    Rho_G = get_scalar(rho$Rho_G),
    Rho_Z_Mu = get_scalar(rho$Rho_Z_Mu),
    Rho_Z_Cov = get_scalar(rho$Rho_Z_Cov)
  )
}


#' Refit a model at zero penalty, reusing its other fitting settings
#'
#' @param model The original (penalized) fitted model, supplying \code{K}
#'   and the fitting controls to reuse.
#' @param lucid_model "early", "parallel", or "serial".
#' @param G,Z,Y,CoG,CoY The data to refit on.
#' @return The zero-penalty refit, or \code{model} unchanged (with a
#'   warning) if the refit fails.
#' @noRd
refit_bootstrap_zero_penalty <- function(model, lucid_model, G, Z, Y, CoG = NULL, CoY = NULL) {
  em_ctrl <- model$em_control
  tol_fit <- if(!is.null(em_ctrl$tol)) em_ctrl$tol else 0.001
  max_itr_fit <- if(!is.null(em_ctrl$max_itr)) em_ctrl$max_itr else 1000
  max_tot_fit <- if(!is.null(em_ctrl$max_tot.itr)) em_ctrl$max_tot.itr else 10000
  seed_fit <- if(!is.null(model$seed) && length(model$seed) > 0 && is.finite(model$seed[1])) {
    as.integer(model$seed[1])
  } else {
    123
  }
  useY_fit <- if(!is.null(model$useY)) model$useY else TRUE

  invisible(capture.output(
    refit_try <- try(
      estimate_lucid(
        lucid_model = lucid_model,
        G = G,
        Z = Z,
        Y = Y,
        CoG = CoG,
        CoY = CoY,
        K = model$K,
        init_omic.data.model = model$init_omic.data.model,
        useY = useY_fit,
        tol = tol_fit,
        max_itr = max_itr_fit,
        max_tot.itr = max_tot_fit,
        Rho_G = 0,
        Rho_Z_Mu = 0,
        Rho_Z_Cov = 0,
        family = model$family,
        seed = seed_fit,
        init_impute = model$init_impute,
        init_par = model$init_par,
        verbose = FALSE
      ),
      silent = TRUE
    )
  ))

  if("try-error" %in% class(refit_try)) {
    stop(
      paste0(
        "Bootstrap CI requires a zero-penalty model and fallback refit failed. ",
        "Please fit estimate_lucid(..., Rho_G = 0, Rho_Z_Mu = 0, Rho_Z_Cov = 0) and retry."
      )
    )
  }
  refit_try
}


#' Recursively detect any deselected feature in a select indicator
#'
#' @param x A logical selection indicator, or a (possibly nested) list of
#'   them (e.g. a parallel model's per-layer \code{selectZ}).
#' @return \code{TRUE} if any leaf has any \code{FALSE} entry, \code{FALSE}
#'   if \code{x} is \code{NULL} or every entry is selected.
#' @noRd
has_unselected_feature <- function(x) {
  if(is.null(x)) {
    return(FALSE)
  }
  if(is.list(x)) {
    return(any(vapply(x, has_unselected_feature, logical(1))))
  }
  x_logical <- as.logical(x)
  any(!x_logical, na.rm = TRUE)
}


#' Extract a fixed-shape parameter vector from a parallel-model fit
#'
#' Concatenates, in order: every layer's non-reference-cluster exposure
#' effects, every layer's cluster means, and the outcome coefficients --
#' this fixed shape (independent of which features happened to be selected
#' in a given replicate) is what \code{\link{align_boot_vector}} aligns
#' every replicate onto.
#'
#' @param model A fitted \code{lucid_parallel} object.
#' @param dimG Number of true exposure columns (excluding covariates).
#' @param dimZ Per-layer number of omics features.
#' @param Gnames_exposure Exposure column names; taken from the model if
#'   \code{NULL}.
#' @return A named numeric vector.
#' @noRd
extract_parallel_boot_vector <- function(model, dimG, dimZ, Gnames_exposure = NULL) {
  K <- as.integer(model$K)
  nOmics <- length(K)

  if(!is.null(Gnames_exposure) && length(Gnames_exposure) == dimG) {
    Gnames <- Gnames_exposure
  } else {
    Gnames <- model$var.names$Gnames
    if(!is.null(Gnames) && length(Gnames) >= dimG) {
      Gnames <- Gnames[seq_len(dimG)]
    } else {
      Gnames <- paste0("G", seq_len(dimG))
    }
  }

  Znames <- model$var.names$Znames
  if(is.null(Znames) || length(Znames) != nOmics) {
    Znames <- lapply(seq_len(nOmics), function(i) paste0("Z", i, "_", seq_len(dimZ[i])))
  }

  beta_all <- NULL
  beta_names <- NULL
  beta_list <- model$res_Beta$Beta
  if(is.null(beta_list) && is.list(model$res_Beta)) {
    beta_list <- model$res_Beta
  }

  for(i in seq_len(nOmics)) {
    beta_mat <- normalize_parallel_beta(beta_i = beta_list[[i]], K_i = K[i], Gnames = Gnames)
    beta_vec <- as.vector(t(beta_mat))
    beta_name_i <- as.vector(sapply(2:K[i], function(k) {
      paste0("Layer", i, ".", colnames(beta_mat), ".cluster", k)
    }))
    beta_all <- c(beta_all, beta_vec)
    beta_names <- c(beta_names, beta_name_i)
  }

  mu_all <- NULL
  mu_names <- NULL
  for(i in seq_len(nOmics)) {
    z_names_i <- Znames[[i]]
    if(is.null(z_names_i) || length(z_names_i) != dimZ[i]) {
      z_names_i <- paste0("Z", i, "_", seq_len(dimZ[i]))
    }
    mu_mat <- normalize_parallel_mu(mu_i = model$res_Mu[[i]], K_i = K[i], Znames_i = z_names_i)
    mu_vec <- as.vector(t(mu_mat))
    mu_name_i <- as.vector(sapply(seq_len(K[i]), function(k) {
      paste0("Layer", i, ".", z_names_i, ".cluster", k)
    }))
    mu_all <- c(mu_all, mu_vec)
    mu_names <- c(mu_names, mu_name_i)
  }

  gamma_all <- extract_parallel_gamma(model)
  gamma_names <- names(gamma_all)
  if(is.null(gamma_names)) {
    gamma_names <- paste0("gamma", seq_along(gamma_all))
  }
  gamma_names <- paste0("Y.", gamma_names)

  par <- c(beta_all, mu_all, as.numeric(gamma_all))
  names(par) <- c(beta_names, mu_names, gamma_names)
  par
}


#' Coerce one layer's exposure coefficients to a fixed-shape matrix
#'
#' A replicate's refit may name, order, or select exposure columns
#' differently than the original model; this maps whatever \code{beta_i}
#' looks like onto a matrix with a known shape and column order (matching by
#' name where possible, falling back to position), so bootstrap replicates
#' can be compared/aligned entry-by-entry.
#'
#' @param beta_i This layer's exposure coefficient matrix (or vector, for a
#'   single non-reference cluster), possibly \code{NULL}.
#' @param K_i This layer's number of clusters.
#' @param Gnames Exposure column names, in the target order.
#' @return A \code{(K_i - 1) x (length(Gnames) + 1)} matrix (intercept plus
#'   exposures), \code{NA} where \code{beta_i} didn't supply a value.
#' @noRd
normalize_parallel_beta <- function(beta_i, K_i, Gnames) {
  K_i <- as.integer(K_i)
  dimG <- length(Gnames)
  target <- matrix(NA_real_, nrow = max(K_i - 1, 0), ncol = dimG + 1)
  colnames(target) <- c("(Intercept)", Gnames)
  if(K_i <= 1 || is.null(beta_i)) {
    return(target)
  }

  if(is.null(dim(beta_i))) {
    beta_i <- matrix(beta_i, nrow = 1)
  } else {
    beta_i <- as.matrix(beta_i)
  }

  # keep non-reference clusters only
  if(nrow(beta_i) == K_i) {
    beta_i <- beta_i[2:K_i, , drop = FALSE]
  } else if(nrow(beta_i) >= (K_i - 1)) {
    beta_i <- beta_i[seq_len(K_i - 1), , drop = FALSE]
  }

  beta_exposure <- matrix(NA_real_, nrow = nrow(beta_i), ncol = dimG + 1)
  colnames(beta_exposure) <- c("(Intercept)", Gnames)

  if(!is.null(colnames(beta_i))) {
    idx_int <- match("(Intercept)", colnames(beta_i))
    if(!is.na(idx_int)) {
      beta_exposure[, 1] <- beta_i[, idx_int, drop = TRUE]
    } else if(ncol(beta_i) >= 1) {
      beta_exposure[, 1] <- beta_i[, 1, drop = TRUE]
    }
    idx <- match(Gnames, colnames(beta_i))
    valid <- which(!is.na(idx))
    if(length(valid) > 0) {
      beta_exposure[, 1 + valid] <- beta_i[, idx[valid], drop = FALSE]
    } else {
      # Fallback for naming mismatches: use positional exposure columns.
      offset <- if(ncol(beta_i) >= 1 && grepl("Intercept", colnames(beta_i)[1], fixed = TRUE)) 2 else 1
      if(offset <= ncol(beta_i)) {
        n_fill <- min(dimG, ncol(beta_i) - offset + 1)
        if(n_fill > 0) {
          beta_exposure[, 1 + seq_len(n_fill)] <- beta_i[, seq.int(offset, length.out = n_fill), drop = FALSE]
        }
      }
    }
  } else {
    if(ncol(beta_i) >= 1) {
      beta_exposure[, 1] <- beta_i[, 1, drop = TRUE]
    }
    if(ncol(beta_i) >= (dimG + 1)) {
      beta_exposure[, 1 + seq_len(dimG)] <- beta_i[, 2:(dimG + 1), drop = FALSE]
    } else {
      n_fill <- min(dimG, ncol(beta_i))
      if(n_fill > 0) {
        beta_exposure[, 1 + seq_len(n_fill)] <- beta_i[, seq_len(n_fill), drop = FALSE]
      }
    }
  }

  n_fill <- min(nrow(target), nrow(beta_exposure))
  if(n_fill > 0) {
    target[seq_len(n_fill), ] <- beta_exposure[seq_len(n_fill), , drop = FALSE]
  }
  target
}


#' Coerce one layer's cluster means to a fixed-shape matrix
#'
#' Same purpose as \code{\link{normalize_parallel_beta}}, for a layer's
#' \code{mu} rather than its exposure coefficients.
#'
#' @param mu_i This layer's cluster-mean matrix, possibly \code{NULL} or
#'   transposed.
#' @param K_i This layer's number of clusters.
#' @param Znames_i This layer's omics feature names, in the target order.
#' @return A \code{K_i x length(Znames_i)} matrix, \code{NA} where
#'   \code{mu_i} didn't supply a value.
#' @noRd
normalize_parallel_mu <- function(mu_i, K_i, Znames_i) {
  K_i <- as.integer(K_i)
  dimZ_i <- length(Znames_i)
  target <- matrix(NA_real_, nrow = K_i, ncol = dimZ_i)
  colnames(target) <- Znames_i
  if(is.null(mu_i)) {
    return(target)
  }

  if(is.null(dim(mu_i))) {
    mu_i <- matrix(mu_i, nrow = K_i)
  } else {
    mu_i <- as.matrix(mu_i)
  }

  if(nrow(mu_i) == dimZ_i && ncol(mu_i) == K_i) {
    mu_i <- t(mu_i)
  }

  n_row <- min(K_i, nrow(mu_i))
  n_col <- min(dimZ_i, ncol(mu_i))
  if(n_row > 0 && n_col > 0) {
    target[seq_len(n_row), seq_len(n_col)] <- mu_i[seq_len(n_row), seq_len(n_col), drop = FALSE]
  }
  target
}


#' Extract outcome coefficients from a parallel-model fit
#'
#' Tries the underlying fitted model object first, falling back to the
#' outcome parameter object's own stored coefficients.
#'
#' @param model A fitted \code{lucid_parallel} object.
#' @return A named numeric coefficient vector.
#' @noRd
extract_parallel_gamma <- function(model) {
  gamma <- NULL
  if(!is.null(model$res_Gamma$fit)) {
    gamma <- try(coef(model$res_Gamma$fit), silent = TRUE)
    if(inherits(gamma, "try-error")) {
      gamma <- NULL
    }
  }
  if(is.null(gamma)) {
    family_parallel <- to_parallel_family(model$family)
    if(family_parallel == "gaussian") {
      gamma <- model$res_Gamma$Gamma$mu
    } else {
      gamma <- model$res_Gamma$fit$coefficients
    }
  }
  gamma
}


#' Align a replicate's parameter vector onto the observed-data template
#'
#' Matches by name when \code{par_raw} has complete names; otherwise falls
#' back to positional alignment (e.g. a refit that dropped names entirely).
#'
#' @param par_raw A replicate's raw parameter vector.
#' @param template_names The observed-data statistic's parameter names, in
#'   the target order.
#' @return A numeric vector of length \code{length(template_names)}, named
#'   \code{template_names}, \code{NA} where \code{par_raw} didn't supply a
#'   value.
#' @noRd
align_boot_vector <- function(par_raw, template_names) {
  par <- rep(NA_real_, length(template_names))
  names(par) <- template_names
  if(length(par_raw) == 0) {
    return(par)
  }

  par_names <- names(par_raw)
  if(!is.null(par_names) && all(!is.na(par_names))) {
    idx <- match(par_names, template_names)
    valid <- which(!is.na(idx))
    # Use name-based alignment only when complete; otherwise use positional fallback.
    if(length(valid) == length(par_raw)) {
      par[idx[valid]] <- as.numeric(par_raw[valid])
      return(par)
    }
  }

  n_fill <- min(length(par), length(par_raw))
  par[seq_len(n_fill)] <- as.numeric(par_raw[seq_len(n_fill)])
  par
}


#' Recursively detect any deselected feature across every serial submodel
#'
#' Unlike \code{\link{has_unselected_feature}}, checks every stage's
#' \code{select}, not just the top-level (stage-1-only) \code{select}
#' field -- so a later stage's own selection is still caught.
#'
#' @param model A fitted \code{lucid_serial} object.
#' @return \code{TRUE} if any stage has any deselected exposure or omics
#'   feature.
#' @noRd
has_unselected_feature_serial <- function(model) {
  if (is.null(model$submodel) || !is.list(model$submodel)) {
    return(FALSE)
  }
  any(vapply(model$submodel, function(sm) {
    if (is.null(sm$select)) {
      return(FALSE)
    }
    has_unselected_feature(sm$select$selectG) || has_unselected_feature(sm$select$selectZ)
  }, logical(1)))
}


#' Flatten a (possibly nested) serial \code{Z} into one matrix for \code{boot::boot()}
#'
#' \code{boot::boot()} resamples rows of a single data frame, but serial
#' \code{Z} can be an arbitrarily nested list of matrices (one leaf per
#' early/parallel stage/layer). Column-binds every leaf matrix and records
#' enough structure (\code{meta}) to reconstruct the original nesting later
#' with \code{\link{restore_serial_Z_from_data}}.
#'
#' @param Z A serial model's (possibly nested list) omics data.
#' @return A list: \code{flat} (one combined matrix) and \code{meta} (the
#'   nesting structure and per-leaf column counts/names).
#' @noRd
flatten_serial_Z <- function(Z) {
  leaf_mats <- list()
  leaf_id <- 0L
  recurse <- function(node) {
    if (is.list(node)) {
      children <- vector("list", length(node))
      for (i in seq_along(node)) {
        children[[i]] <- recurse(node[[i]])
      }
      names(children) <- names(node)
      return(list(kind = "list", children = children, names = names(node)))
    }
    z_mat <- as.matrix(node)
    if (!is.numeric(z_mat)) {
      stop("All serial Z blocks must be numeric.")
    }
    leaf_id <<- leaf_id + 1L
    leaf_mats[[leaf_id]] <<- z_mat
    list(kind = "leaf", leaf_id = leaf_id, ncol = ncol(z_mat), colnames = colnames(z_mat))
  }
  meta <- recurse(Z)
  flat <- do.call(cbind, leaf_mats)
  list(flat = flat, meta = meta)
}


#' Reconstruct a nested serial \code{Z} from flattened bootstrap data
#'
#' Inverse of \code{\link{flatten_serial_Z}}.
#'
#' @param d The combined data frame for one bootstrap replicate.
#' @param col_start First column of \code{d} holding \code{Z} data.
#' @param z_meta The nesting structure from \code{flatten_serial_Z()}.
#' @return A list: \code{Z} (the reconstructed nested structure) and
#'   \code{next_col} (the first column after \code{Z}'s block, for chaining
#'   further column extraction).
#' @noRd
restore_serial_Z_from_data <- function(d, col_start, z_meta) {
  idx <- as.integer(col_start)
  recurse <- function(meta) {
    if (identical(meta$kind, "leaf")) {
      cols <- idx:(idx + meta$ncol - 1L)
      z_mat <- as.matrix(d[, cols, drop = FALSE])
      if (!is.null(meta$colnames) && length(meta$colnames) == meta$ncol) {
        colnames(z_mat) <- meta$colnames
      }
      idx <<- idx + meta$ncol
      return(z_mat)
    }
    out <- vector("list", length(meta$children))
    for (i in seq_along(meta$children)) {
      out[[i]] <- recurse(meta$children[[i]])
    }
    names(out) <- meta$names
    out
  }
  list(Z = recurse(z_meta), next_col = idx)
}


#' Parameter names for each stage's between-stage transition coefficients
#'
#' Stage \code{i}'s transition coefficients (its \eqn{G \to X} model, but
#' fit on the previous stage's cluster/state indicators rather than real
#' exposures) need names derived from the previous stage's cluster
#' structure, which differs for an early vs. parallel previous stage.
#'
#' @param submodels The fitted stage models, in order.
#' @return A list, one character vector of parameter names per stage (empty
#'   for stage 1, which has no previous stage).
#' @noRd
build_serial_transition_labels_boot <- function(submodels) {
  n_stage <- length(submodels)
  out <- vector("list", n_stage)
  out[[1]] <- character(0)
  for (i in seq.int(2, n_stage)) {
    prev <- submodels[[i - 1L]]
    if (inherits(prev, "early_lucid")) {
      k_prev <- as.integer(prev$K)
      if (length(k_prev) > 0 && !is.na(k_prev) && k_prev > 1) {
        out[[i]] <- paste0("Stage", i - 1L, ".cluster", seq.int(2, k_prev))
      } else {
        out[[i]] <- character(0)
      }
    } else if (inherits(prev, "lucid_parallel")) {
      k_prev <- as.integer(prev$K)
      labels <- character(0)
      for (layer_idx in seq_along(k_prev)) {
        if (!is.na(k_prev[layer_idx]) && k_prev[layer_idx] > 1) {
          labels <- c(labels, paste0("Stage", i - 1L, ".Layer", layer_idx,
                                     ".cluster", seq.int(2, k_prev[layer_idx])))
        }
      }
      out[[i]] <- labels
    } else {
      out[[i]] <- character(0)
    }
  }
  out
}


#' Bootstrap parameter vector for one early-integration serial stage
#'
#' Like \code{\link{lucid_early_par_vector}}, but for a serial stage: names
#' the exposure/transition coefficients from \code{transition_labels}
#' rather than assuming real exposure names, and only includes outcome
#' coefficients (\code{gamma}) for the last stage, since only the last
#' stage has a real outcome model.
#'
#' @param stage_model One stage's fitted \code{early_lucid} submodel.
#' @param is_last_stage Whether this is the serial chain's final stage.
#' @param transition_labels Parameter names for this stage's non-reference
#'   incoming cluster/state indicators (from
#'   \code{\link{build_serial_transition_labels_boot}}); falls back to
#'   generic names if not supplied or too short.
#' @return A list: \code{vec} (named numeric parameter vector) and
#'   \code{layout} (component lengths, for
#'   \code{\link{split_serial_boot_ci}} to slice the CI back apart later).
#' @noRd
extract_early_stage_vector <- function(stage_model, is_last_stage, transition_labels = character(0)) {
  K <- as.integer(stage_model$K)
  beta_mat <- as.matrix(stage_model$res_Beta)
  if (is.null(dim(beta_mat))) {
    beta_mat <- matrix(beta_mat, nrow = 1)
  }
  if (nrow(beta_mat) == K) {
    beta_use <- beta_mat[2:K, , drop = FALSE]
    cluster_ids <- 2:K
  } else if (nrow(beta_mat) == (K - 1)) {
    beta_use <- beta_mat
    cluster_ids <- 2:K
  } else {
    beta_use <- beta_mat
    cluster_ids <- seq_len(nrow(beta_use))
  }
  if (nrow(beta_use) == 0) {
    beta_use <- matrix(numeric(0), nrow = 0, ncol = ncol(beta_mat))
    cluster_ids <- integer(0)
  }
  n_feat <- max(0, ncol(beta_use) - 1L)
  if (length(transition_labels) > 0 && n_feat > 0) {
    feat_names <- transition_labels[seq_len(min(length(transition_labels), n_feat))]
    if (length(feat_names) < n_feat) {
      feat_names <- c(feat_names, paste0("PrevStageCluster", seq.int(length(feat_names) + 1L, n_feat)))
    }
    beta_col_names <- c("(Intercept)", feat_names)
  } else {
    beta_col_names <- colnames(beta_use)
    if (is.null(beta_col_names) || length(beta_col_names) != ncol(beta_use)) {
      beta_col_names <- c("(Intercept)", paste0("G", seq_len(n_feat)))
    } else {
      beta_col_names[1] <- "(Intercept)"
    }
  }
  if (ncol(beta_use) == length(beta_col_names)) {
    colnames(beta_use) <- beta_col_names
  }
  beta_vec <- as.numeric(t(beta_use))
  beta_names <- as.vector(sapply(cluster_ids, function(k) paste0(beta_col_names, ".cluster", k)))
  if (length(beta_vec) == length(beta_names)) {
    names(beta_vec) <- beta_names
  }

  mu_mat <- as.matrix(stage_model$res_Mu)
  if (is.null(dim(mu_mat))) {
    mu_mat <- matrix(mu_mat, nrow = K)
  }
  if (nrow(mu_mat) != K && ncol(mu_mat) == K) {
    mu_mat <- t(mu_mat)
  }
  z_names <- stage_model$var.names$Znames
  if (is.null(z_names) || length(z_names) != ncol(mu_mat)) {
    z_names <- paste0("Z", seq_len(ncol(mu_mat)))
  }
  colnames(mu_mat) <- z_names
  mu_vec <- as.numeric(t(mu_mat))
  mu_names <- as.vector(sapply(seq_len(nrow(mu_mat)), function(k) paste0(z_names, ".cluster", k)))
  if (length(mu_vec) == length(mu_names)) {
    names(mu_vec) <- mu_names
  }

  gamma_vec <- numeric(0)
  if (isTRUE(is_last_stage)) {
    gamma_raw <- stage_model$res_Gamma$beta
    gamma_vec <- as.numeric(gamma_raw)
    gamma_names <- names(gamma_raw)
    if (is.null(gamma_names) || length(gamma_names) != length(gamma_vec)) {
      gamma_names <- paste0("gamma", seq_along(gamma_vec))
    }
    names(gamma_vec) <- paste0("Y.", gamma_names)
  }
  vec <- c(beta_vec, mu_vec, gamma_vec)
  list(vec = vec, layout = list(type = "early", n_beta = length(beta_vec),
                                n_mu = length(mu_vec), n_gamma = length(gamma_vec)))
}


#' Bootstrap parameter vector for one parallel-integration serial stage
#'
#' Like \code{\link{extract_parallel_boot_vector}}, but for a serial stage:
#' names the exposure/transition coefficients from \code{transition_labels},
#' and drops the outcome coefficients unless this is the last stage.
#'
#' @param stage_model One stage's fitted \code{lucid_parallel} submodel.
#' @param is_last_stage Whether this is the serial chain's final stage.
#' @param transition_labels Parameter names for this stage's non-reference
#'   incoming cluster/state indicators; falls back to generic names if not
#'   supplied.
#' @return A list: \code{vec} (named numeric parameter vector) and
#'   \code{layout} (component lengths, including per-layer breakdowns, for
#'   \code{\link{split_serial_boot_ci}}).
#' @noRd
extract_parallel_stage_vector <- function(stage_model, is_last_stage, transition_labels = character(0)) {
  K <- as.integer(stage_model$K)
  dimZ <- as.integer(sapply(stage_model$Z, ncol))
  if (length(transition_labels) > 0) {
    g_names <- transition_labels
  } else {
    g_names <- stage_model$var.names$Gnames
    if (is.null(g_names) || length(g_names) == 0) {
      b1 <- stage_model$res_Beta$Beta[[1]]
      p <- if (!is.null(b1)) max(0, ncol(as.matrix(b1)) - 1L) else 0L
      g_names <- paste0("G", seq_len(p))
    }
  }
  dimG_stage <- length(g_names)
  par <- extract_parallel_boot_vector(
    model = stage_model,
    dimG = dimG_stage,
    dimZ = dimZ,
    Gnames_exposure = g_names
  )
  n_beta_layer <- as.integer((K - 1L) * (dimG_stage + 1L))
  n_mu_layer <- as.integer(K * dimZ)
  n_beta <- sum(n_beta_layer)
  n_mu <- sum(n_mu_layer)
  n_keep <- n_beta + n_mu
  if (!isTRUE(is_last_stage)) {
    par <- par[seq_len(n_keep)]
  }
  n_gamma <- max(0L, length(par) - n_keep)
  list(vec = par, layout = list(type = "parallel", n_beta = n_beta, n_mu = n_mu,
                                n_gamma = n_gamma, n_beta_layer = n_beta_layer,
                                n_mu_layer = n_mu_layer))
}


#' Build the observed-data bootstrap template for a serial model
#'
#' Concatenates every stage's parameter vector (via
#' \code{\link{extract_early_stage_vector}}/
#' \code{\link{extract_parallel_stage_vector}}), prefixed with
#' \code{"StageN::"} so names stay unique across stages, and records where
#' each stage's block starts/ends for later slicing.
#'
#' @param model A fitted \code{lucid_serial} object.
#' @return A list: \code{vector} (the full concatenated parameter vector)
#'   and \code{stage_layout} (one element per stage, with that stage's
#'   \code{layout} plus \code{start}/\code{end} indices and names).
#' @noRd
extract_serial_boot_template <- function(model) {
  if (is.null(model$submodel) || !is.list(model$submodel) || length(model$submodel) == 0) {
    stop("Input serial model does not contain valid submodels.")
  }
  submodels <- model$submodel
  n_stage <- length(submodels)
  transition_labels <- build_serial_transition_labels_boot(submodels)
  vec_all <- numeric(0)
  stage_layout <- vector("list", n_stage)
  idx_start <- 1L
  for (i in seq_len(n_stage)) {
    is_last <- (i == n_stage)
    sm <- submodels[[i]]
    stage_obj <- if (inherits(sm, "early_lucid")) {
      extract_early_stage_vector(sm, is_last_stage = is_last, transition_labels = transition_labels[[i]])
    } else if (inherits(sm, "lucid_parallel")) {
      extract_parallel_stage_vector(sm, is_last_stage = is_last, transition_labels = transition_labels[[i]])
    } else {
      stop("Unsupported submodel class in serial bootstrap.")
    }
    stage_vec <- as.numeric(stage_obj$vec)
    local_names <- names(stage_obj$vec)
    if (is.null(local_names) || length(local_names) != length(stage_vec)) {
      local_names <- paste0("param", seq_along(stage_vec))
    }
    prefixed_names <- paste0("Stage", i, "::", local_names)
    names(stage_vec) <- prefixed_names
    idx_end <- idx_start + length(stage_vec) - 1L
    stage_layout[[i]] <- c(stage_obj$layout, list(start = idx_start, end = idx_end,
                                                  local_names = local_names,
                                                  prefixed_names = prefixed_names))
    vec_all <- c(vec_all, stage_vec)
    idx_start <- idx_end + 1L
  }
  list(vector = vec_all, stage_layout = stage_layout)
}


#' Bootstrap replicate statistic for the serial model
#'
#' \code{boot::boot()}'s \code{statistic} function for a serial-model
#' bootstrap: reconstructs the nested \code{Z} (via
#' \code{\link{restore_serial_Z_from_data}}), refits the whole serial
#' chain, and extracts the concatenated parameter vector (via
#' \code{\link{extract_serial_boot_template}}), aligned to
#' \code{template_names}.
#'
#' @param data The combined (flattened) data frame passed to
#'   \code{boot::boot()}.
#' @param indices Row indices for this replicate.
#' @param model The original fitted serial model.
#' @param dimG,dimCoY,dimCoG Column-block widths within \code{data}.
#' @param z_meta The nested-\code{Z} structure from
#'   \code{\link{flatten_serial_Z}}.
#' @param template_names,n_template The observed-data statistic's parameter
#'   names/count.
#' @param prog A \code{progress::progress_bar} to tick.
#' @return A named numeric vector aligned to \code{template_names}, or
#'   all-\code{NA} if the refit failed.
#' @noRd
lucid_par_serial <- function(data, indices, model, dimG, dimCoY, dimCoG,
                             z_meta, template_names, n_template, prog) {
  prog$tick()
  d <- data[indices, , drop = FALSE]
  col_start <- 1L
  G <- as.matrix(d[, col_start:(col_start + dimG - 1L), drop = FALSE])
  col_start <- col_start + dimG
  z_rec <- restore_serial_Z_from_data(d = d, col_start = col_start, z_meta = z_meta)
  Z <- z_rec$Z
  col_start <- z_rec$next_col
  Y <- as.matrix(d[, col_start, drop = FALSE])
  col_start <- col_start + 1L
  CoG <- CoY <- NULL
  if (dimCoG > 0) {
    CoG <- as.matrix(d[, col_start:(col_start + dimCoG - 1L), drop = FALSE])
    col_start <- col_start + dimCoG
  }
  if (dimCoY > 0) {
    CoY <- as.matrix(d[, col_start:(col_start + dimCoY - 1L), drop = FALSE])
  }
  seed <- sample(1:2000, 1)
  rG <- if (!is.null(model$Rho$Rho_G)) model$Rho$Rho_G else 0
  rMu <- if (!is.null(model$Rho$Rho_Z_Mu)) model$Rho$Rho_Z_Mu else 0
  rCov <- if (!is.null(model$Rho$Rho_Z_Cov)) model$Rho$Rho_Z_Cov else 0
  em_ctrl <- model$em_control
  tol_fit <- if (!is.null(em_ctrl$tol)) em_ctrl$tol else 0.001
  max_itr_fit <- if (!is.null(em_ctrl$max_itr)) em_ctrl$max_itr else 1000
  max_tot_fit <- if (!is.null(em_ctrl$max_tot.itr)) em_ctrl$max_tot.itr else 10000
  invisible(capture.output(try_lucid <- try(estimate_lucid(
    G = G, Z = Z, Y = Y, CoY = CoY, CoG = CoG,
    lucid_model = "serial", family = model$family,
    init_omic.data.model = model$init_omic.data.model, K = model$K,
    tol = tol_fit, max_itr = max_itr_fit, max_tot.itr = max_tot_fit,
    init_impute = model$init_impute, init_par = model$init_par,
    useY = model$useY, Rho_G = rG, Rho_Z_Mu = rMu, Rho_Z_Cov = rCov,
    seed = seed
  ), silent = TRUE)))
  if ("try-error" %in% class(try_lucid)) {
    par_lucid <- rep(NA_real_, n_template)
    names(par_lucid) <- template_names
    return(par_lucid)
  }
  try_lucid <- tryCatch(align_replicate_serial(try_lucid, model, indices),
                        error = function(e) try_lucid)
  par_raw <- extract_serial_boot_template(try_lucid)$vector
  align_boot_vector(par_raw = par_raw, template_names = template_names)
}

#' Reorder a serial-model replicate fit's per-stage clusters to match the
#' reference fit, re-referencing each downstream stage's transition
#' coefficients when the upstream stage's reference cluster moves.
#' @noRd
align_replicate_serial <- function(rep_fit, model, indices) {
  ref_sub <- model$submodel
  rep_sub <- rep_fit$submodel
  if (is.null(ref_sub) || is.null(rep_sub) ||
      length(ref_sub) != length(rep_sub)) {
    return(rep_fit)
  }
  n_stage <- length(rep_sub)
  perms <- vector("list", n_stage)

  for (s in seq_len(n_stage)) {
    sm_rep <- rep_sub[[s]]
    sm_ref <- ref_sub[[s]]
    if (inherits(sm_rep, "early_lucid")) {
      perm <- tryCatch({
        P_ref <- as.matrix(sm_ref$inclusion.p)[indices, , drop = FALSE]
        match_boot_clusters(P_ref, sm_rep$inclusion.p)
      }, error = function(e) NULL)
      perms[[s]] <- perm
      if (!is.null(perm) && !identical(as.integer(perm), seq_len(sm_rep$K))) {
        sm_rep <- tryCatch({
          rl <- relabel_early_parameters(sm_rep$res_Beta, sm_rep$res_Mu,
                                         sm_rep$res_Sigma, sm_rep$res_Gamma,
                                         sm_rep$K, index = perm)
          sm_rep$res_Beta <- rl$beta
          sm_rep$res_Mu <- rl$mu
          sm_rep$res_Sigma <- rl$sigma
          sm_rep$res_Gamma <- rl$gamma
          sm_rep$inclusion.p <- sm_rep$inclusion.p[, perm, drop = FALSE]
          sm_rep
        }, error = function(e) { perms[[s]] <<- NULL; rep_sub[[s]] })
      }
    } else if (inherits(sm_rep, "lucid_parallel")) {
      sm_rep <- tryCatch(align_replicate_parallel(sm_rep, sm_ref, indices),
                         error = function(e) sm_rep)
      perms[[s]] <- NULL  # parallel-stage transition re-ref not attempted
    }
    rep_sub[[s]] <- sm_rep
  }

  # Re-reference each non-final stage's downstream transition coefficients for
  # the upstream stage's permutation. Stage s+1's leading non-intercept columns
  # are stage s's non-reference cluster PIPs, so a relabel of stage s is a
  # linear re-parameterisation of those columns (identical algebra to
  # relabel_early_parameters' beta step).
  for (s in seq_len(n_stage - 1L)) {
    perm <- perms[[s]]
    if (is.null(perm) || identical(as.integer(perm), seq_len(length(perm)))) next
    K_prev <- length(perm)
    down <- rep_sub[[s + 1L]]
    if (!inherits(down, "early_lucid")) next
    B <- as.matrix(down$res_Beta)                 # (K_down or K_down-1) x p
    ncols <- ncol(B)
    trans_cols <- seq_len(K_prev - 1L) + 1L       # cols 2 .. K_prev
    if (max(trans_cols) > ncols) next
    for (r in seq_len(nrow(B))) {
      b0 <- B[r, 1L]
      e <- c(0, B[r, trans_cols])                 # per prev-cluster effect, ref = 0
      new_b0 <- b0 + e[perm[1]]
      new_e <- e[perm][-1] - e[perm[1]]           # length K_prev - 1
      B[r, 1L] <- new_b0
      B[r, trans_cols] <- new_e
    }
    down$res_Beta <- B
    rep_sub[[s + 1L]] <- down
    if (!is.null(rep_fit$res_Delta) && length(rep_fit$res_Delta) >= s) {
      rep_fit$res_Delta[[s]] <- B
    }
  }

  # Rebuild the top-level views the extractor and summary read from.
  rep_fit$submodel <- rep_sub
  rep_fit$res_Mu <- lapply(rep_sub, function(x) x$res_Mu)
  rep_fit$res_Sigma <- lapply(rep_sub, function(x) x$res_Sigma)
  rep_fit$inclusion.p <- lapply(rep_sub, function(x) x$inclusion.p)
  rep_fit$res_Beta <- rep_sub[[1L]]$res_Beta
  rep_fit$res_Gamma <- rep_sub[[n_stage]]$res_Gamma
  rep_fit
}


#' Split a serial model's concatenated bootstrap CI back into per-stage tables
#'
#' Inverse of the concatenation \code{\link{extract_serial_boot_template}}
#' performs: slices the full CI matrix back into each stage's
#' \code{beta}/\code{mu}/\code{gamma} tables (further split by layer for a
#' parallel stage).
#'
#' @param ci The full bootstrap CI matrix, rows in the concatenated template
#'   order.
#' @param stage_layout Per-stage layout info from
#'   \code{extract_serial_boot_template()}.
#' @return A list, one element per stage, each a list with \code{beta},
#'   \code{mu}, \code{gamma} (parallel stages: \code{beta}/\code{mu} are
#'   themselves per-layer lists).
#' @noRd
split_serial_boot_ci <- function(ci, stage_layout) {
  out <- vector("list", length(stage_layout))
  for (i in seq_along(stage_layout)) {
    lay <- stage_layout[[i]]
    stage_ci <- ci[lay$start:lay$end, , drop = FALSE]
    rownames(stage_ci) <- lay$local_names
    nb <- lay$n_beta
    nm <- lay$n_mu
    ng <- lay$n_gamma
    beta_ci <- if (nb > 0) stage_ci[seq_len(nb), , drop = FALSE] else NULL
    mu_ci <- if (nm > 0) stage_ci[seq.int(nb + 1L, nb + nm), , drop = FALSE] else NULL
    gamma_ci <- if (ng > 0) stage_ci[seq.int(nb + nm + 1L, nb + nm + ng), , drop = FALSE] else NULL
    if (identical(lay$type, "parallel")) {
      beta_list <- vector("list", length(lay$n_beta_layer))
      mu_list <- vector("list", length(lay$n_mu_layer))
      names(beta_list) <- paste0("Layer", seq_along(beta_list))
      names(mu_list) <- paste0("Layer", seq_along(mu_list))
      if (!is.null(beta_ci)) {
        st <- 1L
        for (j in seq_along(lay$n_beta_layer)) {
          nj <- lay$n_beta_layer[j]
          beta_list[[j]] <- beta_ci[seq.int(st, st + nj - 1L), , drop = FALSE]
          st <- st + nj
        }
      }
      if (!is.null(mu_ci)) {
        st <- 1L
        for (j in seq_along(lay$n_mu_layer)) {
          nj <- lay$n_mu_layer[j]
          mu_list[[j]] <- mu_ci[seq.int(st, st + nj - 1L), , drop = FALSE]
          st <- st + nj
        }
      }
      out[[i]] <- list(beta = beta_list, mu = mu_list, gamma = gamma_ci)
    } else {
      out[[i]] <- list(beta = beta_ci, mu = mu_ci, gamma = gamma_ci)
    }
  }
  out
}


#' Report how many bootstrap replicates produced usable estimates
#'
#' A replicate can fail outright, or -- with missing omics data -- resample
#' too few rows with complete \code{Z} for the model to be estimable. Those
#' replicates previously became silent all-\code{NA} columns that were
#' simply dropped from the interval, quietly shrinking the effective
#' \code{R}. This makes the loss visible: warns on any dropped replicates,
#' warns more strongly (and signals \code{NA} limits) if fewer than
#' \code{min_valid} remain, and gives an advisory warning below the
#' literature-recommended replicate counts (Davison & Hinkley, the reference
#' for Eqs 19-20: R >= 200 for normal intervals, R >= 800 for percentile).
#'
#' @param x A \code{boot::boot} result.
#' @param min_valid Minimum valid replicates required to form an interval.
#' @return A list: \code{R}, \code{n_valid}, \code{enough} (whether
#'   \code{n_valid >= min_valid}).
#' @noRd
boot_replicate_status <- function(x, min_valid = 2L) {
  t <- x$t
  R <- if (is.null(dim(t))) length(t) else nrow(t)
  valid <- if (is.null(dim(t))) is.finite(t) else apply(t, 1, function(r) all(is.finite(r)))
  n_valid <- sum(valid)
  if (n_valid < R) {
    warning(sprintf(
      "%d of %d bootstrap replicates failed to produce finite estimates and were dropped.",
      R - n_valid, R), call. = FALSE)
  }
  if (n_valid < min_valid) {
    warning(sprintf(
      paste0("Only %d valid bootstrap replicate(s) remain (minimum %d needed to ",
             "form an interval); confidence limits are reported as NA. Increase ",
             "R, or check for resamples with too few complete omics rows."),
      n_valid, min_valid), call. = FALSE)
  } else if (n_valid < 200L) {
    # Advisory only.  Davison & Hinkley (the reference for Eqs 19-20) suggest
    # R >= 200 for normal intervals and R >= 800 for percentile intervals.
    # Below that the limits are computable but unstable -- which is a statement
    # about interval QUALITY, not about whether they can be formed, so it must
    # not suppress output.
    warning(sprintf(
      paste0("Only %d bootstrap replicates: confidence limits are unstable. ",
             "R >= 200 is suggested for normal intervals and R >= 800 for ",
             "percentile intervals."),
      n_valid), call. = FALSE)
  }
  list(R = R, n_valid = n_valid, enough = n_valid >= min_valid)
}

#' @title generate bootstrp ci (normal, basic and percentile)
#'
#' @param x an object return by boot function
#' @param conf A numeric scalar between 0 and 1 to specify confidence level(s)
#' of the required interval(s).
#' @param min_valid Minimum number of bootstrap replicates that must yield
#' finite estimates before interval limits can be formed. The default, 2, is the
#' mathematical floor: \code{stats::sd()} needs two finite values and the
#' order-statistic interpolation behind the percentile interval needs two order
#' statistics. Raise it to require more replicates before limits are reported;
#' a small number of replicates produces unstable limits, which is warned about
#' but does not suppress them.
#'
#' @return a matrix, the first column is the point estimate from original model
#'
#' @noRd
gen_ci <- function(x, conf = 0.95, min_valid = 2L) {
  t0 <- x$t0
  status <- boot_replicate_status(x, min_valid = min_valid)
  res_ci <- NULL
  if (!status$enough) {
    # Fewer replicates than can form an interval at all -- the limits are
    # undefined rather than merely imprecise.
    res <- cbind(t0, matrix(NA_real_, length(t0), 4))
    colnames(res) <- c("estimate", "norm_lower", "norm_upper",
                       "perc_lower", "perc_upper")
    attr(res, "boot_status") <- status
    return(res)
  }
  for (i in 1:length(t0)) {
    ci <- try(boot::boot.ci(x,
                            index = i,
                            conf = conf,
                            type = c("norm", "perc")), silent = TRUE)
    if("try-error" %in% class(ci)) {
      temp_ci <- rep(NA_real_, 4)
    } else {
      norm_ci <- if(!is.null(ci$normal) && length(ci$normal) >= 3) ci$normal[2:3] else c(NA_real_, NA_real_)
      perc_ci <- if(!is.null(ci$percent) && length(ci$percent) >= 5) ci$percent[4:5] else c(NA_real_, NA_real_)
      temp_ci <- c(norm_ci, perc_ci)
    }
    res_ci <- rbind(res_ci,
                    temp_ci)
  }
  res <- cbind(t0, res_ci)
  colnames(res) <- c("estimate",
                     "norm_lower", "norm_upper",
                     "perc_lower", "perc_upper")
  attr(res, "boot_status") <- status
  return(res)
}