#' Simulate residual variance for merMod objects
#'
#' Internal helper that reproduces the sigma-drawing logic from the original
#' `predictInterval()` implementation. It returns a numeric vector of length
#' `n.sims` containing the simulated residual standard deviations.
#'
#' @param merMod A merMod object from lme4.
#' @param n.sims Number of simulation draws.
#' @return Numeric vector of length `n.sims`.
#' @keywords internal
simulate_residual_variance <- function(merMod, n.sims) {
  devcomp <- getME(merMod, 'devcomp')
  
  # GLMM: No residual variance to simulate (dispersion = 1 by definition)
  if (devcomp$dims[["GLMM"]] == TRUE) {
    return(NULL)
  }
  
  # NLMM: No residual variance either
  if (devcomp$dims[["NLMM"]] == TRUE) {
    return(NULL)
  }
  
  # LMM: Gamma-based sigma simulation
  if (devcomp$dims[["GLMM"]] == 0 && devcomp$dims[["NLMM"]] == 0) {
    sigma <- sqrt(
      1 / rgamma(
        n.sims,
        shape = 0.5 * residDF.merMod(merMod),
        rate = 0.5 * devcomp$cmp[["pwrss"]]
      )
    )
    return(sigma)
  }
  
  stop("Unknown model type in simulate_residual_variance")
}

#' Simulate fixed-effect predictions
#'
#' Draws `n.sims` samples from the multivariate normal distribution defined by
#' the fixed-effect estimates and their variance-covariance matrix, then
#' multiplies by the model matrix to produce predictions. It also respects
#' `ignore.fixed.terms` and optional intercept variance adjustment.
#'
#' @param merMod A merMod object.
#' @param newdata Data frame containing the prediction covariates.
#' @param n.sims Number of simulations.
#' @param ignore.fixed.terms Vector of fixed-effect names or indices to treat as known.
#' @param fix.intercept.variance Logical flag for intercept variance correction.
#' @param .parallel Logical flag to enable parallel execution via foreach.
#' @return Matrix of dimensions `nrow(newdata)` x `n.sims` containing simulated
#'   fixed-effect predictions.
#' @keywords internal
simulate_fixed_effects <- function(
  merMod,
  newdata,
  n.sims,
  ignore.fixed.terms = NULL,
  fix.intercept.variance = FALSE,
  .parallel = FALSE,
  seed = NULL
) {
  if (!is.null(seed)) set.seed(seed)
  else if (!exists('.Random.seed', envir = .GlobalEnv)) runif(1)
  fe.tmp <- fixef(merMod)
  vcov.tmp <- as.matrix(vcov(merMod))

  # Adjust intercept variance if requested (mirrors original logic)
  if (fix.intercept.variance) {
    rr <- ranef(merMod, condVar = TRUE)
    intercept.var <- vcov.tmp[1, 1]
    groupsizes <- ngrps(merMod)
    for (j in names(groupsizes)) {
      reMatrix <- attr(rr[[j]], which = 'postVar')
      groupVar <- (attr(VarCorr(merMod)[[j]], 'stddev')['(Intercept)'])^2
      extraPrec <- 0
      for (eff in seq_len(dim(reMatrix)[3])) {
        term <- 1 / (reMatrix[1, 1, eff] + groupVar)
        if (term > 0) {
          extraPrec <- extraPrec + term
        } else {
          warning("fix.intercept.variance got negative precision; better turn it off.")
        }
      }
      intercept.var <- intercept.var - 1 / extraPrec
    }
    if (intercept.var < 0) {
      warning("fix.intercept.variance got negative variance; better turn it off.")
    }
    ratio <- intercept.var / vcov.tmp[1, 1]
    prec.tmp <- solve(vcov.tmp)
    prec.tmp[1, 1] <- prec.tmp[1, 1] / ratio
    vcov.tmp[1, ] <- vcov.tmp[1, ] * ratio
    vcov.tmp <- solve(prec.tmp, tol = 1e-50)
  }

  # Inflate precision for ignored fixed terms (effectively zero variance)
  if (!is.null(ignore.fixed.terms)) {
    prec.tmp <- solve(vcov.tmp)
    for (term in ignore.fixed.terms) {
      prec.tmp[term, term] <- prec.tmp[term, term] * 1e15
    }
    vcov.tmp <- solve(prec.tmp, tol = 1e-50)
  }

  # Simulate coefficients - use parallel for large n.sims
  if (n.sims > 2000 || .parallel) {
    if (.parallel) {
      setup_parallel()
    }
    i <- 1:n.sims
    fe_call <- as.call(c(list(quote(foreach::foreach), i = i, .combine = 'rbind')))
    fe <- eval(fe_call)
    betaSim <- foreach::`%dopar%`(
      fe,
      mvtnorm::rmvnorm(n = 1, mean = fe.tmp, sigma = vcov.tmp, method = 'chol')
    )
  } else {
    betaSim <- abind::abind(
      lapply(1:n.sims, function(x) {
        mvtnorm::rmvnorm(n = 1, mean = fe.tmp, sigma = vcov.tmp, method = 'chol')
      }),
      along = 1
    )
  }
  colnames(betaSim) <- names(fe.tmp)
  rownames(betaSim) <- 1:n.sims

  # Build model matrix for newdata
  newdata.modelMatrix <- buildModelMatrix(merMod, newdata = newdata, which = "fixed")

  # Pad betaSim if model matrix has more columns (random slopes not in fixed effects)
  if (ncol(newdata.modelMatrix) > ncol(betaSim)) {
    pad <- matrix(
      rep(0),
      nrow = nrow(betaSim),
      ncol = ncol(newdata.modelMatrix) - ncol(betaSim)
    )
    if (ncol(pad) > 0) {
      message("Fixed effect matrix has been padded with 0 coefficients
            for random slopes not included in the fixed effects and interaction terms.")
    }
    colnames(pad) <- setdiff(colnames(newdata.modelMatrix), colnames(betaSim))
    betaSim <- cbind(betaSim, pad)
  }

  # Align columns between model matrix and coefficient matrix
  keep <- intersect(colnames(newdata.modelMatrix), colnames(betaSim))
  newdata.modelMatrix <- newdata.modelMatrix[, keep, drop = FALSE]
  betaSim <- betaSim[, keep, drop = FALSE]

  # Compute predictions: X %*% t(beta) -> (N x n.sims)
  fixed.xb <- newdata.modelMatrix %*% t(betaSim)
  fixed.xb
}

#' Compute random effect predictions for a single grouping factor
#'
#' Internal helper that multiplies predictor data by simulated random effect
#' coefficients for each grouping level. Handles new levels not present in the
#' original model data by returning zeros for their random effect contributions.
#'
#' @param data data.frame with predictor columns and a grouping factor column.
#'   The last column must be the grouping factor.
#' @param coefs 3D array of simulated coefficients with dimensions
#'   (n.sims x terms x levels).
#' @param group character name of the grouping factor column in data.
#' @return matrix of predicted values with dimensions (nrow(data) x n.sims).
#' @keywords internal
tmp.pred <- function(data, coefs, group) {
  # Identify new levels not in the model data
  new.levels <- unique(
    as.character(data[, group])[
      !as.character(data[, group]) %in% dimnames(coefs)[[3]]
    ]
  )


  if (length(new.levels) > 0) {
    msg <- paste0(
      "     The following levels of ", group, " from newdata \n -- ",
      paste0(new.levels, collapse = ", "),
      " -- are not in the model data. \n",
      "     Currently, predictions for these values are based only on the \n",
      " fixed coefficients and the observation-level error."
    )
    warning(msg, call. = FALSE)
  }

  # Initialize output array

  yhatTmp <- array(data = NA, dim = c(nrow(data), dim(coefs)[1]))
  colIdx <- ncol(data) - 1  # Number of predictor columns (excluding group)

  # Handle coefficient dimension mismatch

  # When predictor columns exceed coefficient terms, expand coefficients
  colLL <- length(1:colIdx)
  if (colLL > dim(coefs)[2]) {
    coefs_new <- array(
      NA,
      dim = c(dim(coefs)[1], colLL, dim(coefs)[3])
    )
    dimnames(coefs_new)[c(1, 3)] <- dimnames(coefs)[c(1, 3)]
    dimnames(coefs_new)[[2]] <- rep(dimnames(coefs)[[2]], dim(coefs_new)[2])
    for (k in 1:colLL) {
      coefs_new[, k, 1:dim(coefs)[3]] <- coefs[, 1, 1:dim(coefs)[3]]
    }
    coefs <- coefs_new
  }

  # Compute predictions for each observation
  for (i in 1:nrow(data)) {
    lvl <- as.character(data[, group][i])
    if (!lvl %in% new.levels) {
      # Known level: multiply predictors by coefficients
      yhatTmp[i, ] <- as.numeric(data[i, 1:colIdx]) %*% t(coefs[, 1:colIdx, lvl])
    } else {
      # New level: zero out the random effect contribution
      yhatTmp[i, ] <- rep(0, colIdx) %*% t(coefs[, 1:colIdx, 1])
    }
  }

  rownames(yhatTmp) <- rownames(data)
  return(yhatTmp)
}

#' Combine a per-term list of postVar arrays into a block-diagonal array
#'
#' When a grouping factor has multiple random-effect term blocks — from the
#' double-bar syntax (`(x + y || g)`), explicit splits (`(1|g) + (0+x|g)`),
#' or mixed correlated + uncorrelated specs — `lme4::ranef(..., condVar =
#' TRUE)` returns `postVar` as a list of arrays, one per term block. Each
#' element has shape `d_m × d_m × n_levels`. This helper concatenates the
#' blocks into a single `d × d × n_levels` array with block-diagonal
#' structure per level (zeros off-diagonal because uncorrelated). The block
#' ordering and column ordering match `colnames(reMeans)` / lme4's `cnms`.
#' @param pv_list List of per-term postVar arrays.
#' @param reMeans The `as.matrix(ranef(m)[[j]])` matrix for the same
#'   grouping factor; used to size and name the combined array.
#' @return 3-D array of shape `ncol(reMeans) × ncol(reMeans) × nrow(reMeans)`.
#' @keywords internal
combine_postvar_blocks <- function(pv_list, reMeans) {
  n_levels <- dim(pv_list[[1L]])[3L]
  d_total  <- ncol(reMeans)
  block_dims <- vapply(pv_list, function(b) dim(b)[1L], integer(1L))
  if (sum(block_dims) != d_total) {
    stop(
      "Unexpected lme4 postVar structure: sum of block dimensions (",
      sum(block_dims), ") does not match ncol(reMeans) (", d_total,
      "). Please report this with your model specification and lme4 version.",
      call. = FALSE
    )
  }
  out <- array(0, dim = c(d_total, d_total, n_levels))
  pos <- 1L
  for (block in pv_list) {
    d_m <- dim(block)[1L]
    idx <- seq.int(pos, length.out = d_m)
    for (lvl in seq_len(n_levels)) {
      out[idx, idx, lvl] <- block[, , lvl]
    }
    pos <- pos + d_m
  }
  dimnames(out) <- list(colnames(reMeans),
                        colnames(reMeans),
                        rownames(reMeans))
  out
}

#' Simulate random‑effect contributions for all grouping factors
#'
#' For each random effect term the function draws `n.sims` samples from the
#' conditional multivariate normal distribution (BLUPs + post‑var). The result is a
#' named list where each element is an array of dimensions `nrow(newdata)` ×
#' `n.sims`.
#' @param merMod A merMod object.
#' @param newdata Data frame containing the prediction covariates.
#' @param n.sims Number of simulations.
#' @param .parallel Logical flag to enable parallel execution via foreach.
#' @param seed Optional random seed for reproducibility.
#' @param new.levels Character; how to treat grouping levels present in
#'   `newdata` but absent from the fitted model. `"zero"` drops the random
#'   effect for such levels; `"draw"` samples each unobserved level's effect
#'   from the estimated random-effect covariance (`VarCorr`).
#' @return List of matrices (rows = observations, cols = simulations).
simulate_random_effects <- function(
  merMod,
  newdata,
  n.sims,
  .parallel = FALSE,
  seed = NULL,
  new.levels = c("zero", "draw")
) {
  new.levels <- match.arg(new.levels, c("zero", "draw"), several.ok = FALSE)
  if (!is.null(seed)) set.seed(seed)
  else if (!exists('.Random.seed', envir = .GlobalEnv)) runif(1)
  rr <- ranef(merMod, condVar = TRUE)
  # Marginal random-effect (co)variance per grouping factor, used to sample
  # effects for unobserved levels when new.levels = "draw".
  reVarCov <- if (identical(new.levels, "draw")) lme4::VarCorr(merMod) else NULL
  re.xb <- vector(getME(merMod, 'n_rfacs'), mode = 'list')
  names(re.xb) <- names(ngrps(merMod))
  newdata.modelMatrix <- buildModelMatrix(model = merMod, newdata = newdata)

  for (j in names(re.xb)) {
    reMeans <- as.matrix(rr[[j]])
    reMatrix <- attr(rr[[j]], which = 'postVar')
    # When a grouping factor has multiple RE term blocks (|| or explicit
    # splits), lme4 returns postVar as a list of per-block arrays. Convert
    # to a single block-diagonal array so downstream indexing works (#118).
    if (is.list(reMatrix)) {
      reMatrix <- combine_postvar_blocks(reMatrix, reMeans)
    }
    # Determine levels to keep based on newdata
    if (j %in% names(newdata)) {
      obslvl <- unique(as.character(newdata[, j]))
      alllvl <- rownames(reMeans)
      keep <- intersect(obslvl, alllvl)
    } else {
      obslvl <- colnames(newdata.modelMatrix)
      alllvl <- rownames(reMeans)
      keep <- intersect(obslvl, alllvl)
    }
    if (length(keep) > 0 && !identical(keep, alllvl)) {
      reMeans <- reMeans[keep, , drop = FALSE]
      dimnames(reMatrix)[[3]] <- alllvl
      reMatrix <- reMatrix[, , keep, drop = FALSE]
    } else if (length(keep) > 0 && identical(keep, alllvl)) {
      dimnames(reMatrix)[[3]] <- alllvl
      reMatrix <- reMatrix[, , keep, drop = FALSE]
    } else {
      reMeans <- reMeans[1, , drop = FALSE]
      reMatrix <- reMatrix[, , 1, drop = FALSE]
    }

    # Simulate draws for each level
    tmpList <- vector('list', nrow(reMeans))
    for (k in seq_len(nrow(reMeans))) {
      meanTmp <- reMeans[k, ]
      matrixTmp <- as.matrix(reMatrix[, , k])
      tmpList[[k]] <- mvtnorm::rmvnorm(
        n = n.sims,
        mean = meanTmp,
        sigma = matrixTmp,
        method = 'chol'
      )
    }
    REcoefs <- sapply(tmpList, identity, simplify = 'array')
    dimnames(REcoefs) <- list(1:n.sims,
                               attr(reMeans, 'dimnames')[[2]],
                               attr(reMeans, 'dimnames')[[1]])

    # new.levels = "draw": for grouping levels in newdata that were not in the
    # fitted model, sample an effect for each from the estimated random-effect
    # covariance (one draw per unobserved level, shared by its observations) and
    # append it to REcoefs. Downstream code then treats these levels as known,
    # so the prediction carries the between-group uncertainty rather than
    # dropping the effect. The default ("zero") skips this and lets such levels
    # fall back to the fixed effects plus residual variation in tmp.pred().
    if (identical(new.levels, "draw") && j %in% names(newdata)) {
      new_lv <- setdiff(unique(as.character(newdata[, j])), alllvl)
      terms_o <- dimnames(REcoefs)[[2]]
      Sigma <- as.matrix(reVarCov[[j]])
      if (length(new_lv) > 0 && all(terms_o %in% rownames(Sigma))) {
        Sigma <- Sigma[terms_o, terms_o, drop = FALSE]
        add <- array(NA_real_, dim = c(n.sims, length(terms_o), length(new_lv)),
                     dimnames = list(1:n.sims, terms_o, new_lv))
        for (k in seq_along(new_lv)) {
          add[, , k] <- mvtnorm::rmvnorm(n.sims, mean = rep(0, length(terms_o)),
                                         sigma = Sigma, method = 'chol')
        }
        REcoefs <- abind::abind(REcoefs, add, along = 3)
      }
    }

    # Build temporary data frame for prediction as in original code
    if (j %in% names(newdata)) {
      newdata.modelMatrix <- as.matrix(newdata.modelMatrix)
      tmp <- cbind(as.data.frame(newdata.modelMatrix), var = newdata[, j])
      tmp <- tmp[, !duplicated(colnames(tmp))]
      keepCols <- names(tmp)[names(tmp) %in% colnames(REcoefs)]
      if (length(keepCols) == 0) {
        keepCols <- grep(dimnames(REcoefs)[[2]], names(tmp), value = TRUE)
      }
      if (length(keepCols) == 0) {
        tmp <- cbind(model.frame(reformulas::subbars(formula(merMod)), data = newdata),
                     var = newdata[, j])
        keepCols <- grep(dimnames(REcoefs)[[2]], names(tmp), value = TRUE)
      }
      if (length(keepCols) == 0) {
        # Add in an intercept for RE purposes
        tmp <- cbind(as.data.frame(newdata.modelMatrix), var = newdata[, j])
        tmp <- tmp[, !duplicated(colnames(tmp))]
        tmp <- cbind(data.frame(1), tmp)
        names(tmp)[1] <- "(Intercept)"
        keepCols <- "(Intercept)"
      }
      tmp <- tmp[, c(keepCols, "var"), drop = FALSE]
      tmp[, "var"] <- as.character(tmp[, "var"])
      colnames(tmp)[which(names(tmp) == "var")] <- names(newdata[, j, drop = FALSE])
      if (all(grepl(":", keepCols))) {
        # Strip out the interaction after
        keepCols <- unique(gsub("(.*):.*", "\\1", keepCols))
      }
    } else {
      # If newdata.modelMatrix is still sparse at this point, we need to convert it safely
      if (is(newdata.modelMatrix, 'dgCMatrix')) {
        newdata.modelMatrix <- as.matrix(newdata.modelMatrix)
        tmp <- as.data.frame(newdata.modelMatrix)
      } else {
        tmp <- as.data.frame(newdata.modelMatrix)
      }
      tmp <- tmp[, !duplicated(colnames(tmp))]
      # Assign the grouping level for each row based on which model-matrix
      # column is active. Rows whose interaction level is not among the
      # observed `keep` levels are all-zero across tmp[keep]; max.col() would
      # otherwise pick a column at random (an all-zero row is a tie), silently
      # borrowing an unrelated level's random effect and making the result
      # depend on the RNG seed (#124). Detect those rows and assign NA so that
      # tmp.pred() treats them as new levels and zeroes the random-effect
      # contribution, matching the documented out-of-sample behavior.
      keepMat <- as.matrix(tmp[keep])
      activeLevel <- names(tmp[keep])[max.col(keepMat, ties.method = "first")]
      activeLevel[rowSums(keepMat != 0) == 0] <- NA_character_
      tmp$var <- activeLevel
      keepCols <- names(tmp)[names(tmp) %in% dimnames(REcoefs)[[2]]]
      tmp <- tmp[, c(keepCols, "var"), drop = FALSE]
      tmp[, "var"] <- as.character(tmp[, "var"])
      colnames(tmp)[which(names(tmp) == "var")] <- j
    }

    # Compute contribution matrix using the internal prediction helper
    # Use parallel execution for large datasets
    if (nrow(tmp) > 1000 || .parallel) {
      if (requireNamespace("foreach", quietly = TRUE)) {
        if (.parallel) {
          setup_parallel()
        }
        # Split data into chunks of 500 rows
        tmp2 <- split(tmp, (seq_len(nrow(tmp)) - 1) %/% 500)
        tmp2 <- tmp2[vapply(tmp2, function(x) nrow(x) > 0, logical(1))]

        fe_call <- as.call(c(
          list(quote(foreach::foreach), i = seq_along(tmp2), .combine = 'rbind')
        ))
        fe <- eval(fe_call)

        re.xb[[j]] <- foreach::`%dopar%`(
          fe,
          tmp.pred(data = tmp2[[i]], coefs = REcoefs[, keepCols, , drop = FALSE], group = j)
        )
      } else {
        warning("foreach package is unavailable, parallel computing not available")
        re.xb[[j]] <- tmp.pred(data = tmp, coefs = REcoefs[, keepCols, , drop = FALSE], group = j)
      }
    } else {
      re.xb[[j]] <- tmp.pred(data = tmp, coefs = REcoefs[, keepCols, , drop = FALSE], group = j)
    }
  }
  re.xb
}

#' Combine fixed, random and residual components into a prediction array
#'
#' @param fixed_mat Matrix of dimensions `nrow(newdata)` × `n.sims`.
#' @param random_list List returned by `simulate_random_effects()`.
#' @param sigma_vec Residual standard deviations (output of `simulate_residual_variance`).
#' @param include.resid.var Logical, whether to add residual variance.
#' @param which Character: "full", "fixed", "random" or "all".
#' @return Array `nrow(newdata)` × `n.sims` (or list when `which="all"`).
#' @keywords internal
combine_components <- function(
  fixed_mat,
  random_list,
  sigma_vec,
  include.resid.var = TRUE,
  which = c('full', 'fixed', 'random', 'all'),
  family = NULL,
  link = NULL,
  weights = NULL,
  use.probability = FALSE
) {
  which <- match.arg(which)
  re.xb <- random_list
  re.xb$fixed <- fixed_mat

  if (which == 'fixed') {
    yhat <- re.xb[['fixed']]
    if (include.resid.var) {
      N <- nrow(yhat)
      
if (use.probability) {
         yhat <- simulate_glmm_response(yhat, family = family, link = link,
                                         weights = weights, use.probability = TRUE)
       } else if (!is.null(sigma_vec)) {
         yhat <- abind::abind(
           lapply(seq_len(length(sigma_vec)), function(x) rnorm(N, yhat[, x], sigma_vec[x])),
           along = 2
         )
       }
     }
     return(yhat)
   }
 
   if (which == 'random') {
    re.xb[['fixed']] <- NULL
    yhat <- Reduce('+', re.xb)
    if (include.resid.var) {
      N <- nrow(yhat)
      
if (use.probability) {
         yhat <- simulate_glmm_response(yhat, family = family, link = link,
                                         weights = weights, use.probability = TRUE)
       } else if (!is.null(sigma_vec)) {
         yhat <- abind::abind(
           lapply(seq_len(length(sigma_vec)), function(x) rnorm(N, yhat[, x], sigma_vec[x])),
           along = 2
         )
       }
     }
     return(yhat)
   }
 
   if (which == 'full') {
    yhat <- Reduce('+', re.xb)
    if (include.resid.var) {
      N <- nrow(yhat)
      
      if (use.probability) {
        yhat <- simulate_glmm_response(yhat, family = family, link = link,
                                        weights = weights, use.probability = TRUE)
      } else if (!is.null(sigma_vec)) {
        yhat <- abind::abind(
          lapply(seq_len(length(sigma_vec)), function(x) rnorm(N, yhat[, x], sigma_vec[x])),
          along = 2
        )
      }
    }
    return(yhat)
  }

  if (which == 'all') {
    yhat <- Reduce('+', re.xb)
    
    for (i in seq_along(re.xb)) {
      if (include.resid.var) {
        N <- nrow(yhat)
        
if (!is.null(family) && family %in% c("binomial", "poisson", "Gamma")) {
           re.xb[[i]] <- simulate_glmm_response(re.xb[[i]], family = family, link = link,
                                                 weights = weights, use.probability = use.probability)
         } else if (!is.null(sigma_vec)) {
           re.xb[[i]] <- abind::abind(
             lapply(seq_len(length(sigma_vec)), function(x) rnorm(N, re.xb[[i]][, x], sigma_vec[x])),
             along = 2
           )
         }
      }
    }
    
    if (include.resid.var) {
      N <- nrow(yhat)
      
      if (use.probability) {
        yhat <- simulate_glmm_response(yhat, family = family, link = link,
                                        weights = weights, use.probability = TRUE)
      } else if (!is.null(sigma_vec)) {
        yhat <- abind::abind(
          lapply(seq_len(length(sigma_vec)), function(x) rnorm(N, yhat[, x], sigma_vec[x])),
          along = 2
        )
      }
    }
    return(list(yhat = yhat, components = re.xb))
  }
}

#' Summarise prediction simulations into intervals
#'
#' Takes the simulated prediction array and computes summary statistics
#' (fit, lower bound, upper bound) based on the specified statistic type
#' and confidence level. Handles probability transformation and the
#' `which = "all"` case with component-wise summaries.
#'
#' @param yhat_arr matrix of simulated predictions with dimensions (N x n.sims),
#'   or for `which = "all"`, the combined yhat from `combine_components()`.
#' @param level numeric confidence level (e.g., 0.8 for 80
#' @param stat.type character either "median" or "mean" for the central estimate.
#' @param predict.type character either "linear.prediction" or "probability".
#' @param N integer number of observations in the original newdata.
#' @param merMod merMod object, required when `predict.type = "probability"` to
#'   access the link function.
#' @param which.eff character the effect type: "full", "fixed", "random", or "all".
#' @param pi.comps list of component prediction matrices (from `combine_components()`
#'   when `which = "all"`), or NULL for other cases.
#' @return data.frame with columns `fit`, `upr`, `lwr`. For `which = "all"`,
#'   includes additional columns `effect` and `obs`.
#' @keywords internal
summarise_predictions <- function(
  yhat_arr,
  level,
  stat.type,
  predict.type,
  N,
  merMod = NULL,
  which.eff = "full",
  pi.comps = NULL,
  is.glmm.with.response = FALSE
) {
  # Calculate quantile bounds

upCI <- 1 - ((1 - level) / 2)
  loCI <- (1 - level) / 2

  # Initialize output data frame
  outs <- data.frame(
    fit = rep(NA_real_, N),
    upr = rep(NA_real_, N),
    lwr = rep(NA_real_, N)
  )

  # Compute summary statistics based on stat.type
  if (stat.type == "median") {
    outs[, 1:3] <- t(apply(yhat_arr, 1, quantile,
                           prob = c(0.5, upCI, loCI),
                           na.rm = TRUE))
  } else if (stat.type == "mean") {
    outs$fit <- apply(yhat_arr, 1, mean, na.rm = TRUE)
    outs[, 2:3] <- t(apply(yhat_arr, 1, quantile,
                           prob = c(upCI, loCI),
                           na.rm = TRUE))
  }

  # Apply link function for probability predictions
  if (predict.type == "probability") {
    if (is.null(merMod)) {
      stop("merMod object required for probability predictions", call. = FALSE)
    }
    
    # Skip linkinv for ALL GLMM families with response-scale simulations
    devcomp <- getME(merMod, "devcomp")
    is_glmm <- (devcomp$dims[["GLMM"]] == TRUE)
    
    if (!(is.glmm.with.response && is_glmm)) {
      if (nrow(outs) == 1) {
        outs <- t(apply(outs, 2, merMod@resp$family$linkinv))
      } else {
        outs <- apply(outs, 2, merMod@resp$family$linkinv)
      }
    }
    
    outs <- as.data.frame(outs)
    names(outs) <- c("fit", "upr", "lwr")
  }

  # Handle which = "all" case with component breakdown
  if (which.eff == "all" && !is.null(pi.comps)) {
    # Process each component
    for (i in seq_along(pi.comps)) {
      if (stat.type == "median") {
        pi.comps[[i]] <- t(apply(pi.comps[[i]], 1, quantile,
                                 prob = c(0.5, upCI, loCI),
                                 na.rm = TRUE))
        pi.comps[[i]] <- as.data.frame(pi.comps[[i]])
        names(pi.comps[[i]]) <- c("fit", "upr", "lwr")
      } else if (stat.type == "mean") {
        tmp <- pi.comps[[i]]
        pi.comps[[i]] <- data.frame(
          fit = rep(NA_real_, N),
          upr = NA_real_,
          lwr = NA_real_
        )
        pi.comps[[i]]$fit <- apply(tmp, 1, mean, na.rm = TRUE)
        pi.comps[[i]][, 2:3] <- t(apply(tmp, 1, quantile,
                                        prob = c(upCI, loCI),
                                        na.rm = TRUE))
      }
      # Apply link function for probability predictions
      if (predict.type == "probability") {
        devcomp <- getME(merMod, "devcomp")
        is_binomial_glmm <- (devcomp$dims[["GLMM"]] == TRUE &&
                             merMod@resp$family$family == "binomial")
        
        if (!(is.glmm.with.response && is_binomial_glmm)) {
          pi.comps[[i]] <- apply(pi.comps[[i]], 2, merMod@resp$family$linkinv)
        }
        pi.comps[[i]] <- as.data.frame(pi.comps[[i]])
        names(pi.comps[[i]]) <- c("fit", "upr", "lwr")
      }
    }

    # Bind components together with effect labels
    componentOut <- dplyr::bind_rows(pi.comps, .id = "effect")
    outs <- cbind(data.frame(effect = "combined"), outs)
    outs <- suppressWarnings(dplyr::bind_rows(outs, componentOut))
    outs$obs <- rep(seq_len(N), nrow(outs) %/% N)
  }

  as.data.frame(outs)
}

#' Inverse link function for GLMM families
#'
#' Internal helper that applies the inverse link function to convert linear
#' predictors to distribution parameters.
#'
#' @param eta Linear predictor (fixed + random effects).
#' @param family Character string specifying the distribution family.
#' @param link Character string specifying the link function.
#' @return Vector of distribution parameters (prob, lambda, or rate).
#' @keywords internal
inverse_link <- function(eta, family, link) {
  if (family == "binomial") {
    if (link == "logit") return(1 / (1 + exp(-eta)))
    if (link == "probit") return(pnorm(eta))
    if (link == "cloglog") return(1 - exp(-exp(eta)))
  }
  
  if (family == "poisson") {
    if (link == "log") return(exp(eta))
    if (link == "sqrt") return(pmax(eta, 0)^2)
    if (link == "identity") return(pmax(eta, 0))
  }
  
  if (family == "Gamma") {
    eps <- .Machine$double.xmin
    if (link == "inverse") return(1 / pmax(eta, eps))
    if (link == "identity") return(pmax(eta, eps))
    if (link == "log") return(exp(eta))
  }
  
  stop("Unsupported family/link combination: ", family, "/", link)
}

#' Simulate GLMM responses from conditional distribution
#'
#' Internal helper that simulates responses from the conditional distribution
#' of Y|b (given random effects) for GLMMs. This implements the correct
#' theoretical approach where residual variance is inherent in the distribution.
#'
#' @param yhat Matrix of linear predictors (fixed + random effects).
#' @param family Character string specifying the distribution family.
#' @param link Character string specifying the link function.
#' @param weights Optional weights (for binomial with n-trial parameter).
#' @param use.probability For binomial, return proportions (0-1) instead of counts?
#' @return Matrix of simulated responses with same dimensions as yhat.
#' @keywords internal
simulate_glmm_response <- function(yhat, family, link, weights = NULL, use.probability = FALSE) {
  N <- nrow(yhat)
  n_sims <- ncol(yhat)
  
  if (family == "binomial") {
    if (is.null(weights)) weights <- rep(1, N)
    if (length(weights) == 1) weights <- rep(weights, N)
    
    eta <- as.matrix(yhat)
    prob <- inverse_link(eta, family, link)
    
    result <- matrix(nrow = N, ncol = n_sims)
    for (i in 1:N) {
      if (use.probability) {
        result[i, ] <- rbinom(n_sims, size = weights[i], prob = prob[i, ]) / weights[i]
      } else {
        result[i, ] <- rbinom(n_sims, size = weights[i], prob = prob[i, ])
      }
    }
    return(result)
  }
  
  if (family == "poisson") {
    lambda <- inverse_link(as.matrix(yhat), family, link)
    result <- matrix(nrow = N, ncol = n_sims)
    for (i in 1:N) {
      result[i, ] <- rpois(n_sims, lambda = lambda[i])
    }
    return(result)
  }
  
  if (family == "Gamma") {
    rate <- inverse_link(as.matrix(yhat), family, link)
    shape <- 1
    result <- matrix(nrow = N, ncol = n_sims)
    for (i in 1:N) {
      result[i, ] <- rgamma(n_sims, shape = shape, rate = rate[i])
    }
    return(result)
  }
  
  stop("Unsupported family for GLMM simulation: ", family)
}
