gam 0.2.3

//! NUTS Sampler using general-mcmc
//!
//! This module provides NUTS (No-U-Turn Sampler) for honest uncertainty
//! quantification after PIRLS convergence.
//!
//! # Design
//!
//! Since general-mcmc's NUTS uses an identity mass matrix, we whiten the
//! parameter space using the Cholesky decomposition of the inverse Hessian:
//!
//! - Transform: β = μ + L @ z  (where L L^T = H^{-1})
//! - The whitened space has unit covariance, so NUTS mixes efficiently
//! - Samples are un-transformed back to the original space
//!
//! # Analytical Gradients
//!
//! We override `unnorm_logp_and_grad` to compute gradients analytically using
//! ndarray, avoiding burn's autodiff overhead. The gradient computation mirrors
//! the true log-posterior gradient (not the PIRLS working gradient).
//!
//! # Memory Efficiency
//!
//! Large data (design matrix, response, etc.) is wrapped in `Arc` to allow
//! sharing across chains without duplication when general-mcmc clones the target.

use super::polya_gamma::PolyaGamma;
use crate::construction::CanonicalPenalty;
use crate::estimate::reml::FirthDenseOperator;
use crate::estimate::reml::penalty_logdet::PenaltyPseudologdet;
use crate::estimate::{
    EstimationError, UnifiedFitResult, validate_explicit_dense_hessian_for_whitening,
};
use crate::faer_ndarray::{FaerCholesky, FaerEigh, fast_ata_into, fast_atv, fast_av_into};
use crate::families::gamlss::monotone_wiggle_basis_with_derivative_order;
use crate::matrix::DesignMatrix;
use crate::solver::mixture_link::{
    InverseLinkKernel, LinkParamPartials, inverse_link_jet_for_inverse_link, softmax_last_fixedzero,
};
use crate::types::{
    InverseLink, LikelihoodSpec, LinkFunction, ResponseFamily, RhoPrior, is_valid_tweedie_power,
};
use crate::visualizer::VisualizerSession;
use faer::Side;
use general_mcmc::generic_hmc::HamiltonianTarget;
pub use general_mcmc::generic_nuts::NUTSMassMatrixConfig;
use general_mcmc::generic_nuts::{GenericNUTS, MassMatrixAdaptation};
use ndarray::{Array1, Array2, Array3, ArrayView1, ArrayView2, Axis};
use rand::{RngExt, SeedableRng, rngs::StdRng};
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::fmt;
use std::sync::{Arc, Mutex};

/// Only Binomial Logit supports Firth bias reduction.
#[inline]
fn likelihood_spec_supports_firth(spec: &LikelihoodSpec) -> bool {
    matches!(
        (&spec.response, &spec.link),
        (
            ResponseFamily::Binomial,
            InverseLink::Standard(LinkFunction::Logit),
        )
    )
}

/// Typed error variants for the HMC / NUTS sampling module.
///
/// External-facing helpers in this module continue to return
/// `Result<_, String>`; this enum is materialized internally and converted
/// at the public boundary via `.map_err(String::from)` so that the error
/// text remains byte-identical to the previous `format!` output.
#[derive(Debug, Clone)]
pub enum HmcError {
    /// Sampler state (penalty / Hessian / mode / posterior values) contains
    /// NaN or Inf where finiteness is required.
    NonFiniteState { reason: String },
    /// Configuration value (e.g. `target_accept`, unit-weight requirement)
    /// is out of range or otherwise invalid.
    InvalidConfig { reason: String },
    /// Dimensions of the supplied matrices / vectors are inconsistent.
    DimensionMismatch { reason: String },
    /// Firth/Jeffreys correction was requested for a family that does not
    /// support it.
    FirthUnsupported { reason: String },
    /// Inverse-link state does not match the requested likelihood family in
    /// the joint (β, ρ) sampler.
    LinkMismatch { reason: String },
    /// Likelihood family is not implemented in the current sampling path.
    UnsupportedFamily { reason: String },
    /// Sampling produced no usable output (empty kept set, non-finite
    /// summary statistic, etc.).
    SamplingFailed { reason: String },
}

impl fmt::Display for HmcError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HmcError::NonFiniteState { reason }
            | HmcError::InvalidConfig { reason }
            | HmcError::DimensionMismatch { reason }
            | HmcError::FirthUnsupported { reason }
            | HmcError::LinkMismatch { reason }
            | HmcError::UnsupportedFamily { reason }
            | HmcError::SamplingFailed { reason } => f.write_str(reason),
        }
    }
}

impl From<HmcError> for String {
    fn from(err: HmcError) -> String {
        err.to_string()
    }
}

struct SamplingVisualizer {
    session: VisualizerSession,
    chain: usize,
    total_chains: usize,
    warmup: usize,
    samples: usize,
}

impl SamplingVisualizer {
    fn new(stage: &str, total_chains: usize, warmup: usize, samples: usize) -> Self {
        let mut session = VisualizerSession::new(true);
        session.set_stage("sample", stage);
        session.start_workflow(&format!("chains completed ({stage})"), total_chains.max(1));
        Self {
            session,
            chain: 0,
            total_chains,
            warmup,
            samples,
        }
    }

    fn begin_chain(&mut self, chain: usize, label: &str) {
        self.chain = chain;
        self.session
            .set_stage("sample", &format!("[Chain {}] {label}", chain + 1));
        self.session
            .start_workflow("chains completed", self.total_chains.max(1));
        self.session.advance_workflow(chain);
        self.session
            .start_secondary_workflow(&format!("[Chain {}] Warmup", chain + 1), self.warmup);
    }

    fn warmup_step(&mut self, iter: usize) {
        self.session
            .advance_secondary_workflow(iter.min(self.warmup));
    }

    fn start_sampling(&mut self) {
        self.session
            .finish_secondary_progress(&format!("chain {} warmup complete", self.chain + 1));
        self.session
            .start_secondary_workflow(&format!("[Chain {}] Sample", self.chain + 1), self.samples);
    }

    fn sample_step(&mut self, iter: usize) {
        self.session
            .advance_secondary_workflow(iter.min(self.samples));
    }

    fn finish_chain(&mut self, accept_rate: f64) {
        self.session.finish_secondary_progress(&format!(
            "chain {} sample complete | accepted {:.1}%",
            self.chain + 1,
            accept_rate * 100.0
        ));
        self.session
            .advance_workflow((self.chain + 1).min(self.total_chains));
    }

    fn finish_all(&mut self, rhat: f64, ess: f64) {
        self.session.push_diagnostic(&format!(
            "sampling diagnostics | rhat={rhat:.3} | ess={ess:.1}"
        ));
        self.session.finish_progress("sampling complete");
    }
}

/// Compute split-chain R-hat and ESS using the Gelman-Rubin diagnostic.
///
/// This is the standard split-chain formulation (no rank normalization).
/// Returns (max_rhat, min_ess) across dimensions.
fn compute_split_rhat_and_ess(samples: &Array3<f64>) -> (f64, f64) {
    let n_chains = samples.shape()[0];
    let n_samples = samples.shape()[1];
    let dim = samples.shape()[2];

    if n_chains < 2 || n_samples < 4 {
        return (1.0, n_chains as f64 * n_samples as f64 * 0.5);
    }

    // Split each chain in half to detect non-stationarity
    let half = n_samples / 2;
    let n_split_chains = n_chains * 2;
    let n_split_samples = half;

    let mut max_rhat = 0.0f64;
    let mut min_ess = f64::INFINITY;

    #[inline]
    fn splitvalue(
        samples: &Array3<f64>,
        n_chains: usize,
        half: usize,
        dim: usize,
        sc: usize,
        t: usize,
    ) -> f64 {
        let chain = sc % n_chains;
        if sc < n_chains {
            samples[[chain, t, dim]]
        } else {
            samples[[chain, half + t, dim]]
        }
    }

    fn ess_from_split_dimension(
        samples: &Array3<f64>,
        n_chains: usize,
        half: usize,
        dim: usize,
    ) -> f64 {
        let m = n_chains * 2;
        let n = half;
        if m == 0 || n < 4 {
            return (m * n).max(1) as f64;
        }

        let mut means = vec![0.0_f64; m];
        let mut gamma0 = vec![0.0_f64; m];
        for sc in 0..m {
            let mut sum = 0.0;
            for t in 0..n {
                sum += splitvalue(samples, n_chains, half, dim, sc, t);
            }
            let mean = sum / n as f64;
            means[sc] = mean;
            let mut g0 = 0.0;
            for t in 0..n {
                let d = splitvalue(samples, n_chains, half, dim, sc, t) - mean;
                g0 += d * d;
            }
            gamma0[sc] = (g0 / n as f64).max(1e-16);
        }

        let max_lag = (n - 1).min(1000);
        let mut tau = 1.0_f64;
        let mut lag = 1usize;
        while lag < max_lag {
            let mut pair = 0.0_f64;
            for l in [lag, lag + 1] {
                if l > max_lag {
                    continue;
                }
                let mut rho_l = 0.0;
                for sc in 0..m {
                    let mu = means[sc];
                    let mut cov = 0.0;
                    let denom = (n - l) as f64;
                    for t in 0..(n - l) {
                        let x0 = splitvalue(samples, n_chains, half, dim, sc, t);
                        let x1 = splitvalue(samples, n_chains, half, dim, sc, t + l);
                        cov += (x0 - mu) * (x1 - mu);
                    }
                    cov /= denom;
                    rho_l += cov / gamma0[sc];
                }
                rho_l /= m as f64;
                pair += rho_l;
            }
            if !pair.is_finite() || pair <= 0.0 {
                break;
            }
            tau += 2.0 * pair;
            lag += 2;
        }
        if !tau.is_finite() || tau <= 0.0 {
            return 1.0;
        }
        let total = (m * n) as f64;
        (total / tau).clamp(1.0, total)
    }

    let mut chain_means = vec![0.0_f64; n_split_chains];
    let mut chainvars = vec![0.0_f64; n_split_chains];
    for d in 0..dim {
        for chain in 0..n_chains {
            // First half
            let mut sum1 = 0.0;
            for i in 0..half {
                sum1 += samples[[chain, i, d]];
            }
            let mean1 = sum1 / half as f64;
            let mut var1 = 0.0;
            for i in 0..half {
                let diff = samples[[chain, i, d]] - mean1;
                var1 += diff * diff;
            }
            var1 /= (half - 1).max(1) as f64;
            let first_idx = chain;
            chain_means[first_idx] = mean1;
            chainvars[first_idx] = var1;

            // Second half
            let mut sum2 = 0.0;
            for i in half..(2 * half) {
                sum2 += samples[[chain, i, d]];
            }
            let mean2 = sum2 / half as f64;
            let mut var2 = 0.0;
            for i in half..(2 * half) {
                let diff = samples[[chain, i, d]] - mean2;
                var2 += diff * diff;
            }
            var2 /= (half - 1).max(1) as f64;
            let second_idx = n_chains + chain;
            chain_means[second_idx] = mean2;
            chainvars[second_idx] = var2;
        }

        // Within-chain variance W
        let w: f64 = chainvars.iter().copied().sum::<f64>() / n_split_chains as f64;

        // Between-chain variance B
        let overall_mean: f64 = chain_means.iter().copied().sum::<f64>() / n_split_chains as f64;
        let b: f64 = chain_means
            .iter()
            .map(|m| (m - overall_mean).powi(2))
            .sum::<f64>()
            * n_split_samples as f64
            / (n_split_chains - 1) as f64;

        // Estimated variance
        let var_hat = (n_split_samples as f64 - 1.0) / n_split_samples as f64 * w
            + b / n_split_samples as f64;

        // R-hat
        let rhat_d = if w > 1e-10 { (var_hat / w).sqrt() } else { 1.0 };
        max_rhat = max_rhat.max(rhat_d);

        // Real ESS via split-chain autocorrelation with Geyer IPS truncation.
        let ess_d = ess_from_split_dimension(samples, n_chains, half, d);
        min_ess = min_ess.min(ess_d);
    }

    (max_rhat, min_ess.max(1.0))
}

/// Solve L^T * X = I where L is lower triangular.
///
/// Returns X = L^{-T} (the inverse transpose of L).
///
/// This is the correct way to compute the whitening transform matrix:
/// Given H = L L^T (Cholesky), we need W where W W^T = H^{-1}
/// Since H^{-1} = L^{-T} L^{-1}, we have W = L^{-T}.
///
/// Implementation strategy (math-equivalent to back-substitution on L^T):
/// We compute L^{-1} column-wise via forward substitution on L, then the
/// result is `L^{-1}` transposed. Forward-substituting column `c` of L^{-1}
/// uses `L`'s rows (which are contiguous in row-major `Array2`), giving
/// stride-1 inner loops instead of the strided `l[[j, i]]` (column-major
/// access pattern) and double-indexed writes of the original. We also
/// exploit the triangular structure of `L^{-1}` (entries above the diagonal
/// are zero), skipping ~half of the inner work compared to the previous
/// version which traversed `i = (0..dim).rev()` for every column.
///
/// Total cost: ~dim^3 / 6 multiply-adds (down from dim^3 / 2), with all
/// inner loops on contiguous slices.
fn solve_upper_triangular_transpose(l: &Array2<f64>, dim: usize) -> Array2<f64> {
    let mut result = Array2::<f64>::zeros((dim, dim));
    if dim == 0 {
        return result;
    }

    // Pull contiguous row slice access from L (row-major standard layout).
    // Falls back to a one-time owned copy if `l` is not standard-layout
    // (e.g. a transposed view); both branches feed the same inner loop.
    let l_owned;
    let l_rows: &[f64] = if let Some(s) = l.as_slice() {
        s
    } else {
        l_owned = l.to_owned();
        l_owned
            .as_slice()
            .expect("owned standard-layout Array2 has contiguous storage")
    };

    // Scratch column for L^{-1}[:, col]; reused across columns.
    let mut y = vec![0.0_f64; dim];

    for col in 0..dim {
        // Forward-substitute L * y = e_col. y[i] = 0 for i < col.
        // Diagonal term:
        let d_col = l_rows[col * dim + col];
        let inv_d_col = if d_col.abs() > 1e-15 {
            1.0 / d_col
        } else {
            0.0
        };
        y[col] = inv_d_col;

        // Below-diagonal entries: y[i] = -(sum_{j=col..i} L[i,j] * y[j]) / L[i,i].
        // Each inner loop is a stride-1 dot product on row `i` of L (contiguous).
        for i in (col + 1)..dim {
            let row_off = i * dim;
            let l_row = &l_rows[row_off + col..row_off + i];
            let y_seg = &y[col..i];
            // Both operands are contiguous slices of equal length; the loop
            // is a straight-line stride-1 reduction the optimizer can
            // auto-vectorize.
            let mut sum = 0.0_f64;
            for k in 0..l_row.len() {
                sum += l_row[k] * y_seg[k];
            }
            let d = l_rows[row_off + i];
            y[i] = if d.abs() > 1e-15 { -sum / d } else { 0.0 };
        }

        // Write the column into result transposed: result[col, i] = y[i] for i >= col.
        // result[i, col] is left at zero for i < col (upper-triangular L^{-T}).
        // That matches `result[col, i]` filling row `col` from column `col` rightward.
        let res_row_start = col * dim + col;
        let res_row = &mut result.as_slice_mut().expect("owned Array2 contiguous")
            [res_row_start..res_row_start + (dim - col)];
        for (k, slot) in res_row.iter_mut().enumerate() {
            *slot = y[col + k];
        }

        // Clear scratch positions we wrote, so the next column starts clean above.
        for slot in &mut y[col..dim] {
            *slot = 0.0;
        }
    }

    result
}

/// Shared data for NUTS posterior (wrapped in Arc to prevent cloning).
///
/// This struct holds read-only data that is shared across all chains.
/// Using Arc prevents memory explosion when general-mcmc clones the target.
#[derive(Clone)]
struct SharedData {
    /// Design matrix X [n_samples, dim]
    x: Arc<Array2<f64>>,
    /// Response vector y [n_samples]
    y: Arc<Array1<f64>>,
    /// Observation/case weights [n_samples]
    weights: Arc<Array1<f64>>,
    /// MAP estimate (mode) μ [dim]
    mode: Arc<Array1<f64>>,
    /// Auxiliary log-link family parameter: Gamma shape, Tweedie power, or NB theta.
    gamma_shape: f64,
    /// Dispersion parameter φ used to scale the likelihood for families
    /// with an estimated scale (Gaussian: σ²; Gamma: 1/shape). For fixed-
    /// scale families this is `Known(1.0)`. The Gaussian log-likelihood
    /// and its gradient multiply through by `1/φ` so that the whitened
    /// sampler targets `φ·H⁻¹` rather than a silently mis-scaled
    /// posterior. See `inference::dispersion_cov` for the ownership
    /// invariants this field encodes.
    dispersion: crate::solver::estimate::Dispersion,
    /// Number of samples
    n_samples: usize,
    /// Number of coefficients
    dim: usize,
}

thread_local! {
    static NUTS_RESIDUAL_SCRATCH: RefCell<Array1<f64>> = RefCell::new(Array1::zeros(0));
}

/// Whitened log-posterior target with analytical gradients.
///
/// Uses Arc for shared data to prevent memory explosion when cloned for chains.
/// Uses faer for numerically stable Cholesky decomposition.
/// Family mode for NUTS log-likelihood computation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NutsFamily {
    Gaussian,
    BinomialLogit,
    BinomialProbit,
    BinomialCLogLog,
    PoissonLog,
    TweedieLog,
    NegativeBinomialLog,
    GammaLog,
}

impl NutsFamily {
    #[inline]
    fn likelihood_spec(self) -> LikelihoodSpec {
        match self {
            Self::Gaussian => LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            Self::BinomialLogit => LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Logit),
            },
            Self::BinomialProbit => LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Probit),
            },
            Self::BinomialCLogLog => LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::CLogLog),
            },
            Self::PoissonLog => LikelihoodSpec {
                response: ResponseFamily::Poisson,
                link: InverseLink::Standard(LinkFunction::Log),
            },
            Self::TweedieLog => LikelihoodSpec {
                response: ResponseFamily::Tweedie { p: 1.5 },
                link: InverseLink::Standard(LinkFunction::Log),
            },
            Self::NegativeBinomialLog => LikelihoodSpec {
                response: ResponseFamily::NegativeBinomial { theta: 1.0 },
                link: InverseLink::Standard(LinkFunction::Log),
            },
            Self::GammaLog => LikelihoodSpec {
                response: ResponseFamily::Gamma,
                link: InverseLink::Standard(LinkFunction::Log),
            },
        }
    }
}

/// Whitened-coordinate target for the No-U-Turn HMC sampler.
///
/// The posterior over β is reparameterized via `β = L z` where `L Lᵀ = H⁻¹`
/// (Cholesky factor of the inverse posterior Hessian at the MAP), so that
/// in `z`-coordinates the local curvature is approximately the identity.
/// The struct holds the shared design, the whitening factor `L` and its
/// transpose (for gradient chain-rule pull-back `∇_z = Lᵀ ∇_β`), the
/// family-specific log-likelihood adapter, and a precomputed
/// `M = Lᵀ S L` so the smoothing penalty `−½ βᵀSβ` becomes the cheap
/// quadratic `−½ zᵀMz` inside the leapfrog hot loop.  Optionally adds
/// the identifiable-subspace Firth/Jeffreys term to keep posterior modes
/// away from infinity under separation.
pub struct NutsPosterior {
    /// Shared read-only data (Arc prevents duplication)
    data: SharedData,
    /// Transform: L where L L^T = H^{-1} (computed from Hessian)
    /// This is the inverse-transpose of the Cholesky of H.
    chol: Array2<f64>,
    /// L^T for gradient chain rule: ∇z = L^T @ ∇_β
    chol_t: Array2<f64>,
    /// Family for log-likelihood computation
    nuts_family: NutsFamily,
    /// Whether to add the identifiable-subspace Jeffreys/Firth term to the
    /// target
    firth_enabled: bool,
    /// Precomputed whitened-penalty operator `M = L^T S L` (dim×dim, symmetric
    /// positive-semidefinite). The penalty term in z-coordinates is
    ///   −0.5 βᵀSβ = −[c0 + (Lᵀ S μ)ᵀ z + 0.5 zᵀ M z],
    /// so its z-gradient is just `−(L^T S μ + M z)` — no per-step `S·β` matvec
    /// or `L^T·∇_β penalty` map is needed.
    penalty_z_quad: Array2<f64>,
    /// Precomputed `Lᵀ S μ` (length dim) — z-space gradient contribution from
    /// the linear-in-z portion of the penalty.
    penalty_z_lin: Array1<f64>,
    /// Precomputed `0.5 μᵀ S μ` (scalar) — constant term of the penalty.
    penalty_z_const: f64,
}

impl NutsPosterior {
    #[inline]
    fn log1pexp(eta: f64) -> f64 {
        if eta > 0.0 {
            eta + (-eta).exp().ln_1p()
        } else {
            eta.exp().ln_1p()
        }
    }

    #[inline]
    fn sigmoid_stable(eta: f64) -> f64 {
        if eta >= 0.0 {
            1.0 / (1.0 + (-eta).exp())
        } else {
            let e = eta.exp();
            e / (1.0 + e)
        }
    }

    /// Creates a new posterior target from ndarray data.
    ///
    /// # Arguments
    /// * `x` - Design matrix [n_samples, dim]
    /// * `y` - Response vector [n_samples]
    /// * `weights` - Observation/case weights [n_samples]
    /// * `penalty_matrix` - Combined penalty S [dim, dim]
    /// * `mode` - MAP estimate μ [dim]
    /// * `hessian` - Hessian H [dim, dim] (NOT the inverse!)
    /// * `nuts_family` - Family for log-likelihood computation
    ///
    /// # Numerical Stability
    /// Accepts the Hessian directly and computes L = (chol(H))^{-T} via
    /// triangular solves, which is more stable than explicitly inverting H.
    pub fn new(
        x: ArrayView2<f64>,
        y: ArrayView1<f64>,
        weights: ArrayView1<f64>,
        penalty_matrix: ArrayView2<f64>,
        mode: ArrayView1<f64>,
        hessian: ArrayView2<f64>,
        nuts_family: NutsFamily,
        gamma_shape: f64,
        dispersion: crate::solver::estimate::Dispersion,
        firth_enabled: bool,
    ) -> Result<Self, String> {
        let n_samples = x.nrows();
        let dim = x.ncols();

        // Validate inputs are finite
        if !penalty_matrix.iter().all(|x| x.is_finite()) {
            return Err(HmcError::NonFiniteState {
                reason: "Penalty matrix contains NaN or Inf values".to_string(),
            }
            .into());
        }
        if !hessian.iter().all(|x| x.is_finite()) {
            return Err(HmcError::NonFiniteState {
                reason: "Hessian matrix contains NaN or Inf values".to_string(),
            }
            .into());
        }
        if !mode.iter().all(|x| x.is_finite()) {
            return Err(HmcError::NonFiniteState {
                reason: "Mode vector contains NaN or Inf values".to_string(),
            }
            .into());
        }

        validate_firth_support(nuts_family, firth_enabled).map_err(String::from)?;
        if nuts_family.likelihood_spec().is_binomial() {
            validate_binary_responses("binomial NUTS", &y, &weights).map_err(String::from)?;
        }
        if matches!(nuts_family, NutsFamily::NegativeBinomialLog) {
            validate_count_responses("negative-binomial NUTS", &y, &weights)
                .map_err(String::from)?;
        }

        // Use faer for numerically stable Cholesky decomposition of H
        // H = L_H L_H^T where L_H is lower triangular
        let hessian_owned = hessian.to_owned();
        let chol_factor = hessian_owned
            .cholesky(Side::Lower)
            .map_err(|e| format!("Hessian Cholesky decomposition failed: {:?}", e))?;

        // We need L where L L^T = H^{-1}
        // Since H = L_H L_H^T, we have H^{-1} = L_H^{-T} L_H^{-1}
        // So L = L_H^{-T} (the inverse transpose of the Cholesky factor)
        //
        // To get L_H^{-T}, we solve L_H^T * X = I using back-substitution
        // Since L_H is lower triangular, L_H^T is upper triangular
        let l_h = chol_factor.lower_triangular();
        let mut chol = solve_upper_triangular_transpose(&l_h, dim);
        // Scale the Cholesky factor by sqrt(phi) so that L Lᵀ = phi · H⁻¹,
        // i.e. the whitened sampler targets the φ-scaled posterior
        // covariance `Vb`. For fixed-scale families (Binomial, Poisson)
        // `phi == 1` and `sqrt_phi() == 1`, so this is a no-op. For
        // Gaussian and Gamma the dispersion is estimated and this is the
        // missing √φ that previously left the sampler whitening against
        // a unit-variance reference instead of `Vb`.
        {
            use crate::inference::dispersion_cov::DispersionExt as _;
            let sqrt_phi = dispersion.sqrt_phi();
            if (sqrt_phi - 1.0).abs() > 0.0 {
                chol.mapv_inplace(|v| v * sqrt_phi);
            }
        }
        let chol_t = chol.t().to_owned();

        // Precompute the whitened penalty operator and constants so that the
        // penalty contribution to logp/grad becomes a single symv against z.
        // Math identity (β = μ + L z, L L^T = H^{-1}):
        //   0.5 β^T S β = 0.5 μ^T S μ + (L^T S μ)^T z + 0.5 z^T (L^T S L) z
        // and ∇_z [0.5 β^T S β] = L^T S μ + (L^T S L) z.
        // This replaces three matvecs per leapfrog step (S·β, L·z used only
        // for that purpose, and L^T·∇_β penalty) with one dim×dim symv.
        let penalty_owned = penalty_matrix.to_owned();
        let mode_owned = mode.to_owned();
        let s_mu = penalty_owned.dot(&mode_owned);
        let penalty_z_const = 0.5 * mode_owned.dot(&s_mu);
        let penalty_z_lin = chol_t.dot(&s_mu);
        // M = L^T S L = chol_t · (S · chol). Computed in two GEMMs at
        // construction time only.
        let s_chol = penalty_owned.dot(&chol);
        let penalty_z_quad = chol_t.dot(&s_chol);

        let data = SharedData {
            x: Arc::new(x.to_owned()),
            y: Arc::new(y.to_owned()),
            weights: Arc::new(weights.to_owned()),
            mode: Arc::new(mode_owned),
            gamma_shape,
            dispersion,
            n_samples,
            dim,
        };

        Ok(Self {
            data,
            chol,
            chol_t,
            nuts_family,
            firth_enabled,
            penalty_z_quad,
            penalty_z_lin,
            penalty_z_const,
        })
    }

    fn compute_logp_and_grad_nd_into(
        &self,
        z: &Array1<f64>,
        residual: &mut Array1<f64>,
        grad: &mut Array1<f64>,
    ) -> f64 {
        // === Step 1: Transform z (whitened) -> β (original) ===
        // β = μ + L @ z
        let beta = self.data.mode.as_ref() + &self.chol.dot(z);

        // === Step 2: Compute η = X @ β ===
        let eta = crate::faer_ndarray::fast_av(self.data.x.as_ref(), &beta);

        // === Step 3: Compute log-likelihood and gradient ===
        let (ll, mut grad_ll_beta) = self.family_logp_and_grad_into(&eta, residual);

        let mut firth_logdet = 0.0;
        if self.firth_enabled {
            match firth_jeffreys_logp_and_grad(self.nuts_family, &self.data, &eta) {
                Ok((value, grad_beta_firth)) => {
                    firth_logdet = value;
                    grad_ll_beta += &grad_beta_firth;
                }
                Err(err) => {
                    log::warn!(
                        "[NUTS/Firth] Jeffreys target became invalid at the current state: {}",
                        err
                    );
                    grad.fill(0.0);
                    return f64::NEG_INFINITY;
                }
            }
        }

        // === Step 4: Penalty in z-coordinates (precomputed; see `new`) ===
        //   −0.5 βᵀ S β  =  −[c0 + lᵀ z + 0.5 zᵀ M z]
        //   ∇_z (−0.5 βᵀ S β) = −(l + M z)
        // where l = L^T S μ, M = L^T S L, c0 = 0.5 μᵀ S μ.
        // This single dim×dim symmetric matvec replaces both the per-step
        // S·β multiply and the L^T·∇_β penalty chain-rule multiply, and lets
        // the penalty value, β-gradient and chain rule fuse into one pass.
        //
        // The module invariant `Vb = φ · H^{-1}` with `H = X'WX + S` requires
        // that the prior is `1/φ`-scaled whenever the likelihood is (Gaussian,
        // GammaLog with the shape ≡ 1/φ scaling at `gamma_log_logp_and_grad`).
        // For fixed-scale families (Bernoulli / Poisson) `φ ≡ 1` so
        // `inv_phi = 1` collapses this to the original expression. This
        // mirrors the `penalty_scale` branching in `LinkWigglePosterior`.
        use crate::inference::dispersion_cov::DispersionExt as _;
        let penalty_scale = match self.nuts_family {
            NutsFamily::Gaussian | NutsFamily::GammaLog => self.data.dispersion.inv_phi(),
            _ => 1.0,
        };
        let mz = self.penalty_z_quad.dot(z);
        let lin_term = self.penalty_z_lin.dot(z);
        let quad_term = 0.5 * z.dot(&mz);
        let penalty = penalty_scale * (self.penalty_z_const + lin_term + quad_term);

        // === Step 5: z-space gradient ===
        // ∇z log p = L^T ∇_β ℓ  −  penalty_scale · (l + M z)
        fast_av_into(&self.chol_t, &grad_ll_beta, grad);
        // gradz -= penalty_scale · (penalty_z_lin + M z); fused parallel update.
        let lin_view = self.penalty_z_lin.view();
        ndarray::Zip::from(grad)
            .and(&lin_view)
            .and(&mz)
            .par_for_each(|g, &l, &m| {
                *g -= penalty_scale * (l + m);
            });

        ll + firth_logdet - penalty
    }

    fn family_logp_and_grad_into(
        &self,
        eta: &Array1<f64>,
        residual: &mut Array1<f64>,
    ) -> (f64, Array1<f64>) {
        nuts_family_logp_and_grad_into(self.nuts_family, &self.data, eta, residual)
    }

    /// Get the Cholesky factor L for un-whitening samples
    pub fn chol(&self) -> &Array2<f64> {
        &self.chol
    }

    /// Get the mode
    pub fn mode(&self) -> &Array1<f64> {
        &self.data.mode
    }

    /// Get dimension
    pub fn dim(&self) -> usize {
        self.data.dim
    }
}

const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_7;

#[inline]
fn standard_normal_log_pdf(x: f64) -> f64 {
    -0.5 * x * x - HALF_LOG_2PI
}

/// Stable log Φ(x) for the standard normal CDF.
#[inline]
fn log_ndtr(x: f64) -> f64 {
    let arg = -x * std::f64::consts::FRAC_1_SQRT_2;
    let erfc_val = statrs::function::erf::erfc(arg);
    if erfc_val > 0.0 {
        erfc_val.ln() - std::f64::consts::LN_2
    } else {
        -0.5 * x * x - (-x).ln() - HALF_LOG_2PI
    }
}

#[inline]
fn validate_firth_support(family: NutsFamily, firth_enabled: bool) -> Result<(), HmcError> {
    let spec = family.likelihood_spec();
    if firth_enabled && !likelihood_spec_supports_firth(&spec) {
        let binomial_logit = LikelihoodSpec {
            response: ResponseFamily::Binomial,
            link: InverseLink::Standard(LinkFunction::Logit),
        };
        return Err(HmcError::FirthUnsupported {
            reason: format!(
                "NUTS with Firth is only supported for {}; {} does not support it",
                binomial_logit.pretty_name(),
                spec.pretty_name()
            ),
        });
    }
    Ok::<(), _>(())
}

#[inline]
fn validate_firth_likelihood_support(
    likelihood: &LikelihoodSpec,
    firth_enabled: bool,
) -> Result<(), HmcError> {
    if firth_enabled && !likelihood_spec_supports_firth(likelihood) {
        let binomial_logit = LikelihoodSpec {
            response: ResponseFamily::Binomial,
            link: InverseLink::Standard(LinkFunction::Logit),
        };
        return Err(HmcError::FirthUnsupported {
            reason: format!(
                "Joint HMC with Firth is only supported for {}; {} does not support it",
                binomial_logit.pretty_name(),
                likelihood.pretty_name()
            ),
        });
    }
    Ok::<(), _>(())
}

#[inline]
fn valid_count_response(y: f64) -> bool {
    y.is_finite() && y >= 0.0 && (y - y.round()).abs() <= 1e-9
}

fn validate_count_responses(
    family: &str,
    y: &ArrayView1<'_, f64>,
    weights: &ArrayView1<'_, f64>,
) -> Result<(), HmcError> {
    for (i, (&yi, &wi)) in y.iter().zip(weights.iter()).enumerate() {
        if wi > 0.0 && !valid_count_response(yi) {
            return Err(HmcError::InvalidConfig {
                reason: format!(
                    "{family} response must be a finite non-negative integer at positive-weight row {i}; got {yi}"
                ),
            });
        }
    }
    Ok(())
}

fn validate_binary_responses(
    family: &str,
    y: &ArrayView1<'_, f64>,
    weights: &ArrayView1<'_, f64>,
) -> Result<(), HmcError> {
    for (i, (&yi, &wi)) in y.iter().zip(weights.iter()).enumerate() {
        if wi > 0.0 && !(yi == 0.0 || yi == 1.0) {
            return Err(HmcError::InvalidConfig {
                reason: format!(
                    "{family} response must be exactly 0 or 1 at positive-weight row {i}; got {yi}"
                ),
            });
        }
    }
    Ok(())
}

/// Compute the identifiable-subspace Jeffreys/Firth contribution and its
/// β-gradient.
///
/// HMC uses the same `FirthDenseOperator` as the REML exact-gradient path.
/// The operator owns the reduced identifiable Fisher factorization, the
/// Jeffreys log-determinant, and the analytic β-gradient.
fn firth_jeffreys_logp_and_grad(
    family: NutsFamily,
    data: &SharedData,
    eta: &Array1<f64>,
) -> Result<(f64, Array1<f64>), HmcError> {
    if eta.len() != data.n_samples {
        return Err(HmcError::DimensionMismatch {
            reason: format!(
                "Firth Jeffreys term eta length {} != number of samples {}",
                eta.len(),
                data.n_samples
            ),
        });
    }
    if data.dim == 0 || data.n_samples == 0 {
        return Ok((0.0, Array1::zeros(data.dim)));
    }
    validate_firth_support(family, true)?;
    if data.weights.iter().all(|w| *w == 0.0) {
        return Ok((0.0, Array1::zeros(data.dim)));
    }

    let op = if data.weights.iter().all(|&w| w == 1.0) {
        FirthDenseOperator::build(data.x.as_ref(), eta)
    } else {
        FirthDenseOperator::build_with_observation_weights(
            data.x.as_ref(),
            eta,
            data.weights.view(),
        )
    }
    .map_err(|e| HmcError::SamplingFailed {
        reason: format!("Firth Jeffreys operator failed: {e}"),
    })?;
    Ok(op.jeffreys_logdet_and_beta_gradient())
}

// ============================================================================
// Shared family log-likelihood helpers
// ============================================================================
//
// Freestanding functions for computing ℓ(y|β) and ∇_β ℓ for each supported
// family. Used by both `NutsPosterior` (fixed-ρ β-only sampling) and
// `JointBetaRhoPosterior` (joint β+ρ sampling).

fn nuts_family_logp_and_grad_into(
    family: NutsFamily,
    data: &SharedData,
    eta: &Array1<f64>,
    residual: &mut Array1<f64>,
) -> (f64, Array1<f64>) {
    match family {
        NutsFamily::BinomialLogit => logit_logp_and_grad_into(data, eta, residual),
        NutsFamily::BinomialProbit => probit_logp_and_grad_into(data, eta, residual),
        NutsFamily::BinomialCLogLog => cloglog_logp_and_grad_into(data, eta, residual),
        NutsFamily::Gaussian => gaussian_logp_and_grad_into(data, eta, residual),
        NutsFamily::PoissonLog => poisson_log_logp_and_grad(data, eta),
        // Family mapping: TweedieLog stores variance power p in data.gamma_shape.
        // Its dispersion phi stays in data.dispersion, matching REML scale ownership.
        NutsFamily::TweedieLog => tweedie_log_quasilogp_and_grad(data, eta, data.gamma_shape),
        NutsFamily::NegativeBinomialLog => {
            // Family mapping: NegativeBinomialLog stores theta in data.gamma_shape.
            // NB has unit REML scale; theta is never sourced from fixed_phi.
            negative_binomial_log_logp_and_grad(data, eta, data.gamma_shape)
        }
        NutsFamily::GammaLog => gamma_log_logp_and_grad(data, eta),
    }
}

#[derive(Clone, Debug)]
struct BinomialLinkTerms {
    log_mu: f64,
    log1m_mu: f64,
    dlog_mu_deta: f64,
    dlog1m_mu_deta: f64,
    dmu_dlink: Vec<f64>,
}

#[inline]
fn log_terms_from_mu_and_dmu(
    mu: f64,
    dmu_deta: f64,
    dmu_dlink: Vec<f64>,
) -> Result<BinomialLinkTerms, String> {
    if !(mu.is_finite() && (0.0..=1.0).contains(&mu) && dmu_deta.is_finite()) {
        return Err(format!(
            "binomial inverse link returned invalid mu/deta derivative: mu={mu}, dmu_deta={dmu_deta}"
        ));
    }
    let log_mu = if mu == 0.0 {
        f64::NEG_INFINITY
    } else {
        mu.ln()
    };
    let one_minus_mu = 1.0 - mu;
    let log1m_mu = if one_minus_mu == 0.0 {
        f64::NEG_INFINITY
    } else {
        one_minus_mu.ln()
    };
    let dlog_mu_deta = if mu == 0.0 {
        f64::INFINITY.copysign(dmu_deta)
    } else {
        dmu_deta / mu
    };
    let dlog1m_mu_deta = if one_minus_mu == 0.0 {
        f64::NEG_INFINITY.copysign(dmu_deta)
    } else {
        -dmu_deta / one_minus_mu
    };
    Ok(BinomialLinkTerms {
        log_mu,
        log1m_mu,
        dlog_mu_deta,
        dlog1m_mu_deta,
        dmu_dlink,
    })
}

#[inline]
fn binomial_link_terms(
    inverse_link: &InverseLink,
    eta: f64,
    n_link_params: usize,
) -> Result<BinomialLinkTerms, String> {
    let jet =
        inverse_link_jet_for_inverse_link(inverse_link, eta).map_err(|err| err.to_string())?;
    let mut dmu_dlink = vec![0.0; n_link_params];
    if n_link_params > 0 {
        match inverse_link
            .param_partials(eta)
            .map_err(|err| err.to_string())?
        {
            Some(LinkParamPartials::Sas(partials)) => {
                if n_link_params != 2 {
                    return Err(format!(
                        "SAS/Beta-Logistic link parameter dimension mismatch: expected 2, got {n_link_params}"
                    ));
                }
                dmu_dlink[0] = partials.djet_depsilon.mu;
                dmu_dlink[1] = partials.djet_dlog_delta.mu;
            }
            Some(LinkParamPartials::Mixture(partials)) => {
                if partials.djet_drho.len() != n_link_params {
                    return Err(format!(
                        "mixture link parameter dimension mismatch: expected {}, got {n_link_params}",
                        partials.djet_drho.len()
                    ));
                }
                for (slot, partial) in dmu_dlink.iter_mut().zip(partials.djet_drho.iter()) {
                    *slot = partial.mu;
                }
            }
            None => {
                return Err(format!(
                    "joint HMC expected {n_link_params} adaptive link parameters, but the inverse link exposes none"
                ));
            }
        }
    }
    log_terms_from_mu_and_dmu(jet.mu, jet.d1, dmu_dlink)
}

fn joint_binomial_logp_grad_and_link_grad(
    inverse_link: &InverseLink,
    data: &SharedData,
    eta: &Array1<f64>,
    n_link_params: usize,
) -> Result<(f64, Array1<f64>, Array1<f64>), String> {
    let n = data.n_samples;
    // Per-row: compute stable log-tail terms and derivatives without endpoint
    // clamping. Positive-weight responses were validated as Bernoulli before
    // target construction, so each row selects exactly one log branch.
    use rayon::iter::{IntoParallelIterator, ParallelIterator};
    let per_row: Result<Vec<(f64, f64, Vec<f64>)>, String> = (0..n)
        .into_par_iter()
        .map(|i| {
            let y_i = data.y[i];
            let w_i = data.weights[i];
            if w_i <= 0.0 {
                return Ok((0.0, 0.0, vec![0.0; n_link_params]));
            }
            let terms = binomial_link_terms(inverse_link, eta[i], n_link_params)?;
            if y_i == 1.0 {
                let inv_mu = terms.log_mu.exp().recip();
                let log_mu = terms.log_mu;
                let dlog_mu_deta = terms.dlog_mu_deta;
                let grad_link = terms
                    .dmu_dlink
                    .into_iter()
                    .map(|dmu| w_i * dmu * inv_mu)
                    .collect();
                Ok((w_i * log_mu, w_i * dlog_mu_deta, grad_link))
            } else if y_i == 0.0 {
                let inv_one_minus_mu = terms.log1m_mu.exp().recip();
                let log1m_mu = terms.log1m_mu;
                let dlog1m_mu_deta = terms.dlog1m_mu_deta;
                let grad_link = terms
                    .dmu_dlink
                    .into_iter()
                    .map(|dmu| -w_i * dmu * inv_one_minus_mu)
                    .collect();
                Ok((w_i * log1m_mu, w_i * dlog1m_mu_deta, grad_link))
            } else {
                Err(format!(
                    "binomial joint HMC response must be exactly 0 or 1 after validation; got {y_i}"
                ))
            }
        })
        .collect();
    let per_row = per_row?;
    let mut residual = Array1::<f64>::zeros(n);
    let mut grad_link = Array1::<f64>::zeros(n_link_params);
    let mut ll = 0.0;
    for (i, (ll_i, residual_i, grad_link_i)) in per_row.into_iter().enumerate() {
        ll += ll_i;
        residual[i] = residual_i;
        for (slot, value) in grad_link.iter_mut().zip(grad_link_i.iter()) {
            *slot += *value;
        }
    }

    Ok((ll, fast_atv(&data.x, &residual), grad_link))
}

fn joint_binomial_logp_and_grad(
    likelihood: &LikelihoodSpec,
    data: &SharedData,
    eta: &Array1<f64>,
) -> Result<(f64, Array1<f64>), String> {
    if !matches!(likelihood.response, ResponseFamily::Binomial) {
        return Err(HmcError::UnsupportedFamily {
            reason: format!(
                "{} is not a binomial joint-HMC family",
                likelihood.pretty_name()
            ),
        }
        .into());
    }
    match &likelihood.link {
        InverseLink::Standard(LinkFunction::Logit) => Ok(logit_logp_and_grad(data, eta)),
        InverseLink::Standard(LinkFunction::Probit) => Ok(probit_logp_and_grad(data, eta)),
        InverseLink::Standard(LinkFunction::CLogLog) => Ok(cloglog_logp_and_grad(data, eta)),
        InverseLink::LatentCLogLog(_)
        | InverseLink::Sas(_)
        | InverseLink::BetaLogistic(_)
        | InverseLink::Mixture(_) => {
            let (ll, grad_beta, _) =
                joint_binomial_logp_grad_and_link_grad(&likelihood.link, data, eta, 0)?;
            Ok((ll, grad_beta))
        }
        InverseLink::Standard(_) => Err(HmcError::UnsupportedFamily {
            reason: format!(
                "{} is not a binomial joint-HMC family",
                likelihood.pretty_name()
            ),
        }
        .into()),
    }
}

fn joint_family_logp_grad_and_link_grad(
    likelihood: &LikelihoodSpec,
    data: &SharedData,
    eta: &Array1<f64>,
    n_link_params: usize,
) -> Result<(f64, Array1<f64>, Array1<f64>), String> {
    match (&likelihood.response, &likelihood.link) {
        (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::Logit)) => {
            let (ll, grad) = logit_logp_and_grad(data, eta);
            Ok((ll, grad, Array1::zeros(n_link_params)))
        }
        (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::Probit)) => {
            let (ll, grad) = probit_logp_and_grad(data, eta);
            Ok((ll, grad, Array1::zeros(n_link_params)))
        }
        (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::CLogLog)) => {
            let (ll, grad) = cloglog_logp_and_grad(data, eta);
            Ok((ll, grad, Array1::zeros(n_link_params)))
        }
        (
            ResponseFamily::Binomial,
            InverseLink::LatentCLogLog(_)
            | InverseLink::Sas(_)
            | InverseLink::BetaLogistic(_)
            | InverseLink::Mixture(_),
        ) => joint_binomial_logp_grad_and_link_grad(&likelihood.link, data, eta, n_link_params),
        _ => {
            let (ll, grad) = joint_family_logp_and_grad(likelihood, data, eta)?;
            Ok((ll, grad, Array1::zeros(n_link_params)))
        }
    }
}

fn joint_family_logp_and_grad(
    likelihood: &LikelihoodSpec,
    data: &SharedData,
    eta: &Array1<f64>,
) -> Result<(f64, Array1<f64>), String> {
    match &likelihood.response {
        ResponseFamily::Binomial => joint_binomial_logp_and_grad(likelihood, data, eta),
        ResponseFamily::Gaussian => Ok(gaussian_logp_and_grad(data, eta)),
        ResponseFamily::Poisson => Ok(poisson_log_logp_and_grad(data, eta)),
        ResponseFamily::Tweedie { p } => {
            // Family mapping: Tweedie payload p is the variance power.
            // Its dispersion phi stays in data.dispersion, matching REML.
            let p = *p;
            if !is_valid_tweedie_power(p) {
                return Err(HmcError::InvalidConfig {
                    reason: format!(
                        "Tweedie variance power must be finite and strictly between 1 and 2; got {p}"
                    ),
                }
                .into());
            }
            Ok(tweedie_log_quasilogp_and_grad(data, eta, p))
        }
        ResponseFamily::NegativeBinomial { theta } => {
            // Family mapping: NegativeBinomial payload theta is overdispersion.
            // NB keeps unit REML scale and never reads fixed_phi for theta.
            Ok(negative_binomial_log_logp_and_grad(data, eta, *theta))
        }
        ResponseFamily::Beta { .. } => Err(HmcError::UnsupportedFamily {
            reason: "Joint HMC fallback is not implemented for BetaLogit".to_string(),
        }
        .into()),
        ResponseFamily::Gamma => Ok(gamma_log_logp_and_grad(data, eta)),
        ResponseFamily::RoystonParmar => Err(HmcError::UnsupportedFamily {
            reason: "Joint HMC fallback is not implemented for RoystonParmar".to_string(),
        }
        .into()),
    }
}

/// Logistic regression log-likelihood and gradient.
///
/// log p(y|η) = y·η − log(1 + exp(η)), gradient = X'(w ⊙ (y − μ))
fn logit_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    let mut residual = Array1::<f64>::zeros(data.n_samples);
    logit_logp_and_grad_into(data, eta, &mut residual)
}

fn logit_logp_and_grad_into(
    data: &SharedData,
    eta: &Array1<f64>,
    residual: &mut Array1<f64>,
) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    assert_eq!(residual.len(), n);
    // Per-row independent: write residual entry into a pre-allocated buffer and
    // reduce the ll contribution in parallel — avoids materialising a
    // Vec<(f64, f64)> and the serial scatter that follows.
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let y_i = data.y[i];
            let w_i = data.weights[i];
            let mu = NutsPosterior::sigmoid_stable(eta_i);
            *slot = w_i * (y_i - mu);
            w_i * (y_i * eta_i - NutsPosterior::log1pexp(eta_i))
        })
        .sum();

    let grad_ll = fast_atv(data.x.as_ref(), &*residual);
    (ll, grad_ll)
}

/// Probit regression log-likelihood and gradient.
///
/// log p(y|η) = Σ [y·log Φ(η) + (1-y)·log(1-Φ(η))],
/// gradient_i = w_i · [y_i · φ(η_i)/Φ(η_i) − (1-y_i) · φ(η_i)/(1−Φ(η_i))]
///
/// Uses erfc-based log Φ for numerical stability.
fn probit_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    let mut residual = Array1::<f64>::zeros(data.n_samples);
    probit_logp_and_grad_into(data, eta, &mut residual)
}

fn probit_logp_and_grad_into(
    data: &SharedData,
    eta: &Array1<f64>,
    residual: &mut Array1<f64>,
) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    assert_eq!(residual.len(), n);
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let y_i = data.y[i];
            let w_i = data.weights[i];
            let log_phi_pos = log_ndtr(eta_i);
            let log_phi_neg = log_ndtr(-eta_i);
            let log_phi_val = standard_normal_log_pdf(eta_i);
            let ratio_pos = (log_phi_val - log_phi_pos).exp();
            let ratio_neg = (log_phi_val - log_phi_neg).exp();
            let grad_i = y_i * ratio_pos - (1.0 - y_i) * ratio_neg;
            *slot = w_i * grad_i;
            w_i * (y_i * log_phi_pos + (1.0 - y_i) * log_phi_neg)
        })
        .sum();

    let grad_ll = fast_atv(data.x.as_ref(), &*residual);
    (ll, grad_ll)
}

/// Complementary log-log regression log-likelihood and gradient.
///
/// CLogLog link: μ = 1 − exp(−exp(η))
/// log p(y|η) = Σ [y·log(1−exp(−exp(η))) + (1−y)·(−exp(η))]
/// gradient_i = w_i · [y_i · exp(η_i)·exp(−exp(η_i)) / (1−exp(−exp(η_i))) − (1−y_i)·exp(η_i)]
#[inline]
fn cloglog_bernoulli_logp_and_residual(eta: f64, y: f64) -> Result<(f64, f64), EstimationError> {
    if !(eta.is_finite() && (-700.0..=700.0).contains(&eta)) {
        return Err(EstimationError::InvalidInput(format!(
            "cloglog eta must be finite and within [-700, 700]; got {eta}"
        )));
    }
    let exp_eta = eta.exp();
    let log_mu = if exp_eta <= std::f64::consts::LN_2 {
        (-(-exp_eta).exp_m1()).ln()
    } else {
        (-(-exp_eta).exp()).ln_1p()
    };
    let log_one_minus_mu = -exp_eta;
    let grad_log_mu = (eta - exp_eta - log_mu).exp();
    let ll_i = y * log_mu + (1.0 - y) * log_one_minus_mu;
    let residual_i = y * grad_log_mu - (1.0 - y) * exp_eta;
    Ok((ll_i, residual_i))
}

fn cloglog_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    let mut residual = Array1::<f64>::zeros(data.n_samples);
    cloglog_logp_and_grad_into(data, eta, &mut residual)
}

fn cloglog_logp_and_grad_into(
    data: &SharedData,
    eta: &Array1<f64>,
    residual: &mut Array1<f64>,
) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    assert_eq!(residual.len(), n);
    if eta
        .iter()
        .any(|&eta_i| !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)))
    {
        return (f64::NEG_INFINITY, Array1::zeros(data.dim));
    }
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let y_i = data.y[i];
            let w_i = data.weights[i];
            let (ll_i, residual_i) =
                cloglog_bernoulli_logp_and_residual(eta[i], y_i).expect("validated cloglog eta");
            *slot = w_i * residual_i;
            w_i * ll_i
        })
        .sum();

    let grad_ll = fast_atv(data.x.as_ref(), &*residual);
    (ll, grad_ll)
}

/// Gaussian log-likelihood and gradient.
///
/// log p(y|η) = −½ (w/φ)·(y − η)²,  gradient = (1/φ)·X'(w ⊙ (y − η))
///
/// Both the log-likelihood and its β-gradient are scaled by `1/φ` so that
/// the working likelihood matches the φ-scaled posterior covariance the
/// HMC whitening transform targets. With `φ == 1` (the only value
/// passed by the pre-refactor call sites) this collapses to the original
/// `−½ w·(y − η)²` expression; with an estimated dispersion (the
/// `Dispersion::Estimated(σ²)` branch) it removes the silent unit-σ
/// approximation the Gaussian NUTS log-density used previously.
fn gaussian_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    let mut weighted_residual = Array1::<f64>::zeros(data.n_samples);
    gaussian_logp_and_grad_into(data, eta, &mut weighted_residual)
}

fn gaussian_logp_and_grad_into(
    data: &SharedData,
    eta: &Array1<f64>,
    weighted_residual: &mut Array1<f64>,
) -> (f64, Array1<f64>) {
    use crate::inference::dispersion_cov::DispersionExt as _;
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    let inv_phi = data.dispersion.inv_phi();
    assert_eq!(weighted_residual.len(), n);
    // Per-row: residual = y - η, weighted_residual = (w/φ)·residual,
    // ll contribution = -0.5·(w/φ)·residual². All independent across rows.
    let ll: f64 = weighted_residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let residual = data.y[i] - eta[i];
            let w_i = data.weights[i];
            let scaled = w_i * inv_phi;
            *slot = scaled * residual;
            -0.5 * scaled * residual * residual
        })
        .sum();

    let grad_ll = fast_atv(data.x.as_ref(), &*weighted_residual);
    (ll, grad_ll)
}

/// Poisson(log) log-likelihood and gradient.
///
/// log p(y|η) = y·η − exp(η), gradient = X'(w ⊙ (y − μ))
fn poisson_log_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    if eta
        .iter()
        .any(|&eta_i| !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)))
    {
        return (f64::NEG_INFINITY, Array1::zeros(data.dim));
    }
    let mut residual = Array1::<f64>::zeros(n);
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let mu_i = eta_i.exp();
            let y_i = data.y[i];
            let w_i = data.weights[i];
            *slot = w_i * (y_i - mu_i);
            w_i * (y_i * eta_i - mu_i)
        })
        .sum();

    let grad_ll = fast_atv(&data.x, &residual);
    (ll, grad_ll)
}

fn tweedie_log_quasilogp_and_grad(
    data: &SharedData,
    eta: &Array1<f64>,
    p: f64,
) -> (f64, Array1<f64>) {
    use crate::inference::dispersion_cov::DispersionExt as _;
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    // Family mapping: Tweedie p is the variant payload; phi is data.dispersion.
    // Invalid payloads invalidate the target instead of falling back to p=1.5.
    if !is_valid_tweedie_power(p) {
        return (f64::NAN, Array1::from_elem(data.dim, f64::NAN));
    }
    if eta
        .iter()
        .any(|&eta_i| !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)))
    {
        return (f64::NEG_INFINITY, Array1::zeros(data.dim));
    }
    let inv_phi = data.dispersion.inv_phi();
    let mut residual = Array1::<f64>::zeros(n);
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let mu_i = eta_i.exp().max(1e-300);
            let y_i = data.y[i];
            let w_i = data.weights[i] * inv_phi;
            *slot = w_i * (y_i - mu_i) * mu_i.powf(1.0 - p);
            let qll = y_i * mu_i.powf(1.0 - p) / (1.0 - p) - mu_i.powf(2.0 - p) / (2.0 - p);
            w_i * qll
        })
        .sum();

    let grad_ll = fast_atv(&data.x, &residual);
    (ll, grad_ll)
}

fn negative_binomial_log_logp_and_grad(
    data: &SharedData,
    eta: &Array1<f64>,
    theta: f64,
) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    if !(theta.is_finite() && theta > 0.0)
        || eta
            .iter()
            .any(|&eta_i| !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)))
        || data
            .y
            .iter()
            .zip(data.weights.iter())
            .any(|(&y_i, &w_i)| w_i > 0.0 && !valid_count_response(y_i))
    {
        return (f64::NEG_INFINITY, Array1::zeros(data.dim));
    }
    let mut residual = Array1::<f64>::zeros(n);
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let mu_i = eta_i.exp().max(1e-12);
            let y_i = data.y[i];
            let w_i = data.weights[i];
            if w_i <= 0.0 {
                *slot = 0.0;
                return 0.0;
            }
            let log_mu_term = if y_i > 0.0 { y_i * mu_i.ln() } else { 0.0 };
            *slot = w_i * theta * (y_i - mu_i) / (theta + mu_i);
            w_i * (statrs::function::gamma::ln_gamma(y_i + theta)
                - statrs::function::gamma::ln_gamma(theta)
                - statrs::function::gamma::ln_gamma(y_i + 1.0)
                + theta * (theta.ln() - (theta + mu_i).ln())
                + log_mu_term
                - y_i * (theta + mu_i).ln())
        })
        .sum();

    let grad_ll = fast_atv(&data.x, &residual);
    (ll, grad_ll)
}

fn gamma_log_logp_and_grad(data: &SharedData, eta: &Array1<f64>) -> (f64, Array1<f64>) {
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
    let n = data.n_samples;
    if eta
        .iter()
        .any(|&eta_i| !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)))
    {
        return (f64::NEG_INFINITY, Array1::zeros(data.dim));
    }
    let shape = data.gamma_shape.max(1e-10);
    // Hoist shape-only constants out of the per-sample loop: ln Γ(shape) and
    // shape · ln(shape) are independent of i, so previously each sample paid
    // an extra `ln_gamma` and `ln` plus a multiply. n is typically biobank-
    // scale, so this collapses Θ(n) gamma-function evaluations to one.
    let shape_ln_shape = shape * shape.ln();
    let log_gamma_shape = statrs::function::gamma::ln_gamma(shape);
    let shape_minus_one = shape - 1.0;
    let mut residual = Array1::<f64>::zeros(n);
    let ll: f64 = residual
        .as_slice_mut()
        .unwrap()
        .par_iter_mut()
        .enumerate()
        .map(|(i, slot)| {
            let eta_i = eta[i];
            let mu_i = eta_i.exp();
            let y_i = data.y[i];
            let w_i = data.weights[i];
            let ll_i = w_i
                * (shape_ln_shape - log_gamma_shape - shape * eta_i
                    + shape_minus_one * y_i.max(1e-12).ln()
                    - shape * y_i / mu_i);
            *slot = w_i * shape * (y_i / mu_i - 1.0);
            ll_i
        })
        .sum();

    let grad_ll = fast_atv(&data.x, &residual);
    (ll, grad_ll)
}

#[cfg(test)]
mod tests {
    use super::{
        FamilyNutsInputs, GlmFlatInputs, JointBetaRhoInputs, JointBetaRhoPosterior,
        LinkWigglePosterior, LinkWiggleSplineArtifacts, NutsConfig, NutsFamily, NutsPosterior,
        SharedData, cloglog_bernoulli_logp_and_residual, firth_jeffreys_logp_and_grad,
        joint_family_logp_and_grad, laplace_directional_cubic_diagnostic,
        run_joint_beta_rho_sampling, run_logit_polya_gamma_gibbs,
        run_nuts_sampling_flattened_family,
    };
    use crate::construction::CanonicalPenalty;
    use crate::estimate::{
        BlockRole, FitGeometry, FitInference, FittedBlock, FittedLinkState, UnifiedFitResult,
        UnifiedFitResultParts,
    };
    use crate::matrix::DesignMatrix;
    use crate::survival::{MonotonicityPenalty, PenaltyBlocks, SurvivalSpec};
    use crate::types::{
        InverseLink, LikelihoodScaleMetadata, LikelihoodSpec, LinkFunction,
        LogLikelihoodNormalization, ResponseFamily, RhoPrior,
    };
    use general_mcmc::generic_hmc::HamiltonianTarget;
    use ndarray::{Array1, Array2, array};
    use std::sync::Arc;

    impl NutsPosterior {
        /// Test-only allocation wrapper around `compute_logp_and_grad_nd_into`.
        pub(super) fn compute_logp_and_grad_nd(&self, z: &Array1<f64>) -> (f64, Array1<f64>) {
            let mut residual = Array1::<f64>::zeros(self.data.n_samples);
            let mut grad = Array1::<f64>::zeros(z.len());
            let logp = self.compute_logp_and_grad_nd_into(z, &mut residual, &mut grad);
            (logp, grad)
        }
    }

    impl LinkWigglePosterior {
        /// Test-only allocation wrapper around `compute_logp_and_grad_into`.
        pub(super) fn compute_logp_and_grad(&self, z: &Array1<f64>) -> (f64, Array1<f64>) {
            let dim = self.p_base + self.p_link;
            let mut grad = Array1::<f64>::zeros(dim);
            let logp = self.compute_logp_and_grad_into(z, &mut grad);
            (logp, grad)
        }
    }

    impl JointBetaRhoPosterior {
        /// Test-only allocation wrapper around `compute_joint_logp_and_grad_into`.
        pub(super) fn compute_joint_logp_and_grad(
            &self,
            params: &Array1<f64>,
        ) -> (f64, Array1<f64>) {
            let total_dim = self.n_beta + self.n_rho + self.n_link_params;
            let mut grad = Array1::<f64>::zeros(total_dim);
            let logp = self.compute_joint_logp_and_grad_into(params, &mut grad);
            (logp, grad)
        }
    }

    fn hmc_test_fit(
        blocks: Vec<FittedBlock>,
        inference: Option<FitInference>,
        geometry: Option<FitGeometry>,
    ) -> UnifiedFitResult {
        let lambdas = Array1::zeros(0);
        UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
            blocks,
            log_lambdas: lambdas.clone(),
            lambdas,
            likelihood_family: Some(LikelihoodSpec::new(
                ResponseFamily::Gaussian,
                InverseLink::Standard(LinkFunction::Identity),
            )),
            likelihood_scale: LikelihoodScaleMetadata::ProfiledGaussian,
            log_likelihood_normalization: LogLikelihoodNormalization::Full,
            log_likelihood: -1.0,
            deviance: 2.0,
            reml_score: 0.0,
            stable_penalty_term: 0.0,
            penalized_objective: 0.0,
            outer_iterations: 1,
            outer_converged: true,
            outer_gradient_norm: None,
            standard_deviation: 1.0,
            covariance_conditional: None,
            covariance_corrected: None,
            inference,
            fitted_link: FittedLinkState::Standard(None),
            geometry,
            block_states: Vec::new(),
            pirls_status: crate::pirls::PirlsStatus::Converged,
            max_abs_eta: 0.0,
            constraint_kkt: None,
            artifacts: Default::default(),
            inner_cycles: 0,
        })
        .expect("valid HMC handoff test fit")
    }

    #[test]
    fn hmc_whitening_consumes_standard_fit_inference_hessian() {
        let hessian = array![[2.0, 0.1], [0.1, 1.6]];
        let fit = hmc_test_fit(
            vec![FittedBlock {
                beta: array![0.05, -0.1],
                role: BlockRole::Mean,
                edf: 2.0,
                lambdas: Array1::zeros(0),
            }],
            Some(FitInference {
                edf_by_block: vec![],
                edf_total: 2.0,
                smoothing_correction: None,
                penalized_hessian: hessian.clone().into(),
                working_weights: array![1.0, 1.0, 1.0],
                working_response: array![0.0, 0.1, -0.2],
                reparam_qs: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                beta_covariance: None,
                beta_standard_errors: None,
                beta_covariance_corrected: None,
                beta_standard_errors_corrected: None,
                beta_covariance_frequentist: None,
                coefficient_influence: None,
                covariance_is_diagonal_only: false,
                bias_correction_beta: None,
            }),
            None,
        );

        let explicit = super::explicit_fit_hessian_for_whitening(&fit, 2, "standard fit")
            .expect("standard fit exports explicit Hessian");
        assert_eq!(explicit, &hessian);

        let x = array![[1.0, 0.0], [1.0, 0.5], [1.0, -0.5]];
        let y = array![0.0, 0.2, -0.1];
        let weights = Array1::ones(3);
        let penalty = Array2::eye(2);
        NutsPosterior::new(
            x.view(),
            y.view(),
            weights.view(),
            penalty.view(),
            fit.beta.view(),
            explicit.view(),
            NutsFamily::Gaussian,
            1.0,
            crate::estimate::Dispersion::Known(1.0),
            false,
        )
        .expect("HMC target whitens with upstream Hessian");
    }

    #[test]
    fn hmc_whitening_consumes_blockwise_geometry_hessian() {
        let hessian = array![[3.0, 0.2], [0.2, 2.0]];
        let fit = hmc_test_fit(
            vec![
                FittedBlock {
                    beta: array![0.1],
                    role: BlockRole::Location,
                    edf: 1.0,
                    lambdas: Array1::zeros(0),
                },
                FittedBlock {
                    beta: array![-0.2],
                    role: BlockRole::Scale,
                    edf: 1.0,
                    lambdas: Array1::zeros(0),
                },
            ],
            None,
            Some(FitGeometry {
                penalized_hessian: hessian.clone().into(),
                working_weights: array![1.0, 0.8],
                working_response: array![0.0, 0.1],
            }),
        );

        let explicit = super::explicit_fit_hessian_for_whitening(&fit, 2, "blockwise fit")
            .expect("blockwise fit exports materialized Hessian");
        assert_eq!(explicit, &hessian);
    }

    #[test]
    fn hmc_whitening_rejects_covariance_only_fit_without_synthesizing_hessian() {
        let fit = UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
            blocks: vec![FittedBlock {
                beta: array![0.0],
                role: BlockRole::Mean,
                edf: 1.0,
                lambdas: Array1::zeros(0),
            }],
            log_lambdas: Array1::zeros(0),
            lambdas: Array1::zeros(0),
            likelihood_family: Some(LikelihoodSpec::new(
                ResponseFamily::Gaussian,
                InverseLink::Standard(LinkFunction::Identity),
            )),
            likelihood_scale: LikelihoodScaleMetadata::ProfiledGaussian,
            log_likelihood_normalization: LogLikelihoodNormalization::Full,
            log_likelihood: -1.0,
            deviance: 2.0,
            reml_score: 0.0,
            stable_penalty_term: 0.0,
            penalized_objective: 0.0,
            outer_iterations: 1,
            outer_converged: true,
            outer_gradient_norm: None,
            standard_deviation: 1.0,
            covariance_conditional: Some(array![[0.5]]),
            covariance_corrected: None,
            inference: None,
            fitted_link: FittedLinkState::Standard(None),
            geometry: None,
            block_states: Vec::new(),
            pirls_status: crate::pirls::PirlsStatus::Converged,
            max_abs_eta: 0.0,
            constraint_kkt: None,
            artifacts: Default::default(),
            inner_cycles: 0,
        })
        .expect("covariance-only fit can exist for prediction");

        let err = super::explicit_fit_hessian_for_whitening(&fit, 1, "covariance-only fit")
            .expect_err("HMC must not invert covariance as a Hessian fallback");
        assert!(
            err.contains("missing an explicit penalized Hessian"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn log1pexp_is_finite_for_extreme_eta() {
        assert!(NutsPosterior::log1pexp(1000.0).is_finite());
        assert!(NutsPosterior::log1pexp(-1000.0).is_finite());
        assert!((NutsPosterior::log1pexp(-1000.0) - 0.0).abs() < 1e-12);
    }

    #[test]
    fn sigmoid_stable_behaves_at_extremes() {
        let hi = NutsPosterior::sigmoid_stable(1000.0);
        let lo = NutsPosterior::sigmoid_stable(-1000.0);
        assert!((1.0 - 1e-12..=1.0).contains(&hi));
        assert!((0.0..=1e-12).contains(&lo));
    }

    #[test]
    fn cloglog_log_mu_uses_complementary_loglog_inverse_link() {
        let eta = -1.0_f64;
        let (ll_y1, residual_y1) =
            cloglog_bernoulli_logp_and_residual(eta, 1.0).expect("valid eta");
        let expected = (1.0 - (-eta.exp()).exp()).ln();
        let wrong_log_one_minus_exp_eta = (1.0 - eta.exp()).ln();

        assert!((ll_y1 - expected).abs() < 1e-14);
        assert!((ll_y1 - wrong_log_one_minus_exp_eta).abs() > 0.5);

        let eps = 1e-6;
        let (lp, _) = cloglog_bernoulli_logp_and_residual(eta + eps, 1.0).expect("valid eta");
        let (lm, _) = cloglog_bernoulli_logp_and_residual(eta - eps, 1.0).expect("valid eta");
        let fd = (lp - lm) / (2.0 * eps);
        assert!(
            (residual_y1 - fd).abs() < 1e-9,
            "cloglog residual is not the derivative of log μ: analytic={residual_y1}, fd={fd}"
        );
    }

    #[test]
    fn link_wiggle_posterior_whitening_uses_supplied_explicit_joint_hessian() {
        let x = array![[1.0], [1.0], [1.0]];
        let y = array![0.0, 1.0, 1.0];
        let weights = Array1::ones(3);
        let penalty_base = Array2::zeros((1, 1));
        let penalty_link = Array2::zeros((1, 1));
        let mode_beta = array![0.2];
        let mode_theta = array![0.05];
        let hessian = array![[4.0, 1.0], [1.0, 3.0]];
        let spline = LinkWiggleSplineArtifacts {
            knot_range: (-1.0, 1.0),
            knot_vector: Array1::from_vec(vec![-1.0, -1.0, -1.0, 1.0, 1.0, 1.0]),
            degree: 2,
        };

        let posterior = LinkWigglePosterior::new(
            x.view(),
            y.view(),
            weights.view(),
            penalty_base.view(),
            penalty_link.view(),
            mode_beta.view(),
            mode_theta.view(),
            hessian.view(),
            spline,
            NutsFamily::BinomialLogit,
            1.0,
        )
        .expect("link-wiggle posterior should accept explicit SPD joint Hessian");

        let reconstructed_cov = posterior.chol().dot(&posterior.chol().t());
        let eye_from_hessian = hessian.dot(&reconstructed_cov);
        for r in 0..2 {
            for c in 0..2 {
                let expected = if r == c { 1.0 } else { 0.0 };
                assert!(
                    (eye_from_hessian[[r, c]] - expected).abs() < 1e-10,
                    "whitening did not use the supplied explicit joint Hessian at ({r},{c}): got {} expected {}",
                    eye_from_hessian[[r, c]],
                    expected
                );
            }
        }
    }

    #[test]
    fn link_wiggle_cloglog_gradient_matches_its_log_likelihood() {
        let x = array![[1.0], [1.0], [1.0], [1.0]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let weights = array![1.0, 1.2, 0.8, 1.4];
        let penalty_base = Array2::zeros((1, 1));
        let penalty_link = Array2::zeros((1, 1));
        let mode_beta = array![-0.8];
        let mode_theta = array![0.04];
        let hessian = Array2::eye(2);
        let spline = LinkWiggleSplineArtifacts {
            knot_range: (-1.5, 0.5),
            knot_vector: Array1::from_vec(vec![-1.5, -1.5, -1.5, 0.5, 0.5, 0.5]),
            degree: 2,
        };

        let posterior = LinkWigglePosterior::new(
            x.view(),
            y.view(),
            weights.view(),
            penalty_base.view(),
            penalty_link.view(),
            mode_beta.view(),
            mode_theta.view(),
            hessian.view(),
            spline,
            NutsFamily::BinomialCLogLog,
            1.0,
        )
        .expect("cloglog link-wiggle posterior");

        let z = array![0.2, -0.03];
        let (_, grad) = posterior.compute_logp_and_grad(&z);
        let eps = 1e-6;
        for j in 0..z.len() {
            let mut z_plus = z.clone();
            let mut z_minus = z.clone();
            z_plus[j] += eps;
            z_minus[j] -= eps;
            let (lp, _) = posterior.compute_logp_and_grad(&z_plus);
            let (lm, _) = posterior.compute_logp_and_grad(&z_minus);
            let fd = (lp - lm) / (2.0 * eps);
            assert!(
                (grad[j] - fd).abs() < 1e-6,
                "link-wiggle cloglog gradient mismatch at {j}: analytic={}, fd={}",
                grad[j],
                fd
            );
        }
    }

    #[test]
    fn nuts_logitgradient_matches_finite_difference() {
        let x = array![[1.0, -0.5], [0.2, 0.7], [-1.0, 0.3], [0.5, -1.2]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.5, 0.8, 1.2];
        let penalty = array![[0.4, 0.0], [0.0, 0.6]];
        let mode = array![0.1, -0.2];
        let hessian = array![[2.0, 0.2], [0.2, 1.7]]; // SPD

        let posterior = NutsPosterior::new(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            hessian.view(),
            NutsFamily::BinomialLogit,
            1.0,
            crate::estimate::Dispersion::Known(1.0),
            true,
        )
        .expect("posterior");

        let z = array![0.15, -0.35];
        let (_, grad) = posterior.compute_logp_and_grad_nd(&z);

        let eps = 1e-6;
        for j in 0..z.len() {
            let mut z_plus = z.clone();
            let mut z_minus = z.clone();
            z_plus[j] += eps;
            z_minus[j] -= eps;
            let (lp, _) = posterior.compute_logp_and_grad_nd(&z_plus);
            let (lm, _) = posterior.compute_logp_and_grad_nd(&z_minus);
            let fd = (lp - lm) / (2.0 * eps);
            assert_eq!(
                grad[j].signum(),
                fd.signum(),
                "gradient sign mismatch at {}: analytic={}, fd={}",
                j,
                grad[j],
                fd
            );
            assert!(
                (grad[j] - fd).abs() < 1e-5,
                "gradient mismatch at {}: analytic={}, fd={}",
                j,
                grad[j],
                fd
            );
        }
    }

    #[test]
    fn gamma_log_logp_and_grad_uses_fitted_shape() {
        let x = array![[1.0_f64], [1.0_f64]];
        let y = array![1.5_f64, 2.5_f64];
        let weights = array![1.0_f64, 2.0_f64];
        let eta = array![0.2_f64, 0.4_f64];
        let shape = 3.5_f64;
        let data = SharedData {
            x: Arc::new(x.clone()),
            y: Arc::new(y.clone()),
            weights: Arc::new(weights.clone()),
            mode: Arc::new(Array1::zeros(1)),
            gamma_shape: shape,
            dispersion: crate::estimate::Dispersion::Known(1.0),
            n_samples: x.nrows(),
            dim: x.ncols(),
        };

        let (ll, grad) = super::gamma_log_logp_and_grad(&data, &eta);

        let mut expected_ll = 0.0;
        let mut expected_score = 0.0;
        for i in 0..eta.len() {
            let mu = eta[i].exp();
            expected_ll += weights[i]
                * (shape * shape.ln() - statrs::function::gamma::ln_gamma(shape) - shape * eta[i]
                    + (shape - 1.0) * y[i].ln()
                    - shape * y[i] / mu);
            expected_score += weights[i] * shape * (y[i] / mu - 1.0);
        }

        assert!((ll - expected_ll).abs() < 1e-12);
        assert_eq!(grad.len(), 1);
        assert!((grad[0] - expected_score).abs() < 1e-12);
    }

    #[test]
    fn firth_jeffreys_logit_is_finite_for_rank_deficient_design() {
        let x = array![
            [1.0, -0.5, 1.0],
            [1.0, 0.3, 1.0],
            [1.0, 0.8, 1.0],
            [1.0, -1.2, 1.0],
        ];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let weights = array![1.0, 2.0, 0.5, 1.5];
        let eta = array![0.2, -0.1, 0.4, -0.3];

        let data = SharedData {
            x: Arc::new(x.clone()),
            y: Arc::new(y),
            weights: Arc::new(weights.clone()),
            mode: Arc::new(Array1::zeros(x.ncols())),
            gamma_shape: 1.0,
            dispersion: crate::estimate::Dispersion::Known(1.0),
            n_samples: x.nrows(),
            dim: x.ncols(),
        };

        let (value, grad) =
            firth_jeffreys_logp_and_grad(NutsFamily::BinomialLogit, &data, &eta).expect("firth");

        assert!(value.is_finite());
        assert_eq!(grad.len(), x.ncols());
        assert!(grad.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn logit_pg_gibbs_returns_finite_samples() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let cfg = NutsConfig {
            n_samples: 30,
            nwarmup: 30,
            n_chains: 2,
            target_accept: 0.8,
            seed: 123,
        };
        let out = run_logit_polya_gamma_gibbs(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            &cfg,
        )
        .expect("pg gibbs should run");
        assert_eq!(out.samples.ncols(), 2);
        assert_eq!(out.samples.nrows(), cfg.n_samples * cfg.n_chains);
        assert!(out.samples.iter().all(|v| v.is_finite()));
        assert!(out.posterior_mean.iter().all(|v| v.is_finite()));
        assert!(out.posterior_std.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn family_pg_dispatch_rejects_non_bernoulli_response() {
        let x = array![[1.0], [1.0]];
        let y = array![2.0, 0.0];
        let w = array![1.0, 1.0];
        let penalty = array![[0.1]];
        let mode = array![0.0];
        let non_spd_hessian = array![[0.0]];
        let cfg = NutsConfig {
            n_samples: 1,
            nwarmup: 1,
            n_chains: 1,
            target_accept: 0.8,
            seed: 321,
        };

        let result = run_nuts_sampling_flattened_family(
            LikelihoodSpec::binomial_logit(),
            FamilyNutsInputs::Glm(GlmFlatInputs {
                x: x.view(),
                y: y.view(),
                weights: w.view(),
                penalty_matrix: penalty.view(),
                mode: mode.view(),
                hessian: non_spd_hessian.view(),
                gamma_shape: None,
                dispersion: crate::solver::estimate::Dispersion::Known(1.0),
                firth_bias_reduction: false,
            }),
            &cfg,
        );

        let err = result.err().expect("PG dispatch should reject count rows");
        assert!(
            err.contains("response must be exactly 0 or 1"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn family_dispatch_uses_pg_gibbs_for_standard_logit() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let non_spdhessian = array![[0.0, 0.0], [0.0, 0.0]];
        let cfg = NutsConfig {
            n_samples: 20,
            nwarmup: 20,
            n_chains: 2,
            target_accept: 0.8,
            seed: 456,
        };
        let out = run_nuts_sampling_flattened_family(
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Logit),
            },
            FamilyNutsInputs::Glm(GlmFlatInputs {
                x: x.view(),
                y: y.view(),
                weights: w.view(),
                penalty_matrix: penalty.view(),
                mode: mode.view(),
                hessian: non_spdhessian.view(),
                gamma_shape: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                firth_bias_reduction: false,
            }),
            &cfg,
        )
        .expect("dispatch should use PG Gibbs and not require Hessian factorization");
        assert_eq!(out.samples.nrows(), cfg.n_samples * cfg.n_chains);
        assert!(out.samples.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn family_dispatch_routes_probit_to_nuts_path() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let non_spdhessian = array![[0.0, 0.0], [0.0, 0.0]];
        let cfg = NutsConfig {
            n_samples: 20,
            nwarmup: 20,
            n_chains: 2,
            target_accept: 0.8,
            seed: 654,
        };

        let err = match run_nuts_sampling_flattened_family(
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Probit),
            },
            FamilyNutsInputs::Glm(GlmFlatInputs {
                x: x.view(),
                y: y.view(),
                weights: w.view(),
                penalty_matrix: penalty.view(),
                mode: mode.view(),
                hessian: non_spdhessian.view(),
                gamma_shape: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                firth_bias_reduction: false,
            }),
            &cfg,
        ) {
            Ok(_) => panic!("non-SPD Hessian should fail after probit routes to the NUTS path"),
            Err(err) => err,
        };

        assert!(
            err.contains("Hessian Cholesky decomposition failed"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn family_dispatch_rejects_nonlogit_firth_family() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 2.0, 0.0, 3.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let hessian = array![[1.5, 0.1], [0.1, 1.2]];
        let cfg = NutsConfig {
            n_samples: 20,
            nwarmup: 20,
            n_chains: 2,
            target_accept: 0.8,
            seed: 111,
        };

        let err = match run_nuts_sampling_flattened_family(
            LikelihoodSpec {
                response: ResponseFamily::Poisson,
                link: InverseLink::Standard(LinkFunction::Log),
            },
            FamilyNutsInputs::Glm(GlmFlatInputs {
                x: x.view(),
                y: y.view(),
                weights: w.view(),
                penalty_matrix: penalty.view(),
                mode: mode.view(),
                hessian: hessian.view(),
                gamma_shape: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                firth_bias_reduction: true,
            }),
            &cfg,
        ) {
            Ok(_) => panic!("Poisson Firth should be rejected explicitly"),
            Err(err) => err,
        };

        assert!(
            err.contains("NUTS with Firth is only supported for Binomial Logit"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn run_nuts_sampling_rejects_invalid_target_accept() {
        let x = array![[1.0], [1.0], [1.0]];
        let y = array![0.5, -0.5, 1.0];
        let weights = array![1.0, 1.0, 1.0];
        let penalty = array![[0.25]];
        let mode = array![0.0];
        let hessian = array![[1.25]];
        let cfg = NutsConfig {
            n_samples: 10,
            nwarmup: 10,
            n_chains: 1,
            target_accept: 1.0,
            seed: 222,
        };

        let err = super::run_nuts_sampling(
            x.view(),
            y.view(),
            weights.view(),
            penalty.view(),
            mode.view(),
            hessian.view(),
            NutsFamily::Gaussian,
            1.0,
            crate::estimate::Dispersion::Known(1.0),
            false,
            &cfg,
        )
        .expect_err("invalid target_accept should be rejected before sampling");

        assert!(
            err.contains("target_accept must be finite and lie in (0, 1)"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn joint_hmc_boundary_rejects_nonlogit_firth_family() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 2.0, 0.0, 3.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let hessian = array![[1.5, 0.1], [0.1, 1.2]];
        let penalty_root = array![[0.4, 0.0], [0.0, 0.6]];
        let mode = array![0.0, 0.0];
        let rho_mode = array![0.0];
        let cfg = NutsConfig {
            n_samples: 20,
            nwarmup: 20,
            n_chains: 2,
            target_accept: 0.8,
            seed: 111,
        };

        let inputs = JointBetaRhoInputs {
            x: x.view(),
            y: y.view(),
            weights: w.view(),
            likelihood: LikelihoodSpec {
                response: ResponseFamily::Poisson,
                link: InverseLink::Standard(LinkFunction::Log),
            },
            gamma_shape: None,
            mode: mode.view(),
            hessian: hessian.view(),
            penalty_roots: vec![CanonicalPenalty::from_dense_root(
                penalty_root.clone(),
                penalty_root.ncols(),
            )],
            rho_mode: rho_mode.view(),
            rho_prior: RhoPrior::default(),
            firth_bias_reduction: true,
            trigger_skewness: 0.75,
        };

        let err = match run_joint_beta_rho_sampling(&inputs, &cfg) {
            Ok(_) => panic!("Poisson joint HMC Firth should be rejected explicitly"),
            Err(err) => err,
        };

        assert!(
            err.contains("Joint HMC with Firth is only supported for Binomial Logit"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn joint_hmc_uses_combined_penalty_logdet_for_overlapping_penalties() {
        let x = array![[0.0, 0.0]];
        let y = array![0.0];
        let w = array![0.0];
        let mode = array![0.0, 0.0];
        let hessian = array![[1.0, 0.0], [0.0, 1.0]];
        let rho_mode = array![0.0, 0.0];
        let penalty_1 = array![[1.0, 0.0], [0.0, 1.0]];
        let penalty_2 = array![[2.0_f64.sqrt(), 0.0], [0.0, 1.0]];
        let target = JointBetaRhoPosterior::new(
            x.view(),
            y.view(),
            w.view(),
            mode.view(),
            hessian.view(),
            vec![
                CanonicalPenalty::from_dense_root(penalty_1, 2),
                CanonicalPenalty::from_dense_root(penalty_2, 2),
            ],
            rho_mode.view(),
            LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            None,
            RhoPrior::Flat,
            false,
        )
        .expect("joint target");

        let params = array![0.0, 0.0, 0.0, 0.0];
        let (_, grad) = target.compute_joint_logp_and_grad(&params);
        assert!(
            (grad[2] - 5.0 / 12.0).abs() < 1.0e-10,
            "expected overlapping-penalty gradient 5/12, got {}",
            grad[2]
        );
        assert!(
            (grad[3] - 7.0 / 12.0).abs() < 1.0e-10,
            "expected overlapping-penalty gradient 7/12, got {}",
            grad[3]
        );
    }

    #[test]
    fn joint_hmc_target_does_not_depend_on_rho_mode_when_prior_is_fixed() {
        let x = array![[0.0]];
        let y = array![0.0];
        let w = array![0.0];
        let mode = array![0.0];
        let hessian = array![[1.0]];
        let penalty = CanonicalPenalty::from_dense_root(array![[1.0]], 1);
        let prior = RhoPrior::Normal {
            mean: 0.25,
            sd: 1.7,
        };

        let target_a = JointBetaRhoPosterior::new(
            x.view(),
            y.view(),
            w.view(),
            mode.view(),
            hessian.view(),
            vec![penalty.clone()],
            array![0.0].view(),
            LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            None,
            prior.clone(),
            false,
        )
        .expect("target a");
        let target_b = JointBetaRhoPosterior::new(
            x.view(),
            y.view(),
            w.view(),
            mode.view(),
            hessian.view(),
            vec![penalty],
            array![2.5].view(),
            LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            None,
            prior,
            false,
        )
        .expect("target b");

        let params = array![0.0, -0.4];
        let (lp_a, grad_a) = target_a.compute_joint_logp_and_grad(&params);
        let (lp_b, grad_b) = target_b.compute_joint_logp_and_grad(&params);
        assert!((lp_a - lp_b).abs() < 1.0e-12);
        for i in 0..grad_a.len() {
            assert!(
                (grad_a[i] - grad_b[i]).abs() < 1.0e-12,
                "rho_mode leaked into target gradient at {}: {} vs {}",
                i,
                grad_a[i],
                grad_b[i]
            );
        }
    }

    #[test]
    fn joint_hmc_binomial_sas_uses_runtime_link_state() {
        let x = array![[1.0], [1.0]];
        let y = array![1.0, 0.0];
        let weights = array![1.0, 1.0];
        let eta = array![0.3, -0.2];
        let sas_state = crate::mixture_link::state_from_sasspec(crate::types::SasLinkSpec {
            initial_epsilon: 0.4,
            initial_log_delta: -0.2,
        })
        .expect("sas state");
        let data = SharedData {
            x: Arc::new(x),
            y: Arc::new(y),
            weights: Arc::new(weights),
            mode: Arc::new(Array1::zeros(1)),
            gamma_shape: 1.0,
            dispersion: crate::estimate::Dispersion::Known(1.0),
            n_samples: 2,
            dim: 1,
        };

        let (ll_sas, _) = joint_family_logp_and_grad(
            &LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Sas(sas_state),
            },
            &data,
            &eta,
        )
        .expect("sas joint logp");
        let (ll_logit, _) = joint_family_logp_and_grad(
            &LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Logit),
            },
            &data,
            &eta,
        )
        .expect("logit joint logp");

        assert!(
            (ll_sas - ll_logit).abs() > 1.0e-6,
            "adaptive SAS link should not collapse to the logit likelihood"
        );
    }

    #[test]
    fn directional_cubic_diagnostic_is_rotation_invariant_for_hessian_eigenvectors() {
        let x = array![[1.0, 0.5], [-0.3, 1.4], [0.8, -1.1]];
        let c = array![0.7, -0.5, 0.2];
        let h = array![[4.0, 0.0], [0.0, 1.0]];
        let theta = std::f64::consts::FRAC_PI_4;
        let q = array![[theta.cos(), -theta.sin()], [theta.sin(), theta.cos()],];
        let x_rot = x.dot(&q);
        let h_rot = q.t().dot(&h).dot(&q);

        let (base_max, base_vals) = laplace_directional_cubic_diagnostic(
            &h,
            &DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(x)),
            &c,
        )
        .expect("base diagnostic");
        let (rot_max, rot_vals) = laplace_directional_cubic_diagnostic(
            &h_rot,
            &DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(x_rot)),
            &c,
        )
        .expect("rotated diagnostic");

        let mut base_abs: Vec<f64> = base_vals.iter().map(|v| v.abs()).collect();
        let mut rot_abs: Vec<f64> = rot_vals.iter().map(|v| v.abs()).collect();
        base_abs.sort_by(|a, b| a.partial_cmp(b).expect("finite compare"));
        rot_abs.sort_by(|a, b| a.partial_cmp(b).expect("finite compare"));

        assert!((base_max - rot_max).abs() < 1.0e-10);
        for i in 0..base_abs.len() {
            assert!(
                (base_abs[i] - rot_abs[i]).abs() < 1.0e-10,
                "directional diagnostic changed under rotation at {}: {} vs {}",
                i,
                base_abs[i],
                rot_abs[i]
            );
        }
    }

    /// Verify that joint HMC and REML compute identical penalty logdet
    /// derivatives for the same penalty system. This catches any divergence
    /// between the two code paths.
    #[test]
    fn joint_hmc_penalty_logdet_agrees_with_reml_path() {
        use crate::estimate::reml::penalty_logdet::PenaltyPseudologdet;

        // Two overlapping 3x3 penalties with non-trivial lambdas.
        let root_1 = array![[1.0, 0.5, 0.0], [0.0, 0.8, 0.3]];
        let root_2 = array![[0.0, 0.7, 0.0], [0.0, 0.0, 1.2]];
        let cp1 = CanonicalPenalty::from_dense_root(root_1, 3);
        let cp2 = CanonicalPenalty::from_dense_root(root_2, 3);
        let lambdas = [2.5_f64, 0.8];
        let penalties = [cp1.clone(), cp2.clone()];

        // REML path: PenaltyPseudologdet directly.
        let pld =
            PenaltyPseudologdet::from_penalties(&penalties, &lambdas, 0.0, 3).expect("reml pld");
        let reml_value = pld.value();
        let (reml_d1, reml_d2) = pld.rho_derivatives_from_penalties(&penalties, &lambdas);

        // Joint HMC path: build a JointBetaRhoPosterior and extract the
        // penalty logdet contribution. We isolate it by using zero data
        // (so likelihood = 0, penalty quadratic = 0) and Flat rho prior.
        let x = Array2::<f64>::zeros((1, 3));
        let y = array![0.0];
        let w = array![0.0];
        let mode = Array1::<f64>::zeros(3);
        let hessian = Array2::<f64>::eye(3);
        let rho = Array1::from_vec(lambdas.iter().map(|l| l.ln()).collect());
        let target = JointBetaRhoPosterior::new(
            x.view(),
            y.view(),
            w.view(),
            mode.view(),
            hessian.view(),
            vec![cp1, cp2],
            rho.view(),
            LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            None,
            RhoPrior::Flat,
            false,
        )
        .expect("joint target");

        // Evaluate at beta=0, rho=ln(lambdas).
        let mut params = Array1::<f64>::zeros(3 + 2);
        params[3] = rho[0];
        params[4] = rho[1];
        let (logp, grad) = target.compute_joint_logp_and_grad(&params);

        // logp should be 0.5 * reml_value (likelihood=0, prior=0, quadratic=0).
        assert!(
            (logp - 0.5 * reml_value).abs() < 1.0e-8,
            "joint HMC logdet value {} vs REML 0.5*{} = {}",
            logp,
            reml_value,
            0.5 * reml_value,
        );

        // grad[3..5] should be 0.5 * reml_d1.
        for k in 0..2 {
            assert!(
                (grad[3 + k] - 0.5 * reml_d1[k]).abs() < 1.0e-8,
                "joint HMC logdet gradient[{}] = {} vs REML 0.5*{} = {}",
                k,
                grad[3 + k],
                reml_d1[k],
                0.5 * reml_d1[k],
            );
        }

        // Sanity: second derivatives are available from REML but not directly
        // from a single HMC gradient call; just verify they're symmetric.
        assert!(
            (reml_d2[[0, 1]] - reml_d2[[1, 0]]).abs() < 1.0e-12,
            "REML penalty logdet Hessian not symmetric"
        );
    }

    /// Verify the family-gating invariant: every LikelihoodSpec that
    /// joint_family_logp_and_grad accepts produces a result (not an error
    /// about missing implementation). Every family it rejects returns an
    /// explicit error. No family is silently remapped to a different one.
    #[test]
    fn joint_hmc_family_gating_never_remaps() {
        let data = SharedData {
            x: Arc::new(array![[1.0], [1.0]]),
            y: Arc::new(array![1.0, 0.0]),
            weights: Arc::new(array![1.0, 1.0]),
            mode: Arc::new(Array1::zeros(1)),
            gamma_shape: 1.0,
            dispersion: crate::estimate::Dispersion::Known(1.0),
            n_samples: 2,
            dim: 1,
        };
        let eta = array![0.1, -0.1];

        // These families must succeed with their own inverse link.
        let accepted = [
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Logit),
            },
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::Probit),
            },
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Standard(LinkFunction::CLogLog),
            },
            LikelihoodSpec {
                response: ResponseFamily::Gaussian,
                link: InverseLink::Standard(LinkFunction::Identity),
            },
            LikelihoodSpec {
                response: ResponseFamily::Poisson,
                link: InverseLink::Standard(LinkFunction::Log),
            },
            LikelihoodSpec {
                response: ResponseFamily::Gamma,
                link: InverseLink::Standard(LinkFunction::Log),
            },
        ];
        for spec in &accepted {
            let result = joint_family_logp_and_grad(spec, &data, &eta);
            assert!(
                result.is_ok(),
                "spec {:?} should be accepted but got error: {:?}",
                spec,
                result.err(),
            );
        }

        // SAS/BetaLogistic/Mixture must succeed with their real link state,
        // NOT be remapped to logit.
        let sas_state = crate::mixture_link::state_from_sasspec(crate::types::SasLinkSpec {
            initial_epsilon: 0.0,
            initial_log_delta: 0.0,
        })
        .expect("sas state");
        let adaptive = [
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::Sas(sas_state),
            },
            LikelihoodSpec {
                response: ResponseFamily::Binomial,
                link: InverseLink::BetaLogistic(
                    crate::mixture_link::state_from_sasspec(crate::types::SasLinkSpec {
                        initial_epsilon: 0.0,
                        initial_log_delta: 0.0,
                    })
                    .expect("bl state"),
                ),
            },
        ];
        for spec in &adaptive {
            let result = joint_family_logp_and_grad(spec, &data, &eta);
            assert!(
                result.is_ok(),
                "adaptive spec {:?} should be accepted with its real link",
                spec,
            );
        }

        // RoystonParmar must be explicitly rejected (not silently remapped).
        let rp_result = joint_family_logp_and_grad(
            &LikelihoodSpec {
                response: ResponseFamily::RoystonParmar,
                link: InverseLink::Standard(LinkFunction::Logit),
            },
            &data,
            &eta,
        );
        assert!(
            rp_result.is_err(),
            "RoystonParmar should be rejected, not silently accepted"
        );
    }

    /// The power-iteration refinement should find non-Gaussianity at least
    /// as large as the eigenvector-only pass (it's a supremum search).
    #[test]
    fn directional_cubic_power_iteration_finds_larger_or_equal_skewness() {
        // Construct a design where the maximum |gamma| occurs off-axis.
        // A single row with asymmetric structure makes the cubic form
        // peak between eigenvectors.
        let x = array![
            [2.0, 1.0],
            [-1.0, 2.0],
            [0.5, -0.5],
            [1.5, 0.3],
            [-0.8, 1.7],
        ];
        let c = array![1.0, -0.5, 0.3, -0.7, 0.4];
        let h = array![[3.0, 1.0], [1.0, 2.0]];

        let (max_val, eigenvector_vals) = laplace_directional_cubic_diagnostic(
            &h,
            &DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(x)),
            &c,
        )
        .expect("diagnostic");

        // max_val should be >= max of eigenvector-only values.
        let eig_max = eigenvector_vals
            .iter()
            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
        assert!(
            max_val >= eig_max - 1.0e-12,
            "power iteration result {} should be >= eigenvector max {}",
            max_val,
            eig_max,
        );
    }

    #[test]
    fn logit_pg_rao_blackwell_returns_finite_terms() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let roots = vec![array![[0.2_f64.sqrt(), 0.0], [0.0, 0.4_f64.sqrt()]]];
        let cfg = NutsConfig {
            n_samples: 30,
            nwarmup: 30,
            n_chains: 2,
            target_accept: 0.8,
            seed: 789,
        };

        let rb = super::estimate_logit_pg_rao_blackwell_terms(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            &roots,
            &cfg,
        )
        .expect("rao-blackwell PG should run");

        assert_eq!(rb.len(), 1);
        assert!(rb[0].is_finite());
        assert!(rb[0] >= 0.0);
    }

    #[test]
    fn logit_pg_rao_blackwell_rejects_non_bernoulli_response() {
        let x = array![[1.0], [1.0]];
        let y = array![0.25, 1.0];
        let w = array![1.0, 1.0];
        let penalty = array![[0.1]];
        let mode = array![0.0];
        let roots = vec![array![[0.1_f64.sqrt()]]];
        let cfg = NutsConfig {
            n_samples: 1,
            nwarmup: 1,
            n_chains: 1,
            target_accept: 0.8,
            seed: 654,
        };

        let result = super::estimate_logit_pg_rao_blackwell_terms(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            &roots,
            &cfg,
        );

        let err = result
            .err()
            .expect("PG Rao-Blackwell should reject proportion rows");
        assert!(
            err.contains("response must be exactly 0 or 1"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn logit_pg_rao_blackwell_matches_beta_quadratic_moment_sanity() {
        let x = array![[1.0, 0.2], [1.0, -0.1], [1.0, 1.2], [1.0, -0.7]];
        let y = array![1.0, 0.0, 1.0, 0.0];
        let w = array![1.0, 1.0, 1.0, 1.0];
        let penalty = array![[0.2, 0.0], [0.0, 0.4]];
        let mode = array![0.0, 0.0];
        let roots = vec![array![[0.2_f64.sqrt(), 0.0], [0.0, 0.4_f64.sqrt()]]];
        let cfg = NutsConfig {
            n_samples: 120,
            nwarmup: 80,
            n_chains: 2,
            target_accept: 0.8,
            seed: 901,
        };

        let gibbs = run_logit_polya_gamma_gibbs(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            &cfg,
        )
        .expect("pg gibbs should run");
        let mc_quad = gibbs
            .samples
            .rows()
            .into_iter()
            .map(|beta| {
                let sb = penalty.dot(&beta.to_owned());
                beta.dot(&sb)
            })
            .sum::<f64>()
            / (gibbs.samples.nrows() as f64);

        let rb = super::estimate_logit_pg_rao_blackwell_terms(
            x.view(),
            y.view(),
            w.view(),
            penalty.view(),
            mode.view(),
            &roots,
            &cfg,
        )
        .expect("rao-blackwell PG should run");

        let diff = (rb[0] - mc_quad).abs();
        assert!(
            diff < 0.35,
            "Rao-Blackwell vs beta-moment mismatch too large: rb={}, mc={}, diff={}",
            rb[0],
            mc_quad,
            diff
        );
    }

    #[test]
    fn survival_hmc_structural_monotonic_returns_finitevalues() {
        let age_entry = array![1.0];
        let age_exit = array![2.0];
        let event_target = array![1u8];
        let event_competing = array![0u8];
        let sampleweight = array![1.0];
        let x_entry = array![[1.0, 0.2]];
        let x_exit = array![[1.0, 0.6]];
        let x_derivative = array![[0.0, 1.0]];
        let penalties = PenaltyBlocks::new(Vec::new());
        let monotonicity = MonotonicityPenalty { tolerance: 3.0 };
        let mode = array![0.0, 0.0];
        let hessian = Array2::<f64>::eye(2);

        let posterior = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            None,
            penalties,
            monotonicity,
            SurvivalSpec::Net,
            true,
            2,
            mode.view(),
            hessian.view(),
        )
        .expect("construct survival posterior");

        let position = array![0.0, 0.0];
        let mut grad = Array1::<f64>::zeros(2);
        let logp = HamiltonianTarget::logp_and_grad(&posterior, &position, &mut grad);
        assert!(logp.is_finite());
        assert!(grad.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn survival_hmc_structural_monotonic_differs_from_linear_geometry() {
        let age_entry = array![1.0];
        let age_exit = array![2.0];
        let event_target = array![1u8];
        let event_competing = array![0u8];
        let sampleweight = array![1.0];
        let x_entry = array![[0.2, 0.1]];
        let x_exit = array![[0.6, 0.3]];
        let x_derivative = array![[1.0, 0.0]];
        let monotonicity = MonotonicityPenalty { tolerance: 3.0 };
        let mode = array![0.0, 0.0];
        let hessian = Array2::<f64>::eye(2);
        let z = array![std::f64::consts::LN_2, 0.0];

        let posterior_linear = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            None,
            PenaltyBlocks::new(Vec::new()),
            monotonicity,
            SurvivalSpec::Net,
            false,
            0,
            mode.view(),
            hessian.view(),
        )
        .expect("construct linear posterior");
        let mut grad_linear = Array1::<f64>::zeros(2);
        HamiltonianTarget::logp_and_grad(&posterior_linear, &z, &mut grad_linear);

        let posterior_struct = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            None,
            PenaltyBlocks::new(Vec::new()),
            monotonicity,
            SurvivalSpec::Net,
            true,
            2,
            mode.view(),
            hessian.view(),
        )
        .expect("construct structural posterior");
        let mut grad_struct = Array1::<f64>::zeros(2);
        HamiltonianTarget::logp_and_grad(&posterior_struct, &z, &mut grad_struct);

        assert!(
            (grad_struct[0] - grad_linear[0]).abs() > 1e-6,
            "expected structural and linear fallback gradients to differ"
        );
        assert!(grad_struct[0].is_finite());
        assert!(grad_linear[0].is_finite());
    }

    #[test]
    fn survival_hmc_fallback_barrier_rejects_offsets_below_monotonicity_threshold() {
        let age_entry = array![1.0];
        let age_exit = array![2.0];
        let event_target = array![1u8];
        let event_competing = array![0u8];
        let sampleweight = array![1.0];
        let x_entry = array![[1.0, 0.0]];
        let x_exit = array![[1.0, 0.0]];
        // Zero derivative design so derivative_offset_exit drives d_eta/dt.
        let x_derivative = array![[0.0, 0.0]];
        let penalties = PenaltyBlocks::new(Vec::new());
        let monotonicity = MonotonicityPenalty { tolerance: 3.0 };
        let mode = array![0.0, 0.0];
        let hessian = Array2::<f64>::eye(2);
        let z = array![0.0, 0.0];

        let posterior_no_offset = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            Some(array![0.0].view()),
            penalties.clone(),
            monotonicity,
            SurvivalSpec::Net,
            false,
            0,
            mode.view(),
            hessian.view(),
        )
        .expect("construct posterior without derivative offset");
        let mut grad_no_offset = Array1::<f64>::zeros(2);
        let logp_no_offset =
            HamiltonianTarget::logp_and_grad(&posterior_no_offset, &z, &mut grad_no_offset);

        let posteriorwith_offset = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            Some(array![2.0].view()),
            penalties,
            monotonicity,
            SurvivalSpec::Net,
            false,
            0,
            mode.view(),
            hessian.view(),
        )
        .expect("construct posterior with derivative offset");
        let mut gradwith_offset = Array1::<f64>::zeros(2);
        let logpwith_offset =
            HamiltonianTarget::logp_and_grad(&posteriorwith_offset, &z, &mut gradwith_offset);

        assert!(!logp_no_offset.is_finite());
        assert!(!logpwith_offset.is_finite());
        assert!(grad_no_offset.iter().all(|v| *v == 0.0));
        assert!(gradwith_offset.iter().all(|v| *v == 0.0));
    }

    #[test]
    fn survival_hmc_fallback_barrier_becomes_finite_once_offset_clears_guard() {
        let age_entry = array![1.0];
        let age_exit = array![2.0];
        let event_target = array![1u8];
        let event_competing = array![0u8];
        let sampleweight = array![1.0];
        let x_entry = array![[1.0, 0.0]];
        let x_exit = array![[1.0, 0.0]];
        let x_derivative = array![[0.0, 0.0]];
        let penalties = PenaltyBlocks::new(Vec::new());
        let monotonicity = MonotonicityPenalty { tolerance: 3.0 };
        let mode = array![0.0, 0.0];
        let hessian = Array2::<f64>::eye(2);
        let z = array![0.0, 0.0];

        let posterior_below_guard = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            Some(array![2.0].view()),
            penalties.clone(),
            monotonicity,
            SurvivalSpec::Net,
            false,
            0,
            mode.view(),
            hessian.view(),
        )
        .expect("construct posterior below derivative guard");
        let mut grad_below_guard = Array1::<f64>::zeros(2);
        let logp_below_guard =
            HamiltonianTarget::logp_and_grad(&posterior_below_guard, &z, &mut grad_below_guard);

        let posterior_above_guard = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            Some(array![3.1].view()),
            penalties,
            monotonicity,
            SurvivalSpec::Net,
            false,
            0,
            mode.view(),
            hessian.view(),
        )
        .expect("construct posterior above derivative guard");
        let mut grad_above_guard = Array1::<f64>::zeros(2);
        let logp_above_guard =
            HamiltonianTarget::logp_and_grad(&posterior_above_guard, &z, &mut grad_above_guard);

        assert!(!logp_below_guard.is_finite());
        assert!(logp_above_guard.is_finite());
        assert!(grad_below_guard.iter().all(|v| *v == 0.0));
        assert!(grad_above_guard.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn survival_hmc_structural_monotonic_handles_sparse_multirow_geometry() {
        let age_entry = array![1.0, 1.2];
        let age_exit = array![2.0, 2.4];
        let event_target = array![1u8, 1u8];
        let event_competing = array![0u8, 0u8];
        let sampleweight = array![1.0, 1.0];
        let x_entry = array![[0.1, 0.0, 0.2], [0.2, 0.1, 0.2]];
        let x_exit = array![[0.4, 0.2, 0.3], [0.6, 0.1, 0.3]];
        // First row constrains only column 0, second row constrains columns 0 and 1.
        let x_derivative = array![[1.0, 0.0, 0.0], [0.5, 1.0, 0.0]];
        let monotonicity = MonotonicityPenalty { tolerance: 3.0 };
        let mode = array![4.0, 2.0, 0.0];
        let hessian = Array2::<f64>::eye(3);
        let z = array![0.05, -0.1, 0.15];

        let posterior = super::survival_hmc::SurvivalPosterior::new(
            age_entry.view(),
            age_exit.view(),
            event_target.view(),
            event_competing.view(),
            sampleweight.view(),
            x_entry.view(),
            x_exit.view(),
            x_derivative.view(),
            None,
            None,
            None,
            PenaltyBlocks::new(Vec::new()),
            monotonicity,
            SurvivalSpec::Net,
            true,
            2,
            mode.view(),
            hessian.view(),
        )
        .expect("construct structural posterior");

        let mut grad = Array1::<f64>::zeros(3);
        let logp = HamiltonianTarget::logp_and_grad(&posterior, &z, &mut grad);
        assert!(logp.is_finite());
        assert!(grad.iter().all(|v| v.is_finite()));
    }
}

/// Implement HamiltonianTarget for NUTS with analytical gradients.
impl HamiltonianTarget<Array1<f64>> for NutsPosterior {
    fn logp_and_grad(&self, position: &Array1<f64>, grad: &mut Array1<f64>) -> f64 {
        NUTS_RESIDUAL_SCRATCH.with(|scratch| {
            let mut residual = scratch.borrow_mut();
            if residual.len() != self.data.n_samples {
                *residual = Array1::<f64>::zeros(self.data.n_samples);
            }
            self.compute_logp_and_grad_nd_into(position, &mut residual, grad)
        })
    }
}

/// Configuration for NUTS sampling.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct NutsConfig {
    /// Number of samples to collect (after warmup)
    pub n_samples: usize,
    /// Number of warmup samples to discard
    pub nwarmup: usize,
    /// Number of parallel chains
    pub n_chains: usize,
    /// Target acceptance probability (0.6-0.9 recommended)
    pub target_accept: f64,
    /// Seed for deterministic chain initialization
    #[serde(default = "default_nuts_seed")]
    pub seed: u64,
}

fn default_nuts_seed() -> u64 {
    42
}

fn validate_nuts_target_accept(target_accept: f64) -> Result<(), HmcError> {
    if target_accept.is_finite() && target_accept > 0.0 && target_accept < 1.0 {
        Ok(())
    } else {
        Err(HmcError::InvalidConfig {
            reason: format!(
                "NUTS target_accept must be finite and lie in (0, 1), got {target_accept}"
            ),
        })
    }
}

#[inline]
fn splitmix64(x: u64) -> u64 {
    crate::linalg::utils::splitmix64_hash(x)
}

#[inline]
fn chain_stream_seed(seed: u64, chain: usize, stream: u64) -> u64 {
    splitmix64(seed ^ stream ^ ((chain as u64).wrapping_mul(0xD1B5_4A32_D192_ED03)))
}

#[inline]
fn nuts_transition_seed(seed: u64, stream: u64) -> u64 {
    splitmix64(seed ^ stream ^ 0xA24B_AED4_963E_E407)
}

#[inline]
fn robust_target_accept(requested: f64, dim: usize) -> f64 {
    let floor = if dim > 50 { 0.92 } else { 0.90 };
    requested.max(floor).min(0.95)
}

fn jittered_initial_positions(
    config: &NutsConfig,
    dim: usize,
    scale: f64,
    stream: u64,
) -> Vec<Array1<f64>> {
    (0..config.n_chains)
        .map(|chain| {
            let mut rng = StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, stream));
            Array1::from_shape_fn(dim, |_| sample_standard_normal(&mut rng) * scale)
        })
        .collect()
}

fn robust_mass_matrix_config(dim: usize, nwarmup: usize) -> NUTSMassMatrixConfig {
    if nwarmup < 80 {
        return NUTSMassMatrixConfig::disabled();
    }
    let start_buffer = (nwarmup / 8).clamp(35, 180);
    let end_buffer = (nwarmup / 5).clamp(50, 250);
    let initial_window = (nwarmup / 20).clamp(10, 60);
    NUTSMassMatrixConfig {
        adaptation: MassMatrixAdaptation::Diagonal,
        start_buffer,
        end_buffer,
        initial_window,
        regularize: if dim > 50 { 0.14 } else { 0.10 },
        jitter: 1e-5,
        dense_max_dim: 75,
    }
}

fn robust_survival_mass_matrix_config(dim: usize, nwarmup: usize) -> NUTSMassMatrixConfig {
    if nwarmup < 80 {
        return NUTSMassMatrixConfig::disabled();
    }
    // Survival posteriors with censoring/rare events are often skewed; this
    // configuration uses diagonal adaptation.
    let start_buffer = (nwarmup / 7).clamp(40, 200);
    let end_buffer = (nwarmup / 4).clamp(60, 280);
    let initial_window = (nwarmup / 20).clamp(10, 60);
    NUTSMassMatrixConfig {
        adaptation: MassMatrixAdaptation::Diagonal,
        start_buffer,
        end_buffer,
        initial_window,
        regularize: if dim > 50 { 0.18 } else { 0.12 },
        jitter: 1e-5,
        dense_max_dim: 75,
    }
}

impl Default for NutsConfig {
    fn default() -> Self {
        Self {
            n_samples: 1000,
            nwarmup: 500,
            n_chains: 4,
            target_accept: 0.9,
            seed: 42,
        }
    }
}

impl NutsConfig {
    /// Create a config with sample counts tuned for the model dimension.
    ///
    /// Higher dimensions need more samples because:
    /// - ESS decreases with dimension (autocorrelation grows)
    /// - Split R-hat needs enough samples per chain to be meaningful
    ///
    /// Rule of thumb: target 100 effective samples per parameter.
    pub fn for_dimension(n_params: usize) -> Self {
        // ESS ≈ n_samples / (1 + 2τ) where τ ≈ sqrt(dim) for well-tuned NUTS
        let effective_autocorr = (n_params as f64).sqrt().max(1.0);

        // Target: at least 100 effective samples per parameter
        let target_ess = 100 * n_params;

        // Samples needed = ESS * (1 + 2τ), with 1.5x safety factor
        let raw_samples = (target_ess as f64 * (1.0 + 2.0 * effective_autocorr) * 1.5) as usize;

        // Clamp to reasonable range [500, 10000]
        let n_samples = raw_samples.clamp(500, 10_000);

        // Warmup ≈ samples (standard practice for adaptation)
        let nwarmup = n_samples;

        // More chains for higher dims (better R-hat estimation)
        let n_chains = if n_params > 50 { 4 } else { 2 };

        Self {
            n_samples,
            nwarmup,
            n_chains,
            target_accept: 0.9,
            seed: 42,
        }
    }
}

/// Result of NUTS sampling.
#[derive(Clone, Debug)]
pub struct NutsResult {
    /// Coefficient samples in ORIGINAL space: shape (n_total_samples, n_coeffs)
    pub samples: Array2<f64>,
    /// Posterior mean
    pub posterior_mean: Array1<f64>,
    /// Posterior standard deviation
    pub posterior_std: Array1<f64>,
    /// R-hat convergence diagnostic
    pub rhat: f64,
    /// Effective sample size
    pub ess: f64,
    /// Whether sampling converged (R-hat < 1.1)
    pub converged: bool,
}

impl NutsResult {
    /// Computes the posterior mean of a function applied to coefficients.
    /// Returns 0.0 if samples is empty to avoid divide-by-zero.
    pub fn posterior_mean_of<F>(&self, f: F) -> f64
    where
        F: Fn(ArrayView1<f64>) -> f64 + Sync,
    {
        let n = self.samples.nrows();
        if n == 0 {
            return 0.0;
        }
        // Posterior mean of a sample-function: parallel reduction over rows.
        // `f: Fn(ArrayView1) -> f64` is shared-access so safe across threads.
        use rayon::iter::{IntoParallelIterator, ParallelIterator};
        let sum: f64 = (0..n).into_par_iter().map(|i| f(self.samples.row(i))).sum();
        sum / n as f64
    }

    /// Computes percentiles of a function applied to coefficients.
    pub fn posterior_interval_of<F>(&self, f: F, lower_pct: f64, upper_pct: f64) -> (f64, f64)
    where
        F: Fn(ArrayView1<f64>) -> f64,
    {
        let n = self.samples.nrows();
        if n == 0 {
            return (0.0, 0.0);
        }
        let mut values: Vec<f64> = (0..n).map(|i| f(self.samples.row(i))).collect();
        values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let lower_idx = ((lower_pct / 100.0) * n as f64).floor() as usize;
        let upper_idx = ((upper_pct / 100.0) * n as f64).ceil() as usize;

        (
            values[lower_idx.min(n.saturating_sub(1))],
            values[upper_idx.min(n.saturating_sub(1))],
        )
    }
}

#[inline]
fn sample_standard_normal<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
    let u1 = rng.random::<f64>().max(1e-16);
    let u2 = rng.random::<f64>();
    (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
}

/// Back-substitution against `L^T x = rhs`, with `L` a lower-triangular factor.
///
/// For `Q = L Lᵀ`, a draw `x ~ N(0, Q⁻¹)` is obtained from `z ~ N(0, I)` via
/// `x = L^{-T} z`, i.e. solving `Lᵀ x = z` by back-substitution. Using the
/// forward solve (`L x = z`) instead would produce `Var(x) = L⁻¹ L^{-T}`,
/// which equals `Q⁻¹` only when `L` is symmetric — wrong in general.
#[inline]
fn back_solve_lower_transposed(l: &Array2<f64>, rhs: &Array1<f64>, out: &mut Array1<f64>) {
    let p = rhs.len();
    assert_eq!(l.nrows(), p);
    assert_eq!(l.ncols(), p);
    assert_eq!(out.len(), p);
    // Solve Lᵀ x = rhs from the bottom row up. Row i of Lᵀ has nonzeros
    // at columns j ≥ i (= column i of L at rows j ≥ i), so
    //   rhs[i] = L[i,i] · x[i] + Σ_{j>i} L[j,i] · x[j].
    for i in (0..p).rev() {
        let mut v = rhs[i];
        for j in (i + 1)..p {
            v -= l[[j, i]] * out[j];
        }
        let d = l[[i, i]];
        out[i] = if d.abs() > 1e-14 { v / d } else { 0.0 };
    }
}

/// Runs a Pólya-Gamma Gibbs sampler for Bernoulli-logit models.
///
/// This sampler is gradient-free: each iteration alternates
/// 1) ω_i | β, y ~ PG(1, x_i^T β), and
/// 2) β | ω, y ~ N(Q^{-1} b, Q^{-1}), with Q = S + X^T diag(ω) X, b = X^T(y - 1/2).
///
/// For weighted data, this implementation is defined for weights ≈ 1.0 because it
/// samples PG(1,·) latent variables.
pub fn run_logit_polya_gamma_gibbs(
    x: ArrayView2<f64>,
    y: ArrayView1<f64>,
    weights: ArrayView1<f64>,
    penalty_matrix: ArrayView2<f64>,
    mode: ArrayView1<f64>,
    config: &NutsConfig,
) -> Result<NutsResult, String> {
    let n = x.nrows();
    let p = x.ncols();
    if y.len() != n || weights.len() != n {
        return Err(HmcError::DimensionMismatch {
            reason: "run_logit_polya_gamma_gibbs: input length mismatch".to_string(),
        }
        .into());
    }
    if mode.len() != p || penalty_matrix.nrows() != p || penalty_matrix.ncols() != p {
        return Err(HmcError::DimensionMismatch {
            reason: "run_logit_polya_gamma_gibbs: coefficient/penalty dimension mismatch"
                .to_string(),
        }
        .into());
    }
    if !weights.iter().all(|w| (*w - 1.0).abs() <= 1e-10) {
        return Err(HmcError::InvalidConfig {
            reason: "run_logit_polya_gamma_gibbs requires unit weights (PG(1,·)); use NUTS for non-unit weights".to_string(),
        }
        .into());
    }
    validate_binary_responses("run_logit_polya_gamma_gibbs", &y, &weights).map_err(String::from)?;

    let n_iter = config.nwarmup + config.n_samples;

    // b = X^T (y - 1/2), constant across iterations.
    let kappa = y.mapv(|v| v - 0.5);
    let rhs_b = fast_atv(&x, &kappa);

    let mut samples_array = Array3::<f64>::zeros((config.n_chains, config.n_samples, p));
    let mut eta = Array1::<f64>::zeros(n);
    let mut omega = Array1::<f64>::ones(n);
    let mut xw = x.to_owned();
    let mut xt_omega_x = Array2::<f64>::zeros((p, p));
    let penalty = penalty_matrix.to_owned();
    let mut q = Array2::<f64>::zeros((p, p));
    let mut mean = Array1::<f64>::zeros(p);
    let mut z = Array1::<f64>::zeros(p);
    let mut noise = Array1::<f64>::zeros(p);
    let mut progress = SamplingVisualizer::new(
        "polya-gamma gibbs",
        config.n_chains,
        config.nwarmup,
        config.n_samples,
    );

    for chain in 0..config.n_chains {
        progress.begin_chain(chain, "polya-gamma gibbs");
        let mut pg_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0x4D94_DF4E_5D72_81AB));
        let mut init_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0xB3C4_5A1F_8E9D_7632));
        let mut draw_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0x17A9_26D5_4C1B_E083));
        let pg = PolyaGamma::new();
        let mut beta = mode.to_owned();
        // Small jitter so chains are not perfectly coupled.
        for j in 0..p {
            beta[j] += 0.05 * sample_standard_normal(&mut init_rng);
        }

        for iter in 0..n_iter {
            eta.assign(&crate::faer_ndarray::fast_av(&x, &beta));
            for i in 0..n {
                omega[i] = pg.draw(&mut pg_rng, eta[i]).max(1e-12);
            }

            // Build Xweighted = diag(sqrt(ω)) X and compute X^T Ω X via faer GEMM.
            // Per-row scaling is fully independent across rows.
            ndarray::Zip::indexed(xw.rows_mut())
                .and(x.rows())
                .and(&omega)
                .par_for_each(|_idx, mut xw_row, x_row, omega_i| {
                    let s = omega_i.sqrt();
                    for j in 0..p {
                        xw_row[j] = x_row[j] * s;
                    }
                });
            fast_ata_into(&xw, &mut xt_omega_x);

            q.assign(&penalty);
            q += &xt_omega_x;

            // β | ω,y ~ N(Q^{-1} b, Q^{-1})
            let factor = q
                .cholesky(Side::Lower)
                .map_err(|e| format!("PG Gibbs failed to factor Q: {:?}", e))?;
            mean.assign(&factor.solvevec(&rhs_b));

            for j in 0..p {
                z[j] = sample_standard_normal(&mut draw_rng);
            }
            let l = factor.lower_triangular();
            back_solve_lower_transposed(&l, &z, &mut noise);
            beta.assign(&(&mean + &noise));

            if iter < config.nwarmup {
                progress.warmup_step(iter + 1);
            } else {
                if iter == config.nwarmup {
                    progress.start_sampling();
                }
                let keep_idx = iter - config.nwarmup;
                progress.sample_step(keep_idx + 1);
                samples_array
                    .slice_mut(ndarray::s![chain, keep_idx, ..])
                    .assign(&beta);
            }
        }
        progress.finish_chain(1.0);
    }

    let total_samples = config.n_chains * config.n_samples;
    let mut samples = Array2::<f64>::zeros((total_samples, p));
    for chain in 0..config.n_chains {
        for s in 0..config.n_samples {
            let idx = chain * config.n_samples + s;
            samples
                .row_mut(idx)
                .assign(&samples_array.slice(ndarray::s![chain, s, ..]));
        }
    }

    let posterior_mean = samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(p));
    let posterior_std = samples.std_axis(Axis(0), 0.0);
    let (rhat, ess) = if config.n_chains >= 2 && config.n_samples >= 4 {
        compute_split_rhat_and_ess(&samples_array)
    } else {
        (1.0, (total_samples as f64) * 0.5)
    };
    let converged = rhat < 1.1 && ess > 100.0;
    progress.finish_all(rhat, ess);

    Ok(NutsResult {
        samples,
        posterior_mean,
        posterior_std,
        rhat,
        ess,
        converged,
    })
}

/// Estimate E_{ω|y,ρ}[ tr(S_k Q^{-1}) + μᵀ S_k μ ] with PG Gibbs + Rao-Blackwellization.
///
/// For each retained Gibbs state ω:
///   Q = S + Xᵀ diag(ω) X,  μ = Q^{-1} Xᵀ(y-1/2),
/// and with S_k = R_kᵀ R_k:
///   tr(S_k Q^{-1}) + μᵀ S_k μ
/// = tr(R_k Q^{-1} R_kᵀ) + ||R_k μ||².
///
/// Returns one expectation per penalty block k, averaged over retained draws.
pub fn estimate_logit_pg_rao_blackwell_terms(
    x: ArrayView2<f64>,
    y: ArrayView1<f64>,
    weights: ArrayView1<f64>,
    penalty_matrix: ArrayView2<f64>,
    mode: ArrayView1<f64>,
    penalty_roots: &[Array2<f64>],
    config: &NutsConfig,
) -> Result<Array1<f64>, String> {
    let n = x.nrows();
    let p = x.ncols();
    if y.len() != n || weights.len() != n {
        return Err(HmcError::DimensionMismatch {
            reason: "estimate_logit_pg_rao_blackwell_terms: input length mismatch".to_string(),
        }
        .into());
    }
    if mode.len() != p || penalty_matrix.nrows() != p || penalty_matrix.ncols() != p {
        return Err(HmcError::DimensionMismatch {
            reason: "estimate_logit_pg_rao_blackwell_terms: coefficient/penalty dimension mismatch"
                .to_string(),
        }
        .into());
    }
    if !weights.iter().all(|w| (*w - 1.0).abs() <= 1e-10) {
        return Err(HmcError::InvalidConfig {
            reason: "estimate_logit_pg_rao_blackwell_terms requires unit weights (PG(1,·))"
                .to_string(),
        }
        .into());
    }
    validate_binary_responses("estimate_logit_pg_rao_blackwell_terms", &y, &weights)
        .map_err(String::from)?;
    if penalty_roots.iter().any(|r| r.ncols() != p) {
        return Err(HmcError::DimensionMismatch {
            reason: "estimate_logit_pg_rao_blackwell_terms: root width mismatch".to_string(),
        }
        .into());
    }
    // Precompute transposed root blocks once:
    //   R_k^T is the RHS used for batched solves Q X = R_k^T.
    let penalty_roots_t: Vec<Array2<f64>> =
        penalty_roots.iter().map(|r| r.t().to_owned()).collect();

    let n_iter = config.nwarmup + config.n_samples;

    // Logistic PG identity uses kappa_i = y_i - 1/2 so that
    // b = X^T kappa in the Gaussian conditional for beta|omega.
    let kappa = y.mapv(|v| v - 0.5);
    let rhs_b = fast_atv(&x, &kappa);

    let penalty = penalty_matrix.to_owned();
    let mut eta = Array1::<f64>::zeros(n);
    let mut omega = Array1::<f64>::ones(n);
    let mut xw = x.to_owned();
    let mut xt_omega_x = Array2::<f64>::zeros((p, p));
    let mut q = Array2::<f64>::zeros((p, p));
    let mut mean = Array1::<f64>::zeros(p);
    let mut rb_sum = Array1::<f64>::zeros(penalty_roots.len());
    let mut z = Array1::<f64>::zeros(p);
    let mut noise = Array1::<f64>::zeros(p);

    let mut kept = 0usize;
    for chain in 0..config.n_chains {
        let mut pg_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0x83F1_56C9_A7E0_2D4B));
        let mut init_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0x28F0_7B65_1A4D_C93E));
        let mut draw_rng =
            StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0xC642_6E35_B5A9_1D80));
        let pg = PolyaGamma::new();
        let mut beta = mode.to_owned();
        for j in 0..p {
            beta[j] += 0.05 * sample_standard_normal(&mut init_rng);
        }

        for iter in 0..n_iter {
            eta.assign(&crate::faer_ndarray::fast_av(&x, &beta));
            for i in 0..n {
                omega[i] = pg.draw(&mut pg_rng, eta[i]).max(1e-12);
            }

            ndarray::Zip::from(xw.rows_mut())
                .and(x.rows())
                .and(&omega)
                .par_for_each(|mut xw_row, x_row, &omega_i| {
                    let s = omega_i.sqrt();
                    for j in 0..p {
                        xw_row[j] = x_row[j] * s;
                    }
                });
            fast_ata_into(&xw, &mut xt_omega_x);

            // Conditional precision:
            //   Q = S + X^T diag(omega) X.
            q.assign(&penalty);
            q += &xt_omega_x;

            let factor = q
                .cholesky(Side::Lower)
                .map_err(|e| format!("PG Rao-Blackwell failed to factor Q: {:?}", e))?;
            // Conditional mean:
            //   mu = Q^{-1} b,  b = X^T(y - 1/2).
            mean.assign(&factor.solvevec(&rhs_b));

            // Draw beta for the next Gibbs state.
            for j in 0..p {
                z[j] = sample_standard_normal(&mut draw_rng);
            }
            let l = factor.lower_triangular();
            back_solve_lower_transposed(&l, &z, &mut noise);
            beta.assign(&(&mean + &noise));

            if iter < config.nwarmup {
                continue;
            }
            kept += 1;

            for (k, r_k) in penalty_roots.iter().enumerate() {
                if r_k.nrows() == 0 {
                    continue;
                }

                // mu^T S_k mu via root form S_k = R_k^T R_k.
                let rmu = r_k.dot(&mean);
                let mu_quad = rmu.dot(&rmu);

                // Batched trace solve:
                //   V_k = Q^{-1} R_k^T  (single multi-RHS solve)
                // then tr(R_k Q^{-1} R_k^T) = <R_k, V_k^T>_F.
                let solved_mat = factor.solve_mat(&penalty_roots_t[k]); // (p, r_k)
                let solved_t = solved_mat.t();
                let mut trace_term = 0.0_f64;
                for (&a, &b) in r_k.iter().zip(solved_t.iter()) {
                    trace_term += a * b;
                }

                rb_sum[k] += trace_term + mu_quad;
            }
        }
    }

    if kept == 0 {
        return Err(HmcError::SamplingFailed {
            reason: "estimate_logit_pg_rao_blackwell_terms: no retained samples".to_string(),
        }
        .into());
    }
    let out = rb_sum.mapv(|v| v / (kept as f64));
    if !out.iter().all(|v| v.is_finite()) {
        return Err(HmcError::NonFiniteState {
            reason: "estimate_logit_pg_rao_blackwell_terms: non-finite expectation".to_string(),
        }
        .into());
    }
    Ok(out)
}

/// Runs NUTS sampling using general-mcmc with whitened parameter space.
///
/// # Arguments
/// * `x` - Design matrix [n_samples, dim]
/// * `y` - Response vector [n_samples]
/// * `weights` - Observation/case weights [n_samples]
/// * `penalty_matrix` - Combined penalty S [dim, dim]
/// * `mode` - MAP estimate μ [dim]
/// * `hessian` - Penalized Hessian H [dim, dim] (NOT the inverse!)
/// * `nuts_family` - Family for log-likelihood computation
/// * `firth_bias_reduction` - Whether Firth bias reduction was used in training
/// * `config` - NUTS configuration
pub(crate) fn run_nuts_sampling(
    x: ArrayView2<f64>,
    y: ArrayView1<f64>,
    weights: ArrayView1<f64>,
    penalty_matrix: ArrayView2<f64>,
    mode: ArrayView1<f64>,
    hessian: ArrayView2<f64>,
    nuts_family: NutsFamily,
    gamma_shape: f64,
    dispersion: crate::solver::estimate::Dispersion,
    firth_bias_reduction: bool,
    config: &NutsConfig,
) -> Result<NutsResult, String> {
    validate_firth_support(nuts_family, firth_bias_reduction).map_err(String::from)?;
    validate_nuts_target_accept(config.target_accept).map_err(String::from)?;
    if nuts_family == NutsFamily::TweedieLog && !is_valid_tweedie_power(gamma_shape) {
        return Err(format!(
            "Tweedie variance power must be finite and strictly between 1 and 2; got {gamma_shape}"
        ));
    }
    let dim = mode.len();

    // Create posterior target with analytical gradients. When Firth is enabled,
    // this target includes the identifiable-subspace Jeffreys term.
    let target = NutsPosterior::new(
        x,
        y,
        weights,
        penalty_matrix,
        mode,
        hessian,
        nuts_family,
        gamma_shape,
        dispersion,
        firth_bias_reduction,
    )?;

    // Get Cholesky factor for un-whitening samples later
    let chol = target.chol().clone();
    let mode_arr = target.mode().clone();

    let initial_positions = jittered_initial_positions(config, dim, 0.1, 0x0F65_83B2_BC71_4D9E);

    // Create GenericNUTS sampler - it auto-tunes step size!
    let mass_cfg = robust_mass_matrix_config(dim, config.nwarmup);
    let mut sampler = GenericNUTS::new_with_mass_matrix(
        target,
        initial_positions,
        robust_target_accept(config.target_accept, dim),
        mass_cfg,
    )
    .set_seed(nuts_transition_seed(config.seed, 0xF1D3_C2B5_A697_804E));

    let (samples_array, run_stats) = sampler
        .run_progress(config.n_samples, config.nwarmup)
        .map_err(|e| format!("NUTS sampling failed: {e}"))?;
    log::info!("NUTS sampling complete: {}", run_stats);

    // Convert samples from whitened space back to original space
    // samples_array has shape [n_chains, n_samples, dim]
    let shape = samples_array.shape();
    let n_chains = shape[0];
    let n_samples_out = shape[1];
    let total_samples = n_chains * n_samples_out;

    let mut samples = Array2::<f64>::zeros((total_samples, dim));
    let mut z_buffer = Array1::<f64>::zeros(dim);
    for chain in 0..n_chains {
        for sample_i in 0..n_samples_out {
            let zview = samples_array.slice(ndarray::s![chain, sample_i, ..]);
            z_buffer.assign(&zview);
            let beta = &mode_arr + &chol.dot(&z_buffer);
            let sample_idx = chain * n_samples_out + sample_i;
            samples.row_mut(sample_idx).assign(&beta);
        }
    }

    // Compute split-chain R-hat and ESS for proper convergence diagnostics
    let posterior_mean = samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(dim));
    let posterior_std = samples.std_axis(Axis(0), 0.0);

    // Split-chain R-hat: compare variance within vs between chains
    // Gelman-Rubin diagnostic with split chains
    let (rhat, ess) = compute_split_rhat_and_ess(&samples_array);

    let converged = rhat < 1.1 && ess > 100.0;

    Ok(NutsResult {
        samples,
        posterior_mean,
        posterior_std,
        rhat,
        ess,
        converged,
    })
}

/// Flattened numeric inputs for GLM-family NUTS sampling.
pub struct GlmFlatInputs<'a> {
    pub x: ArrayView2<'a, f64>,
    pub y: ArrayView1<'a, f64>,
    pub weights: ArrayView1<'a, f64>,
    pub penalty_matrix: ArrayView2<'a, f64>,
    pub mode: ArrayView1<'a, f64>,
    pub hessian: ArrayView2<'a, f64>,
    pub gamma_shape: Option<f64>,
    /// Dispersion parameter φ used to scale the likelihood and the
    /// whitening Cholesky. For fixed-scale families (Binomial, Poisson)
    /// this is `Dispersion::Known(1.0)` and has no numerical effect;
    /// for Gaussian / Gamma it carries the estimated `phi` so that the
    /// sampler targets the φ-scaled posterior covariance `Vb = φ·H⁻¹`.
    /// See `inference::dispersion_cov` for the ownership invariants.
    pub dispersion: crate::solver::estimate::Dispersion,
    pub firth_bias_reduction: bool,
}

/// Flat survival inputs for engine-facing HMC APIs.
pub struct SurvivalFlatInputs<'a> {
    pub age_entry: ArrayView1<'a, f64>,
    pub age_exit: ArrayView1<'a, f64>,
    pub event_target: ArrayView1<'a, u8>,
    pub event_competing: ArrayView1<'a, u8>,
    pub weights: ArrayView1<'a, f64>,
    pub x_entry: ArrayView2<'a, f64>,
    pub x_exit: ArrayView2<'a, f64>,
    pub x_derivative: ArrayView2<'a, f64>,
    pub eta_offset_entry: Option<ArrayView1<'a, f64>>,
    pub eta_offset_exit: Option<ArrayView1<'a, f64>>,
    pub derivative_offset_exit: Option<ArrayView1<'a, f64>>,
}

/// Flattened numeric inputs for Royston-Parmar NUTS sampling.
pub struct SurvivalNutsInputs<'a> {
    pub flat: SurvivalFlatInputs<'a>,
    pub penalties: crate::survival::PenaltyBlocks,
    pub monotonicity: crate::survival::MonotonicityPenalty,
    pub spec: crate::survival::SurvivalSpec,
    pub structurally_monotonic: bool,
    pub structural_time_columns: usize,
    pub mode: ArrayView1<'a, f64>,
    pub hessian: ArrayView2<'a, f64>,
}

/// Family-dispatched flattened NUTS inputs.
pub enum FamilyNutsInputs<'a> {
    Glm(GlmFlatInputs<'a>),
    Survival(Box<SurvivalNutsInputs<'a>>),
}

/// Return the explicit fitted penalized Hessian used for HMC/NUTS whitening.
///
/// This is the only supported upstream-to-HMC curvature handoff: callers must
/// pass a dense Hessian (or an already materialized exact operator stored as a
/// dense Hessian) exported by the fitter. We deliberately do not synthesize a
/// numerical Hessian and do not invert `beta_covariance` as a compatibility
/// fallback, because either path can silently whiten against curvature that the
/// upstream fit never certified.
pub fn explicit_fit_hessian_for_whitening<'a>(
    fit: &'a UnifiedFitResult,
    expected_dim: usize,
    label: &str,
) -> Result<&'a Array2<f64>, String> {
    let hessian = fit.penalized_hessian().ok_or_else(|| {
        format!(
            "{label}: fit result is missing an explicit penalized Hessian for HMC/NUTS whitening"
        )
    })?;
    validate_explicit_dense_hessian_for_whitening(
        &format!("{label} penalized Hessian"),
        hessian,
        expected_dim,
    )
    .map_err(|err| err.to_string())?;
    Ok(hessian)
}

/// Family-agnostic flattened NUTS entrypoint across all supported likelihood families.
pub fn run_nuts_sampling_flattened_family(
    likelihood: LikelihoodSpec,
    inputs: FamilyNutsInputs<'_>,
    config: &NutsConfig,
) -> Result<NutsResult, String> {
    if let FamilyNutsInputs::Glm(glm) = &inputs
        && glm.firth_bias_reduction
        && !likelihood_spec_supports_firth(&likelihood)
    {
        let binomial_logit = LikelihoodSpec {
            response: ResponseFamily::Binomial,
            link: InverseLink::Standard(LinkFunction::Logit),
        };
        return Err(HmcError::FirthUnsupported {
            reason: format!(
                "NUTS with Firth is only supported for {}; {} does not support it",
                binomial_logit.pretty_name(),
                likelihood.pretty_name()
            ),
        }
        .into());
    }

    match (likelihood.response.clone(), likelihood.link.clone(), inputs) {
        (
            ResponseFamily::Gaussian,
            InverseLink::Standard(LinkFunction::Identity),
            FamilyNutsInputs::Glm(glm),
        ) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::Gaussian,
            1.0,
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (
            ResponseFamily::Binomial,
            InverseLink::Standard(LinkFunction::Logit),
            FamilyNutsInputs::Glm(glm),
        ) => {
            // Auto-select PG Gibbs when assumptions hold; otherwise fall back to NUTS.
            // This gives gradient-free posterior draws for standard Bernoulli logit GAMs.
            if !glm.firth_bias_reduction && glm.weights.iter().all(|w| (*w - 1.0).abs() <= 1e-10) {
                run_logit_polya_gamma_gibbs(
                    glm.x,
                    glm.y,
                    glm.weights,
                    glm.penalty_matrix,
                    glm.mode,
                    config,
                )
            } else {
                run_nuts_sampling(
                    glm.x,
                    glm.y,
                    glm.weights,
                    glm.penalty_matrix,
                    glm.mode,
                    glm.hessian,
                    NutsFamily::BinomialLogit,
                    1.0,
                    glm.dispersion,
                    glm.firth_bias_reduction,
                    config,
                )
            }
        }
        (
            ResponseFamily::Binomial,
            InverseLink::Standard(LinkFunction::Probit),
            FamilyNutsInputs::Glm(glm),
        ) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::BinomialProbit,
            1.0,
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (
            ResponseFamily::Binomial,
            InverseLink::Standard(LinkFunction::CLogLog),
            FamilyNutsInputs::Glm(glm),
        ) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::BinomialCLogLog,
            1.0,
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (
            ResponseFamily::Binomial,
            InverseLink::LatentCLogLog(_),
            FamilyNutsInputs::Glm(glm),
        ) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::BinomialCLogLog,
            1.0,
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (ResponseFamily::Binomial, InverseLink::Mixture(_), FamilyNutsInputs::Glm(_)) => Err(
            "BinomialMixture NUTS is not implemented yet; use fit_gam/predict_gam for blended inverse-link models"
                .to_string(),
        ),
        (ResponseFamily::Binomial, InverseLink::Sas(_), FamilyNutsInputs::Glm(_)) => Err(
            "BinomialSas NUTS is not implemented yet; use fit_gam/predict_gam for SAS-link models"
                .to_string(),
        ),
        (ResponseFamily::Binomial, InverseLink::BetaLogistic(_), FamilyNutsInputs::Glm(_)) => Err(
            "BinomialBetaLogistic NUTS is not implemented yet; use fit_gam/predict_gam for beta-logistic-link models"
                .to_string(),
        ),
        (ResponseFamily::Binomial, InverseLink::Standard(_), FamilyNutsInputs::Glm(_)) => Err(
            "NUTS sampling is not implemented for this binomial inverse link".to_string(),
        ),
        (ResponseFamily::RoystonParmar, _, FamilyNutsInputs::Survival(survival)) => {
            survival_hmc::run_survival_nuts_sampling(
                survival.flat.age_entry,
                survival.flat.age_exit,
                survival.flat.event_target,
                survival.flat.event_competing,
                survival.flat.weights,
                survival.flat.x_entry,
                survival.flat.x_exit,
                survival.flat.x_derivative,
                survival.flat.eta_offset_entry,
                survival.flat.eta_offset_exit,
                survival.flat.derivative_offset_exit,
                survival.penalties,
                survival.monotonicity,
                survival.spec,
                survival.structurally_monotonic,
                survival.structural_time_columns,
                survival.mode,
                survival.hessian,
                config,
            )
        }
        (ResponseFamily::RoystonParmar, _, FamilyNutsInputs::Glm(_)) => Err(
            "RoystonParmar family requires FamilyNutsInputs::Survival flattened inputs".to_string(),
        ),
        (_, _, FamilyNutsInputs::Survival(_)) => Err(
            "Survival flattened inputs are only valid for the Royston-Parmar response family"
                .to_string(),
        ),
        (ResponseFamily::Poisson, _, FamilyNutsInputs::Glm(glm)) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::PoissonLog,
            1.0,
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (ResponseFamily::Tweedie { p }, _, FamilyNutsInputs::Glm(glm)) => {
            // Family mapping: Tweedie payload p is passed through the family-parameter slot.
            // The Tweedie dispersion phi remains in glm.dispersion, matching REML.
            if !is_valid_tweedie_power(p) {
                return Err(format!(
                    "Tweedie variance power must be finite and strictly between 1 and 2; got {p}"
                ));
            }
            run_nuts_sampling(
                glm.x,
                glm.y,
                glm.weights,
                glm.penalty_matrix,
                glm.mode,
                glm.hessian,
                NutsFamily::TweedieLog,
                p,
                glm.dispersion,
                glm.firth_bias_reduction,
                config,
            )
        }
        (ResponseFamily::NegativeBinomial { theta }, _, FamilyNutsInputs::Glm(glm)) => {
            // Family mapping: NegativeBinomial payload theta is passed through the family slot.
            // NB dispersion scale is unit; theta is not derived from fixed_phi.
            run_nuts_sampling(
                glm.x,
                glm.y,
                glm.weights,
                glm.penalty_matrix,
                glm.mode,
                glm.hessian,
                NutsFamily::NegativeBinomialLog,
                theta,
                glm.dispersion,
                glm.firth_bias_reduction,
                config,
            )
        }
        (ResponseFamily::Beta { .. }, _, FamilyNutsInputs::Glm(_)) => Err(
            "NUTS sampling is not implemented for beta-regression logit".to_string(),
        ),
        (ResponseFamily::Gamma, _, FamilyNutsInputs::Glm(glm)) => run_nuts_sampling(
            glm.x,
            glm.y,
            glm.weights,
            glm.penalty_matrix,
            glm.mode,
            glm.hessian,
            NutsFamily::GammaLog,
            glm.gamma_shape.unwrap_or(1.0),
            glm.dispersion,
            glm.firth_bias_reduction,
            config,
        ),
        (ResponseFamily::Gaussian, _, FamilyNutsInputs::Glm(_)) => Err(
            "NUTS sampling is only implemented for Gaussian with identity link".to_string(),
        ),
    }
}

// ============================================================================
// Joint (β, θ) Link-Wiggle HMC
// ============================================================================
//
// NUTS sampling over the joint parameter space [β_eta; β_wiggle] for models
// with a structurally monotone I-spline link wiggle. The wiggle introduces a
// nonlinear coupling:
//
//   η(β_eta, β_wiggle) = q₀(β_eta) + B(q₀(β_eta)) · β_wiggle
//
// where B is the shared monotone wiggle basis evaluated at the base linear
// predictor q₀ = X · β_eta. The gradient of log p(y|β_eta, β_wiggle) w.r.t.
// β_eta picks up a chain-rule factor g'(q₀) = 1 + B'(q₀) · β_wiggle / range_width
// from the dependence of B on q₀.
//
// Whitening uses the Cholesky of the joint Hessian at the mode, exactly as for
// the standard NutsPosterior. C^1 linear extension outside the training knot
// range prevents basis evaluation discontinuities.

/// Fixed spline artifacts for link-wiggle posterior sampling.
#[derive(Clone)]
pub struct LinkWiggleSplineArtifacts {
    /// Knot range (min, max) from training (in standardized [0,1] space of q₀)
    pub knot_range: (f64, f64),
    /// Full knot vector for the shared monotone I-spline basis
    pub knot_vector: Array1<f64>,
    /// I-spline degree
    pub degree: usize,
}

/// Whitened log-posterior target for joint (β_eta, β_wiggle) with analytical gradients.
#[derive(Clone)]
pub struct LinkWigglePosterior {
    /// Main design matrix X (n × p_main)
    x: Arc<Array2<f64>>,
    y: Arc<Array1<f64>>,
    weights: Arc<Array1<f64>>,
    /// Penalty for main coefficients (p_main × p_main)
    penalty_base: Arc<Array2<f64>>,
    /// Penalty for wiggle coefficients (p_wiggle × p_wiggle)
    penalty_link: Arc<Array2<f64>>,
    mode_beta: Arc<Array1<f64>>,
    mode_theta: Arc<Array1<f64>>,
    spline: LinkWiggleSplineArtifacts,
    /// L where LL^T = H^{-1} (joint Hessian)
    chol: Array2<f64>,
    /// L^T for gradient chain rule
    chol_t: Array2<f64>,
    p_base: usize,
    p_link: usize,
    n_samples: usize,
    nuts_family: NutsFamily,
    /// Family-specific noise parameter: Gaussian sigma or Gamma shape.
    scale: f64,
}

impl LinkWigglePosterior {
    /// Standardize q₀ values to [0,1] range using training knot bounds.
    #[inline]
    fn standardized_z(&self, u: &Array1<f64>) -> (Array1<f64>, Array1<f64>, f64) {
        let (min_u, max_u) = self.spline.knot_range;
        let rw = (max_u - min_u).max(1e-6);
        let z_raw: Array1<f64> = u.mapv(|v| (v - min_u) / rw);
        let z_c: Array1<f64> = z_raw.mapv(|z| z.clamp(0.0, 1.0));
        (z_raw, z_c, rw)
    }

    /// Creates a new link-wiggle posterior target.
    pub fn new(
        x: ArrayView2<f64>,
        y: ArrayView1<f64>,
        weights: ArrayView1<f64>,
        penalty_base: ArrayView2<f64>,
        penalty_link: ArrayView2<f64>,
        mode_beta: ArrayView1<f64>,
        mode_theta: ArrayView1<f64>,
        hessian: ArrayView2<f64>,
        spline: LinkWiggleSplineArtifacts,
        nuts_family: NutsFamily,
        scale: f64,
    ) -> Result<Self, String> {
        let n_samples = x.nrows();
        let p_base = x.ncols();
        let p_link = mode_theta.len();
        let dim = p_base + p_link;
        if hessian.nrows() != dim || hessian.ncols() != dim {
            return Err(HmcError::DimensionMismatch {
                reason: format!(
                    "LinkWigglePosterior: Hessian dim mismatch: {}x{} vs expected {}x{}",
                    hessian.nrows(),
                    hessian.ncols(),
                    dim,
                    dim,
                ),
            }
            .into());
        }
        if nuts_family.likelihood_spec().is_binomial() {
            validate_binary_responses("binomial link-wiggle NUTS", &y, &weights)
                .map_err(String::from)?;
        }
        if matches!(nuts_family, NutsFamily::NegativeBinomialLog) {
            validate_count_responses("negative-binomial link-wiggle NUTS", &y, &weights)
                .map_err(String::from)?;
        }
        let hessian_owned = hessian.to_owned();
        let chol_factor = hessian_owned
            .cholesky(Side::Lower)
            .map_err(|e| format!("LinkWigglePosterior Cholesky failed: {:?}", e))?;
        let l_h = chol_factor.lower_triangular();
        let mut chol = solve_upper_triangular_transpose(&l_h, dim);
        // Scale the Cholesky factor by √φ for families with estimated
        // dispersion. For Gaussian the relevant scale is σ (so φ = σ²)
        // and the posterior covariance the sampler should target is
        // `Vb = φ·H⁻¹`. Without this scaling the sampler whitened
        // against the unscaled Hessian, silently targeting H⁻¹.
        // Fixed-scale families (Binomial / Poisson) take this branch
        // with `sqrt_phi == 1` and pay no overhead.
        // Gamma uses the shape parameterization: `scale = α = 1/φ`, so the
        // dispersion that scales `Vb = φ·H⁻¹` is `φ = 1/α`. Without this
        // factor the Gamma sampler whitened against an unscaled Hessian
        // exactly the way the Gaussian branch did before this fix.
        let sqrt_phi = match nuts_family {
            NutsFamily::Gaussian => scale.max(0.0),
            NutsFamily::GammaLog => 1.0 / scale.max(1e-10).sqrt(),
            _ => 1.0,
        };
        if (sqrt_phi - 1.0).abs() > 0.0 {
            chol.mapv_inplace(|v| v * sqrt_phi);
        }
        let chol_t = chol.t().to_owned();
        Ok(Self {
            x: Arc::new(x.to_owned()),
            y: Arc::new(y.to_owned()),
            weights: Arc::new(weights.to_owned()),
            penalty_base: Arc::new(penalty_base.to_owned()),
            penalty_link: Arc::new(penalty_link.to_owned()),
            mode_beta: Arc::new(mode_beta.to_owned()),
            mode_theta: Arc::new(mode_theta.to_owned()),
            spline,
            chol,
            chol_t,
            p_base,
            p_link,
            n_samples,
            nuts_family,
            scale,
        })
    }

    /// Evaluate the wiggle basis and compute η = q₀ + B(q₀)·θ with C^1 linear extension.
    fn evaluate_link(&self, u: &Array1<f64>, theta: &Array1<f64>) -> (Array2<f64>, Array1<f64>) {
        let n = u.len();
        if theta.is_empty() {
            return (Array2::zeros((n, 0)), u.clone());
        }

        let (z_raw, z_c, _) = self.standardized_z(u);
        let Ok(mut basis) = monotone_wiggle_basis_with_derivative_order(
            z_c.view(),
            &self.spline.knot_vector,
            self.spline.degree,
            0,
        ) else {
            return (Array2::zeros((n, theta.len())), u.clone());
        };
        if basis.ncols() != theta.len() {
            return (Array2::zeros((n, theta.len())), u.clone());
        }

        // C^1 linear extension outside [0, 1]:
        // B_ext(z_raw) = B(z_c) + (z_raw - z_c) * B'(z_c)
        let mut needs_ext = false;
        for i in 0..n {
            if (z_raw[i] - z_c[i]).abs() > 1e-12 {
                needs_ext = true;
                break;
            }
        }
        if needs_ext
            && let Ok(b_prime) = monotone_wiggle_basis_with_derivative_order(
                z_c.view(),
                &self.spline.knot_vector,
                self.spline.degree,
                1,
            )
        {
            for i in 0..n {
                let dz = z_raw[i] - z_c[i];
                if dz.abs() <= 1e-12 {
                    continue;
                }
                for j in 0..basis.ncols().min(b_prime.ncols()) {
                    basis[[i, j]] += dz * b_prime[[i, j]];
                }
            }
        }
        (
            basis.clone(),
            u + &crate::faer_ndarray::fast_av(&basis, theta),
        )
    }

    /// Compute dη/dq₀ = 1 + B'(q₀)·θ / range_width (chain-rule factor for β_eta gradient).
    fn compute_g_prime(&self, u: &Array1<f64>, theta: &Array1<f64>) -> Array1<f64> {
        let n = u.len();
        let mut g = Array1::<f64>::ones(n);
        let (_, z_c, rw) = self.standardized_z(u);
        if theta.is_empty() {
            return g;
        }

        let Ok(b_prime_constrained) = monotone_wiggle_basis_with_derivative_order(
            z_c.view(),
            &self.spline.knot_vector,
            self.spline.degree,
            1,
        ) else {
            return g;
        };
        if b_prime_constrained.ncols() != theta.len() {
            return g;
        }
        let dwiggle_dz = crate::faer_ndarray::fast_av(&b_prime_constrained, theta);
        ndarray::Zip::from(&mut g)
            .and(&dwiggle_dz)
            .par_for_each(|gi, &dw| *gi = 1.0 + dw / rw);
        g
    }

    fn compute_logp_and_grad_into(&self, z: &Array1<f64>, grad: &mut Array1<f64>) -> f64 {
        let dim = self.p_base + self.p_link;

        // Un-whiten: q = mode + L·z
        let mut mode = Array1::<f64>::zeros(dim);
        mode.slice_mut(ndarray::s![0..self.p_base])
            .assign(&self.mode_beta);
        mode.slice_mut(ndarray::s![self.p_base..])
            .assign(&self.mode_theta);
        let q = &mode + &self.chol.dot(z);
        let beta = q.slice(ndarray::s![0..self.p_base]).to_owned();
        let theta = q.slice(ndarray::s![self.p_base..]).to_owned();

        // Compute η = q₀ + B(q₀)·θ where q₀ = X·β
        let u = crate::faer_ndarray::fast_av(self.x.as_ref(), &beta);
        let (bwiggle, eta) = self.evaluate_link(&u, &theta);

        // Log-likelihood and residuals via family dispatch
        let ll;
        let mut residual = Array1::<f64>::zeros(self.n_samples);
        match self.nuts_family {
            NutsFamily::Gaussian => {
                let inv_scale_sq = 1.0 / (self.scale * self.scale).max(1e-10);
                let mut ll_acc = 0.0;
                for i in 0..self.n_samples {
                    let r = self.y[i] - eta[i];
                    let w = self.weights[i];
                    ll_acc -= 0.5 * w * r * r * inv_scale_sq;
                    residual[i] = w * r * inv_scale_sq;
                }
                ll = ll_acc;
            }
            NutsFamily::BinomialLogit => {
                let mut ll_acc = 0.0;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    ll_acc += w_i * (y_i * eta_i - NutsPosterior::log1pexp(eta_i));
                    let mu = NutsPosterior::sigmoid_stable(eta_i);
                    residual[i] = w_i * (y_i - mu);
                }
                ll = ll_acc;
            }
            NutsFamily::BinomialProbit => {
                let mut ll_acc = 0.0;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    let log_phi_pos = log_ndtr(eta_i);
                    let log_phi_neg = log_ndtr(-eta_i);
                    ll_acc += w_i * (y_i * log_phi_pos + (1.0 - y_i) * log_phi_neg);
                    let log_phi = standard_normal_log_pdf(eta_i);
                    let ratio_pos = (log_phi - log_phi_pos).exp();
                    let ratio_neg = (log_phi - log_phi_neg).exp();
                    residual[i] = w_i * (y_i * ratio_pos - (1.0 - y_i) * ratio_neg);
                }
                ll = ll_acc;
            }
            NutsFamily::BinomialCLogLog => {
                let mut ll_acc = 0.0;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    if !(eta_i.is_finite() && (-700.0..=700.0).contains(&eta_i)) {
                        grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    let (ll_i, residual_i) = match cloglog_bernoulli_logp_and_residual(eta_i, y_i) {
                        Ok(values) => values,
                        Err(_) => {
                            grad.fill(0.0);
                            return f64::NEG_INFINITY;
                        }
                    };
                    ll_acc += w_i * ll_i;
                    residual[i] = w_i * residual_i;
                }
                ll = ll_acc;
            }
            NutsFamily::PoissonLog => {
                let mut ll_acc = 0.0;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    if !(eta_i.is_finite() && (-30.0..=30.0).contains(&eta_i)) {
                        grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    let mu = eta_i.exp();
                    ll_acc += w_i * (y_i * eta_i - mu);
                    residual[i] = w_i * (y_i - mu);
                }
                ll = ll_acc;
            }
            NutsFamily::TweedieLog => {
                let mut ll_acc = 0.0;
                // Family mapping: Tweedie scale carries payload p; phi is not stored here.
                // Invalid p makes the link-wiggle target invalid instead of defaulting.
                if !is_valid_tweedie_power(self.scale) {
                    grad.fill(0.0);
                    return f64::NEG_INFINITY;
                }
                let p = self.scale;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    if !(eta_i.is_finite() && (-30.0..=30.0).contains(&eta_i)) {
                        grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    let mu = eta_i.exp().max(1e-300);
                    ll_acc +=
                        w_i * (y_i * mu.powf(1.0 - p) / (1.0 - p) - mu.powf(2.0 - p) / (2.0 - p));
                    residual[i] = w_i * (y_i - mu) * mu.powf(1.0 - p);
                }
                ll = ll_acc;
            }
            NutsFamily::NegativeBinomialLog => {
                let mut ll_acc = 0.0;
                // Family mapping: NegativeBinomial scale carries payload theta.
                // Invalid theta makes the link-wiggle target invalid instead of clamping.
                if !(self.scale.is_finite() && self.scale > 0.0) {
                    grad.fill(0.0);
                    return f64::NEG_INFINITY;
                }
                let theta = self.scale;
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    if !(eta_i.is_finite() && (-30.0..=30.0).contains(&eta_i)) {
                        grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    if w_i <= 0.0 {
                        residual[i] = 0.0;
                        continue;
                    }
                    let mu = eta_i.exp().max(1e-12);
                    let log_mu_term = if y_i > 0.0 { y_i * mu.ln() } else { 0.0 };
                    ll_acc += w_i
                        * (statrs::function::gamma::ln_gamma(y_i + theta)
                            - statrs::function::gamma::ln_gamma(theta)
                            - statrs::function::gamma::ln_gamma(y_i + 1.0)
                            + theta * (theta.ln() - (theta + mu).ln())
                            + log_mu_term
                            - y_i * (theta + mu).ln());
                    residual[i] = w_i * theta * (y_i - mu) / (theta + mu);
                }
                ll = ll_acc;
            }
            NutsFamily::GammaLog => {
                let mut ll_acc = 0.0;
                let shape = self.scale.max(1e-10);
                for i in 0..self.n_samples {
                    let eta_i = eta[i];
                    if !(eta_i.is_finite() && (-30.0..=30.0).contains(&eta_i)) {
                        grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                    let (y_i, w_i) = (self.y[i], self.weights[i]);
                    let mu = eta_i.exp();
                    ll_acc += w_i * shape * (-y_i / mu - eta_i);
                    residual[i] = w_i * shape * (y_i / mu - 1.0);
                }
                ll = ll_acc;
            }
        }

        // Penalty scaling factor: the Gaussian likelihood block above
        // multiplies through by `inv_scale_sq = 1/φ`, so the penalty
        // quadratic and its gradient must use the same factor — otherwise
        // the prior and likelihood live on incompatible scales and the
        // MAP-anchored posterior used for whitening picks up a hidden
        // φ-mismatch.
        //
        // Gamma uses the shape parameterization: `self.scale = α = 1/φ` and
        // the data-side log-likelihood already carries an explicit `shape`
        // factor (`w_i * shape * (-y_i/μ - eta_i)` and
        // `residual = w_i * shape * (y_i/μ - 1)`). Derivation: for the
        // shape-rate Gamma with `Var(Y) = φ·μ² = μ²/α`, the canonical-form
        // log-density score w.r.t. η on the log link is
        // `dℓ/dη = α·(y/μ − 1)`. The data term is therefore implicitly
        // `1/φ`-scaled via `shape = α`. The penalty must match that scale,
        // so use `1/φ = shape` here.
        //
        // Other non-Gaussian fixed-scale branches (Binomial / Poisson) have
        // `φ ≡ 1` and `1/φ = 1`, so the prior shape is unchanged.
        let penalty_scale = match self.nuts_family {
            NutsFamily::Gaussian => 1.0 / (self.scale * self.scale).max(1e-10),
            NutsFamily::GammaLog => self.scale.max(1e-10),
            _ => 1.0,
        };

        // Gradient w.r.t. θ (wiggle): ∂ℓ/∂θ = B(q₀)^T · residual − S_link · θ
        let s_link_theta = self.penalty_link.dot(&theta);
        let grad_theta = &fast_atv(&bwiggle, &residual) - &(&s_link_theta * penalty_scale);

        // Gradient w.r.t. β_eta: ∂ℓ/∂β = X^T · (residual ⊙ g'(q₀)) − S_base · β
        // where g'(q₀) = dη/dq₀ is the chain-rule factor
        let g_prime = self.compute_g_prime(&u, &theta);
        let r_scaled: Array1<f64> = residual
            .iter()
            .zip(g_prime.iter())
            .map(|(&r, &g)| r * g)
            .collect();
        let s_base_beta = self.penalty_base.dot(&beta);
        let grad_beta = &fast_atv(&self.x, &r_scaled) - &(&s_base_beta * penalty_scale);

        // Penalty (also φ-scaled for Gaussian; see `penalty_scale` above).
        let penalty =
            penalty_scale * (0.5 * beta.dot(&s_base_beta) + 0.5 * theta.dot(&s_link_theta));

        // Assemble joint gradient and transform to whitened space
        let mut grad_q = Array1::<f64>::zeros(dim);
        grad_q
            .slice_mut(ndarray::s![0..self.p_base])
            .assign(&grad_beta);
        grad_q
            .slice_mut(ndarray::s![self.p_base..])
            .assign(&grad_theta);
        fast_av_into(&self.chol_t, &grad_q, grad);
        ll - penalty
    }

    /// Get the Cholesky factor L for un-whitening samples.
    pub fn chol(&self) -> &Array2<f64> {
        &self.chol
    }

    /// Get the mode [β_eta; β_wiggle].
    pub fn mode_joint(&self) -> Array1<f64> {
        let dim = self.p_base + self.p_link;
        let mut mode = Array1::<f64>::zeros(dim);
        mode.slice_mut(ndarray::s![0..self.p_base])
            .assign(&self.mode_beta);
        mode.slice_mut(ndarray::s![self.p_base..])
            .assign(&self.mode_theta);
        mode
    }
}

impl HamiltonianTarget<Array1<f64>> for LinkWigglePosterior {
    fn logp_and_grad(&self, position: &Array1<f64>, grad: &mut Array1<f64>) -> f64 {
        self.compute_logp_and_grad_into(position, grad)
    }
}

/// Runs NUTS sampling for joint (β_eta, β_wiggle) in a link-wiggle model.
pub fn run_link_wiggle_nuts_sampling(
    x: ArrayView2<f64>,
    y: ArrayView1<f64>,
    weights: ArrayView1<f64>,
    penalty_base: ArrayView2<f64>,
    penalty_link: ArrayView2<f64>,
    mode_beta: ArrayView1<f64>,
    mode_theta: ArrayView1<f64>,
    hessian: ArrayView2<f64>,
    spline: LinkWiggleSplineArtifacts,
    nuts_family: NutsFamily,
    scale: f64,
    config: &NutsConfig,
) -> Result<NutsResult, String> {
    validate_nuts_target_accept(config.target_accept).map_err(String::from)?;
    let dim = mode_beta.len() + mode_theta.len();
    let target = LinkWigglePosterior::new(
        x,
        y,
        weights,
        penalty_base,
        penalty_link,
        mode_beta,
        mode_theta,
        hessian,
        spline,
        nuts_family,
        scale,
    )?;
    let chol = target.chol().clone();
    let mode_arr = target.mode_joint();

    let initial_positions = jittered_initial_positions(config, dim, 0.1, 0x8C48_0F65_3A2B_D917);

    let mass_cfg = robust_mass_matrix_config(dim, config.nwarmup);
    let mut sampler = GenericNUTS::new_with_mass_matrix(
        target,
        initial_positions,
        robust_target_accept(config.target_accept, dim),
        mass_cfg,
    )
    .set_seed(nuts_transition_seed(config.seed, 0x2E31_A4B6_C908_F57D));

    let (samples_array, run_stats) = sampler
        .run_progress(config.n_samples, config.nwarmup)
        .map_err(|e| format!("Link-wiggle NUTS sampling failed: {e}"))?;
    log::info!("Link-wiggle NUTS sampling complete: {}", run_stats);

    // Un-whiten samples: β = mode + L·z
    let shape = samples_array.shape();
    let n_chains = shape[0];
    let n_samples_out = shape[1];
    let total_samples = n_chains * n_samples_out;

    let mut samples = Array2::<f64>::zeros((total_samples, dim));
    let mut z_buffer = Array1::<f64>::zeros(dim);
    for chain in 0..n_chains {
        for sample_i in 0..n_samples_out {
            z_buffer.assign(&samples_array.slice(ndarray::s![chain, sample_i, ..]));
            samples
                .row_mut(chain * n_samples_out + sample_i)
                .assign(&(&mode_arr + &chol.dot(&z_buffer)));
        }
    }

    let posterior_mean = samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(dim));
    let posterior_std = samples.std_axis(Axis(0), 0.0);
    let (rhat, ess) = compute_split_rhat_and_ess(&samples_array);
    let converged = rhat < 1.1 && ess > 100.0;

    Ok(NutsResult {
        samples,
        posterior_mean,
        posterior_std,
        rhat,
        ess,
        converged,
    })
}

// ============================================================================
// Joint (β, ρ) HMC for Skewed Posteriors
// ============================================================================
//
// When the Laplace approximation to the marginal likelihood is unreliable
// (high posterior skewness), we bypass LAML entirely and sample from the
// joint posterior p(β, ρ | y) ∝ p(y|β) p(β|ρ) p(ρ).
//
// The joint log-posterior is:
//   log p(β, ρ | y) = ℓ(y|β) + Φ(β) [if Firth]
//                    - 0.5 β'S(ρ)β + 0.5 log|S(ρ)|_+ + log p(ρ) + const
//
// Gradients:
//   ∇_β: ∇_β ℓ + ∇_β Φ(β) [if Firth] - S(ρ) β
//   ∂/∂ρ_k: -0.5 λ_k β'S_k β + 0.5 tr(S_+⁻¹ A_k) + ∂log p(ρ)/∂ρ_k
//
// This completely avoids the Laplace approximation. When Firth bias reduction
// is active, the sampled target also includes the Jeffreys term Φ(β) in
// addition to the smoothing-parameter prior.

/// Directional cubic non-Gaussianity diagnostic for the Laplace approximation.
///
/// For each positive-curvature Hessian eigenpair `(lambda_r, v_r)`, this computes
///
///   gamma_r = T[v_r, v_r, v_r] / lambda_r^(3/2)
///            = Σ_i c_i (x_i^T v_r)^3 / lambda_r^(3/2),
///
/// and reports `max_r |gamma_r|`. This is invariant to arbitrary coordinate
/// relabeling and uses the full directional cubic contraction rather than only
/// diagonal tensor entries.
pub fn laplace_directional_cubic_diagnostic(
    hessian: &Array2<f64>,
    design: &DesignMatrix,
    c_weights: &Array1<f64>,
) -> Result<(f64, Array1<f64>), String> {
    let p = hessian.nrows();
    if p == 0 || hessian.ncols() != p {
        return Ok((0.0, Array1::zeros(0)));
    }

    let sym_h = (hessian + &hessian.t()) * 0.5;
    let (evals, evecs) = sym_h
        .eigh(Side::Lower)
        .map_err(|e| format!("directional cubic diagnostic eigendecomposition failed: {e}"))?;
    let max_eval = evals.iter().fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
    let tol = (max_eval * 1.0e-12).max(1.0e-14);
    let mut directional = Array1::<f64>::zeros(p);
    let mut max_abs = 0.0_f64;

    // Build the whitening transform L^{-1} where H = L L^T, so that
    // the standardized cubic along whitened direction u is:
    //   gamma(u) = T[L^{-T}u, L^{-T}u, L^{-T}u]  for ||u||=1
    // Eigenvector directions v_r satisfy u_r = lambda_r^{1/2} v_r (after
    // appropriate normalization), so gamma_r = T[v_r,v_r,v_r] / lambda_r^{3/2}.

    // Phase 1: evaluate gamma_r for all positive-curvature eigenvectors.
    for r in 0..p {
        let lambda = evals[r];
        if lambda <= tol {
            continue;
        }
        let v = evecs.column(r);
        let gamma = directional_cubic_contraction(design, c_weights, &v) / lambda.powf(1.5);
        directional[r] = if gamma.is_finite() { gamma } else { 0.0 };
        max_abs = max_abs.max(directional[r].abs());
    }

    // Phase 2: power-iteration refinement in whitened space.
    //
    // The supremum of |gamma(u)| over ||u||_H=1 can exceed the max over
    // eigenvectors. We approximate it with a few rounds of cubic power
    // iteration: given current direction v, the gradient of T[v,v,v] w.r.t.
    // v on the H-unit sphere is 3 T[·,v,v] projected onto the tangent space.
    // Since T[·,v,v] = X^T diag(c_i (x_i^T v)^2) which is a matrix-vector
    // product, each iteration is O(np).
    //
    // We seed from the eigenvector with largest |gamma_r| and also from a
    // few random probe directions.
    if p >= 2 {
        // Build H^{-1/2} columns for whitening: H^{-1/2} = V diag(1/sqrt(lam)) V^T
        // We need it to map whitened u -> original v = H^{-1/2} u, and
        // H^{1/2} to project back: H^{1/2} v = V diag(sqrt(lam)) V^T v.
        let positive_mask: Vec<bool> = evals.iter().map(|&ev| ev > tol).collect();
        let n_pos = positive_mask.iter().filter(|&&m| m).count();
        if n_pos >= 2 {
            let max_abs_from_probes = cubic_power_iteration_refinement(
                design,
                c_weights,
                &evals,
                &evecs,
                &positive_mask,
                n_pos,
            );
            if max_abs_from_probes > max_abs {
                max_abs = max_abs_from_probes;
            }
        }
    }

    Ok((max_abs, directional))
}

/// Compute T[v,v,v] = Σ_i c_i (x_i^T v)^3 for a given direction v.
fn directional_cubic_contraction(
    design: &DesignMatrix,
    c_weights: &Array1<f64>,
    v: &ArrayView1<f64>,
) -> f64 {
    match design.as_sparse() {
        Some(x_sparse) => {
            let (symbolic, values) = x_sparse.as_ref().parts();
            let col_ptr = symbolic.col_ptr();
            let row_idx = symbolic.row_idx();
            let mut row_scores = vec![0.0_f64; x_sparse.nrows()];
            for col in 0..x_sparse.ncols() {
                let coeff = v[col];
                for ptr in col_ptr[col]..col_ptr[col + 1] {
                    row_scores[row_idx[ptr]] += values[ptr] * coeff;
                }
            }
            let mut cubic = 0.0_f64;
            for i in 0..row_scores.len().min(c_weights.len()) {
                cubic += c_weights[i] * row_scores[i].powi(3);
            }
            cubic
        }
        None => {
            let x_dense = design.to_dense_cow();
            let x_dense = x_dense.as_ref();
            let mut cubic = 0.0_f64;
            for i in 0..x_dense.nrows().min(c_weights.len()) {
                let proj = x_dense.row(i).dot(v);
                cubic += c_weights[i] * proj.powi(3);
            }
            cubic
        }
    }
}

/// Compute the gradient of T[v,v,v] w.r.t. v:  3 X^T diag(c_i (x_i^T v)^2) 1.
/// More precisely: ∂/∂v T[v,v,v] = 3 Σ_i c_i (x_i^T v)^2 x_i.
fn directional_cubic_gradient(
    design: &DesignMatrix,
    c_weights: &Array1<f64>,
    v: &Array1<f64>,
) -> Array1<f64> {
    let p = v.len();
    match design.as_sparse() {
        Some(x_sparse) => {
            let (symbolic, values) = x_sparse.as_ref().parts();
            let col_ptr = symbolic.col_ptr();
            let row_idx = symbolic.row_idx();
            let n = x_sparse.nrows();
            let mut row_scores = vec![0.0_f64; n];
            for col in 0..x_sparse.ncols() {
                let coeff = v[col];
                for ptr in col_ptr[col]..col_ptr[col + 1] {
                    row_scores[row_idx[ptr]] += values[ptr] * coeff;
                }
            }
            // quadratic weights: 3 c_i (x_i^T v)^2
            let mut quad_weights = vec![0.0_f64; n];
            for i in 0..n.min(c_weights.len()) {
                quad_weights[i] = 3.0 * c_weights[i] * row_scores[i] * row_scores[i];
            }
            // X^T quad_weights
            let mut grad = Array1::<f64>::zeros(p);
            for col in 0..x_sparse.ncols() {
                let mut acc = 0.0_f64;
                for ptr in col_ptr[col]..col_ptr[col + 1] {
                    acc += values[ptr] * quad_weights[row_idx[ptr]];
                }
                grad[col] = acc;
            }
            grad
        }
        None => {
            let x_dense = design.to_dense_cow();
            let x_dense = x_dense.as_ref();
            let n = x_dense.nrows();
            let mut grad = Array1::<f64>::zeros(p);
            for i in 0..n.min(c_weights.len()) {
                let proj = x_dense.row(i).dot(v);
                let w = 3.0 * c_weights[i] * proj * proj;
                // scaled_add works with any ArrayBase reference.
                let row = x_dense.row(i);
                for j in 0..p {
                    grad[j] += w * row[j];
                }
            }
            grad
        }
    }
}

/// Power-iteration refinement for the supremum of |gamma(u)| over ||u||_H = 1.
///
/// Seeds from the best eigenvector direction plus deterministic probe
/// directions constructed from pairs of eigenvectors. Runs a few Riemannian
/// gradient ascent steps on the whitened unit sphere.
fn cubic_power_iteration_refinement(
    design: &DesignMatrix,
    c_weights: &Array1<f64>,
    evals: &Array1<f64>,
    evecs: &Array2<f64>,
    positive_mask: &[bool],
    n_pos: usize,
) -> f64 {
    let p = evals.len();
    let max_probes = 8;
    let max_iters = 5;

    // Helper: convert whitened u -> original v = Σ_r (u_r / sqrt(lam_r)) * evec_r
    // (only over positive eigenspace).
    let to_original = |u: &Array1<f64>| -> Array1<f64> {
        let mut v = Array1::<f64>::zeros(p);
        let mut idx = 0;
        for r in 0..p {
            if positive_mask[r] {
                let scale = u[idx] / evals[r].sqrt();
                let col = evecs.column(r);
                for j in 0..p {
                    v[j] += scale * col[j];
                }
                idx += 1;
            }
        }
        v
    };

    // Helper: project original-space vector to whitened: u_j = sqrt(lam_r) (evec_r^T g)
    let to_whitened = |g: &Array1<f64>| -> Array1<f64> {
        let mut u = Array1::<f64>::zeros(n_pos);
        let mut idx = 0;
        for r in 0..p {
            if positive_mask[r] {
                u[idx] = evals[r].sqrt() * evecs.column(r).dot(g);
                idx += 1;
            }
        }
        u
    };

    // Evaluate |gamma(u)| for whitened direction u.
    let eval_gamma = |u: &Array1<f64>| -> f64 {
        let norm = u.dot(u).sqrt();
        if norm < 1e-30 {
            return 0.0;
        }
        let u_normed: Array1<f64> = u / norm;
        let v = to_original(&u_normed);
        // gamma = T[v,v,v] since v already has ||v||_H = 1
        let cubic = directional_cubic_contraction(design, c_weights, &v.view());
        if cubic.is_finite() { cubic.abs() } else { 0.0 }
    };

    // One step of Riemannian gradient ascent on the whitened sphere for |T[v,v,v]|.
    let refine_step = |u: &Array1<f64>| -> Array1<f64> {
        let norm = u.dot(u).sqrt();
        if norm < 1e-30 {
            return u.clone();
        }
        let u_normed: Array1<f64> = u / norm;
        let v = to_original(&u_normed);
        // Gradient of T[v,v,v] w.r.t. v in original space
        let grad_v = directional_cubic_gradient(design, c_weights, &v);
        // Map to whitened space
        let mut grad_u = to_whitened(&grad_v);
        // Project onto tangent plane of sphere: grad - (grad . u) u
        let dot = grad_u.dot(&u_normed);
        grad_u.scaled_add(-dot, &u_normed);
        // Sign: we want to maximize |T|, so follow sign(T) * grad
        let cubic_val = directional_cubic_contraction(design, c_weights, &v.view());
        let sign = if cubic_val >= 0.0 { 1.0 } else { -1.0 };
        let step_size = 0.3;
        let mut u_new = &u_normed + &(&grad_u * (sign * step_size));
        let new_norm = u_new.dot(&u_new).sqrt();
        if new_norm > 1e-30 {
            u_new /= new_norm;
        }
        u_new
    };

    let mut best = 0.0_f64;

    // Build seed directions:
    // (a) The eigenvector with largest |gamma_r| (already computed by caller,
    //     but we re-derive the whitened form here).
    // (b) Deterministic probe directions from pairs of top eigenvectors:
    //     (e_i + e_j) / sqrt(2) and (e_i - e_j) / sqrt(2) in whitened space.
    let mut seeds: Vec<Array1<f64>> = Vec::with_capacity(max_probes);

    // Seed (a): each eigenvector is a standard basis vector in whitened space.
    // Find the one with largest |gamma|.
    let mut best_eig_idx = 0;
    let mut best_eig_gamma = 0.0_f64;
    for j in 0..n_pos {
        let mut u = Array1::<f64>::zeros(n_pos);
        u[j] = 1.0;
        let g = eval_gamma(&u);
        if g > best_eig_gamma {
            best_eig_gamma = g;
            best_eig_idx = j;
        }
    }
    best = best.max(best_eig_gamma);
    let mut u_best = Array1::<f64>::zeros(n_pos);
    u_best[best_eig_idx] = 1.0;
    seeds.push(u_best);

    // Seed (b): pairwise combinations of the top few eigenvectors.
    let n_top = n_pos.min(4);
    for i in 0..n_top {
        for j in (i + 1)..n_top {
            if seeds.len() >= max_probes {
                break;
            }
            let inv_sqrt2 = std::f64::consts::FRAC_1_SQRT_2;
            let mut u_plus = Array1::<f64>::zeros(n_pos);
            u_plus[i] = inv_sqrt2;
            u_plus[j] = inv_sqrt2;
            seeds.push(u_plus);
            if seeds.len() < max_probes {
                let mut u_minus = Array1::<f64>::zeros(n_pos);
                u_minus[i] = inv_sqrt2;
                u_minus[j] = -inv_sqrt2;
                seeds.push(u_minus);
            }
        }
    }

    // Run power iteration from each seed.
    for seed in &seeds {
        let mut u = seed.clone();
        for _ in 0..max_iters {
            u = refine_step(&u);
        }
        let g = eval_gamma(&u);
        best = best.max(g);
    }

    best
}

/// Result of joint (β, ρ) sampling.
#[derive(Clone, Debug)]
pub struct JointBetaRhoResult {
    /// Coefficient samples: shape (n_total_samples, n_beta)
    pub beta_samples: Array2<f64>,
    /// Log-smoothing parameter samples: shape (n_total_samples, n_rho)
    pub rho_samples: Array2<f64>,
    /// Posterior mean of β
    pub beta_mean: Array1<f64>,
    /// Adaptive inverse-link parameter samples: shape (n_total_samples, n_link_params)
    pub link_param_samples: Array2<f64>,
    /// Posterior mean of adaptive inverse-link parameters
    pub link_param_mean: Array1<f64>,
    /// Posterior mean of ρ
    pub rho_mean: Array1<f64>,
    /// R-hat diagnostic
    pub rhat: f64,
    /// Effective sample size
    pub ess: f64,
    /// Whether sampling converged
    pub converged: bool,
    /// Max skewness that triggered this sampling
    pub trigger_skewness: f64,
}

/// Joint (β, ρ) posterior target for NUTS.
///
/// Samples from p(β, ρ | y) ∝ p(y|β) p(β|ρ) p(ρ) directly,
/// completely bypassing the Laplace approximation.
///
/// The parameter vector is [z_β; ρ] where z_β = L⁻¹(β - μ) is the
/// whitened β, ρ is the raw log-smoothing parameters, and adaptive inverse-link
/// parameters follow when the binomial link has fitted shape/mixing parameters.
struct JointBetaRhoPosterior {
    data: SharedData,
    /// L where LL' = H⁻¹ (whitening for β block)
    chol: Array2<f64>,
    /// L' for chain rule
    chol_t: Array2<f64>,
    /// Joint likelihood specification (response + parameterized link).
    likelihood: LikelihoodSpec,
    /// Dimension of β
    n_beta: usize,
    /// Dimension of ρ
    n_rho: usize,
    /// Dimension of adaptive inverse-link parameters
    n_link_params: usize,
    /// LAML-converged adaptive inverse-link parameters (used only to initialize chains)
    link_param_mode: Array1<f64>,
    /// Canonical penalties in the transformed basis.
    penalty_canonical: Vec<crate::construction::CanonicalPenalty>,
    /// Fixed prior on rho used by the sampled target.
    rho_prior: RhoPrior,
    /// LAML-converged ρ (used only to initialize chains)
    rho_mode: Array1<f64>,
    /// Whether to add the identifiable-subspace Jeffreys/Firth term to the
    /// target
    firth_enabled: bool,
    /// One-deep cache for the structural penalty pseudo-logdet and its
    /// ρ-gradient. NUTS tree-doubling and U-turn checks repeatedly evaluate
    /// the joint log-posterior at the same `rho` bytes, so a single-slot
    /// cache keyed on the exact f64 bit pattern of `rho` avoids redundant
    /// SVD/eigendecompositions inside `PenaltyPseudologdet::from_penalties`.
    /// `Mutex` (not `RefCell`) because chains share the target via
    /// `Arc<Target>` and run in parallel via rayon.
    penalty_logdet_cache: Mutex<Option<(u64, f64, Array1<f64>)>>,
}

impl JointBetaRhoPosterior {
    fn new(
        x: ArrayView2<f64>,
        y: ArrayView1<f64>,
        weights: ArrayView1<f64>,
        mode: ArrayView1<f64>,
        hessian: ArrayView2<f64>,
        penalty_canonical: Vec<crate::construction::CanonicalPenalty>,
        rho_mode: ArrayView1<f64>,
        likelihood: LikelihoodSpec,
        gamma_shape: Option<f64>,
        rho_prior: RhoPrior,
        firth_enabled: bool,
    ) -> Result<Self, String> {
        let n_samples = x.nrows();
        let n_beta = x.ncols();
        let n_rho = penalty_canonical.len();

        if rho_mode.len() != n_rho {
            return Err(HmcError::DimensionMismatch {
                reason: format!(
                    "rho_mode length {} != penalty count {}",
                    rho_mode.len(),
                    n_rho
                ),
            }
            .into());
        }

        match (&likelihood.response, &likelihood.link) {
            (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::Logit)) => {}
            (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::Probit)) => {}
            (ResponseFamily::Binomial, InverseLink::Standard(LinkFunction::CLogLog)) => {}
            (ResponseFamily::Binomial, InverseLink::LatentCLogLog(_)) => {}
            (ResponseFamily::Binomial, InverseLink::Sas(_)) => {}
            (ResponseFamily::Binomial, InverseLink::BetaLogistic(_)) => {}
            (ResponseFamily::Binomial, InverseLink::Mixture(_)) => {}
            (ResponseFamily::Binomial, InverseLink::Standard(other)) => {
                return Err(HmcError::LinkMismatch {
                    reason: format!(
                        "Joint HMC binomial response requires a binomial-compatible inverse link; got {:?}",
                        other
                    ),
                }
                .into());
            }
            (ResponseFamily::Gaussian, InverseLink::Standard(LinkFunction::Identity)) => {}
            (ResponseFamily::Gaussian, _) => {
                return Err(HmcError::LinkMismatch {
                    reason: "Joint HMC Gaussian requires an identity inverse link".to_string(),
                }
                .into());
            }
            (
                ResponseFamily::Poisson
                | ResponseFamily::Tweedie { .. }
                | ResponseFamily::NegativeBinomial { .. }
                | ResponseFamily::Gamma,
                InverseLink::Standard(LinkFunction::Log),
            ) => {}
            (
                ResponseFamily::Poisson
                | ResponseFamily::Tweedie { .. }
                | ResponseFamily::NegativeBinomial { .. }
                | ResponseFamily::Gamma,
                _,
            ) => {
                return Err(HmcError::LinkMismatch {
                    reason: "Joint HMC log-link family requires a log inverse link".to_string(),
                }
                .into());
            }
            (ResponseFamily::Beta { .. }, InverseLink::Standard(LinkFunction::Logit)) => {}
            (ResponseFamily::Beta { .. }, _) => {
                return Err(HmcError::LinkMismatch {
                    reason: "Joint HMC Beta requires a logit inverse link".to_string(),
                }
                .into());
            }
            (ResponseFamily::RoystonParmar, _) => {
                return Err(HmcError::UnsupportedFamily {
                    reason: "Joint HMC fallback is not implemented for RoystonParmar".to_string(),
                }
                .into());
            }
        }

        validate_firth_likelihood_support(&likelihood, firth_enabled).map_err(String::from)?;
        if matches!(likelihood.response, ResponseFamily::NegativeBinomial { .. }) {
            validate_count_responses("negative-binomial joint HMC", &y, &weights)
                .map_err(String::from)?;
        }
        if likelihood.is_binomial() {
            validate_binary_responses("binomial joint HMC", &y, &weights).map_err(String::from)?;
        }

        // Cholesky of H for β-whitening (same as NutsPosterior)
        let hessian_owned = hessian.to_owned();
        let chol_factor = hessian_owned
            .cholesky(Side::Lower)
            .map_err(|e| format!("Joint HMC: Hessian Cholesky failed: {:?}", e))?;
        let l_h = chol_factor.lower_triangular();
        let chol = solve_upper_triangular_transpose(&l_h, n_beta);
        let chol_t = chol.t().to_owned();

        let data = SharedData {
            x: Arc::new(x.to_owned()),
            y: Arc::new(y.to_owned()),
            weights: Arc::new(weights.to_owned()),
            mode: Arc::new(mode.to_owned()),
            gamma_shape: gamma_shape.unwrap_or(1.0),
            // Joint (β, ρ) HMC keeps the likelihood on its native scale;
            // dispersion enters via the per-family scale parameter, not
            // via the whitening transform here. `Known(1.0)` matches the
            // pre-refactor behaviour for this code path.
            dispersion: crate::solver::estimate::Dispersion::Known(1.0),
            n_samples,
            dim: n_beta,
        };
        let link_param_mode = Self::link_param_mode(&likelihood.link);

        Ok(Self {
            data,
            chol,
            chol_t,
            likelihood,
            n_beta,
            n_rho,
            n_link_params: link_param_mode.len(),
            link_param_mode,
            penalty_canonical,
            rho_prior,
            rho_mode: rho_mode.to_owned(),
            firth_enabled,
            penalty_logdet_cache: Mutex::new(None),
        })
    }

    /// FNV-1a hash over the raw f64 bit pattern of `rho`.
    ///
    /// NUTS leapfrog / tree-doubling / U-turn checks revisit identical
    /// position vectors byte-for-byte, so exact-equality on `to_bits()`
    /// captures the dominant repetition pattern without any tolerance.
    #[inline]
    fn hash_rho(rho: ndarray::ArrayView1<f64>) -> u64 {
        let mut h: u64 = 0xcbf2_9ce4_8422_2325;
        for &x in rho.iter() {
            h ^= x.to_bits();
            h = h.wrapping_mul(0x0000_0100_0000_01b3);
        }
        h
    }

    fn link_param_mode(inverse_link: &InverseLink) -> Array1<f64> {
        match inverse_link {
            InverseLink::Sas(state) | InverseLink::BetaLogistic(state) => {
                Array1::from_vec(vec![state.epsilon, state.log_delta])
            }
            InverseLink::Mixture(state) => state.rho.clone(),
            InverseLink::Standard(_) | InverseLink::LatentCLogLog(_) => Array1::zeros(0),
        }
    }

    fn inverse_link_with_params(
        &self,
        link_params: ndarray::ArrayView1<'_, f64>,
    ) -> Result<InverseLink, String> {
        match &self.likelihood.link {
            InverseLink::Sas(_) => {
                if link_params.len() != 2 {
                    return Err(format!(
                        "SAS link parameter length must be 2, got {}",
                        link_params.len()
                    ));
                }
                Ok(InverseLink::Sas(crate::types::SasLinkState::new(
                    link_params[0],
                    link_params[1],
                )?))
            }
            InverseLink::BetaLogistic(_) => {
                if link_params.len() != 2 || !link_params.iter().all(|v| v.is_finite()) {
                    return Err(
                        "Beta-Logistic link parameters must be finite with length 2".to_string()
                    );
                }
                Ok(InverseLink::BetaLogistic(crate::types::SasLinkState {
                    epsilon: link_params[0],
                    log_delta: link_params[1],
                    delta: link_params[1].exp(),
                }))
            }
            InverseLink::Mixture(state) => {
                let rho = link_params.to_owned();
                Ok(InverseLink::Mixture(crate::types::MixtureLinkState {
                    components: state.components.clone(),
                    pi: softmax_last_fixedzero(&rho),
                    rho,
                }))
            }
            InverseLink::Standard(_) | InverseLink::LatentCLogLog(_) => {
                Ok(self.likelihood.link.clone())
            }
        }
    }

    /// Compute the joint log-posterior and gradient.
    ///
    /// The joint log-posterior is:
    ///   log p(β, ρ | y) = ℓ(y|β) + ½ log|I(β)| [if Firth]
    ///                    − ½β'S(ρ)β + ½ log|S(ρ)|₊ + log p(ρ) + const
    ///
    /// This is NOT the REML/LAML objective (which integrates out β). Here β is
    /// an explicit parameter being sampled, evaluated at arbitrary values — not
    /// just at the mode β̂(ρ).
    ///
    /// Parameter vector layout: [z_β (whitened, length n_beta); ρ (length n_rho);
    /// adaptive inverse-link params (length n_link_params)]
    fn compute_joint_logp_and_grad_into(
        &self,
        params: &Array1<f64>,
        out_grad: &mut Array1<f64>,
    ) -> f64 {
        let n_beta = self.n_beta;
        let n_rho = self.n_rho;
        let n_link_params = self.n_link_params;

        // Split parameter vector — keep as views to avoid two per-step
        // `to_owned()` allocations of size n_beta and n_rho.
        let z = params.slice(ndarray::s![..n_beta]);
        let rho = params.slice(ndarray::s![n_beta..n_beta + n_rho]);
        let link_params = params.slice(ndarray::s![n_beta + n_rho..]);
        let lambdas: Array1<f64> = rho.mapv(f64::exp);

        let inverse_link = match self.inverse_link_with_params(link_params) {
            Ok(link) => link,
            Err(err) => {
                log::warn!(
                    "[Joint HMC] adaptive inverse-link parameters are invalid: {}",
                    err
                );
                out_grad.fill(0.0);
                return f64::NEG_INFINITY;
            }
        };

        // Un-whiten: β = μ + L z
        let beta = self.data.mode.as_ref() + &self.chol.dot(&z);

        // η = X β
        let eta = crate::faer_ndarray::fast_av(self.data.x.as_ref(), &beta);

        // ---- Log-likelihood ℓ(y|β) and ∇_β ℓ ----
        let step_likelihood = LikelihoodSpec {
            response: self.likelihood.response.clone(),
            link: inverse_link,
        };
        let (ll, mut grad_ll_beta, grad_link) = match joint_family_logp_grad_and_link_grad(
            &step_likelihood,
            &self.data,
            &eta,
            n_link_params,
        ) {
            Ok(value) => value,
            Err(err) => {
                log::warn!(
                    "[Joint HMC] likelihood target became invalid at the current state: {}",
                    err
                );
                out_grad.fill(0.0);
                return f64::NEG_INFINITY;
            }
        };

        let mut firth_logdet = 0.0;
        if self.firth_enabled {
            match firth_jeffreys_logp_and_grad(NutsFamily::BinomialLogit, &self.data, &eta) {
                Ok((value, grad_beta_firth)) => {
                    firth_logdet = value;
                    grad_ll_beta += &grad_beta_firth;
                }
                Err(err) => {
                    log::warn!(
                        "[Joint HMC/Firth] Jeffreys target became invalid at the current state: {}",
                        err
                    );
                    out_grad.fill(0.0);
                    return f64::NEG_INFINITY;
                }
            }
        }

        // ---- Penalty: -0.5 β'S(ρ)β ----
        // S(ρ) = Σ_k λ_k S_k where S_k = R_k'R_k (precomputed in penalty_matrices).
        // Uses penalty_roots for the efficient ||R_k β||² form.
        let mut penalty_val = 0.0;
        let mut s_beta = Array1::<f64>::zeros(n_beta);
        let mut grad_rho = Array1::<f64>::zeros(n_rho);

        // Reuse one max-rank scratch buffer for r_beta = R_k · β_block across
        // all penalty blocks instead of allocating a fresh Array1 per block
        // per HMC step.
        let max_rank = self
            .penalty_canonical
            .iter()
            .map(|cp| cp.rank())
            .max()
            .unwrap_or(0);
        let mut r_beta_scratch = Array1::<f64>::zeros(max_rank);

        for (k, cp) in self.penalty_canonical.iter().enumerate() {
            // Block-local quadratic: β'S_k β via root
            let r = &cp.col_range;
            let beta_block = beta.slice(ndarray::s![r.start..r.end]);
            let rank_k = cp.rank();
            crate::faer_ndarray::fast_av_view_into(
                &cp.root,
                &beta_block,
                r_beta_scratch.slice_mut(ndarray::s![..rank_k]),
            );
            let r_beta = r_beta_scratch.slice(ndarray::s![..rank_k]);
            let quad_k = r_beta.dot(&r_beta);
            penalty_val += 0.5 * lambdas[k] * quad_k;

            // Accumulate S(ρ)β for β-gradient — block-local
            for a in 0..cp.block_dim() {
                let val: f64 = (0..rank_k).map(|row| cp.root[[row, a]] * r_beta[row]).sum();
                s_beta[r.start + a] += lambdas[k] * val;
            }

            // ρ_k gradient from penalty
            grad_rho[k] = -0.5 * lambdas[k] * quad_k;
        }

        // ---- Structural penalty log-determinant: +0.5 log|S(ρ)|₊ and ρ-derivatives ----
        //
        // One-deep cache keyed on the exact f64 bits of `rho`: NUTS tree
        // doubling revisits identical positions byte-for-byte, so an
        // exact-equality cache eliminates the dominant SVD/eigendecomp
        // cost in `PenaltyPseudologdet::from_penalties` across leapfrog
        // half-steps.
        let log_det_s = if self.penalty_canonical.is_empty() {
            0.0
        } else {
            let rho_hash = Self::hash_rho(rho);
            let cached = self.penalty_logdet_cache.lock().ok().and_then(|guard| {
                guard.as_ref().and_then(|(h, v, g)| {
                    if *h == rho_hash && g.len() == n_rho {
                        for k in 0..n_rho {
                            grad_rho[k] += 0.5 * g[k];
                        }
                        Some(*v)
                    } else {
                        None
                    }
                })
            });
            if let Some(hit) = cached {
                hit
            } else {
                match PenaltyPseudologdet::from_penalties(
                    &self.penalty_canonical,
                    lambdas.as_slice().unwrap_or(&[]),
                    0.0,
                    n_beta,
                ) {
                    Ok(pld) => {
                        let (det1, _) = pld.rho_derivatives_from_penalties(
                            &self.penalty_canonical,
                            lambdas.as_slice().unwrap_or(&[]),
                        );
                        let value = pld.value();
                        if let Ok(mut guard) = self.penalty_logdet_cache.lock() {
                            *guard = Some((rho_hash, value, det1.clone()));
                        }
                        for k in 0..n_rho {
                            grad_rho[k] += 0.5 * det1[k];
                        }
                        value
                    }
                    Err(err) => {
                        log::warn!(
                            "[Joint HMC] structural penalty logdet became invalid at the current state: {}",
                            err
                        );
                        out_grad.fill(0.0);
                        return f64::NEG_INFINITY;
                    }
                }
            }
        };

        // ---- Prior on ρ ----
        let mut rho_prior = 0.0;
        match &self.rho_prior {
            RhoPrior::Flat => {}
            RhoPrior::Normal { mean, sd } => {
                let inv_var = 1.0 / (*sd * *sd);
                for k in 0..n_rho {
                    let d = rho[k] - *mean;
                    rho_prior -= 0.5 * inv_var * d * d;
                    grad_rho[k] -= inv_var * d;
                }
            }
            RhoPrior::GammaPrecision { shape, rate } => {
                for k in 0..n_rho {
                    let lambda = rho[k].exp();
                    rho_prior += (*shape - 1.0) * rho[k] - *rate * lambda;
                    grad_rho[k] += (*shape - 1.0) - *rate * lambda;
                }
            }
            RhoPrior::Independent(priors) => {
                if priors.len() != n_rho {
                    out_grad.fill(0.0);
                    return f64::NEG_INFINITY;
                }
                for k in 0..n_rho {
                    match &priors[k] {
                        RhoPrior::Flat => {}
                        RhoPrior::Normal { mean, sd } => {
                            let inv_var = 1.0 / (*sd * *sd);
                            let d = rho[k] - *mean;
                            rho_prior -= 0.5 * inv_var * d * d;
                            grad_rho[k] -= inv_var * d;
                        }
                        RhoPrior::GammaPrecision { shape, rate } => {
                            let lambda = rho[k].exp();
                            rho_prior += (*shape - 1.0) * rho[k] - *rate * lambda;
                            grad_rho[k] += (*shape - 1.0) - *rate * lambda;
                        }
                        RhoPrior::Independent(_) => {
                            out_grad.fill(0.0);
                            return f64::NEG_INFINITY;
                        }
                    }
                }
            }
        }

        // ---- Assemble ----
        let logp = ll + firth_logdet - penalty_val + 0.5 * log_det_s + rho_prior;

        // β-gradient in original space: ∇_β ℓ - S(ρ)β
        let grad_beta = &grad_ll_beta - &s_beta;

        // Combined gradient: [∇_z; ∇_ρ; ∇_link]
        crate::faer_ndarray::fast_av_view_into(
            &self.chol_t,
            &grad_beta,
            out_grad.slice_mut(ndarray::s![..n_beta]),
        );
        out_grad
            .slice_mut(ndarray::s![n_beta..n_beta + n_rho])
            .assign(&grad_rho);
        out_grad
            .slice_mut(ndarray::s![n_beta + n_rho..])
            .assign(&grad_link);

        logp
    }
}

impl HamiltonianTarget<Array1<f64>> for JointBetaRhoPosterior {
    fn logp_and_grad(&self, position: &Array1<f64>, grad: &mut Array1<f64>) -> f64 {
        self.compute_joint_logp_and_grad_into(position, grad)
    }
}

/// Inputs for joint (β, ρ) sampling.
pub struct JointBetaRhoInputs<'a> {
    pub x: ArrayView2<'a, f64>,
    pub y: ArrayView1<'a, f64>,
    pub weights: ArrayView1<'a, f64>,
    pub likelihood: LikelihoodSpec,
    pub gamma_shape: Option<f64>,
    pub mode: ArrayView1<'a, f64>,
    pub hessian: ArrayView2<'a, f64>,
    pub penalty_roots: Vec<CanonicalPenalty>,
    pub rho_mode: ArrayView1<'a, f64>,
    pub rho_prior: RhoPrior,
    pub firth_bias_reduction: bool,
    /// Max posterior skewness that triggered this sampling
    pub trigger_skewness: f64,
}

/// Run joint (β, ρ) NUTS sampling.
///
/// This is the automatic fallback when the Laplace approximation has high
/// skewness. It samples from the true joint posterior, completely bypassing
/// the Laplace approximation for smoothing parameter selection.
pub fn run_joint_beta_rho_sampling(
    inputs: &JointBetaRhoInputs<'_>,
    config: &NutsConfig,
) -> Result<JointBetaRhoResult, String> {
    validate_firth_likelihood_support(&inputs.likelihood, inputs.firth_bias_reduction)
        .map_err(String::from)?;
    validate_nuts_target_accept(config.target_accept).map_err(String::from)?;
    let n_beta = inputs.mode.len();
    let n_rho = inputs.penalty_roots.len();
    let n_link_params = JointBetaRhoPosterior::link_param_mode(&inputs.likelihood.link).len();
    let total_dim = n_beta + n_rho + n_link_params;

    log::info!(
        "[Joint HMC] Sampling (β, ρ, link) jointly: {} β-params + {} ρ-params + {} link-params = {} total (triggered by skewness {:.3})",
        n_beta,
        n_rho,
        n_link_params,
        total_dim,
        inputs.trigger_skewness,
    );

    let target = JointBetaRhoPosterior::new(
        inputs.x,
        inputs.y,
        inputs.weights,
        inputs.mode,
        inputs.hessian,
        inputs.penalty_roots.clone(),
        inputs.rho_mode,
        inputs.likelihood.clone(),
        inputs.gamma_shape,
        inputs.rho_prior.clone(),
        inputs.firth_bias_reduction,
    )?;

    let chol = target.chol.clone();
    let mode_arr = target.data.mode.clone();
    let rho_mode = target.rho_mode.clone();
    let link_param_mode = target.link_param_mode.clone();

    // Initialize chains: z_β at 0 (= mode), ρ at rho_mode, link params at fitted state.
    let initial_positions: Vec<Array1<f64>> = (0..config.n_chains)
        .map(|chain| {
            let mut rng =
                StdRng::seed_from_u64(chain_stream_seed(config.seed, chain, 0x9B51_6E37_F2D0_A48C));
            let mut pos = Array1::<f64>::zeros(total_dim);
            // Small jitter for β (whitened space)
            for j in 0..n_beta {
                pos[j] = sample_standard_normal(&mut rng) * 0.1;
            }
            // Small jitter for ρ around mode
            for k in 0..n_rho {
                pos[n_beta + k] = rho_mode[k] + sample_standard_normal(&mut rng) * 0.2;
            }
            // Small jitter for adaptive link parameters around fitted state
            for k in 0..n_link_params {
                pos[n_beta + n_rho + k] =
                    link_param_mode[k] + sample_standard_normal(&mut rng) * 0.05;
            }
            pos
        })
        .collect();

    // Keep warmup covariance phase-local: diagonal windows are less likely to
    // encode cross-block covariance from a transient mode switch.
    let mass_cfg = robust_mass_matrix_config(total_dim, config.nwarmup);

    let mut sampler = GenericNUTS::new_with_mass_matrix(
        target,
        initial_positions,
        robust_target_accept(config.target_accept, total_dim),
        mass_cfg,
    )
    .set_seed(nuts_transition_seed(config.seed, 0x63AF_175B_D820_C94E));

    let (samples_array, run_stats) = sampler
        .run_progress(config.n_samples, config.nwarmup)
        .map_err(|e| format!("Joint (β,ρ) NUTS sampling failed: {e}"))?;
    log::info!("[Joint HMC] Sampling complete: {}", run_stats);

    // Unpack samples
    let shape = samples_array.shape();
    let n_chains = shape[0];
    let n_samples_out = shape[1];
    let total_samples = n_chains * n_samples_out;

    let mut beta_samples = Array2::<f64>::zeros((total_samples, n_beta));
    let mut rho_samples = Array2::<f64>::zeros((total_samples, n_rho));
    let mut link_param_samples = Array2::<f64>::zeros((total_samples, n_link_params));

    for chain in 0..n_chains {
        for sample_i in 0..n_samples_out {
            let sample_idx = chain * n_samples_out + sample_i;
            let zview = samples_array.slice(ndarray::s![chain, sample_i, ..]);

            // Un-whiten β: β = μ + L z
            let z_beta = zview.slice(ndarray::s![..n_beta]).to_owned();
            let beta = mode_arr.as_ref() + &chol.dot(&z_beta);
            beta_samples.row_mut(sample_idx).assign(&beta);

            // ρ and adaptive link parameters are stored directly
            let rho_slice = zview.slice(ndarray::s![n_beta..n_beta + n_rho]);
            rho_samples.row_mut(sample_idx).assign(&rho_slice);
            let link_slice = zview.slice(ndarray::s![n_beta + n_rho..]);
            link_param_samples.row_mut(sample_idx).assign(&link_slice);
        }
    }

    let beta_mean = beta_samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(n_beta));
    let rho_mean = rho_samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(n_rho));
    let link_param_mean = link_param_samples
        .mean_axis(Axis(0))
        .unwrap_or_else(|| Array1::zeros(n_link_params));

    let (rhat, ess) = compute_split_rhat_and_ess(&samples_array);

    let converged = rhat < 1.1 && ess > 50.0;
    if !converged {
        log::warn!(
            "[Joint HMC] Convergence warning: R-hat={:.3}, ESS={:.1}",
            rhat,
            ess,
        );
    }

    Ok(JointBetaRhoResult {
        beta_samples,
        rho_samples,
        beta_mean,
        link_param_samples,
        link_param_mean,
        rho_mean,
        rhat,
        ess,
        converged,
        trigger_skewness: inputs.trigger_skewness,
    })
}

// ============================================================================
// Survival Model HMC Support
// ============================================================================

mod survival_hmc {
    use super::*;
    use crate::survival::{
        MonotonicityPenalty, PenaltyBlocks, SurvivalEngineInputs, SurvivalSpec,
        WorkingModelSurvival,
    };

    /// Shared data for survival NUTS posterior (wrapped in Arc to prevent cloning).
    #[derive(Clone)]
    struct SharedSurvivalData {
        /// Exact survival model in original spline coordinates.
        base_model: Arc<WorkingModelSurvival>,
        /// MAP estimate in coefficient coordinates.
        mode: Arc<Array1<f64>>,
    }

    /// Whitened log-posterior target for survival models with analytical gradients.
    #[derive(Clone)]
    pub struct SurvivalPosterior {
        /// Shared read-only data (Arc prevents duplication)
        data: SharedSurvivalData,
        /// Transform: L where L L^T = H^{-1}
        chol: Array2<f64>,
        /// L^T for gradient chain rule: ∇z = L^T @ ∇_β
        chol_t: Array2<f64>,
    }

    impl SurvivalPosterior {
        /// Creates a new survival posterior target.
        pub fn new(
            age_entry: ArrayView1<'_, f64>,
            age_exit: ArrayView1<'_, f64>,
            event_target: ArrayView1<'_, u8>,
            event_competing: ArrayView1<'_, u8>,
            sampleweight: ArrayView1<'_, f64>,
            x_entry: ArrayView2<'_, f64>,
            x_exit: ArrayView2<'_, f64>,
            x_derivative: ArrayView2<'_, f64>,
            offset_eta_entry: Option<ArrayView1<'_, f64>>,
            offset_eta_exit: Option<ArrayView1<'_, f64>>,
            offset_derivative_exit: Option<ArrayView1<'_, f64>>,
            penalties: PenaltyBlocks,
            monotonicity: MonotonicityPenalty,
            spec: SurvivalSpec,
            structurally_monotonic: bool,
            structural_time_columns: usize,
            mode: ArrayView1<f64>,
            hessian: ArrayView2<f64>,
        ) -> Result<Self, String> {
            let n = age_entry.len();
            let off_eta_entry = offset_eta_entry
                .map(|v| v.to_owned())
                .unwrap_or_else(|| Array1::zeros(n));
            let off_eta_exit = offset_eta_exit
                .map(|v| v.to_owned())
                .unwrap_or_else(|| Array1::zeros(n));
            let off_deriv_exit = offset_derivative_exit
                .map(|v| v.to_owned())
                .unwrap_or_else(|| Array1::zeros(n));

            let mut base_model = WorkingModelSurvival::from_engine_inputswith_offsets(
                SurvivalEngineInputs {
                    age_entry,
                    age_exit,
                    event_target,
                    event_competing,
                    sampleweight,
                    x_entry,
                    x_exit,
                    x_derivative,
                    monotonicity_constraint_rows: None,
                    monotonicity_constraint_offsets: None,
                },
                Some(crate::survival::SurvivalBaselineOffsets {
                    eta_entry: off_eta_entry.view(),
                    eta_exit: off_eta_exit.view(),
                    derivative_exit: off_deriv_exit.view(),
                }),
                penalties,
                monotonicity,
                spec,
            )
            .map_err(|e| format!("Survival state construction failed: {:?}", e))?;
            if structurally_monotonic {
                base_model
                    .set_structural_monotonicity(true, structural_time_columns)
                    .map_err(|e| {
                        format!("Failed to enable structural monotonicity in survival HMC: {e}")
                    })?;
            }

            let sampler_mode = mode.to_owned();
            let dim = sampler_mode.len();

            // Compute whitening transform via Cholesky of Hessian
            let chol_factor = hessian
                .to_owned()
                .cholesky(Side::Lower)
                .map_err(|e| format!("Hessian Cholesky decomposition failed: {:?}", e))?;
            let l_h = chol_factor.lower_triangular();
            let chol = solve_upper_triangular_transpose(&l_h, dim);
            let chol_t = chol.t().to_owned();

            let data = SharedSurvivalData {
                base_model: Arc::new(base_model),
                mode: Arc::new(sampler_mode),
            };

            Ok(Self { data, chol, chol_t })
        }

        fn compute_logp_and_grad_into(
            &self,
            z: &Array1<f64>,
            grad: &mut Array1<f64>,
        ) -> Result<f64, String> {
            let sampler_position = self.data.mode.as_ref() + &self.chol.dot(z);
            let state = self
                .data
                .base_model
                .update_state(&sampler_position)
                .map_err(|e| format!("Survival state update failed: {:?}", e))?;
            let logp = state.log_likelihood - state.penalty_term;
            let grad_beta = state.gradient.mapv(|g| -g);
            fast_av_into(&self.chol_t, &grad_beta, grad);
            Ok(logp)
        }

        /// Get the Cholesky factor L for un-whitening samples
        pub fn chol(&self) -> &Array2<f64> {
            &self.chol
        }

        /// Get the mode
        pub fn mode(&self) -> &Array1<f64> {
            &self.data.mode
        }
    }

    impl HamiltonianTarget<Array1<f64>> for SurvivalPosterior {
        fn logp_and_grad(&self, position: &Array1<f64>, grad: &mut Array1<f64>) -> f64 {
            match self.compute_logp_and_grad_into(position, grad) {
                Ok(logp) => logp,
                Err(e) => {
                    log::warn!("Survival posterior evaluation failed: {}", e);
                    grad.fill(0.0);
                    f64::NEG_INFINITY
                }
            }
        }
    }

    /// Runs NUTS sampling for survival models with whitened parameter space.
    pub(crate) fn run_survival_nuts_sampling(
        age_entry: ArrayView1<'_, f64>,
        age_exit: ArrayView1<'_, f64>,
        event_target: ArrayView1<'_, u8>,
        event_competing: ArrayView1<'_, u8>,
        sampleweight: ArrayView1<'_, f64>,
        x_entry: ArrayView2<'_, f64>,
        x_exit: ArrayView2<'_, f64>,
        x_derivative: ArrayView2<'_, f64>,
        eta_offset_entry: Option<ArrayView1<'_, f64>>,
        eta_offset_exit: Option<ArrayView1<'_, f64>>,
        derivative_offset_exit: Option<ArrayView1<'_, f64>>,
        penalties: PenaltyBlocks,
        monotonicity: MonotonicityPenalty,
        spec: SurvivalSpec,
        structurally_monotonic: bool,
        structural_time_columns: usize,
        mode: ArrayView1<f64>,
        hessian: ArrayView2<f64>,
        config: &NutsConfig,
    ) -> Result<NutsResult, String> {
        validate_nuts_target_accept(config.target_accept).map_err(String::from)?;
        // Create posterior target
        let target = SurvivalPosterior::new(
            age_entry,
            age_exit,
            event_target,
            event_competing,
            sampleweight,
            x_entry,
            x_exit,
            x_derivative,
            eta_offset_entry,
            eta_offset_exit,
            derivative_offset_exit,
            penalties,
            monotonicity,
            spec,
            structurally_monotonic,
            structural_time_columns,
            mode,
            hessian,
        )?;

        // Get Cholesky factor for un-whitening samples later
        let chol = target.chol().clone();
        let mode_arr = target.mode().clone();
        let dim = mode_arr.len();

        let initial_positions = jittered_initial_positions(config, dim, 0.1, 0xEC2D_7A9B_4051_F638);

        // Create GenericNUTS sampler
        let mass_cfg = robust_survival_mass_matrix_config(dim, config.nwarmup);
        let mut sampler = GenericNUTS::new_with_mass_matrix(
            target,
            initial_positions,
            robust_target_accept(config.target_accept, dim),
            mass_cfg,
        )
        .set_seed(nuts_transition_seed(config.seed, 0x731B_60D4_AE52_9C8F));

        // Run sampling with progress bar
        let (samples_array, run_stats) = sampler
            .run_progress(config.n_samples, config.nwarmup)
            .map_err(|e| format!("NUTS sampling failed: {}", e))?;

        log::info!("Survival NUTS sampling complete: {}", run_stats);

        // Convert samples from whitened space back to original space
        let shape = samples_array.shape();
        let n_chains = shape[0];
        let n_samples_out = shape[1];
        let total_samples = n_chains * n_samples_out;

        let mut samples = Array2::<f64>::zeros((total_samples, dim));
        let mut z_buffer = Array1::<f64>::zeros(dim);
        for chain in 0..n_chains {
            for sample_i in 0..n_samples_out {
                let zview = samples_array.slice(ndarray::s![chain, sample_i, ..]);
                z_buffer.assign(&zview);

                let beta = &mode_arr + &chol.dot(&z_buffer);

                let sample_idx = chain * n_samples_out + sample_i;
                samples.row_mut(sample_idx).assign(&beta);
            }
        }

        // Compute statistics
        let posterior_mean = samples
            .mean_axis(Axis(0))
            .unwrap_or_else(|| Array1::zeros(dim));
        let posterior_std = samples.std_axis(Axis(0), 0.0);
        let (rhat, ess) = compute_split_rhat_and_ess(&samples_array);
        let converged = rhat < 1.1;

        Ok(NutsResult {
            samples,
            posterior_mean,
            posterior_std,
            rhat,
            ess,
            converged,
        })
    }
}

/// Engine-facing flattened survival NUTS entrypoint.
pub fn run_survival_nuts_sampling_flattened<'a>(
    flat: SurvivalFlatInputs<'a>,
    penalties: crate::survival::PenaltyBlocks,
    monotonicity: crate::survival::MonotonicityPenalty,
    spec: crate::survival::SurvivalSpec,
    structurally_monotonic: bool,
    structural_time_columns: usize,
    mode: ArrayView1<'a, f64>,
    hessian: ArrayView2<'a, f64>,
    config: &NutsConfig,
) -> Result<NutsResult, String> {
    run_nuts_sampling_flattened_family(
        LikelihoodSpec {
            response: ResponseFamily::RoystonParmar,
            link: InverseLink::Standard(LinkFunction::Identity),
        },
        FamilyNutsInputs::Survival(Box::new(SurvivalNutsInputs {
            flat,
            penalties,
            monotonicity,
            spec,
            structurally_monotonic,
            structural_time_columns,
            mode,
            hessian,
        })),
        config,
    )
}