R/plot_diagnostics.R

In tulpa: Templated Unified Library for Posterior Approximation in Bayesian Hierarchical Models

#' Diagnostic plotting functions for tulpa models
#'
#' @description
#' Visual diagnostic tools for MCMC convergence assessment. All functions
#' provide base R fallbacks when ggplot2/bayesplot are unavailable.
#'
#' @name plot_diagnostics
#' @keywords internal
NULL

# Consistent color scheme for diagnostic plots
tulpa_diag_colors <- list(

  good = "#2E7D32",      # Green - pass

  warn = "#F57F17",      # Yellow/orange - caution
  bad = "#C62828",       # Red - fail
  divergent = "#D32F2F", # Red for divergent transitions
  chain = c("#1b9e77", "#d95f02", "#7570b3", "#e7298a",
            "#66a61e", "#e6ab02", "#a6761d", "#666666")
)

# Internal theme: minimal without grid lines (cleaner diagnostic plots)
theme_tulpa <- function() {
  ggplot2::theme_minimal() +
    ggplot2::theme(panel.grid = ggplot2::element_blank())
}


#' Plot Rhat convergence diagnostic
#'
#' @description
#' Creates a visual display of Rhat values for all parameters, with color-coding
#' to highlight convergence issues. Rhat values > 1.01 indicate potential
#' convergence problems.
#'
#' @param fit A `tulpa_fit` object.
#' @param threshold Rhat threshold for warnings (default: 1.01).
#' @param pars Character vector of parameter names to include. If NULL (default),
#'   includes all main parameters (excludes high-dimensional RE/spatial).
#'
#' @return A ggplot object (if ggplot2 available) or base R plot (invisible).
#'
#' @details
#' Color coding:
#' - Green: Rhat < 1.01 (converged)
#' - Yellow: 1.01 <= Rhat < 1.05 (borderline)
#' - Red: Rhat >= 1.05 (not converged)
#'
#' @examples
#' # Diagnostic plots require a fitted tulpa model
#' # See tulpa() examples for fitting models
#'
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_rhat(fit)
#' }
#'
#' @seealso [plot_ess()], [diagnostic_summary()], [diagnostics()]
#' @export
plot_rhat <- function(fit, threshold = 1.01, pars = NULL) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  if (!.tulpa_is_chain(fit)) {
    message(.tulpa_non_chain_msg(fit))
    return(invisible(NULL))
  }

  # Get diagnostics

  diag <- diagnostics(fit, pars = pars)

  # Filter to main parameters if not specified
  if (is.null(pars)) {
    main_pars <- select_main_params(diag$parameter)
    diag <- diag[diag$parameter %in% main_pars, ]
  }

  # Remove NA values

  diag <- diag[!is.na(diag$rhat), ]

  if (nrow(diag) == 0) {
    message("No Rhat values available for selected parameters")
    return(invisible(NULL))
  }

  # Assign colors based on thresholds
  diag$color <- ifelse(
    diag$rhat < threshold, tulpa_diag_colors$good,
    ifelse(diag$rhat < 1.05, tulpa_diag_colors$warn, tulpa_diag_colors$bad)
  )

  # Sort by Rhat value

  diag <- diag[order(diag$rhat, decreasing = TRUE), ]
  diag$parameter <- factor(diag$parameter, levels = diag$parameter)

  # Plot with ggplot2 if available

  if (requireNamespace("ggplot2", quietly = TRUE)) {
    p <- ggplot2::ggplot(diag, ggplot2::aes(x = .data$rhat, y = .data$parameter)) +
      ggplot2::geom_point(ggplot2::aes(color = .data$color), size = 3) +
      ggplot2::geom_vline(xintercept = threshold, linetype = "dashed",
                          color = tulpa_diag_colors$warn, linewidth = 0.8) +
      ggplot2::geom_vline(xintercept = 1, linetype = "solid",
                          color = "gray50", linewidth = 0.5) +
      ggplot2::scale_color_identity() +
      ggplot2::labs(
        title = "Rhat Convergence Diagnostic",
        x = expression(hat(R)),
        y = NULL,
        caption = paste0("Threshold: ", threshold, " (dashed line)")
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"),
        panel.grid.major.y = ggplot2::element_line(color = "gray90")
      )

    return(p)
  }

  # Base R fallback
  plot_rhat_base(diag, threshold)
  invisible(NULL)
}


#' Base R Rhat plot
#' @noRd
plot_rhat_base <- function(diag, threshold) {
  n_pars <- nrow(diag)
  old_par <- graphics::par(mar = c(4, 8, 3, 1))
  on.exit(graphics::par(old_par))

  graphics::plot(
    diag$rhat, seq_len(n_pars),
    xlim = c(min(0.99, min(diag$rhat)), max(1.1, max(diag$rhat))),
    ylim = c(0.5, n_pars + 0.5),
    pch = 19, col = diag$color, cex = 1.5,
    xlab = expression(hat(R)),
    ylab = "",
    yaxt = "n",
    main = "Rhat Convergence Diagnostic"
  )

  graphics::axis(2, at = seq_len(n_pars), labels = diag$parameter, las = 2, cex.axis = 0.8)
  graphics::abline(v = threshold, lty = 2, col = tulpa_diag_colors$warn, lwd = 2)
  graphics::abline(v = 1, lty = 1, col = "gray50")
}


#' Plot effective sample size diagnostic
#'
#' @description
#' Creates a visual display of effective sample size (ESS) for all parameters,
#' expressed as a ratio of ESS to total samples. Low ESS indicates high
#' autocorrelation.
#'
#' @param fit A `tulpa_fit` object.
#' @param type Type of ESS: "bulk" (default) or "tail".
#' @param threshold Minimum acceptable ESS (default: 400).
#' @param pars Character vector of parameter names to include.
#'
#' @return A ggplot object (if ggplot2 available) or base R plot (invisible).
#'
#' @details
#' ESS/iter ratio interpretation:
#' - Green: ESS >= threshold
#' - Yellow: threshold/2 <= ESS < threshold
#' - Red: ESS < threshold/2
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_ess(fit)
#' plot_ess(fit, type = "tail")
#' }
#'
#' @seealso [plot_rhat()], [diagnostic_summary()], [diagnostics()]
#' @export
plot_ess <- function(fit, type = c("bulk", "tail"), threshold = 400, pars = NULL) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  type <- match.arg(type)

  if (!.tulpa_is_chain(fit)) {
    message(.tulpa_non_chain_msg(fit))
    return(invisible(NULL))
  }

  # Get diagnostics
  diag <- diagnostics(fit, pars = pars)

  # Filter to main parameters if not specified
  if (is.null(pars)) {
    main_pars <- select_main_params(diag$parameter)
    diag <- diag[diag$parameter %in% main_pars, ]
  }

  # Select ESS column
  ess_col <- if (type == "bulk") "ess_bulk" else "ess_tail"
  diag$ess <- diag[[ess_col]]

  # Remove NA values
  diag <- diag[!is.na(diag$ess), ]

  if (nrow(diag) == 0) {
    message("No ESS values available for selected parameters")
    return(invisible(NULL))
  }

  # Assign colors based on thresholds
  diag$color <- ifelse(
    diag$ess >= threshold, tulpa_diag_colors$good,
    ifelse(diag$ess >= threshold / 2, tulpa_diag_colors$warn, tulpa_diag_colors$bad)
  )

  # Sort by ESS (ascending so worst are at top)
  diag <- diag[order(diag$ess), ]
  diag$parameter <- factor(diag$parameter, levels = diag$parameter)

  # Plot with ggplot2 if available
  if (requireNamespace("ggplot2", quietly = TRUE)) {
    p <- ggplot2::ggplot(diag, ggplot2::aes(x = .data$ess, y = .data$parameter)) +
      ggplot2::geom_point(ggplot2::aes(color = .data$color), size = 3) +
      ggplot2::geom_vline(xintercept = threshold, linetype = "dashed",
                          color = tulpa_diag_colors$warn, linewidth = 0.8) +
      ggplot2::scale_color_identity() +
      ggplot2::labs(
        title = paste0("Effective Sample Size (", type, ")"),
        x = "ESS",
        y = NULL,
        caption = paste0("Threshold: ", threshold, " (dashed line)")
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"),
        panel.grid.major.y = ggplot2::element_line(color = "gray90")
      )

    return(p)
  }

  # Base R fallback
  plot_ess_base(diag, threshold, type)
  invisible(NULL)
}


#' Base R ESS plot
#' @noRd
plot_ess_base <- function(diag, threshold, type) {
  n_pars <- nrow(diag)
  old_par <- graphics::par(mar = c(4, 8, 3, 1))
  on.exit(graphics::par(old_par))

  graphics::plot(
    diag$ess, seq_len(n_pars),
    xlim = c(0, max(diag$ess) * 1.1),
    ylim = c(0.5, n_pars + 0.5),
    pch = 19, col = diag$color, cex = 1.5,
    xlab = "ESS",
    ylab = "",
    yaxt = "n",
    main = paste0("Effective Sample Size (", type, ")")
  )

  graphics::axis(2, at = seq_len(n_pars), labels = diag$parameter, las = 2, cex.axis = 0.8)
  graphics::abline(v = threshold, lty = 2, col = tulpa_diag_colors$warn, lwd = 2)
}


#' Plot autocorrelation functions
#'
#' @description
#' Creates autocorrelation function (ACF) plots for selected parameters.
#' High autocorrelation indicates slow mixing and low effective sample size.
#'
#' @param fit A `tulpa_fit` object.
#' @param pars Character vector of parameter names. If NULL, selects worst-mixing
#'   parameters based on ESS.
#' @param lags Maximum number of lags to compute (default: 25).
#' @param n_pars Maximum number of parameters to plot (default: 6).
#'
#' @return A ggplot object (if ggplot2 available) or base R plot (invisible).
#'
#' @details
#' Ideal ACF plots show rapid decay to zero. Slow decay indicates high
#' autocorrelation, which reduces effective sample size and may indicate
#' poor mixing.
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_acf(fit)
#' plot_acf(fit, lags = 10)
#' }
#'
#' @seealso [plot_ess()], [diagnostics()]
#' @export
plot_acf <- function(fit, pars = NULL, lags = 25, n_pars = 6) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }
  if (!.tulpa_is_chain(fit)) {
    message(.tulpa_non_chain_msg(fit))
    return(invisible(NULL))
  }

  # Get draws (pooled 2-D with colnames; handles a 3-D [iter, chain, param] fit)
  draws <- .tulpa_pooled_draws(fit)
  if (is.null(draws)) {
    stop("No draws available in fit object", call. = FALSE)
  }

  # Select parameters
  if (is.null(pars)) {
    # Select worst ESS parameters
    diag <- diagnostics(fit)
    main_pars <- select_main_params(diag$parameter)
    diag <- diag[diag$parameter %in% main_pars, ]
    diag <- diag[order(diag$ess_bulk), ]
    pars <- head(diag$parameter, n_pars)
  } else {
    pars <- grep_params(pars, colnames(draws))
    pars <- head(pars, n_pars)
  }

  if (length(pars) == 0) {
    message("No parameters selected for ACF plot")
    return(invisible(NULL))
  }

  # Compute ACF for each parameter
  acf_data <- lapply(pars, function(p) {
    x <- draws[, p]
    acf_obj <- stats::acf(x, lag.max = lags, plot = FALSE)
    data.frame(
      parameter = p,
      lag = acf_obj$lag[-1],
      acf = acf_obj$acf[-1],
      stringsAsFactors = FALSE
    )
  })
  acf_df <- do.call(rbind, acf_data)

  # Plot with ggplot2 if available
  if (requireNamespace("ggplot2", quietly = TRUE)) {
    # Compute confidence bounds (approximate 95% CI)
    n_samples <- nrow(draws)
    ci_bound <- qnorm(0.975) / sqrt(n_samples)

    p <- ggplot2::ggplot(acf_df, ggplot2::aes(x = .data$lag, y = .data$acf)) +
      ggplot2::geom_hline(yintercept = 0, color = "gray50") +
      ggplot2::geom_hline(yintercept = c(-ci_bound, ci_bound),
                          linetype = "dashed", color = "blue", alpha = 0.5) +
      ggplot2::geom_segment(ggplot2::aes(xend = .data$lag, yend = 0),
                            color = tulpa_diag_colors$chain[1], linewidth = 0.8) +
      ggplot2::geom_point(color = tulpa_diag_colors$chain[1], size = 2) +
      ggplot2::facet_wrap(~ parameter, scales = "free_y") +
      ggplot2::labs(
        title = "Autocorrelation Functions",
        x = "Lag",
        y = "ACF"
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"),
        strip.text = ggplot2::element_text(face = "bold")
      )

    return(p)
  }

  # Base R fallback
  plot_acf_base(draws, pars, lags)
  invisible(NULL)
}


#' Base R ACF plot
#' @noRd
plot_acf_base <- function(draws, pars, lags) {
  n_pars <- length(pars)
  n_col <- min(3, n_pars)
  n_row <- ceiling(n_pars / n_col)

  old_par <- graphics::par(mfrow = c(n_row, n_col), mar = c(4, 4, 3, 1))
  on.exit(graphics::par(old_par))

  for (p in pars) {
    x <- draws[, p]
    stats::acf(x, lag.max = lags, main = p, col = tulpa_diag_colors$chain[1])
  }
}


#' Plot bivariate parameter posteriors (pairs plot)
#'
#' @description
#' Creates a pairs plot showing bivariate relationships between parameters.
#' Divergent transitions (if present) are highlighted to help identify
#' problematic posterior regions.
#'
#' @param fit A `tulpa_fit` object.
#' @param pars Character vector of parameter names. If NULL, selects main
#'   variance parameters.
#' @param highlight_divergent Logical; highlight divergent transitions in red
#'   (default: TRUE).
#' @param n_pars Maximum number of parameters (default: 5).
#' @param alpha Point transparency (default: 0.3).
#'
#' @return A ggplot object (if ggplot2/GGally available) or base R plot.
#'
#' @details
#' Pairs plots help identify:
#' - Strong correlations between parameters (potential non-identifiability)
#' - Multimodality
#' - Regions where divergences cluster (indicating problematic geometry)
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_pairs(fit)
#' plot_pairs(fit, n_pars = 2)
#' }
#'
#' @seealso [plot_divergences()], [diagnostics()]
#' @export
plot_pairs <- function(fit, pars = NULL, highlight_divergent = TRUE,
                       n_pars = 5, alpha = 0.3) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  # Get draws (pooled 2-D with colnames; handles a 3-D [iter, chain, param] fit
  # and fills any NULL/blank parameter name positionally, gcol33/tulpa#780)
  draws <- .tulpa_pooled_draws(fit)
  if (is.null(draws)) {
    message(.tulpa_no_draws_note(fit, "plot_pairs"))
    return(invisible(NULL))
  }

  # Select parameters
  if (is.null(pars)) {
    # Select variance parameters and key fixed effects
    all_pars <- colnames(draws)
    var_pars <- grep("^(sigma|phi|tau|rho)", all_pars, value = TRUE)
    # Fixed effects are always the leading columns of `draws` by convention
    # (`.pg_as_chain`, the tgmrf joint samplers, tulpa()'s own layout), so
    # position is tried first: a NULL or blank fixed name that
    # `.tulpa_pooled_draws()` has since filled in positionally (e.g. "param1")
    # no longer name-matches `fit$fixed_names` itself (gcol33/tulpa#780). Fall
    # back to the name match, then a generic beta / intercept pattern --
    # covering ratio's beta_num / beta_denom -- when `n_fixed` is absent.
    beta_pars <- if (!is.null(fit$n_fixed) && fit$n_fixed >= 1L &&
                     fit$n_fixed <= length(all_pars)) {
      all_pars[seq_len(fit$n_fixed)]
    } else {
      intersect(fit$fixed_names %||% character(0), all_pars)
    }
    if (!length(beta_pars)) {
      beta_pars <- grep("^(beta[_[]|\\(Intercept\\))", all_pars, value = TRUE)
    }
    pars <- c(beta_pars, var_pars)
    pars <- head(pars, n_pars)
  } else {
    pars <- grep_params(pars, colnames(draws))
    pars <- head(pars, n_pars)
  }

  if (length(pars) < 2) {
    message("Need at least 2 parameters for pairs plot")
    return(invisible(NULL))
  }

  # Subset draws
  draws_subset <- as.data.frame(draws[, pars, drop = FALSE])

  # Add divergence indicator if available
  has_divergent <- FALSE
  if (highlight_divergent && (fit$backend %||% "") == "hmc") {
    div_idx <- .tulpa_divergent_idx(fit)
    if (!is.null(div_idx) && length(div_idx) > 0) {
      draws_subset$divergent <- FALSE
      draws_subset$divergent[div_idx] <- TRUE
      has_divergent <- any(draws_subset$divergent)
    }
  }

  # Use bayesplot for pairs if available. Fall back to the base R renderer if
  # bayesplot is absent or rejects the inputs.
  if (requireNamespace("bayesplot", quietly = TRUE)) {
    draws_array <- get_draws_array(fit)$draws
    draws_array <- draws_array[, , pars, drop = FALSE]

    np <- if (has_divergent) {
      .tulpa_divergent_np(fit, dim(draws_array)[1L], dim(draws_array)[2L])
    } else {
      NULL
    }

    p <- tryCatch(
      bayesplot::mcmc_pairs(
        draws_array,
        pars = pars,
        np = np,
        off_diag_args = list(size = 0.5, alpha = alpha)
      ),
      error = function(e) NULL
    )
    if (!is.null(p)) return(p)
  }

  # Base R pairs plot fallback
  plot_pairs_base(draws_subset, pars, has_divergent, alpha)
  invisible(NULL)
}


# The one read of a fit's divergence record, so the count and the row indices
# cannot disagree on the same object: `plot_divergences()` calls both in
# sequence and used to be able to believe divergences exist while reporting
# their indices unavailable (gcol33/tulpa#840).
#
# Returns `n` (0 when nothing is recorded) and `idx`, the row indices where the
# record is a per-draw flag vector or an index vector, NULL where it is only a
# count. Every sampler in this package writes the top-level `$divergent`
# vector, row-aligned with `$log_prob` / `$chain_id` (gcol33/tulpa#783); the
# `$diagnostics$*` shapes are the ones a consumer package may assemble instead.
# A vector-shaped record is preferred over a count because it answers both
# questions; among vectors the engine's own field is read first.
.tulpa_divergence_record <- function(fit) {
  d <- fit$diagnostics
  flags <- function(v) {
    i <- which(as.logical(v))
    list(n = length(i), idx = i)
  }
  if (!is.null(fit$divergent))   return(flags(fit$divergent))
  if (!is.null(d$divergent))     return(flags(d$divergent))
  if (!is.null(d$divergent_idx)) {
    i <- as.integer(d$divergent_idx)
    return(list(n = length(i), idx = i))
  }
  if (!is.null(d$n_divergent))   return(list(n = as.integer(d$n_divergent), idx = NULL))
  if (!is.null(fit$n_divergent)) return(list(n = as.integer(fit$n_divergent), idx = NULL))
  list(n = 0L, idx = NULL)
}


# Divergent-transition row indices, or NULL when the fit records only a count.
.tulpa_divergent_idx <- function(fit) {
  .tulpa_divergence_record(fit)$idx
}


#' Energy Bayesian Fraction of Missing Information (E-BFMI)
#'
#' Computes E-BFMI (Betancourt 2016) from the per-iteration Hamiltonian
#' energies an HMC or NUTS sampler records: the variance of the energy
#' transition between successive iterations over the variance of the marginal
#' energy. Each chain is diffed separately, since a transition across a chain
#' boundary compares unrelated trajectories, and the lowest chain value is
#' reported: one badly exploring chain is enough to distrust a fit. Values
#' below 0.3 indicate that momentum resampling moves the sampler too little
#' between energy level sets.
#'
#' @param energy Numeric vector of per-iteration energies.
#' @param chain Chain index of each entry of `energy`. `NULL` (default), or a
#'   vector of the wrong length, treats `energy` as one chain.
#'
#' @return A list with `value` (the lowest per-chain E-BFMI, `NA` when no
#'   chain has two finite, non-constant energies), `per_chain`,
#'   `energy_centred` (energies centred within each chain) and `energy_diff`
#'   (within-chain transitions), the last two pooled across chains for
#'   plotting.
#'
#' @examples
#' set.seed(1)
#' energy <- c(cumsum(rnorm(200, sd = 0.3)), rnorm(200, sd = 1))
#' energy_bfmi(energy, chain = rep(1:2, each = 200))$per_chain
#'
#' @export
energy_bfmi <- function(energy, chain = NULL) {
  energy <- as.numeric(energy)
  if (is.null(chain) || length(chain) != length(energy)) {
    chain <- rep(1L, length(energy))
  }
  by_chain <- split(energy, chain)
  diffs <- lapply(by_chain, function(e) if (length(e) > 1L) diff(e) else numeric(0))
  per_chain <- vapply(seq_along(by_chain), function(i) {
    e <- by_chain[[i]]; ed <- diffs[[i]]
    if (length(ed) < 1L || !is.finite(stats::var(e)) || stats::var(e) == 0) {
      return(NA_real_)
    }
    stats::var(ed) / stats::var(e)
  }, numeric(1))
  names(per_chain) <- names(by_chain)
  list(
    value = if (all(is.na(per_chain))) NA_real_ else min(per_chain, na.rm = TRUE),
    per_chain = per_chain,
    energy_centred = unlist(lapply(by_chain, function(e) e - mean(e)),
                            use.names = FALSE),
    energy_diff = unlist(diffs, use.names = FALSE)
  )
}

# E-BFMI of a tulpa fit, read off the top-level `$energy` vector every HMC/NUTS
# sampler writes and its `$chain_id` (gcol33/tulpa#783).
.tulpa_energy_bfmi <- function(fit) {
  energy <- fit$energy
  if (is.null(energy) || length(energy) < 10) {
    return(list(value = NA_real_, energy_diff = numeric(0)))
  }
  energy_bfmi(energy, fit$chain_id)
}


# Build a bayesplot `np`-style data frame flagging divergent transitions, keyed
# by (Iteration, Chain) to align with an [iter, chain, param] draws array.
# Divergence row indices come from the top-level `fit$divergent` vector every
# sampler writes, mapped from the pooled chain-major draws onto the per-chain
# (iteration, chain) grid.
.tulpa_divergent_np <- function(fit, n_iter, n_chain) {
  np <- data.frame(
    Iteration = rep(seq_len(n_iter), times = n_chain),
    Chain     = rep(seq_len(n_chain), each = n_iter),
    Parameter = factor("divergent__"),
    Value     = 0
  )
  div_idx <- .tulpa_divergent_idx(fit)
  if (is.null(div_idx) || length(div_idx) == 0L) return(np)

  cid <- fit$chain_id
  if (!is.null(cid) && length(cid) >= max(div_idx)) {
    chain <- cid[div_idx]
    iter  <- ave(seq_along(cid), cid, FUN = seq_along)[div_idx]
  } else {
    per   <- if (is.matrix(fit$draws)) nrow(fit$draws) %/% n_chain else n_iter
    chain <- (div_idx - 1L) %/% per + 1L
    iter  <- (div_idx - 1L) %% per + 1L
  }
  hit <- chain >= 1L & chain <= n_chain & iter >= 1L & iter <= n_iter
  for (k in which(hit)) {
    np$Value[np$Chain == chain[k] & np$Iteration == iter[k]] <- 1
  }
  np
}


#' Base R pairs plot
#' @noRd
plot_pairs_base <- function(draws_df, pars, has_divergent, alpha) {
  # Subset to just the parameters
  plot_data <- draws_df[, pars, drop = FALSE]

  # Color vector
  if (has_divergent) {
    cols <- ifelse(draws_df$divergent,
                   grDevices::adjustcolor(tulpa_diag_colors$divergent, alpha.f = 0.8),
                   grDevices::adjustcolor("black", alpha.f = alpha))
    pch <- ifelse(draws_df$divergent, 19, 1)
  } else {
    cols <- grDevices::adjustcolor("black", alpha.f = alpha)
    pch <- 1
  }

  graphics::pairs(
    plot_data,
    pch = pch,
    col = cols,
    main = "Parameter Pairs Plot",
    upper.panel = function(x, y, ...) {
      graphics::points(x, y, ...)
      # Add correlation
      r <- cor(x, y)
      usr <- graphics::par("usr")
      graphics::text(mean(usr[1:2]), mean(usr[3:4]),
                     sprintf("r = %.2f", r), cex = 1.2, col = "gray40")
    }
  )
}


#' Plot divergent transitions
#'
#' @description
#' Creates visualizations to investigate divergent transitions. Parallel
#' coordinates and scatter plots highlight where in parameter space
#' divergences occur.
#'
#' @param fit A `tulpa_fit` object (HMC backend).
#' @param pars Character vector of parameter names. If NULL, uses variance
#'   parameters.
#' @param type Plot type: "parcoord" (parallel coordinates) or "scatter".
#'
#' @return A ggplot object or base R plot (invisible).
#'
#' @details
#' Divergent transitions indicate regions of high posterior curvature that
#' the sampler cannot efficiently explore. Common causes:
#' - Very narrow funnels (hierarchical models)
#' - Strong correlations
#' - Multi-modality
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_divergences(fit)
#' plot_divergences(fit, type = "scatter")
#' }
#'
#' @seealso [plot_pairs()], [n_divergent()]
#' @export
plot_divergences <- function(fit, pars = NULL, type = c("parcoord", "scatter")) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  if ((fit$backend %||% "") != "hmc") {
    message("Divergence plots only available for HMC backend")
    return(invisible(NULL))
  }

  type <- match.arg(type)

  # Check for divergent transitions
  n_div <- n_divergent(fit)
  if (n_div == 0) {
    message("No divergent transitions to plot")
    return(invisible(NULL))
  }

  div_idx <- .tulpa_divergent_idx(fit)
  if (is.null(div_idx) || length(div_idx) == 0) {
    message("Divergent indices not available")
    return(invisible(NULL))
  }

  # Get draws
  draws <- .fit_draws(fit)
  if (is.null(draws)) {
    message(.tulpa_no_draws_note(fit, "plot_divergences"))
    return(invisible(NULL))
  }

  # Select parameters
  if (is.null(pars)) {
    all_pars <- colnames(draws)
    var_pars <- grep("^(sigma|phi|tau|rho)", all_pars, value = TRUE)
    pars <- head(var_pars, 6)
  } else {
    pars <- grep_params(pars, colnames(draws))
    pars <- head(pars, 6)
  }

  if (length(pars) < 2) {
    message("Need at least 2 parameters for divergence plot")
    return(invisible(NULL))
  }

  # Create data frame with divergence indicator
  draws_df <- as.data.frame(draws[, pars, drop = FALSE])
  draws_df$divergent <- FALSE
  draws_df$divergent[div_idx] <- TRUE

  if (type == "parcoord") {
    plot_divergences_parcoord(draws_df, pars)
  } else {
    plot_divergences_scatter(draws_df, pars)
  }
}


#' Parallel coordinates plot for divergences
#' @noRd
plot_divergences_parcoord <- function(draws_df, pars) {
  # Scale parameters to [0, 1] for comparison
  scaled <- draws_df[, pars, drop = FALSE]
  for (p in pars) {
    rng <- range(scaled[[p]])
    if (diff(rng) > 0) {
      scaled[[p]] <- (scaled[[p]] - rng[1]) / diff(rng)
    } else {
      scaled[[p]] <- 0.5
    }
  }
  scaled$divergent <- draws_df$divergent

  if (requireNamespace("ggplot2", quietly = TRUE)) {
    # Reshape for ggplot
    scaled$id <- seq_len(nrow(scaled))
    long_data <- reshape(
      scaled,
      direction = "long",
      varying = pars,
      v.names = "value",
      timevar = "parameter",
      times = pars
    )
    long_data$parameter <- factor(long_data$parameter, levels = pars)

    # Plot non-divergent first, then divergent on top
    p <- ggplot2::ggplot(long_data, ggplot2::aes(x = .data$parameter, y = .data$value,
                                                  group = .data$id)) +
      ggplot2::geom_line(data = long_data[!long_data$divergent, ],
                         alpha = 0.05, color = "gray50") +
      ggplot2::geom_line(data = long_data[long_data$divergent, ],
                         alpha = 0.7, color = tulpa_diag_colors$divergent,
                         linewidth = 0.8) +
      ggplot2::labs(
        title = "Divergent Transitions (Parallel Coordinates)",
        subtitle = paste0(sum(draws_df$divergent), " divergent transitions (red)"),
        x = "Parameter",
        y = "Scaled Value"
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"),
        axis.text.x = ggplot2::element_text(angle = 45, hjust = 1)
      )

    return(p)
  }

  # Base R fallback
  plot_divergences_parcoord_base(scaled, pars, draws_df$divergent)
}


#' Base R parallel coordinates
#' @noRd
plot_divergences_parcoord_base <- function(scaled, pars, divergent) {
  n_pars <- length(pars)
  n_samples <- nrow(scaled)

  old_par <- graphics::par(mar = c(6, 4, 4, 2))
  on.exit(graphics::par(old_par))

  graphics::plot(
    NULL, NULL,
    xlim = c(1, n_pars),
    ylim = c(0, 1),
    xlab = "", ylab = "Scaled Value",
    xaxt = "n",
    main = "Divergent Transitions (Parallel Coordinates)"
  )
  graphics::axis(1, at = seq_len(n_pars), labels = pars, las = 2, cex.axis = 0.7)

  # Draw non-divergent lines
  for (i in which(!divergent)) {
    graphics::lines(seq_len(n_pars), as.numeric(scaled[i, pars]),
                    col = grDevices::adjustcolor("gray50", alpha.f = 0.02))
  }

  # Draw divergent lines on top
  for (i in which(divergent)) {
    graphics::lines(seq_len(n_pars), as.numeric(scaled[i, pars]),
                    col = grDevices::adjustcolor(tulpa_diag_colors$divergent, alpha.f = 0.7),
                    lwd = 1.5)
  }
}


#' Scatter plot matrix for divergences
#' @noRd
plot_divergences_scatter <- function(draws_df, pars) {
  # Use first 2-4 parameters for scatter matrix
  pars <- head(pars, 4)

  if (requireNamespace("ggplot2", quietly = TRUE) && length(pars) >= 2) {
    # Simple 2D scatter of first two parameters
    p <- ggplot2::ggplot(draws_df, ggplot2::aes(x = .data[[pars[1]]], y = .data[[pars[2]]])) +
      ggplot2::geom_point(data = draws_df[!draws_df$divergent, ],
                          alpha = 0.1, size = 1, color = "gray50") +
      ggplot2::geom_point(data = draws_df[draws_df$divergent, ],
                          alpha = 0.8, size = 2, color = tulpa_diag_colors$divergent) +
      ggplot2::labs(
        title = "Divergent Transitions",
        subtitle = paste0(sum(draws_df$divergent), " divergent (red)"),
        x = pars[1],
        y = pars[2]
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold")
      )

    return(p)
  }

  # Base R fallback
  plot_pairs_base(draws_df, pars, has_divergent = TRUE, alpha = 0.2)
}


#' Plot energy diagnostic (E-BFMI)
#'
#' @description
#' Creates overlaid histograms of marginal energy and energy transition,
#' along with the E-BFMI statistic. Low E-BFMI indicates poor exploration
#' of the posterior.
#'
#' @param fit A `tulpa_fit` object (HMC backend).
#'
#' @return A ggplot object (if ggplot2 available) or base R plot.
#'
#' @details
#' E-BFMI (Energy Bayesian Fraction of Missing Information) compares the
#' distribution of energy levels to energy transitions. Values below 0.3
#' indicate the sampler may not be exploring the full posterior.
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_energy(fit)
#' }
#'
#' @seealso [diagnostic_summary()], [check_diagnostics()]
#' @export
plot_energy <- function(fit) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  if ((fit$backend %||% "") != "hmc") {
    message("Energy plots only available for HMC backend")
    return(invisible(NULL))
  }

  energy <- fit$energy
  if (is.null(energy) || length(energy) < 10) {
    message("Energy values not available in fit object")
    return(invisible(NULL))
  }

  # E-BFMI, per chain (a cross-chain diff would compare unrelated
  # trajectories); .tulpa_energy_bfmi() reports the worst chain.
  bfmi <- .tulpa_energy_bfmi(fit)
  e_bfmi <- bfmi$value
  energy_diff <- bfmi$energy_diff

  # Create data for plotting
  n <- length(energy)
  nd <- length(energy_diff)
  plot_data <- data.frame(
    value = c(energy, energy_diff),
    type = rep(c("Marginal E", "E transition"), c(n, nd))
  )

  # Status message
  status <- if (is.na(e_bfmi)) "E-BFMI unavailable" else
    if (e_bfmi < 0.3) "WARNING: Low E-BFMI" else "OK"

  if (requireNamespace("ggplot2", quietly = TRUE)) {
    p <- ggplot2::ggplot(plot_data, ggplot2::aes(x = .data$value, fill = .data$type)) +
      ggplot2::geom_histogram(alpha = 0.6, position = "identity", bins = 50) +
      ggplot2::scale_fill_manual(values = c("Marginal E" = tulpa_diag_colors$chain[1],
                                            "E transition" = tulpa_diag_colors$chain[2])) +
      ggplot2::labs(
        title = "HMC Energy Diagnostic",
        subtitle = sprintf("E-BFMI = %.3f (%s)", e_bfmi, status),
        x = "Energy",
        y = "Count",
        fill = NULL
      ) +
      theme_tulpa() +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"),
        legend.position = "bottom"
      )

    return(p)
  }

  # Base R fallback
  plot_energy_base(energy, energy_diff, e_bfmi, status)
  invisible(NULL)
}


#' Base R energy plot
#' @noRd
plot_energy_base <- function(energy, energy_diff, e_bfmi, status) {
  old_par <- graphics::par(mar = c(4, 4, 4, 2))
  on.exit(graphics::par(old_par))

  # Compute breaks
  all_vals <- c(energy, energy_diff)
  breaks <- seq(min(all_vals), max(all_vals), length.out = 50)

  # Marginal energy histogram
  h1 <- graphics::hist(energy, breaks = breaks, plot = FALSE)
  h2 <- graphics::hist(energy_diff, breaks = breaks, plot = FALSE)

  graphics::plot(
    h1,
    col = grDevices::adjustcolor(tulpa_diag_colors$chain[1], alpha.f = 0.6),
    main = sprintf("HMC Energy Diagnostic\nE-BFMI = %.3f (%s)", e_bfmi, status),
    xlab = "Energy",
    ylim = c(0, max(c(h1$counts, h2$counts)) * 1.1)
  )
  graphics::plot(
    h2,
    col = grDevices::adjustcolor(tulpa_diag_colors$chain[2], alpha.f = 0.6),
    add = TRUE
  )
  graphics::legend(
    "topright",
    legend = c("Marginal E", "E transition"),
    fill = c(tulpa_diag_colors$chain[1], tulpa_diag_colors$chain[2]),
    bty = "n"
  )
}


#' Comprehensive diagnostic summary
#'
#' @description
#' Provides a comprehensive diagnostic report for a tulpa model, combining
#' convergence metrics, divergence information, and actionable recommendations.
#'
#' @param fit A `tulpa_fit` object.
#' @param quiet Logical; if TRUE, suppress printed output (default: FALSE).
#'
#' @return A list with class `tulpa_diagnostic_summary` containing:
#' \describe{
#'   \item{status}{Overall status: "PASS", "WARN", or "FAIL"}
#'   \item{n_divergent}{Number of divergent transitions}
#'   \item{divergent_pct}{Percentage of divergent transitions}
#'   \item{worst_rhat}{Data frame of parameters with worst Rhat}
#'   \item{worst_ess}{Data frame of parameters with worst ESS}
#'   \item{e_bfmi}{E-BFMI value (HMC only)}
#'   \item{pareto_k, quad_ess}{approximation fits only: the outer PSIS k-hat,
#'     or the grid quadrature ESS when no k-hat was produced}
#'   \item{hyper_share_min, hyper_share_max}{approximation fits only: the range
#'     over coefficients of the share of the fixed-effect marginal variance the
#'     hyperparameter integration contributed, `between / (within + between)`
#'     under the law of total variance over the outer grid. Near zero means the
#'     reported interval is what conditioning on a single hyperparameter point
#'     would have given. Reported without a threshold and never scored: a small
#'     share is not by itself a miscalibrated interval. `NA` when the fit
#'     retains no per-cell pieces to decompose}
#'   \item{pareto_k_declined}{approximation fits only, and only when there is no
#'     k-hat: WHY -- `"not_requested"` and
#'     `"unguessable_axis: <axis>"` are benign or permanent,
#'     `"degenerate_proposal"` and `"grid_too_small"` are signals about the fit,
#'     and `"internal_inconsistency"` is an engine bug and raises the status to
#'     `"WARN"`}
#'   \item{inner_skew_max, inner_skew_declined}{approximation fits only: the
#'     largest scored inner-Laplace `|gamma_3|`, or why nothing was scored
#'     (`"coupled_arm"` marks arms the inner layer could
#'     score neither per observation nor through the cell tensor)}
#'   \item{axis_fields_dropped}{data frame of grid axes the fit's own resolved
#'     path could not read and dropped as engine defaults:
#'     one row per field, with the block, family, path and the axis that path
#'     integrated instead. Absent whenever every supplied axis was used}
#'   \item{recommendations}{Character vector of recommendations}
#' }
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' ds <- diagnostic_summary(fit)
#' print(ds)
#' }
#'
#' @seealso [check_diagnostics()], [diagnostics()], [plot_diagnostics()]
#' @export
diagnostic_summary <- function(fit, quiet = FALSE) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  backend <- fit$backend %||% "unknown"
  recommendations <- character(0)
  status <- "PASS"

  # Initialize result
  result <- list(
    backend = backend,
    n_divergent = 0,
    divergent_pct = 0,
    worst_rhat = NULL,
    worst_ess = NULL,
    e_bfmi = NA,
    pareto_k = NA_real_,
    s2z_log_det_fallback_cells = 0L,
    s2z_field_indefinite_cells = 0L,
    unconverged_cells = 0L,
    outer_regime = NA_character_,
    interval_read = NA_character_,
    interval_design_mass = NA_real_,
    quad_ess = NA_real_,
    hyper_share_min = NA_real_,
    hyper_share_max = NA_real_,
    recommendations = character(0),
    status = "PASS"
  )

  # Get MCMC diagnostics (chain fits only; the registry's `emits` property
  # decides, so every chain backend -- hmc, gibbs, mala, imh_laplace, ess, ...
  # -- is covered, and no approximation fit gets a vacuous Rhat/ESS table).
  if (.tulpa_is_chain(fit)) {
    diag <- tryCatch(diagnostics(fit), error = function(e) NULL)

    if (!is.null(diag)) {
      # Filter to main parameters
      main_pars <- select_main_params(diag$parameter)
      diag <- diag[diag$parameter %in% main_pars, ]

      # Worst Rhat
      bad_rhat <- diag[!is.na(diag$rhat) & diag$rhat > 1.01, ]
      if (nrow(bad_rhat) > 0) {
        bad_rhat <- bad_rhat[order(bad_rhat$rhat, decreasing = TRUE), ]
        result$worst_rhat <- head(bad_rhat[, c("parameter", "rhat")], 5)
        recommendations <- c(recommendations,
          "Rhat > 1.01: Run more iterations or chains")
        status <- "WARN"
      }

      # Worst ESS
      low_ess <- diag[!is.na(diag$ess_bulk) & diag$ess_bulk < 400, ]
      if (nrow(low_ess) > 0) {
        low_ess <- low_ess[order(low_ess$ess_bulk), ]
        result$worst_ess <- head(low_ess[, c("parameter", "ess_bulk", "ess_tail")], 5)
        recommendations <- c(recommendations,
          "ESS < 400: Run more iterations or use thinning")
        status <- "WARN"
      }
    }
  }

  # Check divergences (HMC only)
  if (backend == "hmc") {
    n_div <- n_divergent(fit)
    result$n_divergent <- n_div

    # Calculate percentage
    n_samples <- nrow(fit$draws) %||% 0
    if (n_samples > 0) {
      result$divergent_pct <- 100 * n_div / n_samples
    }

    if (n_div > 0) {
      if (result$divergent_pct > 10) {
        status <- "FAIL"
        recommendations <- c(recommendations,
          "Many divergences: Check model specification or reparameterize")
      } else if (result$divergent_pct > 1) {
        status <- "WARN"
        recommendations <- c(recommendations,
          "Some divergences: Use plot_divergences() to investigate")
      } else {
        recommendations <- c(recommendations,
          "Few divergences: May be acceptable, but verify results")
      }
    }

    # E-BFMI
    bfmi <- .tulpa_energy_bfmi(fit)
    if (!is.na(bfmi$value)) {
      result$e_bfmi <- bfmi$value

      if (result$e_bfmi < 0.3) {
        status <- if (status == "FAIL") "FAIL" else "WARN"
        recommendations <- c(recommendations,
          "Low E-BFMI: Sampler may not be exploring the full posterior")
      }
    }
  }

  # Approximation (non-chain) fits: chain diagnostics do not apply. The
  # accuracy gate is the Pareto-k-hat of the nested integration (when the
  # backend computed it); fall back to the grid's quadrature effective sample
  # size for the C++ nested paths that store weights but no sampled k-hat yet.
  if (!.tulpa_is_chain(fit)) {
    # Model-package wrappers carry the nested pieces under `$joint_fit`; unwrap it
    # the same way the shared reliability readers do
    # (.tulpa_psis_reliability / .tulpa_grid_reliability) so the k-hat / weights
    # surface here too, rather than falling through to the generic message.
    jf <- if (!is.null(fit$joint_fit)) fit$joint_fit else fit
    # Outer-integration regime: read off the stored grid weights, so it is
    # available whether or not the k-hat diagnostic ran, and it qualifies what a
    # high k-hat means. Reported BEFORE the k-hat line so a
    # collapsed grid is not read as a failed integration.
    rgm <- .tulpa_outer_regime(fit)
    result$outer_regime <- if (is.null(rgm)) NA_character_ else rgm$regime
    rnote <- .tulpa_outer_regime_note(rgm)
    if (!is.null(rnote)) {
      recommendations <- c(recommendations, paste0(toupper(substring(rnote, 1, 1)),
                                                   substring(rnote, 2), "."))
      if (identical(rgm$regime, "collapsed_edge")) status <- "WARN"
    }
    # What the reported hyperparameter intervals were read off. A homogeneous
    # support is already named by `integration`; the locally refined grid is the
    # one that carries both kinds, and the share on the design part is the regime
    # variable for how much of the interval is a moment rule read as a CDF.
    ir <- .tulpa_interval_read(fit)
    if (!is.null(ir)) {
      result$interval_read <- ir$read
      result$interval_design_mass <- ir$design_mass
      inote <- .tulpa_interval_read_note(ir)
      if (!is.null(inote)) {
        recommendations <- c(recommendations, paste0(
          toupper(substring(inote, 1, 1)), substring(inote, 2), "."))
      }
    }
    k_hat <- jf$pareto_k
    if (!is.null(k_hat) && is.finite(k_hat)) {
      result$pareto_k <- k_hat
      if (k_hat >= .nl_diag("k_usable")) {
        status <- "WARN"
        recommendations <- c(recommendations, sprintf(
          paste("Pareto k-hat = %.2f (>= %s): the %s is misfit by the nested",
                "grid; the approximation is unreliable -- escalate (debias / MCMC)."),
          k_hat, format(.nl_diag("k_usable")), jf$pareto_k_scope %||% "hyperparameter posterior"))
      } else {
        recommendations <- c(recommendations, sprintf(
          "Pareto k-hat = %.2f (< %s): nested approximation is reliable.",
          k_hat, format(.nl_diag("k_usable"))))
      }
      # Opt-in per-arm k-hat: localises which arm's
      # hyperparameter axes drive a tail-heavy joint k.
      bak <- jf$pareto_k_by_arm
      if (!is.null(bak) && length(bak) > 0L) {
        result$pareto_k_by_arm <- bak
        nm <- names(bak) %||% paste0("arm", seq_along(bak))
        parts <- vapply(seq_along(bak), function(i)
          sprintf("%s %s", nm[i],
                  if (is.finite(bak[i])) sprintf("%.2f", bak[i]) else "NA"),
          character(1))
        recommendations <- c(recommendations, sprintf(
          "Per-arm Pareto k-hat (%s): %s.",
          jf$pareto_k_by_arm_scope %||% "other arms fixed at posterior mean",
          paste(parts, collapse = ", ")))
      }
    } else {
      # WHY there is no k-hat: "you turned it off", "this
      # family can never be scored" and "the outer proposal degenerated" all
      # arrived here as a bare NA before, and the last two are signals worth
      # acting on. An internal inconsistency is an engine bug, so it WARNs.
      kdec <- jf$pareto_k_declined
      knote <- .k_decline_note(kdec)
      if (!is.null(knote)) {
        result$pareto_k_declined <- kdec
        recommendations <- c(recommendations, sprintf(
          "No outer Pareto k-hat: %s.", knote))
        if (identical(sub(":.*$", "", kdec), "internal_inconsistency")) status <- "WARN"
      }
      w <- jf$weights
      if (!is.null(w) && length(w) > 1L && is.finite(sum(w))) {
        result$quad_ess <- .nl_grid_ess(w)
        recommendations <- c(recommendations, sprintf(
          paste("Quadrature ESS = %.1f of %d grid nodes (no sampled Pareto",
                "k-hat for this backend); validate with debias / coverage checks."),
          result$quad_ess, length(w)))
      } else {
        recommendations <- c(recommendations,
          paste("Approximation fit: Rhat/ESS are not convergence diagnostics here.",
                "Validate with the debias step (IMH/Gibbs) or coverage/recovery checks."))
      }

      # How much of a fixed-effect marginal the hyperparameter integration
      # carried. `attr(diagnostics(fit), "summary")` reported this from 0.6.0
      # and this door did not, so the one number that quantifies
      # gcol33/tulpa#862's symptom was reachable from one reporting path and
      # not the other. Same helper, so the two cannot disagree.
      #
      # Stated, never scored: what counts as too little share depends on the
      # model, so this carries no threshold and never moves `status` -- a fit
      # whose marginal rides almost entirely on the within-cell Gaussian may be
      # perfectly calibrated, which is what the coverage gate in
      # `test-spatial-beta-coverage.R` measures. Absent rather than zero when
      # the per-cell pieces were not retained.
      hs <- .tulpa_hyper_share(fit)
      if (!is.null(hs) && is.finite(hs$min_share)) {
        result$hyper_share_min <- hs$min_share
        result$hyper_share_max <- hs$max_share
        recommendations <- c(recommendations, sprintf(
          paste("Hyperparameter integration contributes %.1f-%.1f%% of the",
                "fixed-effect marginal variance (law of total variance over the",
                "outer grid); the rest is the within-cell Gaussian."),
          100 * hs$min_share, 100 * hs$max_share))
      }
    }

    # The inner (latent-field) Laplace layer, and why it went unscored when it
    # did -- a model class that can never be scored reads
    # differently from a diagnostic that was switched off.
    isk <- .tulpa_inner_skew_reliability(fit)
    if (!is.null(isk)) {
      if (is.finite(isk$max_abs_gamma3)) {
        result$inner_skew_max <- isk$max_abs_gamma3
        if (identical(isk$band, "unreliable")) status <- "WARN"
        recommendations <- c(recommendations, sprintf(
          "Inner-Laplace max |gamma_3| = %.2f (%s) over %d scored latents.",
          isk$max_abs_gamma3, isk$band, isk$n_scored))
      }
      inote <- .inner_skew_decline_note(isk$declined, isk$arms_declined)
      if (!is.null(inote)) {
        result$inner_skew_declined <- isk$declined
        recommendations <- c(recommendations,
                             sprintf("Inner-Laplace gamma_3 not reported: %s.", inote))
      }
    }

    # The second score on the same inner layer: the
    # importance k-hat, which needs no likelihood derivative and so answers
    # where gamma_3 declines.
    ikk <- .tulpa_inner_k_reliability(fit)
    if (!is.null(ikk)) {
      if (is.finite(ikk$max_pareto_k)) {
        result$inner_pareto_k <- ikk$max_pareto_k
        result$inner_pareto_k_rel_ess <- ikk$min_rel_ess
        if (identical(ikk$band, "unreliable")) status <- "WARN"
        recommendations <- c(recommendations, if (isTRUE(ikk$weights_uniform)) sprintf(
          paste("Inner-Laplace importance weights are uniform on all %d probed",
                "latents (min IS efficiency %.4f): the inner Gaussian needs no",
                "correction over the sampled region."),
          ikk$n_scored, ikk$min_rel_ess) else sprintf(
          paste("Inner-Laplace importance pareto_k = %.2f (%s) over %d scored",
                "latents, min IS efficiency %.3f."),
          ikk$max_pareto_k, ikk$band, ikk$n_scored, ikk$min_rel_ess))
      }
      knote2 <- .inner_k_decline_note(ikk$declined)
      if (!is.null(knote2)) {
        result$inner_pareto_k_declined <- ikk$declined
        recommendations <- c(recommendations,
          sprintf("Inner-Laplace importance k-hat not reported: %s.", knote2))
      }
    }

    converged <- fit$laplace_result$converged %||% fit$.internal$converged
    if (!is.null(converged) && !converged) {
      status <- "WARN"
      recommendations <- c(recommendations,
        "Inner Laplace optimization did not converge")
    }
  }

  # Outer-grid cells whose log-determinant came from the PD-enforced factor.
  # Read outside the approximation-only block for the same reason the axis note
  # below is: the flag rides the grid on every tier that produces one.
  s2z_fb <- .tulpa_s2z_flag_cells(fit, "s2z_log_det_fallback")
  if (!is.null(s2z_fb)) {
    result$s2z_log_det_fallback_cells <- s2z_fb$n
    status <- "WARN"
    recommendations <- c(recommendations, sprintf(
      paste("%d of %d outer-grid cells fell back to the PD-enforced",
            "log-determinant: their weights come from a determinant of",
            "H + lambda I rather than of the pinned sum-to-zero matrix."),
      s2z_fb$n, s2z_fb$n_grid))
  }
  # Cells whose field block was indefinite at the final pass: a count, not a
  # warning, since the inertia route reads the same value wherever the pinned
  # matrix is PD, and the cells where it is not are counted above.
  s2z_ind <- .tulpa_s2z_flag_cells(fit, "s2z_field_indefinite")
  if (!is.null(s2z_ind)) result$s2z_field_indefinite_cells <- s2z_ind$n

  # Outer-grid cells whose inner Newton stopped at the cap without a mode.
  uc <- .tulpa_unconverged_cells(fit)
  if (!is.null(uc)) {
    result$unconverged_cells <- uc$n
    status <- if (status == "FAIL") "FAIL" else "WARN"
    recommendations <- c(recommendations, .tulpa_unconverged_cells_note(uc))
  }

  # Grid axes the resolved path could not read. Recorded on
  # every tier, so it is read outside the approximation-only block above; a fit
  # that used every supplied axis carries nothing and this stays silent.
  axis_dropped <- .tulpa_axis_dropped(fit)
  if (!is.null(axis_dropped)) {
    result$axis_fields_dropped <- axis_dropped
    recommendations <- c(recommendations, .tulpa_axis_dropped_note(axis_dropped))
  }

  # The inner layer's remedy, appended LAST so it follows the bands that
  # motivate it rather than leading them. Same helper the diagnostics table
  # prints, so the two doors cannot say different things about what is
  # available, and it is silent on a fit that already ran a correction.
  idn <- .tulpa_inner_debias_note(fit, .tulpa_inner_skew_reliability(fit),
                                  .tulpa_inner_k_reliability(fit))
  if (!is.null(idn)) {
    recommendations <- c(recommendations,
                         paste0(toupper(substring(idn, 1, 1)),
                                substring(idn, 2), "."))
  }

  result$recommendations <- recommendations
  result$status <- status

  class(result) <- c("tulpa_diagnostic_summary", "list")

  # Print if not quiet

if (!quiet) {
    print(result)
  }

  invisible(result)
}


#' @rdname diagnostic_summary
#' @param x A `tulpa_diagnostic_summary` object.
#' @param ... Ignored.
#' @export
print.tulpa_diagnostic_summary <- function(x, ...) {
  # Status color (ANSI codes)
  status_color <- switch(
    x$status,
    "PASS" = "\033[32m",  # green
    "WARN" = "\033[33m",  # yellow
    "FAIL" = "\033[31m",  # red
    ""
  )
  reset <- "\033[0m"

  cat("\n")
  cat("=== tulpa Diagnostic Summary ===\n")
  cat("\n")
  cat("Backend:", x$backend, "\n")
  cat("Status: ", status_color, x$status, reset, "\n", sep = "")
  cat("\n")

  # Divergences
  if (x$backend == "hmc") {
    cat("Divergent transitions:", x$n_divergent)
    if (x$n_divergent > 0) {
      cat(sprintf(" (%.1f%%)", x$divergent_pct))
    }
    cat("\n")

    if (!is.na(x$e_bfmi)) {
      ebfmi_status <- if (x$e_bfmi >= 0.3) "OK" else "LOW"
      cat(sprintf("E-BFMI: %.3f (%s)\n", x$e_bfmi, ebfmi_status))
    }
    cat("\n")
  }

  # Pareto k-hat / quadrature ESS (approximation fits)
  if (!is.null(x$pareto_k) && is.finite(x$pareto_k)) {
    khat_status <- if (x$pareto_k < .nl_diag("k_usable")) "OK" else "HIGH"
    cat(sprintf("Pareto k-hat: %.3f (%s)\n\n", x$pareto_k, khat_status))
  } else if (!is.null(x$quad_ess) && is.finite(x$quad_ess)) {
    cat(sprintf("Quadrature ESS: %.1f\n\n", x$quad_ess))
  }

  if (!is.null(x$hyper_share_min) && is.finite(x$hyper_share_min)) {
    cat(sprintf("Hyperparameter share of the fixed-effect marginal: %.1f-%.1f%%\n\n",
                100 * x$hyper_share_min, 100 * x$hyper_share_max))
  }

  # Worst Rhat
  if (!is.null(x$worst_rhat) && nrow(x$worst_rhat) > 0) {
    cat("Parameters with Rhat > 1.01:\n")
    x$worst_rhat$rhat <- round(x$worst_rhat$rhat, 3)
    print(x$worst_rhat, row.names = FALSE)
    cat("\n")
  }

  # Worst ESS
  if (!is.null(x$worst_ess) && nrow(x$worst_ess) > 0) {
    cat("Parameters with ESS < 400:\n")
    x$worst_ess$ess_bulk <- round(x$worst_ess$ess_bulk, 0)
    x$worst_ess$ess_tail <- round(x$worst_ess$ess_tail, 0)
    print(x$worst_ess, row.names = FALSE)
    cat("\n")
  }

  # Recommendations
  if (length(x$recommendations) > 0) {
    cat("Recommendations:\n")
    for (rec in x$recommendations) {
      cat("  -", rec, "\n")
    }
  }

  invisible(x)
}


#' Multi-Panel Diagnostic Figure
#'
#' @description
#' Creates a combined diagnostic figure with Rhat, ESS, trace plot, and
#' energy/ACF panels. Requires the `patchwork` package for layout.
#'
#' @param fit A `tulpa_fit` object.
#' @param pars Character vector of parameter names for trace plot. If NULL,
#'   uses the parameter with worst Rhat.
#'
#' @return A combined plot (ggplot + patchwork) or NULL if requirements not met.
#'
#' @details
#' Creates a 2x2 grid:
#' - Top left: Rhat plot
#' - Top right: ESS plot
#' - Bottom left: Trace for worst parameter
#' - Bottom right: Energy (HMC) or ACF (other)
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' plot_diagnostics(fit)
#' }
#'
#' @seealso [diagnostic_summary()], [plot_rhat()], [plot_ess()]
#' @export
plot_diagnostics <- function(fit, pars = NULL) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  if (!.tulpa_is_chain(fit)) {
    message(.tulpa_non_chain_msg(fit))
    return(invisible(NULL))
  }

  # Check for required packages
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    message("plot_diagnostics() requires ggplot2. Using individual plot functions instead.")
    message("Install with: install.packages('ggplot2')")
    return(invisible(NULL))
  }

  if (!requireNamespace("patchwork", quietly = TRUE)) {
    message("plot_diagnostics() requires patchwork for layout.")
    message("Install with: install.packages('patchwork')")
    message("Using sequential plots instead...")

    # Fall back to sequential plots
    print(plot_rhat(fit))
    print(plot_ess(fit))
    return(invisible(NULL))
  }

  # Create individual plots
  p_rhat <- plot_rhat(fit) + ggplot2::ggtitle("Rhat")
  p_ess <- plot_ess(fit) + ggplot2::ggtitle("ESS (bulk)")

  # Trace plot for worst parameter
  diag <- tryCatch(diagnostics(fit), error = function(e) NULL)
  if (!is.null(diag)) {
    main_pars <- select_main_params(diag$parameter)
    diag <- diag[diag$parameter %in% main_pars, ]
    # which.max() returns integer(0) when every rhat is NA (all-constant
    # parameters), which would make worst_par length-0 and error at the `%in%`
    # test below; fall back to the first draws column.
    wi <- which.max(diag$rhat)
    worst_par <- if (length(wi)) diag$parameter[wi] else colnames(fit$draws)[1]
  } else {
    worst_par <- colnames(fit$draws)[1]
  }

  # Create trace plot
  draws_array <- get_draws_array(fit)$draws
  if (length(worst_par) && worst_par %in% dimnames(draws_array)[[3]]) {
    trace_data <- data.frame(
      iteration = seq_len(dim(draws_array)[1]),
      value = draws_array[, 1, worst_par]
    )

    p_trace <- ggplot2::ggplot(trace_data, ggplot2::aes(x = .data$iteration, y = .data$value)) +
      ggplot2::geom_line(color = tulpa_diag_colors$chain[1], alpha = 0.7) +
      ggplot2::labs(title = paste("Trace:", worst_par), x = "Iteration", y = "Value") +
      theme_tulpa() +
      ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"))
  } else {
    p_trace <- ggplot2::ggplot() +
      ggplot2::annotate("text", x = 0.5, y = 0.5, label = "No trace available") +
      ggplot2::theme_void()
  }

  # Bottom right: Energy or ACF
  if (fit$backend == "hmc" && !is.null(fit$energy)) {
    p_br <- plot_energy(fit)
  } else {
    p_br <- plot_acf(fit, n_pars = 1)
  }

  # Handle NULL plots
  if (is.null(p_rhat)) p_rhat <- ggplot2::ggplot() + ggplot2::theme_void()
  if (is.null(p_ess)) p_ess <- ggplot2::ggplot() + ggplot2::theme_void()
  if (is.null(p_br)) p_br <- ggplot2::ggplot() + ggplot2::theme_void()

  # Combine with patchwork
  combined <- (p_rhat + p_ess) / (p_trace + p_br) +
    patchwork::plot_annotation(
      title = "tulpa Model Diagnostics",
      theme = ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5, face = "bold", size = 14))
    )

  combined
}


#' Geweke convergence test
#'
#' @description
#' Performs Geweke's convergence diagnostic, comparing the mean of the first
#' portion of a chain to the last portion. It is a within-chain diagnostic, so
#' on a multi-chain fit it is computed for each chain separately.
#'
#' @param fit A `tulpa_fit` object.
#' @param frac1 Fraction of each chain for the first window (default: 0.1).
#' @param frac2 Fraction of each chain for the second window (default: 0.5).
#'   Both lie in `(0, 1)` and the windows may not overlap:
#'   `frac1 + frac2 <= 1`.
#' @param pars Character vector of parameter names (default: all main params).
#'
#' @return A data frame of class `tulpa_geweke` with columns `chain`,
#'   `parameter`, `z_score` and `p_value`, one row per chain and parameter.
#'
#' @details
#' The Geweke test computes a z-score comparing the means of early and late
#' portions of a chain. Large z-scores (|z| > 2) indicate the chain has not
#' converged. The chains are never pooled: concatenated, the "first" window
#' would come from the first chain and the "last" from another, which compares
#' two chains rather than a chain's start with its end (use [diagnostics()]
#' for the between-chain R-hat). A parameter that is constant within a chain
#' has an undefined z-score, reported as `NA`.
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' geweke_test(fit)
#' }
#'
#' @seealso [diagnostics()], [check_diagnostics()]
#' @export
geweke_test <- function(fit, frac1 = 0.1, frac2 = 0.5, pars = NULL) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }
  .check_unit_interval(frac1, "frac1")
  .check_unit_interval(frac2, "frac2")
  if (frac1 + frac2 > 1) {
    stop("`frac1 + frac2` must not exceed 1: the two windows would overlap.",
         call. = FALSE)
  }
  if (!.tulpa_is_chain(fit)) {
    message(.tulpa_non_chain_msg(fit))
    return(invisible(NULL))
  }

  draws <- .tulpa_pooled_draws(fit)
  if (is.null(draws)) {
    stop("No draws available in fit object", call. = FALSE)
  }
  # One series per chain (gcol33/tulpa#895); the pooled matrix only supplies
  # the resolved parameter names, in the same column order.
  chains <- lapply(.tulpa_chain_list(fit), function(m) {
    colnames(m) <- colnames(draws)
    m
  })

  # Select parameters
  if (is.null(pars)) {
    pars <- select_main_params(colnames(draws))
  } else {
    pars <- grep_params(pars, colnames(draws))
  }

  n <- min(vapply(chains, nrow, integer(1)))
  n1 <- floor(frac1 * n)
  n2 <- floor(frac2 * n)
  if (n1 < 2L || n2 < 2L) {
    stop(sprintf(paste0(
      "The Geweke windows hold %d and %d draws of a %d-draw chain; each needs ",
      "at least 2. Raise `frac1` / `frac2` or run longer chains."),
      as.integer(n1), as.integer(n2), as.integer(n)), call. = FALSE)
  }

  results <- lapply(seq_along(chains), function(k) {
    ch <- chains[[k]]
    nk <- nrow(ch)
    do.call(rbind, lapply(pars, function(p) {
      x <- ch[, p]

      # First window
      x1 <- x[seq_len(floor(frac1 * nk))]
      var1 <- spectrum0_ar(x1) / length(x1)

      # Second window (last portion)
      x2 <- x[(nk - floor(frac2 * nk) + 1):nk]
      var2 <- spectrum0_ar(x2) / length(x2)

      # Z-score; a window with no variation leaves it undefined, not infinite
      z <- if (isTRUE(var1 + var2 > 0)) (mean(x1) - mean(x2)) / sqrt(var1 + var2) else
        NA_real_
      p_val <- 2 * (1 - pnorm(abs(z)))

      data.frame(
        chain = k,
        parameter = p,
        z_score = z,
        p_value = p_val,
        stringsAsFactors = FALSE
      )
    }))
  })

  result <- do.call(rbind, results)
  class(result) <- c("tulpa_geweke", "data.frame")
  result
}


#' Spectrum at frequency zero (AR fit)
#' @noRd
spectrum0_ar <- function(x) {
  # Fit AR model to estimate spectral density at frequency 0
  n <- length(x)
  if (n < 10) return(var(x))

  # Use AR(p) with AIC selection
  ar_fit <- tryCatch(
    ar(x, aic = TRUE, method = "burg"),
    error = function(e) NULL
  )

  if (is.null(ar_fit) || ar_fit$var.pred <= 0) {
    return(var(x))
  }

  # Spectral density at frequency 0
  ar_fit$var.pred / (1 - sum(ar_fit$ar))^2
}


#' @rdname geweke_test
#' @param x A `tulpa_geweke` result.
#' @param ... Ignored.
#' @export
print.tulpa_geweke <- function(x, ...) {
  cat("Geweke Convergence Diagnostic\n")
  cat("=============================\n\n")

  x$z_score <- round(x$z_score, 3)
  x$p_value <- round(x$p_value, 4)

  print.data.frame(x, row.names = FALSE)

  # Flag potential issues; an NA z (a window without variation) is not one
  bad <- !is.na(x$z_score) & abs(x$z_score) > 2
  if (any(bad)) {
    what <- if (!is.null(x$chain) && length(unique(x$chain)) > 1L) {
      "chain/parameter pair(s)"
    } else {
      "parameter(s)"
    }
    cat("\nWarning:", sum(bad), what,
        "have |z| > 2 (potential non-convergence)\n")
  }

  invisible(x)
}


# Resolve a vector of parameter patterns against available draw names. Exact
# names (including bracketed-index entries like `u[3]`) match literally; any
# remaining patterns are treated as regular expressions, falling back to a
# fixed-substring match if the pattern is not valid regex.
grep_params <- function(pattern, names) {
  if (is.null(pattern)) return(names)
  out <- character(0)
  for (p in pattern) {
    exact <- names[names == p]
    if (length(exact)) {
      out <- c(out, exact)
      next
    }
    m <- tryCatch(grep(p, names, value = TRUE), error = function(e) NULL)
    if (is.null(m) || length(m) == 0L) m <- grep(p, names, value = TRUE, fixed = TRUE)
    out <- c(out, m)
  }
  unique(out)
}


#' Number of divergent transitions
#'
#' Counts divergent transitions recorded by an HMC/NUTS fit, reading whichever
#' field the backend populated (the top-level `$divergent` flag vector every
#' sampler in this package writes, or `$diagnostics$divergent`,
#' `$diagnostics$divergent_idx`, `$diagnostics$n_divergent`, `$n_divergent`).
#' It reads the same record [plot_divergences()] locates the divergent rows
#' from, so the two always agree on one fit.
#'
#' @param fit A `tulpa_fit` object.
#' @return Integer count of divergent transitions (0 if none are recorded).
#' @seealso [plot_divergences()], [check_diagnostics()]
#' @examples
#' \donttest{
#' set.seed(1)
#' d <- data.frame(x = rnorm(60))
#' d$y <- rpois(60, exp(0.5 + 0.3 * d$x))
#' fit <- tulpa(y ~ x, data = d, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 400L, warmup = 200L, n_chains = 2L,
#'                             seed = 1L))
#' n_divergent(fit)
#' }
#' @export
n_divergent <- function(fit) {
  .tulpa_divergence_record(fit)$n
}


#' Quick convergence check
#'
#' Runs the core convergence diagnostics on a fit and reports whether Rhat,
#' bulk-ESS, and divergence thresholds are all met. A terse companion to
#' [diagnostic_summary()] for use in scripts and tests.
#'
#' @param fit A `tulpa_fit` object.
#' @param rhat_threshold Maximum acceptable Rhat (default 1.01).
#' @param ess_threshold Minimum acceptable bulk-ESS (default 400).
#' @param quiet Logical; if TRUE, suppress messages (default FALSE).
#'
#' @return Invisibly, `TRUE` if all checks pass, `FALSE` if any fail, or `NA`
#'   for a non-chain (approximation) fit where Rhat/ESS do not apply.
#'
#' @examples
#' \donttest{
#' set.seed(123)
#' df <- data.frame(x = rnorm(60))
#' df$y <- rpois(60, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 500L, warmup = 250L, n_chains = 2L,
#'                             seed = 1L))
#' check_diagnostics(fit)
#' }
#'
#' @seealso [diagnostic_summary()], [diagnostics()], [n_divergent()]
#' @export
check_diagnostics <- function(fit, rhat_threshold = 1.01, ess_threshold = 400,
                              quiet = FALSE) {
  if (!inherits(fit, "tulpa_fit")) {
    stop("fit must be a tulpa_fit object", call. = FALSE)
  }

  # Non-chain (approximation) fits have no Rhat/ESS to check; report "not
  # applicable" rather than a vacuous pass.
  if (!.tulpa_is_chain(fit)) {
    if (!quiet) message(.tulpa_non_chain_msg(fit))
    return(invisible(NA))
  }

  issues <- character(0)
  diag <- tryCatch(diagnostics(fit), error = function(e) NULL)
  if (is.null(diag)) {
    # A chain fit too short to diagnose (< 4 draws/chain) yields no Rhat/ESS;
    # do not fall through to a spurious "checks passed" on divergences alone.
    if (!quiet) message("Insufficient draws to assess convergence ",
                        "(need >= 4 draws per chain); Rhat/ESS not computed.")
    return(invisible(NA))
  }
  if (!is.null(diag)) {
    main_pars <- select_main_params(diag$parameter)
    diag <- diag[diag$parameter %in% main_pars, , drop = FALSE]
    n_bad_rhat <- sum(diag$rhat > rhat_threshold, na.rm = TRUE)
    if (n_bad_rhat > 0) {
      issues <- c(issues, sprintf("%d parameter(s) with Rhat > %s",
                                  n_bad_rhat, format(rhat_threshold)))
    }
    n_low_ess <- sum(diag$ess_bulk < ess_threshold, na.rm = TRUE)
    if (n_low_ess > 0) {
      issues <- c(issues, sprintf("%d parameter(s) with bulk-ESS < %d",
                                  n_low_ess, ess_threshold))
    }
  }

  n_div <- n_divergent(fit)
  if (n_div > 0) {
    issues <- c(issues, sprintf("%d divergent transition(s)", n_div))
  }

  ok <- length(issues) == 0L
  if (!quiet) {
    if (ok) {
      message("Convergence checks passed (Rhat, bulk-ESS, divergences).")
    } else {
      message("Convergence warnings:")
      for (msg in issues) message("  - ", msg)
    }
  }
  invisible(ok)
}