use crate::basis::{
    BSplineBasisSpec, BSplineBoundaryConditions, BSplineEndpointBoundaryCondition,
    BSplineIdentifiability, BSplineKnotSpec, BasisBuildResult, BasisError, BasisMetadata,
    BasisOptions, BasisPsiDerivativeResult, BasisPsiSecondDerivativeResult, BasisWorkspace,
    CenterStrategy, CenterStrategyKind, Dense, DuchonBasisSpec, KnotSource, KroneckerFactoredBasis,
    MaternBasisSpec, MaternIdentifiability, OneDimensionalBoundary, PenaltyCandidate, PenaltyInfo,
    PenaltySource, SpatialIdentifiability, SphericalSplineBasisSpec,
    SphericalSplineIdentifiability, ThinPlateBasisSpec, apply_sum_to_zero_constraint,
    build_bspline_basis_1d, build_duchon_basis, build_duchon_basiswithworkspace,
    build_matern_basis, build_matern_basis_log_kappa_aniso_derivatives,
    build_matern_basis_log_kappa_derivatives, build_matern_basiswithworkspace,
    build_matern_collocation_operator_matrices, build_spherical_spline_basis,
    build_thin_plate_basis, build_thin_plate_basis_log_kappa_derivatives, center_strategy_is_auto,
    center_strategy_kind, center_strategy_num_centers, center_strategy_with_num_centers,
    estimate_penalty_nullity, filter_active_penalty_candidates,
    filter_active_penalty_candidates_with_ops, initial_aniso_contrasts,
    orthogonality_transform_for_design, pairwise_distance_bounds, pairwise_distance_bounds_sampled,
    points_in_aniso_y_space, select_centers_by_strategy,
};
use crate::construction::{
    kronecker_logdet_and_derivatives, kronecker_marginal_eigensystems, kronecker_product,
};
use crate::custom_family::{
    BlockEffectiveJacobian, BlockGeometryDirectionalDerivative, BlockWorkingSet,
    BlockwiseFitOptions, CustomFamily, CustomFamilyBlockPsiDerivative, CustomFamilyWarmStart,
    ExactNewtonJointPsiTerms, ExactNewtonOuterObjective, FamilyEvaluation,
    FamilyLinearizationState, ParameterBlockSpec, ParameterBlockState, PenaltyMatrix,
    evaluate_custom_family_joint_hyper, evaluate_custom_family_joint_hyper_efs, fit_custom_family,
};
use crate::estimate::{
    EstimationError, ExternalOptimOptions, FitInference, FitOptions, FittedLinkState, PenaltySpec,
    UnifiedFitResult, UnifiedFitResultParts, fit_gamwith_heuristic_lambdas,
    reml::DirectionalHyperParam,
};
use crate::faer_ndarray::{fast_ab, fast_atb, fast_atv};
use crate::families::strategy::{FamilyStrategy, strategy_for_spec};
use crate::matrix::{
    BlockDesignOperator, CoefficientTransformOperator, DenseDesignOperator, DesignBlock,
    DesignMatrix, LinearOperator, RandomEffectOperator, SymmetricMatrix,
    TensorProductDesignOperator,
};
use crate::mixture_link::{
    logit_inverse_link_jet5, state_from_beta_logisticspec, state_from_sasspec, state_fromspec,
};
use crate::pirls::LinearInequalityConstraints;
use crate::resource::MatrixMaterializationError;
use crate::types::{
    InverseLink, LatentCLogLogState, LikelihoodSpec, MixtureLinkState, ResponseFamily,
    SasLinkState, StandardLink,
};
use crate::util::quantile::quantile_from_sorted;
use faer::sparse::{SparseColMat, Triplet};
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayViewMut2, Axis, s};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};
use std::f64;
use std::fs::File;
use std::ops::Range;
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::{Arc, Mutex};

// ---------------------------------------------------------------------------
// Typed errors
// ---------------------------------------------------------------------------

/// Typed errors emitted by smooth-term helpers in this module. `Display`
/// reproduces the exact pre-refactor `format!(...)` text byte-for-byte, so
/// string-matching callers (tests, log assertions) keep working unchanged.
/// Public boundaries that interface with `term_builder.rs` / `construction`
/// continue to return `Result<_, String>`; typed `Err(...)` values flow
/// through `From<SmoothError> for String` at those boundaries via `.into()`.
#[derive(Clone, Debug)]
pub enum SmoothError {
    /// Spec-level configuration error: unfrozen knots/centers, invalid
    /// numeric bounds on coefficient priors, length-scale optimizer options
    /// that violate `min > 0 < max` / `min < max` invariants, etc.
    InvalidConfig { reason: String },
    /// Shape / length disagreement between design columns, penalty blocks,
    /// coefficient ranges, directional-derivative vectors, theta length, and
    /// other bookkeeping invariants that are checked at runtime.
    DimensionMismatch { reason: String },
    /// Out-of-range index into adaptive hyperparameter components, psi
    /// derivative blocks, or non-zero block indices for single-block
    /// custom-family impls that only support `block_idx == 0`.
    InvalidIndex { reason: String },
}

impl std::fmt::Display for SmoothError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SmoothError::InvalidConfig { reason }
            | SmoothError::DimensionMismatch { reason }
            | SmoothError::InvalidIndex { reason } => f.write_str(reason),
        }
    }
}

impl std::error::Error for SmoothError {}

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

impl From<crate::util::block_count::BlockCountMismatch> for SmoothError {
    fn from(err: crate::util::block_count::BlockCountMismatch) -> SmoothError {
        SmoothError::invalid_index(err.message())
    }
}

impl SmoothError {
    #[inline]
    fn invalid_config(reason: impl Into<String>) -> Self {
        SmoothError::InvalidConfig {
            reason: reason.into(),
        }
    }
    #[inline]
    fn dimension_mismatch(reason: impl Into<String>) -> Self {
        SmoothError::DimensionMismatch {
            reason: reason.into(),
        }
    }
    #[inline]
    fn invalid_index(reason: impl Into<String>) -> Self {
        SmoothError::InvalidIndex {
            reason: reason.into(),
        }
    }
}

fn describe_thin_plate_center_request(strategy: &CenterStrategy) -> String {
    match strategy {
        CenterStrategy::Auto(inner) => describe_thin_plate_center_request(inner),
        CenterStrategy::UserProvided(centers) => format!("{} centers", centers.nrows()),
        CenterStrategy::EqualMass { num_centers }
        | CenterStrategy::EqualMassCovarRepresentative { num_centers }
        | CenterStrategy::FarthestPoint { num_centers }
        | CenterStrategy::KMeans { num_centers, .. } => format!("{num_centers} centers"),
        CenterStrategy::UniformGrid { points_per_dim } => {
            format!("uniform grid with {points_per_dim} points per dimension")
        }
    }
}

fn rewrite_thin_plate_knots_error(
    err: BasisError,
    termname: &str,
    feature_count: usize,
    spec: &ThinPlateBasisSpec,
) -> BasisError {
    match err {
        // Polynomial-nullspace shortfall reported directly by the kernel
        // builder ("thin-plate spline requires at least N centers to span ...").
        BasisError::InvalidInput(msg)
            if msg.contains("thin-plate spline requires at least")
                && (msg.contains("centers to span") || msg.contains("knots to span")) =>
        {
            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
            let requested = describe_thin_plate_center_request(&spec.center_strategy);
            BasisError::InvalidInput(format!(
                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
            ))
        }
        // Insufficient-rows shortfall raised by `select_thin_plate_knots` when
        // the requested center count exceeds the available row count. Rewrite
        // it in term language so the diagnostic points at the smooth term and
        // the polynomial-nullspace minimum the user needs to satisfy.
        BasisError::InvalidInput(msg)
            if msg.starts_with("requested ") && msg.contains(" knots but only ") =>
        {
            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
            let requested = describe_thin_plate_center_request(&spec.center_strategy);
            BasisError::InvalidInput(format!(
                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
            ))
        }
        other => other,
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ShapeConstraint {
    None,
    MonotoneIncreasing,
    MonotoneDecreasing,
    Convex,
    Concave,
}

/// Parse a shape-constraint string into a [`ShapeConstraint`].
///
/// This is the single source of truth shared by the formula DSL
/// (`s(x, shape=...)`) and the `smooths={...}` override path
/// (`Smooth.shape_constraint`). The accepted spellings cover the canonical
/// Python `ShapeConstraintLiteral` strings exactly
/// (`"none"` / `"monotone_increasing"` / `"monotone_decreasing"` /
/// `"convex"` / `"concave"`) plus a few common aliases. Hyphens and case are
/// normalized, so `"Monotone-Increasing"` and `"mono_inc"` both resolve to
/// [`ShapeConstraint::MonotoneIncreasing`].
pub fn parse_shape_constraint(raw: &str) -> Result<ShapeConstraint, String> {
    let normalized = raw.trim().to_ascii_lowercase().replace('-', "_");
    match normalized.as_str() {
        "" | "none" => Ok(ShapeConstraint::None),
        "monotone_increasing" | "monotonic_increasing" | "increasing" | "mono_inc" | "mpi" => {
            Ok(ShapeConstraint::MonotoneIncreasing)
        }
        "monotone_decreasing" | "monotonic_decreasing" | "decreasing" | "mono_dec" | "mpd" => {
            Ok(ShapeConstraint::MonotoneDecreasing)
        }
        "convex" | "cvx" => Ok(ShapeConstraint::Convex),
        "concave" | "ccv" => Ok(ShapeConstraint::Concave),
        other => Err(format!(
            "unknown shape constraint {other:?}; expected one of \
             \"none\", \"monotone_increasing\", \"monotone_decreasing\", \
             \"convex\", \"concave\""
        )),
    }
}

impl ShapeConstraint {
    /// Canonical formula-DSL spelling, i.e. the text emitted into
    /// `s(x, shape=...)`. Round-trips through [`parse_shape_constraint`].
    pub fn dsl_str(&self) -> &'static str {
        match self {
            ShapeConstraint::None => "none",
            ShapeConstraint::MonotoneIncreasing => "monotone_increasing",
            ShapeConstraint::MonotoneDecreasing => "monotone_decreasing",
            ShapeConstraint::Convex => "convex",
            ShapeConstraint::Concave => "concave",
        }
    }
}

/// Smooth-term head keywords recognised by the formula DSL. A `shape=` option
/// may be attached to any term whose head is one of these.
const SMOOTH_HEAD_KEYWORDS: [&str; 11] = [
    "s",
    "smooth",
    "te",
    "tensor",
    "thinplate",
    "tps",
    "duchon",
    "matern",
    "sphere",
    "bs",
    "bspline",
];

/// Rewrite smooth-term calls in `formula` so each named smooth carries a
/// `shape=<kind>` option understood by the formula DSL.
///
/// `constraints` pairs the smooth-term text as it appears in the formula
/// (e.g. `"s(x)"` or `"s(x, type=duchon, centers=8)"`) with a shape-constraint
/// spelling accepted by [`parse_shape_constraint`]; comparison is exact after
/// whitespace removal. A `"none"` constraint is a no-op. Referencing a term not
/// present in the formula is an error.
///
/// This is the single source of truth for the `gamfit.fit(..., constraints=…)`
/// rewrite — the Python wrapper only marshals the mapping across the FFI and
/// holds no formula-parsing or alias-normalization logic of its own.
pub fn apply_shape_constraints_to_formula(
    formula: &str,
    constraints: &[(String, String)],
) -> Result<String, String> {
    use std::collections::{BTreeMap, BTreeSet};

    if constraints.is_empty() {
        return Ok(formula.to_string());
    }
    let strip_ws = |s: &str| -> String { s.chars().filter(|c| !c.is_whitespace()).collect() };

    // Whitespace-stripped term text -> canonical shape spelling.
    let mut wanted: BTreeMap<String, &'static str> = BTreeMap::new();
    // Whitespace-stripped term text -> original key (for error labels).
    let mut originals: BTreeMap<String, String> = BTreeMap::new();
    for (key, kind_raw) in constraints {
        let kind = parse_shape_constraint(kind_raw)?;
        let nk = strip_ws(key);
        originals.entry(nk.clone()).or_insert_with(|| key.clone());
        if kind != ShapeConstraint::None {
            wanted.insert(nk, kind.dsl_str());
        }
    }
    if wanted.is_empty() {
        return Ok(formula.to_string());
    }

    let chars: Vec<char> = formula.chars().collect();
    let n = chars.len();
    let is_ident = |c: char| c.is_ascii_alphanumeric() || c == '_';

    let mut out = String::with_capacity(formula.len() + 32);
    let mut matched: BTreeSet<String> = BTreeSet::new();
    let mut i = 0usize;
    while i < n {
        // Locate the next smooth-term head (`<keyword> \s* (`) at or after `i`,
        // respecting word boundaries so `abs(` never matches the `s(` head.
        let mut head: Option<(usize, usize)> = None; // (head_start, paren_index)
        let mut p = i;
        while p < n {
            let boundary = p == 0 || !is_ident(chars[p - 1]);
            if boundary {
                for kw in SMOOTH_HEAD_KEYWORDS.iter() {
                    let klen = kw.chars().count();
                    if p + klen > n || chars[p..p + klen].iter().collect::<String>() != **kw {
                        continue;
                    }
                    let mut q = p + klen;
                    while q < n && chars[q].is_whitespace() {
                        q += 1;
                    }
                    if q < n && chars[q] == '(' {
                        head = Some((p, q));
                        break;
                    }
                }
            }
            if head.is_some() {
                break;
            }
            p += 1;
        }
        let (head_start, paren_open) = match head {
            Some(h) => h,
            None => {
                out.extend(chars[i..].iter());
                break;
            }
        };
        out.extend(chars[i..head_start].iter());

        // Find the matching close paren, honoring nesting and string literals.
        let body_start = paren_open + 1;
        let mut depth = 1i32;
        let mut j = body_start;
        let mut in_str: Option<char> = None;
        let mut closed = false;
        while j < n {
            let ch = chars[j];
            if let Some(quote) = in_str {
                if ch == quote {
                    in_str = None;
                }
            } else if ch == '\'' || ch == '"' {
                in_str = Some(ch);
            } else if ch == '(' {
                depth += 1;
            } else if ch == ')' {
                depth -= 1;
                if depth == 0 {
                    closed = true;
                    break;
                }
            }
            j += 1;
        }
        if !closed {
            // Unbalanced — emit the remainder verbatim; the DSL parser will
            // produce the canonical error.
            out.extend(chars[head_start..].iter());
            break;
        }

        let term_text: String = chars[head_start..=j].iter().collect();
        let key_norm = strip_ws(&term_text);
        match wanted.get(&key_norm) {
            None => out.extend(chars[head_start..=j].iter()),
            Some(kind) => {
                let head_paren: String = chars[head_start..body_start].iter().collect();
                let inside: String = chars[body_start..j].iter().collect();
                let inside = inside.trim();
                if inside.is_empty() {
                    out.push_str(&format!("{head_paren}shape={kind})"));
                } else {
                    out.push_str(&format!("{head_paren}{inside}, shape={kind})"));
                }
                matched.insert(key_norm);
            }
        }
        i = j + 1;
    }

    let mut missing: Vec<String> = wanted
        .keys()
        .filter(|k| !matched.contains(*k))
        .map(|k| originals.get(k).cloned().unwrap_or_else(|| k.clone()))
        .collect();
    if !missing.is_empty() {
        missing.sort();
        return Err(format!(
            "shape constraints referenced smooth term(s) not found in formula: {}",
            missing.join(", ")
        ));
    }
    Ok(out)
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BySmoothKind {
    Numeric,
    Level { level_bits: u64 },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SmoothBasisSpec {
    /// Row-gated wrapper used for mgcv-style ``by=`` smooths.
    ///
    /// ``ByNumeric`` multiplies the inner smooth by a numeric column.
    /// ``ByLevel`` keeps the inner smooth active only for rows whose encoded
    /// categorical value has the stored bit pattern.  Unordered factor-by
    /// smooths are represented as one independent ``ByLevel`` term per level.
    ///
    /// `kind` preserves the compact structural discriminator, while `by`
    /// carries the full row-gating spec used to build the local design.
    ByVariable {
        inner: Box<SmoothBasisSpec>,
        by_col: usize,
        kind: BySmoothKind,
        by: ByVariableSpec,
    },
    /// Sum-to-zero factor smooth (`bs="sz"`): with L levels, estimate L-1
    /// deviation coefficient blocks and use the final level as the negative
    /// sum of the others, enforcing coefficient-wise zero sums across levels.
    FactorSumToZero {
        inner: Box<SmoothBasisSpec>,
        by_col: usize,
        levels: Vec<u64>,
    },
    BSpline1D {
        feature_col: usize,
        spec: BSplineBasisSpec,
    },
    /// A smooth modulated by a `by=` variable. Numeric `by` scales one inner
    /// smooth; factor `by` replicates the inner smooth by level.
    BySmooth {
        smooth: Box<SmoothBasisSpec>,
        by_kind: ByVarKind,
    },
    /// Factor-smooth interaction families (`bs="fs"`, `bs="sz"`) and
    /// random slopes (`bs="re"`).
    FactorSmooth { spec: FactorSmoothSpec },
    ThinPlate {
        feature_cols: Vec<usize>,
        spec: ThinPlateBasisSpec,
        /// Per-column standard deviations used to standardize input dimensions
        /// before kernel evaluation when d > 1. `None` means no standardization
        /// (either d == 1 or explicitly disabled).
        #[serde(default)]
        input_scales: Option<Vec<f64>>,
    },
    Sphere {
        feature_cols: Vec<usize>,
        spec: SphericalSplineBasisSpec,
    },
    Matern {
        feature_cols: Vec<usize>,
        spec: MaternBasisSpec,
        #[serde(default)]
        input_scales: Option<Vec<f64>>,
    },
    Duchon {
        feature_cols: Vec<usize>,
        spec: DuchonBasisSpec,
        #[serde(default)]
        input_scales: Option<Vec<f64>>,
    },
    Pca {
        feature_cols: Vec<usize>,
        basis_matrix: Array2<f64>,
        centered: bool,
        #[serde(default = "default_pca_smooth_penalty")]
        smooth_penalty: f64,
        #[serde(default)]
        center_mean: Option<Array1<f64>>,
        #[serde(default)]
        pca_basis_path: Option<PathBuf>,
        #[serde(default = "default_pca_chunk_size")]
        chunk_size: usize,
    },
    /// Tensor-product smooth built from 1D B-spline marginals.
    ///
    /// This is the `te()`-style construction used when axes have different units/scales
    /// (for example, space x time) and isotropic radial kernels are not appropriate.
    TensorBSpline {
        feature_cols: Vec<usize>,
        spec: TensorBSplineSpec,
    },
}

impl SmoothBasisSpec {
    /// Conservative lower bound on the number of sample rows needed for this
    /// smooth basis to have a well-posed REML fit.
    ///
    /// Each basis kind answers the question for itself, so the workflow does
    /// not have to know how many columns a B-spline, tensor product, PCA
    /// projection, or spatial kernel emits. The contract is a *lower bound*:
    /// returning too small a number is permitted (the inner solver will catch
    /// any genuine n-vs-rank failure that slips past); returning too large a
    /// number is a regression because it rejects legitimate fits.
    ///
    /// Rationale: B-spline / tensor / PCA bases have a closed-form column
    /// count, so we use the exact dimension. Radial bases (TPS, Matern,
    /// Duchon, Sphere) and factor smooths choose their column count from the
    /// data (`heuristic_centers`, `unique_count`); we fall back to a small
    /// constant floor because a fit on fewer than five rows cannot stabilise
    /// any radial smooth regardless of the configured kernel scale.
    pub fn min_sample_rows(&self) -> usize {
        // Floor used for data-driven bases whose column count is not known
        // from the spec alone. Five rows is the minimum at which the inner
        // pivot/QR + REML smoothing-parameter search has any chance of being
        // well-posed for a non-parametric smooth.
        const RADIAL_FLOOR: usize = 5;

        match self {
            Self::ByVariable { inner, .. } => inner.min_sample_rows(),
            Self::FactorSumToZero { inner, levels, .. } => {
                // L-1 independent deviation blocks each carrying the inner
                // basis dimension. Skip the levels-multiplier if it doesn't
                // bring more rows; we want the *lower bound* not the rank.
                let inner_min = inner.min_sample_rows();
                let lvls = levels.len().saturating_sub(1).max(1);
                inner_min.saturating_mul(lvls)
            }
            Self::BSpline1D { spec, .. } => bspline_basis_min_rows(spec),
            Self::BySmooth { smooth, .. } => smooth.min_sample_rows(),
            Self::FactorSmooth { spec } => {
                // Replicates the marginal once per level; without a known
                // level count we conservatively require at least the marginal
                // basis dimension.
                bspline_basis_min_rows(&spec.marginal)
            }
            Self::ThinPlate { .. }
            | Self::Sphere { .. }
            | Self::Matern { .. }
            | Self::Duchon { .. } => RADIAL_FLOOR,
            Self::Pca { basis_matrix, .. } => basis_matrix.ncols().max(1),
            Self::TensorBSpline { spec, .. } => {
                // A `te(...)` smooth is *penalized*: each margin carries a
                // difference (wiggliness) penalty and the tensor inherits a
                // Kronecker-sum penalty `S = Σ_i I ⊗ … ⊗ S_i ⊗ … ⊗ I`. The raw
                // column count is the *product* of the per-marginal column
                // counts, but that product is the lower bound for an
                // *unpenalized* tensor regression — it is the number of rows you
                // would need to identify every interaction column with no
                // regularization. The penalty regularizes all of those
                // interaction directions; only the combined penalty *null space*
                // (the tensor product of the per-margin polynomial trends, a
                // handful of columns) must be identified by the data, and the
                // smoothing-parameter search shrinks the rest. The effective
                // degrees of freedom of the fitted `te()` are therefore a small
                // fraction of the column product, which is exactly why mgcv
                // fits a default `te(x, y)` on a couple hundred rows.
                //
                // The honest *penalized* lower bound is the **sum** of the
                // per-marginal column counts, not their product: a row floor of
                // `Σ_i k_i` still guarantees enough data to identify each
                // margin's additive main-effect (the largest sub-block the
                // penalty cannot shrink to zero), while no longer conflating
                // unpenalized column-count identifiability with penalized
                // well-posedness. This accepts moderate-`n` penalized tensors
                // (e.g. a 20×20 default basis on n=200) yet still rejects a
                // genuinely undersized fit where `n < Σ_i k_i` and even the
                // additive part is rank-deficient.
                //
                // Binary / low-cardinality margins (#724): gam will accept a
                // `te(x, badh)` whose `badh ∈ {0, 1}` margin nominally requests
                // more basis columns than `badh` has unique values, where mgcv
                // refuses the unpenalized term as ill-posed ("badh has
                // insufficient unique values to support k knots"). This is
                // correct-by-design, *not* a degenerate fit: the marginal
                // wiggliness penalty on the `badh` axis has a null space that is
                // exactly its identifiable trend (the two cell means of a binary
                // covariate), and the Kronecker-sum penalty shrinks every tensor
                // column outside that null space toward zero. The resulting fit
                // is the well-posed "per-level `x` smooth + binary main effect"
                // that mgcv reaches only after manually collapsing the basis —
                // gam reaches it automatically because the penalty, not the raw
                // column count, sets the effective rank. A genuinely
                // rank-deficient design (penalty null space wider than the data
                // can support) is still caught downstream by the inner pivoted
                // factorization, which owns the exact n-vs-rank decision; this
                // pre-fit gate only refuses the grossly-undersized formula.
                let mut total: usize = 0;
                for marginal in &spec.marginalspecs {
                    let m = bspline_basis_min_rows(marginal);
                    total = total.saturating_add(m.max(1));
                }
                total.max(RADIAL_FLOOR)
            }
        }
    }
}

/// Lower bound on the number of basis columns produced by a 1D B-spline
/// configuration. Used as the per-smooth row floor in
/// [`SmoothBasisSpec::min_sample_rows`].
fn bspline_basis_min_rows(spec: &crate::terms::basis::BSplineBasisSpec) -> usize {
    use crate::terms::basis::BSplineKnotSpec;
    let from_knots = match &spec.knotspec {
        BSplineKnotSpec::Generate {
            num_internal_knots, ..
        } => *num_internal_knots + spec.degree + 1,
        BSplineKnotSpec::Automatic {
            num_internal_knots: Some(k),
            ..
        } => *k + spec.degree + 1,
        BSplineKnotSpec::Automatic {
            num_internal_knots: None,
            ..
        } => {
            // Knot count is data-derived (`default_internal_knot_count_for_data`).
            // A minimal cubic basis is `degree + 2` columns; below that the
            // basis cannot represent a non-parametric smooth.
            spec.degree + 2
        }
        BSplineKnotSpec::Provided(knots) => knots.len().saturating_sub(spec.degree + 1).max(1),
        BSplineKnotSpec::PeriodicUniform { num_basis, .. } => *num_basis,
    };
    from_knots.max(spec.degree + 2)
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ByVariableSpec {
    Numeric,
    Level { value_bits: u64, label: String },
}

/// Tensor-product B-spline smooth specification.
///
/// `marginalspecs[i]` is the 1D B-spline setup for `feature_cols[i]`.
/// The final penalty set is one Kronecker penalty per margin:
/// `S_i = I ⊗ ... ⊗ S_marginal_i ⊗ ... ⊗ I`, plus optional global ridge.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TensorMarginalSpec {
    BSpline(BSplineBasisSpec),
    Categorical {
        feature_col_offset: usize,
        drop_first_level: bool,
        center_for_identifiability: bool,
        frozen_levels: Option<Vec<u64>>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ByVarKind {
    Numeric {
        feature_col: usize,
    },
    Factor {
        feature_col: usize,
        ordered: bool,
        frozen_levels: Option<Vec<u64>>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FactorSmoothSpec {
    pub continuous_cols: Vec<usize>,
    pub group_col: usize,
    pub marginal: BSplineBasisSpec,
    pub flavour: FactorSmoothFlavour,
    pub group_frozen_levels: Option<Vec<u64>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FactorSmoothFlavour {
    Fs { m_null_penalty_orders: Vec<usize> },
    Sz,
    Re,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TensorBSplineSpec {
    pub marginalspecs: Vec<BSplineBasisSpec>,
    #[serde(default)]
    pub periods: Vec<Option<f64>>,
    pub double_penalty: bool,
    #[serde(default)]
    pub identifiability: TensorBSplineIdentifiability,
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub enum TensorBSplineIdentifiability {
    None,
    #[default]
    SumToZero,
    /// mgcv `ti(...)` semantics: a *tensor interaction* smooth that excludes the
    /// marginal main effects. A sum-to-zero constraint is applied to **each
    /// marginal basis independently** before forming the tensor product, so the
    /// resulting column space contains no function of a single variable alone —
    /// only the pure interaction survives. The realized identifiability
    /// transform is the Kronecker product `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}` of the
    /// per-margin sum-to-zero null-space bases, which is exactly the
    /// reparameterization that turns the full-tensor design into the tensor
    /// product of the centered margins.
    MarginalSumToZero,
    FrozenTransform {
        transform: Array2<f64>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmoothTermSpec {
    pub name: String,
    pub basis: SmoothBasisSpec,
    pub shape: ShapeConstraint,
    /// Joint-null absorption rotation captured at fit time. `Some(Q)` means
    /// the fitted coefficient vector lives in `γ`-coordinates with
    /// `β_raw = Q · γ`; prediction must rotate the raw-basis design via
    /// `X_new = X_new_raw · Q` to match. `None` means either the smooth had
    /// no joint null space (penalty already full-rank) or rotation was
    /// suppressed (smooth carries shape constraints whose cone geometry
    /// would not survive an arbitrary orthogonal rotation). Persisted so
    /// `save → load → predict` is bit-equivalent to in-memory prediction.
    #[serde(default)]
    pub joint_null_rotation: Option<crate::terms::basis::JointNullRotation>,
}

#[derive(Debug, Clone)]
pub struct SmoothTerm {
    pub name: String,
    pub coeff_range: Range<usize>,
    pub shape: ShapeConstraint,
    pub penalties_local: Vec<Array2<f64>>,
    pub nullspace_dims: Vec<usize>,
    pub penaltyinfo_local: Vec<PenaltyInfo>,
    pub metadata: BasisMetadata,
    /// Optional term-local lower bounds for constrained coefficients.
    /// `-inf` means unconstrained.
    pub lower_bounds_local: Option<Array1<f64>>,
    /// Optional term-local inequality constraints in local coefficient coordinates.
    /// `A_local * beta_local >= b_local`.
    pub linear_constraints_local: Option<LinearInequalityConstraints>,
    /// Optional factored tensor-product representation preserved for operator-backed
    /// assembly in the main design builder.
    pub kronecker_factored: Option<KroneckerFactoredBasis>,
    /// Joint-null absorption rotation. `Some(Q)` records the orthonormal
    /// `(p_local × p_local)` matrix that was applied to this term's design
    /// and per-block penalties at construction time:
    /// `term_design ← X_raw · Q`, `penalties_local[k] ← Qᵀ · S_raw · Q`.
    /// The smooth's coefficient vector therefore lives in the rotated
    /// (`γ`) coordinate system, with `β_raw = Q · γ` recovering the raw
    /// pre-rotation parameterization. `None` means either no joint null
    /// space (penalty already full-rank) or rotation was suppressed —
    /// suppression fires when the smooth carries shape constraints
    /// (lower bounds or local linear inequalities) that would lose their
    /// cone geometry under a general orthogonal rotation.
    ///
    /// Prediction-side replay: callers building a new-data design `X_new_raw`
    /// from the *raw* basis must call [`SmoothTerm::apply_rotation_to_predict`]
    /// (or equivalent) to obtain `X_new = X_new_raw · Q` matching this
    /// term's coefficient system.
    ///
    /// Persistence replay: `freeze_term_collection_from_design` copies this
    /// rotation into `SmoothTermSpec`, which is serialized with fitted-model
    /// payloads and reused by the predict-time basis builder. Saved models
    /// therefore replay the same `X_new_raw · Q` transform as in-memory
    /// prediction.
    pub joint_null_rotation: Option<crate::terms::basis::JointNullRotation>,
}

impl SmoothTerm {
    /// Apply the joint-null absorption rotation to a raw new-data design
    /// matrix, returning `X_new_raw · Q` when this term was rotated at
    /// fit time, or `X_new_raw` unchanged when no rotation was applied.
    ///
    /// Callers in the prediction path: after building the smooth's basis
    /// at new data via the *raw* basis builder (the same builder used at
    /// fit time, applied to `x_new` instead of the training rows), call
    /// this method on the resulting matrix before forming `X · β`. The
    /// fitted `β` lives in `γ`-coordinates if Q was applied; multiplying
    /// the un-rotated `X_new_raw` by `β` would give a wrong η.
    ///
    /// Returns an error if the raw design's column count does not match
    /// the rotation's `p_local`. The width invariant must hold: the raw
    /// basis builder MUST emit the same `p_local` columns that the
    /// fit-time builder did, and the rotation is `(p_local × p_local)`.
    pub fn apply_rotation_to_predict(
        &self,
        x_new_raw: Array2<f64>,
    ) -> Result<Array2<f64>, BasisError> {
        let Some(rot) = self.joint_null_rotation.as_ref() else {
            return Ok(x_new_raw);
        };
        let p_local = rot.rotation.nrows();
        if x_new_raw.ncols() != p_local {
            crate::bail_dim_basis!(
                "joint-null rotation replay for term '{}': raw design has {} columns, \
                 rotation expects {} (the raw basis builder must emit the same column \
                 count as at fit time)",
                self.name,
                x_new_raw.ncols(),
                p_local,
            );
        }
        Ok(crate::linalg::faer_ndarray::fast_ab(
            &x_new_raw,
            &rot.rotation,
        ))
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PenaltyBlockInfo {
    pub global_index: usize,
    pub termname: Option<String>,
    pub penalty: PenaltyInfo,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DroppedPenaltyBlockInfo {
    pub termname: Option<String>,
    pub penalty: PenaltyInfo,
}

#[derive(Debug, Clone)]
pub struct SmoothDesign {
    pub term_designs: Vec<DesignMatrix>,
    /// Per-term block-local penalties.  Each `col_range` is relative to the
    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
    pub penalties: Vec<BlockwisePenalty>,
    pub nullspace_dims: Vec<usize>,
    pub penaltyinfo: Vec<PenaltyBlockInfo>,
    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
    pub terms: Vec<SmoothTerm>,
    /// Optional smooth-block lower bounds in smooth coefficient coordinates.
    /// Length equals `total_smooth_cols()` when present.
    pub coefficient_lower_bounds: Option<Array1<f64>>,
    /// Optional smooth-block inequality constraints:
    /// `A_smooth * beta_smooth >= b`.
    pub linear_constraints: Option<LinearInequalityConstraints>,
}

impl SmoothDesign {
    pub fn total_smooth_cols(&self) -> usize {
        self.term_designs.iter().map(DesignMatrix::ncols).sum()
    }
    pub fn nrows(&self) -> usize {
        self.term_designs.first().map_or(0, DesignMatrix::nrows)
    }
}

#[derive(Debug, Clone)]
pub struct RawSmoothDesign {
    pub term_designs: Vec<DesignMatrix>,
    /// Per-term block-local penalties.  Each `col_range` is relative to the
    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
    pub penalties: Vec<BlockwisePenalty>,
    pub nullspace_dims: Vec<usize>,
    pub penaltyinfo: Vec<PenaltyBlockInfo>,
    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
    pub terms: Vec<SmoothTerm>,
    pub coefficient_lower_bounds: Option<Array1<f64>>,
    pub linear_constraints: Option<LinearInequalityConstraints>,
}

impl RawSmoothDesign {
    pub fn total_smooth_cols(&self) -> usize {
        self.term_designs.iter().map(DesignMatrix::ncols).sum()
    }
    pub fn nrows(&self) -> usize {
        self.term_designs.first().map_or(0, DesignMatrix::nrows)
    }
}

impl From<RawSmoothDesign> for SmoothDesign {
    fn from(value: RawSmoothDesign) -> Self {
        Self {
            term_designs: value.term_designs,
            penalties: value.penalties,
            nullspace_dims: value.nullspace_dims,
            penaltyinfo: value.penaltyinfo,
            dropped_penaltyinfo: value.dropped_penaltyinfo,
            terms: value.terms,
            coefficient_lower_bounds: value.coefficient_lower_bounds,
            linear_constraints: value.linear_constraints,
        }
    }
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub enum BoundedCoefficientPriorSpec {
    #[default]
    None,
    Uniform,
    Beta {
        a: f64,
        b: f64,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum LinearCoefficientGeometry {
    #[default]
    Unconstrained,
    Bounded {
        min: f64,
        max: f64,
        #[serde(default)]
        prior: BoundedCoefficientPriorSpec,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LinearTermSpec {
    pub name: String,
    /// Primary feature column index. For Wilkinson-Rogers `:` interaction
    /// terms (`a:b[:c...]`) this is the first column in `feature_cols`; the
    /// realized design column is the elementwise product across every entry
    /// of `feature_cols`. Plain (non-interaction) linear terms set
    /// `feature_cols == vec![feature_col]`.
    pub feature_col: usize,
    /// Full list of columns whose elementwise product yields this term's
    /// design column. `len() >= 1`; `len() == 1` is a plain linear effect.
    #[serde(default)]
    pub feature_cols: Vec<usize>,
    /// Optional ridge (`S = I`, REML-selected `λ`) on this linear coefficient.
    /// A parametric linear term carries no wiggliness, so it is **unpenalized by
    /// default** — gam reports the MLE, matching mgcv/glm/survreg/VGAM (which
    /// penalize parametric terms only under an explicit `paraPen`). Set `true`
    /// to opt into an explicit shrinkage ridge (a zero-mean Gaussian prior
    /// `β ~ N(0, λ⁻¹)`); doing so adds one outer REML smoothing coordinate.
    #[serde(default = "default_linear_term_double_penalty")]
    pub double_penalty: bool,
    #[serde(default)]
    pub coefficient_geometry: LinearCoefficientGeometry,
    #[serde(default)]
    pub coefficient_min: Option<f64>,
    #[serde(default)]
    pub coefficient_max: Option<f64>,
}

impl LinearTermSpec {
    /// Return the effective list of feature columns. Backfills from
    /// `feature_col` for legacy specs that predate the multi-column field.
    pub fn effective_feature_cols(&self) -> Vec<usize> {
        if self.feature_cols.is_empty() {
            vec![self.feature_col]
        } else {
            self.feature_cols.clone()
        }
    }

    /// True when this term is a Wilkinson-Rogers `:` interaction (multi-col).
    pub fn is_interaction(&self) -> bool {
        self.feature_cols.len() > 1
    }
}

const fn default_linear_term_double_penalty() -> bool {
    // Parametric/linear terms are unpenalized by default — a single linear
    // coefficient has no roughness for a smoothing penalty to control, so the
    // historical `S = I`, REML-selected `λ` shrank every linear coefficient off
    // the MLE and injected a spurious outer smoothing coordinate (#749). Mature
    // tools (mgcv/glm/survreg/VGAM) leave parametric terms unpenalized; gam now
    // matches that and reports the MLE. An explicit `double_penalty = true`
    // still opts a term into a ridge.
    false
}

const fn default_pca_smooth_penalty() -> f64 {
    1.0
}

const fn default_pca_chunk_size() -> usize {
    4096
}

/// Random-effects term specification.
///
/// The selected feature column is interpreted as a categorical grouping variable.
/// The term contributes a one-hot dummy block with an identity penalty on group
/// coefficients, equivalent to i.i.d. Gaussian random effects.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RandomEffectTermSpec {
    pub name: String,
    pub feature_col: usize,
    /// If true, drop the lexicographically first group level to use treatment coding.
    /// If false, keep all levels (full one-hot block, still identifiable under ridge).
    pub drop_first_level: bool,
    /// If true, add a ridge penalty and estimate this block as a random effect.
    /// If false, leave the one-hot/treatment-coded block unpenalized so it is a
    /// fixed categorical main effect.  The default preserves older saved models.
    #[serde(default = "default_random_effect_penalized")]
    pub penalized: bool,
    /// Optional fixed kept-level set (sorted by f64 bit pattern) captured at fit time.
    /// When present, prediction uses exactly these columns to avoid design drift.
    #[serde(default)]
    pub frozen_levels: Option<Vec<u64>>,
}

fn default_random_effect_penalized() -> bool {
    true
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TermCollectionSpec {
    pub linear_terms: Vec<LinearTermSpec>,
    pub random_effect_terms: Vec<RandomEffectTermSpec>,
    pub smooth_terms: Vec<SmoothTermSpec>,
}

fn validate_smooth_basis_frozen(
    basis: &SmoothBasisSpec,
    label: &str,
    term_name: &str,
) -> Result<(), String> {
    match basis {
        SmoothBasisSpec::ByVariable { inner, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
            validate_smooth_basis_frozen(inner, label, term_name)
        }
        SmoothBasisSpec::BSpline1D { spec, .. } => {
            if !matches!(
                spec.knotspec,
                BSplineKnotSpec::Provided(_) | BSplineKnotSpec::PeriodicUniform { .. }
            ) {
                return Err(format!(
                    "{label} term '{term_name}' is not frozen: BSpline knotspec must be Provided or PeriodicUniform"
                ));
            }
            Ok(())
        }
        SmoothBasisSpec::ThinPlate { spec, .. } => {
            if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
                return Err(format!(
                    "{label} term '{term_name}' is not frozen: ThinPlate centers must be UserProvided"
                ));
            }
            if matches!(
                spec.identifiability,
                SpatialIdentifiability::OrthogonalToParametric
            ) {
                return Err(format!(
                    "{label} term '{term_name}' is not frozen: ThinPlate identifiability must be FrozenTransform or None"
                ));
            }
            Ok(())
        }
        _ => Ok(()),
    }
}

impl TermCollectionSpec {
    /// Validate that a term collection spec represents a fully frozen model
    /// (i.e. all knots/centers are pre-computed, identifiability transforms are
    /// baked in, and random-effect levels are fixed).
    pub fn validate_frozen(&self, label: &str) -> Result<(), String> {
        for linear in &self.linear_terms {
            if let (Some(min), Some(max)) = (linear.coefficient_min, linear.coefficient_max)
                && (!min.is_finite() || !max.is_finite() || min > max)
            {
                return Err(SmoothError::invalid_config(format!(
                    "{label} linear term '{}' has invalid coefficient constraint [{min}, {max}]",
                    linear.name
                ))
                .into());
            }
            if let Some(min) = linear.coefficient_min
                && !min.is_finite()
            {
                return Err(SmoothError::invalid_config(format!(
                    "{label} linear term '{}' has non-finite coefficient minimum {min}",
                    linear.name
                ))
                .into());
            }
            if let Some(max) = linear.coefficient_max
                && !max.is_finite()
            {
                return Err(SmoothError::invalid_config(format!(
                    "{label} linear term '{}' has non-finite coefficient maximum {max}",
                    linear.name
                ))
                .into());
            }
            if let LinearCoefficientGeometry::Bounded { min, max, prior } =
                &linear.coefficient_geometry
            {
                if !min.is_finite() || !max.is_finite() || min >= max {
                    return Err(SmoothError::invalid_config(format!(
                        "{label} bounded term '{}' has invalid bounds [{min}, {max}]",
                        linear.name
                    ))
                    .into());
                }
                match prior {
                    BoundedCoefficientPriorSpec::None | BoundedCoefficientPriorSpec::Uniform => {}
                    BoundedCoefficientPriorSpec::Beta { a, b } => {
                        if !a.is_finite() || !b.is_finite() || *a < 1.0 || *b < 1.0 {
                            return Err(SmoothError::invalid_config(format!(
                                "{label} bounded term '{}' has invalid Beta prior ({a}, {b})",
                                linear.name
                            ))
                            .into());
                        }
                    }
                }
            }
        }
        for st in &self.smooth_terms {
            match &st.basis {
                SmoothBasisSpec::ByVariable { inner, .. } => {
                    validate_smooth_basis_frozen(inner, label, &st.name)?;
                    let nested = SmoothTermSpec {
                        name: st.name.clone(),
                        basis: (**inner).clone(),
                        shape: st.shape,
                        joint_null_rotation: None,
                    };
                    TermCollectionSpec {
                        linear_terms: Vec::new(),
                        random_effect_terms: Vec::new(),
                        smooth_terms: vec![nested],
                    }
                    .validate_frozen(label)?;
                }
                SmoothBasisSpec::FactorSumToZero { inner, levels, .. } => {
                    if levels.len() < 2 {
                        return Err(format!(
                            "{label} term '{}' has invalid frozen sz levels",
                            st.name
                        ));
                    }
                    validate_smooth_basis_frozen(inner, label, &st.name)?;
                }
                SmoothBasisSpec::BSpline1D { spec, .. } => {
                    if !matches!(
                        spec.knotspec,
                        BSplineKnotSpec::Provided(_) | BSplineKnotSpec::PeriodicUniform { .. }
                    ) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: BSpline knotspec must be Provided or PeriodicUniform",
                            st.name
                        ))
                        .into());
                    }
                }
                SmoothBasisSpec::ThinPlate { spec, .. } => {
                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: ThinPlate centers must be UserProvided",
                            st.name
                        ))
                        .into());
                    }
                    if matches!(
                        spec.identifiability,
                        SpatialIdentifiability::OrthogonalToParametric
                    ) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: ThinPlate identifiability must be FrozenTransform or None",
                            st.name
                        ))
                        .into());
                    }
                }
                SmoothBasisSpec::Sphere { spec, .. } => {
                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: Sphere centers must be UserProvided",
                            st.name
                        ))
                        .into());
                    }
                    if matches!(spec.method, crate::basis::SphereMethod::Harmonic)
                        && spec.max_degree.is_none_or(|d| d == 0)
                    {
                        return Err(format!(
                            "{label} term '{}' is not frozen: sphere max_degree must be positive",
                            st.name
                        ));
                    }
                }
                SmoothBasisSpec::Matern { spec, .. } => {
                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: Matern centers must be UserProvided",
                            st.name
                        ))
                        .into());
                    }
                }
                SmoothBasisSpec::Duchon { spec, .. } => {
                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: Duchon centers must be UserProvided",
                            st.name
                        ))
                        .into());
                    }
                    if matches!(
                        spec.identifiability,
                        SpatialIdentifiability::OrthogonalToParametric
                    ) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: Duchon identifiability must be FrozenTransform or None",
                            st.name
                        ))
                        .into());
                    }
                }
                SmoothBasisSpec::Pca {
                    centered,
                    center_mean,
                    pca_basis_path,
                    ..
                } => {
                    if *centered && center_mean.is_none() && pca_basis_path.is_none() {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: centered Pca missing center_mean",
                            st.name
                        ))
                        .into());
                    }
                }
                SmoothBasisSpec::BySmooth { smooth, by_kind } => {
                    if let SmoothBasisSpec::BySmooth { .. } = smooth.as_ref() {
                        return Err(format!("{label} term '{}' has nested by-smooths", st.name));
                    }
                    match by_kind {
                        ByVarKind::Numeric { .. } => {}
                        ByVarKind::Factor { frozen_levels, .. } if frozen_levels.is_none() => {
                            return Err(format!(
                                "{label} term '{}' is not frozen: by-factor levels missing",
                                st.name
                            ));
                        }
                        ByVarKind::Factor { .. } => {}
                    }
                    let nested = TermCollectionSpec {
                        linear_terms: vec![],
                        random_effect_terms: vec![],
                        smooth_terms: vec![SmoothTermSpec {
                            name: st.name.clone(),
                            basis: (**smooth).clone(),
                            shape: st.shape,
                            joint_null_rotation: None,
                        }],
                    };
                    nested.validate_frozen(label)?;
                }
                SmoothBasisSpec::FactorSmooth { spec } => {
                    if spec.group_frozen_levels.is_none() {
                        return Err(format!(
                            "{label} term '{}' is not frozen: factor-smooth levels missing",
                            st.name
                        ));
                    }
                    if !matches!(
                        spec.marginal.knotspec,
                        BSplineKnotSpec::Provided(_) | BSplineKnotSpec::PeriodicUniform { .. }
                    ) {
                        return Err(format!(
                            "{label} term '{}' is not frozen: factor-smooth marginal knots missing",
                            st.name
                        ));
                    }
                }
                SmoothBasisSpec::TensorBSpline { spec, .. } => {
                    for (dim, marginal) in spec.marginalspecs.iter().enumerate() {
                        if !matches!(
                            marginal.knotspec,
                            BSplineKnotSpec::Provided(_) | BSplineKnotSpec::PeriodicUniform { .. }
                        ) {
                            return Err(SmoothError::invalid_config(format!(
                                "{label} term '{}' dim {} is not frozen: tensor marginal knotspec must be Provided or PeriodicUniform",
                                st.name, dim
                            ))
                            .into());
                        }
                    }
                    if matches!(
                        spec.identifiability,
                        TensorBSplineIdentifiability::SumToZero
                            | TensorBSplineIdentifiability::MarginalSumToZero
                    ) {
                        return Err(SmoothError::invalid_config(format!(
                            "{label} term '{}' is not frozen: tensor identifiability must be FrozenTransform or None",
                            st.name
                        ))
                        .into());
                    }
                }
            }
        }

        for rt in &self.random_effect_terms {
            if rt.frozen_levels.is_none() {
                return Err(SmoothError::invalid_config(format!(
                    "{label} random-effect term '{}' is not frozen: missing frozen_levels",
                    rt.name
                ))
                .into());
            }
        }

        Ok(())
    }

    /// Re-resolve every stored feature-column index through `remap`, returning a
    /// spec that addresses a different column layout.
    ///
    /// A frozen `TermCollectionSpec` stores feature columns as *absolute indices
    /// into the training table*. To replay it on a fresh dataset whose columns
    /// sit at different positions — the common case at prediction time, where
    /// the response column is unknown and may be absent entirely — every index
    /// must be re-resolved against the new layout. `remap` receives each
    /// training-table index and returns its position in the runtime table;
    /// callers typically implement it as "look the name up in the training
    /// headers, then resolve that name against the prediction dataset".
    ///
    /// This is the single authority on *which* fields carry a column index
    /// across every basis variant (linear, random-effect, the `by=` column of
    /// `ByVariable`/`FactorSumToZero`/`BySmooth`, the continuous and group
    /// columns of a `FactorSmooth`, and the multi-axis `feature_cols` of every
    /// spatial/tensor basis), so a predict-time realignment cannot silently miss
    /// one and dereference a stale training index.
    pub fn remap_feature_columns<E, F>(&self, mut remap: F) -> Result<TermCollectionSpec, E>
    where
        F: FnMut(usize) -> Result<usize, E>,
    {
        let mut out = self.clone();
        for lt in &mut out.linear_terms {
            lt.feature_col = remap(lt.feature_col)?;
        }
        for rt in &mut out.random_effect_terms {
            rt.feature_col = remap(rt.feature_col)?;
        }
        for st in &mut out.smooth_terms {
            remap_smooth_basis_feature_columns(&mut st.basis, &mut remap)?;
        }
        Ok(out)
    }
}

/// Walk a `SmoothBasisSpec` tree, re-resolving every column index through
/// `remap`. Shared by all predict-time column realignment (see
/// [`TermCollectionSpec::remap_feature_columns`]); kept exhaustive so a newly
/// added index-bearing variant fails to compile until it is handled here.
fn remap_smooth_basis_feature_columns<E, F>(
    basis: &mut SmoothBasisSpec,
    remap: &mut F,
) -> Result<(), E>
where
    F: FnMut(usize) -> Result<usize, E>,
{
    match basis {
        SmoothBasisSpec::ByVariable { inner, by_col, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
            *by_col = remap(*by_col)?;
            remap_smooth_basis_feature_columns(inner, remap)?;
        }
        SmoothBasisSpec::BSpline1D { feature_col, .. } => {
            *feature_col = remap(*feature_col)?;
        }
        SmoothBasisSpec::BySmooth { smooth, by_kind } => {
            let by_feature_col = match by_kind {
                ByVarKind::Numeric { feature_col } | ByVarKind::Factor { feature_col, .. } => {
                    feature_col
                }
            };
            *by_feature_col = remap(*by_feature_col)?;
            remap_smooth_basis_feature_columns(smooth, remap)?;
        }
        SmoothBasisSpec::FactorSmooth { spec } => {
            for fc in spec.continuous_cols.iter_mut() {
                *fc = remap(*fc)?;
            }
            spec.group_col = remap(spec.group_col)?;
        }
        SmoothBasisSpec::ThinPlate { feature_cols, .. }
        | SmoothBasisSpec::Sphere { feature_cols, .. }
        | SmoothBasisSpec::Matern { feature_cols, .. }
        | SmoothBasisSpec::Duchon { feature_cols, .. }
        | SmoothBasisSpec::Pca { feature_cols, .. }
        | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
            for fc in feature_cols.iter_mut() {
                *fc = remap(*fc)?;
            }
        }
    }
    Ok(())
}

#[derive(Debug, Clone)]
pub enum PenaltyStructureHint {
    Ridge(f64),
    Kronecker(Vec<Array2<f64>>),
}

/// A penalty matrix stored at its natural block size together with the
/// column range it occupies in the global coefficient vector.
///
/// Instead of embedding every penalty into a full `p_total × p_total` dense
/// matrix filled with zeros, we keep the compact local matrix and reconstruct
/// the global view only when a downstream consumer explicitly requires it.
#[derive(Clone)]
pub struct BlockwisePenalty {
    /// Column range in the global coefficient vector that this penalty covers.
    pub col_range: Range<usize>,
    /// The local penalty matrix — dimensions `block_p × block_p` where
    /// `block_p = col_range.len()`.
    pub local: Array2<f64>,
    /// Optional nonzero centering vector for this coefficient block.
    pub prior_mean: crate::solver::estimate::CoefficientPriorMean,
    /// Optional structural hint so downstream spectral/logdet code can stay
    /// block-local or factorized without reverse-engineering the matrix.
    pub structure_hint: Option<PenaltyStructureHint>,
    /// Optional operator-form handle bit-equivalent to `local`. Populated when
    /// the originating closed-form factory emitted an op-form penalty so exact
    /// operator algebra can use matvec instead of materializing the dense
    /// `block_p × block_p` Gram. `None` for ordinary dense penalties.
    pub op: Option<std::sync::Arc<dyn crate::terms::penalty_op::PenaltyOp>>,
}

impl std::fmt::Debug for BlockwisePenalty {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BlockwisePenalty")
            .field("col_range", &self.col_range)
            .field(
                "local",
                &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
            )
            .field("prior_mean", &self.prior_mean)
            .field("structure_hint", &self.structure_hint)
            .field("op", &self.op.as_ref().map(|o| o.dim()))
            .finish()
    }
}

impl BlockwisePenalty {
    /// Create a new blockwise penalty.
    pub fn new(col_range: Range<usize>, local: Array2<f64>) -> Self {
        assert_eq!(col_range.len(), local.nrows());
        assert_eq!(col_range.len(), local.ncols());
        Self {
            col_range,
            local,
            prior_mean: crate::solver::estimate::CoefficientPriorMean::Zero,
            structure_hint: None,
            op: None,
        }
    }

    pub fn with_prior_mean(
        mut self,
        prior_mean: crate::solver::estimate::CoefficientPriorMean,
    ) -> Self {
        self.prior_mean = prior_mean;
        self
    }

    /// Attach an op-form penalty handle bit-equivalent to `local`.
    pub fn with_op(
        mut self,
        op: Option<std::sync::Arc<dyn crate::terms::penalty_op::PenaltyOp>>,
    ) -> Self {
        self.op = op;
        self
    }

    pub fn ridge(col_range: Range<usize>, scale: f64) -> Self {
        let block_size = col_range.len();
        let mut local = Array2::<f64>::zeros((block_size, block_size));
        for i in 0..block_size {
            local[[i, i]] = scale;
        }
        Self {
            col_range,
            local,
            prior_mean: crate::solver::estimate::CoefficientPriorMean::Zero,
            structure_hint: Some(PenaltyStructureHint::Ridge(scale)),
            op: None,
        }
    }

    pub fn kronecker(
        col_range: Range<usize>,
        local: Array2<f64>,
        factors: Vec<Array2<f64>>,
    ) -> Self {
        assert_eq!(col_range.len(), local.nrows());
        assert_eq!(col_range.len(), local.ncols());
        Self {
            col_range,
            local,
            prior_mean: crate::solver::estimate::CoefficientPriorMean::Zero,
            structure_hint: Some(PenaltyStructureHint::Kronecker(factors)),
            op: None,
        }
    }

    /// Expand this blockwise penalty into a full `p_total × p_total` dense
    /// matrix (mostly zeros). Use sparingly — the whole point of blockwise
    /// storage is to avoid this allocation.
    pub fn to_global(&self, p_total: usize) -> Array2<f64> {
        let mut g = Array2::<f64>::zeros((p_total, p_total));
        let r = &self.col_range;
        assert!(
            r.end <= p_total && self.local.nrows() == r.len() && self.local.ncols() == r.len(),
            "BlockwisePenalty::to_global shape invariant violated: \
             col_range={}..{}, local={}x{}, p_total={}",
            r.start,
            r.end,
            self.local.nrows(),
            self.local.ncols(),
            p_total,
        );
        g.slice_mut(s![r.start..r.end, r.start..r.end])
            .assign(&self.local);
        g
    }

    /// The block size of this penalty.
    #[inline]
    pub fn block_size(&self) -> usize {
        self.col_range.len()
    }
}

/// Compute `Σ_k λ_k S_k` directly from blockwise penalties, accumulating
/// into a pre-allocated `p_total × p_total` output without ever materializing
/// individual global matrices.
pub fn weighted_blockwise_penalty_sum(
    penalties: &[BlockwisePenalty],
    lambdas: &[f64],
    p_total: usize,
) -> Array2<f64> {
    assert_eq!(penalties.len(), lambdas.len());
    // Smoothing parameters λ_k must be non-negative and finite. A negative
    // λ would flip the sign of the corresponding block S_k, turning the
    // total penalty matrix indefinite and silently corrupting every
    // downstream Cholesky / PIRLS / REML / pseudo-logdet computation that
    // assumes S_λ ⪰ 0. Catch this at the boundary rather than after it
    // has propagated.
    for (idx, &lam) in lambdas.iter().enumerate() {
        assert!(
            lam.is_finite() && lam >= 0.0,
            "weighted_blockwise_penalty_sum: lambdas[{idx}] = {lam} is invalid (must be finite and non-negative; negative smoothing parameters violate S_λ ⪰ 0)",
        );
    }
    // Block column ranges must also fit inside the declared total parameter
    // dimension; an out-of-bounds slice would otherwise panic from ndarray
    // with a far less informative message.
    for (idx, bp) in penalties.iter().enumerate() {
        let r = &bp.col_range;
        assert!(
            r.end <= p_total,
            "weighted_blockwise_penalty_sum: penalties[{idx}] col_range {:?} exceeds p_total = {p_total}",
            r,
        );
    }
    let mut out = Array2::<f64>::zeros((p_total, p_total));
    for (bp, &lam) in penalties.iter().zip(lambdas.iter()) {
        let r = &bp.col_range;
        let mut slice = out.slice_mut(s![r.start..r.end, r.start..r.end]);
        slice.scaled_add(lam, &bp.local);
    }
    out
}

// ---------------------------------------------------------------------------
// KroneckerPenaltySystem — factored tensor-product penalty representation
// ---------------------------------------------------------------------------

/// Factored representation of tensor-product penalties with precomputed
/// marginal eigensystems for O(∏q_j) logdet and penalty operations.
#[derive(Debug, Clone)]
pub struct KroneckerPenaltySystem {
    /// Marginal penalty matrices: `marginal_penalties[k]` is `(q_k, q_k)`.
    pub marginal_penalties: Vec<Array2<f64>>,
    /// Precomputed eigensystems: `(eigenvalues, eigenvectors)` per marginal.
    pub marginal_eigensystems: Vec<(Array1<f64>, Array2<f64>)>,
    /// Marginal basis dimensions.
    pub marginal_dims: Vec<usize>,
    /// Whether a global ridge (double) penalty is present.
    pub has_double_penalty: bool,
}

impl KroneckerPenaltySystem {
    pub fn new(
        marginal_penalties: Vec<Array2<f64>>,
        marginal_dims: Vec<usize>,
        has_double_penalty: bool,
    ) -> Result<Self, BasisError> {
        if marginal_penalties.len() != marginal_dims.len() {
            crate::bail_dim_basis!(
                "KroneckerPenaltySystem: {} penalties vs {} dims",
                marginal_penalties.len(),
                marginal_dims.len()
            );
        }
        let eigensystems =
            kronecker_marginal_eigensystems(&marginal_penalties, "KroneckerPenaltySystem")
                .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
        Ok(Self {
            marginal_penalties,
            marginal_eigensystems: eigensystems,
            marginal_dims,
            has_double_penalty,
        })
    }

    pub fn p_total(&self) -> usize {
        self.marginal_dims.iter().copied().product()
    }

    pub fn ndim(&self) -> usize {
        self.marginal_dims.len()
    }

    pub fn num_penalties(&self) -> usize {
        self.marginal_dims.len() + if self.has_double_penalty { 1 } else { 0 }
    }

    /// Compute `log|S|₊` and its first/second derivatives w.r.t. `ρ_k = log(λ_k)`.
    ///
    /// Iterates over the ∏q_j multi-index grid. Cost: O(d · ∏q_j), no O(p²) storage.
    pub fn logdet_and_derivatives(
        &self,
        lambdas: &[f64],
        ridge: f64,
    ) -> (f64, Array1<f64>, Array2<f64>) {
        let n_pen = self.num_penalties();
        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
        let marginal_evals: Vec<_> = self
            .marginal_eigensystems
            .iter()
            .map(|(evals, _)| evals.view())
            .collect();
        kronecker_logdet_and_derivatives(
            &marginal_evals,
            &self.marginal_dims,
            lambdas,
            self.has_double_penalty,
            ridge,
        )
    }

    pub fn logdet_rank_and_derivatives(
        &self,
        lambdas: &[f64],
        ridge: f64,
    ) -> (f64, usize, Array1<f64>, Array2<f64>) {
        let n_pen = self.num_penalties();
        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
        let d = self.marginal_dims.len();
        let mut logdet = 0.0;
        let mut rank = 0usize;
        let mut grad = Array1::<f64>::zeros(n_pen);
        let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
        let tol = 1e-12;
        let structural_zero_tol = 1e-12;
        let mut multi_idx = vec![0usize; d];
        loop {
            let mut sigma = 0.0;
            let mut structural_sigma = 0.0;
            for k in 0..d {
                let marginal_eigenvalue = self.marginal_eigensystems[k].0[multi_idx[k]];
                structural_sigma += marginal_eigenvalue;
                sigma += lambdas[k] * marginal_eigenvalue;
            }
            if self.has_double_penalty {
                structural_sigma += 1.0;
                sigma += lambdas[d];
            }
            if structural_sigma > structural_zero_tol {
                sigma += ridge;
            }

            if sigma > tol {
                rank += 1;
                logdet += sigma.ln();
                let inv_sigma = 1.0 / sigma;
                let inv_sigma2 = inv_sigma * inv_sigma;
                for k in 0..n_pen {
                    let ck = if k < d {
                        lambdas[k] * self.marginal_eigensystems[k].0[multi_idx[k]]
                    } else {
                        lambdas[d]
                    };
                    grad[k] += ck * inv_sigma;
                    hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
                    for l in (k + 1)..n_pen {
                        let cl = if l < d {
                            lambdas[l] * self.marginal_eigensystems[l].0[multi_idx[l]]
                        } else {
                            lambdas[d]
                        };
                        let off = -ck * cl * inv_sigma2;
                        hess[[k, l]] += off;
                        hess[[l, k]] += off;
                    }
                }
            }

            let mut carry = true;
            for dim in (0..d).rev() {
                if carry {
                    multi_idx[dim] += 1;
                    if multi_idx[dim] < self.marginal_dims[dim] {
                        carry = false;
                    } else {
                        multi_idx[dim] = 0;
                    }
                }
            }
            if carry {
                break;
            }
        }
        (logdet, rank, grad, hess)
    }
}

#[derive(Clone, Debug)]
pub struct TermCollectionDesign {
    /// The full design matrix.
    ///
    /// Prefer a true sparse matrix when every block is sparse-compatible.
    /// If the collection already contains intrinsically sparse blocks, preserve
    /// that storage and let PIRLS decide later whether the penalized system is
    /// sparse-native eligible. Purely dense materialized blocks still fall back
    /// to the lazy block operator when sparse storage would just re-encode a
    /// dense matrix.
    pub design: DesignMatrix,
    pub penalties: Vec<BlockwisePenalty>,
    pub nullspace_dims: Vec<usize>,
    pub penaltyinfo: Vec<PenaltyBlockInfo>,
    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
    /// Optional global coefficient lower bounds for constrained fitting.
    /// Length equals `design.ncols()` when present. Unconstrained entries are `-inf`.
    pub coefficient_lower_bounds: Option<Array1<f64>>,
    /// Optional global inequality constraints:
    /// `A * beta >= b`.
    pub linear_constraints: Option<LinearInequalityConstraints>,
    pub intercept_range: Range<usize>,
    pub linear_ranges: Vec<(String, Range<usize>)>,
    pub random_effect_ranges: Vec<(String, Range<usize>)>,
    pub random_effect_levels: Vec<(String, Vec<u64>)>,
    pub smooth: SmoothDesign,
}

impl TermCollectionDesign {
    /// Convert blockwise penalties to `PenaltyMatrix::Blockwise` without
    /// expanding to `p_total × p_total`. This is the preferred path for
    /// family modules that accept `Vec<PenaltyMatrix>`.
    pub fn penalties_as_penalty_matrix(&self) -> Vec<crate::custom_family::PenaltyMatrix> {
        let p = self.design.ncols();
        self.penalties
            .iter()
            .map(|bp| crate::custom_family::PenaltyMatrix::from_blockwise(bp.clone(), p))
            .collect()
    }

    /// Number of penalty blocks.
    #[inline]
    pub fn num_penalties(&self) -> usize {
        self.penalties.len()
    }

    /// Resolve coefficient groups against this design's global coefficient
    /// layout and append their penalties after the existing term penalties.
    pub fn realize_coefficient_groups(
        &self,
        groups: &[CoefficientGroupSpec],
        base_prior: &crate::types::RhoPrior,
    ) -> Result<RealizedCoefficientGroups, BasisError> {
        realize_coefficient_groups(self, groups, base_prior)
    }

    /// Extract a `KroneckerPenaltySystem` if exactly one smooth term has
    /// Kronecker structure and it accounts for all penalties.
    ///
    /// Returns `None` if no Kronecker structure is present or if the model
    /// has multiple smooth terms with mixed structure.
    pub fn kronecker_penalty_system(&self) -> Option<KroneckerPenaltySystem> {
        let kron_terms: Vec<&KroneckerFactoredBasis> = self
            .smooth
            .terms
            .iter()
            .filter_map(|t| t.kronecker_factored.as_ref())
            .collect();
        if kron_terms.len() != 1 {
            return None; // 0 or multiple Kronecker terms — not supported yet
        }
        let kron = kron_terms[0];
        // Only use the Kronecker path when the model is purely this tensor term
        // (no other smooth terms with separate penalties).
        let has_non_kron_smooth_terms = self
            .smooth
            .terms
            .iter()
            .any(|t| t.kronecker_factored.is_none());
        if has_non_kron_smooth_terms {
            return None; // mixed model — fall back to standard path
        }
        KroneckerPenaltySystem::new(
            kron.marginal_penalties.clone(),
            kron.marginal_dims.clone(),
            kron.has_double_penalty,
        )
        .ok()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CoefficientSelector {
    /// Explicit global coefficient indices in the realized design matrix.
    GlobalColumns(Vec<usize>),
    /// A half-open global coefficient range.
    GlobalRange(Range<usize>),
    LinearTerm(String),
    RandomEffectTerm(String),
    SmoothTerm(String),
    /// Selected basis columns within one smooth term.
    SmoothTermColumns {
        term: String,
        columns: Vec<usize>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CoefficientGroupPrior {
    Flat,
    NormalLogPrecision {
        mean: f64,
        sd: f64,
    },
    GammaPrecision {
        shape: f64,
        rate: f64,
    },
    /// Penalized-complexity prior calibrated by `P(exp(-ρ/2) > upper) =
    /// tail_prob`; see [`crate::types::RhoPrior::PenalizedComplexity`].
    PenalizedComplexity {
        upper: f64,
        tail_prob: f64,
    },
}

impl CoefficientGroupPrior {
    fn to_rho_prior(&self) -> crate::types::RhoPrior {
        match *self {
            Self::Flat => crate::types::RhoPrior::Flat,
            Self::NormalLogPrecision { mean, sd } => crate::types::RhoPrior::Normal { mean, sd },
            Self::GammaPrecision { shape, rate } => {
                crate::types::RhoPrior::GammaPrecision { shape, rate }
            }
            Self::PenalizedComplexity { upper, tail_prob } => {
                crate::types::RhoPrior::PenalizedComplexity { upper, tail_prob }
            }
        }
    }

    fn validate(&self, context: &str) -> Result<(), BasisError> {
        match *self {
            Self::Flat => Ok(()),
            Self::NormalLogPrecision { mean, sd } => {
                if !mean.is_finite() {
                    crate::bail_invalid_basis!(
                        "{context} Normal log-precision prior requires finite mean, got {mean}"
                    );
                }
                if !sd.is_finite() || sd <= 0.0 {
                    crate::bail_invalid_basis!(
                        "{context} Normal log-precision prior requires sd > 0, got {sd}"
                    );
                }
                Ok(())
            }
            Self::GammaPrecision { shape, rate } => {
                validate_gamma_precision_prior(context, shape, rate)
            }
            Self::PenalizedComplexity { upper, tail_prob } => {
                validate_penalized_complexity_prior(context, upper, tail_prob)
            }
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CoefficientGroupSpec {
    pub name: String,
    pub selectors: Vec<CoefficientSelector>,
    pub parent: Option<String>,
    pub prior: Option<CoefficientGroupPrior>,
    #[serde(skip, default)]
    pub prior_mean: crate::solver::estimate::CoefficientPriorMean,
}

#[derive(Debug, Clone)]
pub struct RealizedCoefficientGroups {
    pub penalty_specs: Vec<PenaltySpec>,
    pub nullspace_dims: Vec<usize>,
    pub rho_prior: crate::types::RhoPrior,
    pub group_column_indices: Vec<(String, Vec<usize>)>,
}

#[derive(Debug, Clone)]
pub struct PenaltyBlockGammaPriorMetadata<'a> {
    pub label: String,
    pub global_index: usize,
    pub termname: Option<&'a str>,
    pub source: String,
    pub effective_rank: usize,
    pub nullspace_dim_hint: usize,
}

fn penalty_block_label_candidates(info: &PenaltyBlockInfo) -> Vec<String> {
    let mut labels = Vec::<String>::new();
    labels.push(format!("penalty:{}", info.global_index));
    labels.push(info.global_index.to_string());
    if let Some(termname) = info.termname.as_ref() {
        labels.push(termname.clone());
        labels.push(format!("{termname}:{}", info.penalty.original_index));
    }
    if let PenaltySource::Other(label) = &info.penalty.source {
        labels.push(label.clone());
    }
    labels.push(format!("{:?}", info.penalty.source));
    labels.sort();
    labels.dedup();
    labels
}

fn penalty_block_metadata(info: &PenaltyBlockInfo) -> PenaltyBlockGammaPriorMetadata<'_> {
    PenaltyBlockGammaPriorMetadata {
        label: info
            .termname
            .clone()
            .unwrap_or_else(|| format!("penalty:{}", info.global_index)),
        global_index: info.global_index,
        termname: info.termname.as_deref(),
        source: format!("{:?}", info.penalty.source),
        effective_rank: info.penalty.effective_rank,
        nullspace_dim_hint: info.penalty.nullspace_dim_hint,
    }
}

fn validate_gamma_precision_prior(label: &str, shape: f64, rate: f64) -> Result<(), BasisError> {
    if !shape.is_finite() || shape <= 0.0 {
        crate::bail_invalid_basis!(
            "Gamma precision hyperprior for penalty block '{label}' requires shape > 0, got {shape}"
        );
    }
    if !rate.is_finite() || rate < 0.0 {
        crate::bail_invalid_basis!(
            "Gamma precision hyperprior for penalty block '{label}' requires rate >= 0, got {rate}"
        );
    }
    Ok::<(), _>(())
}

fn validate_penalized_complexity_prior(
    label: &str,
    upper: f64,
    tail_prob: f64,
) -> Result<(), BasisError> {
    if !upper.is_finite() || upper <= 0.0 {
        crate::bail_invalid_basis!(
            "Penalized-complexity hyperprior for '{label}' requires upper > 0, got {upper}"
        );
    }
    if !tail_prob.is_finite() || tail_prob <= 0.0 || tail_prob >= 1.0 {
        crate::bail_invalid_basis!(
            "Penalized-complexity hyperprior for '{label}' requires tail probability in (0, 1), got {tail_prob}"
        );
    }
    Ok::<(), _>(())
}

fn realize_penalty_block_gamma_priors<F>(
    design: &TermCollectionDesign,
    mut callback: F,
) -> Result<crate::types::RhoPrior, BasisError>
where
    F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
{
    let mut priors = Vec::<crate::types::RhoPrior>::with_capacity(design.penaltyinfo.len());
    for info in &design.penaltyinfo {
        let metadata = penalty_block_metadata(info);
        if let Some((shape, rate)) = callback(&metadata) {
            validate_gamma_precision_prior(&metadata.label, shape, rate)?;
            priors.push(crate::types::RhoPrior::GammaPrecision { shape, rate });
        } else {
            priors.push(crate::types::RhoPrior::Flat);
        }
    }
    Ok(crate::types::RhoPrior::Independent(priors))
}

fn realize_keyed_penalty_block_gamma_priors(
    design: &TermCollectionDesign,
    priors: &[(String, f64, f64)],
) -> Result<crate::types::RhoPrior, BasisError> {
    let mut keyed = BTreeMap::<String, (f64, f64)>::new();
    for (label, shape, rate) in priors {
        if keyed.insert(label.clone(), (*shape, *rate)).is_some() {
            crate::bail_invalid_basis!(
                "duplicate Gamma precision hyperprior for penalty block label '{label}'"
            );
        }
    }
    let mut consumed = BTreeSet::<String>::new();
    let prior = realize_penalty_block_gamma_priors(design, |metadata| {
        let info = design
            .penaltyinfo
            .iter()
            .find(|info| info.global_index == metadata.global_index)
            .expect("metadata global index should match penaltyinfo");
        for label in penalty_block_label_candidates(info) {
            if let Some(value) = keyed.get(&label) {
                consumed.insert(label);
                return Some(*value);
            }
        }
        None
    })?;
    let unknown: Vec<String> = keyed
        .keys()
        .filter(|label| !consumed.contains(*label))
        .cloned()
        .collect();
    if !unknown.is_empty() {
        let available = design
            .penaltyinfo
            .iter()
            .flat_map(penalty_block_label_candidates)
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect::<Vec<_>>()
            .join(", ");
        crate::bail_invalid_basis!(
            "unknown Gamma precision hyperprior penalty block label(s): {}; available labels: {available}",
            unknown.join(", ")
        );
    }
    Ok(prior)
}

fn validate_rho_prior_coordinate(
    prior: &crate::types::RhoPrior,
    context: &str,
) -> Result<(), BasisError> {
    match prior {
        crate::types::RhoPrior::Flat => Ok(()),
        crate::types::RhoPrior::Normal { mean, sd } => {
            if !mean.is_finite() {
                crate::bail_invalid_basis!(
                    "{context} Normal log-precision prior requires finite mean, got {mean}"
                );
            }
            if !sd.is_finite() || *sd <= 0.0 {
                crate::bail_invalid_basis!(
                    "{context} Normal log-precision prior requires sd > 0, got {sd}"
                );
            }
            Ok(())
        }
        crate::types::RhoPrior::GammaPrecision { shape, rate } => {
            validate_gamma_precision_prior(context, *shape, *rate)
        }
        crate::types::RhoPrior::PenalizedComplexity { upper, tail_prob } => {
            validate_penalized_complexity_prior(context, *upper, *tail_prob)
        }
        crate::types::RhoPrior::Independent(_) => Err(BasisError::InvalidInput(format!(
            "{context} must be a scalar rho prior, not a nested Independent prior"
        ))),
    }
}

fn expand_base_rho_prior(
    base_prior: &crate::types::RhoPrior,
    base_count: usize,
    context: &str,
) -> Result<Vec<crate::types::RhoPrior>, BasisError> {
    match base_prior {
        crate::types::RhoPrior::Independent(priors) => {
            if priors.len() != base_count {
                crate::bail_invalid_basis!(
                    "{context} base Independent rho prior length mismatch: got {}, expected {base_count}",
                    priors.len()
                );
            }
            for (idx, prior) in priors.iter().enumerate() {
                validate_rho_prior_coordinate(prior, &format!("{context} base prior {idx}"))?;
            }
            Ok(priors.clone())
        }
        prior => {
            validate_rho_prior_coordinate(prior, context)?;
            Ok((0..base_count).map(|_| prior.clone()).collect())
        }
    }
}

fn combine_group_rho_prior(
    base_prior: &crate::types::RhoPrior,
    base_count: usize,
    groups: &[CoefficientGroupSpec],
) -> Result<crate::types::RhoPrior, BasisError> {
    let mut priors = Vec::with_capacity(base_count + groups.len());
    priors.extend(expand_base_rho_prior(
        base_prior,
        base_count,
        "coefficient groups",
    )?);
    for group in groups {
        let context = format!("coefficient group '{}'", group.name);
        let prior = match group.prior.as_ref() {
            Some(prior) => {
                prior.validate(&context)?;
                prior.to_rho_prior()
            }
            None => {
                validate_rho_prior_coordinate(base_prior, &context)?;
                base_prior.clone()
            }
        };
        priors.push(prior);
    }
    Ok(crate::types::RhoPrior::Independent(priors))
}

fn insert_range(
    cols: &mut BTreeSet<usize>,
    range: Range<usize>,
    p: usize,
    context: &str,
) -> Result<(), BasisError> {
    if range.end > p {
        crate::bail_dim_basis!(
            "{context} coefficient range {}..{} exceeds design width {p}",
            range.start,
            range.end
        );
    }
    cols.extend(range);
    Ok(())
}

fn resolve_group_columns(
    design: &TermCollectionDesign,
    group: &CoefficientGroupSpec,
) -> Result<BTreeSet<usize>, BasisError> {
    let p = design.design.ncols();
    let mut cols = BTreeSet::<usize>::new();
    for selector in &group.selectors {
        match selector {
            CoefficientSelector::GlobalColumns(indices) => {
                for &idx in indices {
                    if idx >= p {
                        crate::bail_dim_basis!(
                            "coefficient group '{}' references global column {idx}, but design width is {p}",
                            group.name
                        );
                    }
                    cols.insert(idx);
                }
            }
            CoefficientSelector::GlobalRange(range) => insert_range(
                &mut cols,
                range.clone(),
                p,
                &format!("coefficient group '{}'", group.name),
            )?,
            CoefficientSelector::LinearTerm(name) => {
                let (_, range) = design
                    .linear_ranges
                    .iter()
                    .find(|(term, _)| term == name)
                    .ok_or_else(|| {
                        BasisError::InvalidInput(format!(
                            "coefficient group '{}' references unknown linear term '{name}'",
                            group.name
                        ))
                    })?;
                insert_range(&mut cols, range.clone(), p, &group.name)?;
            }
            CoefficientSelector::RandomEffectTerm(name) => {
                let (_, range) = design
                    .random_effect_ranges
                    .iter()
                    .find(|(term, _)| term == name)
                    .ok_or_else(|| {
                        BasisError::InvalidInput(format!(
                            "coefficient group '{}' references unknown random-effect term '{name}'",
                            group.name
                        ))
                    })?;
                insert_range(&mut cols, range.clone(), p, &group.name)?;
            }
            CoefficientSelector::SmoothTerm(name) => {
                let term = design
                    .smooth
                    .terms
                    .iter()
                    .find(|term| &term.name == name)
                    .ok_or_else(|| {
                        BasisError::InvalidInput(format!(
                            "coefficient group '{}' references unknown smooth term '{name}'",
                            group.name
                        ))
                    })?;
                let start = p - design.smooth.total_smooth_cols() + term.coeff_range.start;
                insert_range(
                    &mut cols,
                    start..(start + term.coeff_range.len()),
                    p,
                    &group.name,
                )?;
            }
            CoefficientSelector::SmoothTermColumns { term, columns } => {
                let smooth_term = design
                    .smooth
                    .terms
                    .iter()
                    .find(|smooth_term| &smooth_term.name == term)
                    .ok_or_else(|| {
                        BasisError::InvalidInput(format!(
                            "coefficient group '{}' references unknown smooth term '{term}'",
                            group.name
                        ))
                    })?;
                let smooth_start = p - design.smooth.total_smooth_cols();
                for &local_col in columns {
                    if local_col >= smooth_term.coeff_range.len() {
                        crate::bail_dim_basis!(
                            "coefficient group '{}' references smooth term '{term}' local column {local_col}, but the term has {} columns",
                            group.name,
                            smooth_term.coeff_range.len()
                        );
                    }
                    cols.insert(smooth_start + smooth_term.coeff_range.start + local_col);
                }
            }
        }
    }
    if cols.is_empty() {
        crate::bail_invalid_basis!(
            "coefficient group '{}' contains no coefficients",
            group.name
        );
    }
    Ok(cols)
}

fn realize_coefficient_groups(
    design: &TermCollectionDesign,
    groups: &[CoefficientGroupSpec],
    base_prior: &crate::types::RhoPrior,
) -> Result<RealizedCoefficientGroups, BasisError> {
    use crate::terms::coefficient_group_resolver::{ResolvedGroup, ResolvedGroupHierarchy};

    let p = design.design.ncols();
    // Carrier-specific validation and selector resolution. The standard-term
    // carrier is columns of the realized design matrix; `resolve_group_columns`
    // turns each declared selector into a `BTreeSet<usize>` and rejects empty
    // selector lists. The prior is validated here because its diagnostic
    // context uses the standard-term label.
    for group in groups {
        if group.selectors.is_empty() {
            crate::bail_invalid_basis!("coefficient group '{}' contains no selectors", group.name);
        }
        if let Some(prior) = group.prior.as_ref() {
            prior.validate(&format!("coefficient group '{}'", group.name))?;
        }
    }

    let resolved_groups = groups
        .iter()
        .map(|group| {
            Ok(ResolvedGroup {
                label: group.name.clone(),
                parent: group.parent.clone(),
                coordinates: resolve_group_columns(design, group)?,
            })
        })
        .collect::<Result<Vec<_>, BasisError>>()?;
    // Carrier-agnostic policy: unique non-empty labels, acyclic hierarchy,
    // child ⊆ parent, interior == union of children.
    let hierarchy =
        ResolvedGroupHierarchy::build(resolved_groups).map_err(BasisError::InvalidInput)?;

    let mut penalty_specs: Vec<PenaltySpec> = design
        .penalties
        .iter()
        .map(PenaltySpec::from_blockwise_ref)
        .collect();
    let mut nullspace_dims = design.nullspace_dims.clone();
    let mut group_column_indices = Vec::<(String, Vec<usize>)>::with_capacity(groups.len());
    for (group, resolved) in groups.iter().zip(hierarchy.groups()) {
        let cols = &resolved.coordinates;
        let mut penalty = Array2::<f64>::zeros((p, p));
        let penalty_components = hierarchy.concatenated_penalty_components(&group.name);
        let active_cols = penalty_components
            .iter()
            .flat_map(|component| component.iter().copied())
            .collect::<BTreeSet<_>>();
        let local_mean = group
            .prior_mean
            .evaluate(
                active_cols.len(),
                &format!("coefficient group '{}'", group.name),
            )
            .map_err(|err| BasisError::InvalidInput(err.to_string()))?;
        let mut prior_mean = Array1::<f64>::zeros(p);
        // Hierarchical Gamma precision update.
        //
        // For a leaf group,
        //
        //   p(beta_g | lambda_g) p(lambda_g)
        //     ∝ lambda_g^{|g|/2}
        //       exp[-lambda_g (beta_g - mu_g)' S_g (beta_g - mu_g) / 2]
        //       lambda_g^{a_g-1} exp[-b_g lambda_g],
        //
        // so fixed-beta MAP gives
        //
        //   lambda_g* = (a_g + |g|/2 - 1)
        //               / (b_g + (beta_g - mu_g)' S_g (beta_g - mu_g) / 2).
        //
        // Interior nodes use the same identity with beta_g formed by
        // concatenating child beta vectors.  Equivalently, |g| and the
        // quadratic are sums over the recursively expanded child factors.  In
        // the standard term-collection path there is one rho coordinate per
        // group, so we materialize that summed child penalty into the group's
        // dense S_g.  Leaves reduce to the ordinary identity penalty.
        for component in &penalty_components {
            for &col in component {
                penalty[[col, col]] += 1.0;
            }
        }
        for (mean_idx, &col) in active_cols.iter().enumerate() {
            prior_mean[col] = local_mean[mean_idx];
        }
        penalty_specs.push(PenaltySpec::DenseWithMean {
            matrix: penalty,
            prior_mean: crate::solver::estimate::CoefficientPriorMean::constant(prior_mean),
        });
        nullspace_dims.push(p.saturating_sub(active_cols.len()));
        group_column_indices.push((group.name.clone(), cols.iter().copied().collect()));
    }

    Ok(RealizedCoefficientGroups {
        penalty_specs,
        nullspace_dims,
        rho_prior: combine_group_rho_prior(base_prior, design.penalties.len(), groups)?,
        group_column_indices,
    })
}

#[derive(Clone)]
pub struct FittedTermCollection {
    pub fit: UnifiedFitResult,
    pub design: TermCollectionDesign,
    pub adaptive_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
}

#[derive(Clone)]
pub struct FittedTermCollectionWithSpec {
    pub fit: UnifiedFitResult,
    pub design: TermCollectionDesign,
    pub resolvedspec: TermCollectionSpec,
    pub adaptive_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
}

#[derive(Clone)]
pub struct StandardLatentCoordConfig {
    pub values: std::sync::Arc<crate::terms::latent_coord::LatentCoordValues>,
    pub term_index: crate::types::SmoothTermIdx,
    pub feature_cols: Vec<usize>,
    pub manifold: crate::terms::latent_coord::LatentManifold,
    pub manifold_auto: bool,
    pub retraction_registry: crate::solver::latent_cache::LatentRetractionRegistry,
    pub analytic_penalties: Option<std::sync::Arc<crate::terms::AnalyticPenaltyRegistry>>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AdaptiveSpatialMap {
    pub termname: String,
    pub feature_cols: Vec<usize>,
    pub collocation_points: Array2<f64>,
    pub inv_magweight: Array1<f64>,
    pub invgradweight: Array1<f64>,
    pub inv_lapweight: Array1<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AdaptiveRegularizationDiagnostics {
    pub epsilon_0: f64,
    pub epsilon_g: f64,
    pub epsilon_c: f64,
    pub epsilon_outer_iterations: usize,
    pub mm_iterations: usize,
    pub converged: bool,
    pub maps: Vec<AdaptiveSpatialMap>,
}

#[derive(Debug, Clone)]
struct LinearColumnConditioning {
    col_idx: usize,
    mean: f64,
    scale: f64,
}

#[derive(Debug, Clone, Default)]
struct LinearFitConditioning {
    intercept_idx: usize,
    columns: Vec<LinearColumnConditioning>,
}

#[derive(Clone)]
pub(crate) struct SpatialPsiDerivative {
    // These are derivatives with respect to psi = log(kappa), not log(length_scale).
    pub penalty_index: usize,
    pub penalty_indices: Vec<usize>,
    pub global_range: Range<usize>,
    pub total_p: usize,
    pub x_psi_local: Array2<f64>,
    pub s_psi_components_local: Vec<Array2<f64>>,
    pub x_psi_psi_local: Array2<f64>,
    pub s_psi_psi_components_local: Vec<Array2<f64>>,
    pub aniso_group_id: Option<usize>,
    /// Pre-computed cross-derivative design matrices for other axes
    /// in the same aniso group: Vec of (axis_offset_in_group, matrix).
    pub aniso_cross_designs: Option<Vec<(usize, Array2<f64>)>>,
    /// On-demand cross-penalty second derivatives ∂²S_m/∂ψ_a∂ψ_b for axes in
    /// the same anisotropy group. The input is the other axis offset in the
    /// group, and the output is one local penalty matrix per active penalty.
    pub aniso_cross_penalty_provider: Option<
        std::sync::Arc<
            dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError> + Send + Sync + 'static,
        >,
    >,
    /// Optional implicit design-derivative operator (shared across all axes
    /// in the same aniso group). When present, `x_psi_local` and
    /// `x_psi_psi_local` may be zero-sized, and design-derivative matvecs
    /// should go through this operator using `implicit_axis` as the axis index.
    pub implicit_operator: Option<std::sync::Arc<crate::terms::basis::ImplicitDesignPsiDerivative>>,
    /// Which axis in the implicit operator this entry corresponds to.
    pub implicit_axis: usize,
}

#[derive(Debug, Clone)]
pub(crate) struct SpatialLogKappaCoords {
    /// Flattened ψ values. For isotropic terms, one entry per term.
    /// For anisotropic terms, d entries per term (one ψ_a per axis).
    values: Array1<f64>,
    /// Dimensionality of each term: 1 for isotropic, d for anisotropic.
    dims_per_term: Vec<usize>,
}

/// Which end of the ψ bound the shared `aniso_bounds_from_data` helper is
/// computing. The lower end uses `-max_length_scale.ln()` as the pure-Duchon
/// fallback and the `.0` element of `spatial_term_psi_bounds`; the upper end
/// uses `-min_length_scale.ln()` and `.1`. Everything else is identical.
#[derive(Clone, Copy)]
enum AnisoBoundEnd {
    Lower,
    Upper,
}

impl SpatialLogKappaCoords {
    /// Construct from an explicit dims layout plus values.
    pub(crate) fn new_with_dims(values: Array1<f64>, dims_per_term: Vec<usize>) -> Self {
        assert_eq!(
            values.len(),
            dims_per_term.iter().sum::<usize>(),
            "SpatialLogKappaCoords: values length {} != sum of dims_per_term {}",
            values.len(),
            dims_per_term.iter().sum::<usize>(),
        );
        Self {
            values,
            dims_per_term,
        }
    }

    /// Isotropic initialization (backward-compatible path).
    pub(crate) fn from_length_scales(
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        let mut out = Array1::<f64>::zeros(term_indices.len());
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            let length_scale = get_spatial_length_scale(spec, term_idx)
                .unwrap_or(options.min_length_scale)
                .clamp(options.min_length_scale, options.max_length_scale);
            out[slot] = -length_scale.ln();
        }
        Self {
            values: out,
            dims_per_term: vec![1; term_indices.len()],
        }
    }

    /// Anisotropic-aware initialization.
    ///
    /// Initialization strategy (per math team recommendation): standardize the
    /// knot cloud axiswise, then run the existing isotropic κ initializer in
    /// the standardized space. This reuses the trusted isotropic initializer
    /// and gives initial η_a = −ln(σ_a) + mean(ln(σ_a)), which satisfies
    /// Ση_a = 0 by construction.
    ///
    /// For each term, checks whether it has `aniso_log_scales` set on its basis spec.
    /// - If isotropic (no aniso_log_scales, or 1-D): 1 entry = −ln(length_scale).
    /// - If anisotropic with a scalar length scale: d entries, one ψ_a per axis.
    ///   Initialized as ψ_a = −ln(length_scale) + η_a  where η_a are the existing
    ///   aniso_log_scales (which sum to zero). Multi-dimensional terms without
    ///   explicit anisotropy stay scalar here so the seed dimensionality matches
    ///   `spatial_dims_per_term`.
    /// - If pure Duchon anisotropic: d - 1 free entries store the leading η_a
    ///   values directly; the final axis is reconstructed to keep Ση_a = 0.
    pub(crate) fn from_length_scales_aniso(
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        let mut vals = Vec::new();
        let mut dims = Vec::new();
        for &term_idx in term_indices {
            let length_scale = get_spatial_length_scale(spec, term_idx)
                .unwrap_or(options.min_length_scale)
                .clamp(options.min_length_scale, options.max_length_scale);
            let psi_bar = -length_scale.ln(); // global scale = −ln(length_scale)

            if spatial_term_uses_per_axis_psi(spec, term_idx) {
                // Per-axis anisotropy is enrolled in the joint outer vector:
                // ψ_a = ψ̄ + η_a, one slot per axis. The hyper_dirs builder
                // produces matching per-axis derivatives in
                // `try_build_spatial_term_log_kappa_aniso_derivativeinfos`.
                let d = get_spatial_feature_dim(spec, term_idx).unwrap_or(1);
                let eta_raw = get_spatial_aniso_log_scales(spec, term_idx)
                    .expect("predicate guarantees aniso_log_scales is Some");
                let eta = center_aniso_log_scales(&eta_raw);
                for &eta_a in &eta {
                    vals.push(psi_bar + eta_a);
                }
                dims.push(d);
            } else {
                // Isotropic enrollment — either a 1-D term, a multi-D term
                // without explicit anisotropy, or a basis (e.g. Duchon) whose
                // η is a fixed geometry parameter rather than a REML hyper
                // axis. Exactly one ψ̄ slot, matching the single
                // `SpatialPsiDerivative` produced by
                // `try_build_spatial_term_log_kappa_derivativeinfo`.
                vals.push(psi_bar);
                dims.push(1);
            }
        }
        Self {
            values: Array1::from_vec(vals),
            dims_per_term: dims,
        }
    }

    /// Isotropic lower bounds derived from per-term data geometry.
    /// Each entry gets the ψ_lo bound returned by `spatial_term_psi_bounds`
    /// for the corresponding term, intersected with the options window.
    pub(crate) fn lower_bounds_from_data(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        let mut values = Array1::<f64>::zeros(term_indices.len());
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options).0;
        }
        Self {
            values,
            dims_per_term: vec![1; term_indices.len()],
        }
    }

    /// Isotropic upper bounds derived from per-term data geometry.
    pub(crate) fn upper_bounds_from_data(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        let mut values = Array1::<f64>::zeros(term_indices.len());
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options).1;
        }
        Self {
            values,
            dims_per_term: vec![1; term_indices.len()],
        }
    }

    /// Anisotropic-aware lower bounds derived from per-term data geometry.
    /// For hybrid anisotropic terms the scalar ψ_lo bound applies to the
    /// mean `ψ̄`, not directly to every raw axis coordinate `ψ_a = ψ̄ + η_a`.
    /// Shift each axis by the current centered `η_a` so projecting/clamping
    /// the seed moves only the global scale direction and does not silently
    /// shrink anisotropy that is already consistent with the current
    /// `length_scale`.
    ///
    /// Pure Duchon anisotropy is structurally different: its stored
    /// coordinates are (d-1) free η_a values representing log axis-scale
    /// ratios, NOT log-κ. For those terms the κ-range geometry bound is
    /// over-restrictive (η_a = ±5 is normal, but that corresponds to 7+
    /// orders of magnitude in κ-space and would be rejected by the data
    /// window). Fall back to the options window `[-ln(max_ls), -ln(min_ls)]`
    /// for those coordinates — that's the same bound the pre-data-geometry
    /// code used, which is calibrated to allow legitimate anisotropy.
    pub(crate) fn lower_bounds_aniso_from_data(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        dims_per_term: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        Self::aniso_bounds_from_data(
            data,
            spec,
            term_indices,
            dims_per_term,
            options,
            AnisoBoundEnd::Lower,
        )
    }

    /// Anisotropic-aware upper bounds derived from per-term data geometry.
    /// See `lower_bounds_aniso_from_data` for the hybrid-aniso offsetting and
    /// pure-Duchon dispatch rationale.
    pub(crate) fn upper_bounds_aniso_from_data(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        dims_per_term: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        Self::aniso_bounds_from_data(
            data,
            spec,
            term_indices,
            dims_per_term,
            options,
            AnisoBoundEnd::Upper,
        )
    }

    /// Shared implementation for the lower/upper aniso bounds. The bound end
    /// only changes which options scale (`max_length_scale` vs
    /// `min_length_scale`) becomes the pure-Duchon fallback bound and which
    /// element of the `(lo, hi)` data-geometry tuple is consumed; the
    /// per-term cursor walk and aniso-offset handling are identical.
    fn aniso_bounds_from_data(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        dims_per_term: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
        end: AnisoBoundEnd,
    ) -> Self {
        assert_eq!(term_indices.len(), dims_per_term.len());
        let total: usize = dims_per_term.iter().sum();
        let mut values = Array1::<f64>::zeros(total);
        let mut cursor = 0;
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            let d = dims_per_term[slot];
            let psi_bound = {
                let (lo, hi) = spatial_term_psi_bounds(data, spec, term_idx, options);
                match end {
                    AnisoBoundEnd::Lower => lo,
                    AnisoBoundEnd::Upper => hi,
                }
            };
            let axis_offsets = if d <= 1 {
                vec![0.0; d]
            } else {
                get_spatial_aniso_log_scales(spec, term_idx)
                    .filter(|eta| eta.len() == d)
                    .map(|eta| center_aniso_log_scales(&eta))
                    .unwrap_or_else(|| vec![0.0; d])
            };
            for offset in 0..d {
                values[cursor + offset] = psi_bound + axis_offsets[offset];
            }
            cursor += d;
        }
        Self {
            values,
            dims_per_term: dims_per_term.to_vec(),
        }
    }

    /// Rewrite any ψ entries whose originating term lacks an explicit
    /// `length_scale` so they sit at the midpoint of the per-term data-derived
    /// ψ window. Used so the outer optimizer starts inside the physically
    /// meaningful region instead of at an arbitrary `options.max_length_scale`
    /// derived seed. For terms with an explicit length_scale, the user's
    /// choice is respected. Anisotropy offsets η_a (those stored by
    /// `from_length_scales_aniso`) are preserved: we re-center around the new
    /// ψ̄, keeping Ση_a = 0.
    pub(crate) fn reseed_from_data(
        mut self,
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
    ) -> Self {
        assert_eq!(term_indices.len(), self.dims_per_term.len());
        let mut cursor = 0;
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            let d = self.dims_per_term[slot];
            let Some(psi_bar_new) = spatial_term_psi_seed(data, spec, term_idx, options) else {
                cursor += d;
                continue;
            };
            if d == 0 {
                continue;
            }
            let current: Vec<f64> = self.values.slice(s![cursor..cursor + d]).to_vec();
            let psi_bar_old = current.iter().sum::<f64>() / d as f64;
            for (offset, &old_value) in current.iter().enumerate() {
                self.values[cursor + offset] = psi_bar_new + (old_value - psi_bar_old);
            }
            cursor += d;
        }
        self
    }

    /// Project ψ values into `[lower, upper]` element-wise. Used after
    /// `from_length_scales*` + `reseed_from_data` when a user-supplied
    /// `spec.length_scale` falls outside the data-derived ψ window set by
    /// `{lower,upper}_bounds*_from_data`. BFGS requires theta0 ∈ [lower,
    /// upper]; projecting is the unique closest feasible seed. The user's
    /// length_scale was always a hint for the outer optimizer (the optimizer
    /// is authoritative for κ), not a hard constraint — so clipping preserves
    /// their intent as far as the geometry allows. Emits `log::info!` when
    /// any coordinate moves, so the outside-window case is diagnostically
    /// visible (not silent).
    pub(crate) fn clamp_to_bounds(
        mut self,
        lower: &SpatialLogKappaCoords,
        upper: &SpatialLogKappaCoords,
    ) -> Self {
        assert_eq!(self.values.len(), lower.values.len());
        assert_eq!(self.values.len(), upper.values.len());
        let mut n_projected = 0usize;
        let mut worst_delta = 0.0_f64;
        for idx in 0..self.values.len() {
            let lo = lower.values[idx];
            let hi = upper.values[idx];
            if !(lo.is_finite() && hi.is_finite()) {
                continue;
            }
            let v = self.values[idx];
            if v < lo {
                worst_delta = worst_delta.max(lo - v);
                self.values[idx] = lo;
                n_projected += 1;
            } else if v > hi {
                worst_delta = worst_delta.max(v - hi);
                self.values[idx] = hi;
                n_projected += 1;
            }
        }
        if n_projected > 0 {
            log::info!(
                "[spatial-kappa] projected {n_projected}/{} ψ seed coords into data-derived bounds \
                 (worst excess={worst_delta:.3} log units); user length_scale falls outside \
                 [2/r_max, 1e2/r_min] geometry window",
                self.values.len()
            );
        }
        self
    }

    /// Reconstruct from theta tail with known dimensionality layout.
    pub(crate) fn from_theta_tail_with_dims(
        theta: &Array1<f64>,
        start: usize,
        dims_per_term: Vec<usize>,
    ) -> Self {
        let total: usize = dims_per_term.iter().sum();
        Self {
            values: theta.slice(s![start..start + total]).to_owned(),
            dims_per_term,
        }
    }

    /// Total number of ψ values in the flat array (= sum of dims_per_term).
    pub(crate) fn len(&self) -> usize {
        self.values.len()
    }

    /// Dimensionality layout: how many ψ values each term contributes.
    pub(crate) fn dims_per_term(&self) -> &[usize] {
        &self.dims_per_term
    }

    /// Get the offset into the flat array for logical term i.
    fn term_offset(&self, term_idx: usize) -> usize {
        self.dims_per_term[..term_idx].iter().sum()
    }

    /// Get the slice of ψ values for logical term i.
    fn term_slice(&self, term_idx: usize) -> &[f64] {
        let offset = self.term_offset(term_idx);
        let d = self.dims_per_term[term_idx];
        &self.values.as_slice().unwrap()[offset..offset + d]
    }

    pub(crate) fn as_array(&self) -> &Array1<f64> {
        &self.values
    }

    /// Split at a logical-term boundary. `mid` is the number of terms in the
    /// first half (not a flat-array index).
    pub(crate) fn split_at(&self, mid: usize) -> (Self, Self) {
        let flat_mid: usize = self.dims_per_term[..mid].iter().sum();
        (
            Self {
                values: self.values.slice(s![0..flat_mid]).to_owned(),
                dims_per_term: self.dims_per_term[..mid].to_vec(),
            },
            Self {
                values: self.values.slice(s![flat_mid..]).to_owned(),
                dims_per_term: self.dims_per_term[mid..].to_vec(),
            },
        )
    }

    /// Apply optimized ψ values back to the spec.
    ///
    /// For isotropic terms (dims=1): sets scalar length_scale = exp(−ψ).
    /// For anisotropic terms (dims=d): hybrid/isotropic families set
    /// length_scale = exp(−ψ̄) with centered η_a = ψ_a − ψ̄, while pure Duchon
    /// writes only centered η_a and leaves length_scale = None.
    pub(crate) fn apply_tospec(
        &self,
        spec: &TermCollectionSpec,
        term_indices: &[usize],
    ) -> Result<TermCollectionSpec, EstimationError> {
        if term_indices.len() != self.dims_per_term.len() {
            crate::bail_invalid_estim!(
                "SpatialLogKappaCoords::apply_tospec: term count mismatch: \
                 term_indices={} dims_per_term={}",
                term_indices.len(),
                self.dims_per_term.len()
            );
        }
        let mut updated = spec.clone();
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            let psi = self.term_slice(slot);
            let d = self.dims_per_term[slot];
            let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
            if (d == 1 || next_length_scale.is_some())
                && let Some(length_scale) = next_length_scale
            {
                set_spatial_length_scale(&mut updated, term_idx, length_scale)?;
            }
            if let Some(eta) = next_aniso {
                set_spatial_aniso_log_scales(&mut updated, term_idx, eta)?;
            }
        }
        Ok(updated)
    }
}

fn center_aniso_log_scales(eta: &[f64]) -> Vec<f64> {
    if eta.len() <= 1 {
        return eta.to_vec();
    }
    let mean = eta.iter().sum::<f64>() / eta.len() as f64;
    eta.iter()
        .map(|&v| {
            let centered = v - mean;
            if centered.abs() <= 1e-15 {
                0.0
            } else {
                centered
            }
        })
        .collect()
}

fn spatial_term_supports_hyper_optimization(spec: &TermCollectionSpec, term_idx: usize) -> bool {
    // Ordinary penalized thin-plate regression splines do not have an
    // identifiable kernel scale once REML is already learning the smoothing
    // penalty. Treat the resolved length scale as fixed geometry; enrolling a
    // scalar TPS kappa axis creates the flat ρ/κ valleys reported in #718,
    // #721, #731, and #732.
    if let Some(term) = spec.smooth_terms.get(term_idx)
        && let SmoothBasisSpec::ThinPlate { .. } = &term.basis
    {
        return false;
    }

    // Duchon anisotropy η is a FIXED, geometry-derived basis parameter, NOT a
    // REML hyper axis: the metric is estimated once from the knot-cloud spread
    // (`maybe_initialize_aniso_contrasts`, applied on every basis build) and the
    // Hilbert-scale λ's carry all learned smoothness. So a pure Duchon (no κ)
    // contributes no outer optimization axis even when `scale_dims` is on —
    // "standardize the geometry, then learn the smoothness." Only an explicit
    // kernel length scale κ (the Matérn / hybrid path) is optimized here.
    //
    // ISOTROPIC Matérn: the *default* `matern(x1, x2)` is isotropic
    // (`scale_dims=false` → `aniso_log_scales = None`). It contributes exactly
    // ONE κ optimization axis — its scalar log-κ. The shared GAMLSS /
    // location-scale exact-joint ψ engine and the spatial-κ joint outer solver
    // both require an isotropic Matérn block to expose this single isotropic κ
    // axis (#822/#851); without it the per-block ψ-derivative lists are empty
    // and the joint-ψ hooks degenerate to `None`. The isotropic κ is the lone
    // kernel hyper axis here, mirroring the per-axis ψ ARD that the anisotropic
    // path exposes (just collapsed to one dimension).
    //
    // ANISOTROPIC Matérn (`scale_dims=true` → `aniso_log_scales = Some`) keeps
    // its per-axis kernel-η ARD: the d-dimensional ψ search is the *point* of
    // the anisotropic request ("Matérn keeps its kernel-η ARD").
    //
    // Either way a Matérn term always enrolls a κ/ψ axis (1 isotropic, or d
    // anisotropic), so `spatial_dims_per_term` reports the correct count.
    if let Some(term) = spec.smooth_terms.get(term_idx)
        && let SmoothBasisSpec::Matern { .. } = &term.basis
    {
        return true;
    }
    get_spatial_length_scale(spec, term_idx).is_some()
}

/// Returns `true` when a spatial term has NO outer optimization axes — i.e.
/// the user provided an explicit `length_scale` and the term does not enroll
/// REML-side per-axis ψ contrasts, so both the scalar κ and any fixed geometry
/// anisotropy are anchored.
///
/// This is the per-term predicate that distinguishes "fixed kernel scale"
/// from "optimize the kernel scale" within the family entry points that
/// want to honor an explicit user-supplied scale (e.g. Bernoulli
/// marginal-slope, where the joint-spatial outer solver otherwise spends
/// ~80 iters stalled on the user's chosen ρ at high gradient).
pub fn spatial_term_has_locked_kappa(spec: &TermCollectionSpec, term_idx: usize) -> bool {
    get_spatial_length_scale(spec, term_idx).is_some()
        && !spatial_term_uses_per_axis_psi(spec, term_idx)
}

/// Per-term data-derived ψ = log κ bounds.
///
/// Uses the same safe operating range documented in
/// [`crate::basis::build_matern_basis`] / [`crate::basis::build_duchon_basis`]:
///   κ ∈ [2 / r_max, 1e2 / r_min]
/// where (r_min, r_max) are pairwise-distance extrema of the term's resolved
/// centers (post-fit) or the standardized feature data columns (pre-fit).
/// Returns ψ-space bounds (ψ_lo = ln(κ_lo), ψ_hi = ln(κ_hi)).
///
/// When geometry is unavailable (e.g., fewer than 2 distinct points), falls
/// back to the scalar `options.min_length_scale` / `options.max_length_scale`
/// window so the outer optimizer never sees NaN bounds.
///
/// The returned window is intersected with the options window so user-set
/// `min_length_scale` / `max_length_scale` remain hard limits.
fn spatial_term_psi_bounds(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
    term_idx: usize,
    options: &SpatialLengthScaleOptimizationOptions,
) -> (f64, f64) {
    let fallback = (
        -options.max_length_scale.ln(),
        -options.min_length_scale.ln(),
    );
    let Some(term) = spec.smooth_terms.get(term_idx) else {
        return fallback;
    };
    // Prefer resolved centers (post-fit) since they live in the same standardized
    // space the kernel actually sees. Centers are capped at `default_num_centers`
    // (<=2000), so exact pairwise bounds are cheap (<4M ops). If centers are
    // not yet UserProvided, fall back to the standardized feature data columns
    // with the capped-sample path (O(K²·d), K=1024) — the sample is
    // conservative for κ bounds (see `pairwise_distance_bounds_sampled`
    // docs): it never excludes a feasible κ the exact method would include.
    //
    // Under anisotropy the kernel metric is y-space (y_a = exp(η_a) x_a),
    // so r_min/r_max must be y-space distances. This matters only when the
    // spec already carries calibrated η_a at setup time (e.g., warm-start
    // or refit paths); for fresh optimization η_a starts at 0 and y = x.
    let aniso = get_spatial_aniso_log_scales(spec, term_idx);
    let r_bounds = match spatial_term_center_strategy(term) {
        Some(CenterStrategy::UserProvided(centers)) if centers.nrows() >= 2 => {
            match aniso.as_deref() {
                Some(eta) if eta.len() == centers.ncols() => {
                    let y = points_in_aniso_y_space(centers.view(), eta);
                    pairwise_distance_bounds(y.view())
                }
                _ => pairwise_distance_bounds(centers.view()),
            }
        }
        _ => standardized_spatial_term_data(data, term)
            .ok()
            .and_then(|x| match aniso.as_deref() {
                Some(eta) if eta.len() == x.ncols() => {
                    let y = points_in_aniso_y_space(x.view(), eta);
                    pairwise_distance_bounds_sampled(y.view())
                }
                _ => pairwise_distance_bounds_sampled(x.view()),
            }),
    };
    let Some((r_min, r_max)) = r_bounds else {
        return fallback;
    };
    // Length scales substantially larger than the data diameter make radial
    // TPS/Matern columns nearly collinear with their polynomial nullspace.
    // The nullspace already carries constant/linear low-frequency structure,
    // so cap the kernel range at the diameter scale instead of letting the
    // optimizer enter a numerically degenerate basis geometry.
    let psi_lo_data = (2.0 / r_max).ln();
    let psi_hi_data = (1e2 / r_min).ln();
    // Intersect with the options window so min/max_length_scale remain hard caps.
    let psi_lo = psi_lo_data.max(fallback.0);
    let psi_hi = psi_hi_data.min(fallback.1);
    if psi_lo >= psi_hi {
        // Degenerate intersection — fall back to the options window to keep the
        // outer optimizer from collapsing to a point.
        return fallback;
    }
    (psi_lo, psi_hi)
}

/// Data-derived ψ seed for a spatial term when the user has not set an
/// explicit length_scale on its basis spec. Uses the geometric mean of the
/// data-informed kappa range (i.e., the midpoint of the ψ window).
fn spatial_term_psi_seed(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
    term_idx: usize,
    options: &SpatialLengthScaleOptimizationOptions,
) -> Option<f64> {
    if get_spatial_length_scale(spec, term_idx).is_some() {
        return None; // user/spec-provided length_scale wins
    }
    let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options);
    Some(0.5 * (psi_lo + psi_hi))
}

fn spatial_term_psi_to_length_scale_and_aniso(psi: &[f64]) -> (Option<f64>, Option<Vec<f64>>) {
    if psi.len() <= 1 {
        (Some((-psi.first().copied().unwrap_or(0.0)).exp()), None)
    } else {
        let psi_bar = psi.iter().sum::<f64>() / psi.len() as f64;
        (
            Some((-psi_bar).exp()),
            Some(psi.iter().map(|&value| value - psi_bar).collect()),
        )
    }
}

/// Get the `aniso_log_scales` from a spatial term, if present.
pub fn get_spatial_aniso_log_scales(
    spec: &TermCollectionSpec,
    term_idx: usize,
) -> Option<Vec<f64>> {
    spec.smooth_terms
        .get(term_idx)
        .and_then(|term| match &term.basis {
            SmoothBasisSpec::Matern { spec, .. } => spec.aniso_log_scales.clone(),
            SmoothBasisSpec::Duchon { spec, .. } => spec.aniso_log_scales.clone(),
            _ => None,
        })
}

/// Get the number of feature columns (spatial dimensionality) for a spatial term.
fn get_spatial_feature_dim(spec: &TermCollectionSpec, term_idx: usize) -> Option<usize> {
    spec.smooth_terms
        .get(term_idx)
        .and_then(|term| match &term.basis {
            SmoothBasisSpec::ThinPlate { feature_cols, .. } => Some(feature_cols.len()),
            SmoothBasisSpec::Matern { feature_cols, .. } => Some(feature_cols.len()),
            SmoothBasisSpec::Duchon { feature_cols, .. } => Some(feature_cols.len()),
            _ => None,
        })
}

/// Log the learned per-axis spatial anisotropy for all spatial terms that
/// have `aniso_log_scales` set after optimization.
///
/// For scalar-scale families this reports eta, effective per-axis length
/// scales, and per-axis kappa values. For pure Duchon it reports the centered
/// eta contrasts only.
pub fn log_spatial_aniso_scales(spec: &TermCollectionSpec) {
    for (term_idx, term) in spec.smooth_terms.iter().enumerate() {
        let (aniso, length_scale) = match &term.basis {
            SmoothBasisSpec::Matern { spec, .. } => {
                (spec.aniso_log_scales.as_ref(), Some(spec.length_scale))
            }
            SmoothBasisSpec::Duchon { spec, .. } => {
                (spec.aniso_log_scales.as_ref(), spec.length_scale)
            }
            _ => (None, None),
        };
        let Some(eta) = aniso else { continue };
        if eta.is_empty() {
            continue;
        }
        let mut lines = match length_scale {
            Some(ls) => format!(
                "[spatial-kappa] term {} (\"{}\"): anisotropic length scales optimized (global length_scale={:.4})",
                term_idx, term.name, ls
            ),
            None => format!(
                "[spatial-kappa] term {} (\"{}\"): pure Duchon shape anisotropy optimized",
                term_idx, term.name
            ),
        };
        for (a, &eta_a) in eta.iter().enumerate() {
            if let Some(ls) = length_scale {
                let length_a = ls * (-eta_a).exp();
                let kappa_a = (1.0 / ls) * eta_a.exp();
                lines.push_str(&format!(
                    "\n  axis {}: eta={:+.4}, length={:.4}, kappa={:.4}",
                    a, eta_a, length_a, kappa_a
                ));
            } else {
                lines.push_str(&format!("\n  axis {}: eta={:+.4}", a, eta_a));
            }
        }
        log::info!("{}", lines);
    }
}

/// Set `aniso_log_scales` on a spatial term's basis spec.
fn set_spatial_aniso_log_scales(
    spec: &mut TermCollectionSpec,
    term_idx: usize,
    eta: Vec<f64>,
) -> Result<(), EstimationError> {
    let eta = center_aniso_log_scales(&eta);
    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
        crate::bail_invalid_estim!("spatial aniso_log_scales term index {term_idx} out of range");
    };
    match &mut term.basis {
        SmoothBasisSpec::Matern { spec, .. } => {
            spec.aniso_log_scales = Some(eta);
            Ok(())
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            spec.aniso_log_scales = Some(eta);
            Ok(())
        }
        _ => Err(EstimationError::InvalidInput(format!(
            "term '{}' does not support aniso_log_scales",
            term.name
        ))),
    }
}

/// Sync knot-cloud-derived anisotropy contrasts from basis metadata back into
/// the mutable spec so the optimizer starts from the correct eta values.
///
/// Call this after building the smooth design but before initializing the
/// optimizer's psi coordinates. For each spatial term whose metadata contains
/// computed `aniso_log_scales`, this writes them into the spec.
pub(crate) fn sync_aniso_contrasts_from_metadata(
    spec: &mut TermCollectionSpec,
    design: &SmoothDesign,
) {
    for (term_idx, term) in design.terms.iter().enumerate() {
        let meta_aniso = match &term.metadata {
            BasisMetadata::Matern {
                aniso_log_scales, ..
            } => aniso_log_scales.clone(),
            BasisMetadata::Duchon {
                aniso_log_scales, ..
            } => aniso_log_scales.clone(),
            _ => None,
        };
        if let Some(eta) = meta_aniso
            && eta.len() > 1
        {
            set_spatial_aniso_log_scales(spec, term_idx, eta).ok();
        }
    }
}

#[derive(Debug, Clone)]
pub struct SpatialLengthScaleOptimizationOptions {
    /// Enable outer-loop optimization over spatial κ (= 1 / length_scale)
    /// for supported radial-kernel smooths.
    /// This applies to ThinPlate, Matérn, and Duchon terms.
    pub enabled: bool,
    /// Maximum number of outer iterations in the exact joint [rho, psi] solve.
    pub max_outer_iter: usize,
    /// Relative improvement threshold for terminating the outer solve.
    pub rel_tol: f64,
    /// Initial log(length_scale) perturbation used for seed construction.
    pub log_step: f64,
    /// Minimum allowed length_scale during κ search.
    pub min_length_scale: f64,
    /// Maximum allowed length_scale during κ search.
    pub max_length_scale: f64,
    /// Automatic geometry-initializer threshold for biobank-scale spatial fits.
    ///
    /// When n exceeds twice this value, the fitter uses a spatially stratified
    /// subsample only to seed κ/anisotropy geometry: centers are resolved,
    /// axis contrasts are initialized from center/data spread, and one or two
    /// cheap ψ reseeding updates are applied. It never runs PIRLS, REML, ARC,
    /// BFGS, or any recursive optimizer on the pilot.
    ///
    /// The final coefficients, smoothing parameters, and spatial geometry are
    /// always optimized on the full dataset.
    ///
    /// Set to 0 to skip the pilot geometry initializer.
    pub pilot_subsample_threshold: usize,
}

impl Default for SpatialLengthScaleOptimizationOptions {
    fn default() -> Self {
        Self {
            enabled: true,
            max_outer_iter: 80,
            rel_tol: 1e-4,
            log_step: std::f64::consts::LN_2,
            min_length_scale: 1e-3,
            max_length_scale: 1e3,
            pilot_subsample_threshold: 10_000,
        }
    }
}

impl SpatialLengthScaleOptimizationOptions {
    /// Validate the struct's invariants. Callers that construct these options
    /// from external input (CLI, config, Python API) should call this before
    /// passing the options into the fitter. Returns `Err` with a descriptive
    /// message when an invariant is violated; the fitter then panics or
    /// returns `EstimationError` at its own boundary.
    ///
    /// Invariants:
    ///   * `min_length_scale > 0`, finite
    ///   * `max_length_scale > 0`, finite
    ///   * `min_length_scale < max_length_scale`
    ///   * `rel_tol > 0`, finite
    ///   * `log_step > 0`, finite
    ///
    /// These invariants are what the downstream κ-bound and ψ-window code
    /// assumes (`-log(max_ls)` must be finite, `(min,max)` must not be
    /// inverted, etc.). Without validation, invalid options produce silent
    /// NaN-propagation inside the outer optimizer.
    pub fn validate(&self) -> Result<(), String> {
        if !self.min_length_scale.is_finite() || self.min_length_scale <= 0.0 {
            return Err(SmoothError::invalid_config(format!(
                "SpatialLengthScaleOptimizationOptions::min_length_scale must be > 0 and finite, got {}",
                self.min_length_scale
            ))
            .into());
        }
        if !self.max_length_scale.is_finite() || self.max_length_scale <= 0.0 {
            return Err(SmoothError::invalid_config(format!(
                "SpatialLengthScaleOptimizationOptions::max_length_scale must be > 0 and finite, got {}",
                self.max_length_scale
            ))
            .into());
        }
        if self.min_length_scale >= self.max_length_scale {
            return Err(SmoothError::invalid_config(format!(
                "SpatialLengthScaleOptimizationOptions requires min_length_scale < max_length_scale, got min={} max={}",
                self.min_length_scale, self.max_length_scale
            ))
            .into());
        }
        if !self.rel_tol.is_finite() || self.rel_tol <= 0.0 {
            return Err(SmoothError::invalid_config(format!(
                "SpatialLengthScaleOptimizationOptions::rel_tol must be > 0 and finite, got {}",
                self.rel_tol
            ))
            .into());
        }
        if !self.log_step.is_finite() || self.log_step <= 0.0 {
            return Err(SmoothError::invalid_config(format!(
                "SpatialLengthScaleOptimizationOptions::log_step must be > 0 and finite, got {}",
                self.log_step
            ))
            .into());
        }
        Ok(())
    }
}

#[derive(Debug, Clone)]
struct RandomEffectBlock {
    name: String,
    /// O(n) group-label vector: group_ids[i] = column index in [0, num_groups).
    /// `None` if the observation's level is not in the kept set.
    group_ids: Vec<Option<usize>>,
    num_groups: usize,
    kept_levels: Vec<u64>,
}

const BLOCK_SPARSE_ZERO_EPS: f64 = 1e-12;
const BLOCK_SPARSE_MAX_DENSITY: f64 = 0.20;

fn blocks_have_intrinsic_sparse_structure(blocks: &[DesignBlock]) -> bool {
    blocks
        .iter()
        .any(|block| matches!(block, DesignBlock::Sparse(_) | DesignBlock::RandomEffect(_)))
}

fn sparse_compatible_block_nnz(block: &DesignBlock) -> Option<usize> {
    match block {
        DesignBlock::Intercept(n) => Some(*n),
        DesignBlock::RandomEffect(op) => {
            Some(op.group_ids.iter().filter(|gid| gid.is_some()).count())
        }
        DesignBlock::Sparse(sparse) => Some(sparse.val().len()),
        DesignBlock::Dense(dense) => dense.as_dense_ref().map(|matrix| {
            matrix
                .iter()
                .filter(|&&value| value.abs() > BLOCK_SPARSE_ZERO_EPS)
                .count()
        }),
    }
}

fn try_build_sparse_design_from_blocks(
    blocks: &[DesignBlock],
) -> Result<Option<DesignMatrix>, BasisError> {
    if blocks.is_empty() {
        return Ok(None);
    }
    let nrows = blocks[0].nrows();
    let ncols: usize = blocks.iter().map(DesignBlock::ncols).sum();
    if nrows == 0 || ncols == 0 || ncols <= 32 {
        return Ok(None);
    }

    let preserve_sparse_storage = blocks_have_intrinsic_sparse_structure(blocks);
    let sparse_nnz_limit = if preserve_sparse_storage {
        usize::MAX
    } else {
        let total_cells = nrows.saturating_mul(ncols);
        ((total_cells as f64) * BLOCK_SPARSE_MAX_DENSITY).floor() as usize
    };
    let mut nnz = 0usize;
    for block in blocks {
        let block_nnz = if let Some(block_nnz) = sparse_compatible_block_nnz(block) {
            block_nnz
        } else {
            return Ok(None);
        };
        nnz = nnz.saturating_add(block_nnz);
        if nnz > sparse_nnz_limit {
            return Ok(None);
        }
    }

    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(nnz);
    let mut col_offset = 0usize;
    for block in blocks {
        match block {
            DesignBlock::Intercept(n) => {
                for row in 0..*n {
                    triplets.push(Triplet::new(row, col_offset, 1.0));
                }
            }
            DesignBlock::RandomEffect(op) => {
                for (row, group_id) in op.group_ids.iter().enumerate() {
                    if let Some(group) = group_id {
                        triplets.push(Triplet::new(row, col_offset + group, 1.0));
                    }
                }
            }
            DesignBlock::Sparse(sparse) => {
                let (symbolic, values) = sparse.parts();
                let col_ptr = symbolic.col_ptr();
                let row_idx = symbolic.row_idx();
                for col in 0..sparse.ncols() {
                    for idx in col_ptr[col]..col_ptr[col + 1] {
                        let value = values[idx];
                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
                            triplets.push(Triplet::new(row_idx[idx], col_offset + col, value));
                        }
                    }
                }
            }
            DesignBlock::Dense(dense) => {
                let matrix = dense.as_dense_ref().ok_or_else(|| {
                    BasisError::InvalidInput(
                        "sparse-compatible block assembly requires materialized dense blocks"
                            .to_string(),
                    )
                })?;
                for row in 0..matrix.nrows() {
                    for col in 0..matrix.ncols() {
                        let value = matrix[[row, col]];
                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
                            triplets.push(Triplet::new(row, col_offset + col, value));
                        }
                    }
                }
            }
        }
        col_offset += block.ncols();
    }

    let sparse = SparseColMat::try_new_from_triplets(nrows, ncols, &triplets).map_err(|_| {
        BasisError::SparseCreation("failed to assemble sparse term-collection design".to_string())
    })?;
    Ok(Some(DesignMatrix::Sparse(
        crate::matrix::SparseDesignMatrix::new(sparse),
    )))
}

fn assemble_term_collection_design_matrix(
    blocks: Vec<DesignBlock>,
) -> Result<DesignMatrix, BasisError> {
    if let Some(sparse) = try_build_sparse_design_from_blocks(&blocks)? {
        return Ok(sparse);
    }
    let block_op = BlockDesignOperator::new(blocks).map_err(|e| {
        BasisError::InvalidInput(format!("failed to build block design operator: {e}"))
    })?;
    Ok(DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(
        Arc::new(block_op),
    )))
}

/// Compute per-column standard deviations for multivariate spatial inputs (d > 1).
/// Returns `None` when d == 1 (standardization unnecessary) or when the caller
/// already supplies frozen scales (prediction path).
fn compute_spatial_input_scales(x: ArrayView2<'_, f64>) -> Option<Vec<f64>> {
    let d = x.ncols();
    if d <= 1 {
        return None;
    }
    let n = x.nrows() as f64;
    if n < 2.0 {
        return None;
    }
    let mut scales = Vec::with_capacity(d);
    for j in 0..d {
        let col = x.column(j);
        let mean = col.sum() / n;
        let var = col.iter().map(|&v| (v - mean).powi(2)).sum::<f64>() / (n - 1.0);
        scales.push(var.sqrt().max(1e-12));
    }
    Some(scales)
}

/// Apply per-column standardization to a data matrix using precomputed scales.
fn apply_input_standardization(x: &mut Array2<f64>, scales: &[f64]) {
    for j in 0..x.ncols() {
        let inv = 1.0 / scales[j];
        x.column_mut(j).mapv_inplace(|v| v * inv);
    }
}

/// Geometric mean of strictly positive scales: `(∏ s_a)^(1/d)`.
///
/// Computed via log-sum-divide to avoid overflow / underflow when d is large
/// or when individual scales are small. The Matérn / Duchon / thin-plate
/// auto-standardization paths use this to compensate the user's
/// `length_scale` so the kernel range remains expressed in *original* data
/// coordinates after per-axis division by σ_a:
///
///   ‖x_std − c_std‖ / L_eff with L_eff = L_user / σ_geom
///
/// matches `‖x − c‖ / L_user` exactly for uniform σ_a (= σ_geom) and reduces
/// to the natural anisotropic-Mahalanobis preconditioning when σ_a vary —
/// the convention σ_geom = (∏σ_a)^(1/d) preserves the kernel volume scale.
fn geometric_mean_scale(scales: &[f64]) -> f64 {
    if scales.is_empty() {
        return 1.0;
    }
    let log_mean: f64 = scales.iter().map(|&s| s.ln()).sum::<f64>() / scales.len() as f64;
    log_mean.exp()
}

fn compensate_length_scale_for_standardization(length_scale: f64, scales: &[f64]) -> f64 {
    let sigma_geom = geometric_mean_scale(scales);
    if sigma_geom > 0.0 && sigma_geom.is_finite() {
        length_scale / sigma_geom
    } else {
        length_scale
    }
}

fn compensate_optional_length_scale_for_standardization(
    length_scale: Option<f64>,
    scales: &[f64],
) -> Option<f64> {
    length_scale.map(|l| compensate_length_scale_for_standardization(l, scales))
}

fn select_columns(data: ArrayView2<'_, f64>, cols: &[usize]) -> Result<Array2<f64>, BasisError> {
    let n = data.nrows();
    let p = data.ncols();
    for &c in cols {
        if c >= p {
            crate::bail_dim_basis!("feature column {c} is out of bounds for data with {p} columns");
        }
    }
    let mut out = Array2::<f64>::zeros((n, cols.len()));
    for (j, &c) in cols.iter().enumerate() {
        out.column_mut(j).assign(&data.column(c));
    }
    Ok(out)
}

fn nonfinite_value_label(value: f64) -> &'static str {
    if value.is_nan() {
        "NaN"
    } else if value.is_sign_positive() {
        "+Inf"
    } else {
        "-Inf"
    }
}

fn validate_term_feature_column_finite(
    data: ArrayView2<'_, f64>,
    term_kind: &str,
    term_name: &str,
    feature_col: usize,
) -> Result<(), BasisError> {
    let p = data.ncols();
    if feature_col >= p {
        crate::bail_dim_basis!(
            "{term_kind} term '{term_name}' feature column {feature_col} out of bounds for {p} columns"
        );
    }
    for (row, &value) in data.column(feature_col).iter().enumerate() {
        if !value.is_finite() {
            crate::bail_invalid_basis!(
                "{term_kind} term '{term_name}' feature column {feature_col} row {row} contains non-finite value {}",
                nonfinite_value_label(value)
            );
        }
    }
    Ok(())
}

fn validate_smooth_terms_finite_inputs(
    data: ArrayView2<'_, f64>,
    terms: &[SmoothTermSpec],
) -> Result<(), BasisError> {
    for term in terms {
        for feature_col in smooth_term_feature_cols(term) {
            validate_term_feature_column_finite(data, "smooth", &term.name, feature_col)?;
        }
    }
    Ok(())
}

fn validate_term_collection_finite_inputs(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
) -> Result<(), BasisError> {
    for term in &spec.linear_terms {
        validate_term_feature_column_finite(data, "linear", &term.name, term.feature_col)?;
    }
    for term in &spec.random_effect_terms {
        validate_term_feature_column_finite(data, "random-effect", &term.name, term.feature_col)?;
    }
    validate_smooth_terms_finite_inputs(data, &spec.smooth_terms)
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct JointSpatialCenterGroupKey {
    feature_cols: Vec<usize>,
    strategy_kind: CenterStrategyKind,
    strategy_aux: usize,
    requested_num_centers: usize,
    input_scale_bits: Option<Vec<u64>>,
}

fn spatial_term_min_center_count(term: &SmoothTermSpec) -> usize {
    match &term.basis {
        SmoothBasisSpec::ThinPlate { feature_cols, .. } => feature_cols.len() + 1,
        SmoothBasisSpec::Duchon {
            feature_cols, spec, ..
        } => match spec.nullspace_order {
            crate::basis::DuchonNullspaceOrder::Zero => 1,
            crate::basis::DuchonNullspaceOrder::Linear => feature_cols.len() + 1,
            crate::basis::DuchonNullspaceOrder::Degree(degree) => {
                crate::basis::duchon_nullspace_dimension(feature_cols.len(), degree)
            }
        },
        SmoothBasisSpec::Matern { .. } => 1,
        _ => 1,
    }
}

fn spatial_term_group_key(term: &SmoothTermSpec) -> Option<JointSpatialCenterGroupKey> {
    let (feature_cols, strategy, input_scales) = match &term.basis {
        SmoothBasisSpec::ThinPlate {
            feature_cols,
            spec,
            input_scales,
        } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
        SmoothBasisSpec::Matern {
            feature_cols,
            spec,
            input_scales,
        } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
        SmoothBasisSpec::Duchon {
            feature_cols,
            spec,
            input_scales,
        } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
        _ => return None,
    };
    let strategy_kind = center_strategy_kind(strategy);
    let strategy_aux = match strategy {
        CenterStrategy::Auto(inner) => match inner.as_ref() {
            CenterStrategy::KMeans { max_iter, .. } => *max_iter,
            CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
            _ => 0,
        },
        CenterStrategy::KMeans { max_iter, .. } => *max_iter,
        CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
        _ => 0,
    };
    Some(JointSpatialCenterGroupKey {
        feature_cols: feature_cols.clone(),
        strategy_kind,
        strategy_aux,
        requested_num_centers: center_strategy_num_centers(strategy)?,
        input_scale_bits: input_scales
            .map(|values| values.iter().map(|value| value.to_bits()).collect()),
    })
}

fn spatial_term_center_strategy(term: &SmoothTermSpec) -> Option<&CenterStrategy> {
    match &term.basis {
        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.center_strategy),
        SmoothBasisSpec::Matern { spec, .. } => Some(&spec.center_strategy),
        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.center_strategy),
        _ => None,
    }
}

fn set_spatial_term_centers(
    term: &mut SmoothTermSpec,
    centers: Array2<f64>,
) -> Result<(), BasisError> {
    match &mut term.basis {
        SmoothBasisSpec::ThinPlate { spec, .. } => {
            spec.center_strategy = CenterStrategy::UserProvided(centers);
            Ok(())
        }
        SmoothBasisSpec::Matern { spec, .. } => {
            spec.center_strategy = CenterStrategy::UserProvided(centers);
            Ok(())
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            spec.center_strategy = CenterStrategy::UserProvided(centers);
            Ok(())
        }
        _ => Err(BasisError::InvalidInput(format!(
            "term '{}' does not support spatial center planning",
            term.name
        ))),
    }
}

fn standardized_spatial_term_data(
    data: ArrayView2<'_, f64>,
    term: &SmoothTermSpec,
) -> Result<Array2<f64>, BasisError> {
    let (feature_cols, input_scales) = match &term.basis {
        SmoothBasisSpec::ThinPlate {
            feature_cols,
            input_scales,
            ..
        }
        | SmoothBasisSpec::Matern {
            feature_cols,
            input_scales,
            ..
        }
        | SmoothBasisSpec::Duchon {
            feature_cols,
            input_scales,
            ..
        } => (feature_cols, input_scales.as_ref()),
        _ => {
            crate::bail_invalid_basis!("term '{}' is not a spatial smooth", term.name);
        }
    };
    let mut x = select_columns(data, feature_cols)?;
    if let Some(scales) = input_scales {
        apply_input_standardization(&mut x, scales);
    } else if let Some(scales) = compute_spatial_input_scales(x.view()) {
        apply_input_standardization(&mut x, &scales);
    }
    Ok(x)
}

fn plan_joint_spatial_centers_for_term_blocks(
    data: ArrayView2<'_, f64>,
    term_blocks: &[Vec<SmoothTermSpec>],
) -> Result<Vec<Vec<SmoothTermSpec>>, BasisError> {
    let mut planned_blocks = term_blocks.to_vec();
    let n = data.nrows();
    let mut groups: BTreeMap<JointSpatialCenterGroupKey, Vec<(usize, usize)>> = BTreeMap::new();

    for (block_idx, terms) in planned_blocks.iter().enumerate() {
        for (term_idx, term) in terms.iter().enumerate() {
            let Some(strategy) = spatial_term_center_strategy(term) else {
                continue;
            };
            if !center_strategy_is_auto(strategy) {
                continue;
            }
            let Some(group_key) = spatial_term_group_key(term) else {
                continue;
            };
            if !matches!(
                group_key.strategy_kind,
                CenterStrategyKind::EqualMass
                    | CenterStrategyKind::EqualMassCovarRepresentative
                    | CenterStrategyKind::FarthestPoint
                    | CenterStrategyKind::KMeans
            ) {
                continue;
            }
            if center_strategy_num_centers(strategy).is_none() {
                continue;
            }
            groups
                .entry(group_key)
                .or_default()
                .push((block_idx, term_idx));
        }
    }

    for (group_key, members) in groups {
        if members.len() < 2 {
            continue;
        }
        let min_required = members
            .iter()
            .map(|&(block_idx, term_idx)| {
                spatial_term_min_center_count(&planned_blocks[block_idx][term_idx])
            })
            .max()
            .unwrap_or(1);
        let joint_centers = group_key
            .requested_num_centers
            .max(min_required)
            .min(n.max(1));
        let (first_block_idx, first_term_idx) = members[0];
        let prototype = &planned_blocks[first_block_idx][first_term_idx];
        let standardized = standardized_spatial_term_data(data, prototype)?;
        let strategy = spatial_term_center_strategy(prototype).ok_or_else(|| {
            BasisError::InvalidInput(format!(
                "term '{}' lost its spatial center strategy during joint planning",
                prototype.name
            ))
        })?;
        let joint_strategy = center_strategy_with_num_centers(strategy, joint_centers)?;
        let shared_centers = select_centers_by_strategy(standardized.view(), &joint_strategy)?;
        log::info!(
            "sharing {} spatial centers across {} smooth terms over columns {:?} (requested {} centers)",
            shared_centers.nrows(),
            members.len(),
            group_key.feature_cols,
            group_key.requested_num_centers,
        );
        for (block_idx, term_idx) in members {
            set_spatial_term_centers(
                &mut planned_blocks[block_idx][term_idx],
                shared_centers.clone(),
            )?;
        }
    }

    // Sentinel auto-init: Matern and thin-plate builders write length_scale =
    // 0.0 when the user didn't pass `length_scale=...`. Replace those with a
    // data-driven initialization here so REML starts in a regime where it can
    // escape; the hard-coded 1.0 default was a basin from which ν ≥ 5/2 Matern
    // could not recover on high-frequency truths, silently collapsing the fit
    // to a near-constant prediction.
    for block in planned_blocks.iter_mut() {
        for term in block.iter_mut() {
            auto_init_length_scale_in_place(data, term);
        }
    }

    Ok(planned_blocks)
}

/// Compute a data-driven initial length scale from the per-axis range of the
/// feature columns. The heuristic `max_range / sqrt(n)` puts the kernel on
/// the wiggly side of REML's basin so the optimizer can grow it back if the
/// signal is smooth, but is small enough that high-frequency truths remain
/// reachable for smoother kernels (ν ≥ 5/2). Clamped to a tiny positive
/// floor so degenerate constant-input columns can't produce 0.
fn auto_initial_length_scale(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> f64 {
    let n = data.nrows();
    if n == 0 || feature_cols.is_empty() {
        return 1.0;
    }
    let mut max_range = 0.0_f64;
    for &c in feature_cols {
        if c >= data.ncols() {
            continue;
        }
        let col = data.column(c);
        let mut lo = f64::INFINITY;
        let mut hi = f64::NEG_INFINITY;
        for &v in col.iter() {
            if v.is_finite() {
                if v < lo {
                    lo = v;
                }
                if v > hi {
                    hi = v;
                }
            }
        }
        if hi > lo {
            let r = hi - lo;
            if r > max_range {
                max_range = r;
            }
        }
    }
    if !max_range.is_finite() || max_range <= 0.0 {
        return 1.0;
    }
    let init = max_range / (n as f64).sqrt();
    init.max(1e-6).min(max_range)
}

/// Walk a term and, if it is a Matern or thin-plate smooth whose length_scale
/// was left at the auto sentinel (`0.0`), overwrite it with
/// [`auto_initial_length_scale`].
fn auto_init_length_scale_in_place(data: ArrayView2<'_, f64>, term: &mut SmoothTermSpec) {
    auto_init_length_scale_in_basis(data, &mut term.basis);
}

/// Replace the `0.0` auto-init length-scale sentinel with a data-derived value
/// for any Matern / thin-plate kernel reachable from this basis — including the
/// inner kernel of a `by=`/factor-smooth wrapper.
///
/// `by=<factor>` and the sum-to-zero factor smooth wrap a spatial kernel inside
/// `SmoothBasisSpec::ByVariable` / `SmoothBasisSpec::FactorSumToZero` /
/// `SmoothBasisSpec::BySmooth`, so the wrapper variant is what the planner sees.
/// Without recursing into the wrapped basis the inner ThinPlate/Matern keeps the
/// `0.0` sentinel (the post-`1605b3a6e` builder default), which makes the kernel
/// distance divide by `length_scale² = 0`, producing a non-finite design at both
/// fit and predict time. Recurse so the inner kernel is initialized identically
/// to a top-level one.
fn auto_init_length_scale_in_basis(data: ArrayView2<'_, f64>, basis: &mut SmoothBasisSpec) {
    match basis {
        SmoothBasisSpec::Matern {
            feature_cols, spec, ..
        } => {
            if spec.length_scale == 0.0 {
                spec.length_scale = auto_initial_length_scale(data, feature_cols);
            }
        }
        SmoothBasisSpec::ThinPlate {
            feature_cols, spec, ..
        } => {
            if spec.length_scale == 0.0 {
                spec.length_scale = auto_initial_length_scale(data, feature_cols);
            }
        }
        SmoothBasisSpec::ByVariable { inner, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
            auto_init_length_scale_in_basis(data, inner);
        }
        SmoothBasisSpec::BySmooth { smooth, .. } => {
            auto_init_length_scale_in_basis(data, smooth);
        }
        _ => {}
    }
}

fn linear_conditioning_chunk_rows(n: usize, p: usize) -> usize {
    const TARGET_BYTES: usize = 8 * 1024 * 1024;
    const MIN_ROWS: usize = 256;
    const MAX_ROWS: usize = 65_536;
    if p == 0 {
        return n.max(1);
    }
    (TARGET_BYTES / (p * 8))
        .clamp(MIN_ROWS, MAX_ROWS)
        .min(n.max(1))
}

impl LinearFitConditioning {
    fn from_columns(design: &TermCollectionDesign, selected_cols: &[usize]) -> Self {
        const SCALE_EPS: f64 = 1e-12;
        let n = design.design.nrows();
        let p = design.design.ncols();
        let mut columns = Vec::with_capacity(selected_cols.len());
        if n == 0 || selected_cols.is_empty() {
            return Self {
                intercept_idx: design.intercept_range.start,
                columns,
            };
        }
        let chunk_rows = linear_conditioning_chunk_rows(n, p);
        // Two-pass mean/variance so operator-backed designs don't need to
        // materialize the full dense matrix. Pass 1 accumulates per-column
        // sums; pass 2 accumulates the sum of squared deviations from the
        // pass-1 mean. This matches the original `Σ (x − mean)² / n` formula
        // without the catastrophic cancellation of `E[X²] − E[X]²`.
        let mut sums = vec![0.0_f64; selected_cols.len()];
        for start in (0..n).step_by(chunk_rows) {
            let end = (start + chunk_rows).min(n);
            let chunk = design
                .design
                .try_row_chunk(start..end)
                .expect("LinearFitConditioning::from_columns row chunk failed");
            for (k, &col_idx) in selected_cols.iter().enumerate() {
                let column = chunk.column(col_idx);
                for &v in column.iter() {
                    sums[k] += v;
                }
            }
        }
        let inv_n = 1.0_f64 / n as f64;
        let means: Vec<f64> = sums.iter().map(|&s| s * inv_n).collect();
        let mut sq_devs = vec![0.0_f64; selected_cols.len()];
        for start in (0..n).step_by(chunk_rows) {
            let end = (start + chunk_rows).min(n);
            let chunk = design
                .design
                .try_row_chunk(start..end)
                .expect("LinearFitConditioning::from_columns row chunk failed");
            for (k, &col_idx) in selected_cols.iter().enumerate() {
                let mean_k = means[k];
                let column = chunk.column(col_idx);
                for &v in column.iter() {
                    let d = v - mean_k;
                    sq_devs[k] += d * d;
                }
            }
        }
        for (k, &col_idx) in selected_cols.iter().enumerate() {
            let mean = means[k];
            let var = sq_devs[k] * inv_n;
            let (mean, scale) = if var.is_finite() && var > SCALE_EPS * SCALE_EPS {
                (mean, var.sqrt())
            } else {
                // Leave nearly-constant columns untouched; centering them would collapse
                // the design column to ~0 and change the model rather than just condition it.
                (0.0, 1.0)
            };
            columns.push(LinearColumnConditioning {
                col_idx,
                mean,
                scale,
            });
        }
        Self {
            intercept_idx: design.intercept_range.start,
            columns,
        }
    }

    fn apply_to_design(&self, design: &Array2<f64>) -> Array2<f64> {
        let mut out = design.clone();
        for col in &self.columns {
            {
                let mut dst = out.column_mut(col.col_idx);
                dst -= col.mean;
            }
            if col.scale != 1.0 {
                out.column_mut(col.col_idx).mapv_inplace(|v| v / col.scale);
            }
        }
        out
    }

    fn transform_matrix_columnswith_a(&self, mat: &Array2<f64>) -> Array2<f64> {
        let mut out = mat.clone();
        let intercept = self.intercept_idx;
        for col in &self.columns {
            let intercept_col = out.column(intercept).to_owned();
            let mut target = out.column_mut(col.col_idx);
            target -= &(intercept_col * col.mean);
            if col.scale != 1.0 {
                target.mapv_inplace(|v| v / col.scale);
            }
        }
        out
    }

    fn transform_matrixrowswith_a_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
        let mut out = mat.clone();
        let intercept = self.intercept_idx;
        for col in &self.columns {
            let interceptrow = out.row(intercept).to_owned();
            let mut target = out.row_mut(col.col_idx);
            target -= &(interceptrow * col.mean);
            if col.scale != 1.0 {
                target.mapv_inplace(|v| v / col.scale);
            }
        }
        out
    }

    fn transform_matrix_columnswith_b(&self, mat: &Array2<f64>) -> Array2<f64> {
        let mut out = mat.clone();
        let intercept = self.intercept_idx;
        for col in &self.columns {
            let intercept_col = out.column(intercept).to_owned();
            let mut target = out.column_mut(col.col_idx);
            if col.mean != 0.0 {
                target += &(intercept_col * col.mean);
            }
            if col.scale != 1.0 {
                target.mapv_inplace(|v| v * col.scale);
            }
        }
        out
    }

    fn transform_matrixrowswith_b_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
        let mut out = mat.clone();
        let intercept = self.intercept_idx;
        for col in &self.columns {
            let interceptrow = out.row(intercept).to_owned();
            let mut target = out.row_mut(col.col_idx);
            if col.mean != 0.0 {
                target += &(interceptrow * col.mean);
            }
            if col.scale != 1.0 {
                target.mapv_inplace(|v| v * col.scale);
            }
        }
        out
    }

    /// Transform blockwise penalties through the conditioning.
    ///
    /// For block-local penalties whose `col_range` does not overlap with any
    /// conditioning column, the transform is identity (the conditioning only
    /// affects unpenalized linear columns). In that common case the penalty
    /// passes through unchanged, avoiding O(p²) materialization entirely.
    fn transform_blockwise_penalties_to_internal(
        &self,
        penalties: &[BlockwisePenalty],
        p: usize,
    ) -> Vec<PenaltySpec> {
        let conditioning_cols: std::collections::HashSet<usize> =
            self.columns.iter().map(|c| c.col_idx).collect();
        penalties
            .iter()
            .map(|bp| {
                let overlaps =
                    (bp.col_range.start..bp.col_range.end).any(|j| conditioning_cols.contains(&j));
                if overlaps {
                    // Rare: penalty block overlaps conditioning columns.
                    // Fall back to dense transform.
                    let global = bp.to_global(p);
                    let right = self.transform_matrix_columnswith_a(&global);
                    let transformed = self.transform_matrixrowswith_a_transpose(&right);
                    PenaltySpec::Dense(transformed)
                } else {
                    // Common: smooth penalty block doesn't touch linear columns.
                    // The conditioning is identity on this block.
                    PenaltySpec::from_blockwise(bp.clone())
                }
            })
            .collect()
    }

    fn backtransform_beta(&self, beta_internal: &Array1<f64>) -> Array1<f64> {
        let mut beta = beta_internal.clone();
        let intercept = self.intercept_idx;
        for col in &self.columns {
            beta[intercept] -= beta_internal[col.col_idx] * col.mean / col.scale;
            beta[col.col_idx] = beta_internal[col.col_idx] / col.scale;
        }
        beta
    }

    fn transform_penalized_hessian_to_original(&self, h_internal: &Array2<f64>) -> Array2<f64> {
        let right = self.transform_matrix_columnswith_b(h_internal);
        self.transform_matrixrowswith_b_transpose(&right)
    }

    fn backtransform_covariance(&self, cov_internal: &Array2<f64>) -> Array2<f64> {
        let right = self.transform_matrix_columnswith_a(cov_internal);
        self.transform_matrixrowswith_a_transpose(&right)
    }

    fn internal_bounds_for(&self, col_idx: usize, min: f64, max: f64) -> (f64, f64) {
        if let Some(col) = self.columns.iter().find(|c| c.col_idx == col_idx) {
            (min * col.scale, max * col.scale)
        } else {
            (min, max)
        }
    }
}

fn cumulative_exp(values: &Array1<f64>, sign: f64) -> Array1<f64> {
    let mut out = Array1::<f64>::zeros(values.len());
    let mut run = 0.0;
    for i in 0..values.len() {
        run += values[i].exp();
        out[i] = sign * run;
    }
    out
}

fn second_cumulative_exp(values: &Array1<f64>, sign: f64) -> Array1<f64> {
    let first = cumulative_exp(values, sign);
    let mut out = Array1::<f64>::zeros(values.len());
    let mut run = 0.0;
    for i in 0..values.len() {
        run += first[i];
        out[i] = run;
    }
    out
}

fn cumulative_sum_transform_matrix(dim: usize, order: usize, sign: f64) -> Array2<f64> {
    // The loop form `t = next^order` where `next` is the unit lower-triangular
    // matrix of ones (`next[i,j] = 1` iff `i >= j`) has the closed form
    //     T[i,j] = C(i - j + order - 1, order - 1)   for i >= j, else 0.
    // Order 0 yields the identity (consistent with the loop, which leaves
    // t = I when `order == 0`).
    let mut t = Array2::<f64>::zeros((dim, dim));
    if order == 0 {
        for i in 0..dim {
            t[[i, i]] = 1.0;
        }
    } else {
        let k = order - 1;
        for i in 0..dim {
            for j in 0..=i {
                t[[i, j]] = binomial(i - j + k, k) as f64;
            }
        }
    }
    if sign < 0.0 {
        t.mapv_inplace(|v| -v);
    }
    t
}

/// Small integer binomial coefficient C(n, k). Used to build the
/// cumulative-sum transform; n stays small (dim plus a tiny order).
fn binomial(n: usize, k: usize) -> u64 {
    if k > n {
        return 0;
    }
    let k = k.min(n - k);
    let mut num: u64 = 1;
    for i in 0..k {
        num = num * (n - i) as u64 / (i + 1) as u64;
    }
    num
}

#[cfg(test)]
mod cumulative_sum_transform_tests {
    use super::cumulative_sum_transform_matrix;
    use ndarray::Array2;

    fn reference(dim: usize, order: usize, sign: f64) -> Array2<f64> {
        let mut t = Array2::<f64>::eye(dim);
        for _ in 0..order {
            let mut next = Array2::<f64>::zeros((dim, dim));
            for i in 0..dim {
                for j in 0..=i {
                    next[[i, j]] = 1.0;
                }
            }
            t = t.dot(&next);
        }
        if sign < 0.0 {
            t.mapv_inplace(|v| -v);
        }
        t
    }

    #[test]
    fn closed_form_matches_loop() {
        for &dim in &[5usize, 10] {
            for &order in &[1usize, 2, 3, 4] {
                for &sign in &[1.0_f64, -1.0] {
                    let got = cumulative_sum_transform_matrix(dim, order, sign);
                    let want = reference(dim, order, sign);
                    for i in 0..dim {
                        for j in 0..dim {
                            assert!(
                                (got[[i, j]] - want[[i, j]]).abs() < 1e-12,
                                "mismatch at dim={dim} order={order} sign={sign} ({i},{j}): \
                                 got {} want {}",
                                got[[i, j]],
                                want[[i, j]],
                            );
                        }
                    }
                }
            }
        }
    }
}

fn shape_order_and_sign(shape: ShapeConstraint) -> Option<(usize, f64)> {
    match shape {
        ShapeConstraint::None => None,
        ShapeConstraint::MonotoneIncreasing => Some((1, 1.0)),
        ShapeConstraint::MonotoneDecreasing => Some((1, -1.0)),
        ShapeConstraint::Convex => Some((2, 1.0)),
        ShapeConstraint::Concave => Some((2, -1.0)),
    }
}

fn shape_lower_bounds_local(shape: ShapeConstraint, dim: usize) -> Option<Array1<f64>> {
    let (order, _) = shape_order_and_sign(shape)?;
    let mut lb = Array1::<f64>::from_elem(dim, f64::NEG_INFINITY);
    for j in order..dim {
        lb[j] = 0.0;
    }
    Some(lb)
}

fn shape_supports_basis(term: &SmoothTermSpec) -> bool {
    !matches!(
        term.basis,
        SmoothBasisSpec::TensorBSpline { .. } | SmoothBasisSpec::Pca { .. }
    )
}

fn freeze_raw_spatial_metadata(metadata: BasisMetadata, raw_cols: usize) -> BasisMetadata {
    match metadata {
        BasisMetadata::ThinPlate {
            centers,
            length_scale,
            periodic,
            identifiability_transform: None,
            input_scales,
            radial_reparam,
        } => BasisMetadata::ThinPlate {
            centers,
            length_scale,
            periodic,
            identifiability_transform: Some(Array2::eye(raw_cols)),
            input_scales,
            radial_reparam,
        },
        BasisMetadata::Duchon {
            centers,
            length_scale,
            periodic,
            power,
            nullspace_order,
            identifiability_transform: None,
            input_scales,
            aniso_log_scales,
            operator_collocation_points,
        } => BasisMetadata::Duchon {
            centers,
            length_scale,
            periodic,
            power,
            nullspace_order,
            identifiability_transform: Some(Array2::eye(raw_cols)),
            input_scales,
            aniso_log_scales,
            operator_collocation_points,
        },
        other => other,
    }
}

fn matern_operator_penalty_triplet_from_metadata(
    metadata: &BasisMetadata,
) -> Result<(Vec<Array2<f64>>, Vec<usize>, Vec<PenaltyInfo>), BasisError> {
    let BasisMetadata::Matern {
        centers,
        length_scale,
        periodic,
        nu,
        include_intercept,
        identifiability_transform,
        aniso_log_scales,
        input_scales,
    } = metadata
    else {
        crate::bail_invalid_basis!("Matérn operator penalties require Matérn metadata");
    };
    // The metadata records `length_scale` in *original* (un-standardized) data
    // coordinates, while `centers` live in the *standardized* coordinate frame
    // (per-axis division by `input_scales`). The realized design built the
    // kernel against those standardized centers using the σ_geom-compensated
    // effective length scale `length_scale / σ_geom`. The collocation operators
    // here are evaluated on the same standardized centers, so they must use the
    // SAME effective length scale — otherwise the penalty regularizes a
    // different RKHS range than the design lives in, leaving rough coefficient
    // directions effectively unpenalized. That mismatch is benign in 1-D
    // (no standardization) but produces a catastrophic out-of-sample blow-up in
    // d ≥ 2 where σ_geom ≠ 1 (#706).
    let penalty_length_scale = match input_scales.as_deref() {
        Some(scales) => compensate_length_scale_for_standardization(*length_scale, scales),
        None => *length_scale,
    };
    let penalty_centers = crate::basis::expand_periodic_centers(centers, periodic.as_deref())?;
    let ops = build_matern_collocation_operator_matrices(
        penalty_centers.view(),
        None,
        penalty_length_scale,
        *nu,
        *include_intercept,
        identifiability_transform.as_ref().map(|z| z.view()),
        aniso_log_scales.as_deref(),
    )?;
    // Gate the operator dials on the Matérn-ν RKHS Sobolev order m = ν + d/2:
    // mass (j=0) is always on, tension (j=1) is on for m > 1, stiffness (j=2)
    // is on for m > 2. The threshold is strict so the roughest kernel ν=1/2 in
    // d=1 (m=1, the exponential/OU H¹ process) sheds both higher operators —
    // its kernel already encodes the H¹ control, so adding an extra tension
    // dial over-smooths the oscillation it is meant to track (#707). The
    // matching gate lives at `DuchonOperatorPenaltySpec::matern_for_smoothness`.
    const ORDER_EPS: f64 = 1e-9;
    let d = penalty_centers.ncols();
    let m = nu.half_integer_value() + 0.5 * d as f64;
    let mut candidates = Vec::with_capacity(3);
    for (raw, source, min_order) in [
        (ops.d0.t().dot(&ops.d0), PenaltySource::OperatorMass, 0.0),
        (ops.d1.t().dot(&ops.d1), PenaltySource::OperatorTension, 1.0),
        (
            ops.d2.t().dot(&ops.d2),
            PenaltySource::OperatorStiffness,
            2.0,
        ),
    ] {
        if min_order > 0.0 && m <= min_order + ORDER_EPS {
            continue;
        }
        let sym = (&raw + &raw.t()) * 0.5;
        let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&sym);
        candidates.push(PenaltyCandidate {
            matrix,
            nullspace_dim_hint: 0,
            source,
            normalization_scale,
            kronecker_factors: None,
            op: None,
        });
    }
    filter_active_penalty_candidates(candidates)
}

fn shape_uses_box_reparameterization(basis: &SmoothBasisSpec) -> bool {
    matches!(basis, SmoothBasisSpec::BSpline1D { .. })
}

fn build_shape_constraint_grid_1d(x: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
    if x.is_empty() {
        crate::bail_invalid_basis!("shape-constrained smooth requires non-empty covariate values");
    }
    if x.iter().any(|v| !v.is_finite()) {
        crate::bail_invalid_basis!("shape-constrained smooth requires finite covariate values");
    }

    let mut x_sorted: Vec<f64> = x.iter().copied().collect();
    x_sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let mut x_unique: Vec<f64> = Vec::with_capacity(x_sorted.len());
    let mut last: Option<f64> = None;
    for v in x_sorted {
        let take = match last {
            None => true,
            Some(prev) => (v - prev).abs() > 1e-12 * prev.abs().max(v.abs()).max(1.0),
        };
        if take {
            x_unique.push(v);
            last = Some(v);
        }
    }
    if x_unique.len() < 2 {
        crate::bail_invalid_basis!(
            "shape-constrained smooth requires at least two unique covariate values"
        );
    }

    let min_x = x_unique[0];
    let max_x = *x_unique
        .last()
        .expect("x_unique has at least two elements by construction");
    if (max_x - min_x).abs() <= 1e-12 {
        crate::bail_invalid_basis!(
            "shape-constrained smooth requires non-degenerate covariate range"
        );
    }

    let target_points = x_unique.len().clamp(96, 320);
    let mut grid = Array1::<f64>::zeros(target_points);
    let denom = (target_points - 1) as f64;
    for i in 0..target_points {
        let t = i as f64 / denom;
        grid[i] = min_x + t * (max_x - min_x);
    }
    Ok(grid)
}

fn build_shape_constraint_design_1d(
    data: ArrayView2<'_, f64>,
    term: &SmoothTermSpec,
    metadata: &BasisMetadata,
    axis_col: usize,
) -> Result<(Array1<f64>, Array2<f64>), BasisError> {
    let x_grid = build_shape_constraint_grid_1d(data.column(axis_col))?;
    let grid_2d = x_grid
        .clone()
        .into_shape_with_order((x_grid.len(), 1))
        .map_err(|e| {
            BasisError::InvalidInput(format!(
                "failed to construct 1D shape grid matrix for term '{}': {e}",
                term.name
            ))
        })?;

    let design = match (&term.basis, metadata) {
        (
            SmoothBasisSpec::BSpline1D { spec, .. },
            BasisMetadata::BSpline1D {
                knots,
                identifiability_transform,
                periodic,
                degree: meta_degree,
                ..
            },
        ) => {
            // Issue #340: predict against the metadata-recorded effective
            // degree so fit-time auto-shrink (cubic → linear for small n) is
            // honoured at prediction time too.
            let effective_degree = meta_degree.unwrap_or(spec.degree);
            let evalspec = BSplineBasisSpec {
                degree: effective_degree,
                penalty_order: spec.penalty_order,
                knotspec: periodic
                    .map(
                        |(domain_start, period, num_basis)| BSplineKnotSpec::PeriodicUniform {
                            data_range: (domain_start, domain_start + period),
                            num_basis,
                        },
                    )
                    .unwrap_or_else(|| BSplineKnotSpec::Provided(knots.clone())),
                double_penalty: false,
                identifiability: identifiability_transform
                    .as_ref()
                    .map(|z| BSplineIdentifiability::FrozenTransform {
                        transform: z.clone(),
                    })
                    .unwrap_or(BSplineIdentifiability::None),
                boundary: spec.boundary.clone(),
                boundary_conditions: BSplineBoundaryConditions::default(),
            };
            build_bspline_basis_1d(x_grid.view(), &evalspec)?
                .design
                .to_dense()
        }
        (
            SmoothBasisSpec::ThinPlate { .. },
            BasisMetadata::ThinPlate {
                centers,
                length_scale,
                identifiability_transform,
                radial_reparam,
                ..
            },
        ) => {
            let evalspec = ThinPlateBasisSpec {
                periodic: None,
                center_strategy: crate::basis::CenterStrategy::UserProvided(centers.clone()),
                length_scale: *length_scale,
                double_penalty: false,
                identifiability: identifiability_transform
                    .as_ref()
                    .map(|z| SpatialIdentifiability::FrozenTransform {
                        transform: z.clone(),
                    })
                    .unwrap_or(SpatialIdentifiability::None),
                radial_reparam: radial_reparam.clone(),
            };
            build_thin_plate_basis(grid_2d.view(), &evalspec)?
                .design
                .to_dense()
        }
        (
            SmoothBasisSpec::Matern { .. },
            BasisMetadata::Matern {
                centers,
                length_scale,
                nu,
                include_intercept,
                identifiability_transform,
                aniso_log_scales,
                ..
            },
        ) => {
            let ident = identifiability_transform
                .as_ref()
                .map(|z| MaternIdentifiability::FrozenTransform {
                    transform: z.clone(),
                })
                .unwrap_or(MaternIdentifiability::None);
            let evalspec = MaternBasisSpec {
                periodic: None,
                center_strategy: crate::basis::CenterStrategy::UserProvided(centers.clone()),
                length_scale: *length_scale,
                nu: *nu,
                include_intercept: *include_intercept,
                double_penalty: false,
                identifiability: ident,
                aniso_log_scales: aniso_log_scales.clone(),
            };
            build_matern_basis(grid_2d.view(), &evalspec)?
                .design
                .to_dense()
        }
        (
            SmoothBasisSpec::Duchon { spec, .. },
            BasisMetadata::Duchon {
                centers,
                length_scale,
                power,
                nullspace_order,
                identifiability_transform,
                aniso_log_scales,
                ..
            },
        ) => {
            let evalspec = DuchonBasisSpec {
                periodic: None,
                center_strategy: crate::basis::CenterStrategy::UserProvided(centers.clone()),
                length_scale: *length_scale,
                power: *power,
                nullspace_order: *nullspace_order,
                identifiability: identifiability_transform
                    .as_ref()
                    .map(|z| SpatialIdentifiability::FrozenTransform {
                        transform: z.clone(),
                    })
                    .unwrap_or_else(|| spec.identifiability.clone()),
                aniso_log_scales: aniso_log_scales.clone(),
                operator_penalties: spec.operator_penalties.clone(),
                boundary: spec.boundary.clone(),
            };
            build_duchon_basis(grid_2d.view(), &evalspec)?
                .design
                .to_dense()
        }
        _ => {
            crate::bail_invalid_basis!(
                "shape-constraint grid reconstruction metadata mismatch for term '{}'",
                term.name
            );
        }
    };

    Ok((x_grid, design))
}

fn build_shape_linear_constraints_1d(
    x: ArrayView1<'_, f64>,
    design_local: ArrayView2<'_, f64>,
    shape: ShapeConstraint,
) -> Result<Option<LinearInequalityConstraints>, BasisError> {
    let Some((order, sign)) = shape_order_and_sign(shape) else {
        return Ok(None);
    };
    let n = x.len();
    let p = design_local.ncols();
    if n == 0 || p == 0 {
        return Ok(None);
    }
    if x.iter().any(|v| !v.is_finite()) {
        crate::bail_invalid_basis!("shape-constrained smooth requires finite covariate values");
    }

    let mut idx: Vec<usize> = (0..n).collect();
    idx.sort_by(|&i, &j| x[i].partial_cmp(&x[j]).unwrap_or(std::cmp::Ordering::Equal));

    let x_scale = x.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs())).max(1.0);
    let x_tol = 1e-12 * x_scale;
    let mut collapsedrows: Vec<Array1<f64>> = Vec::new();
    let mut group_sum = Array1::<f64>::zeros(p);
    let mut group_count = 0usize;
    let mut last_x: Option<f64> = None;
    for &r in &idx {
        let xr = x[r];
        let start_new = match last_x {
            None => false,
            Some(prev) => (xr - prev).abs() > x_tol,
        };
        if start_new {
            if group_count > 0 {
                collapsedrows.push(group_sum.mapv(|v| v / group_count as f64));
            }
            group_sum.fill(0.0);
            group_count = 0;
        }
        group_sum.scaled_add(1.0, &design_local.row(r));
        group_count += 1;
        last_x = Some(xr);
    }
    if group_count > 0 {
        collapsedrows.push(group_sum.mapv(|v| v / group_count as f64));
    }

    let m = collapsedrows.len();
    if m <= order {
        crate::bail_invalid_basis!(
            "shape-constrained smooth requires at least {} unique covariate locations; found {}",
            order + 1,
            m
        );
    }

    let q_raw = m - order;
    let mut arows: Vec<Array1<f64>> = Vec::with_capacity(q_raw);
    for i in 0..q_raw {
        let row = if order == 1 {
            &collapsedrows[i + 1] - &collapsedrows[i]
        } else {
            &collapsedrows[i + 2] - &collapsedrows[i + 1].mapv(|v| 2.0 * v) + &collapsedrows[i]
        };
        let mut row_signed = row;
        if sign < 0.0 {
            row_signed.mapv_inplace(|v| -v);
        }
        let norm = row_signed.dot(&row_signed).sqrt();
        if norm > 1e-12 {
            arows.push(row_signed);
        }
    }
    if arows.is_empty() {
        return Ok(None);
    }

    let mut a = Array2::<f64>::zeros((arows.len(), p));
    for (i, row) in arows.iter().enumerate() {
        a.row_mut(i).assign(row);
    }
    let b = Array1::<f64>::zeros(a.nrows());
    Ok(Some(LinearInequalityConstraints { a, b }))
}

fn linear_constraints_from_lower_bounds_global(
    lower_bounds: &Array1<f64>,
) -> Option<LinearInequalityConstraints> {
    LinearInequalityConstraints::from_per_coordinate_lower_bounds(lower_bounds)
}

fn merge_linear_constraints_global(
    first: Option<LinearInequalityConstraints>,
    second: Option<LinearInequalityConstraints>,
) -> Option<LinearInequalityConstraints> {
    match (first, second) {
        (None, None) => None,
        (Some(c), None) | (None, Some(c)) => Some(c),
        (Some(a), Some(b)) => {
            if a.a.ncols() != b.a.ncols() {
                return None;
            }
            let m1 = a.a.nrows();
            let m2 = b.a.nrows();
            let p = a.a.ncols();
            let mut mat = Array2::<f64>::zeros((m1 + m2, p));
            mat.slice_mut(s![0..m1, ..]).assign(&a.a);
            mat.slice_mut(s![m1..(m1 + m2), ..]).assign(&b.a);
            let mut rhs = Array1::<f64>::zeros(m1 + m2);
            rhs.slice_mut(s![0..m1]).assign(&a.b);
            rhs.slice_mut(s![m1..(m1 + m2)]).assign(&b.b);
            Some(LinearInequalityConstraints { a: mat, b: rhs })
        }
    }
}

fn normalize_penalty_in_constrained_space(matrix: &Array2<f64>) -> (Array2<f64>, f64) {
    // Constrained-space normalization:
    //   c = ||S_con||_F,  S_tilde = S_con / c.
    // This is the only normalization coherent with a REML objective that is
    // evaluated entirely in constrained coordinates.
    let matrix = (matrix + &matrix.t().to_owned()) * 0.5;
    // Clamp noise-floor negative eigenvalues so β'Sβ is non-negative as a contract, not just in exact arithmetic.
    let matrix = crate::terms::basis::project_penalty_to_psd_cone(&matrix);
    let c = matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
    if c.is_finite() && c > 0.0 {
        (matrix.mapv(|v| v / c), c)
    } else {
        (matrix, 1.0)
    }
}

fn build_periodic_fourier_margin(
    x: ArrayView1<'_, f64>,
    period: f64,
    requested_cols: usize,
    penalty_order: usize,
) -> Result<(Array2<f64>, Array2<f64>, Array1<f64>), BasisError> {
    if !period.is_finite() || period <= 0.0 {
        crate::bail_invalid_basis!(
            "periodic tensor margin requires finite positive period, got {period}"
        );
    }
    let q = requested_cols.max(3);
    let harmonics = q / 2;
    let has_nyquist_cos = q.is_multiple_of(2);
    let mut basis = Array2::<f64>::zeros((x.len(), q));
    basis.column_mut(0).fill(1.0);
    for (i, &xi) in x.iter().enumerate() {
        let angle = 2.0 * std::f64::consts::PI * xi / period;
        let mut col = 1usize;
        for h in 1..=harmonics {
            if col >= q {
                break;
            }
            basis[[i, col]] = (h as f64 * angle).cos();
            col += 1;
            if col >= q {
                break;
            }
            basis[[i, col]] = (h as f64 * angle).sin();
            col += 1;
        }
        if has_nyquist_cos && q > 1 {
            basis[[i, q - 1]] = (harmonics as f64 * angle).cos();
        }
    }
    let mut penalty = Array2::<f64>::zeros((q, q));
    let order = penalty_order.max(1) as i32;
    let mut col = 1usize;
    for h in 1..=harmonics {
        let w = (h as f64).powi(2 * order);
        if col < q {
            penalty[[col, col]] = w;
            col += 1;
        }
        if col < q {
            penalty[[col, col]] = w;
            col += 1;
        }
    }
    if has_nyquist_cos && q > 1 {
        penalty[[q - 1, q - 1]] = (harmonics as f64).powi(2 * order);
    }
    // The Fourier margin has no B-spline knots; this vector is a placeholder
    // that downstream code (the tensor freeze) treats as carrying the periodic
    // control-site count in its *length* and the domain start in `[0]`. Its
    // length MUST equal the basis column count `q` (= `basis.ncols()`): the
    // freeze records `num_basis = knots.len()` and the predict-time rebuild
    // re-derives `q` columns from it, so a `q + 1`-length vector reconstructs
    // one extra column per periodic axis and breaks the frozen identifiability
    // transform (issue #498).
    let knots = Array1::linspace(0.0, period, q);
    Ok((basis, penalty, knots))
}

fn tensor_product_design_from_sparse_marginals(
    marginal_sparse: &[&SparseColMat<usize, f64>],
) -> Result<SparseColMat<usize, f64>, BasisError> {
    if marginal_sparse.is_empty() {
        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
    }
    let n = marginal_sparse[0].nrows();
    for (i, m) in marginal_sparse.iter().enumerate().skip(1) {
        if m.nrows() != n {
            crate::bail_dim_basis!(
                "tensor sparse marginal row mismatch at dim {i}: expected {n}, got {}",
                m.nrows()
            );
        }
    }
    let dims: Vec<usize> = marginal_sparse.iter().map(|m| m.ncols()).collect();
    let total_cols = dims.iter().try_fold(1usize, |acc, &q| {
        acc.checked_mul(q)
            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
    })?;
    let mut strides = vec![1usize; dims.len()];
    for d in (0..dims.len().saturating_sub(1)).rev() {
        strides[d] = strides[d + 1]
            .checked_mul(dims[d + 1])
            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))?;
    }

    use faer::sparse::SparseRowMat;
    let csrs: Vec<SparseRowMat<usize, f64>> = marginal_sparse
        .iter()
        .enumerate()
        .map(|(d, m)| {
            m.as_ref().to_row_major().map_err(|e| {
                BasisError::SparseCreation(format!(
                    "tensor sparse marginal {d} CSR conversion failed: {e:?}"
                ))
            })
        })
        .collect::<Result<Vec<_>, _>>()?;
    let row_ptrs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().row_ptr()).collect();
    let col_idxs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().col_idx()).collect();
    let vals: Vec<&[f64]> = csrs.iter().map(|c| c.val()).collect();

    use rayon::prelude::*;
    const CHUNK: usize = 1024;
    let num_chunks = n.div_ceil(CHUNK);
    let per_chunk: Vec<Vec<Triplet<usize, usize, f64>>> = (0..num_chunks)
        .into_par_iter()
        .map(|chunk_idx| {
            let row_start = chunk_idx * CHUNK;
            let row_end = (row_start + CHUNK).min(n);
            let mut chunk_triplets = Vec::<Triplet<usize, usize, f64>>::new();
            let mut cur_cols = Vec::<usize>::with_capacity(64);
            let mut cur_vals = Vec::<f64>::with_capacity(64);
            let mut next_cols = Vec::<usize>::with_capacity(64);
            let mut next_vals = Vec::<f64>::with_capacity(64);
            for i in row_start..row_end {
                cur_cols.clear();
                cur_vals.clear();
                cur_cols.push(0);
                cur_vals.push(1.0);
                let mut row_is_zero = false;
                for d in 0..dims.len() {
                    let row_start_d = row_ptrs[d][i];
                    let row_end_d = row_ptrs[d][i + 1];
                    if row_start_d == row_end_d {
                        row_is_zero = true;
                        break;
                    }
                    let stride = strides[d];
                    next_cols.clear();
                    next_vals.clear();
                    next_cols.reserve(cur_cols.len() * (row_end_d - row_start_d));
                    next_vals.reserve(cur_vals.len() * (row_end_d - row_start_d));
                    for (&prev_col, &prev_val) in cur_cols.iter().zip(cur_vals.iter()) {
                        for ptr in row_start_d..row_end_d {
                            let cj = col_idxs[d][ptr];
                            let vj = vals[d][ptr];
                            next_cols.push(prev_col + cj * stride);
                            next_vals.push(prev_val * vj);
                        }
                    }
                    std::mem::swap(&mut cur_cols, &mut next_cols);
                    std::mem::swap(&mut cur_vals, &mut next_vals);
                }
                if row_is_zero {
                    continue;
                }
                for (&col, &val) in cur_cols.iter().zip(cur_vals.iter()) {
                    chunk_triplets.push(Triplet::new(i, col, val));
                }
            }
            chunk_triplets
        })
        .collect();
    let total_nnz: usize = per_chunk.iter().map(Vec::len).sum();
    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(total_nnz);
    for chunk in per_chunk {
        triplets.extend(chunk);
    }
    SparseColMat::try_new_from_triplets(n, total_cols, &triplets).map_err(|e| {
        BasisError::SparseCreation(format!(
            "failed to assemble sparse tensor product design: {e:?}"
        ))
    })
}

fn build_tensor_bspline_basis(
    data: ArrayView2<'_, f64>,
    feature_cols: &[usize],
    spec: &TensorBSplineSpec,
) -> Result<BasisBuildResult, BasisError> {
    if feature_cols.is_empty() {
        crate::bail_invalid_basis!("TensorBSpline requires at least one feature column");
    }
    if feature_cols.len() != spec.marginalspecs.len() {
        crate::bail_dim_basis!(
            "TensorBSpline feature/spec mismatch: feature_cols={}, marginalspecs={}",
            feature_cols.len(),
            spec.marginalspecs.len()
        );
    }
    if !spec.periods.is_empty() && spec.periods.len() != feature_cols.len() {
        crate::bail_dim_basis!(
            "TensorBSpline periods length {} does not match feature count {}",
            spec.periods.len(),
            feature_cols.len()
        );
    }
    let p = data.ncols();
    for &c in feature_cols {
        if c >= p {
            crate::bail_dim_basis!(
                "tensor feature column {c} is out of bounds for data with {p} columns"
            );
        }
    }

    let mut marginal_knots = Vec::<Array1<f64>>::with_capacity(feature_cols.len());
    let mut marginal_degrees = Vec::<usize>::with_capacity(feature_cols.len());
    let mut marginalnum_basis = Vec::<usize>::with_capacity(feature_cols.len());
    let mut marginal_penalties = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
    let mut marginal_designs = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
    // Per-margin effective period: either user-set via `spec.periods` (forcing
    // the Fourier path) or implied by a `PeriodicUniform` marginal knotspec
    // (which the 1D B-spline builder already realizes as a periodic basis).
    // Captured here so freeze→reload round-trips both routes back to a
    // `PeriodicUniform` marginal knotspec; otherwise a `PeriodicUniform`
    // margin specified without `spec.periods` would freeze as a plain
    // `Provided(knots)` open spline and lose its wrap-around at predict time.
    let mut marginal_effective_periods = Vec::<Option<f64>>::with_capacity(feature_cols.len());
    // Per-marginal sparse representation, populated when the 1D builder returned
    // a `DesignMatrix::Sparse`. Used to assemble the Khatri-Rao tensor product
    // sparsely (only ∏(degree+1) nonzeros per row) instead of densifying to
    // shape (n, ∏ q_j) up front. When any marginal lacks a sparse form (e.g.
    // periodic B-splines currently realize a dense Array2), we fall back to the
    // existing dense Khatri-Rao path.
    let mut marginal_sparse =
        Vec::<Option<SparseColMat<usize, f64>>>::with_capacity(feature_cols.len());

    // Reuse the robust 1D builder to ensure the same knot validation and
    // marginal difference-penalty construction as standalone smooth terms.
    for (dim, (&col, marginalspec)) in feature_cols
        .iter()
        .zip(spec.marginalspecs.iter())
        .enumerate()
    {
        // Tensor basis uses raw marginal knot-product columns. Applying 1D
        // identifiability constraints here would change marginal penalty sizes
        // without changing the tensor design construction, causing dimension
        // mismatch. Keep marginal builders unconstrained at this stage.
        if let Some(period) = spec.periods.get(dim).and_then(|p| *p) {
            let requested_cols = match marginalspec.knotspec {
                BSplineKnotSpec::Generate {
                    num_internal_knots, ..
                } => num_internal_knots + marginalspec.degree + 1,
                BSplineKnotSpec::Provided(ref knots) => {
                    knots.len().saturating_sub(marginalspec.degree + 1)
                }
                BSplineKnotSpec::Automatic {
                    num_internal_knots, ..
                } => num_internal_knots.unwrap_or(8) + marginalspec.degree + 1,
                BSplineKnotSpec::PeriodicUniform { num_basis, .. } => num_basis,
            };
            let (basis, penalty, knots) = build_periodic_fourier_margin(
                data.column(col),
                period,
                requested_cols,
                marginalspec.penalty_order,
            )?;
            marginal_knots.push(knots);
            marginal_degrees.push(marginalspec.degree);
            marginalnum_basis.push(basis.ncols());
            marginal_designs.push(basis);
            marginal_penalties.push(penalty);
            // Periodic Fourier margins are realized densely; no sparse form
            // is available, so record `None` and force the dense fall-back
            // for the tensor product if any dimension is periodic.
            marginal_sparse.push(None);
            marginal_effective_periods.push(Some(period));
        } else {
            let mut marginal_unconstrained = marginalspec.clone();
            marginal_unconstrained.identifiability = BSplineIdentifiability::None;
            let built = build_bspline_basis_1d(data.column(col), &marginal_unconstrained)?;
            let knots = match built.metadata {
                BasisMetadata::BSpline1D { knots, .. } => knots,
                _ => {
                    crate::bail_invalid_basis!(
                        "internal TensorBSpline error at dim {dim}: expected BSpline1D metadata"
                    );
                }
            };
            marginal_knots.push(knots);
            marginal_degrees.push(marginalspec.degree);
            marginalnum_basis.push(built.design.ncols());
            // Capture the sparse representation of this marginal (when the
            // 1D builder produced one) before densifying for the dense
            // marginal cache used by `tensor_product_design_from_marginals`
            // and `TensorProductDesignOperator`.
            let sparse_view: Option<SparseColMat<usize, f64>> =
                built.design.as_sparse().map(|sd| {
                    let inner: &SparseColMat<usize, f64> = sd;
                    inner.clone()
                });
            marginal_sparse.push(sparse_view);
            marginal_designs.push(built.design.to_dense());
            marginal_penalties.push(
                built
                    .penalties
                    .first()
                    .ok_or_else(|| {
                        BasisError::InvalidInput(format!(
                            "internal TensorBSpline error at dim {dim}: missing marginal penalty"
                        ))
                    })?
                    .clone(),
            );
            built.nullspace_dims.first().ok_or_else(|| {
                BasisError::InvalidInput(format!(
                    "internal TensorBSpline error at dim {dim}: missing marginal nullspace dim"
                ))
            })?;
            // A `PeriodicUniform` marginal knotspec implies the margin is
            // wrap-around: the 1D builder already realized it as a periodic
            // basis, so the tensor product inherits that periodicity. Record
            // the period derived from the knotspec's data range so freeze
            // restores `PeriodicUniform` on the marginal — otherwise the
            // round-trip downgrades it to `Provided(knots)` (an open spline)
            // and predict-time wraps disappear.
            let implied_period = match marginalspec.knotspec {
                BSplineKnotSpec::PeriodicUniform { data_range, .. } => {
                    Some(data_range.1 - data_range.0)
                }
                _ => None,
            };
            marginal_effective_periods.push(implied_period);
        }
    }

    let total_cols: usize = marginalnum_basis.iter().product();
    let mut dense_design = (!matches!(spec.identifiability, TensorBSplineIdentifiability::None))
        .then(|| tensor_product_design_from_marginals(&marginal_designs))
        .transpose()?;
    let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
        marginal_penalties.len() + if spec.double_penalty { 1 } else { 0 },
    );

    // Tensor-product smoothing parameters are one-per-margin.  Therefore the
    // physical penalty attached to a margin must be normalized in that margin's
    // own working coordinates before it is embedded in the full tensor product.
    // Normalizing only the already-Kroneckered matrix would fold arbitrary
    // dimension-dependent identity factors into the margin's lambda and would
    // make anisotropic REML/LAML smoothing depend on the other margins' basis
    // sizes rather than on the marginal roughness operator itself.
    let normalized_marginal_penalties: Vec<(Array2<f64>, f64)> = marginal_penalties
        .iter()
        .map(normalize_penalty_in_constrained_space)
        .collect();
    let mut kronecker_marginal_penalties =
        Vec::<Array2<f64>>::with_capacity(normalized_marginal_penalties.len());

    // Accumulate the Kronecker-sum of the per-margin penalties, `Σ_dim S_dim`,
    // whose null space is exactly the *joint* null space of all marginal
    // penalties — the tensor of the marginal polynomial null spaces, i.e. the
    // bilinear (low-order) directions that no marginal roughness operator
    // touches. The tensor double penalty (below) shrinks only this joint null,
    // never the already-penalized interaction range.
    let mut marginal_kron_sum = Array2::<f64>::zeros((total_cols, total_cols));

    for dim in 0..normalized_marginal_penalties.len() {
        let mut s_dim = Array2::<f64>::eye(1);
        let mut factors = Vec::<Array2<f64>>::with_capacity(marginalnum_basis.len());
        for (j, &qj) in marginalnum_basis.iter().enumerate() {
            let factor = if j == dim {
                normalized_marginal_penalties[j].0.clone()
            } else {
                Array2::<f64>::eye(qj)
            };
            factors.push(factor.clone());
            s_dim = kronecker_product(&s_dim, &factor);
        }
        if dim == kronecker_marginal_penalties.len() {
            kronecker_marginal_penalties.push(normalized_marginal_penalties[dim].0.clone());
        }
        marginal_kron_sum += &s_dim;

        candidates.push(PenaltyCandidate {
            matrix: s_dim,
            nullspace_dim_hint: 0,
            source: PenaltySource::TensorMarginal { dim },
            normalization_scale: normalized_marginal_penalties[dim].1,
            kronecker_factors: Some(factors),
            op: None,
        });
    }

    if spec.double_penalty {
        // mgcv `select=TRUE` semantics, mirrored from the 1D double penalty:
        // the extra penalty shrinks ONLY the directions left unpenalized by the
        // primary (marginal) penalties — here the joint null space of the
        // Kronecker-sum `Σ_dim S_dim` (the bilinear tensor null `Z₀ ⊗ Z₁`).
        // A full identity ridge `I` over *all* tensor coefficients (the prior
        // behavior) instead penalizes the already-penalized interaction range
        // as well, and REML/LAML then places positive weight on it and
        // systematically over-smooths the recovered surface; mgcv's `te`/`ti`
        // carry no such global ridge. Penalizing the null subspace alone keeps
        // the interaction range governed solely by the per-margin λ's.
        if let Some(shrink) =
            crate::terms::basis::build_nullspace_shrinkage_penalty(&marginal_kron_sum)?
        {
            let (matrix, normalization_scale) =
                normalize_penalty_in_constrained_space(&shrink.sym_penalty);
            candidates.push(PenaltyCandidate {
                matrix,
                nullspace_dim_hint: 0,
                source: PenaltySource::TensorGlobalRidge,
                normalization_scale,
                kronecker_factors: None,
                op: None,
            });
        }
    }

    let z_opt = match &spec.identifiability {
        TensorBSplineIdentifiability::None => None,
        TensorBSplineIdentifiability::SumToZero => {
            if total_cols < 2 {
                crate::bail_invalid_basis!(
                    "TensorBSpline requires at least 2 basis coefficients to enforce sum-to-zero identifiability"
                );
            }
            let dense_design_ref = dense_design.as_ref().ok_or_else(|| {
                BasisError::InvalidInput(
                    "tensor sum-to-zero identifiability requires a realized basis".to_string(),
                )
            })?;
            let (_, z) = apply_sum_to_zero_constraint(dense_design_ref.view(), None)?;
            Some(z)
        }
        TensorBSplineIdentifiability::MarginalSumToZero => {
            // `ti(...)`: drop the marginal main effects by centering every
            // margin independently, then form the tensor product of the
            // centered margins. Concretely, each margin `j` is reparameterized
            // by its own sum-to-zero null basis `Z_j` (so the constant — i.e.
            // the marginal intercept — is removed from that axis), and the
            // combined reparameterization is the Kronecker product
            // `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}`. Applying `Z` to the full-tensor
            // design `B = B₀ ⊗ … ⊗ B_{d-1}` yields `B Z = (B₀ Z₀) ⊗ … ⊗
            // (B_{d-1} Z_{d-1})`, the tensor product of the centered margins,
            // which by construction contains no pure main effect.
            if marginal_designs.len() < 2 {
                crate::bail_invalid_basis!(
                    "tensor interaction (ti) identifiability requires at least 2 margins"
                );
            }
            let mut z = Array2::<f64>::eye(1);
            for (dim, marginal) in marginal_designs.iter().enumerate() {
                if marginal.ncols() < 2 {
                    crate::bail_invalid_basis!(
                        "tensor interaction (ti) margin {dim} has fewer than 2 basis functions; \
                         cannot remove its marginal main effect"
                    );
                }
                let (_, z_dim) = apply_sum_to_zero_constraint(marginal.view(), None)?;
                z = kronecker_product(&z, &z_dim);
            }
            Some(z)
        }
        TensorBSplineIdentifiability::FrozenTransform { transform } => {
            if transform.nrows() != total_cols {
                crate::bail_dim_basis!(
                    "frozen tensor identifiability transform mismatch: design has {} columns but transform has {} rows",
                    total_cols,
                    transform.nrows()
                );
            }
            Some(transform.clone())
        }
    };

    if let Some(z) = z_opt.as_ref() {
        let dense = dense_design.as_mut().ok_or_else(|| {
            BasisError::InvalidInput(
                "tensor identifiability transform requires a realized basis".to_string(),
            )
        })?;
        *dense = dense.dot(z);
        candidates = candidates
            .into_iter()
            .map(|candidate| -> Result<PenaltyCandidate, BasisError> {
                let zt_s = fast_atb(z, &candidate.matrix);
                let matrix = fast_ab(&zt_s, z);
                let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
                let preserve_margin_scale =
                    matches!(&candidate.source, PenaltySource::TensorMarginal { .. });
                let (matrix, normalization_scale) = if preserve_margin_scale {
                    (matrix.mapv(|v| v * c_new), candidate.normalization_scale)
                } else {
                    (matrix, candidate.normalization_scale * c_new)
                };
                Ok(PenaltyCandidate {
                    nullspace_dim_hint: candidate.nullspace_dim_hint,
                    matrix,
                    source: candidate.source,
                    normalization_scale,
                    kronecker_factors: candidate.kronecker_factors.clone(),
                    op: candidate.op.clone(),
                })
            })
            .collect::<Result<Vec<_>, _>>()?;
    }

    let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
        filter_active_penalty_candidates_with_ops(candidates)?;
    let identifiability_is_none =
        matches!(spec.identifiability, TensorBSplineIdentifiability::None);
    // All marginals expose a sparse representation iff each `marginal_sparse`
    // slot is `Some(...)`. Currently this is true when every marginal is a
    // free-boundary, non-periodic 1D B-spline returned as
    // `DesignMatrix::Sparse` from `build_bspline_basis_1d`. Periodic B-splines
    // and other dense-only marginals leave a `None` and trigger the fall-back
    // path. Identifiability transforms (`SumToZero`, `FrozenTransform`) make
    // the tensor design dense in general, so we also gate on that.
    let all_marginals_sparse = marginal_sparse.iter().all(Option::is_some);
    let design = if let Some(dense_design) = dense_design {
        DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(dense_design))
    } else if identifiability_is_none && all_marginals_sparse {
        // Sparse Khatri-Rao path: assemble the (n, ∏ q_j) tensor product
        // directly as a SparseColMat, preserving the ∏(degree_j+1) nonzero
        // structure per row instead of densifying to ∏ q_j columns. This is
        // mathematically identical to `tensor_product_design_from_marginals`
        // applied to the corresponding dense marginals.
        let sparse_marginals: Vec<&SparseColMat<usize, f64>> = marginal_sparse
            .iter()
            .map(|m| m.as_ref().expect("all_marginals_sparse just verified"))
            .collect();
        let sparse_design = tensor_product_design_from_sparse_marginals(&sparse_marginals)?;
        DesignMatrix::Sparse(crate::matrix::SparseDesignMatrix::new(sparse_design))
    } else {
        let marginals: Vec<Arc<Array2<f64>>> = marginal_designs
            .iter()
            .map(|m| Arc::new(m.clone()))
            .collect();
        let op = TensorProductDesignOperator::new(marginals).map_err(|e| {
            BasisError::InvalidInput(format!("TensorProductDesignOperator build failed: {e}"))
        })?;
        DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(Arc::new(op)))
    };

    Ok(BasisBuildResult {
        design,
        penalties,
        nullspace_dims,
        penaltyinfo,
        ops,
        null_eigenvectors,
        joint_null_rotation: None,
        metadata: BasisMetadata::TensorBSpline {
            feature_cols: feature_cols.to_vec(),
            knots: marginal_knots,
            degrees: marginal_degrees,
            // Prefer the per-margin effective period derived in the loop —
            // it captures both the explicit `spec.periods` route and the
            // implied period from a `PeriodicUniform` marginal knotspec.
            // Falling back to `spec.periods` when populated keeps any
            // user-supplied explicit period authoritative even if the
            // marginal knotspec carried no periodicity hint.
            periods: marginal_effective_periods,
            identifiability_transform: z_opt,
        },
        kronecker_factored: if matches!(spec.identifiability, TensorBSplineIdentifiability::None) {
            Some(KroneckerFactoredBasis {
                marginal_designs,
                marginal_penalties: kronecker_marginal_penalties,
                marginal_dims: marginalnum_basis.clone(),
                has_double_penalty: spec.double_penalty,
            })
        } else {
            None
        },
    })
}

fn tensor_product_design_from_marginals(
    marginal_designs: &[Array2<f64>],
) -> Result<Array2<f64>, BasisError> {
    if marginal_designs.is_empty() {
        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
    }
    let n = marginal_designs[0].nrows();
    for (i, b) in marginal_designs.iter().enumerate().skip(1) {
        if b.nrows() != n {
            crate::bail_dim_basis!(
                "tensor marginal row mismatch at dim {i}: expected {n}, got {}",
                b.nrows()
            );
        }
    }
    let total_cols = marginal_designs.iter().try_fold(1usize, |acc, b| {
        acc.checked_mul(b.ncols())
            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
    })?;
    // Tensor-product Khatri-Rao: design[i, j] = Π_d marginal_d[i, j_d]
    // where j is the multi-index (j_1, ..., j_D) flattened. Independent
    // across rows; parallelize row chunks and fill the pre-allocated
    // contiguous Array2 in place (no Vec-flatten-collect intermediate,
    // which doubled the peak memory at biobank N).
    use ndarray::parallel::prelude::*;
    use rayon::iter::{IntoParallelIterator, ParallelIterator};
    let mut design = Array2::<f64>::zeros((n, total_cols));
    design
        .axis_chunks_iter_mut(ndarray::Axis(0), 1024)
        .into_par_iter()
        .enumerate()
        .for_each(|(chunk_idx, mut block)| {
            let row_offset = chunk_idx * 1024;
            // Scratch buffers reused across rows in this chunk.
            let mut cur = Vec::<f64>::with_capacity(total_cols);
            let mut next = Vec::<f64>::with_capacity(total_cols);
            for (local_i, mut out_row) in block.outer_iter_mut().enumerate() {
                let i = row_offset + local_i;
                cur.clear();
                cur.push(1.0);
                for b in marginal_designs {
                    let q = b.ncols();
                    next.clear();
                    next.resize(cur.len() * q, 0.0);
                    // Hoist the row view out of the inner `col` loop so the
                    // q reads per `a_idx` reuse a single contiguous slice
                    // instead of recomputing `b[[i, col]]` strides per cell.
                    let b_row = b.row(i);
                    let b_slice = b_row
                        .as_slice()
                        .expect("Array2 row from outer_iter is contiguous");
                    for (a_idx, &aval) in cur.iter().enumerate() {
                        let off = a_idx * q;
                        let dst = &mut next[off..off + q];
                        for col in 0..q {
                            dst[col] = aval * b_slice[col];
                        }
                    }
                    std::mem::swap(&mut cur, &mut next);
                }
                // `out_row` is a row of the contiguous C-major `design`
                // Array2, so it is backed by a contiguous slice. Use a
                // bulk slice copy instead of an element-by-element write
                // loop.
                let out_slice = out_row
                    .as_slice_mut()
                    .expect("design row is contiguous in C-major Array2");
                out_slice.copy_from_slice(&cur);
            }
        });
    Ok(design)
}

fn build_random_effect_block(
    data: ArrayView2<'_, f64>,
    spec: &RandomEffectTermSpec,
) -> Result<RandomEffectBlock, BasisError> {
    let n = data.nrows();
    let p = data.ncols();
    if spec.feature_col >= p {
        crate::bail_dim_basis!(
            "random-effect term '{}' feature column {} out of bounds for {} columns",
            spec.name,
            spec.feature_col,
            p
        );
    }

    let col = data.column(spec.feature_col);
    if col.iter().any(|v| !v.is_finite()) {
        crate::bail_invalid_basis!(
            "random-effect term '{}' contains non-finite group values",
            spec.name
        );
    }

    let kept_levels: Vec<u64> = if let Some(levels) = spec.frozen_levels.as_ref() {
        if levels.is_empty() {
            crate::bail_invalid_basis!(
                "random-effect term '{}' has empty frozen_levels",
                spec.name
            );
        }
        levels.clone()
    } else {
        let mut levels_set = BTreeSet::<u64>::new();
        for &v in col {
            levels_set.insert(v.to_bits());
        }
        if levels_set.is_empty() {
            crate::bail_invalid_basis!("random-effect term '{}' has no observed levels", spec.name);
        }
        let levels: Vec<u64> = levels_set.into_iter().collect();
        let start_idx = if spec.drop_first_level && levels.len() > 1 {
            1usize
        } else {
            0usize
        };
        levels[start_idx..].to_vec()
    };

    if kept_levels.is_empty() {
        crate::bail_invalid_basis!(
            "random-effect term '{}' drops all levels; keep at least one level",
            spec.name
        );
    }

    let q = kept_levels.len();
    let mut level_to_col = BTreeMap::<u64, usize>::new();
    for (idx, &bits) in kept_levels.iter().enumerate() {
        if level_to_col.insert(bits, idx).is_some() {
            crate::bail_invalid_basis!(
                "random-effect term '{}' has duplicate frozen level bits {bits}",
                spec.name
            );
        }
    }
    let mut group_ids = Vec::with_capacity(n);
    for &v in col {
        let bits = v.to_bits();
        group_ids.push(level_to_col.get(&bits).copied());
    }

    Ok(RandomEffectBlock {
        name: spec.name.clone(),
        group_ids,
        num_groups: q,
        kept_levels,
    })
}

impl SmoothDesign {
    /// Map an unconstrained term coefficient vector to its constrained shape space.
    /// This is useful for nonlinear fits that optimize unconstrained parameters.
    pub fn map_term_coefficients(
        unconstrained: &Array1<f64>,
        shape: ShapeConstraint,
    ) -> Result<Array1<f64>, BasisError> {
        if unconstrained.is_empty() {
            crate::bail_invalid_basis!("unconstrained coefficient vector cannot be empty");
        }
        let mapped = match shape {
            ShapeConstraint::None => unconstrained.clone(),
            ShapeConstraint::MonotoneIncreasing => cumulative_exp(unconstrained, 1.0),
            ShapeConstraint::MonotoneDecreasing => cumulative_exp(unconstrained, -1.0),
            ShapeConstraint::Convex => second_cumulative_exp(unconstrained, 1.0),
            ShapeConstraint::Concave => second_cumulative_exp(unconstrained, -1.0),
        };
        Ok(mapped)
    }
}

fn bspline_boundary_endpoint(
    knots: &Array1<f64>,
    degree: usize,
    right: bool,
) -> Result<f64, BasisError> {
    if knots.len() <= degree + 1 {
        crate::bail_invalid_basis!("B-spline boundary condition requires a valid knot vector");
    }
    let n_basis = knots.len() - degree - 1;
    Ok(if right { knots[n_basis] } else { knots[degree] })
}

fn bspline_endpoint_row(
    knots: &Array1<f64>,
    degree: usize,
    condition: BSplineEndpointBoundaryCondition,
    right: bool,
    identifiability_transform: Option<&Array2<f64>>,
    coefficient_transform: Option<&Array2<f64>>,
) -> Result<Option<(Array1<f64>, f64)>, BasisError> {
    let derivative_order = match condition {
        BSplineEndpointBoundaryCondition::Free => return Ok(None),
        BSplineEndpointBoundaryCondition::Clamped => 1,
        BSplineEndpointBoundaryCondition::Anchored { .. } => 0,
    };
    let target = match condition {
        BSplineEndpointBoundaryCondition::Free | BSplineEndpointBoundaryCondition::Clamped => 0.0,
        BSplineEndpointBoundaryCondition::Anchored { value } => value,
    };
    if !target.is_finite() {
        crate::bail_invalid_basis!("anchored B-spline boundary value must be finite");
    }
    let endpoint = bspline_boundary_endpoint(knots, degree, right)?;
    let point = Array1::from_vec(vec![endpoint]);
    let options = if derivative_order == 1 {
        BasisOptions::first_derivative()
    } else {
        BasisOptions::value()
    };
    let (raw, _) = crate::basis::create_basis::<Dense>(
        point.view(),
        KnotSource::Provided(knots.view()),
        degree,
        options,
    )?;
    let mut row = raw.row(0).to_owned();
    if let Some(z) = identifiability_transform {
        if row.len() != z.nrows() {
            crate::bail_dim_basis!(
                "B-spline boundary constraint transform mismatch: row has {} columns but transform has {} rows",
                row.len(),
                z.nrows()
            );
        }
        row = row.dot(z);
    }
    if let Some(t) = coefficient_transform {
        if row.len() != t.nrows() {
            crate::bail_dim_basis!(
                "B-spline boundary constraint coefficient transform mismatch: row has {} columns but transform has {} rows",
                row.len(),
                t.nrows()
            );
        }
        row = row.dot(t);
    }
    Ok(Some((row, target)))
}

fn bspline_boundary_linear_constraints(
    boundary_conditions: BSplineBoundaryConditions,
    metadata: &BasisMetadata,
    degree: usize,
    coefficient_transform: Option<&Array2<f64>>,
) -> Result<Option<LinearInequalityConstraints>, BasisError> {
    if boundary_conditions.is_free() {
        return Ok(None);
    }
    let BasisMetadata::BSpline1D {
        knots,
        identifiability_transform,
        periodic: _,
        ..
    } = metadata
    else {
        crate::bail_invalid_basis!("B-spline boundary constraints require B-spline metadata");
    };

    let mut eq_rows = Vec::<Array1<f64>>::new();
    let mut eq_targets = Vec::<f64>::new();
    for (cond, right) in [
        (boundary_conditions.left, false),
        (boundary_conditions.right, true),
    ] {
        if let Some((row, target)) = bspline_endpoint_row(
            knots,
            degree,
            cond,
            right,
            identifiability_transform.as_ref(),
            coefficient_transform,
        )? {
            let norm = row.dot(&row).sqrt();
            if norm <= 1e-12 {
                if target.abs() > 1e-12 {
                    crate::bail_invalid_basis!(
                        "anchored B-spline boundary value {target} is infeasible because the endpoint row is zero"
                    );
                }
                continue;
            }
            eq_rows.push(row);
            eq_targets.push(target);
        }
    }
    if eq_rows.is_empty() {
        return Ok(None);
    }
    let p = eq_rows[0].len();
    let mut a = Array2::<f64>::zeros((2 * eq_rows.len(), p));
    let mut b = Array1::<f64>::zeros(2 * eq_rows.len());
    for (i, (row, &target)) in eq_rows.iter().zip(eq_targets.iter()).enumerate() {
        a.row_mut(2 * i).assign(row);
        b[2 * i] = target;
        a.row_mut(2 * i + 1).assign(&row.mapv(|v| -v));
        b[2 * i + 1] = -target;
    }
    Ok(Some(LinearInequalityConstraints { a, b }))
}

struct LocalSmoothTermBuild {
    dim: usize,
    design: DesignMatrix,
    penalties: Vec<Array2<f64>>,
    ops: Vec<Option<std::sync::Arc<dyn crate::terms::penalty_op::PenaltyOp>>>,
    nullspaces: Vec<usize>,
    /// Per-active-penalty null-space eigenvector matrices, parallel to
    /// `penalties` / `ops` / `nullspaces`. `Some(U_null)` when
    /// `nullspaces[k] > 0`, with `U_null` orthonormal columns spanning
    /// `null(penalties[k])` in this smooth's local coordinate system; `None`
    /// when the active block is full-rank. Stage 1 plumbing; Stage 2
    /// consumes this to absorb the smooth's null space into the parametric
    /// block at `TermCollectionDesign` construction.
    null_eigenvectors: Vec<Option<Array2<f64>>>,
    /// Joint-null absorption rotation for this smooth. `Some(rotation)`
    /// records `Q = [U_range | U_null]` spanning `null(Σ_k penalties[k])`,
    /// the joint null across all active penalty blocks on this smooth.
    /// `None` means the joint penalty is full-rank (joint nullity = 0) or
    /// there are no penalties. Stage-2 commit A: plumbing only — populated
    /// by commit B, applied by commit D.
    joint_null_rotation: Option<crate::terms::basis::JointNullRotation>,
    penaltyinfo: Vec<PenaltyInfo>,
    pre_dropped_penaltyinfo: Vec<PenaltyInfo>,
    metadata: BasisMetadata,
    linear_constraints: Option<LinearInequalityConstraints>,
    box_reparam: bool,
    kronecker_factored: Option<KroneckerFactoredBasis>,
}

#[derive(Clone)]
struct PcaScoresMemmapDesignOperator {
    mmap: Arc<memmap2::Mmap>,
    data_offset: usize,
    nrows: usize,
    ncols: usize,
    chunk_size: usize,
}

impl PcaScoresMemmapDesignOperator {
    fn open(path: PathBuf, chunk_size: usize) -> Result<Self, BasisError> {
        let file = File::open(&path).map_err(|err| {
            BasisError::InvalidInput(format!(
                "failed to open lazy Pca .npy scores '{}': {err}",
                path.display()
            ))
        })?;
        // The .npy scores file is read-only training-cache data; this
        // module never mutates it. The error path below converts mmap
        // failure to a typed `BasisError::InvalidInput`.
        // SAFETY: `memmap2::Mmap::map` requires no concurrent writers; the
        // contract is held by this module's read-only access pattern.
        let mmap = unsafe {
            memmap2::Mmap::map(&file).map_err(|err| {
                BasisError::InvalidInput(format!(
                    "failed to memmap lazy Pca .npy scores '{}': {err}",
                    path.display()
                ))
            })?
        };
        let (data_offset, nrows, ncols) = parse_f64_2d_npy_header(&mmap, &path)?;
        let expected = data_offset
            .checked_add(nrows.saturating_mul(ncols).saturating_mul(8))
            .ok_or_else(|| {
                BasisError::InvalidInput(format!(
                    "lazy Pca .npy scores '{}' shape is too large",
                    path.display()
                ))
            })?;
        if mmap.len() < expected {
            crate::bail_invalid_basis!(
                "lazy Pca .npy scores '{}' is truncated: header expects {} bytes, file has {}",
                path.display(),
                expected,
                mmap.len()
            );
        }
        Ok(Self {
            mmap: Arc::new(mmap),
            data_offset,
            nrows,
            ncols,
            chunk_size: chunk_size.max(1),
        })
    }

    fn value(&self, row: usize, col: usize) -> f64 {
        let offset = self.data_offset + (row * self.ncols + col) * 8;
        let mut bytes = [0_u8; 8];
        bytes.copy_from_slice(&self.mmap[offset..offset + 8]);
        f64::from_le_bytes(bytes)
    }

    fn chunk_rows(&self) -> usize {
        self.chunk_size.min(self.nrows.max(1))
    }
}

impl LinearOperator for PcaScoresMemmapDesignOperator {
    fn nrows(&self) -> usize {
        self.nrows
    }

    fn ncols(&self) -> usize {
        self.ncols
    }

    fn apply(&self, vector: &Array1<f64>) -> Array1<f64> {
        assert_eq!(
            vector.len(),
            self.ncols,
            "lazy Pca apply vector length mismatch"
        );
        let mut out = Array1::<f64>::zeros(self.nrows);
        for start in (0..self.nrows).step_by(self.chunk_rows()) {
            let end = (start + self.chunk_rows()).min(self.nrows);
            for row in start..end {
                let mut acc = 0.0;
                for col in 0..self.ncols {
                    acc += self.value(row, col) * vector[col];
                }
                out[row] = acc;
            }
        }
        out
    }

    fn apply_transpose(&self, vector: &Array1<f64>) -> Array1<f64> {
        assert_eq!(
            vector.len(),
            self.nrows,
            "lazy Pca apply_transpose vector length mismatch"
        );
        let mut out = Array1::<f64>::zeros(self.ncols);
        for start in (0..self.nrows).step_by(self.chunk_rows()) {
            let end = (start + self.chunk_rows()).min(self.nrows);
            for row in start..end {
                let scale = vector[row];
                if scale == 0.0 {
                    continue;
                }
                for col in 0..self.ncols {
                    out[col] += scale * self.value(row, col);
                }
            }
        }
        out
    }

    fn diag_xtw_x(&self, weights: &Array1<f64>) -> Result<Array2<f64>, String> {
        if weights.len() != self.nrows {
            return Err(format!(
                "lazy Pca diag_xtw_x weight length mismatch: weights={}, nrows={}",
                weights.len(),
                self.nrows
            ));
        }
        let mut gram = Array2::<f64>::zeros((self.ncols, self.ncols));
        for start in (0..self.nrows).step_by(self.chunk_rows()) {
            let end = (start + self.chunk_rows()).min(self.nrows);
            for row in start..end {
                let w = weights[row];
                if w == 0.0 {
                    continue;
                }
                for a in 0..self.ncols {
                    let xa = self.value(row, a);
                    if xa == 0.0 {
                        continue;
                    }
                    for b in a..self.ncols {
                        gram[[a, b]] += w * xa * self.value(row, b);
                    }
                }
            }
        }
        for a in 0..self.ncols {
            for b in 0..a {
                gram[[a, b]] = gram[[b, a]];
            }
        }
        Ok(gram)
    }

    fn apply_weighted_normal(
        &self,
        weights: &Array1<f64>,
        vector: &Array1<f64>,
        penalty: Option<&Array2<f64>>,
        ridge: f64,
    ) -> Array1<f64> {
        assert_eq!(
            weights.len(),
            self.nrows,
            "lazy Pca weighted-normal weight mismatch"
        );
        assert_eq!(
            vector.len(),
            self.ncols,
            "lazy Pca weighted-normal vector mismatch"
        );
        let mut out = Array1::<f64>::zeros(self.ncols);
        for start in (0..self.nrows).step_by(self.chunk_rows()) {
            let end = (start + self.chunk_rows()).min(self.nrows);
            for row in start..end {
                let w = weights[row].max(0.0);
                if w == 0.0 {
                    continue;
                }
                let mut row_dot = 0.0;
                for col in 0..self.ncols {
                    row_dot += self.value(row, col) * vector[col];
                }
                if row_dot == 0.0 {
                    continue;
                }
                let scaled = w * row_dot;
                for col in 0..self.ncols {
                    out[col] += scaled * self.value(row, col);
                }
            }
        }
        if let Some(pen) = penalty {
            out += &pen.dot(vector);
        }
        if ridge > 0.0 {
            out += &vector.mapv(|x| ridge * x);
        }
        out
    }
}

impl DenseDesignOperator for PcaScoresMemmapDesignOperator {
    fn compute_xtwy(&self, weights: &Array1<f64>, y: &Array1<f64>) -> Result<Array1<f64>, String> {
        if weights.len() != self.nrows || y.len() != self.nrows {
            return Err(format!(
                "lazy Pca compute_xtwy dimension mismatch: weights={}, y={}, nrows={}",
                weights.len(),
                y.len(),
                self.nrows
            ));
        }
        let mut out = Array1::<f64>::zeros(self.ncols);
        for start in (0..self.nrows).step_by(self.chunk_rows()) {
            let end = (start + self.chunk_rows()).min(self.nrows);
            for row in start..end {
                let scale = weights[row] * y[row];
                if scale == 0.0 {
                    continue;
                }
                for col in 0..self.ncols {
                    out[col] += scale * self.value(row, col);
                }
            }
        }
        Ok(out)
    }

    fn row_chunk_into(
        &self,
        rows: Range<usize>,
        mut out: ArrayViewMut2<'_, f64>,
    ) -> Result<(), MatrixMaterializationError> {
        if rows.end > self.nrows || rows.start > rows.end {
            return Err(MatrixMaterializationError::MissingRowChunk {
                context: "lazy Pca row range out of bounds",
            });
        }
        if out.nrows() != rows.end - rows.start || out.ncols() != self.ncols {
            return Err(MatrixMaterializationError::MissingRowChunk {
                context: "lazy Pca row_chunk_into shape mismatch",
            });
        }
        for (local, row) in (rows.start..rows.end).enumerate() {
            for col in 0..self.ncols {
                out[[local, col]] = self.value(row, col);
            }
        }
        Ok(())
    }

    fn to_dense(&self) -> Array2<f64> {
        let mut out = Array2::<f64>::zeros((self.nrows, self.ncols));
        self.row_chunk_into(0..self.nrows, out.view_mut())
            .expect("lazy Pca full materialization failed");
        out
    }
}

fn parse_f64_2d_npy_header(
    bytes: &[u8],
    path: &PathBuf,
) -> Result<(usize, usize, usize), BasisError> {
    if bytes.len() < 10 || &bytes[0..6] != b"\x93NUMPY" {
        crate::bail_invalid_basis!("lazy Pca scores '{}' is not a .npy file", path.display());
    }
    let major = bytes[6];
    let header_len = match major {
        1 => u16::from_le_bytes([bytes[8], bytes[9]]) as usize,
        2 | 3 => {
            if bytes.len() < 12 {
                crate::bail_invalid_basis!(
                    "lazy Pca scores '{}' has a truncated .npy header",
                    path.display()
                );
            }
            u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize
        }
        other => {
            crate::bail_invalid_basis!(
                "lazy Pca scores '{}' uses unsupported .npy version {}",
                path.display(),
                other
            );
        }
    };
    let header_start = if major == 1 { 10 } else { 12 };
    let data_offset = header_start + header_len;
    if bytes.len() < data_offset {
        crate::bail_invalid_basis!(
            "lazy Pca scores '{}' has a truncated .npy header",
            path.display()
        );
    }
    let header = std::str::from_utf8(&bytes[header_start..data_offset]).map_err(|err| {
        BasisError::InvalidInput(format!(
            "lazy Pca scores '{}' has a non-UTF8 .npy header: {err}",
            path.display()
        ))
    })?;
    if !(header.contains("'descr': '<f8'")
        || header.contains("\"descr\": \"<f8\"")
        || header.contains("'descr': '|f8'")
        || header.contains("\"descr\": \"|f8\""))
    {
        crate::bail_invalid_basis!(
            "lazy Pca scores '{}' must be float64 little-endian .npy",
            path.display()
        );
    }
    if header.contains("True") {
        crate::bail_invalid_basis!(
            "lazy Pca scores '{}' must be C-contiguous, not Fortran-ordered",
            path.display()
        );
    }
    let shape_pos = header.find("shape").ok_or_else(|| {
        BasisError::InvalidInput(format!(
            "lazy Pca scores '{}' .npy header is missing shape",
            path.display()
        ))
    })?;
    let open = header[shape_pos..].find('(').ok_or_else(|| {
        BasisError::InvalidInput(format!(
            "lazy Pca scores '{}' .npy header has malformed shape",
            path.display()
        ))
    })? + shape_pos;
    let close = header[open..].find(')').ok_or_else(|| {
        BasisError::InvalidInput(format!(
            "lazy Pca scores '{}' .npy header has malformed shape",
            path.display()
        ))
    })? + open;
    let dims = header[open + 1..close]
        .split(',')
        .map(str::trim)
        .filter(|part| !part.is_empty())
        .map(|part| part.parse::<usize>())
        .collect::<Result<Vec<_>, _>>()
        .map_err(|err| {
            BasisError::InvalidInput(format!(
                "lazy Pca scores '{}' .npy shape is not integral: {err}",
                path.display()
            ))
        })?;
    if dims.len() != 2 {
        crate::bail_invalid_basis!(
            "lazy Pca scores '{}' must have shape (N, K), got {:?}",
            path.display(),
            dims
        );
    }
    Ok((data_offset, dims[0], dims[1]))
}

fn pca_center_mean(x: ArrayView2<'_, f64>) -> Result<Array1<f64>, BasisError> {
    if x.nrows() == 0 {
        crate::bail_invalid_basis!("Pca basis requires at least one row to compute center mean");
    }
    let mut mean = Array1::<f64>::zeros(x.ncols());
    for row in x.rows() {
        mean += &row;
    }
    mean.mapv_inplace(|v| v / x.nrows() as f64);
    Ok(mean)
}

fn build_pca_smooth_basis(
    data: ArrayView2<'_, f64>,
    feature_cols: &[usize],
    basis_matrix: &Array2<f64>,
    centered: bool,
    smooth_penalty: f64,
    center_mean: Option<&Array1<f64>>,
    pca_basis_path: Option<&PathBuf>,
    chunk_size: usize,
) -> Result<BasisBuildResult, BasisError> {
    if let Some(path) = pca_basis_path {
        let op = PcaScoresMemmapDesignOperator::open(path.clone(), chunk_size)?;
        if op.nrows != data.nrows() {
            crate::bail_dim_basis!(
                "lazy Pca scores row mismatch: .npy has {}, data has {}",
                op.nrows,
                data.nrows()
            );
        }
        let k = op.ncols;
        let mut penalty = Array2::<f64>::eye(k);
        penalty.mapv_inplace(|v| v * smooth_penalty);
        let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
            filter_active_penalty_candidates_with_ops(vec![PenaltyCandidate {
                matrix: penalty,
                nullspace_dim_hint: 0,
                source: PenaltySource::Other("PcaRidge".to_string()),
                normalization_scale: 1.0,
                kronecker_factors: None,
                op: None,
            }])?;
        return Ok(BasisBuildResult {
            design: DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(Arc::new(op))),
            penalties,
            nullspace_dims,
            penaltyinfo,
            ops,
            null_eigenvectors,
            joint_null_rotation: None,
            metadata: BasisMetadata::Pca {
                feature_cols: feature_cols.to_vec(),
                basis_matrix: basis_matrix.clone(),
                centered,
                smooth_penalty,
                center_mean: center_mean.cloned(),
                pca_basis_path: Some(path.clone()),
                chunk_size: chunk_size.max(1),
            },
            kronecker_factored: None,
        });
    }
    if basis_matrix.nrows() != feature_cols.len() {
        crate::bail_dim_basis!(
            "Pca basis row mismatch: basis rows={}, feature columns={}",
            basis_matrix.nrows(),
            feature_cols.len()
        );
    }
    let mut x = select_columns(data, feature_cols)?;
    let mean = if centered {
        match center_mean {
            Some(mean) => mean.clone(),
            None => pca_center_mean(x.view())?,
        }
    } else {
        Array1::<f64>::zeros(feature_cols.len())
    };
    if centered {
        for mut row in x.rows_mut() {
            row -= &mean;
        }
    }
    let design = fast_ab(&x, basis_matrix);
    let k = basis_matrix.ncols();
    let mut penalty = Array2::<f64>::eye(k);
    penalty.mapv_inplace(|v| v * smooth_penalty);
    let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
        filter_active_penalty_candidates_with_ops(vec![PenaltyCandidate {
            matrix: penalty,
            nullspace_dim_hint: 0,
            source: PenaltySource::Other("PcaRidge".to_string()),
            normalization_scale: 1.0,
            kronecker_factors: None,
            op: None,
        }])?;
    Ok(BasisBuildResult {
        design: DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(design)),
        penalties,
        nullspace_dims,
        penaltyinfo,
        ops,
        null_eigenvectors,
        joint_null_rotation: None,
        metadata: BasisMetadata::Pca {
            feature_cols: feature_cols.to_vec(),
            basis_matrix: basis_matrix.clone(),
            centered,
            smooth_penalty,
            center_mean: centered.then_some(mean),
            pca_basis_path: None,
            chunk_size: chunk_size.max(1),
        },
        kronecker_factored: None,
    })
}

fn apply_by_variable_to_local_build(
    mut built: LocalSmoothTermBuild,
    data: ArrayView2<'_, f64>,
    by_col: usize,
    by: &ByVariableSpec,
    term_name: &str,
) -> Result<LocalSmoothTermBuild, BasisError> {
    if by_col >= data.ncols() {
        crate::bail_dim_basis!(
            "by-variable smooth term '{term_name}' references column {by_col}, but data has {} columns",
            data.ncols()
        );
    }
    let weights = match by {
        ByVariableSpec::Numeric => data.column(by_col).to_owned(),
        ByVariableSpec::Level { value_bits, .. } => data.column(by_col).mapv(|value| {
            if value.to_bits() == *value_bits {
                1.0
            } else {
                0.0
            }
        }),
    };
    if weights.iter().any(|value| !value.is_finite()) {
        crate::bail_invalid_basis!(
            "by-variable smooth term '{term_name}' has non-finite by-column values"
        );
    }

    let mut dense = built
        .design
        .try_to_dense_by_chunks("by-variable smooth row gating")
        .map_err(BasisError::InvalidInput)?;
    for (mut row, &weight) in dense.rows_mut().into_iter().zip(weights.iter()) {
        row.mapv_inplace(|value| value * weight);
    }
    built.design = DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(dense));
    built.kronecker_factored = None;
    Ok(built)
}

fn ensure_by_variable_specs_match(
    kind: &BySmoothKind,
    by: &ByVariableSpec,
    term_name: &str,
) -> Result<(), BasisError> {
    match (kind, by) {
        (BySmoothKind::Numeric, ByVariableSpec::Numeric) => Ok(()),
        (BySmoothKind::Level { level_bits }, ByVariableSpec::Level { value_bits, .. })
            if level_bits == value_bits =>
        {
            Ok(())
        }
        _ => Err(BasisError::InvalidInput(format!(
            "by-variable smooth term '{term_name}' has inconsistent by-variable specifications"
        ))),
    }
}

/// Build a factor-smooth interaction basis (`bs="fs"`/`"sz"`/`"re"`).
///
/// A factor smooth replicates a shared marginal smooth in the continuous
/// covariate(s) once per level of a grouping factor, coupling all level blocks
/// through a *single* set of smoothing parameters (one per marginal penalty).
/// This is mgcv's `smooth.construct.fs.smooth.spec` realization and the
/// random-effect interpretation of a smooth: the per-level deviations are an
/// exchangeable family whose joint wiggliness/shrinkage is governed by the
/// shared λ, so the construction scales to many levels with a fixed parameter
/// count.
///
/// Flavours:
/// * `Fs` — full random factor-smooth. The marginal carries its wiggliness
///   penalty *and* a null-space ridge (double penalty), so the replicated
///   design is a proper full-rank random effect: each level's curve is shrunk
///   toward zero (intercept + linear trend included), recovering the mgcv
///   `bs="fs"` penalty structure `I_L ⊗ S_j` for every marginal penalty `S_j`.
/// * `Sz` — sum-to-zero factor smooth. Delegates to the existing
///   [`SmoothBasisSpec::FactorSumToZero`] construction (`L-1` deviation blocks,
///   coefficient-wise zero sum across levels).
/// * `Re` — pure random effect / random slope (`bs="re"`). A degree-1 marginal
///   gives the per-level `[1, x]` span; the penalty is the identity over each
///   level block (iid Gaussian coefficients), matching mgcv's `bs="re"` ridge.
///
/// The grouping levels are resolved once at fit time (sorted unique bit
/// patterns of the factor column) and frozen into the returned metadata so the
/// predict-time rebuild evaluates every row against its own level's block.
fn build_factor_smooth(
    data: ArrayView2<'_, f64>,
    spec: &FactorSmoothSpec,
    term_name: &str,
    workspace: &mut crate::basis::BasisWorkspace,
) -> Result<LocalSmoothTermBuild, BasisError> {
    if spec.continuous_cols.len() != 1 {
        crate::bail_invalid_basis!(
            "factor smooth term '{}' currently supports exactly one continuous covariate; found {}",
            term_name,
            spec.continuous_cols.len()
        );
    }
    let feature_col = spec.continuous_cols[0];
    let group_col = spec.group_col;
    if feature_col >= data.ncols() || group_col >= data.ncols() {
        crate::bail_dim_basis!(
            "factor smooth term '{}' references columns ({}, {}) out of bounds for {} columns",
            term_name,
            feature_col,
            group_col,
            data.ncols()
        );
    }

    // `Sz` is exactly the existing sum-to-zero factor smooth: reuse it verbatim
    // so there is a single source of truth for the zero-sum construction.
    if matches!(spec.flavour, FactorSmoothFlavour::Sz) {
        let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
        let inner = SmoothBasisSpec::BSpline1D {
            feature_col,
            spec: factor_smooth_marginal_for_replay(&spec.marginal),
        };
        let sz_term = SmoothTermSpec {
            name: term_name.to_string(),
            basis: SmoothBasisSpec::FactorSumToZero {
                inner: Box::new(inner),
                by_col: group_col,
                levels,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        return build_single_local_smooth_term(data, &sz_term, workspace);
    }

    let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
    let n_levels = levels.len();
    if n_levels < 2 {
        crate::bail_invalid_basis!(
            "factor smooth term '{}' requires at least two grouping levels; found {}",
            term_name,
            n_levels
        );
    }

    // Build the shared marginal design + penalties from the 1-D B-spline.
    // `Re` forces a degree-1 marginal (linear span) and replaces the marginal
    // wiggliness with an identity ridge below; `Fs` keeps the user's marginal
    // (cubic by default) with its double penalty so the null space is shrunk.
    let marginal_spec = factor_smooth_marginal_for_replay(&spec.marginal);
    let inner_term = SmoothTermSpec {
        name: format!("{term_name}::marginal"),
        basis: SmoothBasisSpec::BSpline1D {
            feature_col,
            spec: marginal_spec,
        },
        shape: ShapeConstraint::None,
        joint_null_rotation: None,
    };
    let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
    let base = inner
        .design
        .try_to_dense_by_chunks("factor smooth marginal")
        .map_err(BasisError::InvalidInput)?;
    let n = base.nrows();
    let p = base.ncols();
    let q = p * n_levels;

    // Block-diagonal replicated design: row i contributes its marginal row to
    // the column block owned by its grouping level, zeros elsewhere.
    let mut dense = Array2::<f64>::zeros((n, q));
    for i in 0..n {
        let bits = data[[i, group_col]].to_bits();
        let level_idx = levels.iter().position(|b| *b == bits).ok_or_else(|| {
            BasisError::InvalidInput(format!(
                "factor smooth term '{term_name}' saw an unseen grouping level at row {}",
                i + 1
            ))
        })?;
        let start = level_idx * p;
        dense
            .slice_mut(s![i, start..start + p])
            .assign(&base.row(i));
    }

    // Penalties: replicate each marginal penalty into a block-diagonal
    // `I_L ⊗ S_j` so every level shares the same smoothing parameter λ_j (one
    // λ per marginal penalty), the defining feature of a factor smooth. For
    // `Re` the marginal penalty is replaced by an identity ridge so each
    // per-level coefficient is an iid Gaussian random effect.
    let marginal_penalties: Vec<Array2<f64>> = if matches!(spec.flavour, FactorSmoothFlavour::Re) {
        vec![Array2::<f64>::eye(p)]
    } else {
        inner.penalties.clone()
    };
    let marginal_penaltyinfo: Vec<PenaltyInfo> = if matches!(spec.flavour, FactorSmoothFlavour::Re)
    {
        vec![PenaltyInfo {
            source: PenaltySource::Primary,
            original_index: 0,
            active: true,
            effective_rank: p,
            dropped_reason: None,
            nullspace_dim_hint: 0,
            normalization_scale: 1.0,
            kronecker_factors: None,
        }]
    } else {
        inner.penaltyinfo.clone()
    };
    if marginal_penalties.len() != marginal_penaltyinfo.len() {
        crate::bail_invalid_basis!(
            "internal factor-smooth penalty metadata mismatch for term '{}': penalties={}, infos={}",
            term_name,
            marginal_penalties.len(),
            marginal_penaltyinfo.len()
        );
    }

    let mut penalties = Vec::<Array2<f64>>::with_capacity(marginal_penalties.len());
    let mut penaltyinfo = Vec::<PenaltyInfo>::with_capacity(marginal_penalties.len());
    for (penalty_pos, s_inner) in marginal_penalties.iter().enumerate() {
        let mut s_big = Array2::<f64>::zeros((q, q));
        for level in 0..n_levels {
            let start = level * p;
            s_big
                .slice_mut(s![start..start + p, start..start + p])
                .assign(s_inner);
        }
        let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
        let mut info = marginal_penaltyinfo[penalty_pos].clone();
        info.original_index = penalty_pos;
        info.normalization_scale *= factor_smooth_scale;
        info.nullspace_dim_hint = info.nullspace_dim_hint.saturating_mul(n_levels);
        info.kronecker_factors = None;
        penalties.push(s_big);
        penaltyinfo.push(info);
    }

    let mut nullspaces: Vec<usize> = if matches!(spec.flavour, FactorSmoothFlavour::Re) {
        vec![0]
    } else {
        inner
            .nullspaces
            .iter()
            .map(|ns| ns.saturating_mul(n_levels))
            .collect()
    };

    // `Fs` is the random-effect flavour of a smooth: the per-group curve is an
    // exchangeable Gaussian *function*, so its WHOLE coefficient vector — not
    // only the {const, linear} null space the double penalty already ridges —
    // must be shrinkable toward zero under one shared variance. The marginal's
    // derivative (wiggliness) penalty `S_wiggle` shapes the curvature but leaves
    // its overall MAGNITUDE governed only by the curvature eigenvalues; the
    // softest curvature mode (just past linear) is then barely penalized, so
    // REML can park ~0.8 residual curvature DF per group on within-window noise
    // and the held-out forecast curves away from the true per-group line
    // (gam#712 real arm, gam#713). Capping the basis dimension (dc912b30) shrank
    // the count of those modes but not their variance.
    //
    // Add a range-space identity ridge `I_L ⊗ (I_p − Z Zᵀ)` (Z = the marginal
    // wiggliness null basis = the {const, linear} directions), shared across all
    // groups under a single λ, so the curvature MAGNITUDE of every group's curve
    // is a proper zero-mean random effect. With linear data REML drives this λ
    // up and the per-group curvature variance → 0, degrading `fs` to the linear
    // random slope the DGP is (edf → ≈2/group); with genuine curvature the same
    // λ stays small and the wiggle survives (it is data-adaptive, not a cap).
    // The ridge is gated by `m_null_penalty_orders`: order ≥ 1 enables it, the
    // explicit `m=0` form opts out and keeps the legacy null-only double penalty.
    if let FactorSmoothFlavour::Fs {
        m_null_penalty_orders,
    } = &spec.flavour
        && m_null_penalty_orders.iter().copied().max().unwrap_or(0) >= 1
    {
        // Range-space projector for one marginal block: `P = I_p − Z Zᵀ`, the
        // orthogonal complement of the wiggliness penalty's null space (the
        // curvature subspace). `inner.null_eigenvectors[0]` is the orthonormal
        // null basis of the primary marginal penalty `S_wiggle`; when the
        // marginal is full-rank (`None`) the whole block is curvature and the
        // projector is the identity.
        let mut p_range = Array2::<f64>::eye(p);
        if let Some(Some(z)) = inner.null_eigenvectors.first()
            && z.nrows() == p
        {
            p_range -= &z.dot(&z.t());
        }
        let mut s_range = Array2::<f64>::zeros((q, q));
        for level in 0..n_levels {
            let start = level * p;
            s_range
                .slice_mut(s![start..start + p, start..start + p])
                .assign(&p_range);
        }
        let (s_range, range_scale) = normalize_penalty_in_constrained_space(&s_range);
        let range_block = crate::terms::basis::analyze_penalty_block_with_op(&s_range, None)?;
        if range_block.rank > 0 {
            let original_index = penalties.len();
            penalties.push(range_block.sym_penalty);
            nullspaces.push(range_block.nullity);
            penaltyinfo.push(PenaltyInfo {
                source: PenaltySource::Primary,
                original_index,
                active: true,
                effective_rank: range_block.rank,
                dropped_reason: None,
                nullspace_dim_hint: range_block.nullity,
                normalization_scale: range_scale,
                kronecker_factors: None,
            });
        }
    }
    let null_eigenvectors = crate::terms::basis::recompute_null_eigenvectors(&penalties)?;
    let joint_null_rotation = crate::terms::basis::compute_joint_null_rotation(&penalties)?;

    // Metadata: carry the marginal knot geometry + frozen levels so prediction
    // reconstructs an identical replicated design.
    let (knots, degree, periodic) = match &inner.metadata {
        BasisMetadata::BSpline1D {
            knots,
            periodic,
            degree,
            ..
        } => (
            knots.clone(),
            degree.unwrap_or(spec.marginal.degree),
            *periodic,
        ),
        other => {
            crate::bail_invalid_basis!(
                "factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
                term_name,
                other
            );
        }
    };
    let flavour_tag = match &spec.flavour {
        FactorSmoothFlavour::Fs { .. } => "fs",
        FactorSmoothFlavour::Sz => "sz",
        FactorSmoothFlavour::Re => "re",
    }
    .to_string();
    let metadata = BasisMetadata::FactorSmooth {
        continuous_cols: spec.continuous_cols.clone(),
        group_col,
        knots,
        degree,
        periodic,
        group_levels: levels,
        flavour: flavour_tag,
    };

    let ops = vec![None; penalties.len()];
    Ok(LocalSmoothTermBuild {
        dim: q,
        design: DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(dense)),
        penalties,
        ops,
        nullspaces,
        null_eigenvectors,
        joint_null_rotation,
        penaltyinfo,
        pre_dropped_penaltyinfo: Vec::new(),
        metadata,
        linear_constraints: None,
        box_reparam: false,
        kronecker_factored: None,
    })
}

/// Resolve the grouping levels for a factor smooth: replay the frozen level
/// list when present (predict path), otherwise discover the sorted unique bit
/// patterns of the factor column (fit path).
fn resolve_factor_smooth_levels(
    data: ArrayView2<'_, f64>,
    group_col: usize,
    spec: &FactorSmoothSpec,
    term_name: &str,
) -> Result<Vec<u64>, BasisError> {
    if let Some(frozen) = &spec.group_frozen_levels {
        if frozen.is_empty() {
            crate::bail_invalid_basis!(
                "factor smooth term '{}' has an empty frozen level list",
                term_name
            );
        }
        return Ok(frozen.clone());
    }
    let mut bits: Vec<u64> = data.column(group_col).iter().map(|v| v.to_bits()).collect();
    bits.sort_by(|a, b| {
        f64::from_bits(*a)
            .partial_cmp(&f64::from_bits(*b))
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    bits.dedup();
    Ok(bits)
}

/// Marginal B-spline spec for a factor-smooth block. The marginal always builds
/// without an identifiability constraint (the per-level replication, not a
/// sum-to-zero side constraint, provides identifiability against the parametric
/// block). At predict time the marginal's knot geometry has already been pinned
/// into `marginal.knotspec` by the metadata replay, so the spec is used
/// verbatim aside from clearing the identifiability transform.
fn factor_smooth_marginal_for_replay(marginal: &BSplineBasisSpec) -> BSplineBasisSpec {
    let mut m = marginal.clone();
    m.identifiability = BSplineIdentifiability::None;
    m
}

fn build_single_local_smooth_term(
    data: ArrayView2<'_, f64>,
    term: &SmoothTermSpec,
    workspace: &mut crate::basis::BasisWorkspace,
) -> Result<LocalSmoothTermBuild, BasisError> {
    if term.shape != ShapeConstraint::None && !shape_supports_basis(term) {
        crate::bail_invalid_basis!(
            "ShapeConstraint::{:?} is unsupported for term '{}'",
            term.shape,
            term.name
        );
    }
    if let SmoothBasisSpec::ByVariable {
        inner,
        by_col,
        kind,
        by,
    } = &term.basis
    {
        ensure_by_variable_specs_match(kind, by, &term.name)?;
        let inner_term = SmoothTermSpec {
            name: term.name.clone(),
            basis: (**inner).clone(),
            shape: term.shape,
            joint_null_rotation: None,
        };
        let built = build_single_local_smooth_term(data, &inner_term, workspace)?;
        return apply_by_variable_to_local_build(built, data, *by_col, by, &term.name);
    }

    let mut shape_axis_col: Option<usize> = None;
    let mut built: BasisBuildResult = match &term.basis {
        SmoothBasisSpec::FactorSumToZero {
            inner,
            by_col,
            levels,
        } => {
            if *by_col >= data.ncols() {
                crate::bail_dim_basis!(
                    "term '{}' by column {} out of bounds for {} columns",
                    term.name,
                    by_col,
                    data.ncols()
                );
            }
            if levels.len() < 2 {
                crate::bail_invalid_basis!(
                    "sum-to-zero factor smooth term '{}' requires at least two levels",
                    term.name
                );
            }
            if term.shape != ShapeConstraint::None {
                crate::bail_invalid_basis!(
                    "ShapeConstraint::{:?} is unsupported for sum-to-zero factor smooth term '{}'",
                    term.shape,
                    term.name
                );
            }
            let inner_term = SmoothTermSpec {
                name: format!("{}::inner", term.name),
                basis: (**inner).clone(),
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            };
            let mut inner_built = build_single_local_smooth_term(data, &inner_term, workspace)?;
            let base = inner_built
                .design
                .try_to_dense_by_chunks("sum-to-zero factor smooth")
                .map_err(BasisError::InvalidInput)?;
            let n = base.nrows();
            let p = base.ncols();
            let l_minus_one = levels.len() - 1;
            let mut dense = Array2::<f64>::zeros((n, p * l_minus_one));
            for i in 0..n {
                let bits = data[[i, *by_col]].to_bits();
                let level_idx = levels.iter().position(|b| *b == bits).ok_or_else(|| {
                    BasisError::InvalidInput(format!(
                        "sum-to-zero factor smooth term '{}' saw an unseen level at row {}",
                        term.name,
                        i + 1
                    ))
                })?;
                if level_idx < l_minus_one {
                    let start = level_idx * p;
                    dense
                        .slice_mut(s![i, start..start + p])
                        .assign(&base.row(i));
                } else {
                    for level in 0..l_minus_one {
                        let start = level * p;
                        dense
                            .slice_mut(s![i, start..start + p])
                            .assign(&base.row(i).mapv(|v| -v));
                    }
                }
            }
            let mut penalties = Vec::<Array2<f64>>::with_capacity(inner_built.penalties.len());
            let active_penalty_indices = inner_built
                .penaltyinfo
                .iter()
                .enumerate()
                .filter_map(|(idx, info)| info.active.then_some(idx))
                .collect::<Vec<_>>();
            if active_penalty_indices.len() != inner_built.penalties.len() {
                crate::bail_invalid_basis!(
                    "internal sz penalty metadata mismatch: activeinfos={}, penalties={}",
                    active_penalty_indices.len(),
                    inner_built.penalties.len()
                );
            }
            for (penalty_pos, s_inner) in inner_built.penalties.iter().enumerate() {
                let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
                for a in 0..l_minus_one {
                    for b in 0..l_minus_one {
                        let factor = if a == b { 2.0 } else { 1.0 };
                        let mut block = s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
                        block.assign(&s_inner.mapv(|v| v * factor));
                    }
                }
                let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
                let info_idx = active_penalty_indices[penalty_pos];
                inner_built.penaltyinfo[info_idx].normalization_scale *= factor_smooth_scale;
                penalties.push(s_big);
            }
            inner_built.dim = p * l_minus_one;
            inner_built.design = DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(dense));
            inner_built.penalties = penalties;
            inner_built.ops = vec![None; inner_built.penalties.len()];
            inner_built.nullspaces = inner_built
                .nullspaces
                .iter()
                .map(|ns| ns.saturating_mul(l_minus_one))
                .collect();
            // Invariant: `null_eigenvectors[k]` must mirror `penalties[k]`'s
            // spectral null space. We just rebuilt `inner_built.penalties` from
            // Kronecker-like `S_big` blocks, so the previously-plumbed
            // `null_eigenvectors` (still parallel to the OLD per-level penalty)
            // is stale. Recompute from the rebuilt penalties to restore the
            // invariant; ditto for the joint-null absorption rotation.
            inner_built.null_eigenvectors =
                crate::terms::basis::recompute_null_eigenvectors(&inner_built.penalties)?;
            inner_built.joint_null_rotation =
                crate::terms::basis::compute_joint_null_rotation(&inner_built.penalties)?;
            inner_built.kronecker_factored = None;
            return Ok(inner_built);
        }
        SmoothBasisSpec::BSpline1D { feature_col, spec } => {
            if *feature_col >= data.ncols() {
                crate::bail_dim_basis!(
                    "term '{}' feature column {} out of bounds for {} columns",
                    term.name,
                    feature_col,
                    data.ncols()
                );
            }
            let mut spec_local = spec.clone();
            if term.shape != ShapeConstraint::None {
                // Shape-constrained B-splines are anchored by construction.
                // Sum-to-zero side constraints conflict with monotonic/convex cones.
                spec_local.identifiability = BSplineIdentifiability::None;
            }
            // Boundary conditions are emitted by the smooth-level paired
            // linear-constraint path (`bspline_boundary_linear_constraints`),
            // which supports non-zero anchors and composes them with the frozen
            // identifiability transform. Clear them here so the basis builder
            // does not also bake them into the basis null space — the legacy
            // basis-level path rejected non-zero anchors and dropped columns
            // before the frozen transform, conflicting with the constraint path.
            spec_local.boundary_conditions = BSplineBoundaryConditions::default();
            build_bspline_basis_1d(data.column(*feature_col), &spec_local)?
        }
        SmoothBasisSpec::ThinPlate {
            feature_cols,
            spec,
            input_scales,
        } => {
            if term.shape != ShapeConstraint::None {
                if feature_cols.len() != 1 {
                    crate::bail_invalid_basis!(
                        "ShapeConstraint::{:?} for term '{}' on ThinPlate basis requires exactly 1 feature axis; found {}",
                        term.shape,
                        term.name,
                        feature_cols.len()
                    );
                }
                shape_axis_col = Some(feature_cols[0]);
            }
            let mut x = select_columns(data, feature_cols)?;
            // Auto-standardize multivariate inputs: use stored scales (prediction)
            // or compute fresh ones (training). Same standardization-vs-
            // length-scale compensation as Matérn / hybrid Duchon: divide
            // the user's L by σ_geom so kernel(‖x_std − c_std‖/L_eff)
            // matches the original-coord kernel for uniform σ.
            let (scales, length_scale_eff) = if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                (
                    Some(s.clone()),
                    compensate_length_scale_for_standardization(spec.length_scale, s),
                )
            } else if let Some(s) = compute_spatial_input_scales(x.view()) {
                apply_input_standardization(&mut x, &s);
                let l_eff = compensate_length_scale_for_standardization(spec.length_scale, &s);
                (Some(s), l_eff)
            } else {
                (None, spec.length_scale)
            };
            let mut spec_local = spec.clone();
            spec_local.length_scale = length_scale_eff;
            if matches!(
                spec_local.identifiability,
                SpatialIdentifiability::OrthogonalToParametric
            ) {
                spec_local.identifiability = SpatialIdentifiability::None;
            }
            let mut result = build_thin_plate_basis(x.view(), &spec_local).map_err(|err| {
                rewrite_thin_plate_knots_error(err, &term.name, feature_cols.len(), spec)
            })?;
            // Inject input scales into metadata; also restore the user's
            // original length_scale (not the σ_geom-compensated one) so a
            // metadata-driven rebuild that re-applies compensation does not
            // double-divide. The build may auto-promote to Duchon when
            // canonical TPS is infeasible (k < polynomial-nullspace size);
            // in that case patch the Duchon metadata variant so predict-time
            // round-trips through the same standardized data path.
            match &mut result.metadata {
                BasisMetadata::ThinPlate {
                    input_scales: ms,
                    length_scale,
                    ..
                } => {
                    *ms = scales;
                    *length_scale = spec.length_scale;
                }
                BasisMetadata::Duchon {
                    input_scales: ms,
                    length_scale,
                    ..
                } => {
                    *ms = scales;
                    // The ThinPlate auto-promotion path delegates to
                    // `build_duchon_basis` with `Some(spec_local.length_scale)`,
                    // which is the σ_geom-compensated value. The metadata
                    // therefore records the compensated kernel range, but the
                    // freeze→replay round trip plugs that value back into a
                    // user-facing `DuchonBasisSpec.length_scale` whose builder
                    // applies the σ_geom compensation a second time. Restore
                    // the user-facing scale here so replay re-compensates
                    // exactly once and reproduces the realized fit-time basis.
                    *length_scale = Some(spec.length_scale);
                }
                _ => {}
            }
            result
        }
        SmoothBasisSpec::Sphere { feature_cols, spec } => {
            if term.shape != ShapeConstraint::None {
                crate::bail_invalid_basis!(
                    "ShapeConstraint::{:?} for term '{}' is not supported on spherical splines",
                    term.shape,
                    term.name
                );
            }
            let x = select_columns(data, feature_cols)?;
            build_spherical_spline_basis(x.view(), spec)?
        }
        SmoothBasisSpec::Matern {
            feature_cols,
            spec,
            input_scales,
        } => {
            if term.shape != ShapeConstraint::None {
                if feature_cols.len() != 1 {
                    crate::bail_invalid_basis!(
                        "ShapeConstraint::{:?} for term '{}' on Matern basis requires exactly 1 feature axis; found {}",
                        term.shape,
                        term.name,
                        feature_cols.len()
                    );
                }
                shape_axis_col = Some(feature_cols[0]);
            }
            let mut x = select_columns(data, feature_cols)?;
            // Auto-standardization (per-axis division by σ_a) reinterprets
            // the user's `length_scale` from original data coordinates
            // into post-standardization coordinates: for uniform σ_a = σ,
            // `kernel(‖x_std − c_std‖/L)` equals `kernel(‖x − c‖/(σ·L))`,
            // so the effective kernel range shrinks by σ. To keep
            // `length_scale` consistently expressed in *original* data
            // coordinates regardless of axis variances, we standardize
            // and divide L by σ_geom = (∏σ_a)^(1/d). For uniform σ this
            // recovers the user's kernel exactly; for anisotropic data
            // the resulting per-axis effective scales σ_a / σ_geom are
            // the standard Mahalanobis preconditioning and preserve the
            // geometric-mean kernel range. Storing the σ vector in
            // metadata.input_scales makes the same transformation
            // replayable at predict time.
            let (scales, length_scale_eff) = if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                (
                    Some(s.clone()),
                    compensate_length_scale_for_standardization(spec.length_scale, s),
                )
            } else if let Some(s) = compute_spatial_input_scales(x.view()) {
                apply_input_standardization(&mut x, &s);
                let l_eff = compensate_length_scale_for_standardization(spec.length_scale, &s);
                (Some(s), l_eff)
            } else {
                (None, spec.length_scale)
            };
            let mut spec_local = spec.clone();
            spec_local.length_scale = length_scale_eff;
            let mut result = build_matern_basiswithworkspace(x.view(), &spec_local, workspace)?;
            if let BasisMetadata::Matern {
                input_scales,
                length_scale,
                ..
            } = &mut result.metadata
            {
                *input_scales = scales;
                *length_scale = spec.length_scale;
            }
            result
        }
        SmoothBasisSpec::Duchon {
            feature_cols,
            spec,
            input_scales,
        } => {
            if term.shape != ShapeConstraint::None {
                if feature_cols.len() != 1 {
                    crate::bail_invalid_basis!(
                        "ShapeConstraint::{:?} for term '{}' on Duchon basis requires exactly 1 feature axis; found {}",
                        term.shape,
                        term.name,
                        feature_cols.len()
                    );
                }
                shape_axis_col = Some(feature_cols[0]);
            }
            let mut x = select_columns(data, feature_cols)?;
            // Hybrid Duchon (length_scale=Some) is governed by the same
            // standardization-vs-length-scale equivalence as Matérn: the
            // user's `length_scale` is interpreted in original data
            // coordinates, but auto-standardization (per-axis division by
            // σ_a) reinterprets it as σ_geom · L. Pre-multiply by 1/σ_geom
            // so kernel(‖x_std − c_std‖/L_eff) reproduces the user's
            // original-coord kernel exactly for uniform σ_a, and reduces
            // to standard Mahalanobis preconditioning for anisotropic σ.
            // Pure Duchon (length_scale=None) is scale-free and needs no
            // compensation.
            let (scales, length_scale_eff) = if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                (
                    Some(s.clone()),
                    compensate_optional_length_scale_for_standardization(spec.length_scale, s),
                )
            } else if let Some(s) = compute_spatial_input_scales(x.view()) {
                apply_input_standardization(&mut x, &s);
                let l_eff =
                    compensate_optional_length_scale_for_standardization(spec.length_scale, &s);
                (Some(s), l_eff)
            } else {
                (None, spec.length_scale)
            };
            let mut spec_local = spec.clone();
            spec_local.length_scale = length_scale_eff;
            if matches!(
                spec_local.identifiability,
                SpatialIdentifiability::OrthogonalToParametric
            ) {
                spec_local.identifiability = SpatialIdentifiability::None;
            }
            let mut result = build_duchon_basiswithworkspace(x.view(), &spec_local, workspace)?;
            if let BasisMetadata::Duchon {
                input_scales,
                length_scale,
                ..
            } = &mut result.metadata
            {
                *input_scales = scales;
                *length_scale = spec.length_scale;
            }
            result
        }
        SmoothBasisSpec::Pca {
            feature_cols,
            basis_matrix,
            centered,
            smooth_penalty,
            center_mean,
            pca_basis_path,
            chunk_size,
        } => {
            if term.shape != ShapeConstraint::None {
                crate::bail_invalid_basis!(
                    "ShapeConstraint::{:?} for term '{}' is not supported on Pca basis",
                    term.shape,
                    term.name
                );
            }
            build_pca_smooth_basis(
                data,
                feature_cols,
                basis_matrix,
                *centered,
                *smooth_penalty,
                center_mean.as_ref(),
                pca_basis_path.as_ref(),
                *chunk_size,
            )?
        }
        SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
            build_tensor_bspline_basis(data, feature_cols, spec)?
        }
        SmoothBasisSpec::ByVariable { .. } => {
            crate::bail_invalid_basis!(
                "internal: ByVariable smooths must return before inner basis dispatch"
            );
        }
        SmoothBasisSpec::BySmooth { .. } => {
            crate::bail_invalid_basis!("internal: BySmooth smooths must be lowered to ByVariable before inner basis dispatch"
                    .to_string(),);
        }
        SmoothBasisSpec::FactorSmooth { spec } => {
            if term.shape != ShapeConstraint::None {
                crate::bail_invalid_basis!(
                    "ShapeConstraint::{:?} is unsupported for factor smooth term '{}'",
                    term.shape,
                    term.name
                );
            }
            return build_factor_smooth(data, spec, &term.name, workspace);
        }
    };

    if let SmoothBasisSpec::Matern { .. } = &term.basis {
        let (penalties, nullspace_dims, penaltyinfo) =
            matern_operator_penalty_triplet_from_metadata(&built.metadata)?;
        built.penalties = penalties;
        built.nullspace_dims = nullspace_dims;
        built.penaltyinfo = penaltyinfo;
    }

    let p_local = built.design.ncols();
    let mut metadata = built.metadata.clone();
    // Extract factored Kronecker representation before consuming fields.
    // Invalidate it if shape transforms will be applied (they break structure).
    let kron_factored = if term.shape == ShapeConstraint::None {
        built.kronecker_factored
    } else {
        None
    };
    let mut design_t = built.design;
    let mut penalties_t: Vec<Array2<f64>> = built.penalties;
    // Ops vector parallel to `penalties_t`. Survives unchanged through the
    // identity path; nulled element-wise when `T^T S T` reparametrization
    // is applied (operator no longer bit-equivalent to the transformed
    // matrix); wrapped in `ScaledPenaltyOp` after Frobenius normalization.
    let mut ops_t: Vec<Option<std::sync::Arc<dyn crate::terms::penalty_op::PenaltyOp>>> = built.ops;
    if matches!(
        spatial_identifiability_policy(term),
        Some(SpatialIdentifiability::OrthogonalToParametric)
    ) {
        metadata = freeze_raw_spatial_metadata(metadata, design_t.ncols());
    }

    let active_penaltyinfo_t = built
        .penaltyinfo
        .iter()
        .filter(|info| info.active)
        .cloned()
        .collect::<Vec<_>>();
    let pre_dropped_penaltyinfo_t = built
        .penaltyinfo
        .iter()
        .filter(|info| !info.active)
        .cloned()
        .collect::<Vec<_>>();
    let use_box_reparam =
        term.shape != ShapeConstraint::None && shape_uses_box_reparameterization(&term.basis);
    let mut coefficient_transform_for_constraints: Option<Array2<f64>> = None;
    if let Some((order, sign)) = shape_order_and_sign(term.shape)
        && use_box_reparam
    {
        let t = cumulative_sum_transform_matrix(p_local, order, sign);
        coefficient_transform_for_constraints = Some(t.clone());
        // Coefficient-side transform: wrap the design in an operator that
        // applies T on the coefficient side, preserving sparsity/operator
        // structure of the inner design.
        let inner_dense = match design_t {
            DesignMatrix::Dense(d) => d,
            DesignMatrix::Sparse(sp) => crate::matrix::DenseDesignMatrix::from(
                sp.try_to_dense_arc("shape-constrained coefficient transform")
                    .map_err(BasisError::InvalidInput)?,
            ),
        };
        let coeff_op = crate::matrix::CoefficientTransformOperator::new(inner_dense, t.clone())
            .map_err(|e| BasisError::InvalidInput(format!("CoefficientTransformOperator: {e}")))?;
        design_t = DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(Arc::new(coeff_op)));
        if penalties_t.len() != active_penaltyinfo_t.len() {
            crate::bail_invalid_basis!(
                "internal box-reparam penalty/info mismatch for term '{}': penalties={}, infos={}",
                term.name,
                penalties_t.len(),
                active_penaltyinfo_t.len()
            );
        }
        // Wiggliness penalties undergo the exact congruence `S → TᵀST` (PSD
        // preserving). The double-penalty *nullspace shrinkage* ridge must NOT:
        // it is a unit-eigenvalue projector `ZZᵀ` onto null(S_wiggle) in the
        // β (B-spline coefficient) coordinates, and the congruence
        // `Tᵀ(ZZᵀ)T = (TᵀZ)(TᵀZ)ᵀ` is no longer a projector — its eigenvalues
        // blow up by the conditioning of the cumulative-sum `T` (cond(T) grows
        // with the basis dim), concentrating an enormous penalty on the leading
        // γ₀ "level" coordinate. REML then drives the shared λ to its ceiling
        // and the smooth collapses to a flat constant (#509, the over-smoothing
        // face). The principled fix keeps mgcv's double-penalty semantics in the
        // *reparametrized* space: rebuild the ridge as the unit-eigenvalue
        // nullspace projector of the transformed wiggliness penalty `TᵀST`, so
        // the double penalty shrinks exactly the unpenalized polynomial
        // directions of the γ-space smooth with eigenvalue 1, identical in
        // conditioning to the unconstrained fit.
        let transformed_wiggliness = penalties_t
            .iter()
            .zip(active_penaltyinfo_t.iter())
            .find(|(_, info)| !matches!(info.source, PenaltySource::DoublePenaltyNullspace))
            .map(|(s_local, _)| {
                let tt_s = fast_atb(&t, s_local);
                fast_ab(&tt_s, &t)
            });
        let mut rebuilt = Vec::with_capacity(penalties_t.len());
        for (s_local, info) in penalties_t.iter().zip(active_penaltyinfo_t.iter()) {
            if matches!(info.source, PenaltySource::DoublePenaltyNullspace) {
                let s_wiggle_t = transformed_wiggliness.as_ref().ok_or_else(|| {
                    BasisError::InvalidInput(format!(
                        "box-reparam term '{}' has a double-penalty ridge but no primary wiggliness penalty to derive its nullspace from",
                        term.name
                    ))
                })?;
                let ridge = crate::terms::basis::build_nullspace_shrinkage_penalty(s_wiggle_t)?
                    .map(|shrink| shrink.sym_penalty)
                    .unwrap_or_else(|| Array2::<f64>::zeros((p_local, p_local)));
                rebuilt.push(ridge);
            } else {
                let tt_s = fast_atb(&t, s_local);
                rebuilt.push(fast_ab(&tt_s, &t));
            }
        }
        penalties_t = rebuilt;
        // T^T S T (and the rebuilt γ-space ridge) invalidate op-form
        // bit-equivalence; drop ops here.
        ops_t = vec![None; penalties_t.len()];
    }
    if penalties_t.len() != active_penaltyinfo_t.len() {
        crate::bail_invalid_basis!(
            "internal penalty metadata mismatch for term '{}': active penalties={}, active infos={}",
            term.name,
            penalties_t.len(),
            active_penaltyinfo_t.len()
        );
    }
    if ops_t.len() != penalties_t.len() {
        ops_t = vec![None; penalties_t.len()];
    }
    let penalty_candidates = penalties_t
        .into_iter()
        .zip(active_penaltyinfo_t.into_iter())
        .zip(ops_t.into_iter())
        .map(
            |((matrix, info), op_in)| -> Result<PenaltyCandidate, BasisError> {
                let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
                let preserve_margin_scale =
                    matches!(&info.source, PenaltySource::TensorMarginal { .. });
                let (matrix, normalization_scale, op_scale, kronecker_scale) =
                    if preserve_margin_scale {
                        (
                            matrix.mapv(|v| v * c_new),
                            info.normalization_scale,
                            1.0,
                            1.0,
                        )
                    } else {
                        (
                            matrix,
                            info.normalization_scale * c_new,
                            1.0 / c_new,
                            1.0 / c_new,
                        )
                    };
                // Frobenius rescale: wrap inner op in `ScaledPenaltyOp(1/c_new)`
                // so `op.as_dense() == matrix` post-normalization.
                let scaled_op = if op_scale > 0.0 && op_scale.is_finite() {
                    op_in.map(|op| {
                        std::sync::Arc::new(crate::terms::penalty_op::ScaledPenaltyOp::new(
                            op, op_scale,
                        ))
                            as std::sync::Arc<dyn crate::terms::penalty_op::PenaltyOp>
                    })
                } else {
                    None
                };
                let kronecker_factors = info.kronecker_factors.map(|mut factors| {
                    if let Some(first) = factors.first_mut() {
                        first.mapv_inplace(|v| v * kronecker_scale);
                    }
                    factors
                });
                Ok(PenaltyCandidate {
                    nullspace_dim_hint: info.nullspace_dim_hint,
                    matrix,
                    source: info.source,
                    normalization_scale,
                    kronecker_factors,
                    op: scaled_op,
                })
            },
        )
        .collect::<Result<Vec<_>, _>>()?;
    let (penalties_t, nullspaces_t, penaltyinfo_t, null_eigenvectors_t, ops_t) =
        crate::terms::basis::filter_active_penalty_candidates_with_ops(penalty_candidates)?;
    let shape_linear_constraints = if term.shape != ShapeConstraint::None && !use_box_reparam {
        let axis = shape_axis_col.ok_or_else(|| {
            BasisError::InvalidInput(format!(
                "internal shape-constraint axis missing for term '{}'",
                term.name
            ))
        })?;
        let (x_shape_eval, design_shape_eval) =
            build_shape_constraint_design_1d(data, term, &metadata, axis)?;
        build_shape_linear_constraints_1d(
            x_shape_eval.view(),
            design_shape_eval.view(),
            term.shape,
        )?
    } else {
        None
    };
    let boundary_linear_constraints = match &term.basis {
        SmoothBasisSpec::BSpline1D { spec, .. } => bspline_boundary_linear_constraints(
            spec.boundary_conditions,
            &metadata,
            spec.degree,
            coefficient_transform_for_constraints.as_ref(),
        )?,
        _ => None,
    };
    let linear_constraints_local =
        merge_linear_constraints_global(shape_linear_constraints, boundary_linear_constraints);

    // Joint-null absorption rotation. Fresh fit specs compute Q from the final
    // per-smooth penalty set (after all in-smooth reparameterizations have
    // already been applied). Frozen specs already carry the complete realized
    // coefficient chart in their `FrozenTransform`; recomputing Q there would
    // rotate an already-frozen chart a second time and desynchronize value
    // rebuilds from derivative operators.
    //
    // Kronecker-factored smooths (tensor B-splines under `TensorBSplineIdentifiability::None`)
    // carry their joint penalty as `Σ_d S_d` with `S_d = I ⊗ … ⊗ S_d^{1D} ⊗ … ⊗ I`.
    // The joint null space is the tensor of marginal nulls and is handled directly
    // by the REML runtime's `kronecker_penalty_system` path (see
    // `runtime.rs:8334-8344`). Applying a dense (p × p) Q here would densify
    // `X_raw = mx ⊗ my` into `X_raw · Q`, destroying the Kronecker product
    // structure that the runtime relies on for fast log-det/derivative
    // assembly — and the rotation block at the wrapper site also unconditionally
    // wipes `kronecker_factored`, leaving the runtime to fall back to the
    // dense per-block log-det. Skip the rotation for Kronecker-factored terms
    // so the factored representation survives end-to-end.
    let joint_null_rotation = match term.joint_null_rotation.clone() {
        Some(persisted) => Some(persisted),
        None if smooth_has_frozen_identifiability(term) => None,
        None if kron_factored.is_some() => None,
        None => crate::terms::basis::compute_joint_null_rotation(&penalties_t)?,
    };

    Ok(LocalSmoothTermBuild {
        dim: p_local,
        design: design_t,
        penalties: penalties_t,
        ops: ops_t,
        nullspaces: nullspaces_t,
        null_eigenvectors: null_eigenvectors_t,
        joint_null_rotation,
        penaltyinfo: penaltyinfo_t,
        pre_dropped_penaltyinfo: pre_dropped_penaltyinfo_t,
        metadata,
        linear_constraints: linear_constraints_local,
        box_reparam: use_box_reparam,
        kronecker_factored: kron_factored,
    })
}

pub fn build_smooth_design(
    data: ArrayView2<'_, f64>,
    terms: &[SmoothTermSpec],
) -> Result<RawSmoothDesign, BasisError> {
    let mut ws = crate::basis::BasisWorkspace::new();
    build_smooth_design_withworkspace(data, terms, &mut ws)
}

/// Like `build_smooth_design`, but honors the caller workspace policy while
/// building each planned smooth term with an independent per-term workspace.
///
/// Independent workspaces avoid shared mutable distance-cache state during the
/// parallel term build; the final design, penalties, and metadata are assembled
/// in the original smooth-term order.
pub fn build_smooth_design_withworkspace(
    data: ArrayView2<'_, f64>,
    terms: &[SmoothTermSpec],
    workspace: &mut crate::basis::BasisWorkspace,
) -> Result<RawSmoothDesign, BasisError> {
    validate_smooth_terms_finite_inputs(data, terms)?;
    build_smooth_design_withworkspace_unvalidated(data, terms, workspace)
}

fn build_smooth_design_withworkspace_unvalidated(
    data: ArrayView2<'_, f64>,
    terms: &[SmoothTermSpec],
    workspace: &mut crate::basis::BasisWorkspace,
) -> Result<RawSmoothDesign, BasisError> {
    let mut planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &[terms.to_vec()])?;
    let planned_terms = planned_blocks.pop().ok_or_else(|| {
        BasisError::InvalidInput(
            "joint spatial center planner returned no smooth blocks".to_string(),
        )
    })?;
    let policy = workspace.policy().clone();
    let local_builds: Vec<LocalSmoothTermBuild> = {
        use rayon::iter::{IntoParallelIterator, ParallelIterator};
        planned_terms
            .into_par_iter()
            .map(|term| {
                let mut term_workspace = crate::basis::BasisWorkspace::with_policy(policy.clone());
                build_single_local_smooth_term(data, &term, &mut term_workspace)
            })
            .collect::<Result<Vec<_>, _>>()?
    };

    let total_p: usize = local_builds.iter().map(|built| built.dim).sum();

    let mut local_designs: Vec<DesignMatrix> = Vec::with_capacity(local_builds.len());
    let mut terms_out = Vec::<SmoothTerm>::with_capacity(terms.len());
    let mut penalties_global = Vec::<BlockwisePenalty>::new();
    let mut nullspace_dims_global = Vec::<usize>::new();
    let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
    let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
    let mut any_bounds = false;
    // Each linear-constraint row only touches the current term's column slice.
    // Track `(col_start, col_end, local_row_values)` and assemble the final
    // dense `Array2` in one pass, avoiding per-row `Array1::zeros(total_p)`
    // allocation plus a row-by-row copy at the end.
    let mut linear_constraintsrows: Vec<(usize, usize, Array1<f64>)> = Vec::new();
    let mut linear_constraints_b: Vec<f64> = Vec::new();

    let mut col_start = 0usize;
    for (term, mut built) in terms.iter().zip(local_builds.into_iter()) {
        let p_local = built.dim;
        let col_end = col_start + p_local;
        let lb_local = if built.box_reparam {
            shape_lower_bounds_local(term.shape, p_local)
        } else {
            None
        };

        // Stage-2 joint-null absorption rotation. Fired *before* the
        // penalty / design / global aggregation loops below so that every
        // subsequent reference to `built.penalties`, `built.design`, and
        // `built.ops` sees the post-rotation values.
        //
        // The math: when the smooth's joint penalty `Σ_k S_k` has a
        // non-trivial null space, eigh selects `Q = [U_range | U_null]`
        // with null columns at the tail. Setting `β_raw = Q · γ` and
        // applying:
        //     design        ← X · Q
        //     penalties[k]  ← Qᵀ · S_k · Q   (block-diag, zero null tail)
        // yields a model whose fitted γ is invariant to the rotation
        // (since likelihood depends only on `X · β_raw = X · Q · γ`), but
        // whose penalty is full-rank on the range columns. The biobank
        // failing case (cert refusal in the joint-Newton inner solve)
        // resolves because `H_pen = H_loglik + S` becomes full rank on
        // the smooth's range columns.
        //
        // Rotation is suppressed when the smooth carries coordinate-wise
        // shape constraints (`lb_local` or `built.linear_constraints`):
        // those encode a cone in the original coordinate system and a
        // general orthogonal rotation breaks the cone geometry. Smooths
        // with shape constraints typically have full-rank joint penalty
        // (their structural shape comes from the cone, not from null
        // directions in the penalty), so suppression is rarely a loss.
        //
        // `applied_rotation` carries the Q that was applied (or `None`
        // if no rotation fired). It is persisted onto `SmoothTerm` below
        // so prediction-side `X_new_raw · Q` replay can reproduce the
        // exact rotation. Persistence through the saved-model artifact
        // is a follow-up — see the doc on `SmoothTerm.joint_null_rotation`.
        let applied_rotation: Option<crate::terms::basis::JointNullRotation> = match (
            built.joint_null_rotation.take(),
            lb_local.is_some(),
            built.linear_constraints.is_some(),
        ) {
            (Some(rot), false, false) => {
                let q = &rot.rotation;
                let dense = built
                    .design
                    .try_to_dense_by_chunks("joint-null absorption rotation")
                    .map_err(BasisError::InvalidInput)?;
                let rotated = crate::linalg::faer_ndarray::fast_ab(&dense, q);
                built.design = DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(rotated));
                built.penalties = built
                    .penalties
                    .into_iter()
                    .map(|s_local| {
                        let qt_s = crate::linalg::faer_ndarray::fast_atb(q, &s_local);
                        crate::linalg::faer_ndarray::fast_ab(&qt_s, q)
                    })
                    .collect();
                built.ops = vec![None; built.penalties.len()];
                built.kronecker_factored = None;
                Some(rot)
            }
            (Some(_), _, _) => None,
            (None, _, _) => None,
        };

        let activeinfos = built
            .penaltyinfo
            .iter()
            .filter(|info| info.active)
            .collect::<Vec<_>>();
        if activeinfos.len() != built.penalties.len() {
            crate::bail_invalid_basis!(
                "internal penalty info mismatch for term '{}': activeinfos={}, penalties={}",
                term.name,
                activeinfos.len(),
                built.penalties.len()
            );
        }
        for (((s_local, &ns), info), op_local) in built
            .penalties
            .iter()
            .zip(built.nullspaces.iter())
            .zip(activeinfos.into_iter())
            .zip(built.ops.iter())
        {
            let global_index = penalties_global.len();
            penalties_global.push(
                BlockwisePenalty::new(col_start..col_end, s_local.clone())
                    .with_op(op_local.clone()),
            );
            nullspace_dims_global.push(ns);
            let mut penalty = info.clone();
            penalty.nullspace_dim_hint = ns;
            penaltyinfo_global.push(PenaltyBlockInfo {
                global_index,
                termname: Some(term.name.clone()),
                penalty,
            });
        }
        for info in built.penaltyinfo.iter().filter(|info| !info.active) {
            dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
                termname: Some(term.name.clone()),
                penalty: info.clone(),
            });
        }
        for info in &built.pre_dropped_penaltyinfo {
            dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
                termname: Some(term.name.clone()),
                penalty: info.clone(),
            });
        }

        if let Some(lin_local) = &built.linear_constraints {
            for r in 0..lin_local.a.nrows() {
                linear_constraintsrows.push((col_start, col_end, lin_local.a.row(r).to_owned()));
                linear_constraints_b.push(lin_local.b[r]);
            }
        }
        if let Some(lb_local) = &lb_local {
            coefficient_lower_bounds
                .slice_mut(s![col_start..col_end])
                .assign(lb_local);
            any_bounds = true;
        }

        // Move the per-term design out of `built` rather than cloning it.
        local_designs.push(built.design);

        terms_out.push(SmoothTerm {
            name: term.name.clone(),
            coeff_range: col_start..col_end,
            shape: term.shape,
            penalties_local: built.penalties,
            nullspace_dims: built.nullspaces,
            penaltyinfo_local: built.penaltyinfo,
            metadata: built.metadata,
            lower_bounds_local: lb_local,
            linear_constraints_local: built.linear_constraints,
            kronecker_factored: built.kronecker_factored.take(),
            joint_null_rotation: applied_rotation,
        });

        col_start = col_end;
    }

    assert_eq!(
        penalties_global.len(),
        nullspace_dims_global.len(),
        "global smooth penalty/nullspace bookkeeping diverged"
    );
    assert_eq!(
        penalties_global.len(),
        penaltyinfo_global.len(),
        "global smooth penalty metadata bookkeeping diverged"
    );

    Ok(RawSmoothDesign {
        term_designs: local_designs,
        penalties: penalties_global,
        nullspace_dims: nullspace_dims_global,
        penaltyinfo: penaltyinfo_global,
        dropped_penaltyinfo: dropped_penaltyinfo_global,
        terms: terms_out,
        coefficient_lower_bounds: if any_bounds {
            Some(coefficient_lower_bounds)
        } else {
            None
        },
        linear_constraints: if linear_constraintsrows.is_empty() {
            None
        } else {
            let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
            for (i, (cs, ce, values)) in linear_constraintsrows.iter().enumerate() {
                a.row_mut(i).slice_mut(s![*cs..*ce]).assign(values);
            }
            Some(LinearInequalityConstraints {
                a,
                b: Array1::from_vec(linear_constraints_b),
            })
        },
    })
}

fn build_term_collection_design_inner(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
) -> Result<TermCollectionDesign, BasisError> {
    use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};

    let n = data.nrows();
    let p_data = data.ncols();
    let p_intercept = 1usize;
    let p_lin = spec.linear_terms.len();

    // Smooth construction, random-effect construction, and linear-column
    // extraction are independent at this stage. Run them concurrently, but keep
    // each result in spec order so the final global layout remains stable:
    // [intercept | linear | random_effects | smooth].
    let (smooth_raw_result, (random_blocks_result, linear_block_result)) = rayon::join(
        || {
            let mut ws = crate::basis::BasisWorkspace::new();
            build_smooth_design_withworkspace_unvalidated(data, &spec.smooth_terms, &mut ws)
        },
        || {
            rayon::join(
                || {
                    spec.random_effect_terms
                        .par_iter()
                        .map(|term| build_random_effect_block(data, term))
                        .collect::<Result<Vec<_>, _>>()
                },
                || -> Result<Option<Array2<f64>>, BasisError> {
                    if p_lin == 0 {
                        return Ok(None);
                    }

                    let linear_columns = (0..p_lin)
                        .into_par_iter()
                        .map(|j| {
                            let linear = &spec.linear_terms[j];
                            if linear.feature_col >= p_data {
                                crate::bail_dim_basis!(
                                    "linear term '{}' feature column {} out of bounds for {} columns",
                                    linear.name, linear.feature_col, p_data
                                );
                            }
                            Ok(data.column(linear.feature_col).to_owned())
                        })
                        .collect::<Result<Vec<_>, _>>()?;

                    let mut out = Array2::<f64>::zeros((n, p_lin));
                    for (j, column) in linear_columns.iter().enumerate() {
                        out.column_mut(j).assign(column);
                    }
                    Ok(Some(out))
                },
            )
        },
    );

    let smooth_raw = smooth_raw_result?;
    let random_blocks = random_blocks_result?;
    let linear_block = linear_block_result?;

    let smooth = apply_global_smooth_identifiability(
        smooth_raw,
        data,
        &spec.linear_terms,
        &spec.smooth_terms,
    )?;

    let p_rand: usize = random_blocks.iter().map(|b| b.num_groups).sum();
    let p_smooth = smooth.total_smooth_cols();
    let p_total = p_intercept + p_lin + p_rand + p_smooth;

    let mut linear_ranges = Vec::<(String, Range<usize>)>::with_capacity(p_lin);
    for (j, linear) in spec.linear_terms.iter().enumerate() {
        let col = p_intercept + j;
        // Column ranges are in the global (full) coordinate system:
        // [intercept | linear | random_effects | smooth]
        linear_ranges.push((linear.name.clone(), col..(col + 1)));
    }

    // Track random-effect column ranges in the global coordinate system.
    // Global layout: [intercept(1) | linear(p_lin) | RE_0(q0) | RE_1(q1) | … | smooth(p_smooth)]
    let mut random_effect_ranges =
        Vec::<(String, Range<usize>)>::with_capacity(random_blocks.len());
    let mut random_effect_levels = Vec::<(String, Vec<u64>)>::with_capacity(random_blocks.len());
    let mut col_cursor = p_intercept + p_lin;
    for block in &random_blocks {
        let q = block.num_groups;
        let end = col_cursor + q;
        random_effect_ranges.push((block.name.clone(), col_cursor..end));
        random_effect_levels.push((block.name.clone(), block.kept_levels.clone()));
        col_cursor = end;
    }

    // ── Assemble the full DesignMatrix ────────────────────────────────
    //
    // Always use a BlockDesignOperator with per-term blocks.  The full
    // (n, p_total) dense matrix is NEVER materialized:
    //
    //   Block 0:     Intercept — zero storage, implicit all-ones column
    //   Block 1:     Linear terms — (n, p_lin) extracted from data
    //   Blocks 2..k: Random-effect operators — O(n) one-hot, no dense storage
    //   Blocks k+1..: Per-smooth-term dense blocks — each (n, p_term)
    //
    // Splitting smooth terms into per-term blocks means cross-block grams
    // are small O(p_i × p_j) BLAS operations.  Tensor product terms with
    // Kronecker structure become DesignBlock::Operator, avoiding the full
    // n × ∏q_j materialization.

    let mut blocks = Vec::<DesignBlock>::new();

    // Block 0: intercept — zero storage.
    blocks.push(DesignBlock::Intercept(n));

    // Block 1: linear terms.
    if let Some(lin_block) = linear_block {
        blocks.push(DesignBlock::Dense(crate::matrix::DenseDesignMatrix::from(
            lin_block,
        )));
    }

    // Blocks: random-effect operators — O(n) implicit one-hot.
    for block in &random_blocks {
        let re_op = RandomEffectOperator::new(block.group_ids.clone(), block.num_groups);
        blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
    }

    // Blocks: per-smooth-term.  Each smooth term gets its own block so that
    // cross-block grams are tiny. Tensor terms can stay operator-backed all
    // the way from basis construction; dense smooth terms stay dense.
    if p_smooth > 0 {
        for term_design in &smooth.term_designs {
            match term_design {
                DesignMatrix::Dense(dense) => blocks.push(DesignBlock::Dense(dense.clone())),
                DesignMatrix::Sparse(sparse) => blocks.push(DesignBlock::Sparse(sparse.clone())),
            }
        }
    }

    let design = assemble_term_collection_design_matrix(blocks)?;

    let mut penalties = Vec::<BlockwisePenalty>::new();
    let mut nullspace_dims = Vec::<usize>::new();
    let mut penaltyinfo = Vec::<PenaltyBlockInfo>::new();
    let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
    let mut any_bounds = false;
    let mut linear_constraintrows = Vec::<Array1<f64>>::new();
    let mut linear_constraint_b = Vec::<f64>::new();

    for (j, linear) in spec.linear_terms.iter().enumerate() {
        let col = p_intercept + j;
        if let Some(lb) = linear.coefficient_min {
            let mut row = Array1::<f64>::zeros(p_total);
            row[col] = 1.0;
            linear_constraintrows.push(row);
            linear_constraint_b.push(lb);
        }
        if let Some(ub) = linear.coefficient_max {
            let mut row = Array1::<f64>::zeros(p_total);
            row[col] = -1.0;
            linear_constraintrows.push(row);
            linear_constraint_b.push(-ub);
        }
    }

    // Linear terms with `double_penalty=true` are aggregated into a single
    // shared ridge block named "linear" that spans the union of their
    // coefficient columns. Two reasons this is the only correct shape:
    //
    //   (i) Per-term redundancy: a penalized linear term is a single scalar
    //   coefficient β_j; the only shrinkage direction is β_j² itself. The
    //   earlier per-term form emitted two identical 1×1 ridges per term
    //   (`LinearTermRidge` + `LinearTermDoubleRidge`); for the two λ's that
    //   the outer REML/LAML places on those blocks, the total penalty is
    //   `(λ₁ + λ₂) β_j²` and the outer score is flat along
    //   `λ₁ + λ₂ = const`. Unidentifiable hyperparameters by construction.
    //
    //   (ii) Cross-term aliasing: each penalized linear term in its own 1×1
    //   block gives the outer optimizer a separate λ per term, when the
    //   mgcv-equivalent "ridge over the parametric linear columns" is one
    //   shared λ that shrinks every penalized linear coefficient with the
    //   same scale. The shared block emits exactly one outer hyperparameter
    //   and pins each penalized column's diagonal entry to 1.0; unpenalized
    //   linear columns sit at diagonal 0, so they are mathematically
    //   unaffected (S β = 0 on those rows) while keeping the block range
    //   contiguous over the linear coefficient run.
    //
    // The block's `termname` is the literal "linear" so keyed hyperpriors
    // / coefficient groups address the shared lever directly. The
    // per-term identity is preserved through `linear_ranges`, which still
    // names each individual column.
    let any_double_penalty_linear = spec.linear_terms.iter().any(|t| t.double_penalty);
    if any_double_penalty_linear && p_lin > 0 {
        let block_range = p_intercept..(p_intercept + p_lin);
        let mut s = Array2::<f64>::zeros((p_lin, p_lin));
        let mut effective_rank = 0usize;
        for (j, linear) in spec.linear_terms.iter().enumerate() {
            if linear.double_penalty {
                s[[j, j]] = 1.0;
                effective_rank += 1;
            }
        }
        let global_index = penalties.len();
        penalties.push(BlockwisePenalty::new(block_range, s));
        nullspace_dims.push(0);
        penaltyinfo.push(PenaltyBlockInfo {
            global_index,
            termname: Some("linear".to_string()),
            penalty: PenaltyInfo {
                source: PenaltySource::Other("LinearTermRidge".to_string()),
                original_index: 0,
                active: true,
                effective_rank,
                dropped_reason: None,
                nullspace_dim_hint: 0,
                normalization_scale: 1.0,
                kronecker_factors: None,
            },
        });
    }

    for (re_idx, (name, range)) in random_effect_ranges.iter().enumerate() {
        if range.is_empty() || !spec.random_effect_terms[re_idx].penalized {
            continue;
        }
        let block_size = range.len();
        let global_index = penalties.len();
        penalties.push(BlockwisePenalty::ridge(range.clone(), 1.0));
        nullspace_dims.push(0);
        penaltyinfo.push(PenaltyBlockInfo {
            global_index,
            termname: Some(name.clone()),
            penalty: PenaltyInfo {
                source: PenaltySource::Other(format!("RandomEffectRidge({name})")),
                original_index: re_idx,
                active: true,
                effective_rank: block_size,
                dropped_reason: None,
                nullspace_dim_hint: 0,
                normalization_scale: 1.0,
                kronecker_factors: None,
            },
        });
    }

    if smooth.penaltyinfo.len() != smooth.penalties.len() {
        crate::bail_invalid_basis!(
            "smooth penalty metadata mismatch: penalties={}, metadata={}",
            smooth.penalties.len(),
            smooth.penaltyinfo.len()
        );
    }
    let smooth_start = p_intercept + p_lin + p_rand;
    for ((bp_smooth, &ns), localinfo) in smooth
        .penalties
        .iter()
        .zip(smooth.nullspace_dims.iter())
        .zip(smooth.penaltyinfo.iter())
    {
        let global_index = penalties.len();
        // Offset the per-term block range from smooth-local to model-global.
        let offset_range =
            (bp_smooth.col_range.start + smooth_start)..(bp_smooth.col_range.end + smooth_start);
        let bp = if let Some(factors) = localinfo.penalty.kronecker_factors.as_ref() {
            BlockwisePenalty::kronecker(offset_range, bp_smooth.local.clone(), factors.clone())
                .with_op(bp_smooth.op.clone())
        } else if matches!(
            localinfo.penalty.source,
            PenaltySource::Other(ref s) if s.starts_with("RandomEffectRidge")
        ) {
            BlockwisePenalty::ridge(offset_range, 1.0)
        } else {
            BlockwisePenalty::new(offset_range, bp_smooth.local.clone())
                .with_op(bp_smooth.op.clone())
        };
        penalties.push(bp);
        nullspace_dims.push(ns);
        let mut penalty = localinfo.penalty.clone();
        penalty.nullspace_dim_hint = ns;
        penaltyinfo.push(PenaltyBlockInfo {
            global_index,
            termname: localinfo.termname.clone(),
            penalty,
        });
    }
    dropped_penaltyinfo.extend(smooth.dropped_penaltyinfo.iter().cloned());

    assert_eq!(
        penalties.len(),
        nullspace_dims.len(),
        "term-collection penalty/nullspace bookkeeping diverged"
    );
    assert_eq!(
        penalties.len(),
        penaltyinfo.len(),
        "term-collection penalty metadata bookkeeping diverged"
    );

    if let Some(lb_smooth) = smooth.coefficient_lower_bounds.as_ref() {
        let start = p_intercept + p_lin + p_rand;
        coefficient_lower_bounds
            .slice_mut(s![start..(start + p_smooth)])
            .assign(lb_smooth);
        any_bounds = true;
    }
    if let Some(lin_smooth) = smooth.linear_constraints.as_ref() {
        let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
        let start = p_intercept + p_lin + p_rand;
        a_global
            .slice_mut(s![.., start..(start + p_smooth)])
            .assign(&lin_smooth.a);
        for r in 0..a_global.nrows() {
            linear_constraintrows.push(a_global.row(r).to_owned());
            linear_constraint_b.push(lin_smooth.b[r]);
        }
    }

    // Canonical constraint path: convert any explicit lower bounds into linear
    // inequalities and merge into the global constraint matrix. This keeps fitting
    // behavior independent of user-facing lower-bound options.
    let lower_bound_constraints = if any_bounds {
        linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
    } else {
        None
    };
    let explicit_linear_constraints = if linear_constraintrows.is_empty() {
        None
    } else {
        let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
        for (i, row) in linear_constraintrows.iter().enumerate() {
            a.row_mut(i).assign(row);
        }
        Some(LinearInequalityConstraints {
            a,
            b: Array1::from_vec(linear_constraint_b),
        })
    };
    let linear_constraints =
        merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints);

    Ok(TermCollectionDesign {
        design,
        penalties,
        nullspace_dims,
        penaltyinfo,
        dropped_penaltyinfo,
        coefficient_lower_bounds: if any_bounds {
            Some(coefficient_lower_bounds)
        } else {
            None
        },
        linear_constraints,
        intercept_range: 0..1,
        linear_ranges,
        random_effect_ranges,
        random_effect_levels,
        smooth,
    })
}

pub fn build_term_collection_design(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
) -> Result<TermCollectionDesign, BasisError> {
    validate_term_collection_finite_inputs(data, spec)?;
    let mut planned_specs =
        plan_joint_spatial_centers_for_term_blocks(data, &[spec.smooth_terms.clone()])?;
    let planned_smooth_terms = planned_specs.pop().ok_or_else(|| {
        BasisError::InvalidInput(
            "joint spatial center planner returned no smooth terms for single-spec build"
                .to_string(),
        )
    })?;
    let mut planned_spec = spec.clone();
    planned_spec.smooth_terms = planned_smooth_terms;
    build_term_collection_design_inner(data, &planned_spec)
}

pub fn build_term_collection_designs_joint(
    data: ArrayView2<'_, f64>,
    specs: &[TermCollectionSpec],
) -> Result<Vec<TermCollectionDesign>, BasisError> {
    for spec in specs {
        validate_term_collection_finite_inputs(data, spec)?;
    }
    let smooth_blocks = specs
        .iter()
        .map(|spec| spec.smooth_terms.clone())
        .collect::<Vec<_>>();
    let planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &smooth_blocks)?;
    let mut out = Vec::with_capacity(specs.len());
    for (spec, planned_terms) in specs.iter().zip(planned_blocks.into_iter()) {
        let mut planned_spec = spec.clone();
        planned_spec.smooth_terms = planned_terms;
        out.push(build_term_collection_design_inner(data, &planned_spec)?);
    }
    Ok(out)
}

pub fn build_term_collection_designs_and_freeze_joint(
    data: ArrayView2<'_, f64>,
    specs: &[TermCollectionSpec],
) -> Result<(Vec<TermCollectionDesign>, Vec<TermCollectionSpec>), EstimationError> {
    let designs = build_term_collection_designs_joint(data, specs)?;
    let mut resolved_specs = Vec::with_capacity(specs.len());
    for (spec, design) in specs.iter().zip(designs.iter()) {
        resolved_specs.push(freeze_term_collection_from_design(spec, design)?);
    }
    Ok((designs, resolved_specs))
}

fn smooth_basis_feature_cols(basis: &SmoothBasisSpec) -> Vec<usize> {
    match basis {
        SmoothBasisSpec::ByVariable { inner, by_col, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
            let mut cols = smooth_basis_feature_cols(inner);
            cols.push(*by_col);
            cols.sort_unstable();
            cols.dedup();
            cols
        }
        SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_feature_cols(smooth),
        SmoothBasisSpec::BSpline1D { feature_col, .. } => vec![*feature_col],
        SmoothBasisSpec::ThinPlate { feature_cols, .. }
        | SmoothBasisSpec::Sphere { feature_cols, .. }
        | SmoothBasisSpec::Matern { feature_cols, .. }
        | SmoothBasisSpec::Duchon { feature_cols, .. }
        | SmoothBasisSpec::Pca { feature_cols, .. }
        | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => feature_cols.clone(),
        SmoothBasisSpec::FactorSmooth { spec } => {
            let mut cols = spec.continuous_cols.clone();
            cols.push(spec.group_col);
            cols.sort_unstable();
            cols.dedup();
            cols
        }
    }
}

pub fn smooth_term_feature_cols(term: &SmoothTermSpec) -> Vec<usize> {
    smooth_basis_feature_cols(&term.basis)
}

fn smooth_basis_family_rank(term: &SmoothTermSpec) -> u8 {
    match &term.basis {
        SmoothBasisSpec::ByVariable { inner, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
            smooth_basis_family_rank(&SmoothTermSpec {
                name: term.name.clone(),
                basis: (**inner).clone(),
                shape: term.shape,
                joint_null_rotation: None,
            })
        }
        SmoothBasisSpec::BSpline1D { .. } => 0,
        SmoothBasisSpec::TensorBSpline { .. } => 1,
        SmoothBasisSpec::ThinPlate { .. } => 2,
        SmoothBasisSpec::Sphere { .. } => 3,
        SmoothBasisSpec::Matern { .. } => 4,
        SmoothBasisSpec::Duchon { .. } => 5,
        SmoothBasisSpec::Pca { .. } => 6,
        SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_family_rank(&SmoothTermSpec {
            name: term.name.clone(),
            basis: (**smooth).clone(),
            shape: term.shape,
            joint_null_rotation: None,
        }),
        SmoothBasisSpec::FactorSmooth { .. } => 7,
    }
}

fn smooth_has_frozen_identifiability(term: &SmoothTermSpec) -> bool {
    match &term.basis {
        SmoothBasisSpec::ByVariable { inner, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
            smooth_has_frozen_identifiability(&SmoothTermSpec {
                name: term.name.clone(),
                basis: (**inner).clone(),
                shape: term.shape,
                joint_null_rotation: None,
            })
        }
        SmoothBasisSpec::BSpline1D { spec, .. } => {
            matches!(
                spec.identifiability,
                BSplineIdentifiability::FrozenTransform { .. }
            )
        }
        SmoothBasisSpec::ThinPlate { spec, .. } => matches!(
            spec.identifiability,
            SpatialIdentifiability::FrozenTransform { .. }
        ),
        SmoothBasisSpec::Sphere { spec, .. } => {
            matches!(spec.center_strategy, CenterStrategy::UserProvided(_))
                || matches!(
                    spec.identifiability,
                    SphericalSplineIdentifiability::FrozenTransform { .. }
                )
        }
        SmoothBasisSpec::Matern { spec, .. } => matches!(
            spec.identifiability,
            MaternIdentifiability::FrozenTransform { .. }
        ),
        SmoothBasisSpec::BySmooth { by_kind, .. } => match by_kind {
            ByVarKind::Factor { frozen_levels, .. } => frozen_levels.is_some(),
            ByVarKind::Numeric { .. } => true,
        },
        SmoothBasisSpec::FactorSmooth { spec } => spec.group_frozen_levels.is_some(),
        SmoothBasisSpec::Duchon { spec, .. } => matches!(
            spec.identifiability,
            SpatialIdentifiability::FrozenTransform { .. }
        ),
        SmoothBasisSpec::Pca {
            centered,
            center_mean,
            pca_basis_path,
            ..
        } => !*centered || center_mean.is_some() || pca_basis_path.is_some(),
        SmoothBasisSpec::TensorBSpline { spec, .. } => matches!(
            spec.identifiability,
            TensorBSplineIdentifiability::FrozenTransform { .. }
        ),
    }
}

fn compare_smooth_ownership_priority(
    lhs_idx: usize,
    lhs: &SmoothTermSpec,
    rhs_idx: usize,
    rhs: &SmoothTermSpec,
) -> std::cmp::Ordering {
    let lhs_cols = smooth_term_feature_cols(lhs);
    let rhs_cols = smooth_term_feature_cols(rhs);
    lhs_cols
        .len()
        .cmp(&rhs_cols.len())
        .then_with(|| lhs_cols.cmp(&rhs_cols))
        .then_with(|| smooth_basis_family_rank(lhs).cmp(&smooth_basis_family_rank(rhs)))
        .then_with(|| lhs.name.cmp(&rhs.name))
        .then(lhs_idx.cmp(&rhs_idx))
}

fn smooth_is_owned_by_prior_term(owner: &SmoothTermSpec, target: &SmoothTermSpec) -> bool {
    let owner_features = smooth_term_feature_cols(owner)
        .into_iter()
        .collect::<BTreeSet<_>>();
    let target_features = smooth_term_feature_cols(target)
        .into_iter()
        .collect::<BTreeSet<_>>();
    owner_features.is_subset(&target_features)
}

/// Static (spec-only) description of the hierarchical smooth-ownership decomposition.
///
/// This is the single source of truth for the deterministic ownership policy that
/// [`apply_global_smooth_identifiability`] uses during the fit: the processing order of
/// smooth terms, the feature columns each term spans, the candidate lower-order owners of
/// each term (nested/duplicate feature sets), and the basis-family rank used as a
/// tie-breaker. The fit engine consumes this structure and additionally applies a numerical
/// cross-residual overlap test on the realized design columns; the CLI structure-warning
/// path consumes the same structure for diagnostic messages, so both paths agree on which
/// smooths own which subspaces.
pub struct SmoothStructureAnalysis {
    /// Smooth-term indices sorted into ownership-processing order (lowest priority first):
    /// lower-order / narrower smooths come first and own their subspaces.
    pub ownership_order: Vec<usize>,
    /// `term_feature_cols[idx]` are the sorted, deduplicated feature columns that smooth term
    /// `idx` spans (indexed by the original smooth-term index, not by `ownership_order`).
    pub term_feature_cols: Vec<Vec<usize>>,
    /// `term_owners[idx]` are the indices of prior (in `ownership_order`) smooth terms whose
    /// feature set is a subset of term `idx`'s feature set, i.e. candidate owners of `idx`.
    /// The list is given in ownership-processing order.
    pub term_owners: Vec<Vec<usize>>,
    /// `basis_family_ranks[idx]` is the basis-family ordering rank of smooth term `idx`.
    pub basis_family_ranks: Vec<u8>,
}

/// Compute the static hierarchical smooth-ownership decomposition from the smooth-term specs.
///
/// `smoothspecs` is the same slice that [`apply_global_smooth_identifiability`] receives.
pub fn analyze_smooth_ownership(smoothspecs: &[SmoothTermSpec]) -> SmoothStructureAnalysis {
    let term_feature_cols: Vec<Vec<usize>> =
        smoothspecs.iter().map(smooth_term_feature_cols).collect();
    let basis_family_ranks: Vec<u8> = smoothspecs.iter().map(smooth_basis_family_rank).collect();

    let mut ownership_order: Vec<usize> = (0..smoothspecs.len()).collect();
    ownership_order.sort_by(|&lhs, &rhs| {
        compare_smooth_ownership_priority(lhs, &smoothspecs[lhs], rhs, &smoothspecs[rhs])
    });

    let mut term_owners = vec![Vec::<usize>::new(); smoothspecs.len()];
    for (pos, &target_idx) in ownership_order.iter().enumerate() {
        let target = &smoothspecs[target_idx];
        term_owners[target_idx] = ownership_order[..pos]
            .iter()
            .copied()
            .filter(|&owner_idx| smooth_is_owned_by_prior_term(&smoothspecs[owner_idx], target))
            .collect();
    }

    SmoothStructureAnalysis {
        ownership_order,
        term_feature_cols,
        term_owners,
        basis_family_ranks,
    }
}

fn build_constraint_block(
    n: usize,
    parametric_block: Option<&Array2<f64>>,
    owner_blocks: &[&DesignMatrix],
) -> Result<Array2<f64>, BasisError> {
    let param_cols = parametric_block.map_or(0, |mat| mat.ncols());
    let owner_cols: usize = owner_blocks.iter().map(|design| design.ncols()).sum();
    let mut block = Array2::<f64>::zeros((n, param_cols + owner_cols));
    let mut col_start = 0usize;
    if let Some(parametric) = parametric_block {
        let col_end = col_start + parametric.ncols();
        block
            .slice_mut(s![.., col_start..col_end])
            .assign(parametric);
        col_start = col_end;
    }
    const CHUNK: usize = 1024;
    for owner in owner_blocks {
        let col_end = col_start + owner.ncols();
        for row_start in (0..n).step_by(CHUNK) {
            let row_end = (row_start + CHUNK).min(n);
            let chunk = (*owner)
                .try_row_chunk(row_start..row_end)
                .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
            block
                .slice_mut(s![row_start..row_end, col_start..col_end])
                .assign(&chunk);
        }
        col_start = col_end;
    }
    Ok(block)
}

fn design_cross_relative_residual(
    lhs: &DesignMatrix,
    rhs: &DesignMatrix,
) -> Result<f64, BasisError> {
    let n = lhs.nrows();
    if rhs.nrows() != n {
        return Err(BasisError::ConstraintMatrixRowMismatch {
            basisrows: n,
            constraintrows: rhs.nrows(),
        });
    }
    const CHUNK: usize = 1024;
    let mut cross = Array2::<f64>::zeros((lhs.ncols(), rhs.ncols()));
    let mut lhs_sumsq = 0.0;
    let mut rhs_sumsq = 0.0;
    for start in (0..n).step_by(CHUNK) {
        let end = (start + CHUNK).min(n);
        let lhs_chunk = lhs
            .try_row_chunk(start..end)
            .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
        let rhs_chunk = rhs
            .try_row_chunk(start..end)
            .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
        cross += &crate::faer_ndarray::fast_atb(&lhs_chunk, &rhs_chunk);
        lhs_sumsq += lhs_chunk.iter().map(|v| v * v).sum::<f64>();
        rhs_sumsq += rhs_chunk.iter().map(|v| v * v).sum::<f64>();
    }
    let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
    let denom = (lhs_sumsq.sqrt() * rhs_sumsq.sqrt()).max(1e-300);
    Ok(num / denom)
}

fn smooth_has_overlapping_linear_terms(
    linear_terms: &[LinearTermSpec],
    termspec: &SmoothTermSpec,
) -> bool {
    let feature_cols = smooth_term_feature_cols(termspec);
    linear_terms
        .iter()
        .any(|linear| feature_cols.contains(&linear.feature_col))
}

fn smooth_intrinsic_parametric_feature_cols(
    linear_terms: &[LinearTermSpec],
    term: &SmoothTermSpec,
) -> Vec<usize> {
    // Returns the data columns that should appear in the smooth's parametric
    // constraint block `C = [1, …]`, alongside which the smooth is then
    // ORTHOGONALIZED via `apply_smooth_transform_to_design`.  Every column
    // returned here is therefore a direction that gets projected OUT of the
    // smooth's basis.  The constant intercept is always included by
    // `build_parametric_constraint_block_for_term`, so this function only
    // controls the polynomial axes added on top of that intercept.
    //
    // Ownership rule: explicit linear terms claim their matching axes — and
    // only those axes — for projection.  A standalone smooth (no overlapping
    // linear term) keeps its full polynomial nullspace and is centered only
    // against the implicit intercept; this matches the canonical thin-plate
    // / Duchon model surface where the linear component is part of the
    // smooth itself.
    let feature_cols = smooth_term_feature_cols(term);
    let mut owned = Vec::new();
    for linear in linear_terms {
        if feature_cols.contains(&linear.feature_col) && !owned.contains(&linear.feature_col) {
            owned.push(linear.feature_col);
        }
    }
    owned
}

fn apply_global_smooth_identifiability(
    smooth: RawSmoothDesign,
    data: ArrayView2<'_, f64>,
    linear_terms: &[LinearTermSpec],
    smoothspecs: &[SmoothTermSpec],
) -> Result<SmoothDesign, BasisError> {
    // Global smooth identifiability policy:
    //
    // 1. Any smooth that overlaps explicit linear terms is residualized against
    //    [intercept | overlapping linear columns]. Spatial smooths also keep
    //    their existing parametric orthogonality policy when requested.
    // 2. Higher-order / duplicate smooths are orthogonalized to the realized
    //    design columns of lower-order owned smooths over nested feature sets.
    //
    // This yields a deterministic hierarchical decomposition: lower-order smooths
    // own their subspaces, and broader smooths fit only the residual structure.
    if smoothspecs.len() != smooth.terms.len() {
        crate::bail_dim_basis!(
            "smooth spec count ({}) does not match built term count ({})",
            smoothspecs.len(),
            smooth.terms.len()
        );
    }

    if smooth.terms.is_empty() {
        return Ok(smooth.into());
    }

    let mut local_designs = vec![None; smooth.terms.len()];
    let mut local_penalties = vec![Vec::<Array2<f64>>::new(); smooth.terms.len()];
    let mut local_nullspaces = vec![Vec::<usize>::new(); smooth.terms.len()];
    let mut local_penaltyinfo = vec![Vec::<PenaltyInfo>::new(); smooth.terms.len()];
    let mut local_metadata = vec![None; smooth.terms.len()];
    let mut local_dims = vec![0usize; smooth.terms.len()];
    let mut local_linear_constraints = vec![None; smooth.terms.len()];

    let SmoothStructureAnalysis {
        ownership_order,
        term_owners,
        ..
    } = analyze_smooth_ownership(smoothspecs);

    use rayon::iter::{
        IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, ParallelIterator,
    };

    for &idx in &ownership_order {
        let term = &smooth.terms[idx];
        let termspec = &smoothspecs[idx];
        let design_local = smooth.term_designs[idx].clone();
        let skip_global_transform =
            smooth_has_frozen_identifiability(termspec) || term.lower_bounds_local.is_some();
        let owner_indices = if skip_global_transform {
            Vec::new()
        } else {
            let overlap_tol = 1e-10;
            let owner_cross_checks = term_owners[idx]
                .clone()
                .into_par_iter()
                .map(|owner_idx| {
                    let owner_design = local_designs[owner_idx]
                        .as_ref()
                        .expect("owner design must be available before dependent smooth");
                    design_cross_relative_residual(&design_local, owner_design)
                        .map(|rel| (owner_idx, rel))
                })
                .collect::<Vec<_>>();
            let mut out = Vec::new();
            for check in owner_cross_checks {
                let (owner_idx, rel) = check?;
                if rel > overlap_tol {
                    out.push(owner_idx);
                }
            }
            out
        };
        let owner_blocks = owner_indices
            .iter()
            .map(|owner_idx| {
                local_designs[*owner_idx]
                    .as_ref()
                    .expect("owner design must be available before dependent smooth")
            })
            .collect::<Vec<_>>();
        let needs_parametric_block = !skip_global_transform
            && (smooth_has_overlapping_linear_terms(linear_terms, termspec)
                || !smooth_intrinsic_parametric_feature_cols(linear_terms, termspec).is_empty()
                || smooth_requires_parametric_orthogonality(termspec));
        let parametric_block = if !needs_parametric_block {
            None
        } else {
            Some(build_parametric_constraint_block_for_term(
                data,
                linear_terms,
                termspec,
            )?)
        };
        let c_local =
            if skip_global_transform || (parametric_block.is_none() && owner_blocks.is_empty()) {
                None
            } else {
                Some(build_constraint_block(
                    data.nrows(),
                    parametric_block.as_ref(),
                    &owner_blocks,
                )?)
            };
        let z_opt = if skip_global_transform {
            None
        } else {
            match maybe_smooth_identifiability_transform(
                termspec,
                &design_local,
                c_local.as_ref().map(|mat| mat.view()),
            ) {
                Ok(z_opt) => z_opt,
                Err(BasisError::ConstraintNullspaceCollapsed { .. })
                    if !owner_blocks.is_empty() =>
                {
                    Some(Array2::zeros((design_local.ncols(), 0)))
                }
                Err(err) => return Err(err),
            }
        };
        let design_constrained = if let Some(z) = z_opt.as_ref() {
            apply_smooth_transform_to_design(design_local, z, &term.name)?
        } else {
            design_local
        };

        if let Some(c_ref) = c_local.as_ref() {
            let rel =
                orthogonality_relative_residual_for_design(&design_constrained, c_ref.view())?;
            let tol = 1e-8;
            if rel > tol {
                crate::bail_invalid_basis!(
                    "smooth orthogonality residual too large for term '{}': {:.3e} > {:.1e}",
                    term.name,
                    rel,
                    tol
                );
            }
        }

        let active_penaltyinfo = term
            .penaltyinfo_local
            .iter()
            .filter(|info| info.active)
            .cloned()
            .collect::<Vec<_>>();
        if active_penaltyinfo.len() != term.penalties_local.len() {
            crate::bail_invalid_basis!(
                "internal penalty metadata mismatch for term '{}': activeinfos={}, penalties={}",
                term.name,
                active_penaltyinfo.len(),
                term.penalties_local.len()
            );
        }
        let penalties_constrained: Vec<Array2<f64>> = term
            .penalties_local
            .par_iter()
            .map(|s_local| {
                if let Some(z) = z_opt.as_ref() {
                    let zt_s = fast_atb(z, s_local);
                    fast_ab(&zt_s, z)
                } else {
                    s_local.clone()
                }
            })
            .collect();
        let penalty_candidates = penalties_constrained
            .into_par_iter()
            .zip(active_penaltyinfo.into_par_iter())
            .map(|(matrix, info)| {
                let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
                PenaltyCandidate {
                    nullspace_dim_hint: info.nullspace_dim_hint,
                    matrix,
                    source: info.source,
                    normalization_scale: info.normalization_scale * c_new,
                    kronecker_factors: None,
                    op: None,
                }
            })
            .collect::<Vec<_>>();
        let (penalties_constrained, nullspace_constrained, penaltyinfo_constrained) =
            filter_active_penalty_candidates(penalty_candidates)?;
        let linear_constraints_constrained =
            if let Some(lin_local) = term.linear_constraints_local.as_ref() {
                if let Some(z) = z_opt.as_ref() {
                    Some(LinearInequalityConstraints {
                        a: lin_local.a.dot(z),
                        b: lin_local.b.clone(),
                    })
                } else {
                    Some(lin_local.clone())
                }
            } else {
                None
            };

        local_dims[idx] = design_constrained.ncols();
        local_designs[idx] = Some(design_constrained);
        local_penalties[idx] = penalties_constrained;
        local_nullspaces[idx] = nullspace_constrained;
        local_penaltyinfo[idx] = penaltyinfo_constrained;
        local_linear_constraints[idx] = linear_constraints_constrained;
        let realized_transform = match (term.joint_null_rotation.as_ref(), z_opt.as_ref()) {
            (Some(rotation), Some(z)) => {
                Some(crate::linalg::faer_ndarray::fast_ab(&rotation.rotation, z))
            }
            (Some(rotation), None) => Some(rotation.rotation.clone()),
            (None, Some(z)) => Some(z.clone()),
            (None, None) => None,
        };
        // Block-replicated factor smooths (`bs="sz"` → `FactorSumToZero`) carry a
        // PER-MARGINAL metadata (predict rebuilds the single inner marginal then
        // re-stacks the `L-1` sum-to-zero deviation blocks). Their realized
        // transform — the joint-null absorption rotation `Q` (and any global
        // orthogonality `Z`) — lives in the FULL `p·(L-1)`-column design space, so
        // it cannot be folded into the per-marginal metadata (the dimensions don't
        // compose: `existing pxr` vs `extra (p·(L-1))x(p·(L-1))`). The raw design
        // builder already applies `Q` to the re-stacked design and recomputes it
        // deterministically from the (frozen) penalties at predict time, so the
        // per-marginal metadata must be left untouched here; folding it in both
        // crashed basis generation and would double-count `Q` on rebuild (#700).
        let metadata_is_per_marginal_block =
            matches!(termspec.basis, SmoothBasisSpec::FactorSumToZero { .. });
        local_metadata[idx] = if metadata_is_per_marginal_block {
            Some(term.metadata.clone())
        } else {
            Some(with_identifiability_transform(
                &term.metadata,
                realized_transform.as_ref(),
            )?)
        };
    }

    let total_p: usize = local_dims.iter().sum();
    let mut terms_out = Vec::<SmoothTerm>::with_capacity(smooth.terms.len());
    let mut penalties_global = Vec::<BlockwisePenalty>::new();
    let mut nullspace_dims_global = Vec::<usize>::new();
    let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
    let mut dropped_penaltyinfo_global = smooth.dropped_penaltyinfo.clone();
    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
    let mut any_bounds = false;
    let mut linear_constraintsrows: Vec<Array1<f64>> = Vec::new();
    let mut linear_constraints_b: Vec<f64> = Vec::new();

    let mut col_start = 0usize;
    for idx in 0..smooth.terms.len() {
        let p_local = local_dims[idx];
        let col_end = col_start + p_local;

        let activeinfos = local_penaltyinfo[idx]
            .iter()
            .filter(|info| info.active)
            .collect::<Vec<_>>();
        if activeinfos.len() != local_penalties[idx].len() {
            crate::bail_invalid_basis!(
                "internal penalty info mismatch for term '{}': activeinfos={}, penalties={}",
                smooth.terms[idx].name,
                activeinfos.len(),
                local_penalties[idx].len()
            );
        }
        for ((s_local, &ns), info) in local_penalties[idx]
            .iter()
            .zip(local_nullspaces[idx].iter())
            .zip(activeinfos.into_iter())
        {
            let global_index = penalties_global.len();
            penalties_global.push(BlockwisePenalty::new(col_start..col_end, s_local.clone()));
            nullspace_dims_global.push(ns);
            let mut penalty = info.clone();
            penalty.nullspace_dim_hint = ns;
            penaltyinfo_global.push(PenaltyBlockInfo {
                global_index,
                termname: Some(smooth.terms[idx].name.clone()),
                penalty,
            });
        }
        for info in local_penaltyinfo[idx].iter().filter(|info| !info.active) {
            dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
                termname: Some(smooth.terms[idx].name.clone()),
                penalty: info.clone(),
            });
        }

        terms_out.push(SmoothTerm {
            name: smooth.terms[idx].name.clone(),
            coeff_range: col_start..col_end,
            shape: smooth.terms[idx].shape,
            penalties_local: local_penalties[idx].clone(),
            nullspace_dims: local_nullspaces[idx].clone(),
            penaltyinfo_local: local_penaltyinfo[idx].clone(),
            metadata: local_metadata[idx]
                .clone()
                .expect("local metadata must exist for every smooth term"),
            lower_bounds_local: smooth.terms[idx].lower_bounds_local.clone(),
            linear_constraints_local: local_linear_constraints[idx].clone(),
            // Global orthogonality transforms break Kronecker structure.
            kronecker_factored: None,
            // The final raw-basis → coefficient chart, including any
            // stage-2 joint-null Q and global orthogonality Z, is embedded in
            // `metadata` above. Keeping Q separately here would apply it twice
            // on frozen rebuilds and would put derivative operators in a
            // different chart from the value path.
            joint_null_rotation: None,
        });
        if let Some(lin_local) = &local_linear_constraints[idx] {
            for r in 0..lin_local.a.nrows() {
                let mut row = Array1::<f64>::zeros(total_p);
                row.slice_mut(s![col_start..col_end])
                    .assign(&lin_local.a.row(r));
                linear_constraintsrows.push(row);
                linear_constraints_b.push(lin_local.b[r]);
            }
        }
        if let Some(lb_local) = smooth.terms[idx].lower_bounds_local.as_ref()
            && lb_local.len() == p_local
        {
            coefficient_lower_bounds
                .slice_mut(s![col_start..col_end])
                .assign(lb_local);
            any_bounds = true;
        }

        col_start = col_end;
    }

    assert_eq!(
        penalties_global.len(),
        nullspace_dims_global.len(),
        "globally reparameterized smooth penalty/nullspace bookkeeping diverged"
    );
    assert_eq!(
        penalties_global.len(),
        penaltyinfo_global.len(),
        "globally reparameterized smooth penalty metadata bookkeeping diverged"
    );

    Ok(SmoothDesign {
        term_designs: local_designs
            .into_iter()
            .map(|design| design.expect("local design must exist for every smooth term"))
            .collect(),
        penalties: penalties_global,
        nullspace_dims: nullspace_dims_global,
        penaltyinfo: penaltyinfo_global,
        dropped_penaltyinfo: dropped_penaltyinfo_global,
        terms: terms_out,
        coefficient_lower_bounds: if any_bounds {
            Some(coefficient_lower_bounds)
        } else {
            None
        },
        linear_constraints: if linear_constraintsrows.is_empty() {
            None
        } else {
            let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
            for (i, row) in linear_constraintsrows.iter().enumerate() {
                a.row_mut(i).assign(row);
            }
            Some(LinearInequalityConstraints {
                a,
                b: Array1::from_vec(linear_constraints_b),
            })
        },
    })
}

fn build_parametric_constraint_block_for_term(
    data: ArrayView2<'_, f64>,
    linear_terms: &[LinearTermSpec],
    termspec: &SmoothTermSpec,
) -> Result<Array2<f64>, BasisError> {
    let n = data.nrows();
    let p_data = data.ncols();
    let feature_cols = smooth_term_feature_cols(termspec);
    let mut parametric_cols = smooth_intrinsic_parametric_feature_cols(linear_terms, termspec);
    for &feature_col in &parametric_cols {
        if feature_col >= p_data {
            crate::bail_dim_basis!(
                "smooth term feature column {feature_col} out of bounds for {p_data} columns"
            );
        }
    }
    for linear in linear_terms
        .iter()
        .filter(|linear| feature_cols.contains(&linear.feature_col))
    {
        if linear.feature_col >= p_data {
            crate::bail_dim_basis!(
                "linear term '{}' feature column {} out of bounds for {} columns",
                linear.name,
                linear.feature_col,
                p_data
            );
        }
        if !parametric_cols.contains(&linear.feature_col) {
            parametric_cols.push(linear.feature_col);
        }
    }

    let mut c = Array2::<f64>::zeros((n, 1 + parametric_cols.len()));
    c.column_mut(0).fill(1.0);
    for (j, &feature_col) in parametric_cols.iter().enumerate() {
        c.column_mut(j + 1).assign(&data.column(feature_col));
    }
    Ok(c)
}

fn apply_smooth_transform_to_design(
    design_local: DesignMatrix,
    transform: &Array2<f64>,
    termname: &str,
) -> Result<DesignMatrix, BasisError> {
    match design_local {
        DesignMatrix::Dense(inner) => {
            let op = CoefficientTransformOperator::new(inner, transform.clone()).map_err(|e| {
                BasisError::InvalidInput(format!(
                    "smooth identifiability transform failed for term '{termname}': {e}"
                ))
            })?;
            Ok(DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(
                Arc::new(op),
            )))
        }
        DesignMatrix::Sparse(inner) => {
            let dense = inner
                .try_to_dense_arc("smooth identifiability sparse transform")
                .map_err(BasisError::InvalidInput)?
                .as_ref()
                .dot(transform);
            Ok(DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(
                dense,
            )))
        }
    }
}

fn design_constraint_cross(
    design: &DesignMatrix,
    constraint_matrix: ArrayView2<'_, f64>,
) -> Result<Array2<f64>, BasisError> {
    let n = design.nrows();
    if constraint_matrix.nrows() != n {
        return Err(BasisError::ConstraintMatrixRowMismatch {
            basisrows: n,
            constraintrows: constraint_matrix.nrows(),
        });
    }
    let mut cross = Array2::<f64>::zeros((design.ncols(), constraint_matrix.ncols()));
    const CHUNK: usize = 1024;
    for start in (0..n).step_by(CHUNK) {
        let end = (start + CHUNK).min(n);
        let design_chunk = design
            .try_row_chunk(start..end)
            .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
        let constraint_chunk = constraint_matrix.slice(s![start..end, ..]).to_owned();
        cross += &crate::faer_ndarray::fast_atb(&design_chunk, &constraint_chunk);
    }
    Ok(cross)
}

fn design_frobenius_norm(design: &DesignMatrix) -> Result<f64, BasisError> {
    let n = design.nrows();
    const CHUNK: usize = 1024;
    let mut sumsq = 0.0;
    for start in (0..n).step_by(CHUNK) {
        let end = (start + CHUNK).min(n);
        let chunk = design
            .try_row_chunk(start..end)
            .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
        sumsq += chunk.iter().map(|v| v * v).sum::<f64>();
    }
    Ok(sumsq.sqrt())
}

fn maybe_smooth_identifiability_transform(
    termspec: &SmoothTermSpec,
    design_local: &DesignMatrix,
    constraint_block: Option<ArrayView2<'_, f64>>,
) -> Result<Option<Array2<f64>>, BasisError> {
    if let Some(SpatialIdentifiability::FrozenTransform { transform }) =
        spatial_identifiability_policy(termspec)
    {
        if design_local.ncols() != transform.nrows() {
            crate::bail_dim_basis!(
                "frozen spatial identifiability transform mismatch: design has {} columns but transform has {} rows",
                design_local.ncols(),
                transform.nrows()
            );
        }
        return Ok(Some(transform.clone()));
    }

    if let Some(c) = constraint_block {
        if c.ncols() == 0 {
            Ok(None)
        } else {
            Ok(Some(orthogonality_transform_for_design(
                design_local,
                c,
                None, // fixed subspace: do not use iteration-varying PIRLS weights
            )?))
        }
    } else {
        Ok(None)
    }
}

fn spatial_identifiability_policy(termspec: &SmoothTermSpec) -> Option<&SpatialIdentifiability> {
    match &termspec.basis {
        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.identifiability),
        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.identifiability),
        _ => None,
    }
}

/// Whether this smooth's *realized* design (the basis evaluated at the n data
/// rows) must be residualized against the model's parametric block (intercept +
/// any overlapping linear columns) by `apply_global_smooth_identifiability`.
///
/// This is the universal identifiability invariant for **kernel / radial**
/// spatial smooths (#531): their realized column span contains the constant
/// (and, at `Linear` null-space order, the linear monomials), so without this
/// step the smooth and the parametric intercept fight over the same direction —
/// a structural rank-1 collision. The collision is invisible to the kernels'
/// *own* identifiability constraints because those act in **coefficient space at
/// the K centers**, not on the realized design rows:
///   - Matérn `CenterSumToZero` enforces `1ᵀα = 0` over the centers, so
///     `Kα` evaluated at the data rows still spans the constant.
///   - Duchon / TPS `OrthogonalToParametric` *defers* its centering to this very
///     step, which is why it is listed here too.
///
/// Tensor-product and B-spline bases instead apply a realized-design sum-to-zero
/// at basis-build time (`apply_sum_to_zero_constraint`), so they already satisfy
/// the invariant and must NOT be double-constrained — they return `false`.
///
/// The remaining bases are excluded, each for a concrete reason:
///   - **Sphere, Harmonic method**: the real-spherical-harmonic basis starts at
///     degree `l = 1` (`build_spherical_harmonic_basis`), so it never spans the
///     degree-0 constant — no centering is needed.
///   - **Sphere, Wahba method**: INCLUDED (#532). Its
///     `weighted_coefficient_sum_to_zero_transform` is a *center*-space
///     constraint, so the realized design still spans the constant — same class
///     as Matérn `CenterSumToZero`. The composed parametric transform is frozen
///     onto `SphericalSplineBasisSpec::identifiability`
///     (`SphericalSplineIdentifiability::FrozenTransform`) and replayed by
///     `build_spherical_spline_basis` at predict time, so the orthogonalization
///     survives save → reload exactly as it does for Matérn.
///   - **PCA**: its `with_identifiability_transform` arm rejects a post-hoc
///     transform (the constraint lives inside the orthonormal loadings), and its
///     constant content is governed by the `centered` flag, not a residualizable
///     design.
///
/// `FrozenTransform` bases are excluded: a transform frozen by *this* pipeline
/// already has the parametric orthogonalization composed in (see
/// `with_identifiability_transform`), and they are gated out upstream by
/// `skip_global_transform` regardless.
fn smooth_requires_parametric_orthogonality(termspec: &SmoothTermSpec) -> bool {
    match &termspec.basis {
        SmoothBasisSpec::ByVariable { inner, .. }
        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
            smooth_requires_parametric_orthogonality(&SmoothTermSpec {
                name: termspec.name.clone(),
                basis: (**inner).clone(),
                shape: termspec.shape,
                joint_null_rotation: None,
            })
        }
        SmoothBasisSpec::BySmooth { smooth, .. } => {
            smooth_requires_parametric_orthogonality(&SmoothTermSpec {
                name: termspec.name.clone(),
                basis: (**smooth).clone(),
                shape: termspec.shape,
                joint_null_rotation: None,
            })
        }
        SmoothBasisSpec::ThinPlate { spec, .. } => {
            matches!(
                spec.identifiability,
                SpatialIdentifiability::OrthogonalToParametric
            )
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            matches!(
                spec.identifiability,
                SpatialIdentifiability::OrthogonalToParametric
            )
        }
        SmoothBasisSpec::Matern { spec, .. } => matches!(
            spec.identifiability,
            MaternIdentifiability::CenterSumToZero | MaternIdentifiability::CenterLinearOrthogonal
        ),
        // Wahba sphere (`bs="sos"`, method=Wahba): the area-weighted center
        // sum-to-zero `z` is a *coefficient*-space constraint, so the realized
        // `K·z` design still spans the constant — the same #531 collision class
        // as Matérn `CenterSumToZero`. It requires the global parametric
        // orthogonalization (#532). The Harmonic method starts at degree l=1
        // and never spans the constant, so it is excluded.
        SmoothBasisSpec::Sphere { spec, .. } => {
            matches!(spec.method, crate::basis::SphereMethod::Wahba)
                && matches!(
                    spec.identifiability,
                    SphericalSplineIdentifiability::CenterSumToZero
                )
        }
        SmoothBasisSpec::BSpline1D { .. }
        | SmoothBasisSpec::TensorBSpline { .. }
        | SmoothBasisSpec::Pca { .. }
        | SmoothBasisSpec::FactorSmooth { .. } => false,
    }
}

fn compose_identifiability_transforms(
    existing: Option<&Array2<f64>>,
    extra: Option<&Array2<f64>>,
) -> Result<Option<Array2<f64>>, BasisError> {
    match (existing, extra) {
        (Some(lhs), Some(rhs)) => {
            if lhs.ncols() == rhs.nrows() {
                Ok(Some(lhs.dot(rhs)))
            } else if lhs.nrows() == rhs.nrows() && lhs.ncols() == rhs.ncols() {
                // Rebuilding from an already-frozen spec can surface the same
                // raw->frozen transform twice. Treat that as idempotent
                // metadata, not a sequential Z_left * Z_right composition.
                Ok(Some(rhs.clone()))
            } else {
                Err(BasisError::DimensionMismatch(format!(
                    "identifiability transform mismatch: existing is {}x{}, extra is {}x{}",
                    lhs.nrows(),
                    lhs.ncols(),
                    rhs.nrows(),
                    rhs.ncols(),
                )))
            }
        }
        (Some(lhs), None) => Ok(Some(lhs.clone())),
        (None, Some(rhs)) => Ok(Some(rhs.clone())),
        (None, None) => Ok(None),
    }
}

fn with_identifiability_transform(
    metadata: &BasisMetadata,
    transform: Option<&Array2<f64>>,
) -> Result<BasisMetadata, BasisError> {
    match metadata {
        BasisMetadata::BSpline1D {
            knots,
            identifiability_transform,
            periodic,
            degree,
            auto_shrink_note,
        } => Ok(BasisMetadata::BSpline1D {
            knots: knots.clone(),
            periodic: *periodic,
            identifiability_transform: compose_identifiability_transforms(
                identifiability_transform.as_ref(),
                transform,
            )?,
            degree: *degree,
            auto_shrink_note: auto_shrink_note.clone(),
        }),
        BasisMetadata::ThinPlate {
            centers,
            length_scale,
            periodic,
            identifiability_transform,
            input_scales,
            radial_reparam,
        } => Ok(BasisMetadata::ThinPlate {
            centers: centers.clone(),
            length_scale: *length_scale,
            periodic: periodic.clone(),
            identifiability_transform: compose_identifiability_transforms(
                identifiability_transform.as_ref(),
                transform,
            )?,
            input_scales: input_scales.clone(),
            radial_reparam: radial_reparam.clone(),
        }),
        BasisMetadata::Sphere {
            centers,
            penalty_order,
            method,
            max_degree,
            wahba_kernel,
            constraint_transform,
        } => Ok(BasisMetadata::Sphere {
            centers: centers.clone(),
            penalty_order: *penalty_order,
            method: *method,
            max_degree: *max_degree,
            wahba_kernel: *wahba_kernel,
            constraint_transform: compose_identifiability_transforms(
                constraint_transform.as_ref(),
                transform,
            )?,
        }),
        BasisMetadata::Matern {
            centers,
            length_scale,
            periodic,
            nu,
            include_intercept,
            identifiability_transform,
            input_scales,
            aniso_log_scales,
        } => Ok(BasisMetadata::Matern {
            centers: centers.clone(),
            length_scale: *length_scale,
            periodic: periodic.clone(),
            nu: *nu,
            include_intercept: *include_intercept,
            identifiability_transform: compose_identifiability_transforms(
                identifiability_transform.as_ref(),
                transform,
            )?,
            input_scales: input_scales.clone(),
            aniso_log_scales: aniso_log_scales.clone(),
        }),
        BasisMetadata::Duchon {
            centers,
            length_scale,
            periodic,
            power,
            nullspace_order,
            identifiability_transform,
            input_scales,
            aniso_log_scales,
            operator_collocation_points,
        } => Ok(BasisMetadata::Duchon {
            centers: centers.clone(),
            length_scale: *length_scale,
            periodic: periodic.clone(),
            power: *power,
            nullspace_order: *nullspace_order,
            input_scales: input_scales.clone(),
            aniso_log_scales: aniso_log_scales.clone(),
            operator_collocation_points: operator_collocation_points.clone(),
            identifiability_transform: compose_identifiability_transforms(
                identifiability_transform.as_ref(),
                transform,
            )?,
        }),
        BasisMetadata::SphereHarmonics {
            max_degree,
            radians,
        } => Ok(BasisMetadata::SphereHarmonics {
            max_degree: *max_degree,
            radians: *radians,
        }),
        BasisMetadata::TensorBSpline {
            feature_cols,
            knots,
            degrees,
            periods,
            identifiability_transform,
        } => Ok(BasisMetadata::TensorBSpline {
            feature_cols: feature_cols.clone(),
            knots: knots.clone(),
            degrees: degrees.clone(),
            periods: periods.clone(),
            identifiability_transform: compose_identifiability_transforms(
                identifiability_transform.as_ref(),
                transform,
            )?,
        }),
        BasisMetadata::BySmooth {
            inner,
            by_col,
            levels,
            ordered,
        } => Ok(BasisMetadata::BySmooth {
            inner: Box::new(with_identifiability_transform(inner, transform)?),
            by_col: *by_col,
            levels: levels.clone(),
            ordered: *ordered,
        }),
        BasisMetadata::FactorSmooth {
            continuous_cols,
            group_col,
            knots,
            degree,
            periodic,
            group_levels,
            flavour,
        } => Ok(BasisMetadata::FactorSmooth {
            continuous_cols: continuous_cols.clone(),
            group_col: *group_col,
            knots: knots.clone(),
            degree: *degree,
            periodic: *periodic,
            group_levels: group_levels.clone(),
            flavour: flavour.clone(),
        }),
        BasisMetadata::Pca {
            feature_cols,
            basis_matrix,
            centered,
            smooth_penalty,
            center_mean,
            pca_basis_path,
            chunk_size,
        } => {
            // PCA bases carry an orthonormal projection matrix and do not
            // expose an identifiability transform that can be re-composed
            // (the constraint, if any, lives inside the PCA loadings
            // themselves), so the caller cannot meaningfully attach a
            // post-hoc Z transform here.
            if transform.is_some() {
                crate::bail_invalid_basis!(
                    "PCA bases do not expose a composable identifiability transform"
                );
            }
            Ok(BasisMetadata::Pca {
                feature_cols: feature_cols.clone(),
                basis_matrix: basis_matrix.clone(),
                centered: *centered,
                smooth_penalty: *smooth_penalty,
                center_mean: center_mean.clone(),
                pca_basis_path: pca_basis_path.clone(),
                chunk_size: *chunk_size,
            })
        }
    }
}
fn orthogonality_relative_residual_for_design(
    design: &DesignMatrix,
    constraint_matrix: ArrayView2<'_, f64>,
) -> Result<f64, BasisError> {
    let cross = design_constraint_cross(design, constraint_matrix)?;
    let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
    let b_norm = design_frobenius_norm(design)?;
    let c_norm = constraint_matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
    let denom = (b_norm * c_norm).max(1e-300);
    Ok(num / denom)
}

pub fn fit_term_collection_forspec(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    fit_term_collection_forspecwith_heuristic_lambdas(
        data, y, weights, offset, spec, None, family, options,
    )
}

pub fn fit_term_collection_with_coefficient_groups(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    groups: &[CoefficientGroupSpec],
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    if groups.is_empty() {
        return fit_term_collection_forspec(data, y, weights, offset, spec, family, options);
    }
    let design = build_term_collection_design(data, spec)?;
    let base_fit_opts = adaptive_fit_options_base(options, &design);
    let realized = design
        .realize_coefficient_groups(groups, &base_fit_opts.rho_prior)
        .map_err(EstimationError::BasisError)?;
    let mut grouped_options = base_fit_opts.clone();
    grouped_options.rho_prior = realized.rho_prior;
    let fitted = FittedTermCollection {
        fit: crate::estimate::fit_gam_with_penalty_specs(
            design.design.clone(),
            y,
            weights,
            offset,
            realized.penalty_specs,
            realized.nullspace_dims,
            family.clone(),
            &grouped_options,
        )?,
        design,
        adaptive_diagnostics: None,
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
    Ok(fitted)
}

pub fn fit_term_collection_with_penalty_block_gamma_prior_callback<F>(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    callback: F,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError>
where
    F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
{
    let design = build_term_collection_design(data, spec)?;
    let mut fit_opts = adaptive_fit_options_base(options, &design);
    fit_opts.rho_prior = realize_penalty_block_gamma_priors(&design, callback)
        .map_err(EstimationError::BasisError)?;
    let fitted = FittedTermCollection {
        fit: fit_gamwith_heuristic_lambdas(
            design.design.clone(),
            y,
            weights,
            offset,
            &design.penalties,
            None,
            family.clone(),
            &fit_opts,
        )?,
        design,
        adaptive_diagnostics: None,
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
    Ok(fitted)
}

pub fn fit_term_collection_with_penalty_block_gamma_priors(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    priors: &[(String, f64, f64)],
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    let design = build_term_collection_design(data, spec)?;
    let mut fit_opts = adaptive_fit_options_base(options, &design);
    fit_opts.rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
        .map_err(EstimationError::BasisError)?;
    let fitted = FittedTermCollection {
        fit: fit_gamwith_heuristic_lambdas(
            design.design.clone(),
            y,
            weights,
            offset,
            &design.penalties,
            None,
            family.clone(),
            &fit_opts,
        )?,
        design,
        adaptive_diagnostics: None,
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
    Ok(fitted)
}

pub fn fit_term_collection_with_coefficient_groups_and_penalty_block_gamma_priors(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    groups: &[CoefficientGroupSpec],
    priors: &[(String, f64, f64)],
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    if groups.is_empty() {
        return fit_term_collection_with_penalty_block_gamma_priors(
            data, y, weights, offset, spec, priors, family, options,
        );
    }
    if priors.is_empty() {
        return fit_term_collection_with_coefficient_groups(
            data, y, weights, offset, spec, groups, family, options,
        );
    }

    let design = build_term_collection_design(data, spec)?;
    let base_fit_opts = adaptive_fit_options_base(options, &design);
    let base_rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
        .map_err(EstimationError::BasisError)?;
    let realized = design
        .realize_coefficient_groups(groups, &base_rho_prior)
        .map_err(EstimationError::BasisError)?;
    let mut grouped_options = base_fit_opts.clone();
    grouped_options.rho_prior = realized.rho_prior;
    let fitted = FittedTermCollection {
        fit: crate::estimate::fit_gam_with_penalty_specs(
            design.design.clone(),
            y,
            weights,
            offset,
            realized.penalty_specs,
            realized.nullspace_dims,
            family.clone(),
            &grouped_options,
        )?,
        design,
        adaptive_diagnostics: None,
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
    Ok(fitted)
}

fn fit_term_collection_forspecwith_heuristic_lambdas(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    heuristic_lambdas: Option<&[f64]>,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    let base_design = build_term_collection_design(data, spec)?;
    fit_term_collection_on_realized_design(
        y,
        weights,
        offset,
        spec,
        &base_design,
        heuristic_lambdas,
        family,
        options,
    )
}

fn has_bounded_linear_terms(spec: &TermCollectionSpec) -> bool {
    spec.linear_terms.iter().any(|term| {
        matches!(
            term.coefficient_geometry,
            LinearCoefficientGeometry::Bounded { .. }
        )
    })
}

fn fit_term_collection_on_realized_design(
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    heuristic_lambdas: Option<&[f64]>,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    if has_bounded_linear_terms(spec) {
        return fit_bounded_term_collection_with_design(
            y,
            weights,
            offset,
            spec,
            design,
            heuristic_lambdas,
            family,
            options,
        );
    }
    let base_fit_opts = adaptive_fit_options_base(options, design);
    let fitted = FittedTermCollection {
        fit: fit_gamwith_heuristic_lambdas(
            design.design.clone(),
            y,
            weights,
            offset,
            &design.penalties,
            heuristic_lambdas,
            family.clone(),
            &base_fit_opts,
        )?,
        design: design.clone(),
        adaptive_diagnostics: None,
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;

    let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
    if !adaptive_opts.enabled {
        return Ok(fitted);
    }
    let runtime_caches = extract_spatial_operator_runtime_caches(spec, &fitted.design)?;
    if runtime_caches.is_empty() {
        return Ok(fitted);
    }
    // Spatial-adaptive overlay always runs when the operator caches are
    // non-empty. Catastrophic-overfit protection lives in the operator-log-λ
    // box bound (Fix B at the BFGS bounds construction), which caps maximum
    // unpenalization regardless of n. Production fits at n≈300K must run the
    // overlay; the previous n-gate (n < max(4·p_total, 200)) silently skipped
    // it for any small-n test, contradicting that contract.
    fit_term_collectionwith_exact_spatial_adaptive_regularization(
        fitted,
        y,
        weights,
        offset,
        family,
        options,
        &runtime_caches,
    )
}

#[derive(Clone)]
struct SpatialOperatorRuntimeCache {
    termname: String,
    feature_cols: Vec<usize>,
    coeff_global_range: Range<usize>,
    mass_penalty_global_idx: usize,
    tension_penalty_global_idx: usize,
    stiffness_penalty_global_idx: usize,
    d0: Array2<f64>,
    d1: Array2<f64>,
    d2: Array2<f64>,
    collocation_points: Array2<f64>,
    dimension: usize,
}

#[derive(Clone)]
struct SpatialAdaptiveWeights {
    inv_magweight: Array1<f64>,
    invgradweight: Array1<f64>,
    inv_lapweight: Array1<f64>,
}

#[derive(Clone)]
struct CharbonnierScalarBlockState {
    signal: Array1<f64>,
    radius: Array1<f64>,
    epsilon: f64,
}

impl CharbonnierScalarBlockState {
    fn from_signal(signal: Array1<f64>, epsilon: f64) -> Self {
        let eps = epsilon.max(1e-12);
        let radius = signal.mapv(|t| (t * t + eps * eps).sqrt());
        Self {
            signal,
            radius,
            epsilon: eps,
        }
    }

    fn absolute_signal(&self) -> Array1<f64> {
        self.signal.mapv(f64::abs)
    }

    fn penalty_value(&self) -> f64 {
        self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
    }

    fn betagradient_coeff(&self) -> Array1<f64> {
        Array1::from_iter(
            self.signal
                .iter()
                .zip(self.radius.iter())
                .map(|(t, r)| t / r),
        )
    }

    fn betahessian_diag(&self) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        self.radius.mapv(|r| eps2 / r.powi(3))
    }

    fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
        let epsilon = self.epsilon;
        let eps2 = epsilon * epsilon;
        self.radius.mapv(|r| eps2 / r - epsilon)
    }

    fn log_epsilon_betagradient_coeff(&self) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        Array1::from_iter(
            self.signal
                .iter()
                .zip(self.radius.iter())
                .map(|(t, r)| -eps2 * t / r.powi(3)),
        )
    }

    fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
        let epsilon = self.epsilon;
        let eps2 = epsilon * epsilon;
        let eps4 = eps2 * eps2;
        self.radius
            .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
    }

    fn surrogateweights_posterior_snr(
        &self,
        variance: &Array1<f64>,
        weight_floor: f64,
        weight_ceiling: f64,
    ) -> (Array1<f64>, Array1<f64>) {
        // Posterior-SNR (credible-magnitude) reweighting of the scalar MM
        // majorizer.
        //
        // The magnitude-only surrogate weight uses the *point-estimate* radius
        //
        //   r_k^mag = sqrt( t_k^2 + eps^2 ),   t_k = (D0 beta_hat)_k,
        //   w_k     = 1 / r_k^mag.
        //
        // The weight multiplies the local quadratic surrogate penalty
        // w_k (D0 beta)^2, so a *small* w_k leaves the response un-penalized
        // (treated as a genuine feature) and a *large* w_k pulls it toward zero
        // (enforces flatness). The failure of the point-estimate radius is that
        // a response t_k which is large only because it is poorly determined
        // gets a tiny weight and is left un-penalized — the weight chases noise
        // in low-information regions.
        //
        // Resolution via the posterior second moment under the working-Laplace
        // posterior beta ~ N(beta_hat, Sigma_beta), Sigma_beta = H^{-1}: the
        // variance of the response is
        //
        //   Var( (D0 beta)_k ) = (D0 Sigma_beta D0^T)_kk >= 0,
        //
        // and the *credible* (noise-floor-corrected) squared magnitude is
        //
        //   t_k^credible^2 = max( t_k^2 - Var(...)_k , 0 ),
        //   r_k^snr        = sqrt( t_k^credible^2 + eps^2 ),
        //   w_k            = 1 / r_k^snr.
        //
        // This realizes the intended behavior exactly: a derivative is left
        // un-penalized (small w) only when its magnitude is *credibly* large
        // (t_k^2 >> Var, real feature); the penalty pulls hardest toward zero
        // (large w) where the derivative is credibly near zero — small estimate
        // AND small variance, OR a large-but-poorly-determined estimate whose
        // signal is swamped by its variance. (Note: the additive `+Var` radius
        // sketched in the design notes points the *opposite* way — it would
        // un-penalize noisy responses even harder — so we use the variance-
        // *corrected* second moment, which is the direction that actually
        // suppresses noise while preserving credible edges.) With
        // `variance == 0` everywhere this degrades exactly to `surrogateweights`,
        // so any covariance-unavailable path is unchanged.
        let eps2 = self.epsilon * self.epsilon;
        let weight = Array1::from_iter(self.signal.iter().zip(variance.iter()).map(|(&t, &v)| {
            let credible_sq = (t * t - v.max(0.0)).max(0.0);
            let r = (credible_sq + eps2).sqrt();
            (1.0 / r).clamp(weight_floor, weight_ceiling)
        }));
        let invweight = weight.mapv(|u| 1.0 / u);
        (weight, invweight)
    }

    fn directionalhessian_diag(&self, direction_signal: &Array1<f64>) -> Array1<f64> {
        // Scalar-image directional third derivative:
        //
        // If t(beta) = A beta and
        //   H(beta) = A^T diag( eps^2 / (t_k(beta)^2 + eps^2)^(3/2) ) A,
        // then for q = A u,
        //
        //   D(H)[u]
        //   = A^T diag( -3 eps^2 t_k q_k / (t_k^2 + eps^2)^(5/2) ) A.
        //
        // This is one of the exact P_{beta,beta,beta}[u] terms needed by the
        // Laplace hypergradient
        //
        //   d/dtheta log det H = tr(H^{-1} Hdot_theta),
        //   Hdot_theta = J_{beta,beta,theta} + D_beta(H)[beta_theta].
        let eps2 = self.epsilon * self.epsilon;
        Array1::from_iter(
            self.signal
                .iter()
                .zip(direction_signal.iter())
                .zip(self.radius.iter())
                .map(|((t, q), r)| -3.0 * eps2 * t * q / r.powi(5)),
        )
    }

    fn log_epsilon_betahessian_diag(&self) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        let eps4 = eps2 * eps2;
        Array1::from_iter(
            self.signal
                .iter()
                .zip(self.radius.iter())
                .map(|(_, r)| 2.0 * eps2 / r.powi(3) - 3.0 * eps4 / r.powi(5)),
        )
    }

    fn log_epsilon_beta_mixed_second_coeff(&self) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        Array1::from_iter(
            self.signal
                .iter()
                .zip(self.radius.iter())
                .map(|(t, r)| eps2 * t * (eps2 - 2.0 * t * t) / r.powi(5)),
        )
    }

    fn log_epsilon_betahessian_second_diag(&self) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        let eps4 = eps2 * eps2;
        let eps6 = eps4 * eps2;
        Array1::from_iter(
            self.radius.iter().map(|r| {
                4.0 * eps2 / r.powi(3) - 18.0 * eps4 / r.powi(5) + 15.0 * eps6 / r.powi(7)
            }),
        )
    }

    fn log_epsilon_betahessian_directional_diag(
        &self,
        direction_signal: &Array1<f64>,
    ) -> Array1<f64> {
        let eps2 = self.epsilon * self.epsilon;
        let eps4 = eps2 * eps2;
        Array1::from_iter(
            self.signal
                .iter()
                .zip(direction_signal.iter())
                .zip(self.radius.iter())
                .map(|((t, q), r)| (-6.0 * eps2 * t / r.powi(5) + 15.0 * eps4 * t / r.powi(7)) * q),
        )
    }
}

#[derive(Clone)]
struct CharbonnierGroupedBlockState {
    norm: Array1<f64>,
    radius: Array1<f64>,
    signal_blocks: Array2<f64>,
    epsilon: f64,
}

impl CharbonnierGroupedBlockState {
    fn from_signal_blocks(signal_blocks: Array2<f64>, epsilon: f64) -> Self {
        let eps = epsilon.max(1e-12);
        let norm = Array1::from_iter(
            signal_blocks
                .rows()
                .into_iter()
                .map(|row| row.iter().map(|v| v * v).sum::<f64>().sqrt()),
        );
        let radius = norm.mapv(|g| (g * g + eps * eps).sqrt());
        Self {
            norm,
            radius,
            signal_blocks,
            epsilon: eps,
        }
    }

    fn penalty_value(&self) -> f64 {
        self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
    }

    fn norm_signal(&self) -> Array1<f64> {
        self.norm.clone()
    }

    fn betagradient_blocks(&self) -> Array2<f64> {
        let mut out = self.signal_blocks.clone();
        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
            let scale = 1.0 / self.radius[k];
            row.mapv_inplace(|v| v * scale);
        }
        out
    }

    fn betahessian_blocks(&self) -> Vec<Array2<f64>> {
        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
            let dim = row.len();
            let mut block = Array2::<f64>::eye(dim);
            block.mapv_inplace(|v| v / self.radius[k]);
            for i in 0..dim {
                for j in 0..dim {
                    block[[i, j]] -= row[i] * row[j] / self.radius[k].powi(3);
                }
            }
            out.push(block);
        }
        out
    }

    fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
        let epsilon = self.epsilon;
        let eps2 = epsilon * epsilon;
        self.radius.mapv(|r| eps2 / r - epsilon)
    }

    fn log_epsilon_betagradient_blocks(&self) -> Array2<f64> {
        let mut out = self.signal_blocks.clone();
        let eps2 = self.epsilon * self.epsilon;
        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
            let scale = -eps2 / self.radius[k].powi(3);
            row.mapv_inplace(|v| v * scale);
        }
        out
    }

    fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
        let epsilon = self.epsilon;
        let eps2 = epsilon * epsilon;
        let eps4 = eps2 * eps2;
        self.radius
            .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
    }

    fn surrogateweights_posterior_snr(
        &self,
        variance: &Array1<f64>,
        weight_floor: f64,
        weight_ceiling: f64,
    ) -> (Array1<f64>, Array1<f64>) {
        // Grouped posterior-SNR (credible-magnitude) reweighting.
        //
        // The magnitude-only grouped surrogate weight uses the point-estimate
        // block norm
        //
        //   g_k     = ||v_k||_2,   v_k = G_k beta_hat,
        //   r_k^mag = sqrt( g_k^2 + eps^2 ),
        //   w_k     = 1 / r_k^mag.
        //
        // The posterior covariance of the *block* response v_k = G_k beta under
        // beta ~ N(beta_hat, Sigma_beta), Sigma_beta = H^{-1}, has total trace
        //
        //   Cov(v_k)     = G_k Sigma_beta G_k^T   (a block_dim x block_dim block),
        //   variance[k]  = tr(Cov(v_k)) = sum_axis ( G_k[axis] Sigma_beta G_k[axis]^T ),
        //
        // i.e. the variance aggregated over the axis-block in the same way
        // `norm` aggregates ||v_k||^2. As for the scalar block, we deflate the
        // squared block norm by this noise floor to obtain the credible squared
        // magnitude and shrink poorly-determined responses toward zero:
        //
        //   g_k^credible^2 = max( g_k^2 - tr(Cov(v_k)) , 0 ),
        //   r_k^snr        = sqrt( g_k^credible^2 + eps^2 ),   w_k = 1 / r_k^snr.
        //
        // A block whose norm is credibly large (g_k^2 >> tr Cov) keeps a small
        // weight (real feature, left un-penalized); a block whose norm is
        // dominated by posterior variance is pulled toward zero by a large
        // weight (noise suppressed). With `variance == 0` this recovers
        // `surrogateweights` exactly.
        let eps2 = self.epsilon * self.epsilon;
        let weight = Array1::from_iter(self.norm.iter().zip(variance.iter()).map(|(&g, &v)| {
            let credible_sq = (g * g - v.max(0.0)).max(0.0);
            let r = (credible_sq + eps2).sqrt();
            (1.0 / r).clamp(weight_floor, weight_ceiling)
        }));
        let invweight = weight.mapv(|u| 1.0 / u);
        (weight, invweight)
    }

    fn directionalhessian_blocks(&self, direction_blocks: &Array2<f64>) -> Vec<Array2<f64>> {
        // Exact grouped directional third derivative for the slope penalty.
        //
        // For each collocation block k:
        //   v_k = G_k beta,
        //   q_k = G_k u,
        //   r_k = sqrt(||v_k||^2 + eps^2),
        //
        // the exact Hessian block for psi(g; eps) = sqrt(g^2 + eps^2) - eps is
        //   B_k,
        //   B_k = (1 / r_k) I - v_k v_k^T / r_k^3.
        //
        // Differentiating B_k along u gives
        //   M_k(u)
        //   = -(v_k^T q_k / r_k^3) I
        //     - (q_k v_k^T + v_k q_k^T) / r_k^3
        //     + 3 (v_k^T q_k) v_k v_k^T / r_k^5.
        //
        // This expression must be symmetric because it is the directional
        // derivative of the symmetric matrix
        //
        //   B_k = (1 / r_k) I - v_k v_k^T / r_k^3.
        //
        // The full directional penalty Hessian map is then
        //   D(H_g)[u] = lambda_g * sum_k G_k^T M_k(u) G_k.
        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
        for (k, (v, q)) in self
            .signal_blocks
            .rows()
            .into_iter()
            .zip(direction_blocks.rows().into_iter())
            .enumerate()
        {
            let dim = v.len();
            let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
            let r3 = self.radius[k].powi(3);
            let r5 = self.radius[k].powi(5);
            let mut block = Array2::<f64>::eye(dim);
            block.mapv_inplace(|x| -dot * x / r3);
            for i in 0..dim {
                for j in 0..dim {
                    block[[i, j]] -= (q[i] * v[j] + v[i] * q[j]) / r3;
                    block[[i, j]] += 3.0 * dot * v[i] * v[j] / r5;
                }
            }
            out.push(block);
        }
        out
    }

    fn log_epsilon_betahessian_blocks(&self) -> Vec<Array2<f64>> {
        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
            let dim = row.len();
            let r3 = self.radius[k].powi(3);
            let r5 = self.radius[k].powi(5);
            let mut block = Array2::<f64>::eye(dim);
            let eps2 = self.epsilon * self.epsilon;
            block.mapv_inplace(|v| -eps2 * v / r3);
            for i in 0..dim {
                for j in 0..dim {
                    block[[i, j]] += 3.0 * eps2 * row[i] * row[j] / r5;
                }
            }
            out.push(block);
        }
        out
    }

    fn log_epsilon_beta_mixed_second_blocks(&self) -> Array2<f64> {
        let mut out = self.signal_blocks.clone();
        let eps2 = self.epsilon * self.epsilon;
        for (k, mut row) in out.rows_mut().into_iter().enumerate() {
            let norm2 = self.norm[k] * self.norm[k];
            let scale = eps2 * (eps2 - 2.0 * norm2) / self.radius[k].powi(5);
            row.mapv_inplace(|v| v * scale);
        }
        out
    }

    fn log_epsilon_betahessian_second_blocks(&self) -> Vec<Array2<f64>> {
        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
        let eps2 = self.epsilon * self.epsilon;
        for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
            let dim = row.len();
            let norm2 = self.norm[k] * self.norm[k];
            let r5 = self.radius[k].powi(5);
            let r7 = self.radius[k].powi(7);
            let mut block = Array2::<f64>::eye(dim);
            block.mapv_inplace(|v| eps2 * (eps2 - 2.0 * norm2) * v / r5);
            for i in 0..dim {
                for j in 0..dim {
                    block[[i, j]] += 3.0 * eps2 * (2.0 * norm2 - 3.0 * eps2) * row[i] * row[j] / r7;
                }
            }
            out.push(block);
        }
        out
    }

    fn log_epsilon_betahessian_directional_blocks(
        &self,
        direction_blocks: &Array2<f64>,
    ) -> Vec<Array2<f64>> {
        let mut out = Vec::with_capacity(self.signal_blocks.nrows());
        let eps2 = self.epsilon * self.epsilon;
        for (k, (v, q)) in self
            .signal_blocks
            .rows()
            .into_iter()
            .zip(direction_blocks.rows().into_iter())
            .enumerate()
        {
            let dim = v.len();
            let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
            let r5 = self.radius[k].powi(5);
            let r7 = self.radius[k].powi(7);
            let mut block = Array2::<f64>::eye(dim);
            block.mapv_inplace(|x| 3.0 * eps2 * dot * x / r5);
            for i in 0..dim {
                for j in 0..dim {
                    block[[i, j]] += 3.0 * eps2 * (q[i] * v[j] + v[i] * q[j]) / r5;
                    block[[i, j]] -= 15.0 * eps2 * dot * v[i] * v[j] / r7;
                }
            }
            out.push(block);
        }
        out
    }
}

fn scalar_operatorgradient(operator: &Array2<f64>, coeff: &Array1<f64>) -> Array1<f64> {
    operator.t().dot(coeff)
}

fn scalar_operatorhessian(operator: &Array2<f64>, diag: &Array1<f64>) -> Array2<f64> {
    let mut weighted = operator.clone();
    for (k, &w) in diag.iter().enumerate() {
        weighted.row_mut(k).mapv_inplace(|v| v * w);
    }
    let gram = operator.t().dot(&weighted);
    (&gram + &gram.t().to_owned()) * 0.5
}

fn grouped_operatorgradient(
    d1: &Array2<f64>,
    dimension: usize,
    blocks: &Array2<f64>,
) -> Result<Array1<f64>, EstimationError> {
    if blocks.ncols() != dimension {
        crate::bail_invalid_estim!(
            "grouped gradient block dimension mismatch: got {}, expected {dimension}",
            blocks.ncols()
        );
    }
    if d1.nrows() != blocks.nrows() * dimension {
        crate::bail_invalid_estim!(
            "grouped gradient row mismatch: D1 has {} rows, blocks imply {}",
            d1.nrows(),
            blocks.nrows() * dimension
        );
    }
    let mut out = Array1::<f64>::zeros(d1.ncols());
    for k in 0..blocks.nrows() {
        let gk = d1
            .slice(s![k * dimension..(k + 1) * dimension, ..])
            .to_owned();
        out += &gk.t().dot(&blocks.row(k));
    }
    Ok(out)
}

fn grouped_operatorhessian(
    d1: &Array2<f64>,
    dimension: usize,
    blocks: &[Array2<f64>],
) -> Result<Array2<f64>, EstimationError> {
    if d1.nrows() != blocks.len() * dimension {
        crate::bail_invalid_estim!(
            "grouped Hessian row mismatch: D1 has {} rows, blocks imply {}",
            d1.nrows(),
            blocks.len() * dimension
        );
    }
    let p = d1.ncols();
    let mut out = Array2::<f64>::zeros((p, p));
    for (k, block) in blocks.iter().enumerate() {
        if block.nrows() != dimension || block.ncols() != dimension {
            crate::bail_invalid_estim!(
                "grouped Hessian block {k} has shape {}x{}, expected {}x{}",
                block.nrows(),
                block.ncols(),
                dimension,
                dimension
            );
        }
        let gk = d1
            .slice(s![k * dimension..(k + 1) * dimension, ..])
            .to_owned();
        out += &gk.t().dot(&block.dot(&gk));
    }
    Ok((&out + &out.t().to_owned()) * 0.5)
}

#[derive(Clone)]
struct SpatialPenaltyExactState {
    magnitude: CharbonnierScalarBlockState,
    gradient: CharbonnierGroupedBlockState,
    curvature: CharbonnierGroupedBlockState,
}

fn collocationgradient_blocks(
    gradrows: &Array1<f64>,
    dimension: usize,
) -> Result<Array2<f64>, EstimationError> {
    if dimension == 0 || !gradrows.len().is_multiple_of(dimension) {
        crate::bail_invalid_estim!(
            "invalid collocation gradient layout: rows={}, dimension={dimension}",
            gradrows.len()
        );
    }
    let p = gradrows.len() / dimension;
    let mut out = Array2::<f64>::zeros((p, dimension));
    for k in 0..p {
        for axis in 0..dimension {
            out[[k, axis]] = gradrows[k * dimension + axis];
        }
    }
    Ok(out)
}

fn collocationhessian_blocks(
    hessianrows: &Array1<f64>,
    dimension: usize,
) -> Result<Array2<f64>, EstimationError> {
    let block_dim = dimension.checked_mul(dimension).ok_or_else(|| {
        EstimationError::InvalidInput("invalid collocation Hessian dimension overflow".to_string())
    })?;
    if block_dim == 0 || !hessianrows.len().is_multiple_of(block_dim) {
        crate::bail_invalid_estim!(
            "invalid collocation Hessian layout: rows={}, dimension={dimension}",
            hessianrows.len()
        );
    }
    let p = hessianrows.len() / block_dim;
    let mut out = Array2::<f64>::zeros((p, block_dim));
    for k in 0..p {
        for idx in 0..block_dim {
            out[[k, idx]] = hessianrows[k * block_dim + idx];
        }
    }
    Ok(out)
}

impl SpatialPenaltyExactState {
    fn from_beta_local(
        beta_local: ArrayView1<'_, f64>,
        cache: &SpatialOperatorRuntimeCache,
        epsilons: [f64; 3],
    ) -> Result<Self, EstimationError> {
        // Exact collocation-state extraction for the three Charbonnier penalty blocks.
        //
        // For one spatial smooth term with coefficient vector beta_local, the exact
        // operator-decomposition penalty is built from three collocation images:
        //
        //   magnitude:  f = D0 beta_local
        //   slope:      v_k = G_k beta_local
        //   curvature:  H_k = D2_k beta_local
        //
        // where the gradient operator is stored in row-stacked form:
        //
        //   D1 beta_local in R^(P * d),
        //   row layout = (point 0, axis 0..d-1), (point 1, axis 0..d-1), ...
        //   D2 beta_local in R^(P * d * d),
        //   row layout = (point, Hessian axis_a, Hessian axis_b).
        //
        // so we first reshape that stacked vector into the grouped block array
        //
        //   [v_0^T
        //    ...
        //    v_(P-1)^T]  in R^(P x d).
        //
        // The three exact Charbonnier block states then carry:
        //   - the raw operator signals,
        //   - their radii sqrt(signal^2 + eps^2) or sqrt(||v_k||^2 + eps^2),
        //   - and all exact derivatives derived from those radii.
        //
        // This is the canonical translation from coefficient-space beta to the
        // penalty-side mathematical objects used throughout the implementation.
        let gradientrows = cache.d1.dot(&beta_local);
        let hessianrows = cache.d2.dot(&beta_local);
        Ok(Self {
            magnitude: CharbonnierScalarBlockState::from_signal(
                cache.d0.dot(&beta_local),
                epsilons[0],
            ),
            gradient: CharbonnierGroupedBlockState::from_signal_blocks(
                collocationgradient_blocks(&gradientrows, cache.dimension)?,
                epsilons[1],
            ),
            curvature: CharbonnierGroupedBlockState::from_signal_blocks(
                collocationhessian_blocks(&hessianrows, cache.dimension)?,
                epsilons[2],
            ),
        })
    }

    fn absolute_collocation_magnitudes(&self) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
        (
            self.magnitude.absolute_signal(),
            self.gradient.norm_signal(),
            self.curvature.norm_signal(),
        )
    }
}

fn robust_epsilon_from_samples(values: &[f64], min_epsilon_cfg: f64) -> f64 {
    if values.is_empty() {
        return min_epsilon_cfg.max(1e-12);
    }
    let mut clean = values
        .iter()
        .copied()
        .filter(|v| v.is_finite() && *v >= 0.0)
        .collect::<Vec<_>>();
    if clean.is_empty() {
        return min_epsilon_cfg.max(1e-12);
    }
    clean.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let n = clean.len();
    let median = quantile_from_sorted(&clean, 0.5);
    let q75 = quantile_from_sorted(&clean, 0.75);
    let q95 = quantile_from_sorted(&clean, 0.95);

    let mut abs_dev = clean
        .iter()
        .map(|v| (v - median).abs())
        .filter(|v| v.is_finite())
        .collect::<Vec<_>>();
    abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let mad = 1.4826 * quantile_from_sorted(&abs_dev, 0.5);

    // Charbonnier/MM requires eps bounded away from zero:
    //   u(t0) = 1 / (2*sqrt(t0^2 + eps^2)) ~ 1/(2*eps) near t0=0.
    // Use robust pilot scale:
    //   s = max(median(z), 1.4826*MAD(z), Q75(z)).
    // If s is tiny (<= delta), fallback to:
    //   s <- max(Q95(z), RMS(z)).
    // If still tiny, fallback to absolute floor s_min.
    // Then eps = kappa * s.
    // Primary robust scale: s = max(median, 1.4826*MAD, Q75).
    let mut scale = median.max(mad).max(q75);

    // Safety threshold delta and absolute floor s_min.
    let delta = (f64::EPSILON.sqrt() * q95.max(1.0))
        .max(min_epsilon_cfg)
        .max(1e-12);
    let s_min = min_epsilon_cfg.max(1e-12);

    // If robust scale is tiny, use high-quantile / RMS fallback.
    if scale <= delta {
        let rms = (clean.iter().map(|v| v * v).sum::<f64>() / n as f64).sqrt();
        scale = q95.max(rms);
    }
    if scale <= delta {
        scale = s_min;
    }

    // Start near the observed operator scale so the optimizer begins in a
    // neutral regime where both quadratic and linear behavior are reachable.
    let kappa = 1.0_f64;
    (kappa * scale).max(s_min)
}

fn extract_spatial_operator_runtime_caches(
    spec: &TermCollectionSpec,
    design: &TermCollectionDesign,
) -> Result<Vec<SpatialOperatorRuntimeCache>, EstimationError> {
    let smooth_start = design
        .design
        .ncols()
        .saturating_sub(design.smooth.total_smooth_cols());
    let mut out = Vec::<SpatialOperatorRuntimeCache>::new();
    for (term_idx, (termspec, term_fit)) in spec
        .smooth_terms
        .iter()
        .zip(design.smooth.terms.iter())
        .enumerate()
    {
        let Some(global_base_idx) = smooth_term_penalty_index(spec, design, term_idx) else {
            continue;
        };
        let mut active_local_idx = 0usize;
        let mut mass_local_idx = None;
        let mut tension_local_idx = None;
        let mut stiffness_local_idx = None;
        let mut primary_local_idx = None;
        let mut mass_norm = None;
        let mut tension_norm = None;
        let mut stiffness_norm = None;
        let mut primary_norm = None;
        for info in &term_fit.penaltyinfo_local {
            if !info.active {
                continue;
            }
            match info.source {
                PenaltySource::OperatorMass => {
                    mass_local_idx = Some(active_local_idx);
                    mass_norm = Some(info.normalization_scale);
                }
                PenaltySource::OperatorTension => {
                    tension_local_idx = Some(active_local_idx);
                    tension_norm = Some(info.normalization_scale);
                }
                PenaltySource::OperatorStiffness => {
                    stiffness_local_idx = Some(active_local_idx);
                    stiffness_norm = Some(info.normalization_scale);
                }
                PenaltySource::Primary => {
                    primary_local_idx = Some(active_local_idx);
                    primary_norm = Some(info.normalization_scale);
                }
                _ => {}
            }
            active_local_idx += 1;
        }
        // The curvature channel maps onto the explicit D2 Stiffness operator when
        // the term ships one (Matérn collocation / Duchon `all_active()`); for the
        // *default* Duchon penalty Stiffness is deliberately disabled because the
        // always-on Primary RKHS Gram is the exact curvature (#858). In that
        // {mass, tension, Primary} layout the Primary penalty is the curvature
        // channel: the adaptive Charbonnier surrogate (built from the D2
        // collocation operator below) reweights and replaces it, so the curvature
        // global index / normalization come from Primary.
        let (curvature_local_idx, curvature_norm) = match (stiffness_local_idx, stiffness_norm) {
            (Some(idx), Some(norm)) => (Some(idx), Some(norm)),
            _ => (primary_local_idx, primary_norm),
        };
        let (
            Some(mass_local),
            Some(tension_local),
            Some(stiffness_local),
            Some(mass_scale),
            Some(tension_scale),
            Some(stiffness_scale),
        ) = (
            mass_local_idx,
            tension_local_idx,
            curvature_local_idx,
            mass_norm,
            tension_norm,
            curvature_norm,
        )
        else {
            continue;
        };
        let mass_global_idx = global_base_idx + mass_local;
        let tension_global_idx = global_base_idx + tension_local;
        let stiffness_global_idx = global_base_idx + stiffness_local;

        let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
            match (&termspec.basis, &term_fit.metadata) {
                (
                    SmoothBasisSpec::Matern { feature_cols, .. },
                    BasisMetadata::Matern {
                        centers,
                        length_scale,
                        nu,
                        include_intercept,
                        identifiability_transform,
                        aniso_log_scales,
                        input_scales,
                        ..
                    },
                ) => {
                    // Match the σ_geom-compensated effective length scale the
                    // design (and shipped penalties) use against the standardized
                    // centers; the raw metadata length_scale lives in original
                    // coordinates and would put this overlay on a different kernel
                    // range than the penalties it scales (#706).
                    let collocation_length_scale = match input_scales.as_deref() {
                        Some(scales) => {
                            compensate_length_scale_for_standardization(*length_scale, scales)
                        }
                        None => *length_scale,
                    };
                    let ops = build_matern_collocation_operator_matrices(
                        centers.view(),
                        None,
                        collocation_length_scale,
                        *nu,
                        *include_intercept,
                        identifiability_transform.as_ref().map(|z| z.view()),
                        aniso_log_scales.as_deref(),
                    )?;
                    (
                        feature_cols.clone(),
                        ops.d0,
                        ops.d1,
                        ops.d2,
                        ops.collocation_points,
                        centers.ncols(),
                        false,
                    )
                }
                (
                    SmoothBasisSpec::Duchon { feature_cols, .. },
                    BasisMetadata::Duchon {
                        centers,
                        length_scale,
                        power,
                        nullspace_order,
                        identifiability_transform,
                        input_scales,
                        aniso_log_scales,
                        operator_collocation_points: Some(collocation_points),
                        ..
                    },
                ) => {
                    let collocation_length_scale = match (length_scale, input_scales.as_deref()) {
                        (Some(ls), Some(scales)) => {
                            Some(compensate_length_scale_for_standardization(*ls, scales))
                        }
                        (Some(ls), None) => Some(*ls),
                        (None, _) => None,
                    };
                    let ops =
                        crate::basis::build_duchon_collocation_operator_matriceswithworkspace(
                            centers.view(),
                            collocation_points.view(),
                            None,
                            collocation_length_scale,
                            *power,
                            *nullspace_order,
                            aniso_log_scales.as_deref(),
                            identifiability_transform.as_ref().map(|z| z.view()),
                            2,
                            &mut BasisWorkspace::default(),
                        )?;
                    (
                        feature_cols.clone(),
                        ops.d0,
                        ops.d1,
                        ops.d2,
                        ops.collocation_points,
                        centers.ncols(),
                        true,
                    )
                }
                _ => continue,
            };
        if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
            let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
            for mut row in d0.rows_mut() {
                row -= &means;
            }
        }

        // Runtime operator caches must live on the same normalized penalty scale as the
        // shipped design penalties. The basis builders normalize S0=D0'D0, S1=D1'D1, and
        // S2=D2'D2 before exposing them as smoothing blocks, recording the corresponding
        // Frobenius norms in penaltyinfo_local.normalization_scale. If the exact adaptive
        // path uses raw collocation operators here, then its Charbonnier penalties live on a
        // different geometry from the ordinary Matérn/Duchon penalties:
        //
        //   raw quadratic limit:        beta' (D'D) beta
        //   shipped design penalty:     beta' (D'D / c) beta
        //
        // The correct operator-level normalization is therefore
        //
        //   D_norm = D / sqrt(c),
        //
        // so that D_norm' D_norm = (D'D)/c matches the design penalty exactly. Without this,
        // adaptive lambdas compensate for hidden operator-scale mismatches and are no longer
        // comparable to the baseline smoothing parameters.
        let mass_scale = mass_scale.max(1e-12).sqrt();
        let tension_scale = tension_scale.max(1e-12).sqrt();
        let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
        d0.mapv_inplace(|v| v / mass_scale);
        d1.mapv_inplace(|v| v / tension_scale);
        d2.mapv_inplace(|v| v / stiffness_scale);

        let coeff_global_range =
            (smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
        if d0.ncols() != coeff_global_range.len()
            || d1.ncols() != coeff_global_range.len()
            || d2.ncols() != coeff_global_range.len()
        {
            crate::bail_invalid_estim!(
                "spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
                term_fit.name,
                d0.ncols(),
                d1.ncols(),
                d2.ncols(),
                coeff_global_range.len()
            );
        }
        out.push(SpatialOperatorRuntimeCache {
            termname: term_fit.name.clone(),
            feature_cols,
            coeff_global_range,
            mass_penalty_global_idx: mass_global_idx,
            tension_penalty_global_idx: tension_global_idx,
            stiffness_penalty_global_idx: stiffness_global_idx,
            d0,
            d1,
            d2,
            collocation_points,
            dimension: dim,
        });
    }
    Ok(out)
}

/// Posterior variance of a scalar collocation operator response under the
/// working-Laplace posterior `beta ~ N(beta_hat, Sigma_local)`.
///
/// For operator row `D_k` (one row of `D0`) acting on the term-local coefficient
/// block, `Var((D beta)_k) = D_k Sigma_local D_k^T = (D Sigma_local D^T)_kk`.
/// We compute it without forming `D Sigma D^T` densely: for each row we evaluate
/// `s_k = Sigma_local D_k^T` (one matrix-vector product) and then `D_k . s_k`.
/// `Sigma_local` is the sub-block of the global conditional covariance
/// `Sigma_beta = H^{-1}` indexed by the term's `coeff_global_range`, i.e. the
/// covariance proxy is the already-materialized inner working-Laplace inverse;
/// no second factorization is formed.
fn scalar_operator_response_variance(
    operator: &Array2<f64>,
    cov_local: &Array2<f64>,
) -> Array1<f64> {
    Array1::from_iter(operator.rows().into_iter().map(|row| {
        let s = cov_local.dot(&row);
        row.dot(&s).max(0.0)
    }))
}

/// Posterior second-moment variance aggregated over each grouped collocation
/// block (gradient/curvature). The grouped operator is stored row-stacked with
/// `block_dim` rows per collocation point (`d` axes for the gradient, `d*d` for
/// the Hessian). For block `k`,
///
///   v_k = G_k beta,   Cov(v_k) = G_k Sigma_local G_k^T   (block_dim x block_dim),
///   variance_k = tr(Cov(v_k)) = sum_axis ( G_k[axis] Sigma_local G_k[axis]^T ),
///
/// which matches how `CharbonnierGroupedBlockState::norm` aggregates
/// `||v_k||^2 = sum_axis (G_k[axis] beta)^2` across the axis-block.
fn grouped_operator_response_variance(
    operator: &Array2<f64>,
    block_dim: usize,
    cov_local: &Array2<f64>,
) -> Result<Array1<f64>, EstimationError> {
    if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
        crate::bail_invalid_estim!(
            "grouped variance row layout invalid: rows={}, block_dim={block_dim}",
            operator.nrows()
        );
    }
    let p = operator.nrows() / block_dim;
    let mut out = Array1::<f64>::zeros(p);
    for k in 0..p {
        let mut acc = 0.0;
        for axis in 0..block_dim {
            let row = operator.row(k * block_dim + axis);
            let s = cov_local.dot(&row);
            acc += row.dot(&s);
        }
        out[k] = acc.max(0.0);
    }
    Ok(out)
}

fn compute_spatial_adaptiveweights_for_beta(
    beta: &Array1<f64>,
    caches: &[SpatialOperatorRuntimeCache],
    epsilon_0: f64,
    epsilon_g: f64,
    epsilon_c: f64,
    weight_floor: f64,
    weight_ceiling: f64,
    beta_covariance: Option<&Array2<f64>>,
) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
    // Charbonnier / pseudo-Huber MM derivation (per collocation scalar t):
    //   psi(t; eps) = sqrt(t^2 + eps^2) - eps
    // and for reference t0 the tangent majorizer in t^2 gives:
    //   psi(t) <= 0.5 * w(t0) * t^2 + const(t0),
    //   w(t0) = 1 / sqrt(t0^2 + eps^2).
    //
    // We apply this to:
    //   t = f_k = |f(z_k)|             (magnitude),
    //   t = g_k = ||nabla f(z_k)||_2   (gradient magnitude),
    //   t = c_k = ||D²f(z_k)||_F       (full Hessian curvature),
    // both computed from beta^(t-1).
    //
    // These w values define the quadratic surrogate penalties:
    //   K0 = D0_con^T W_0 D0_con,  W_0 = diag(w_0)
    //   K1 = D1_con^T W_g D1_con,  W_g = diag(w_g) \otimes I_d  (k,axis order)
    //   K2 = D2_con^T W_c D2_con,  W_c = diag(w_c) \otimes I_(d*d).
    //
    // We clamp w directly, then derive inv_w=1/w for diagnostics and row scaling.
    //
    // Posterior-SNR reweighting (magic by default): when the inner working-Laplace
    // conditional covariance `Sigma_beta = H^{-1}` is available we replace the
    // squared point-estimate radius `t_k^2 + eps^2` by the credible (noise-floor-
    // corrected) second moment `max(t_k^2 - Var((D beta)_k), 0) + eps^2`, with
    // `Var = (D Sigma_beta D^T)_kk`. This stops the weight from leaving derivatives
    // un-penalized just because they are large but poorly determined: such
    // responses are shrunk toward zero (large weight, strong smoothing), while
    // credibly large derivatives (real edges) keep their small weight. `Sigma_beta`
    // here is the already-formed inner Hessian inverse from the final exact-family
    // solve — no second factorization is built; we only reuse the materialized
    // covariance. When the covariance is unavailable (`None`) the variance is zero
    // and this degrades *exactly* to the old magnitude-only radius.
    caches
        .iter()
        .map(|cache| {
            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
            let exact = SpatialPenaltyExactState::from_beta_local(
                beta_local,
                cache,
                [epsilon_0, epsilon_g, epsilon_c],
            )?;
            let cov_local = beta_covariance.map(|cov| {
                cov.slice(s![
                    cache.coeff_global_range.clone(),
                    cache.coeff_global_range.clone()
                ])
                .to_owned()
            });
            let dim = cache.dimension;
            let (var_0, var_g, var_c) = match cov_local.as_ref() {
                Some(cov) => (
                    scalar_operator_response_variance(&cache.d0, cov),
                    grouped_operator_response_variance(&cache.d1, dim, cov)?,
                    grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
                ),
                None => (
                    Array1::<f64>::zeros(exact.magnitude.signal.len()),
                    Array1::<f64>::zeros(exact.gradient.norm.len()),
                    Array1::<f64>::zeros(exact.curvature.norm.len()),
                ),
            };
            let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
                &var_0,
                weight_floor,
                weight_ceiling,
            );
            let (_, inv_g) =
                exact
                    .gradient
                    .surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
            let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
                &var_c,
                weight_floor,
                weight_ceiling,
            );
            Ok(SpatialAdaptiveWeights {
                inv_magweight: inv_0,
                invgradweight: inv_g,
                inv_lapweight: inv_c,
            })
        })
        .collect()
}

fn compute_initial_epsilons(
    beta: &Array1<f64>,
    caches: &[SpatialOperatorRuntimeCache],
    min_epsilon: f64,
) -> Result<(f64, f64, f64), EstimationError> {
    let mut fvals = Vec::<f64>::new();
    let mut gvals = Vec::<f64>::new();
    let mut cvals = Vec::<f64>::new();
    for cache in caches {
        let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
        let exact = SpatialPenaltyExactState::from_beta_local(
            beta_local,
            cache,
            [min_epsilon, min_epsilon, min_epsilon],
        )?;
        let (f, g, c) = exact.absolute_collocation_magnitudes();
        fvals.extend(f.iter().copied());
        gvals.extend(g.iter().copied());
        cvals.extend(c.iter().copied());
    }
    // Robust epsilon initialization from pilot magnitudes:
    //   s = max(median(z), 1.4826*MAD(z), Q75(z)),
    //   if s is tiny then fallback to max(Q95(z), RMS(z)),
    //   if still tiny then use absolute floor min_epsilon.
    // Epsilon is then kappa * s.
    let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
    let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
    let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
    Ok((eps_0, eps_g, eps_c))
}

fn exact_spatial_adaptive_penalty_index_set(
    caches: &[SpatialOperatorRuntimeCache],
) -> BTreeSet<usize> {
    let mut out = BTreeSet::new();
    for cache in caches {
        out.insert(cache.mass_penalty_global_idx);
        out.insert(cache.tension_penalty_global_idx);
        out.insert(cache.stiffness_penalty_global_idx);
    }
    out
}

fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
    let mut out = Vec::with_capacity(cache_count * 3 + 3);
    for cache_index in 0..cache_count {
        out.push(SpatialAdaptiveHyperSpec {
            cache_index,
            kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
        });
        out.push(SpatialAdaptiveHyperSpec {
            cache_index,
            kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
        });
        out.push(SpatialAdaptiveHyperSpec {
            cache_index,
            kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
        });
    }
    out.push(SpatialAdaptiveHyperSpec {
        cache_index: 0,
        kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
    });
    out.push(SpatialAdaptiveHyperSpec {
        cache_index: 0,
        kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
    });
    out.push(SpatialAdaptiveHyperSpec {
        cache_index: 0,
        kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
    });
    out
}

fn penalty_matrixwith_local_block(
    total_dim: usize,
    coeff_range: Range<usize>,
    local: &Array2<f64>,
) -> Array2<f64> {
    let mut out = Array2::<f64>::zeros((total_dim, total_dim));
    out.slice_mut(s![coeff_range.clone(), coeff_range])
        .assign(local);
    out
}

fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
    baseline: FittedTermCollection,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    family: LikelihoodSpec,
    options: &FitOptions,
    runtime_caches: &[SpatialOperatorRuntimeCache],
) -> Result<FittedTermCollection, EstimationError> {
    // Exact adaptive-regularization hyperfit.
    //
    // This replaces the old MM-plus-approximate hyperfit with the
    // exact pseudo-Laplace objective agreed in the math notes:
    //
    //   L_tilde(theta)
    //   = J(beta_hat(theta); theta) + 0.5 log det H(beta_hat(theta), theta),
    //
    // where:
    //   - beta_hat(theta) is the exact inner mode of the true nonquadratic
    //     Charbonnier-penalized objective,
    //   - theta contains:
    //       * retained quadratic log-lambdas for non-adaptive penalties,
    //       * one log-lambda per adaptive operator block,
    //       * three global log-epsilons shared by every adaptive spatial term,
    //   - H is the exact beta-Hessian of the true objective at the mode.
    //
    // Implementation structure:
    //   1. keep ordinary quadratic penalties that are unrelated to adaptive
    //      spatial terms in the standard outer-rho path;
    //   2. move the adaptive Charbonnier penalties into a one-block exact-Newton
    //      custom family so the inner solve uses the real model rather than an
    //      MM surrogate;
    //   3. expose exact psi-gradients for adaptive log-lambda / log-epsilon
    //      coordinates through the custom-family pseudo-Laplace hook;
    //   4. refit once at the optimized hyperparameters with all penalties frozen
    //      inside the exact family, so covariance and final diagnostics are
    //      computed on the same exact surface.
    let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
    let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
    let p_total = baseline.design.design.ncols();
    struct RetainedPenaltySetup {
        global_idx: usize,
        global_penalty: Array2<f64>,
        nullspace_dim: usize,
        log_lambda: f64,
        col_range: Range<usize>,
        hessian_piece: Array2<f64>,
    }
    use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
    let retained_setups = baseline
        .design
        .penalties
        .par_iter()
        .enumerate()
        .map(|(idx, bp)| {
            if adaptive_penalty_indices.contains(&idx) {
                return None;
            }
            let lambda = baseline.fit.lambdas[idx];
            Some(RetainedPenaltySetup {
                global_idx: idx,
                global_penalty: bp.to_global(p_total),
                nullspace_dim: baseline
                    .design
                    .nullspace_dims
                    .get(idx)
                    .copied()
                    .unwrap_or(0),
                log_lambda: lambda.max(1e-12).ln(),
                col_range: bp.col_range.clone(),
                hessian_piece: bp.local.mapv(|v| lambda * v),
            })
        })
        .collect::<Vec<_>>();
    let retained_count = retained_setups
        .iter()
        .filter(|setup| setup.is_some())
        .count();
    let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
    let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
    let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
    let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
    let mut fixed_quadratichessian = Array2::<f64>::zeros((p_total, p_total));
    for setup in retained_setups.into_iter().flatten() {
        retained_penalties.push(setup.global_penalty);
        retained_nullspace_dims.push(setup.nullspace_dim);
        retained_log_lambdas.push(setup.log_lambda);
        retained_global_indices.push(setup.global_idx);
        fixed_quadratichessian
            .slice_mut(s![setup.col_range.clone(), setup.col_range])
            .scaled_add(1.0, &setup.hessian_piece);
    }

    let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
        &baseline.fit.beta,
        runtime_caches,
        adaptive_opts.min_epsilon,
    )?;
    let mut initial_theta =
        Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
    for (idx, value) in retained_log_lambdas.iter().enumerate() {
        initial_theta[idx] = *value;
    }
    let adaptive_log_lambda_components = runtime_caches
        .par_iter()
        .map(|cache| {
            [
                baseline.fit.lambdas[cache.mass_penalty_global_idx]
                    .max(1e-12)
                    .ln(),
                baseline.fit.lambdas[cache.tension_penalty_global_idx]
                    .max(1e-12)
                    .ln(),
                baseline.fit.lambdas[cache.stiffness_penalty_global_idx]
                    .max(1e-12)
                    .ln(),
            ]
        })
        .collect::<Vec<_>>();
    let mut at = retained_penalties.len();
    for logs in &adaptive_log_lambda_components {
        initial_theta[at] = logs[0];
        initial_theta[at + 1] = logs[1];
        initial_theta[at + 2] = logs[2];
        at += 3;
    }
    initial_theta[at] = eps_0_init.max(adaptive_opts.min_epsilon).ln();
    initial_theta[at + 1] = eps_g_init.max(adaptive_opts.min_epsilon).ln();
    initial_theta[at + 2] = eps_c_init.max(adaptive_opts.min_epsilon).ln();

    let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
    let zero_psi_op: std::sync::Arc<dyn crate::custom_family::CustomFamilyPsiDerivativeOperator> =
        std::sync::Arc::new(crate::custom_family::ZeroPsiDerivativeOperator::new(
            baseline.design.design.nrows(),
            baseline.design.design.ncols(),
        ));
    let derivative_blocks = vec![
        hyperspecs
            .par_iter()
            .map(|_| CustomFamilyBlockPsiDerivative {
                penalty_index: None,
                x_psi: Array2::<f64>::zeros((0, 0)),
                s_psi: Array2::<f64>::zeros((0, 0)),
                s_psi_components: None,
                s_psi_penalty_components: None,
                x_psi_psi: None,
                s_psi_psi: None,
                s_psi_psi_components: None,
                s_psi_psi_penalty_components: None,
                implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
                implicit_axis: 0,
                implicit_group_id: None,
            })
            .collect::<Vec<_>>(),
    ];

    let mixture_link_state = options
        .mixture_link
        .clone()
        .as_ref()
        .map(state_fromspec)
        .transpose()
        .map_err(EstimationError::InvalidInput)?;
    let sas_link_state = options
        .sas_link
        .map(|spec| {
            if family.is_binomial_beta_logistic() {
                state_from_beta_logisticspec(spec)
            } else {
                state_from_sasspec(spec)
            }
        })
        .transpose()
        .map_err(EstimationError::InvalidInput)?;
    let latent_cloglog_state = options.latent_cloglog;
    let shared_y = Arc::new(y.to_owned());
    let sharedweights = Arc::new(weights.to_owned());
    let shared_design = baseline
        .design
        .design
        .try_to_dense_arc("spatial adaptive exact hyperfit design")
        .map_err(EstimationError::InvalidInput)?;
    let shared_offset = Arc::new(offset.to_owned());
    let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
    let shared_hyperspecs = Arc::new(hyperspecs.clone());
    let zero_quadratic = Arc::new(Array2::<f64>::zeros((
        baseline.design.design.ncols(),
        baseline.design.design.ncols(),
    )));
    let base_family = SpatialAdaptiveExactFamily {
        family: family.clone(),
        latent_cloglog_state,
        mixture_link_state: mixture_link_state.clone(),
        sas_link_state,
        y: shared_y.clone(),
        weights: sharedweights.clone(),
        design: shared_design.clone(),
        offset: shared_offset.clone(),
        linear_constraints: baseline.design.linear_constraints.clone(),
        runtime_caches: shared_runtime_caches.clone(),
        adaptive_params: Vec::new(),
        fixed_quadratichessian: zero_quadratic.clone(),
        hyperspecs: shared_hyperspecs.clone(),
        exact_eval_cache: Arc::new(Mutex::new(None)),
    };

    let rho_dim = retained_penalties.len();
    let operator_slots_end = rho_dim + runtime_caches.len() * 3;
    // Every slot's box is `initial_theta[idx] ± WINDOW` clamped into a
    // per-slot [floor, cap]. Retained-λ previously used a scale-blind
    // ±30 absolute interval, which on small-n / weakly-identified Duchon
    // fits let those lambdas wander to the exp(-30) floor and produce
    // near-interpolant solutions. Anchoring on baseline log-λ inherits the
    // baseline REML's scale calibration so the overlay can only refine
    // within an exp(±6) ≈ 400× band of the well-posed baseline regime,
    // matching the discipline already applied to operator and epsilon
    // slots.
    const UNIFIED_LOG_WINDOW: f64 = 6.0;
    const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
    const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
    const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
    const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
    let epsilon_floor_log = adaptive_opts.min_epsilon.max(1e-12).ln();
    let anchored_bound = |idx: usize, sign: f64| -> f64 {
        let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
        if idx < rho_dim {
            raw.clamp(
                RETAINED_LAMBDA_LOG_LOWER_FLOOR,
                RETAINED_LAMBDA_LOG_UPPER_CAP,
            )
        } else if idx < operator_slots_end {
            raw.clamp(
                OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
                OPERATOR_LAMBDA_LOG_UPPER_CAP,
            )
        } else {
            raw.max(epsilon_floor_log)
        }
    };
    let eps_lower =
        Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
    let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
    let blockspec = ParameterBlockSpec {
        name: "eta".to_string(),
        design: baseline.design.design.clone(),
        offset: offset.to_owned(),
        penalties: retained_penalties
            .iter()
            .cloned()
            .map(PenaltyMatrix::Dense)
            .collect(),
        nullspace_dims: retained_nullspace_dims.clone(),
        initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
        initial_beta: Some(baseline.fit.beta.clone()),
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    let screening_cap = Arc::new(AtomicUsize::new(0));
    let outer_opts = BlockwiseFitOptions {
        inner_max_cycles: options.max_iter,
        inner_tol: options.tol,
        outer_max_iter: options.max_iter,
        outer_tol: options.tol,
        compute_covariance: false,
        screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
        ..BlockwiseFitOptions::default()
    };

    use crate::solver::outer_strategy::{
        DeclaredHessianForm, Derivative, HessianResult, OuterEval, OuterProblem,
    };

    struct SpatialAdaptiveOuterState {
        warm_cache: Option<CustomFamilyWarmStart>,
        last_eval: Option<(
            Array1<f64>,
            f64,
            Array1<f64>,
            HessianResult,
            CustomFamilyWarmStart,
        )>,
    }

    let n_theta = initial_theta.len();

    // Clamp theta to the asymmetric epsilon bounds that run_outer's symmetric
    // rho_bound cannot express directly.
    let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
    let clamp_theta = {
        let lo = eps_lower;
        let hi = eps_upper;
        move |theta: &Array1<f64>| -> Array1<f64> {
            let mut clamped = theta.clone();
            for i in 0..clamped.len() {
                clamped[i] = clamped[i].clamp(lo[i], hi[i]);
            }
            clamped
        }
    };

    let decode_theta = |theta: &Array1<f64>| -> (Array1<f64>, Vec<SpatialAdaptiveTermHyperParams>) {
        let rho = theta.slice(s![..rho_dim]).to_owned();
        let adaptive_lambda_start = rho_dim;
        let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
        let eps = [
            theta[adaptive_lambda_end].exp(),
            theta[adaptive_lambda_end + 1].exp(),
            theta[adaptive_lambda_end + 2].exp(),
        ];
        let adaptive_params = runtime_caches
            .iter()
            .enumerate()
            .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
                lambda: [
                    theta[adaptive_lambda_start + cache_idx * 3].exp(),
                    theta[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
                    theta[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
                ],
                epsilon: eps,
            })
            .collect::<Vec<_>>();
        (rho, adaptive_params)
    };
    let analytic_outer_hessian_available =
        crate::custom_family::joint_exact_analytic_outer_hessian_available()
            && base_family
                .exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
                .has_hessian()
            && crate::custom_family::exact_newton_outer_geometry_supports_second_order_solver(
                &base_family,
            );
    let outer_max_iter = crate::custom_family::cost_gated_first_order_max_iter(
        options.max_iter,
        base_family.coefficient_gradient_cost(std::slice::from_ref(&blockspec)),
        analytic_outer_hessian_available,
    );
    if outer_max_iter < options.max_iter {
        log::info!(
            "[OUTER] exact spatial adaptive regularization: first-order work gate reduced outer_max_iter {} -> {}",
            options.max_iter,
            outer_max_iter,
        );
    }
    // Keep the exact outer Hessian whenever the adaptive family can provide it.
    // The Charbonnier pseudo-Laplace surface mixes ordinary log-lambda
    // coordinates with adaptive λ/ε coordinates; exact curvature is the best
    // route when available. If a family cannot provide exact curvature, this
    // builder declares only the true first-order capability.
    let problem = OuterProblem::new(n_theta)
        .with_gradient(Derivative::Analytic)
        .with_hessian(if analytic_outer_hessian_available {
            DeclaredHessianForm::Either
        } else {
            DeclaredHessianForm::Unavailable
        })
        .with_fallback_policy(crate::solver::outer_strategy::FallbackPolicy::Disabled)
        .with_psi_dim(n_theta.saturating_sub(rho_dim))
        .with_tolerance(options.tol)
        .with_max_iter(outer_max_iter)
        .with_seed_config(crate::seeding::SeedConfig::default())
        .with_screening_cap(Arc::clone(&screening_cap))
        .with_initial_rho(initial_theta.clone());
    let problem = if let Some((lo, hi)) = theta_bounds {
        problem.with_bounds(lo, hi)
    } else {
        problem
    };

    let eval_outer = |st: &mut SpatialAdaptiveOuterState,
                      theta: &Array1<f64>,
                      order: crate::solver::outer_strategy::OuterEvalOrder|
     -> Result<OuterEval, EstimationError> {
        let theta = clamp_theta(theta);

        if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
            &st.last_eval
            && cached_theta.len() == theta.len()
            && cached_theta
                .iter()
                .zip(theta.iter())
                .all(|(&a, &b)| (a - b).abs() <= 1e-12)
            && (!matches!(
                order,
                crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian
            ) || analytic_outer_hessian_available)
        {
            st.warm_cache = Some(cached_warm.clone());
            return Ok(OuterEval {
                cost: *cached_cost,
                gradient: cached_grad.clone(),
                hessian: if matches!(
                    order,
                    crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian
                ) && analytic_outer_hessian_available
                {
                    cached_hess.clone()
                } else {
                    HessianResult::Unavailable
                },
                inner_beta_hint: None,
            });
        }

        let (rho, adaptive_params) = decode_theta(&theta);
        let family_eval = base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
        let need_hessian = matches!(
            order,
            crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian
        ) && analytic_outer_hessian_available;
        let result = evaluate_custom_family_joint_hyper(
            &family_eval,
            std::slice::from_ref(&blockspec),
            &outer_opts,
            &rho,
            &derivative_blocks,
            st.warm_cache.as_ref(),
            if need_hessian {
                crate::solver::estimate::reml::unified::EvalMode::ValueGradientHessian
            } else {
                crate::solver::estimate::reml::unified::EvalMode::ValueAndGradient
            },
        )
        .map_err(|e| {
            EstimationError::RemlOptimizationFailed(format!("spatial adaptive eval failed: {e}"))
        })?;
        if !result.inner_converged {
            st.warm_cache = Some(result.warm_start.clone());
            return Err(EstimationError::RemlOptimizationFailed(
                "exact spatial adaptive inner solve did not converge".to_string(),
            ));
        }
        if !result.objective.is_finite() || result.gradient.iter().any(|v| !v.is_finite()) {
            return Err(EstimationError::RemlOptimizationFailed(
                "exact spatial adaptive objective returned non-finite values".to_string(),
            ));
        }
        let hessian_result = if need_hessian {
            if !result.outer_hessian.is_analytic() {
                return Err(EstimationError::RemlOptimizationFailed(
                    "exact spatial adaptive objective did not return an exact outer Hessian"
                        .to_string(),
                ));
            }
            match result.outer_hessian.dim() {
                Some(dim) if dim == theta.len() => {}
                Some(dim) => {
                    return Err(EstimationError::RemlOptimizationFailed(format!(
                        "exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
                        theta.len(),
                    )));
                }
                None => {
                    return Err(EstimationError::RemlOptimizationFailed(
                        "exact spatial adaptive objective did not report an outer Hessian dimension"
                            .to_string(),
                    ));
                }
            }
            st.last_eval = Some((
                theta.clone(),
                result.objective,
                result.gradient.clone(),
                result.outer_hessian.clone(),
                result.warm_start.clone(),
            ));
            result.outer_hessian
        } else {
            HessianResult::Unavailable
        };
        st.warm_cache = Some(result.warm_start);
        Ok(OuterEval {
            cost: result.objective,
            gradient: result.gradient,
            hessian: hessian_result,
            inner_beta_hint: None,
        })
    };

    let mut obj = problem.build_objective_with_eval_order(
        SpatialAdaptiveOuterState {
            warm_cache: None,
            last_eval: None,
        },
        |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
            let theta = clamp_theta(theta);
            let (rho, adaptive_params) = decode_theta(&theta);
            let family_eval =
                base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
            let result = evaluate_custom_family_joint_hyper(
                &family_eval,
                std::slice::from_ref(&blockspec),
                &outer_opts,
                &rho,
                &derivative_blocks,
                st.warm_cache.as_ref(),
                crate::solver::estimate::reml::unified::EvalMode::ValueOnly,
            )
            .map_err(|e| {
                EstimationError::RemlOptimizationFailed(format!(
                    "spatial adaptive cost eval failed: {e}"
                ))
            })?;
            if !result.inner_converged {
                st.warm_cache = Some(result.warm_start);
                return Err(EstimationError::RemlOptimizationFailed(
                    "exact spatial adaptive cost inner solve did not converge".to_string(),
                ));
            }
            st.warm_cache = Some(result.warm_start);
            Ok(result.objective)
        },
        |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
            eval_outer(
                st,
                theta,
                if analytic_outer_hessian_available {
                    crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian
                } else {
                    crate::solver::outer_strategy::OuterEvalOrder::ValueAndGradient
                },
            )
        },
        |st: &mut SpatialAdaptiveOuterState,
         theta: &Array1<f64>,
         order: crate::solver::outer_strategy::OuterEvalOrder| {
            eval_outer(st, theta, order)
        },
        Some(|st: &mut SpatialAdaptiveOuterState| {
            st.warm_cache = None;
            st.last_eval = None;
        }),
        Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
            let theta = clamp_theta(theta);
            let (rho, adaptive_params) = decode_theta(&theta);
            let family_eval =
                base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
            let result = evaluate_custom_family_joint_hyper_efs(
                &family_eval,
                std::slice::from_ref(&blockspec),
                &outer_opts,
                &rho,
                &derivative_blocks,
                st.warm_cache.as_ref(),
            )
            .map_err(|e| {
                EstimationError::RemlOptimizationFailed(format!(
                    "spatial adaptive EFS eval failed: {e}"
                ))
            })?;
            if !result.inner_converged {
                st.warm_cache = Some(result.warm_start);
                return Err(EstimationError::RemlOptimizationFailed(
                    "exact spatial adaptive EFS inner solve did not converge".to_string(),
                ));
            }
            st.warm_cache = Some(result.warm_start);
            Ok(result.efs_eval)
        }),
    );

    let outer_result = problem
        .run(&mut obj, "exact spatial adaptive regularization")
        .map_err(|e| {
            EstimationError::InvalidInput(format!(
                "exact spatial adaptive outer optimization failed: {e}"
            ))
        })?;
    if !outer_result.converged {
        // The strict absolute-floor gradient criterion (`‖g‖_proj ≤ options.tol`)
        // is too tight near the box-constrained boundary of the adaptive
        // Charbonnier pseudo-Laplace objective: as the optimizer pushes ε → ∞
        // (overlay-disabled corner), λ → λ_min, the Hessian's nearly-null
        // direction lets Cauchy/Newton accept ~e-3-magnitude probe steps that
        // give cost changes well below 6-digit precision, and the projected
        // gradient floors at numerical-noise-scale (≈ 5e-6 for n≈500, cost≈
        // 3e2 fits in double precision) rather than at 0. Accept the iterate
        // when the mgcv-style relative-to-cost criterion ‖g‖_proj ≤ τ·(1+|f|)
        // is satisfied — that is the textbook REML convergence rule and is
        // exactly what `opt::GradientTolerance::relative_to_cost(τ)` would
        // have enforced if this OuterProblem path had wired it through. The
        // strict absolute floor is retained as the primary check; the
        // rel-to-cost form only kicks in once the absolute one has timed out
        // at `max_iter`, so unconverged divergent runs (which have large |g|)
        // still surface as errors.
        let rel_to_cost_threshold = options.tol * (1.0_f64 + outer_result.final_value.abs());
        // Rel-to-cost acceptance requires an actual gradient measurement;
        // `None` (cache-hit short-circuit, gradient-free path) cannot satisfy
        // the mgcv-style criterion regardless of magnitude.
        if let Some(final_grad) = outer_result
            .final_grad_norm
            .filter(|v| v.is_finite() && *v <= rel_to_cost_threshold)
        {
            log::info!(
                "[spatial-adaptive] outer optimization hit max_iter={} but \
                 projected gradient norm {:.3e} ≤ τ·(1+|f|) = {:.3e} \
                 (τ={:.3e}, |f|={:.3e}); accepting iterate under the mgcv-style \
                 relative-to-cost REML convergence criterion.",
                outer_result.iterations,
                final_grad,
                rel_to_cost_threshold,
                options.tol,
                outer_result.final_value.abs(),
            );
        } else {
            crate::bail_invalid_estim!(
                "exact spatial adaptive outer optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
                outer_result.iterations,
                outer_result.final_value,
                outer_result.final_grad_norm_report(),
            );
        }
    }
    let outer_iterations = outer_result.iterations;
    // `None` = no gradient measurement (cache-hit / gradient-free); the
    // authoritative convergence signal is `outer_converged`.
    let outer_grad_norm: Option<f64> = outer_result.final_grad_norm;
    let theta_star = outer_result.rho;
    let rho_star = theta_star.slice(s![..rho_dim]).to_owned();
    let adaptive_lambda_start = rho_dim;
    let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
    let eps_star = [
        theta_star[adaptive_lambda_end].exp(),
        theta_star[adaptive_lambda_end + 1].exp(),
        theta_star[adaptive_lambda_end + 2].exp(),
    ];
    let adaptive_params = runtime_caches
        .iter()
        .enumerate()
        .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
            lambda: [
                theta_star[adaptive_lambda_start + cache_idx * 3].exp(),
                theta_star[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
                theta_star[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
            ],
            epsilon: eps_star,
        })
        .collect::<Vec<_>>();
    let mut fixed_total = Array2::<f64>::zeros((
        baseline.design.design.ncols(),
        baseline.design.design.ncols(),
    ));
    for (idx, penalty) in retained_penalties.iter().enumerate() {
        fixed_total.scaled_add(rho_star[idx].exp(), penalty);
    }
    let final_family =
        base_family.with_adaptive_params(adaptive_params.clone(), Arc::new(fixed_total.clone()));
    let final_blockspec = ParameterBlockSpec {
        name: "eta".to_string(),
        design: baseline.design.design.clone(),
        offset: offset.to_owned(),
        penalties: vec![],
        nullspace_dims: vec![],
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: Some(baseline.fit.beta.clone()),
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    let final_fit = fit_custom_family(
        &final_family,
        &[final_blockspec],
        &BlockwiseFitOptions {
            inner_max_cycles: options.max_iter,
            inner_tol: options.tol,
            outer_max_iter: 1,
            outer_tol: options.tol,
            compute_covariance: true,
            ..BlockwiseFitOptions::default()
        },
    )
    .map_err(EstimationError::CustomFamily)?;
    let beta = final_fit.block_states[0].beta.clone();
    let final_eval = final_family
        .exact_evaluation(&beta)
        .map_err(EstimationError::InvalidInput)?;
    let penalized_hessian = final_eval
        .totalobjectivehessian(&final_family.design)
        .map_err(EstimationError::InvalidInput)?;
    let beta_covariance = final_fit.covariance_conditional.clone();
    let beta_standard_errors = beta_covariance
        .as_ref()
        .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));

    let mut full_lambdas = baseline.fit.lambdas.clone();
    for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
        full_lambdas[global_idx] = rho_star[idx].exp();
    }
    for (cache_idx, cache) in runtime_caches.iter().enumerate() {
        full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
        full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
        full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
    }

    let deviance = if family.is_gaussian_identity() {
        y.iter()
            .zip(final_eval.obs.mu.iter())
            .zip(weights.iter())
            .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
            .sum()
    } else {
        -2.0 * final_eval.obs.log_likelihood
    };
    let mut local_penalty_blocks =
        Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
    for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
        if adaptive_penalty_indices.contains(&global_idx) {
            let cache = runtime_caches
                .iter()
                .find(|cache| {
                    cache.mass_penalty_global_idx == global_idx
                        || cache.tension_penalty_global_idx == global_idx
                        || cache.stiffness_penalty_global_idx == global_idx
                })
                .ok_or_else(|| {
                    EstimationError::InvalidInput(format!(
                        "missing runtime cache for adaptive penalty index {global_idx}"
                    ))
                })?;
            let cache_idx = runtime_caches
                .iter()
                .position(|c| {
                    c.mass_penalty_global_idx == global_idx
                        || c.tension_penalty_global_idx == global_idx
                        || c.stiffness_penalty_global_idx == global_idx
                })
                .ok_or_else(|| {
                    EstimationError::InvalidInput(format!(
                        "missing adaptive cache position for penalty index {global_idx}"
                    ))
                })?;
            let state = &final_eval.adaptive_states[cache_idx];
            let local = if cache.mass_penalty_global_idx == global_idx {
                scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
                    .mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
            } else if cache.tension_penalty_global_idx == global_idx {
                grouped_operatorhessian(
                    &cache.d1,
                    cache.dimension,
                    &state.gradient.betahessian_blocks(),
                )?
                .mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
            } else {
                grouped_operatorhessian(
                    &cache.d2,
                    cache.dimension * cache.dimension,
                    &state.curvature.betahessian_blocks(),
                )?
                .mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
            };
            // Wrap the pre-scaled global penalty matrix as PenaltySpec::Dense.
            local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
                baseline.design.design.ncols(),
                cache.coeff_global_range.clone(),
                &local,
            )));
        } else {
            local_penalty_blocks.push(PenaltySpec::Dense(
                bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
            ));
        }
    }
    let (edf_by_block, edf_total) = if let Some(cov) = beta_covariance.as_ref() {
        exact_bounded_edf(
            &local_penalty_blocks,
            &Array1::from_elem(local_penalty_blocks.len(), 1.0),
            cov,
        )?
    } else {
        (vec![0.0; local_penalty_blocks.len()], 0.0)
    };
    let stable_penalty_term =
        2.0 * final_eval.adaptive_penalty_value + beta.dot(&fixed_total.dot(&beta));
    let standard_deviation = if family.is_gaussian_identity() {
        let denom = (y.len() as f64 - edf_total).max(1.0);
        (deviance / denom).sqrt()
    } else {
        1.0
    };
    let maps = compute_spatial_adaptiveweights_for_beta(
        &beta,
        runtime_caches,
        eps_star[0],
        eps_star[1],
        eps_star[2],
        adaptive_opts.weight_floor,
        adaptive_opts.weight_ceiling,
        // Working-Laplace conditional covariance Sigma_beta = H^{-1} from the
        // final exact-family solve, reused here as the posterior-SNR variance
        // source (no second factorization is formed).
        beta_covariance.as_ref(),
    )?
    .into_iter()
    .zip(runtime_caches.iter())
    .map(|(w, cache)| AdaptiveSpatialMap {
        termname: cache.termname.clone(),
        feature_cols: cache.feature_cols.clone(),
        collocation_points: cache.collocation_points.clone(),
        inv_magweight: w.inv_magweight,
        invgradweight: w.invgradweight,
        inv_lapweight: w.inv_lapweight,
    })
    .collect::<Vec<_>>();
    let fitted_link = if family.is_latent_cloglog() {
        FittedLinkState::LatentCLogLog {
            state: latent_cloglog_state
                .expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
        }
    } else if family.is_binomial_mixture() {
        mixture_link_state
            .clone()
            .map(|state| FittedLinkState::Mixture {
                state,
                covariance: None,
            })
            .unwrap_or(FittedLinkState::Standard(None))
    } else if family.is_binomial_sas() {
        sas_link_state
            .map(|state| FittedLinkState::Sas {
                state,
                covariance: None,
            })
            .unwrap_or(FittedLinkState::Standard(None))
    } else if family.is_binomial_beta_logistic() {
        sas_link_state
            .map(|state| FittedLinkState::BetaLogistic {
                state,
                covariance: None,
            })
            .unwrap_or(FittedLinkState::Standard(None))
    } else {
        FittedLinkState::Standard(None)
    };
    let max_abs_eta = final_eval
        .obs
        .eta
        .iter()
        .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
    let fitted = FittedTermCollection {
        fit: {
            let log_lambdas = full_lambdas.mapv(|v| v.max(1e-300).ln());
            let inf = FitInference {
                edf_by_block,
                edf_total,
                smoothing_correction: None,
                // Boundary adapter: wrap the raw `Array2<f64>` Hessian as
                // `UnscaledPrecision` for the newtype storage.
                penalized_hessian: penalized_hessian.clone().into(),
                working_weights: final_eval.obs.fisherweight.clone(),
                working_response: {
                    let mut out = final_eval.obs.eta.clone();
                    for i in 0..out.len() {
                        let wi = final_eval.obs.fisherweight[i].max(1e-12);
                        out[i] += final_eval.obs.score[i] / wi;
                    }
                    out
                },
                reparam_qs: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                beta_covariance: beta_covariance
                    .clone()
                    .map(crate::inference::dispersion_cov::PhiScaledCovariance::from),
                beta_standard_errors,
                beta_covariance_corrected: None,
                beta_standard_errors_corrected: None,
                beta_covariance_frequentist: None,
                coefficient_influence: None,
                bias_correction_beta: None,
            };
            let geometry = Some(crate::estimate::FitGeometry {
                penalized_hessian: penalized_hessian.into(),
                working_weights: inf.working_weights.clone(),
                working_response: inf.working_response.clone(),
            });
            let covariance_conditional = beta_covariance;
            let pirls_status_val = if final_fit.outer_converged {
                crate::pirls::PirlsStatus::Converged
            } else {
                crate::pirls::PirlsStatus::StalledAtValidMinimum
            };
            UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
                blocks: vec![crate::estimate::FittedBlock {
                    beta: beta.clone(),
                    role: crate::estimate::BlockRole::Mean,
                    edf: edf_total,
                    lambdas: full_lambdas.clone(),
                }],
                log_lambdas,
                lambdas: full_lambdas,
                likelihood_scale: family.default_scale_metadata(),
                likelihood_family: Some(family),
                log_likelihood_normalization:
                    crate::types::LogLikelihoodNormalization::UserProvided,
                log_likelihood: final_eval.obs.log_likelihood,
                deviance,
                reml_score: final_fit.penalized_objective,
                stable_penalty_term,
                penalized_objective: final_fit.penalized_objective,
                outer_iterations,
                outer_converged: final_fit.outer_converged,
                outer_gradient_norm: outer_grad_norm,
                standard_deviation,
                covariance_conditional,
                covariance_corrected: None,
                inference: Some(inf),
                fitted_link,
                geometry,
                block_states: Vec::new(),
                pirls_status: pirls_status_val,
                max_abs_eta,
                constraint_kkt: None,
                artifacts: crate::estimate::FitArtifacts {
                    pirls: None,
                    ..Default::default()
                },
                inner_cycles: 0,
            })?
        },
        design: baseline.design,
        adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
            epsilon_0: eps_star[0],
            epsilon_g: eps_star[1],
            epsilon_c: eps_star[2],
            epsilon_outer_iterations: outer_iterations,
            mm_iterations: 0,
            converged: final_fit.outer_converged,
            maps,
        }),
    };
    enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
    Ok(fitted)
}

fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
    FitOptions {
        latent_cloglog: options.latent_cloglog,
        mixture_link: options.mixture_link.clone(),
        optimize_mixture: options.optimize_mixture,
        sas_link: options.sas_link,
        optimize_sas: options.optimize_sas,
        compute_inference: options.compute_inference,
        max_iter: options.max_iter,
        tol: options.tol,
        nullspace_dims: design.nullspace_dims.clone(),
        linear_constraints: design.linear_constraints.clone(),
        firth_bias_reduction: options.firth_bias_reduction,
        adaptive_regularization: None,
        penalty_shrinkage_floor: options.penalty_shrinkage_floor,
        // Propagate user-supplied rho_prior so the baseline/refit and the
        // joint optimizer minimize the same REML objective.
        rho_prior: options.rho_prior.clone(),
        kronecker_penalty_system: design.kronecker_penalty_system(),
        kronecker_factored: design
            .smooth
            .terms
            .iter()
            .find_map(|t| t.kronecker_factored.clone()),
    }
}

#[derive(Clone)]
struct BoundedLinearTermMeta {
    col_idx: usize,
    min: f64,
    max: f64,
    prior: BoundedCoefficientPriorSpec,
}

/// β-dependent effective Jacobian for the bounded-linear fit block.
///
/// Each bounded coefficient enters the linear predictor non-linearly, as
/// `β = min + width·σ(θ)`, and is supplied to the solver through the family
/// adapter's offset rather than the linear design. To keep that contribution
/// out of the *linear* design the fit places a deliberately **zeroed**
/// placeholder column for every bounded term in the block design
/// (see `fit_bounded_term_collection_with_design`). The pre-fit
/// identifiability audit, however, assesses block rank by reading each block's
/// effective Jacobian — and a zeroed column reads as a structural rank
/// deficiency, so without this callback the audit refuses *every* bounded
/// model before fitting begins.
///
/// This callback reports the model's true Jacobian column for each bounded
/// term, `∂η_i/∂θ = (dβ/dθ)·x_i`, so the audit inspects the same geometry the
/// solver actually fits. Because `dβ/dθ = width·σ(θ)(1−σ(θ))` is strictly
/// positive for finite θ and `width > 0`, a bounded column is rank-deficient
/// in the audit exactly when its underlying covariate is genuinely collinear
/// with the rest of the design — never merely because the placeholder was
/// zeroed. The callback is consumed only by the identifiability audit /
/// canonicalisation; the inner PIRLS solve drives η through the
/// [`BoundedLinearFamily`] adapter, so reporting the non-zeroed Jacobian here
/// does not double-count the bounded contribution.
struct BoundedEffectiveJacobian {
    design: Array2<f64>,
    bounded_terms: Vec<BoundedLinearTermMeta>,
}

impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
    fn effective_jacobian_at(
        &self,
        state: &FamilyLinearizationState<'_>,
    ) -> Result<Array2<f64>, String> {
        let p = self.design.ncols();
        if !state.beta.is_empty() {
            if state.beta.len() != p {
                return Err(format!(
                    "BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
                     ncols {p}",
                    state.beta.len(),
                ));
            }
            if state.beta.iter().any(|v| v.is_nan()) {
                return Err(
                    "BoundedEffectiveJacobian::effective_jacobian_at: beta contains NaN"
                        .to_string(),
                );
            }
        }
        let mut jac = self.design.clone();
        for term in &self.bounded_terms {
            let theta = if state.beta.is_empty() {
                0.0
            } else {
                state.beta[term.col_idx]
            };
            let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
            jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
        }
        Ok(jac)
    }
}

#[derive(Clone)]
struct BoundedLinearFamily {
    family: LikelihoodSpec,
    latent_cloglog_state: Option<LatentCLogLogState>,
    mixture_link_state: Option<MixtureLinkState>,
    sas_link_state: Option<SasLinkState>,
    y: Array1<f64>,
    weights: Array1<f64>,
    design: Array2<f64>,
    designzeroed: Array2<f64>,
    offset: Array1<f64>,
    bounded_terms: Vec<BoundedLinearTermMeta>,
}

#[derive(Clone)]
struct StandardFamilyObservationState {
    eta: Array1<f64>,
    mu: Array1<f64>,
    score: Array1<f64>,
    fisherweight: Array1<f64>,
    neghessian_eta: Array1<f64>,
    neghessian_eta_derivative: Array1<f64>,
    log_likelihood: f64,
}

fn bounded_logit(z: f64) -> f64 {
    let zc = z.clamp(1e-12, 1.0 - 1e-12);
    (zc / (1.0 - zc)).ln()
}

fn stable_sigmoid(theta: f64) -> f64 {
    if theta >= 0.0 {
        let exp_neg = (-theta).exp();
        1.0 / (1.0 + exp_neg)
    } else {
        let exp_pos = theta.exp();
        exp_pos / (1.0 + exp_pos)
    }
}

fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
    let z = stable_sigmoid(theta);
    let width = max - min;
    let beta = min + width * z;
    let db_dtheta = width * z * (1.0 - z);
    (beta, z, db_dtheta)
}

fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
    let z = stable_sigmoid(theta);
    let width = max - min;
    let s = z * (1.0 - z);
    let beta = min + width * z;
    let db_dtheta = width * s;
    let d2b_dtheta2 = width * s * (1.0 - 2.0 * z);
    let d3b_dtheta3 = width * s * (1.0 - 6.0 * z + 6.0 * z * z);
    (beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
}

fn bounded_prior_terms(theta: f64, prior: &BoundedCoefficientPriorSpec) -> (f64, f64, f64, f64) {
    let (a, b) = match prior {
        // `None` means constrained MLE with no extra prior term on the bounded coefficient.
        BoundedCoefficientPriorSpec::None => return (0.0, 0.0, 0.0, 0.0),
        // Uniform on the normalized user-scale coefficient z in (0, 1). In latent space this is
        // exactly the Jacobian term for the logistic transform, up to an additive width constant.
        BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
        BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
    };
    let z = stable_sigmoid(theta).clamp(1e-12, 1.0 - 1e-12);
    let logp = a * z.ln() + b * (1.0 - z).ln();
    let grad = a - (a + b) * z;
    let neghess = (a + b) * z * (1.0 - z);
    let neghess_derivative = (a + b) * z * (1.0 - z) * (1.0 - 2.0 * z);
    (logp, grad, neghess, neghess_derivative)
}

fn evaluate_standard_familyobservations(
    family: LikelihoodSpec,
    latent_cloglog_state: Option<&LatentCLogLogState>,
    mixture_link_state: Option<&MixtureLinkState>,
    sas_link_state: Option<&SasLinkState>,
    y: &Array1<f64>,
    weights: &Array1<f64>,
    eta: &Array1<f64>,
) -> Result<StandardFamilyObservationState, EstimationError> {
    const PROB_EPS: f64 = 1e-10;
    const MU_DERIV_EPS: f64 = 1e-12;
    let n = y.len();
    if weights.len() != n || eta.len() != n {
        crate::bail_invalid_estim!("bounded family observation size mismatch");
    }

    let mut mu = Array1::<f64>::zeros(n);
    let mut score = Array1::<f64>::zeros(n);
    let mut fisherweight = Array1::<f64>::zeros(n);
    let mut neghessian_eta = Array1::<f64>::zeros(n);
    let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
    let mut log_likelihood = 0.0;

    for i in 0..n {
        let w = weights[i].max(0.0);
        let yi = y[i];
        let eta_i = eta[i];
        match (&family.response, &family.link) {
            (ResponseFamily::Gaussian, _) => {
                let resid = yi - eta_i;
                mu[i] = eta_i;
                score[i] = w * resid;
                fisherweight[i] = w.max(MU_DERIV_EPS);
                neghessian_eta[i] = w;
                neghessian_eta_derivative[i] = 0.0;
                log_likelihood += -0.5 * w * resid * resid;
            }
            (ResponseFamily::Binomial, InverseLink::Standard(StandardLink::Logit)) => {
                let jet = logit_inverse_link_jet5(eta_i);
                mu[i] = jet.mu;
                score[i] = w * (yi - jet.mu);
                fisherweight[i] = jet.d1.max(MU_DERIV_EPS);
                neghessian_eta[i] = jet.d1;
                neghessian_eta_derivative[i] = jet.d2;
                let logmu = -crate::linalg::utils::stable_softplus(-eta_i);
                let log_one_minusmu = -crate::linalg::utils::stable_softplus(eta_i);
                log_likelihood += w * (yi * logmu + (1.0 - yi) * log_one_minusmu);
            }
            (ResponseFamily::Binomial, _) => {
                let inverse_link = if let Some(state) = latent_cloglog_state {
                    Some(InverseLink::LatentCLogLog(*state))
                } else if let Some(state) = mixture_link_state {
                    Some(InverseLink::Mixture(state.clone()))
                } else {
                    sas_link_state.map(|state| {
                        if family.is_binomial_beta_logistic() {
                            InverseLink::BetaLogistic(*state)
                        } else {
                            InverseLink::Sas(*state)
                        }
                    })
                };
                let strategy_spec = LikelihoodSpec {
                    response: family.response.clone(),
                    link: inverse_link.clone().unwrap_or_else(|| family.link.clone()),
                };
                let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
                let mu_i_raw = jet.mu;
                let dmu_deta_raw = jet.d1;
                let mu_i: f64 = mu_i_raw.clamp(PROB_EPS, 1.0 - PROB_EPS);
                let dmu_deta = dmu_deta_raw.max(MU_DERIV_EPS);
                let d2mu_deta2 = jet.d2;
                let d3mu_deta3 = jet.d3;
                let var = (mu_i * (1.0 - mu_i)).max(PROB_EPS);
                let lmu = (yi - mu_i) / var;
                let lmumu = -(yi / (mu_i * mu_i)) - ((1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i)));
                let lmumumu = 2.0 * yi / (mu_i * mu_i * mu_i)
                    - 2.0 * (1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i) * (1.0 - mu_i));
                mu[i] = mu_i;
                score[i] = w * lmu * dmu_deta;
                fisherweight[i] = (w * dmu_deta * dmu_deta / var).max(MU_DERIV_EPS);
                neghessian_eta[i] = -w * (lmumu * dmu_deta * dmu_deta + lmu * d2mu_deta2);
                neghessian_eta_derivative[i] = -w
                    * (lmumumu * dmu_deta * dmu_deta * dmu_deta
                        + 3.0 * lmumu * dmu_deta * d2mu_deta2
                        + lmu * d3mu_deta3);
                log_likelihood += w * (yi * mu_i.ln() + (1.0 - yi) * (1.0 - mu_i).ln());
            }
            (ResponseFamily::Poisson, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for PoissonLog fits"
                );
            }
            (ResponseFamily::Tweedie { .. }, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for Tweedie fits"
                );
            }
            (ResponseFamily::NegativeBinomial { .. }, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for NegativeBinomial fits"
                );
            }
            (ResponseFamily::Beta { .. }, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for BetaLogit fits"
                );
            }
            (ResponseFamily::Gamma, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for GammaLog fits"
                );
            }
            (ResponseFamily::RoystonParmar, _) => {
                crate::bail_invalid_estim!(
                    "bounded linear terms are not supported for survival model fits"
                );
            }
        }
    }

    Ok(StandardFamilyObservationState {
        eta: eta.clone(),
        mu,
        score,
        fisherweight,
        neghessian_eta,
        neghessian_eta_derivative,
        log_likelihood,
    })
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SpatialAdaptiveHyperKind {
    LogLambdaMagnitude,
    LogLambdaGradient,
    LogLambdaCurvature,
    LogEpsilonMagnitude,
    LogEpsilonGradient,
    LogEpsilonCurvature,
}

impl SpatialAdaptiveHyperKind {
    fn component_index(self) -> usize {
        match self {
            SpatialAdaptiveHyperKind::LogLambdaMagnitude
            | SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
            SpatialAdaptiveHyperKind::LogLambdaGradient
            | SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
            SpatialAdaptiveHyperKind::LogLambdaCurvature
            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
        }
    }

    fn is_log_lambda(self) -> bool {
        matches!(
            self,
            SpatialAdaptiveHyperKind::LogLambdaMagnitude
                | SpatialAdaptiveHyperKind::LogLambdaGradient
                | SpatialAdaptiveHyperKind::LogLambdaCurvature
        )
    }

    fn is_log_epsilon(self) -> bool {
        matches!(
            self,
            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
                | SpatialAdaptiveHyperKind::LogEpsilonGradient
                | SpatialAdaptiveHyperKind::LogEpsilonCurvature
        )
    }
}

#[derive(Clone, Copy, Debug)]
struct SpatialAdaptiveHyperSpec {
    cache_index: usize,
    kind: SpatialAdaptiveHyperKind,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SpatialAdaptiveExplicitSecondOrderKind {
    StructuralZero,
    LocalAlphaAlpha,
    LocalAlphaEta,
    SharedEtaEta,
}

/// Penalty family selected within one adaptive smooth cache. The component index
/// (0/1/2) used throughout the runtime caches maps onto these three operators:
/// the scalar magnitude operator `d0`, the grouped gradient operator `d1`, and
/// the grouped curvature operator `d2`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum AdaptiveComponent {
    Magnitude,
    Gradient,
    Curvature,
}

impl AdaptiveComponent {
    fn from_index(index: usize) -> Result<Self, String> {
        match index {
            0 => Ok(AdaptiveComponent::Magnitude),
            1 => Ok(AdaptiveComponent::Gradient),
            2 => Ok(AdaptiveComponent::Curvature),
            other => Err(SmoothError::invalid_index(format!(
                "invalid adaptive component index {}",
                other
            ))
            .into()),
        }
    }
}

/// Which hyper-derivative of the adaptive penalty's local pieces to assemble.
/// Each variant selects one accessor triple (objective scalar, beta-mixed
/// gradient, beta hessian) on the per-component exact state; the operator
/// embedding around those accessors is identical across variants.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HyperDerivativeKind {
    /// First derivative in `log lambda` (rho): the bare penalty pieces.
    Rho,
    /// First derivative in `log epsilon`.
    LogEpsilonFirst,
    /// Second derivative in `log epsilon`.
    LogEpsilonSecond,
}

/// Which directional-drift hyper-derivative of the adaptive penalty Hessian to
/// assemble: the bare rho drift, or the shared-`log epsilon` drift. Both share
/// the per-component direction projection, operator embedding, and global
/// embedding; only the directional state accessor differs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HyperDriftKind {
    Rho,
    LogEpsilon,
}

impl SpatialAdaptiveHyperSpec {
    fn component_index(self) -> usize {
        self.kind.component_index()
    }

    fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
        if self.component_index() != other.component_index() {
            return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
        }
        match (
            self.kind.is_log_lambda(),
            other.kind.is_log_lambda(),
            self.kind.is_log_epsilon(),
            other.kind.is_log_epsilon(),
        ) {
            (true, true, false, false) if self.cache_index == other.cache_index => {
                SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
            }
            (true, false, false, true) | (false, true, true, false) => {
                SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
            }
            (false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
            _ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
        }
    }
}

#[derive(Clone, Debug)]
struct SpatialAdaptiveTermHyperParams {
    lambda: [f64; 3],
    epsilon: [f64; 3],
}

#[derive(Clone)]
struct SpatialAdaptiveExactEvaluation {
    obs: StandardFamilyObservationState,
    adaptive_states: Vec<SpatialPenaltyExactState>,
    adaptive_penalty_value: f64,
    adaptive_penaltygradient: Array1<f64>,
    adaptive_penaltyhessian: Array2<f64>,
    fixed_quadraticvalue: f64,
    fixed_quadraticgradient: Array1<f64>,
    fixed_quadratichessian: Array2<f64>,
}

#[derive(Clone)]
struct CachedSpatialAdaptiveExactEvaluation {
    beta: Array1<f64>,
    eval: Arc<SpatialAdaptiveExactEvaluation>,
}

impl SpatialAdaptiveExactEvaluation {
    fn total_penalty_value(&self) -> f64 {
        self.adaptive_penalty_value + self.fixed_quadraticvalue
    }

    fn total_penaltygradient(&self) -> Array1<f64> {
        &self.adaptive_penaltygradient + &self.fixed_quadraticgradient
    }

    fn total_penaltyhessian(&self) -> Array2<f64> {
        &self.adaptive_penaltyhessian + &self.fixed_quadratichessian
    }

    fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
        let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
        out += &self.total_penaltyhessian();
        Ok(out)
    }
}

#[derive(Clone)]
struct SpatialAdaptiveExactFamily {
    family: LikelihoodSpec,
    latent_cloglog_state: Option<LatentCLogLogState>,
    mixture_link_state: Option<MixtureLinkState>,
    sas_link_state: Option<SasLinkState>,
    y: Arc<Array1<f64>>,
    weights: Arc<Array1<f64>>,
    design: Arc<Array2<f64>>,
    offset: Arc<Array1<f64>>,
    linear_constraints: Option<LinearInequalityConstraints>,
    runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
    adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
    fixed_quadratichessian: Arc<Array2<f64>>,
    hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
    exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
}

impl SpatialAdaptiveExactFamily {
    fn with_adaptive_params(
        &self,
        adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
        fixed_quadratichessian: Arc<Array2<f64>>,
    ) -> Self {
        Self {
            family: self.family.clone(),
            latent_cloglog_state: self.latent_cloglog_state,
            mixture_link_state: self.mixture_link_state.clone(),
            sas_link_state: self.sas_link_state,
            y: self.y.clone(),
            weights: self.weights.clone(),
            design: self.design.clone(),
            offset: self.offset.clone(),
            linear_constraints: self.linear_constraints.clone(),
            runtime_caches: self.runtime_caches.clone(),
            adaptive_params,
            fixed_quadratichessian,
            hyperspecs: self.hyperspecs.clone(),
            exact_eval_cache: Arc::new(Mutex::new(None)),
        }
    }

    fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
        crate::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
    }

    fn fixed_quadratic_terms(&self, beta: &Array1<f64>) -> (f64, Array1<f64>) {
        let grad = self.fixed_quadratichessian.dot(beta);
        let value = 0.5 * beta.dot(&grad);
        (value, grad)
    }

    fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
        let mut penalty_value = 0.0;
        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
                format!(
                    "missing adaptive parameter block for cache {}",
                    cache.termname
                )
            })?;
            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
            let state =
                SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
                    .map_err(|e| e.to_string())?;
            penalty_value += params.lambda[0] * state.magnitude.penalty_value();
            penalty_value += params.lambda[1] * state.gradient.penalty_value();
            penalty_value += params.lambda[2] * state.curvature.penalty_value();
        }
        Ok(penalty_value)
    }

    fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
        let total_dim = self.design.ncols();
        (
            Array1::<f64>::zeros(total_dim),
            Array2::<f64>::zeros((total_dim, total_dim)),
        )
    }

    fn embed_local_hyper_parts(
        &self,
        coeff_range: &Range<usize>,
        local_grad: &Array1<f64>,
        local_hess: &Array2<f64>,
    ) -> (Array1<f64>, Array2<f64>) {
        let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
        beta_mixed
            .slice_mut(s![coeff_range.clone()])
            .assign(local_grad);
        betahessian
            .slice_mut(s![coeff_range.clone(), coeff_range.clone()])
            .assign(local_hess);
        (beta_mixed, betahessian)
    }

    fn embed_local_hyper_hessian(
        &self,
        coeff_range: &Range<usize>,
        local_hess: &Array2<f64>,
    ) -> Array2<f64> {
        let total_dim = self.design.ncols();
        let mut out = Array2::<f64>::zeros((total_dim, total_dim));
        out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
            .assign(local_hess);
        out
    }

    /// Unified per-block hyper-derivative assembly. Owns the shared cache /
    /// hyperparameter / exact-state lookup, the component -> operator selection
    /// (scalar magnitude `d0`, grouped gradient `d1`, grouped curvature `d2`),
    /// and the global embedding via [`Self::embed_local_hyper_parts`]. The only
    /// piece that varies with `derivative` is the per-component accessor triple
    /// (objective scalar, beta-mixed gradient, beta hessian) read off the exact
    /// state. Returns `(objective, beta_mixed, betahessian)`, each already
    /// scaled by the component's penalty weight `lambda`.
    fn adaptive_block_eval(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        cache_idx: usize,
        component: AdaptiveComponent,
        derivative: HyperDerivativeKind,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        let cache = self
            .runtime_caches
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
        let params = self
            .adaptive_params
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
        let state = eval
            .adaptive_states
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;

        let (objective_local, beta_mixed_local, betahessian_local) = match component {
            AdaptiveComponent::Magnitude => {
                let lambda = params.lambda[0];
                let mag = &state.magnitude;
                let (objective, gradient_coeff, hessian_diag) = match derivative {
                    HyperDerivativeKind::Rho => (
                        mag.penalty_value(),
                        mag.betagradient_coeff(),
                        mag.betahessian_diag(),
                    ),
                    HyperDerivativeKind::LogEpsilonFirst => (
                        mag.log_epsilon_gradient_terms().sum(),
                        mag.log_epsilon_betagradient_coeff(),
                        mag.log_epsilon_betahessian_diag(),
                    ),
                    HyperDerivativeKind::LogEpsilonSecond => (
                        mag.log_epsilon_hessian_terms().sum(),
                        mag.log_epsilon_beta_mixed_second_coeff(),
                        mag.log_epsilon_betahessian_second_diag(),
                    ),
                };
                (
                    lambda * objective,
                    lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
                    lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
                )
            }
            AdaptiveComponent::Gradient => {
                let lambda = params.lambda[1];
                let grad = &state.gradient;
                let (objective, gradient_blocks, hessian_blocks) = match derivative {
                    HyperDerivativeKind::Rho => (
                        grad.penalty_value(),
                        grad.betagradient_blocks(),
                        grad.betahessian_blocks(),
                    ),
                    HyperDerivativeKind::LogEpsilonFirst => (
                        grad.log_epsilon_gradient_terms().sum(),
                        grad.log_epsilon_betagradient_blocks(),
                        grad.log_epsilon_betahessian_blocks(),
                    ),
                    HyperDerivativeKind::LogEpsilonSecond => (
                        grad.log_epsilon_hessian_terms().sum(),
                        grad.log_epsilon_beta_mixed_second_blocks(),
                        grad.log_epsilon_betahessian_second_blocks(),
                    ),
                };
                (
                    lambda * objective,
                    lambda
                        * grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
                            .map_err(|e| e.to_string())?,
                    lambda
                        * grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
                            .map_err(|e| e.to_string())?,
                )
            }
            AdaptiveComponent::Curvature => {
                let lambda = params.lambda[2];
                let group = cache.dimension * cache.dimension;
                let curv = &state.curvature;
                let (objective, gradient_blocks, hessian_blocks) = match derivative {
                    HyperDerivativeKind::Rho => (
                        curv.penalty_value(),
                        curv.betagradient_blocks(),
                        curv.betahessian_blocks(),
                    ),
                    HyperDerivativeKind::LogEpsilonFirst => (
                        curv.log_epsilon_gradient_terms().sum(),
                        curv.log_epsilon_betagradient_blocks(),
                        curv.log_epsilon_betahessian_blocks(),
                    ),
                    HyperDerivativeKind::LogEpsilonSecond => (
                        curv.log_epsilon_hessian_terms().sum(),
                        curv.log_epsilon_beta_mixed_second_blocks(),
                        curv.log_epsilon_betahessian_second_blocks(),
                    ),
                };
                (
                    lambda * objective,
                    lambda
                        * grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
                            .map_err(|e| e.to_string())?,
                    lambda
                        * grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
                            .map_err(|e| e.to_string())?,
                )
            }
        };

        let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
            &cache.coeff_global_range,
            &beta_mixed_local,
            &betahessian_local,
        );
        Ok((objective_local, beta_mixed, betahessian))
    }

    fn adaptive_shared_log_epsilon_parts(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        component: usize,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        // Exact shared-log-epsilon first-order pieces:
        //
        //   J_{eta_p}         = sum_m lambda_{m,p} U_{m,p,eta},
        //   J_{beta,eta_p}    = sum_m lambda_{m,p} U_{m,p,beta eta},
        //   J_{beta,beta,eta} = sum_m lambda_{m,p} U_{m,p,beta beta eta}.
        self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
    }

    fn adaptive_shared_log_epsilon_second_parts(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        component: usize,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        // Exact shared-log-epsilon second-order pieces:
        //
        //   J_{eta_p,eta_p}            = sum_m lambda_{m,p} U_{m,p,eta eta},
        //   J_{beta,eta_p,eta_p}       = sum_m lambda_{m,p} U_{m,p,beta eta eta},
        //   J_{beta,beta,eta_p,eta_p}  = sum_m lambda_{m,p} U_{m,p,beta beta eta eta}.
        self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
    }

    /// Sum a per-block hyper-derivative across every adaptive term for one shared
    /// `log epsilon` coordinate (selected by `component`). The three log-epsilon
    /// coordinates are shared globally by penalty type, so each contributes the
    /// matching component's block from every cache.
    fn adaptive_shared_block_eval(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        component: usize,
        derivative: HyperDerivativeKind,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        let component = AdaptiveComponent::from_index(component)?;
        let (mut score, mut hessian) = self.zero_hyper_parts();
        let mut objective = 0.0;
        for cache_idx in 0..self.runtime_caches.len() {
            let (local_objective, local_score, local_hessian) =
                self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
            objective += local_objective;
            score += &local_score;
            hessian += &local_hessian;
        }
        Ok((objective, score, hessian))
    }

    fn adaptive_shared_log_epsilon_drift(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        component: usize,
        direction: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        // Exact shared-log-epsilon Hessian drift:
        //
        //   T_{eta_p}[u] = sum_m lambda_{m,p} D_beta(U_{m,p,beta beta eta})[u].
        let component = AdaptiveComponent::from_index(component)?;
        let total_dim = self.design.ncols();
        let mut total = Array2::<f64>::zeros((total_dim, total_dim));
        for cache_idx in 0..self.runtime_caches.len() {
            total += &self.adaptive_block_drift_eval(
                eval,
                cache_idx,
                component,
                HyperDriftKind::LogEpsilon,
                direction,
            )?;
        }
        Ok(total)
    }

    fn adaptive_explicit_second_order_parts(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        left: SpatialAdaptiveHyperSpec,
        right: SpatialAdaptiveHyperSpec,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        // Structural sparsity from the adaptive penalty algebra:
        //
        //   - alpha_{m,p} / alpha_{n,r} is nonzero only when (m,p) = (n,r),
        //   - alpha_{m,p} / eta_r is nonzero only when p = r,
        //   - eta_p / eta_r is nonzero only when p = r,
        //
        // with eta_p contributions summed over all adaptive terms m because the
        // three log-epsilon coordinates are shared globally by penalty type.
        match left.explicit_second_order_kind(right) {
            SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
                let (score, hessian) = self.zero_hyper_parts();
                Ok((0.0, score, hessian))
            }
            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
                eval,
                left.cache_index,
                AdaptiveComponent::from_index(left.component_index())?,
                HyperDerivativeKind::Rho,
            ),
            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
                let local_alpha = if left.kind.is_log_lambda() {
                    left
                } else {
                    right
                };
                self.adaptive_block_eval(
                    eval,
                    local_alpha.cache_index,
                    AdaptiveComponent::from_index(local_alpha.component_index())?,
                    HyperDerivativeKind::LogEpsilonFirst,
                )
            }
            SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
                self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
            }
        }
    }

    /// Unified per-block directional-drift assembly. Owns the shared cache /
    /// hyperparameter / exact-state lookup, the per-component direction
    /// projection through the collocation operators, the operator embedding, and
    /// the global embedding via [`Self::embed_local_hyper_hessian`]. The only
    /// piece that varies with `drift` is the directional state accessor:
    /// [`HyperDriftKind::Rho`] takes the bare directional Hessian drift, while
    /// [`HyperDriftKind::LogEpsilon`] takes its `log epsilon` derivative.
    fn adaptive_block_drift_eval(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        cache_idx: usize,
        component: AdaptiveComponent,
        drift: HyperDriftKind,
        direction: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        let cache = self
            .runtime_caches
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
        let params = self
            .adaptive_params
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
        let state = eval
            .adaptive_states
            .get(cache_idx)
            .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
        let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);

        let local_hessian = match component {
            AdaptiveComponent::Magnitude => {
                let d0_u = cache.d0.dot(&direction_local);
                let mag = &state.magnitude;
                let diag = match drift {
                    HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
                    HyperDriftKind::LogEpsilon => {
                        mag.log_epsilon_betahessian_directional_diag(&d0_u)
                    }
                };
                params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
            }
            AdaptiveComponent::Gradient => {
                let d1_u = cache.d1.dot(&direction_local);
                let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
                    .map_err(|e| e.to_string())?;
                let grad = &state.gradient;
                let blocks = match drift {
                    HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
                    HyperDriftKind::LogEpsilon => {
                        grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
                    }
                };
                params.lambda[1]
                    * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
                        .map_err(|e| e.to_string())?
            }
            AdaptiveComponent::Curvature => {
                let group = cache.dimension * cache.dimension;
                let d2_u = cache.d2.dot(&direction_local);
                let direction_blocks =
                    collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
                let curv = &state.curvature;
                let blocks = match drift {
                    HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
                    HyperDriftKind::LogEpsilon => {
                        curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
                    }
                };
                params.lambda[2]
                    * grouped_operatorhessian(&cache.d2, group, &blocks)
                        .map_err(|e| e.to_string())?
            }
        };

        Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
    }

    fn adaptive_hyper_parts(
        &self,
        eval: &SpatialAdaptiveExactEvaluation,
        hyper: SpatialAdaptiveHyperSpec,
    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
        match hyper.kind {
            // Per-term `log lambda` (rho) hyper-derivative: the bare penalty
            // pieces for this cache's selected component.
            SpatialAdaptiveHyperKind::LogLambdaMagnitude
            | SpatialAdaptiveHyperKind::LogLambdaGradient
            | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
                eval,
                hyper.cache_index,
                AdaptiveComponent::from_index(hyper.component_index())?,
                HyperDerivativeKind::Rho,
            ),
            // Shared `log epsilon` hyper-derivative: summed across all terms.
            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
            | SpatialAdaptiveHyperKind::LogEpsilonGradient
            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
                self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
            }
        }
    }

    fn exact_evaluation_uncached(
        &self,
        beta: &Array1<f64>,
    ) -> Result<SpatialAdaptiveExactEvaluation, String> {
        let eta = self.total_eta(beta);
        let obs = evaluate_standard_familyobservations(
            self.family.clone(),
            self.latent_cloglog_state.as_ref(),
            self.mixture_link_state.as_ref(),
            self.sas_link_state.as_ref(),
            &self.y,
            &self.weights,
            &eta,
        )
        .map_err(|e| e.to_string())?;
        let p = beta.len();
        let mut penalty_value = 0.0;
        let mut penaltygradient = Array1::<f64>::zeros(p);
        let mut penaltyhessian = Array2::<f64>::zeros((p, p));
        let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());

        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
                format!(
                    "missing adaptive parameter block for cache {}",
                    cache.termname
                )
            })?;
            let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
            let state =
                SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
                    .map_err(|e| e.to_string())?;

            let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
            let gg = grouped_operatorgradient(
                &cache.d1,
                cache.dimension,
                &state.gradient.betagradient_blocks(),
            )
            .map_err(|e| e.to_string())?;
            let gc = grouped_operatorgradient(
                &cache.d2,
                cache.dimension * cache.dimension,
                &state.curvature.betagradient_blocks(),
            )
            .map_err(|e| e.to_string())?;
            let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
            let hg = grouped_operatorhessian(
                &cache.d1,
                cache.dimension,
                &state.gradient.betahessian_blocks(),
            )
            .map_err(|e| e.to_string())?;
            let hc = grouped_operatorhessian(
                &cache.d2,
                cache.dimension * cache.dimension,
                &state.curvature.betahessian_blocks(),
            )
            .map_err(|e| e.to_string())?;

            let lambda0 = params.lambda[0];
            let lambdag = params.lambda[1];
            let lambdac = params.lambda[2];

            penalty_value += lambda0 * state.magnitude.penalty_value();
            penalty_value += lambdag * state.gradient.penalty_value();
            penalty_value += lambdac * state.curvature.penalty_value();

            let range = cache.coeff_global_range.clone();
            {
                let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
                grad_local += &(g0.mapv(|v| lambda0 * v));
                grad_local += &(gg.mapv(|v| lambdag * v));
                grad_local += &(gc.mapv(|v| lambdac * v));
            }
            {
                let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
                h_local += &h0.mapv(|v| lambda0 * v);
                h_local += &hg.mapv(|v| lambdag * v);
                h_local += &hc.mapv(|v| lambdac * v);
            }

            adaptive_states.push(state);
        }

        let (fixed_quadraticvalue, fixed_quadraticgradient) = self.fixed_quadratic_terms(beta);
        Ok(SpatialAdaptiveExactEvaluation {
            obs,
            adaptive_states,
            adaptive_penalty_value: penalty_value,
            adaptive_penaltygradient: penaltygradient,
            adaptive_penaltyhessian: penaltyhessian,
            fixed_quadraticvalue,
            fixed_quadraticgradient,
            fixed_quadratichessian: self.fixed_quadratichessian.as_ref().clone(),
        })
    }

    fn exact_evaluation(
        &self,
        beta: &Array1<f64>,
    ) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
        {
            let cache = self
                .exact_eval_cache
                .lock()
                .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
            if let Some(cached) = cache.as_ref()
                && cached.beta.len() == beta.len()
                && cached
                    .beta
                    .iter()
                    .zip(beta.iter())
                    .all(|(&left, &right)| left == right)
            {
                return Ok(Arc::clone(&cached.eval));
            }
        }

        let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
        let mut cache = self
            .exact_eval_cache
            .lock()
            .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
        *cache = Some(CachedSpatialAdaptiveExactEvaluation {
            beta: beta.clone(),
            eval: Arc::clone(&eval),
        });
        Ok(eval)
    }

    fn exacthessian_directional_derivative_from_evaluation(
        &self,
        beta: &Array1<f64>,
        eval: &SpatialAdaptiveExactEvaluation,
        direction: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        assert_eq!(
            beta.len(),
            direction.len(),
            "beta/direction length mismatch",
        );
        let d_eta = crate::faer_ndarray::fast_av(self.design.as_ref(), direction);
        let mut total = xt_diag_x_dense(
            self.design.view(),
            (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
        )?;
        for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
            let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
                format!(
                    "missing adaptive parameter block for cache {}",
                    cache.termname
                )
            })?;
            let state = eval
                .adaptive_states
                .get(cache_idx)
                .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
            let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
            let d0_u = cache.d0.dot(&direction_local);
            let d1_u = cache.d1.dot(&direction_local);
            let d2_u = cache.d2.dot(&direction_local);
            let h0 =
                scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
                    .mapv(|v| params.lambda[0] * v);
            let hg = grouped_operatorhessian(
                &cache.d1,
                cache.dimension,
                &state.gradient.directionalhessian_blocks(
                    &collocationgradient_blocks(&d1_u, cache.dimension)
                        .map_err(|e| e.to_string())?,
                ),
            )
            .map_err(|e| e.to_string())?
            .mapv(|v| params.lambda[1] * v);
            let hc = grouped_operatorhessian(
                &cache.d2,
                cache.dimension * cache.dimension,
                &state.curvature.directionalhessian_blocks(
                    &collocationhessian_blocks(&d2_u, cache.dimension)
                        .map_err(|e| e.to_string())?,
                ),
            )
            .map_err(|e| e.to_string())?
            .mapv(|v| params.lambda[2] * v);
            let range = cache.coeff_global_range.clone();
            let mut local = total.slice_mut(s![range.clone(), range]);
            local += &h0;
            local += &hg;
            local += &hc;
        }
        Ok(total)
    }
}

impl CustomFamily for SpatialAdaptiveExactFamily {
    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
        let eval = self.exact_evaluation(beta)?;
        let mut gradient = fast_atv(&self.design, &eval.obs.score);
        gradient -= &eval.total_penaltygradient();
        let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
        hessian += &eval.total_penaltyhessian();
        Ok(FamilyEvaluation {
            log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
            blockworking_sets: vec![BlockWorkingSet::ExactNewton {
                gradient,
                hessian: SymmetricMatrix::Dense(hessian),
            }],
        })
    }

    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
        let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
        let beta = &state.beta;
        let obs = evaluate_standard_familyobservations(
            self.family.clone(),
            self.latent_cloglog_state.as_ref(),
            self.mixture_link_state.as_ref(),
            self.sas_link_state.as_ref(),
            &self.y,
            &self.weights,
            &state.eta,
        )
        .map_err(|e| e.to_string())?;
        let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
        let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta);
        Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
    }

    fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
        ExactNewtonOuterObjective::StrictPseudoLaplace
    }

    fn exact_newton_joint_hessian(
        &self,
        block_states: &[ParameterBlockState],
    ) -> Result<Option<Array2<f64>>, String> {
        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
        let eval = self.exact_evaluation(beta)?;
        Ok(Some(eval.totalobjectivehessian(&self.design)?))
    }

    fn exact_newton_hessian_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        block_idx: usize,
        d_beta: &Array1<f64>,
    ) -> Result<Option<Array2<f64>>, String> {
        expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
        self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
    }

    fn exact_newton_joint_hessian_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        d_beta_flat: &Array1<f64>,
    ) -> Result<Option<Array2<f64>>, String> {
        let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
        if d_beta_flat.len() != beta.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "spatial adaptive exact family direction length mismatch: got {}, expected {}",
                d_beta_flat.len(),
                beta.len()
            ))
            .into());
        }
        let eval = self.exact_evaluation(beta)?;
        Ok(Some(
            self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
        ))
    }

    fn block_linear_constraints(
        &self,
        block_states: &[ParameterBlockState],
        block_idx: usize,
        block_spec: &ParameterBlockSpec,
    ) -> Result<Option<LinearInequalityConstraints>, String> {
        assert!(!block_states.is_empty(), "block_states must be non-empty");
        assert!(
            !block_spec.name.is_empty(),
            "block spec name must be non-empty",
        );
        expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
        Ok(self.linear_constraints.clone())
    }

    fn exact_newton_joint_psi_terms(
        &self,
        block_states: &[ParameterBlockState],
        specs: &[ParameterBlockSpec],
        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
        psi_index: usize,
    ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
            return Err(SmoothError::dimension_mismatch(format!(
                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
                block_states.len(),
                specs.len(),
                derivative_blocks.len()
            ))
            .into());
        }
        derivative_blocks[0]
            .get(psi_index)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
        let hyper = self
            .hyperspecs
            .get(psi_index)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
        let beta = &block_states[0].beta;
        let eval = self.exact_evaluation(beta)?;
        let (direct, beta_mixed, betahessian_explicit) =
            self.adaptive_hyper_parts(&eval, *hyper)?;

        // Exact pseudo-Laplace psi-gradient.
        //
        // For one hyperparameter coordinate a we use the exact formula
        //
        //   d/da L_tilde
        //   = J_a + 0.5 tr(H^{-1} Hdot_a),
        //
        // with
        //
        //   H u_a   = J_{beta,a},
        //   beta_a  = -u_a,
        //   Hdot_a  = J_{beta,beta,a} + D_beta(H)[beta_a]
        //           = J_{beta,beta,a} - D_beta(H)[u_a].
        //
        // Here:
        //   - `direct` is J_a,
        //   - `beta_mixed` is J_{beta,a},
        //   - `betahessian_explicit` is J_{beta,beta,a},
        //   - `exacthessian_directional_derivative_from_evaluation(..., u)` returns
        //     D_beta(H)[u] for the exact likelihood-plus-Charbonnier model.
        Ok(Some(ExactNewtonJointPsiTerms {
            objective_psi: direct,
            score_psi: beta_mixed,
            hessian_psi: betahessian_explicit,
            hessian_psi_operator: None,
        }))
    }

    fn exact_newton_joint_psisecond_order_terms(
        &self,
        block_states: &[ParameterBlockState],
        specs: &[ParameterBlockSpec],
        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
        psi_i: usize,
        psi_j: usize,
    ) -> Result<Option<crate::custom_family::ExactNewtonJointPsiSecondOrderTerms>, String> {
        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
            return Err(SmoothError::dimension_mismatch(format!(
                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
                block_states.len(),
                specs.len(),
                derivative_blocks.len()
            ))
            .into());
        }
        derivative_blocks[0]
            .get(psi_i)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
        derivative_blocks[0]
            .get(psi_j)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
        let hyper_i = self
            .hyperspecs
            .get(psi_i)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
        let hyper_j = self
            .hyperspecs
            .get(psi_j)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
        let beta = &block_states[0].beta;
        let eval = self.exact_evaluation(beta)?;
        let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
            self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;

        Ok(Some(
            crate::custom_family::ExactNewtonJointPsiSecondOrderTerms {
                objective_psi_psi,
                score_psi_psi,
                hessian_psi_psi,
                hessian_psi_psi_operator: None,
            },
        ))
    }

    fn exact_newton_joint_psihessian_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        specs: &[ParameterBlockSpec],
        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
        psi_index: usize,
        direction: &Array1<f64>,
    ) -> Result<Option<Array2<f64>>, String> {
        if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
            return Err(SmoothError::dimension_mismatch(format!(
                "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
                block_states.len(),
                specs.len(),
                derivative_blocks.len()
            ))
            .into());
        }
        let beta = &block_states[0].beta;
        if direction.len() != beta.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "spatial adaptive exact family direction length mismatch: got {}, expected {}",
                direction.len(),
                beta.len()
            ))
            .into());
        }
        derivative_blocks[0]
            .get(psi_index)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
        let hyper = self
            .hyperspecs
            .get(psi_index)
            .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
        let eval = self.exact_evaluation(beta)?;
        let drift = match hyper.kind {
            SpatialAdaptiveHyperKind::LogLambdaMagnitude
            | SpatialAdaptiveHyperKind::LogLambdaGradient
            | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
                &eval,
                hyper.cache_index,
                AdaptiveComponent::from_index(hyper.kind.component_index())?,
                HyperDriftKind::Rho,
                direction,
            )?,
            SpatialAdaptiveHyperKind::LogEpsilonMagnitude
            | SpatialAdaptiveHyperKind::LogEpsilonGradient
            | SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
                .adaptive_shared_log_epsilon_drift(
                    &eval,
                    hyper.kind.component_index(),
                    direction,
                )?,
        };
        Ok(Some(drift))
    }
}

fn expect_single_block_state<'a>(
    block_states: &'a [ParameterBlockState],
    family_name: &str,
) -> Result<&'a ParameterBlockState, String> {
    crate::util::block_count::validate_block_count::<SmoothError>(
        family_name,
        1,
        block_states.len(),
    )?;
    Ok(&block_states[0])
}

fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
    if block_idx != 0 {
        return Err(SmoothError::invalid_index(format!(
            "{family_name} expects block_idx 0{context}, got {block_idx}"
        ))
        .into());
    }
    Ok::<(), _>(())
}

impl BoundedLinearFamily {
    fn bounded_term_derivative_data(
        &self,
        latent_beta: &Array1<f64>,
    ) -> (
        Array1<f64>,
        Array1<f64>,
        Array1<f64>,
        Array1<f64>,
        Array1<f64>,
    ) {
        let p = latent_beta.len();
        let mut beta_user = latent_beta.clone();
        let mut jac_diag = Array1::<f64>::ones(p);
        let mut second_diag = Array1::<f64>::zeros(p);
        let mut third_diag = Array1::<f64>::zeros(p);
        let mut priorthird = Array1::<f64>::zeros(p);
        for term in &self.bounded_terms {
            let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
                bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
            beta_user[term.col_idx] = beta;
            jac_diag[term.col_idx] = db_dtheta;
            second_diag[term.col_idx] = d2b_dtheta2;
            third_diag[term.col_idx] = d3b_dtheta3;
            let (_, _, _, prior_neghess_derivative) =
                bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
            priorthird[term.col_idx] = prior_neghess_derivative;
        }
        (beta_user, jac_diag, second_diag, third_diag, priorthird)
    }

    fn user_beta_and_jacobian(&self, latent_beta: &Array1<f64>) -> (Array1<f64>, Array1<f64>) {
        let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
        (beta_user, jac_diag)
    }

    fn nonlinear_offset_from_latent(&self, latent_beta: &Array1<f64>) -> Array1<f64> {
        let mut offset = self.offset.clone();
        for term in &self.bounded_terms {
            let (beta, _, _) =
                bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
            offset.scaled_add(beta, &self.design.column(term.col_idx));
        }
        offset
    }

    fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
        let mut x_eff = self.design.clone();
        for term in &self.bounded_terms {
            x_eff
                .column_mut(term.col_idx)
                .mapv_inplace(|v| v * jac_diag[term.col_idx]);
        }
        x_eff
    }

    fn exacthessian_andgradient(
        &self,
        latent_beta: &Array1<f64>,
    ) -> Result<
        (
            StandardFamilyObservationState,
            Array2<f64>,
            Array1<f64>,
            f64,
            Array1<f64>,
            Array1<f64>,
            Array1<f64>,
        ),
        String,
    > {
        let (_, jac_diag, second_diag, third_diag, priorthird) =
            self.bounded_term_derivative_data(latent_beta);
        let x_eff = self.effective_design_for_latent(&jac_diag);
        let eta =
            self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta);
        let obs = evaluate_standard_familyobservations(
            self.family.clone(),
            self.latent_cloglog_state.as_ref(),
            self.mixture_link_state.as_ref(),
            self.sas_link_state.as_ref(),
            &self.y,
            &self.weights,
            &eta,
        )
        .map_err(|e| e.to_string())?;

        let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
        let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
        let mut prior_loglik = 0.0;
        for term in &self.bounded_terms {
            let (logp, grad, neghess, _) =
                bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
            prior_loglik += logp;
            priorgrad[term.col_idx] += grad;
            prior_neghess[[term.col_idx, term.col_idx]] += neghess;
        }

        let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
        let mut gradient = fast_atv(&x_eff, &obs.score);
        for term in &self.bounded_terms {
            let score_beta = self.design.column(term.col_idx).dot(&obs.score);
            hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
        }
        hessian += &prior_neghess;
        gradient += &priorgrad;

        Ok((
            obs,
            hessian,
            gradient,
            prior_loglik,
            second_diag,
            third_diag,
            priorthird,
        ))
    }

    fn evaluation_from_latent(
        &self,
        latent_beta: &Array1<f64>,
    ) -> Result<
        (
            StandardFamilyObservationState,
            Array2<f64>,
            Array1<f64>,
            f64,
        ),
        String,
    > {
        let (obs, hessian, gradient, prior_loglik, _, _, _) =
            self.exacthessian_andgradient(latent_beta)?;
        Ok((obs, hessian, gradient, prior_loglik))
    }
}

impl CustomFamily for BoundedLinearFamily {
    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
        let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
        Ok(FamilyEvaluation {
            log_likelihood: obs.log_likelihood + prior_loglik,
            blockworking_sets: vec![BlockWorkingSet::ExactNewton {
                gradient,
                hessian: SymmetricMatrix::Dense(hessian),
            }],
        })
    }

    fn exact_newton_joint_hessian(
        &self,
        block_states: &[ParameterBlockState],
    ) -> Result<Option<Array2<f64>>, String> {
        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
        let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
        Ok(Some(hessian))
    }

    fn exact_newton_hessian_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        block_idx: usize,
        d_beta: &Array1<f64>,
    ) -> Result<Option<Array2<f64>>, String> {
        expect_block_idx_zero(block_idx, "bounded linear family", "")?;
        self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
    }

    fn exact_newton_joint_hessian_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        d_beta_flat: &Array1<f64>,
    ) -> Result<Option<Array2<f64>>, String> {
        let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
        if d_beta_flat.len() != latent_beta.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "bounded linear family directional derivative length mismatch: got {}, expected {}",
                d_beta_flat.len(),
                latent_beta.len()
            ))
            .into());
        }

        let (obs, _, _, _, second_diag, third_diag, priorthird) =
            self.exacthessian_andgradient(latent_beta)?;

        let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
        let x_eff = self.effective_design_for_latent(&jac_diag);
        let deta = x_eff.dot(d_beta_flat);
        let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;

        let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
        for term in &self.bounded_terms {
            let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
            if scale != 0.0 {
                let mut col = dx_eff.column_mut(term.col_idx);
                col.assign(&self.design.column(term.col_idx));
                col.mapv_inplace(|v| v * scale);
            }
        }

        let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
        let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
        for i in 0..x_eff.nrows() {
            let wi = obs.neghessian_eta[i];
            if wi == 0.0 {
                continue;
            }
            for a in 0..x_eff.ncols() {
                let xa = x_eff[[i, a]];
                for b in 0..x_eff.ncols() {
                    wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
                }
            }
        }
        dhessian += &wxdx;

        let d_score = -&obs.neghessian_eta * &deta;
        for term in &self.bounded_terms {
            let score_beta = self.design.column(term.col_idx).dot(&obs.score);
            let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
            dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
                + score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
            dhessian[[term.col_idx, term.col_idx]] +=
                priorthird[term.col_idx] * d_beta_flat[term.col_idx];
        }

        Ok(Some(dhessian))
    }

    fn block_geometry(
        &self,
        block_states: &[ParameterBlockState],
        spec: &ParameterBlockSpec,
    ) -> Result<(DesignMatrix, Array1<f64>), String> {
        if block_states.is_empty() {
            return Ok((
                DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(
                    self.designzeroed.clone(),
                )),
                self.offset.clone(),
            ));
        }
        let offset = self.nonlinear_offset_from_latent(
            &expect_single_block_state(block_states, "bounded linear family")?.beta,
        );
        let x = if spec.design.ncols() == self.designzeroed.ncols() {
            self.designzeroed.clone()
        } else {
            return Err(SmoothError::dimension_mismatch(
                "bounded linear family design column mismatch",
            )
            .into());
        };
        Ok((
            DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(x)),
            offset,
        ))
    }

    fn block_geometry_is_dynamic(&self) -> bool {
        true
    }

    fn block_geometry_directional_derivative(
        &self,
        block_states: &[ParameterBlockState],
        block_idx: usize,
        spec: &ParameterBlockSpec,
        d_beta: &Array1<f64>,
    ) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
        expect_block_idx_zero(
            block_idx,
            "bounded linear family",
            " for geometry derivative",
        )?;
        expect_single_block_state(block_states, "bounded linear family")?;
        if d_beta.len() != spec.design.ncols() {
            return Err(SmoothError::dimension_mismatch(format!(
                "bounded linear family geometry derivative direction mismatch: got {}, expected {}",
                d_beta.len(),
                spec.design.ncols()
            ))
            .into());
        }
        let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta);
        let mut d_offset = Array1::<f64>::zeros(self.offset.len());
        let has_drift = self
            .bounded_terms
            .iter()
            .any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
        if !has_drift {
            return Ok(Some(BlockGeometryDirectionalDerivative {
                d_design: None,
                d_offset,
            }));
        }
        for term in &self.bounded_terms {
            let col = term.col_idx;
            let drift = jac_diag[col] * d_beta[col];
            if drift != 0.0 {
                d_offset.scaled_add(drift, &self.design.column(col));
            }
        }
        Ok(Some(BlockGeometryDirectionalDerivative {
            d_design: None,
            d_offset,
        }))
    }
}

#[inline]
fn dense_diag_gram_chunkrows(p: usize) -> usize {
    const MIN_ROWS: usize = 512;
    const MAX_ROWS: usize = 2048;
    const TARGET_BYTES: usize = 2 * 1024 * 1024;
    let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
    (TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
}

fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
    if x.nrows() != w.len() {
        return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
    }
    let (n, p) = x.dim();
    if n == 0 || p == 0 {
        return Ok(Array2::<f64>::zeros((p, p)));
    }

    const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
    let dense_work_bytes = n
        .checked_mul(p)
        .and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
        .unwrap_or(usize::MAX);
    if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
        let mut weighted = x.to_owned();
        ndarray::Zip::from(weighted.rows_mut())
            .and(w)
            .par_for_each(|mut row, wi| row *= *wi);
        return Ok(fast_atb(&x, &weighted));
    }

    let chunkrows = dense_diag_gram_chunkrows(p).min(n);
    let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
    let mut out = Array2::<f64>::zeros((p, p));
    for row_start in (0..n).step_by(chunkrows) {
        let rows = (n - row_start).min(chunkrows);
        let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
        {
            let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
            for local_row in 0..rows {
                let scale = w[row_start + local_row];
                if scale == 0.0 {
                    chunk.row_mut(local_row).fill(0.0);
                    continue;
                }
                for col in 0..p {
                    chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
                }
            }
        }
        out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
    }
    Ok(out)
}

fn trace_of_dense_product(a: &Array2<f64>, b: &Array2<f64>) -> Result<f64, String> {
    if a.nrows() != a.ncols() || b.nrows() != b.ncols() || a.nrows() != b.nrows() {
        return Err(
            SmoothError::dimension_mismatch("trace_of_dense_product dimension mismatch").into(),
        );
    }
    let mut trace = 0.0;
    for i in 0..a.nrows() {
        for j in 0..a.ncols() {
            trace += a[[i, j]] * b[[j, i]];
        }
    }
    Ok(trace)
}

fn exact_bounded_edf(
    penalties: &[PenaltySpec],
    lambdas: &Array1<f64>,
    latent_cov: &Array2<f64>,
) -> Result<(Vec<f64>, f64), EstimationError> {
    if penalties.len() != lambdas.len() {
        crate::bail_invalid_estim!(
            "bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
            penalties.len(),
            lambdas.len()
        );
    }
    if latent_cov.nrows() != latent_cov.ncols() {
        crate::bail_invalid_estim!("bounded EDF covariance must be square");
    }

    let p = latent_cov.nrows();
    let mut s_lambda = Array2::<f64>::zeros((p, p));
    let mut edf_by_block = Vec::with_capacity(penalties.len());
    let mut trace_sum = 0.0;

    for (k, ps) in penalties.iter().enumerate() {
        let lambda_k = lambdas[k];
        match ps {
            PenaltySpec::Block {
                local, col_range, ..
            } => {
                s_lambda
                    .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
                    .scaled_add(lambda_k, local);
                // Compute penalty rank from the block-local matrix directly.
                let penalty_rank =
                    local
                        .nrows()
                        .saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
                            EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
                        })?);
                // Trace only involves the block slice of latent_cov.
                let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
                let trace_k = lambda_k
                    * trace_of_dense_product(&cov_block.to_owned(), local)
                        .map_err(EstimationError::InvalidInput)?;
                trace_sum += trace_k;
                let p_k = penalty_rank as f64;
                edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
            }
            PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
                s_lambda.scaled_add(lambda_k, m);
                let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
                    EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
                })?);
                let trace_k = lambda_k
                    * trace_of_dense_product(latent_cov, m)
                        .map_err(EstimationError::InvalidInput)?;
                trace_sum += trace_k;
                let p_k = penalty_rank as f64;
                edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
            }
        }
    }

    let nullity_total = estimate_penalty_nullity(&s_lambda)
        .map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
        as f64;
    let edf_total = (p as f64 - trace_sum).clamp(nullity_total, p as f64);
    Ok((edf_by_block, edf_total))
}

fn transform_bounded_latent_precision_to_user_internal(
    latent_precision: &Array2<f64>,
    jac_diag: &Array1<f64>,
) -> Result<Array2<f64>, EstimationError> {
    let p = latent_precision.nrows();
    if latent_precision.ncols() != p || jac_diag.len() != p {
        crate::bail_invalid_estim!(
            "bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
            latent_precision.nrows(),
            latent_precision.ncols(),
            jac_diag.len()
        );
    }
    let mut out = latent_precision.clone();
    for i in 0..p {
        let scale = jac_diag[i];
        if !scale.is_finite() || scale <= 0.0 {
            crate::bail_invalid_estim!(
                "bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
            );
        }
        if scale != 1.0 {
            out.row_mut(i).mapv_inplace(|v| v / scale);
            out.column_mut(i).mapv_inplace(|v| v / scale);
        }
    }
    Ok(out)
}

fn fit_bounded_term_collection_with_design(
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    spec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    heuristic_lambdas: Option<&[f64]>,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollection, EstimationError> {
    let conditioning_cols: Vec<usize> = spec
        .linear_terms
        .iter()
        .enumerate()
        .filter_map(|(j, linear)| {
            (!linear.double_penalty).then_some(design.intercept_range.end + j)
        })
        .collect();
    let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
    let dense_design = design.design.to_dense_cow();
    let fit_design = conditioning.apply_to_design(&dense_design);
    let fit_penalties = conditioning
        .transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
    if design.linear_constraints.is_some() {
        crate::bail_invalid_estim!(
            "bounded() terms are not yet compatible with explicit linear constraints"
        );
    }
    let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
    for (j, term) in spec.linear_terms.iter().enumerate() {
        if term.double_penalty
            && matches!(
                term.coefficient_geometry,
                LinearCoefficientGeometry::Bounded { .. }
            )
        {
            crate::bail_invalid_estim!(
                "bounded linear term '{}' cannot also use double_penalty",
                term.name
            );
        }
        if let LinearCoefficientGeometry::Bounded { min, max, prior } =
            term.coefficient_geometry.clone()
        {
            let col_idx = design.intercept_range.end + j;
            let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
            bounded_terms.push(BoundedLinearTermMeta {
                col_idx,
                min: min_internal,
                max: max_internal,
                prior,
            });
        }
    }
    if bounded_terms.is_empty() {
        crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
    }

    let mut designzeroed = fit_design.clone();
    let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
    for term in &bounded_terms {
        designzeroed.column_mut(term.col_idx).fill(0.0);
        initial_beta[term.col_idx] = bounded_logit(0.5);
    }

    let initial_log_lambdas = heuristic_lambdas
        .map(|vals| Array1::from_vec(vals.to_vec()))
        .unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
    if initial_log_lambdas.len() != fit_penalties.len() {
        crate::bail_invalid_estim!(
            "heuristic lambda length mismatch for bounded model: got {}, expected {}",
            initial_log_lambdas.len(),
            fit_penalties.len()
        );
    }

    let is_beta_logistic = family.is_binomial_beta_logistic();
    let family_adapter = BoundedLinearFamily {
        family: family.clone(),
        latent_cloglog_state: options.latent_cloglog,
        mixture_link_state: options
            .mixture_link
            .clone()
            .as_ref()
            .map(state_fromspec)
            .transpose()
            .map_err(EstimationError::InvalidInput)?,
        sas_link_state: options
            .sas_link
            .map(|spec| {
                if is_beta_logistic {
                    state_from_beta_logisticspec(spec)
                } else {
                    state_from_sasspec(spec)
                }
            })
            .transpose()
            .map_err(EstimationError::InvalidInput)?,
        y: y.to_owned(),
        weights: weights.to_owned(),
        design: fit_design.clone(),
        designzeroed: designzeroed.clone(),
        offset: offset.to_owned(),
        bounded_terms: bounded_terms.clone(),
    };
    let blockspec = ParameterBlockSpec {
        name: "eta".to_string(),
        design: DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(designzeroed)),
        offset: offset.to_owned(),
        penalties: fit_penalties
            .iter()
            .map(|ps| match ps {
                PenaltySpec::Block {
                    local, col_range, ..
                } => PenaltyMatrix::Blockwise {
                    local: local.clone(),
                    col_range: col_range.clone(),
                    total_dim: design.design.ncols(),
                },
                PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
                    PenaltyMatrix::Dense(m.clone())
                }
            })
            .collect(),
        nullspace_dims: design.nullspace_dims.clone(),
        initial_log_lambdas,
        initial_beta: Some(initial_beta),
        gauge_priority: 100,
        // Report the true β-dependent Jacobian (bounded columns scaled by
        // dβ/dθ) to the identifiability audit so it does not mistake the
        // deliberately-zeroed placeholder columns for a structural rank
        // deficiency. The inner solve still drives η through the family
        // adapter, so this does not affect the fit geometry.
        jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
            design: fit_design.clone(),
            bounded_terms: bounded_terms.clone(),
        })),
        stacked_design: None,
        stacked_offset: None,
    };
    let fit = fit_custom_family(
        &family_adapter,
        &[blockspec],
        &BlockwiseFitOptions {
            inner_max_cycles: options.max_iter,
            inner_tol: options.tol,
            outer_max_iter: options.max_iter,
            outer_tol: options.tol,
            ..BlockwiseFitOptions::default()
        },
    )
    .map_err(EstimationError::CustomFamily)?;

    let latent_beta = fit.block_states[0].beta.clone();
    let (beta_user_internal, jac_diag) = family_adapter.user_beta_and_jacobian(&latent_beta);
    let beta_user = conditioning.backtransform_beta(&beta_user_internal);
    let latent_cov = fit.covariance_conditional.clone();
    let beta_covariance = latent_cov.as_ref().map(|cov| {
        let mut out = cov.clone();
        for i in 0..out.nrows() {
            out.row_mut(i).mapv_inplace(|v| v * jac_diag[i]);
        }
        for j in 0..out.ncols() {
            out.column_mut(j).mapv_inplace(|v| v * jac_diag[j]);
        }
        conditioning.backtransform_covariance(&out)
    });
    let beta_standard_errors = beta_covariance
        .as_ref()
        .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));

    let (eta_state, h_data, _, _) = family_adapter
        .evaluation_from_latent(&latent_beta)
        .map_err(EstimationError::InvalidInput)?;
    let p_fit = fit_design.ncols();
    let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
    for (k, penalty) in fit_penalties.iter().enumerate() {
        match penalty {
            PenaltySpec::Block {
                local, col_range, ..
            } => {
                s_lambda_internal
                    .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
                    .scaled_add(fit.lambdas[k], local);
            }
            PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
                s_lambda_internal.scaled_add(fit.lambdas[k], m);
            }
        }
    }
    let mut penalized_hessian = h_data.clone();
    penalized_hessian += &s_lambda_internal;
    let penalized_hessian =
        transform_bounded_latent_precision_to_user_internal(&penalized_hessian, &jac_diag)?;
    let penalized_hessian =
        conditioning.transform_penalized_hessian_to_original(&penalized_hessian);
    let s_lambda_original = weighted_blockwise_penalty_sum(
        &design.penalties,
        fit.lambdas.as_slice().unwrap(),
        design.design.ncols(),
    );
    let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
    let deviance = if family.is_gaussian_identity() {
        y.iter()
            .zip(eta_state.mu.iter())
            .zip(weights.iter())
            .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
            .sum()
    } else {
        -2.0 * eta_state.log_likelihood
    };
    let (edf_by_block, edf_total) = if let Some(cov) = latent_cov.as_ref() {
        exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
    } else {
        (vec![0.0; fit_penalties.len()], 0.0)
    };
    let geometry = Some(crate::estimate::FitGeometry {
        penalized_hessian: penalized_hessian.clone().into(),
        working_weights: eta_state.fisherweight.clone(),
        working_response: {
            let mut working_response = eta_state.eta.clone();
            for i in 0..working_response.len() {
                let wi = eta_state.fisherweight[i].max(1e-12);
                working_response[i] += eta_state.score[i] / wi;
            }
            working_response
        },
    });
    let max_abs_eta = eta_state
        .eta
        .iter()
        .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
    Ok(FittedTermCollection {
        fit: {
            let log_lambdas = fit.lambdas.mapv(|v| v.max(1e-300).ln());
            let inf = FitInference {
                edf_by_block,
                edf_total,
                smoothing_correction: None,
                // Boundary adapter: `penalized_hessian` storage is now
                // `UnscaledPrecision`.
                penalized_hessian: penalized_hessian.clone().into(),
                working_weights: eta_state.fisherweight.clone(),
                working_response: {
                    let mut working_response = eta_state.eta.clone();
                    for i in 0..working_response.len() {
                        let wi = eta_state.fisherweight[i].max(1e-12);
                        working_response[i] += eta_state.score[i] / wi;
                    }
                    working_response
                },
                reparam_qs: None,
                dispersion: crate::estimate::Dispersion::Known(1.0),
                beta_covariance: beta_covariance
                    .clone()
                    .map(crate::inference::dispersion_cov::PhiScaledCovariance::from),
                beta_standard_errors,
                beta_covariance_corrected: None,
                beta_standard_errors_corrected: None,
                beta_covariance_frequentist: None,
                coefficient_influence: None,
                bias_correction_beta: None,
            };
            let covariance_conditional = beta_covariance;
            let pirls_status_val = if fit.outer_converged {
                crate::pirls::PirlsStatus::Converged
            } else {
                crate::pirls::PirlsStatus::StalledAtValidMinimum
            };
            UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
                blocks: vec![crate::estimate::FittedBlock {
                    beta: beta_user.clone(),
                    role: crate::estimate::BlockRole::Mean,
                    edf: edf_total,
                    lambdas: fit.lambdas.clone(),
                }],
                log_lambdas,
                lambdas: fit.lambdas,
                likelihood_scale: family.default_scale_metadata(),
                likelihood_family: Some(family),
                log_likelihood_normalization:
                    crate::types::LogLikelihoodNormalization::UserProvided,
                log_likelihood: eta_state.log_likelihood,
                deviance,
                reml_score: fit.penalized_objective,
                stable_penalty_term: penalty_term,
                penalized_objective: fit.penalized_objective,
                outer_iterations: fit.outer_iterations,
                outer_converged: fit.outer_converged,
                outer_gradient_norm: fit.outer_gradient_norm,
                standard_deviation: 1.0,
                covariance_conditional,
                covariance_corrected: None,
                inference: Some(inf),
                fitted_link: crate::estimate::FittedLinkState::Standard(None),
                geometry,
                block_states: Vec::new(),
                pirls_status: pirls_status_val,
                max_abs_eta,
                constraint_kkt: None,
                artifacts: crate::estimate::FitArtifacts {
                    pirls: None,
                    ..Default::default()
                },
                inner_cycles: 0,
            })?
        },
        design: design.clone(),
        adaptive_diagnostics: None,
    })
}

fn enforce_term_constraint_feasibility(
    design: &TermCollectionDesign,
    fit: &UnifiedFitResult,
) -> Result<(), EstimationError> {
    // Geometric (per-row-scaled) tolerance, matching the public contract on
    // `ACTIVE_SET_PRIMAL_FEASIBILITY_TOL` and the diagnostic that
    // `compute_constraint_kkt_diagnostics` exposes via `fit.constraint_kkt`.
    // Lower-bound rows are unit-norm (a_i = e_i) so the scale-invariant and
    // raw checks coincide there. Linear-inequality rows generally are NOT
    // unit-norm — e.g. a B-spline endpoint-derivative clamp at k = 12 carries
    // ‖a_i‖ ≈ 38, so a 1e-6 raw residual is only 2.6e-8 in geometric units.
    // Holding this gate to raw 1e-7 while the in-solver acceptance gate
    // measures geometric 1e-8 is the inconsistency that made well-conditioned
    // clamped fits get rejected after they completed cleanly.
    let tol = 1e-7;
    let smooth_start = design
        .design
        .ncols()
        .saturating_sub(design.smooth.total_smooth_cols());
    let mut violations: Vec<String> = Vec::new();
    for term in &design.smooth.terms {
        let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
        let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
        if let Some(lb) = term.lower_bounds_local.as_ref() {
            let mut worst = 0.0_f64;
            let mut worst_idx = 0usize;
            for i in 0..lb.len().min(beta_local.len()) {
                if lb[i].is_finite() {
                    let viol = (lb[i] - beta_local[i]).max(0.0);
                    if viol > worst {
                        worst = viol;
                        worst_idx = i;
                    }
                }
            }
            if worst > tol {
                violations.push(format!(
                    "term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
                    term.name, worst, worst_idx
                ));
            }
        }
        if let Some(lin) = term.linear_constraints_local.as_ref() {
            let mut worst = 0.0_f64;
            let mut worstrow = 0usize;
            for i in 0..lin.a.nrows() {
                let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
                let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
                let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
                let viol = (-s).max(0.0);
                if viol > worst {
                    worst = viol;
                    worstrow = i;
                }
            }
            if worst > tol {
                violations.push(format!(
                    "term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
                    term.name, worst, worstrow
                ));
            }
        }
    }

    if !violations.is_empty() {
        let mut msg = format!(
            "constraint violation after fit ({} violating term constraints): {}",
            violations.len(),
            violations.join(" | ")
        );
        if let Some(kkt) = fit.constraint_kkt.as_ref() {
            msg.push_str(&format!(
                "; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]",
                kkt.primal_feasibility, kkt.dual_feasibility, kkt.complementarity, kkt.stationarity
            ));
        }
        return Err(EstimationError::ParameterConstraintViolation(msg));
    }
    Ok(())
}

fn stratified_spatial_subsample(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
    target_size: usize,
) -> Vec<usize> {
    use rand::SeedableRng;
    use rand::rngs::StdRng;
    use rand::seq::SliceRandom;

    let n = data.nrows();
    if n <= target_size {
        return (0..n).collect();
    }

    let spatial_cols: Option<Vec<usize>> =
        spec.smooth_terms.iter().find_map(|term| match &term.basis {
            SmoothBasisSpec::ThinPlate { feature_cols, .. }
            | SmoothBasisSpec::Matern { feature_cols, .. }
            | SmoothBasisSpec::Duchon { feature_cols, .. } => {
                if !feature_cols.is_empty() {
                    Some(feature_cols.clone())
                } else {
                    None
                }
            }
            _ => None,
        });

    let cols = match spatial_cols {
        Some(c) if !c.is_empty() => c,
        _ => {
            let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
            let mut indices: Vec<usize> = (0..n).collect();
            indices.shuffle(&mut rng);
            indices.truncate(target_size);
            indices.sort_unstable();
            return indices;
        }
    };
    let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));

    let d = cols.len();
    let mut mins = vec![f64::INFINITY; d];
    let mut maxs = vec![f64::NEG_INFINITY; d];
    for i in 0..n {
        for (ax, &col) in cols.iter().enumerate() {
            let v = data[[i, col]];
            if v < mins[ax] {
                mins[ax] = v;
            }
            if v > maxs[ax] {
                maxs[ax] = v;
            }
        }
    }

    let total_cells_target = (target_size / 5).max(1);
    let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
    let cells_per_axis = cells_per_axis.max(1);

    let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
        std::collections::HashMap::new();
    for i in 0..n {
        let mut cell_key = Vec::with_capacity(d);
        for (ax, &col) in cols.iter().enumerate() {
            let range = maxs[ax] - mins[ax];
            let cell = if range <= 0.0 {
                0
            } else {
                let frac = (data[[i, col]] - mins[ax]) / range;
                (frac * cells_per_axis as f64).floor() as usize
            };
            cell_key.push(cell.min(cells_per_axis - 1));
        }
        cell_members.entry(cell_key).or_default().push(i);
    }

    let mut selected: Vec<usize> = Vec::with_capacity(target_size);
    let mut remaining_budget = target_size;
    let mut remaining_population = n;

    let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
    cells.sort_by(|a, b| a.0.cmp(&b.0));

    for (_, members) in &mut cells {
        if remaining_budget == 0 {
            break;
        }
        let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
            .round() as usize;
        let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
        members.shuffle(&mut rng);
        selected.extend_from_slice(&members[..alloc]);
        remaining_budget = remaining_budget.saturating_sub(alloc);
        remaining_population = remaining_population.saturating_sub(members.len());
    }

    if selected.len() > target_size {
        selected.shuffle(&mut rng);
        selected.truncate(target_size);
    }

    selected.sort_unstable();
    selected
}

fn spatial_subsample_seed(
    data: ArrayView2<'_, f64>,
    spatial_cols: &[usize],
    target_size: usize,
) -> u64 {
    let mut state = 0x5350_4154_4941_4C53_u64;
    spatial_seed_mix(&mut state, data.nrows() as u64);
    spatial_seed_mix(&mut state, data.ncols() as u64);
    spatial_seed_mix(&mut state, target_size as u64);
    spatial_seed_mix(&mut state, spatial_cols.len() as u64);
    for &col in spatial_cols {
        spatial_seed_mix(&mut state, col as u64);
    }

    if data.nrows() > 0 {
        let mid = data.nrows() / 2;
        let last = data.nrows() - 1;
        for &row in &[0usize, mid, last] {
            for &col in spatial_cols {
                let value = data[[row, col]];
                spatial_seed_mix(&mut state, value.to_bits());
            }
        }
    }
    state
}

#[inline]
fn spatial_seed_mix(state: &mut u64, value: u64) {
    let mut z = value
        .wrapping_add(0x9E37_79B9_7F4A_7C15)
        .wrapping_add(*state);
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    *state ^= z ^ (z >> 31);
    *state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
}

fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
    let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
    for (new_row, &orig_row) in indices.iter().enumerate() {
        sampled.row_mut(new_row).assign(&data.row(orig_row));
    }
    sampled
}

fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
    match spatial_term_center_strategy(term) {
        Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
        _ => None,
    }
}

fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
    if values.len() != expected_dim || expected_dim <= 1 {
        return None;
    }
    if values.iter().any(|value| !value.is_finite()) {
        return None;
    }
    Some(center_aniso_log_scales(values))
}

fn blended_pilot_axis_contrasts(
    pilot_data: ArrayView2<'_, f64>,
    term: &SmoothTermSpec,
    centers: ArrayView2<'_, f64>,
) -> Option<Vec<f64>> {
    let d = centers.ncols();
    if d <= 1 {
        return None;
    }
    let center_eta = initial_aniso_contrasts(centers);
    let data_eta = standardized_spatial_term_data(pilot_data, term)
        .ok()
        .and_then(|x| finite_centered_axis_contrasts(&initial_aniso_contrasts(x.view()), d));
    let center_eta = finite_centered_axis_contrasts(&center_eta, d)?;
    let blended = match data_eta {
        Some(data_eta) => center_eta
            .iter()
            .zip(data_eta.iter())
            .map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
            .collect::<Vec<_>>(),
        None => center_eta,
    };
    finite_centered_axis_contrasts(&blended, d)
}

fn apply_pilot_spatial_psi_reseed(
    pilot_data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
    spatial_terms: &[usize],
    kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<TermCollectionSpec, EstimationError> {
    let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
    let use_aniso = has_aniso_terms(spec, spatial_terms);
    let log_kappa0 = if use_aniso {
        SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
    } else {
        SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
    };
    let log_kappa0 = log_kappa0.reseed_from_data(pilot_data, spec, spatial_terms, kappa_options);
    let log_kappa_lower = if use_aniso {
        SpatialLogKappaCoords::lower_bounds_aniso_from_data(
            pilot_data,
            spec,
            spatial_terms,
            &dims_per_term,
            kappa_options,
        )
    } else {
        SpatialLogKappaCoords::lower_bounds_from_data(
            pilot_data,
            spec,
            spatial_terms,
            kappa_options,
        )
    };
    let log_kappa_upper = if use_aniso {
        SpatialLogKappaCoords::upper_bounds_aniso_from_data(
            pilot_data,
            spec,
            spatial_terms,
            &dims_per_term,
            kappa_options,
        )
    } else {
        SpatialLogKappaCoords::upper_bounds_from_data(
            pilot_data,
            spec,
            spatial_terms,
            kappa_options,
        )
    };
    log_kappa0
        .clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
        .apply_tospec(spec, spatial_terms)
}

pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
    data: ArrayView2<'_, f64>,
    spec: &mut TermCollectionSpec,
    spatial_terms: &[usize],
    target_size: usize,
    kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> usize {
    if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
    {
        return 0;
    }
    if !has_aniso_terms(spec, spatial_terms) {
        return 0;
    }
    let indices = stratified_spatial_subsample(data, spec, target_size);
    let pilot_data = sampled_rows(data, &indices);
    let mut working = spec.clone();
    let mut updated_terms = 0usize;
    const GEOMETRY_UPDATES: usize = 2;

    for pass in 0..GEOMETRY_UPDATES {
        let planned_terms = match plan_joint_spatial_centers_for_term_blocks(
            pilot_data.view(),
            &[working.smooth_terms.clone()],
        )
        .and_then(|mut blocks| {
            blocks.pop().ok_or_else(|| {
                BasisError::InvalidInput(
                    "pilot geometry initializer produced no smooth-term block".to_string(),
                )
            })
        }) {
            Ok(terms) => terms,
            Err(err) => {
                log::warn!(
                    "[spatial-kappa] pilot geometry initializer skipped after center planning failed: {err}"
                );
                return updated_terms;
            }
        };

        for &term_idx in spatial_terms {
            let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
                continue;
            };
            let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
                continue;
            };
            if d <= 1 || current_eta.len() != d {
                continue;
            }
            let Some(planned_term) = planned_terms.get(term_idx) else {
                continue;
            };
            let Some(centers) = spatial_term_user_centers(planned_term) else {
                continue;
            };
            let Some(eta) = blended_pilot_axis_contrasts(pilot_data.view(), planned_term, centers)
            else {
                continue;
            };
            if set_spatial_aniso_log_scales(&mut working, term_idx, eta).is_ok() {
                updated_terms += usize::from(pass == 0);
            }
        }

        match apply_pilot_spatial_psi_reseed(
            pilot_data.view(),
            &working,
            spatial_terms,
            kappa_options,
        ) {
            Ok(updated) => {
                working = updated;
            }
            Err(err) => {
                log::warn!(
                    "[spatial-kappa] pilot geometry ψ reseed skipped after deterministic initializer error: {err}"
                );
                break;
            }
        }
    }

    if updated_terms > 0 {
        log::info!(
            "[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
            indices.len(),
            updated_terms
        );
        *spec = working;
    }
    updated_terms
}

pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
    spec.smooth_terms
        .iter()
        .enumerate()
        .filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
        .collect()
}

/// Returns `true` when every spatial term in `spec` has a locked kernel
/// scale (explicit `length_scale=X` without anisotropy) and therefore
/// contributes no outer ψ/κ optimization axis. Empty term collections
/// also return `true` — there are no kappas to optimize.
///
/// Used by family entry points that want to honor a user-supplied scalar
/// length scale exactly: when all spatial terms are locked the n-block
/// joint-spatial outer solver has nothing to optimize, and routing
/// through it merely spends ~80 outer iters chasing a stalled ARC at the
/// user's chosen ρ. Skipping straight to the rho-only path avoids that
/// waste and respects the user's explicit kernel-scale input.
pub fn all_spatial_terms_kappa_fixed(spec: &TermCollectionSpec) -> bool {
    spec.smooth_terms.iter().enumerate().all(|(idx, _)| {
        !spatial_term_supports_hyper_optimization(spec, idx)
            || spatial_term_has_locked_kappa(spec, idx)
    })
}

fn fit_score(fit: &UnifiedFitResult) -> f64 {
    if fit.reml_score.is_finite() {
        return fit.reml_score;
    }
    let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
    if score.is_finite() {
        score
    } else {
        f64::INFINITY
    }
}

/// Classify an outer-evaluation error as a *recoverable trial-point
/// infeasibility* versus a genuine fatal failure.
///
/// The spatial-κ / anisotropy outer optimizer probes a sequence of trial
/// hyperparameters. At an extreme trial point the realized kernel design or
/// its ψ-derivatives may simply be non-constructible — e.g. a learned
/// per-axis log-scale stretches the anisotropic distance `r = |Λh|` until the
/// Duchon polyharmonic blocks `r^(2m−d)` overflow, or a degenerate metric
/// collapses two centers onto a non-C² collision. Those points lie outside
/// the model's feasible domain; the principled response is to treat them like
/// the cost-only path already does (objective `+∞`) so the line-search /
/// trust-region solver retreats, rather than aborting the entire REML fit.
///
/// A `BasisError` is exactly this class: it means "the basis/design cannot be
/// built at this hyperparameter". Everything else (PIRLS divergence wired to a
/// distinct variant, layout/topology invariants, over-parameterization) stays
/// fatal so genuine bugs are never masked.
fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
    matches!(err, EstimationError::BasisError(_))
}

fn require_successful_spatial_optimization_result<T>(
    initial_score: f64,
    result: Result<Option<(T, f64)>, EstimationError>,
) -> Result<T, EstimationError> {
    match result {
        Ok(Some((value, exact_score))) => {
            // Allow rounding-level worsening: REML scores accumulate
            // log-determinant terms whose finite-precision re-evaluation
            // can drift well past 1e-10 absolute near a converged optimum
            // (we have seen ~1e-6 between two evaluations whose printed
            // values round to identical 6-digit scientific). Reject genuine
            // worsenings (>1 unit) but admit anything within ~1e-6
            // absolute / 1e-8 relative — meaningful REML gains are
            // orders of magnitude larger.
            let tol = 1e-6_f64.max(initial_score.abs() * 1e-8);
            if exact_score <= initial_score + tol {
                Ok(value)
            } else {
                Err(EstimationError::RemlOptimizationFailed(format!(
                    "spatial kappa optimization made REML score worse ({initial_score:.6e} -> {exact_score:.6e})"
                )))
            }
        }
        Ok(None) => Err(EstimationError::RemlOptimizationFailed(
            "spatial kappa optimization is unavailable for one or more eligible spatial terms"
                .to_string(),
        )),
        Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
            "spatial kappa optimization failed: {err}"
        ))),
    }
}

fn external_opts_for_design(
    family: &LikelihoodSpec,
    design: &TermCollectionDesign,
    options: &FitOptions,
) -> ExternalOptimOptions {
    ExternalOptimOptions {
        family: family.clone(),
        latent_cloglog: options.latent_cloglog,
        mixture_link: options.mixture_link.clone(),
        optimize_mixture: options.optimize_mixture,
        sas_link: options.sas_link,
        optimize_sas: options.optimize_sas,
        compute_inference: options.compute_inference,
        max_iter: options.max_iter,
        tol: options.tol,
        nullspace_dims: design.nullspace_dims.clone(),
        linear_constraints: design.linear_constraints.clone(),
        firth_bias_reduction: Some(options.firth_bias_reduction),
        penalty_shrinkage_floor: options.penalty_shrinkage_floor,
        rho_prior: options.rho_prior.clone(),
        // Propagate Kronecker structure so the joint optimizer minimizes the
        // same REML surface as the baseline/refit (adaptive_fit_options_base).
        kronecker_penalty_system: design.kronecker_penalty_system(),
        kronecker_factored: design
            .smooth
            .terms
            .iter()
            .find_map(|t| t.kronecker_factored.clone()),
    }
}

/// Evaluate the joint REML cost, gradient, and Hessian result at a given θ = [ρ, ψ]
/// for a single-block term collection with spatial hyperparameters.
///
/// This provides a direct evaluation of the profiled REML objective using the
/// external-caller interface, which exposes exact cost/gradient/Hessian without
/// running the full outer smoothing loop. The returned tuple is
/// `(cost, gradient, hessian)` in the joint [ρ, ψ] space.
fn evaluate_joint_reml_outer_eval_at_theta(
    evaluator: &mut crate::estimate::ExternalJointHyperEvaluator<'_>,
    design: &TermCollectionDesign,
    theta: &Array1<f64>,
    rho_dim: usize,
    hyper_dirs: Vec<crate::estimate::reml::DirectionalHyperParam>,
    warm_start_beta: Option<ArrayView1<'_, f64>>,
    order: crate::solver::outer_strategy::OuterEvalOrder,
    design_revision: Option<u64>,
) -> Result<
    (
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    ),
    EstimationError,
> {
    evaluator.evaluate_with_order(
        &design.design,
        &design.penalties,
        &design.nullspace_dims,
        design.linear_constraints.clone(),
        theta,
        rho_dim,
        hyper_dirs,
        warm_start_beta,
        "evaluate_joint_reml_outer_eval_at_theta",
        order,
        design_revision,
    )
}

fn evaluate_joint_reml_efs_at_theta(
    evaluator: &mut crate::estimate::ExternalJointHyperEvaluator<'_>,
    design: &TermCollectionDesign,
    theta: &Array1<f64>,
    rho_dim: usize,
    hyper_dirs: Vec<crate::estimate::reml::DirectionalHyperParam>,
    warm_start_beta: Option<ArrayView1<'_, f64>>,
    design_revision: Option<u64>,
) -> Result<crate::solver::outer_strategy::EfsEval, EstimationError> {
    evaluator.evaluate_efs(
        &design.design,
        &design.penalties,
        &design.nullspace_dims,
        design.linear_constraints.clone(),
        theta,
        rho_dim,
        hyper_dirs,
        warm_start_beta,
        "evaluate_joint_reml_efs_at_theta",
        design_revision,
    )
}

fn exact_joint_spatial_outer_hessian_available(
    family: &LikelihoodSpec,
    design: &TermCollectionDesign,
) -> bool {
    // Every `LikelihoodSpec` variant (Gaussian, Binomial-*, Poisson, Gamma,
    // Royston-Parmar) routes through the unified evaluator's outer-Hessian
    // path: Gaussian Identity uses the no-correction dense form, all GLM
    // variants supply scalar-GLM derivative ingredients consumed by
    // `compute_outer_hessian` / `build_outer_hessian_operator`, and the
    // (n, p, K) crossover in `prefer_outer_hessian_operator` chooses the
    // matrix-free `HessianResult::Operator` representation at biobank scale
    // for dense-lazy designs.  The previous `Identity || sparse_design`
    // gate predates that operator routing and forced binomial+logit+Matern
    // (and any other non-Gaussian dense-lazy spatial design) onto the
    // gradient-only BFGS path even though analytic Hessian is fully
    // available — capability check, not cost.  Match every variant
    // explicitly so any future family addition (which may not yet provide
    // outer-Hessian ingredients) forces an authoring decision here rather
    // than silently inheriting `true`.
    // Every supported response (Gaussian, Binomial-*, Poisson, Tweedie,
    // NegativeBinomial, Beta, Gamma, Royston-Parmar) routes through the
    // unified evaluator's outer-Hessian path; the spec-level capability
    // check therefore always succeeds. Match every response explicitly so
    // any future family addition (which may not yet provide outer-Hessian
    // ingredients) forces an authoring decision here rather than silently
    // inheriting `true`.
    let family_supported = match &family.response {
        ResponseFamily::Gaussian
        | ResponseFamily::Binomial
        | ResponseFamily::Poisson
        | ResponseFamily::Tweedie { .. }
        | ResponseFamily::NegativeBinomial { .. }
        | ResponseFamily::Beta { .. }
        | ResponseFamily::Gamma
        | ResponseFamily::RoystonParmar => true,
    };
    // A design with zero columns has no joint outer-Hessian to compute;
    // the analytic path is only meaningful for non-empty parameter blocks.
    family_supported && design.design.ncols() > 0
}

fn smooth_term_penalty_index(
    spec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    term_idx: usize,
) -> Option<usize> {
    if term_idx >= design.smooth.terms.len() || term_idx >= spec.smooth_terms.len() {
        return None;
    }
    if design.smooth.terms[term_idx].penalties_local.is_empty() {
        return None;
    }
    let linear_penalties = spec
        .linear_terms
        .iter()
        .filter(|t| t.double_penalty)
        .count()
        * 2;
    let random_penalties = design
        .random_effect_ranges
        .iter()
        .filter(|(_, range)| !range.is_empty())
        .count();
    let smooth_offset = linear_penalties + random_penalties;
    let local_offset = design
        .smooth
        .terms
        .iter()
        .take(term_idx)
        .map(|term| term.penalties_local.len())
        .sum::<usize>();
    Some(smooth_offset + local_offset)
}

fn try_build_spatial_term_log_kappa_derivativeinfo(
    data: ArrayView2<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    term_idx: usize,
) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
    let Some((
        global_range,
        total_p,
        x_psi_local,
        s_psi_local_check,
        x_psi_psi_local,
        s_psi_psi_local,
        s_psi_components_local,
        s_psi_psi_components_local,
        implicit_operator,
    )) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
    else {
        return Ok(None);
    };
    let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
        return Ok(None);
    };
    if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
        return Ok(None);
    }
    if s_psi_components_local.len() != s_psi_psi_components_local.len() {
        return Ok(None);
    }
    let penalty_indices = (0..s_psi_components_local.len())
        .map(|j| penalty_start + j)
        .collect::<Vec<_>>();
    let penalty_index = penalty_indices[0];
    if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
        return Ok(None);
    }
    Ok(Some(SpatialPsiDerivative {
        penalty_index,
        penalty_indices,
        global_range,
        total_p,
        x_psi_local,
        s_psi_components_local,
        x_psi_psi_local,
        s_psi_psi_components_local,
        aniso_group_id: None,
        aniso_cross_designs: None,
        aniso_cross_penalty_provider: None,
        implicit_operator,
        implicit_axis: 0,
    }))
}

pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
    data: ArrayView2<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    spatial_terms: &[usize],
) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
    let mut out = Vec::new();
    let mut aniso_gid = 0usize;
    for &term_idx in spatial_terms {
        if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
            if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
                data,
                resolvedspec,
                design,
                term_idx,
                aniso_gid,
            )? {
                aniso_gid += 1;
                out.extend(entries);
                continue;
            } else {
                return Ok(None);
            }
        }
        let Some(info) =
            try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
        else {
            return Ok(None);
        };
        out.push(info);
    }
    Ok(Some(out))
}

/// For an aniso term with d axes, produce d `SpatialPsiDerivative` entries.
fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
    data: ArrayView2<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    term_idx: usize,
    aniso_group_id: usize,
) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
    let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
        return Ok(None);
    };
    let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
        return Ok(None);
    };
    let mut aniso_result = match &termspec.basis {
        SmoothBasisSpec::Sphere { .. } => return Ok(None),
        SmoothBasisSpec::Matern {
            feature_cols,
            spec,
            input_scales,
        } => {
            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
            if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
            }
            build_matern_basis_log_kappa_aniso_derivatives(x.view(), spec)
                .map_err(EstimationError::from)?
        }
        _ => return Ok(None),
    };
    // Get number of axes from the shared operator when available; otherwise
    // fall back to the dense design list.
    let d = if let Some(ref op) = aniso_result.implicit_operator {
        op.n_axes()
    } else if !aniso_result.design_first.is_empty() {
        aniso_result.design_first.len()
    } else {
        0
    };
    if d == 0 {
        return Ok(None);
    }
    let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
        return Ok(None);
    };
    let p_total = design.design.ncols();
    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
    let global_range = (smooth_start + smooth_term.coeff_range.start)
        ..(smooth_start + smooth_term.coeff_range.end);
    let num_penalties = aniso_result.penalties_first[0].len();
    let penalty_indices: Vec<usize> = (0..num_penalties).map(|j| penalty_start + j).collect();
    let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();

    // Dense first/diagonal-second matrices may be present even when the shared
    // operator is available. The operator remains the canonical source for
    // exact cross-axis second derivatives.
    let use_implicit_design = aniso_result.design_first.is_empty();
    let implicit_op_arc = aniso_result
        .implicit_operator
        .as_ref()
        .map(|op| std::sync::Arc::new(op.clone()));

    let mut entries = Vec::with_capacity(d);
    for a in 0..d {
        let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
            // Implicit path: design-derivative matvecs will be dispatched through
            // the ImplicitDerivativeOp inside HyperDesignDerivative, so we do NOT
            // need to materialize the dense (n x p) matrices here.  Store empty
            // placeholders — they are never read when the implicit operator is
            // present (spatial_log_kappa_hyper_dirs_frominfo_list uses from_implicit).
            (Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
        } else {
            // Move the dense (n × p) matrices out of aniso_result instead of
            // cloning. Each axis index `a` is read exactly once across the
            // loop, and aniso_result is dropped at function exit, so leaving
            // empty placeholders behind in those vec slots is safe.
            let x_first = std::mem::take(&mut aniso_result.design_first[a]);
            let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
            if x_first.ncols() != smooth_term.coeff_range.len() {
                return Ok(None);
            }
            (x_first, x_second)
        };
        let s_psi_components = std::mem::take(&mut aniso_result.penalties_first[a]);
        let s_psi_psi_components = std::mem::take(&mut aniso_result.penalties_second_diag[a]);
        // Build cross-design entries for other axes b != a in this group.
        // These will be indexed by (b, cross_matrix) where b is the axis
        // offset within the d-entry block.
        // Cross-axis second derivatives are sourced from the shared operator,
        // so we only need placeholder entries to preserve the axis layout.
        let cross_designs = if implicit_op_arc.is_some() {
            let mut cd = Vec::with_capacity(d - 1);
            for b in 0..d {
                if b == a {
                    continue;
                }
                cd.push((b, Array2::<f64>::zeros((0, 0))));
            }
            cd
        } else if !aniso_result.design_second_cross.is_empty() {
            let mut cd = Vec::new();
            for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
            {
                if pa == a {
                    cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
                } else if pb == a {
                    cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
                }
            }
            cd
        } else {
            Vec::new()
        };
        let cross_penalty_provider = if d > 1 {
            let penalties_cross_provider = penalties_cross_provider.clone();
            Some(std::sync::Arc::new(
                move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
                    if b_axis == a {
                        return Ok(Vec::new());
                    }
                    let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
                    if let Some(provider) = penalties_cross_provider.as_ref() {
                        provider
                            .evaluate(axis_lo, axis_hi)
                            .map_err(EstimationError::from)
                    } else {
                        // No provider: either the pair is unregistered, or it
                        // was registered without data (early-return raw-operator
                        // paths). Both cases contribute no cross penalties.
                        Ok(Vec::new())
                    }
                },
            )
                as std::sync::Arc<
                    dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
                        + Send
                        + Sync
                        + 'static,
                >)
        } else {
            None
        };

        entries.push(SpatialPsiDerivative {
            penalty_index: penalty_indices[0],
            penalty_indices: penalty_indices.clone(),
            global_range: global_range.clone(),
            total_p: p_total,
            x_psi_local,
            s_psi_components_local: s_psi_components,
            x_psi_psi_local,
            s_psi_psi_components_local: s_psi_psi_components,
            aniso_group_id: Some(aniso_group_id),
            aniso_cross_designs: if cross_designs.is_empty() {
                None
            } else {
                Some(cross_designs)
            },
            aniso_cross_penalty_provider: cross_penalty_provider,
            implicit_operator: implicit_op_arc.clone(),
            implicit_axis: a,
        });
    }
    Ok(Some(entries))
}

fn try_build_spatial_term_log_kappa_derivative(
    data: ArrayView2<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    term_idx: usize,
) -> Result<
    Option<(
        Range<usize>,
        usize,
        Array2<f64>,
        Array2<f64>,
        Array2<f64>,
        Array2<f64>,
        Vec<Array2<f64>>,
        Vec<Array2<f64>>,
        Option<std::sync::Arc<crate::terms::basis::ImplicitDesignPsiDerivative>>,
    )>,
    EstimationError,
> {
    let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
        return Ok(None);
    };
    let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
        return Ok(None);
    };

    let derivative_bundle = match &termspec.basis {
        SmoothBasisSpec::ThinPlate {
            feature_cols,
            spec,
            input_scales,
        } => {
            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
            let mut spec_local = spec.clone();
            if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                spec_local.length_scale =
                    compensate_length_scale_for_standardization(spec.length_scale, s);
            }
            build_thin_plate_basis_log_kappa_derivatives(x.view(), &spec_local)
                .map_err(EstimationError::from)?
        }
        SmoothBasisSpec::Sphere { .. } => return Ok(None),
        SmoothBasisSpec::Matern {
            feature_cols,
            spec,
            input_scales,
        } => {
            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
            let mut spec_local = spec.clone();
            if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                spec_local.length_scale =
                    compensate_length_scale_for_standardization(spec.length_scale, s);
            }
            build_matern_basis_log_kappa_derivatives(x.view(), &spec_local)
                .map_err(EstimationError::from)?
        }
        SmoothBasisSpec::Duchon {
            feature_cols,
            spec,
            input_scales,
        } => {
            let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
            let mut spec_local = spec.clone();
            if let Some(s) = input_scales {
                apply_input_standardization(&mut x, s);
                spec_local.length_scale =
                    compensate_optional_length_scale_for_standardization(spec.length_scale, s);
            }
            let BasisMetadata::Duchon {
                centers,
                identifiability_transform,
                operator_collocation_points,
                ..
            } = &smooth_term.metadata
            else {
                return Ok(None);
            };
            crate::basis::build_duchon_basis_log_kappa_derivativeswith_collocationwithworkspace(
                x.view(),
                &spec_local,
                centers.view(),
                identifiability_transform.as_ref(),
                operator_collocation_points
                    .as_ref()
                    .map(|points| points.view()),
                &mut BasisWorkspace::default(),
            )
            .map_err(EstimationError::from)?
        }
        SmoothBasisSpec::BSpline1D { .. }
        | SmoothBasisSpec::TensorBSpline { .. }
        | SmoothBasisSpec::ByVariable { .. }
        | SmoothBasisSpec::FactorSumToZero { .. }
        | SmoothBasisSpec::BySmooth { .. }
        | SmoothBasisSpec::FactorSmooth { .. }
        | SmoothBasisSpec::Pca { .. } => {
            return Ok(None);
        }
    };
    let mut implicit_operator = derivative_bundle.implicit_operator;
    let BasisPsiDerivativeResult {
        design_derivative: mut local_x_psi,
        penalties_derivative: mut local_s_psi,
        implicit_operator: local_implicit_first_unused,
    } = derivative_bundle.first;
    let BasisPsiSecondDerivativeResult {
        designsecond_derivative: mut local_x_psi_psi,
        penaltiessecond_derivative: mut local_s_psi_psi,
        implicit_operator: local_implicit_second_unused,
    } = derivative_bundle.second;
    assert!(local_implicit_first_unused.is_none());
    assert!(local_implicit_second_unused.is_none());

    if let Some(rotation) = smooth_term.joint_null_rotation.as_ref() {
        let q = &rotation.rotation;
        if let Some(op) = implicit_operator.take() {
            implicit_operator = Some(op.append_full_transform(q).map_err(EstimationError::from)?);
        } else {
            if local_x_psi.ncols() != q.nrows() || local_x_psi_psi.ncols() != q.nrows() {
                return Ok(None);
            }
            local_x_psi = fast_ab(&local_x_psi, q);
            local_x_psi_psi = fast_ab(&local_x_psi_psi, q);
        }
        let rotate_penalty = |s_local: Array2<f64>| -> Option<Array2<f64>> {
            if s_local.nrows() != q.nrows() || s_local.ncols() != q.nrows() {
                return None;
            }
            let qt_s = crate::linalg::faer_ndarray::fast_atb(q, &s_local);
            Some(crate::linalg::faer_ndarray::fast_ab(&qt_s, q))
        };
        let Some(rotated_s_psi) = local_s_psi
            .into_iter()
            .map(|s| rotate_penalty(s))
            .collect::<Option<Vec<_>>>()
        else {
            return Ok(None);
        };
        local_s_psi = rotated_s_psi;
        let Some(rotated_s_psi_psi) = local_s_psi_psi
            .into_iter()
            .map(|s| rotate_penalty(s))
            .collect::<Option<Vec<_>>>()
        else {
            return Ok(None);
        };
        local_s_psi_psi = rotated_s_psi_psi;
    }
    let implicit_operator = implicit_operator.map(std::sync::Arc::new);

    if let Some(ref op) = implicit_operator {
        if op.p_out() != smooth_term.coeff_range.len() {
            return Ok(None);
        }
    } else {
        if local_x_psi.ncols() != smooth_term.coeff_range.len() {
            return Ok(None);
        }
        if local_x_psi_psi.ncols() != smooth_term.coeff_range.len() {
            return Ok(None);
        }
    }
    if local_s_psi.is_empty() || local_s_psi.len() != local_s_psi_psi.len() {
        return Ok(None);
    }
    if local_s_psi.iter().any(|s| {
        s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
    }) {
        return Ok(None);
    }
    if local_s_psi_psi.iter().any(|s| {
        s.nrows() != smooth_term.coeff_range.len() || s.ncols() != smooth_term.coeff_range.len()
    }) {
        return Ok(None);
    }

    let p_total = design.design.ncols();
    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
    let global_range = (smooth_start + smooth_term.coeff_range.start)
        ..(smooth_start + smooth_term.coeff_range.end);

    Ok(Some((
        global_range,
        p_total,
        local_x_psi,
        local_s_psi.iter().fold(
            Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
            |acc, m| acc + m,
        ),
        local_x_psi_psi,
        local_s_psi_psi.iter().fold(
            Array2::<f64>::zeros((smooth_term.coeff_range.len(), smooth_term.coeff_range.len())),
            |acc, m| acc + m,
        ),
        local_s_psi,
        local_s_psi_psi,
        implicit_operator,
    )))
}

fn try_build_spatial_log_kappa_hyper_dirs(
    data: ArrayView2<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    spatial_terms: &[usize],
) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
    // Each spatial term contributes one continuous scale hyperparameter
    //   psi = log(kappa) = -log(length_scale),
    // while rho = log(lambda) still indexes the smoothing parameters of the
    // three operator penalties. The joint outer vector is therefore
    //   theta = (rho_0, ..., rho_{K-1}, psi_1, ..., psi_q)
    // for q spatial terms participating in exact joint optimization.
    let Some(info_list) =
        try_build_spatial_log_kappa_derivativeinfo_list(data, resolvedspec, design, spatial_terms)?
    else {
        return Ok(None);
    };
    Ok(Some(spatial_log_kappa_hyper_dirs_frominfo_list(info_list)?))
}

pub(crate) fn try_build_latent_coord_hyper_dirs(
    latent: std::sync::Arc<crate::terms::latent_coord::LatentCoordValues>,
    resolvedspec: &TermCollectionSpec,
    design: &TermCollectionDesign,
    latent_terms: &[crate::types::SmoothTermIdx],
    analytic_rho_count: usize,
) -> Result<Option<Vec<DirectionalHyperParam>>, EstimationError> {
    if latent_terms.is_empty() || latent.is_empty() {
        return Ok(None);
    }
    if latent_terms.len() != 1 {
        crate::bail_invalid_estim!(
            "LatentCoord standard-fit hyper_dirs currently require exactly one latent smooth term"
                .to_string(),
        );
    }
    let term_idx = latent_terms[0];
    let smooth_term = design.smooth.terms.get(term_idx.get()).ok_or_else(|| {
        EstimationError::InvalidInput(format!(
            "LatentCoord term index {term_idx} out of bounds for realized smooth design"
        ))
    })?;
    let termspec = resolvedspec
        .smooth_terms
        .get(term_idx.get())
        .ok_or_else(|| {
            EstimationError::InvalidInput(format!(
                "LatentCoord term index {term_idx} out of bounds for resolved smooth spec"
            ))
        })?;
    let p_total = design.design.ncols();
    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
    let global_range = (smooth_start + smooth_term.coeff_range.start)
        ..(smooth_start + smooth_term.coeff_range.end);

    // Spline bases do not add a separate continuous basis-scale ψ coordinate
    // here. When they are latent-coordinate terms, their ψ directions are the
    // latent-coordinate axes below, using the same DirectionalHyperParam layout
    // as Matérn and Duchon.
    let operator = match (&termspec.basis, &smooth_term.metadata) {
        (
            SmoothBasisSpec::Matern { .. },
            BasisMetadata::Matern {
                centers,
                length_scale,
                nu,
                include_intercept,
                identifiability_transform,
                ..
            },
        ) => crate::terms::basis::LatentCoordDesignDerivative::new_matern(
            latent.clone(),
            std::sync::Arc::new(centers.clone()),
            *length_scale,
            *nu,
            *include_intercept,
            identifiability_transform.clone(),
        )
        .map_err(EstimationError::from)?,
        (
            SmoothBasisSpec::Duchon { .. },
            BasisMetadata::Duchon {
                centers,
                length_scale,
                power,
                nullspace_order,
                identifiability_transform,
                ..
            },
        ) => crate::terms::basis::LatentCoordDesignDerivative::new_duchon(
            latent.clone(),
            std::sync::Arc::new(centers.clone()),
            *length_scale,
            *power,
            *nullspace_order,
            identifiability_transform.clone(),
        )
        .map_err(EstimationError::from)?,
        (
            SmoothBasisSpec::Sphere { .. },
            BasisMetadata::Sphere {
                centers,
                penalty_order,
                method,
                constraint_transform,
                ..
            },
        ) if matches!(*method, crate::basis::SphereMethod::Wahba) => {
            crate::terms::basis::LatentCoordDesignDerivative::new_sphere(
                latent.clone(),
                std::sync::Arc::new(centers.clone()),
                *penalty_order,
                constraint_transform.clone(),
            )
            .map_err(EstimationError::from)?
        }
        (
            SmoothBasisSpec::BSpline1D { spec, .. },
            BasisMetadata::BSpline1D {
                knots,
                identifiability_transform,
                periodic,
                degree: meta_degree,
                ..
            },
        ) => {
            // Issue #340: use the metadata-recorded effective degree so the
            // latent-design Jacobian matches what `build_bspline_basis_1d`
            // actually built at fit time after auto-shrink.
            let effective_degree = meta_degree.unwrap_or(spec.degree);
            if let Some((domain_start, period, num_basis)) = periodic {
                crate::terms::basis::LatentCoordDesignDerivative::new_periodic_bspline(
                    latent.clone(),
                    (*domain_start, *domain_start + *period),
                    effective_degree,
                    *num_basis,
                    identifiability_transform.clone(),
                )
                .map_err(EstimationError::from)?
            } else {
                crate::terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
                    latent.clone(),
                    vec![knots.clone()],
                    vec![effective_degree],
                    identifiability_transform.clone(),
                )
                .map_err(EstimationError::from)?
            }
        }
        (
            SmoothBasisSpec::TensorBSpline { .. },
            BasisMetadata::TensorBSpline {
                knots,
                degrees,
                identifiability_transform,
                ..
            },
        ) => crate::terms::basis::LatentCoordDesignDerivative::new_tensor_bspline(
            latent.clone(),
            knots.clone(),
            degrees.clone(),
            identifiability_transform.clone(),
        )
        .map_err(EstimationError::from)?,
        (SmoothBasisSpec::Pca { .. }, BasisMetadata::Pca { basis_matrix, .. }) => {
            crate::terms::basis::LatentCoordDesignDerivative::new_pca(
                latent.clone(),
                std::sync::Arc::new(basis_matrix.clone()),
            )
            .map_err(EstimationError::from)?
        }
        _ => return Ok(None),
    };
    if operator.p_out() != global_range.len() {
        crate::bail_invalid_estim!(
            "LatentCoord derivative width mismatch for term '{}': operator p={}, coeff range={}",
            smooth_term.name,
            operator.p_out(),
            global_range.len()
        );
    }
    let operator = std::sync::Arc::new(operator);
    let mut hyper_dirs = Vec::with_capacity(operator.n_axes());
    for flat_axis in 0..operator.n_axes() {
        let dir = DirectionalHyperParam::new_compact(
            crate::estimate::reml::HyperDesignDerivative::from_latent_coord(
                operator.clone(),
                flat_axis,
                global_range.clone(),
                p_total,
            ),
            Vec::new(),
            None,
            None,
        )?
        .not_penalty_like();
        hyper_dirs.push(dir);
    }
    let direct_dim = latent_coord_direct_hyper_count(latent.id_mode(), latent.latent_dim());
    if analytic_rho_count + direct_dim > 0 {
        let zero_x = crate::estimate::reml::HyperDesignDerivative::from(Array2::<f64>::zeros((
            design.design.nrows(),
            p_total,
        )));
        for _ in 0..analytic_rho_count {
            hyper_dirs.push(
                DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
                    .not_penalty_like(),
            );
        }
        for _ in 0..direct_dim {
            hyper_dirs.push(
                DirectionalHyperParam::new_compact(zero_x.clone(), Vec::new(), None, None)?
                    .not_penalty_like(),
            );
        }
    }
    Ok(Some(hyper_dirs))
}

fn latent_coord_direct_hyper_count(
    id_mode: &crate::terms::latent_coord::LatentIdMode,
    latent_dim: usize,
) -> usize {
    use crate::terms::latent_coord::{AuxPriorStrength, LatentIdMode};
    match id_mode {
        LatentIdMode::AuxPrior { strength, .. } => match strength {
            AuxPriorStrength::Auto => 1,
            AuxPriorStrength::Fixed(_) => 0,
        },
        LatentIdMode::AuxPriorDimSelection { strength, .. } => {
            latent_dim
                + match strength {
                    AuxPriorStrength::Auto => 1,
                    AuxPriorStrength::Fixed(_) => 0,
                }
        }
        LatentIdMode::DimSelection { .. } => latent_dim,
        LatentIdMode::None => 0,
    }
}

fn latent_coord_initial_direct_hypers(
    id_mode: &crate::terms::latent_coord::LatentIdMode,
    latent_dim: usize,
) -> Result<Array1<f64>, EstimationError> {
    use crate::terms::latent_coord::{AuxPriorStrength, LatentIdMode};
    let mut values = Vec::with_capacity(latent_coord_direct_hyper_count(id_mode, latent_dim));
    match id_mode {
        LatentIdMode::AuxPrior { strength, .. } => {
            if matches!(strength, AuxPriorStrength::Auto) {
                values.push(0.0);
            }
        }
        LatentIdMode::AuxPriorDimSelection {
            strength,
            init_log_precision,
            ..
        } => {
            if matches!(strength, AuxPriorStrength::Auto) {
                values.push(0.0);
            }
            append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
        }
        LatentIdMode::DimSelection { init_log_precision } => {
            append_latent_ard_seed(&mut values, init_log_precision.as_ref(), latent_dim)?;
        }
        LatentIdMode::None => {}
    }
    Ok(Array1::from_vec(values))
}

fn append_latent_ard_seed(
    values: &mut Vec<f64>,
    init: Option<&Array1<f64>>,
    latent_dim: usize,
) -> Result<(), EstimationError> {
    if let Some(init) = init {
        if init.len() != latent_dim {
            crate::bail_invalid_estim!(
                "latent dim_selection init_log_precision length mismatch: got {}, expected {}",
                init.len(),
                latent_dim
            );
        }
        values.extend(init.iter().copied());
    } else {
        values.extend(std::iter::repeat_n(0.0, latent_dim));
    }
    Ok(())
}

struct LatentIdObjectiveContribution {
    cost: f64,
    gradient: Array1<f64>,
}

fn latent_id_objective_contribution(
    theta: &Array1<f64>,
    rho_dim: usize,
    analytic_rho_count: usize,
    latent: &crate::terms::latent_coord::LatentCoordValues,
) -> Result<LatentIdObjectiveContribution, EstimationError> {
    use crate::terms::latent_coord::{AuxPriorStrength, LatentIdMode, aux_prior_targets};
    let n_obs = latent.n_obs();
    let latent_dim = latent.latent_dim();
    let flat_len = latent.len();
    let mut gradient = Array1::<f64>::zeros(theta.len());
    let t_start = rho_dim;
    let direct_start = t_start + flat_len + analytic_rho_count;
    if theta.len() < direct_start {
        crate::bail_invalid_estim!(
            "latent-coordinate theta too short for id objective: got {}, need at least {}",
            theta.len(),
            direct_start
        );
    }
    let t = latent.as_matrix();
    let mut cost = 0.0;
    let mut cursor = direct_start;

    match latent.id_mode() {
        LatentIdMode::AuxPrior {
            u,
            family,
            strength,
        }
        | LatentIdMode::AuxPriorDimSelection {
            u,
            family,
            strength,
            ..
        } => {
            let (log_mu, mu) = match strength {
                AuxPriorStrength::Fixed(mu) => (mu.ln(), *mu),
                AuxPriorStrength::Auto => {
                    let log_mu = theta[cursor];
                    cursor += 1;
                    (log_mu, log_mu.exp())
                }
            };
            let targets = aux_prior_targets(t.view(), u.view(), *family)
                .map_err(EstimationError::InvalidInput)?;
            let residual = &t - &targets;
            let q = residual.iter().map(|v| v * v).sum::<f64>();
            cost += 0.5 * mu * q - 0.5 * n_obs as f64 * log_mu;

            let projected_residual = aux_prior_targets(residual.view(), u.view(), *family)
                .map_err(EstimationError::InvalidInput)?;
            let grad_base = residual - projected_residual;
            for n in 0..n_obs {
                for axis in 0..latent_dim {
                    gradient[t_start + n * latent_dim + axis] += mu * grad_base[[n, axis]];
                }
            }
            if matches!(strength, AuxPriorStrength::Auto) {
                gradient[direct_start] += 0.5 * mu * q - 0.5 * n_obs as f64;
            }
        }
        LatentIdMode::DimSelection { .. } | LatentIdMode::None => {}
    }

    match latent.id_mode() {
        LatentIdMode::AuxPriorDimSelection { .. } | LatentIdMode::DimSelection { .. } => {
            for axis in 0..latent_dim {
                let log_alpha = theta[cursor + axis];
                let alpha = log_alpha.exp();
                let mut q_axis = 0.0;
                for n in 0..n_obs {
                    let flat_idx = n * latent_dim + axis;
                    let value = latent.as_flat()[flat_idx];
                    q_axis += value * value;
                    gradient[t_start + flat_idx] += alpha * value;
                }
                cost += 0.5 * alpha * q_axis - 0.5 * n_obs as f64 * log_alpha;
                gradient[cursor + axis] += 0.5 * alpha * q_axis - 0.5 * n_obs as f64;
            }
            cursor += latent_dim;
        }
        LatentIdMode::AuxPrior { .. } | LatentIdMode::None => {}
    }

    if cursor != theta.len() {
        crate::bail_invalid_estim!(
            "latent-coordinate direct hyperparameter length mismatch: consumed {}, theta len {}",
            cursor,
            theta.len()
        );
    }
    Ok(LatentIdObjectiveContribution { cost, gradient })
}

fn add_latent_id_objective_to_eval(
    theta: &Array1<f64>,
    rho_dim: usize,
    analytic_rho_count: usize,
    latent: &crate::terms::latent_coord::LatentCoordValues,
    eval: &mut (
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    ),
) -> Result<(), EstimationError> {
    let contribution =
        latent_id_objective_contribution(theta, rho_dim, analytic_rho_count, latent)?;
    eval.0 += contribution.cost;
    if eval.1.len() != contribution.gradient.len() {
        crate::bail_invalid_estim!(
            "latent-coordinate REML gradient length mismatch: base={}, id={}",
            eval.1.len(),
            contribution.gradient.len()
        );
    }
    eval.1 += &contribution.gradient;
    if eval.2.is_analytic() {
        eval.2 = crate::solver::outer_strategy::HessianResult::Unavailable;
    }
    Ok(())
}

fn analytic_penalty_objective_contribution(
    theta: &Array1<f64>,
    rho_dim: usize,
    latent: &crate::terms::latent_coord::LatentCoordValues,
    registry: &crate::terms::AnalyticPenaltyRegistry,
) -> Result<LatentIdObjectiveContribution, EstimationError> {
    let flat_len = latent.len();
    let t_start = rho_dim;
    let t_end = t_start + flat_len;
    let rho_start = t_end;
    let rho_end = rho_start + registry.total_rho_count();
    if theta.len() < rho_end {
        crate::bail_invalid_estim!(
            "latent-coordinate theta too short for analytic penalties: got {}, need at least {}",
            theta.len(),
            rho_end
        );
    }
    let target_t = theta.slice(s![t_start..t_end]);
    let rho = theta.slice(s![rho_start..rho_end]);
    let mut cost = 0.0_f64;
    let mut gradient = Array1::<f64>::zeros(theta.len());
    for (penalty, (rho_slice, tier, name)) in registry.penalties.iter().zip(registry.rho_layout()) {
        let rho_local = rho.slice(s![rho_slice.clone()]);
        match tier {
            crate::terms::PenaltyTier::Psi => {
                cost += penalty.value(target_t.view(), rho_local);
                let grad = penalty.grad_target(target_t.view(), rho_local);
                if grad.len() != flat_len {
                    crate::bail_invalid_estim!(
                        "analytic penalty {name:?} gradient length mismatch: got {}, expected {}",
                        grad.len(),
                        flat_len
                    );
                }
                for i in 0..flat_len {
                    gradient[t_start + i] += grad[i];
                }
                let grad_rho_local = penalty.grad_rho(target_t.view(), rho_local);
                if grad_rho_local.len() != rho_slice.len() {
                    crate::bail_invalid_estim!(
                        "analytic penalty {name:?} rho-gradient length mismatch: got {}, expected {}",
                        grad_rho_local.len(),
                        rho_slice.len()
                    );
                }
                for local_idx in 0..grad_rho_local.len() {
                    gradient[rho_start + rho_slice.start + local_idx] += grad_rho_local[local_idx];
                }
            }
            crate::terms::PenaltyTier::Beta => {}
            crate::terms::PenaltyTier::Rho => {}
        }
    }
    Ok(LatentIdObjectiveContribution { cost, gradient })
}

fn add_analytic_penalty_hessian_to_eval(
    theta: &Array1<f64>,
    rho_dim: usize,
    latent: &crate::terms::latent_coord::LatentCoordValues,
    registry: &crate::terms::AnalyticPenaltyRegistry,
    eval: &mut (
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    ),
) -> Result<(), EstimationError> {
    let flat_len = latent.len();
    let t_start = rho_dim;
    let t_end = t_start + flat_len;
    let rho_start = t_end;
    let rho_end = rho_start + registry.total_rho_count();
    if theta.len() < rho_end {
        crate::bail_invalid_estim!(
            "latent-coordinate theta too short for analytic penalty Hessian: got {}, need at least {}",
            theta.len(),
            rho_end
        );
    }
    let crate::solver::outer_strategy::HessianResult::Analytic(hessian) = &mut eval.2 else {
        if eval.2.is_analytic() {
            eval.2 = crate::solver::outer_strategy::HessianResult::Unavailable;
        }
        return Ok(());
    };
    if hessian.dim() != (theta.len(), theta.len()) {
        crate::bail_invalid_estim!(
            "analytic penalty Hessian target shape mismatch: got {}x{}, expected {}x{}",
            hessian.nrows(),
            hessian.ncols(),
            theta.len(),
            theta.len()
        );
    }
    let target_t = theta.slice(s![t_start..t_end]);
    let rho = theta.slice(s![rho_start..rho_end]);
    for (penalty, (rho_slice, tier, _name)) in registry.penalties.iter().zip(registry.rho_layout())
    {
        let rho_local = rho.slice(s![rho_slice]);
        if !matches!(tier, crate::terms::PenaltyTier::Psi) {
            continue;
        }
        if let Some(diag) = penalty.hessian_diag(target_t.view(), rho_local) {
            if diag.len() != flat_len {
                crate::bail_invalid_estim!(
                    "analytic penalty Hessian diagonal length mismatch: got {}, expected {}",
                    diag.len(),
                    flat_len
                );
            }
            for i in 0..flat_len {
                hessian[[t_start + i, t_start + i]] += diag[i];
            }
            continue;
        }
        let mut probe = Array1::<f64>::zeros(flat_len);
        for col in 0..flat_len {
            probe[col] = 1.0;
            let hv = penalty.hvp(target_t.view(), rho_local, probe.view());
            if hv.len() != flat_len {
                crate::bail_invalid_estim!(
                    "analytic penalty Hessian-vector length mismatch: got {}, expected {}",
                    hv.len(),
                    flat_len
                );
            }
            for row in 0..flat_len {
                hessian[[t_start + row, t_start + col]] += hv[row];
            }
            probe[col] = 0.0;
        }
    }
    Ok(())
}

fn add_analytic_penalty_objective_to_eval(
    theta: &Array1<f64>,
    rho_dim: usize,
    latent: &crate::terms::latent_coord::LatentCoordValues,
    registry: &crate::terms::AnalyticPenaltyRegistry,
    eval: &mut (
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    ),
) -> Result<(), EstimationError> {
    let contribution = analytic_penalty_objective_contribution(theta, rho_dim, latent, registry)?;
    eval.0 += contribution.cost;
    if eval.1.len() != contribution.gradient.len() {
        crate::bail_invalid_estim!(
            "latent-coordinate REML gradient length mismatch: base={}, analytic_penalty={}",
            eval.1.len(),
            contribution.gradient.len()
        );
    }
    eval.1 += &contribution.gradient;
    add_analytic_penalty_hessian_to_eval(theta, rho_dim, latent, registry, eval)?;
    Ok(())
}

fn spatial_log_kappa_hyper_dirs_frominfo_list(
    info_list: Vec<SpatialPsiDerivative>,
) -> Result<Vec<DirectionalHyperParam>, EstimationError> {
    use crate::estimate::reml::ImplicitDerivLevel;
    use std::collections::HashMap;

    let log_kappa_dim = info_list.len();
    // Layout-only metadata (group_id per axis) is cheap to snapshot up front so
    // the consumption loop below can MOVE the dense (n × p) derivative arrays
    // out of each entry instead of cloning. At biobank scale (n≈3×10⁵, 16-axis
    // CTN) the prior `.clone()` sites doubled peak working memory for the
    // psi-derivative pass through several GiB.
    let group_ids: Vec<Option<usize>> = info_list.iter().map(|e| e.aniso_group_id).collect();
    let mut group_indices_map: HashMap<usize, Vec<usize>> = HashMap::new();
    for (idx, gid) in group_ids.iter().enumerate() {
        if let Some(g) = gid {
            group_indices_map.entry(*g).or_default().push(idx);
        }
    }

    let mut hyper_dirs = Vec::with_capacity(log_kappa_dim);
    for (i, info) in info_list.into_iter().enumerate() {
        let SpatialPsiDerivative {
            penalty_index: _,
            penalty_indices,
            global_range,
            total_p,
            x_psi_local,
            s_psi_components_local,
            x_psi_psi_local,
            s_psi_psi_components_local,
            aniso_group_id,
            aniso_cross_designs,
            aniso_cross_penalty_provider,
            implicit_operator,
            implicit_axis,
        } = info;

        let mut xsecond = vec![None; log_kappa_dim];
        // Diagonal second derivative (same axis).
        xsecond[i] = Some(if let Some(ref op) = implicit_operator {
            crate::estimate::reml::HyperDesignDerivative::from_implicit(
                op.clone(),
                ImplicitDerivLevel::SecondDiag(implicit_axis),
                global_range.clone(),
                total_p,
            )
        } else {
            crate::estimate::reml::HyperDesignDerivative::from_embedded(
                x_psi_psi_local,
                global_range.clone(),
                total_p,
            )
        });
        // Cross second derivatives for axes in the same aniso group.
        if let Some(cross_designs) = aniso_cross_designs {
            // Use the base index of this aniso group in the original info_list.
            // Entries for the same group are contiguous: the first index in the
            // group gives the base, and axis b is at base+b.
            if let Some(gid) = aniso_group_id {
                let base = group_indices_map
                    .get(&gid)
                    .and_then(|v| v.first().copied())
                    .unwrap_or(i);
                for (b_axis, cross_mat) in cross_designs.into_iter() {
                    let j = base + b_axis;
                    if j < log_kappa_dim {
                        xsecond[j] = Some(if let Some(ref op) = implicit_operator {
                            crate::estimate::reml::HyperDesignDerivative::from_implicit(
                                op.clone(),
                                ImplicitDerivLevel::SecondCross(implicit_axis, b_axis),
                                global_range.clone(),
                                total_p,
                            )
                        } else {
                            crate::estimate::reml::HyperDesignDerivative::from_embedded(
                                cross_mat,
                                global_range.clone(),
                                total_p,
                            )
                        });
                    }
                }
            }
        }
        let s_components = penalty_indices
            .iter()
            .copied()
            .zip(s_psi_components_local.into_iter().map(|local| {
                crate::estimate::reml::HyperPenaltyDerivative::from_embedded(
                    local,
                    global_range.clone(),
                    total_p,
                )
            }))
            .collect::<Vec<_>>();
        let s2_components = penalty_indices
            .iter()
            .copied()
            .zip(s_psi_psi_components_local.into_iter().map(|local| {
                crate::estimate::reml::HyperPenaltyDerivative::from_embedded(
                    local,
                    global_range.clone(),
                    total_p,
                )
            }))
            .collect::<Vec<_>>();
        let mut ssecond_components = vec![None; log_kappa_dim];
        ssecond_components[i] = Some(s2_components);
        let mut penaltysecond_partner_indices: Option<Vec<usize>> = None;
        let penaltysecond_component_provider =
            if let (Some(provider), Some(gid)) = (aniso_cross_penalty_provider, aniso_group_id) {
                let group_indices = group_indices_map.get(&gid).cloned().unwrap_or_default();
                let axis_in_group =
                    group_indices
                        .iter()
                        .position(|&idx| idx == i)
                        .ok_or_else(|| {
                            EstimationError::InvalidInput(format!(
                                "missing spatial hyper axis {} in anisotropy group {}",
                                i, gid
                            ))
                        })?;
                penaltysecond_partner_indices = Some(
                    group_indices
                        .iter()
                        .copied()
                        .filter(|&idx| idx != i)
                        .collect(),
                );
                let penalty_indices_inner = penalty_indices.clone();
                let global_range_inner = global_range.clone();
                let total_p_inner = total_p;
                let group_indices_inner = group_indices;
                Some(std::sync::Arc::new(
                    move |j: usize| -> Result<
                        Option<Vec<crate::estimate::reml::PenaltyDerivativeComponent>>,
                        EstimationError,
                    > {
                        let Some(other_axis_in_group) =
                            group_indices_inner.iter().position(|&idx| idx == j)
                        else {
                            return Ok(None);
                        };
                        if other_axis_in_group == axis_in_group {
                            return Ok(None);
                        }
                        let cross_pens = provider(other_axis_in_group)?;
                        if cross_pens.is_empty() {
                            return Ok(None);
                        }
                        Ok(Some(
                            penalty_indices_inner
                                .iter()
                                .copied()
                                .zip(cross_pens.into_iter().map(|local| {
                                    crate::estimate::reml::HyperPenaltyDerivative::from_embedded(
                                        local,
                                        global_range_inner.clone(),
                                        total_p_inner,
                                    )
                                }))
                                .map(|(penalty_index, matrix)| {
                                    crate::estimate::reml::PenaltyDerivativeComponent {
                                        penalty_index,
                                        matrix,
                                    }
                                })
                                .collect(),
                        ))
                    },
                )
                    as std::sync::Arc<
                        dyn Fn(
                                usize,
                            ) -> Result<
                                Option<Vec<crate::estimate::reml::PenaltyDerivativeComponent>>,
                                EstimationError,
                            > + Send
                            + Sync
                            + 'static,
                    >)
            } else {
                None
            };
        // First derivative: use implicit operator when available to avoid
        // storing dense (n x p) matrices for all D axes simultaneously.
        let x_first_hyper = if let Some(ref op) = implicit_operator {
            crate::estimate::reml::HyperDesignDerivative::from_implicit(
                op.clone(),
                ImplicitDerivLevel::First(implicit_axis),
                global_range.clone(),
                total_p,
            )
        } else {
            crate::estimate::reml::HyperDesignDerivative::from_embedded(
                x_psi_local,
                global_range.clone(),
                total_p,
            )
        };
        let mut dir = DirectionalHyperParam::new_compact(
            x_first_hyper,
            s_components,
            Some(xsecond),
            Some(ssecond_components),
        )?
        .not_penalty_like();
        if let Some(provider) = penaltysecond_component_provider {
            dir = dir.with_penaltysecond_component_provider(provider);
        }
        if let Some(partner_indices) = penaltysecond_partner_indices {
            dir = dir.with_penaltysecond_partner_indices(partner_indices);
        }
        hyper_dirs.push(dir);
    }
    Ok(hyper_dirs)
}

/// Whether a spatial term contributes per-axis ψ entries to the outer joint
/// hyperparameter vector.
///
/// A term enrolls per-axis log-κ ψ when (a) it carries `aniso_log_scales`
/// matching its feature dimension, (b) its dimension is `> 1`, and (c) its
/// basis supports REML-side anisotropic κ optimization. Duchon is explicitly
/// excluded: its anisotropy η is a fixed, geometry-derived basis parameter,
/// never a REML hyper axis. A hybrid Duchon (explicit κ) still optimizes its
/// scalar length scale through the regular isotropic-κ path, so it
/// contributes a single ψ entry, not one per axis.
///
/// This predicate is the *single source of truth* for the joint-θ layout:
/// `has_aniso_terms`, `spatial_dims_per_term`, `from_length_scales_aniso`,
/// and the hyper_dirs builder `try_build_spatial_log_kappa_derivativeinfo_list`
/// all consult it, so the outer optimizer's `n_params = rho_dim + Σ ψ_per_term`
/// always matches the gradient length produced by the inner unified evaluator.
fn spatial_term_uses_per_axis_psi(resolvedspec: &TermCollectionSpec, term_idx: usize) -> bool {
    let Some(d) = get_spatial_feature_dim(resolvedspec, term_idx) else {
        return false;
    };
    if d <= 1 {
        return false;
    }
    let Some(eta) = get_spatial_aniso_log_scales(resolvedspec, term_idx) else {
        return false;
    };
    if eta.len() != d {
        return false;
    }
    !matches!(
        resolvedspec.smooth_terms.get(term_idx).map(|t| &t.basis),
        Some(SmoothBasisSpec::Duchon { .. })
    )
}

/// Compute `dims_per_term` for a list of spatial term indices.
///
/// Returns a vector where entry i is the number of stored ψ values for
/// spatial term i: `d` for terms that enroll per-axis anisotropy in the
/// REML joint vector (`spatial_term_uses_per_axis_psi`), `1` otherwise.
pub(crate) fn spatial_dims_per_term(
    resolvedspec: &TermCollectionSpec,
    spatial_terms: &[usize],
) -> Vec<usize> {
    spatial_terms
        .iter()
        .map(|&term_idx| {
            if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
                get_spatial_feature_dim(resolvedspec, term_idx).unwrap_or(1)
            } else {
                1
            }
        })
        .collect()
}

/// Check whether any spatial terms enroll per-axis anisotropic ψ in the joint
/// outer vector. Mirrors the hyper_dirs builder's enrollment predicate so the
/// outer θ-layout cannot drift from the inner evaluator's ψ count.
fn has_aniso_terms(resolvedspec: &TermCollectionSpec, spatial_terms: &[usize]) -> bool {
    spatial_terms
        .iter()
        .any(|&term_idx| spatial_term_uses_per_axis_psi(resolvedspec, term_idx))
}

/// Emits the `theta`-keyed memoization accessors shared verbatim by the
/// single-block and n-block exact-joint design caches. Both carry the same
/// `current_theta` / `last_cost` / `last_eval` fields, so the cost/eval
/// lookups and the `store_eval` writer are identical; this macro is the single
/// source so the two inherent impls cannot drift.
macro_rules! impl_exact_joint_theta_memo {
    () => {
        fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
            if self
                .current_theta
                .as_ref()
                .is_some_and(|cached| theta_values_match(cached, theta))
            {
                self.last_eval
                    .as_ref()
                    .map(|cached| cached.0)
                    .or(self.last_cost)
            } else {
                None
            }
        }

        fn memoized_eval(
            &self,
            theta: &Array1<f64>,
        ) -> Option<(
            f64,
            Array1<f64>,
            crate::solver::outer_strategy::HessianResult,
        )> {
            if self
                .current_theta
                .as_ref()
                .is_some_and(|cached| theta_values_match(cached, theta))
            {
                self.last_eval.clone()
            } else {
                None
            }
        }

        fn store_eval(
            &mut self,
            eval: (
                f64,
                Array1<f64>,
                crate::solver::outer_strategy::HessianResult,
            ),
        ) {
            self.last_cost = Some(eval.0);
            self.last_eval = Some(eval);
        }
    };
}

#[derive(Debug)]
struct SingleBlockExactJointDesignCache<'d> {
    realizer: FrozenTermCollectionIncrementalRealizer<'d>,
    current_theta: Option<Array1<f64>>,
    last_cost: Option<f64>,
    last_eval: Option<(
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    )>,
    spatial_terms: Vec<usize>,
    rho_dim: usize,
    dims_per_term: Vec<usize>,
}

impl<'d> SingleBlockExactJointDesignCache<'d> {
    fn new(
        data: ArrayView2<'d, f64>,
        spec: TermCollectionSpec,
        design: TermCollectionDesign,
        spatial_terms: Vec<usize>,
        rho_dim: usize,
        dims_per_term: Vec<usize>,
    ) -> Result<Self, String> {
        Ok(Self {
            realizer: FrozenTermCollectionIncrementalRealizer::new(data, spec, design)?,
            current_theta: None,
            last_cost: None,
            last_eval: None,
            spatial_terms,
            rho_dim,
            dims_per_term,
        })
    }

    fn design_revision(&self) -> u64 {
        self.realizer.design_revision()
    }

    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
        {
            return Ok(());
        }
        let t_ensure = std::time::Instant::now();
        let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
            theta,
            self.rho_dim,
            self.dims_per_term.clone(),
        );
        self.realizer
            .apply_log_kappa(&log_kappa, &self.spatial_terms)?;
        log::info!(
            "[STAGE] ensure_theta (apply_log_kappa, {} terms): {:.3}s",
            self.spatial_terms.len(),
            t_ensure.elapsed().as_secs_f64(),
        );
        self.current_theta = Some(theta.clone());
        self.last_cost = None;
        self.last_eval = None;
        Ok(())
    }

    impl_exact_joint_theta_memo!();

    fn store_cost(&mut self, cost: f64) {
        self.last_cost = Some(cost);
    }

    fn spec(&self) -> &TermCollectionSpec {
        self.realizer.spec()
    }

    fn design(&self) -> &TermCollectionDesign {
        self.realizer.design()
    }
}

struct SingleBlockLatentCoordDesignCache {
    data: Array2<f64>,
    spec: TermCollectionSpec,
    design: TermCollectionDesign,
    current_theta: Option<Array1<f64>>,
    current_latent: Option<std::sync::Arc<crate::terms::latent_coord::LatentCoordValues>>,
    current_hyper_dirs: Option<Vec<crate::estimate::reml::DirectionalHyperParam>>,
    current_design_cache_id: Option<u64>,
    latent_design_cache: crate::solver::latent_cache::LatentDesignCache,
    last_cost: Option<f64>,
    last_eval: Option<(
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    )>,
    term_index: crate::types::SmoothTermIdx,
    feature_cols: Vec<usize>,
    rho_dim: usize,
    n_obs: usize,
    latent_dim: usize,
    id_mode: crate::terms::latent_coord::LatentIdMode,
    manifold: crate::terms::latent_coord::LatentManifold,
    retraction_registry: crate::solver::latent_cache::LatentRetractionRegistry,
    latent_id: u64,
    analytic_penalties: Option<std::sync::Arc<crate::terms::AnalyticPenaltyRegistry>>,
    analytic_rho_count: usize,
    design_revision: u64,
    // Stamp the outer-iter the cached cost/eval was computed under; analytic
    // penalty weight schedules advance with this counter, so a stale stamp
    // invalidates the memo even at unchanged θ.
    last_outer_iter: Option<u64>,
}

impl SingleBlockLatentCoordDesignCache {
    fn new(
        data: Array2<f64>,
        spec: TermCollectionSpec,
        design: TermCollectionDesign,
        latent: &StandardLatentCoordConfig,
        rho_dim: usize,
    ) -> Result<Self, String> {
        if latent.term_index.get() >= spec.smooth_terms.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "latent-coordinate term index {} out of bounds for {} smooth terms",
                latent.term_index,
                spec.smooth_terms.len()
            ))
            .into());
        }
        if latent.feature_cols.len() != latent.values.latent_dim() {
            return Err(SmoothError::dimension_mismatch(format!(
                "latent-coordinate feature width mismatch: feature_cols={}, latent_dim={}",
                latent.feature_cols.len(),
                latent.values.latent_dim()
            ))
            .into());
        }
        if latent.values.n_obs() != data.nrows() {
            return Err(SmoothError::dimension_mismatch(format!(
                "latent-coordinate row mismatch: latent n={}, data n={}",
                latent.values.n_obs(),
                data.nrows()
            ))
            .into());
        }
        let analytic_rho_count = latent
            .analytic_penalties
            .as_ref()
            .map_or(0, |registry| registry.total_rho_count());
        Ok(Self {
            data,
            spec,
            design,
            current_theta: None,
            current_latent: None,
            current_hyper_dirs: None,
            current_design_cache_id: None,
            latent_design_cache: crate::solver::latent_cache::LatentDesignCache::default(),
            last_cost: None,
            last_eval: None,
            term_index: latent.term_index,
            feature_cols: latent.feature_cols.clone(),
            rho_dim,
            n_obs: latent.values.n_obs(),
            latent_dim: latent.values.latent_dim(),
            id_mode: latent.values.id_mode().clone(),
            manifold: latent.values.manifold().clone(),
            retraction_registry: latent.values.retraction_registry().clone(),
            latent_id: latent.values.latent_id(),
            analytic_penalties: latent.analytic_penalties.clone(),
            analytic_rho_count,
            design_revision: 0,
            last_outer_iter: None,
        })
    }

    fn design_revision(&self) -> u64 {
        self.design_revision
    }

    fn design(&self) -> &TermCollectionDesign {
        &self.design
    }

    fn latent(
        &self,
    ) -> Result<std::sync::Arc<crate::terms::latent_coord::LatentCoordValues>, String> {
        self.current_latent
            .as_ref()
            .cloned()
            .ok_or_else(|| "latent-coordinate cache has not been realized".to_string())
    }

    fn analytic_penalties(&self) -> Option<std::sync::Arc<crate::terms::AnalyticPenaltyRegistry>> {
        self.analytic_penalties.clone()
    }

    fn analytic_penalty_rho_count(&self) -> usize {
        self.analytic_rho_count
    }

    fn hyper_dirs(&self) -> Result<Vec<crate::estimate::reml::DirectionalHyperParam>, String> {
        self.current_hyper_dirs
            .as_ref()
            .cloned()
            .ok_or_else(|| "latent-coordinate hyper_dirs cache has not been realized".to_string())
    }

    fn latent_basis_kind(&self) -> Result<crate::solver::latent_cache::LatentBasisKind, String> {
        let smooth_term = self
            .design
            .smooth
            .terms
            .get(self.term_index.get())
            .ok_or_else(|| {
                SmoothError::dimension_mismatch(format!(
                    "LatentCoord term index {} out of bounds for realized smooth design",
                    self.term_index
                ))
            })?;
        let termspec = self
            .spec
            .smooth_terms
            .get(self.term_index.get())
            .ok_or_else(|| {
                SmoothError::dimension_mismatch(format!(
                    "LatentCoord term index {} out of bounds for resolved smooth spec",
                    self.term_index
                ))
            })?;
        match (&termspec.basis, &smooth_term.metadata) {
            (
                SmoothBasisSpec::Matern { .. },
                BasisMetadata::Matern {
                    centers,
                    length_scale,
                    nu,
                    aniso_log_scales,
                    ..
                },
            ) => Ok(crate::solver::latent_cache::LatentBasisKind::Matern {
                centers: centers.clone(),
                length_scale: *length_scale,
                nu: *nu,
                aniso_log_scales: aniso_log_scales
                    .clone()
                    .unwrap_or_else(|| vec![0.0; centers.ncols()]),
                chunk_size: crate::basis::auto_streaming_chunk_size_for_dense(
                    self.n_obs,
                    centers.nrows(),
                ),
            }),
            (
                SmoothBasisSpec::Duchon { .. },
                BasisMetadata::Duchon {
                    centers,
                    length_scale,
                    power,
                    nullspace_order,
                    aniso_log_scales,
                    ..
                },
            ) => Ok(crate::solver::latent_cache::LatentBasisKind::Duchon {
                centers: centers.clone(),
                length_scale: *length_scale,
                power: *power,
                nullspace_order: *nullspace_order,
                aniso_log_scales: aniso_log_scales
                    .clone()
                    .unwrap_or_else(|| vec![0.0; centers.ncols()]),
            }),
            (
                SmoothBasisSpec::Sphere { .. },
                BasisMetadata::Sphere {
                    centers,
                    penalty_order,
                    method,
                    ..
                },
            ) if matches!(*method, crate::basis::SphereMethod::Wahba) => {
                Ok(crate::solver::latent_cache::LatentBasisKind::Sphere {
                    centers: centers.clone(),
                    penalty_order: *penalty_order,
                    chunk_size: crate::basis::auto_streaming_chunk_size_for_dense(
                        self.n_obs,
                        centers.nrows(),
                    ),
                })
            }
            (
                SmoothBasisSpec::BSpline1D { spec, .. },
                BasisMetadata::BSpline1D {
                    knots,
                    periodic,
                    degree: meta_degree,
                    ..
                },
            ) => {
                // Issue #340: prefer the metadata-recorded effective degree
                // (which reflects fit-time auto-shrink) over the upstream
                // user-requested `spec.degree`.
                let effective_degree = meta_degree.unwrap_or(spec.degree);
                if let Some((domain_start, period, num_basis)) = periodic {
                    Ok(
                        crate::solver::latent_cache::LatentBasisKind::PeriodicBspline {
                            domain_start: *domain_start,
                            period: *period,
                            degree: effective_degree,
                            num_basis: *num_basis,
                            chunk_size: crate::basis::auto_streaming_chunk_size_for_dense(
                                self.n_obs, *num_basis,
                            ),
                        },
                    )
                } else {
                    let num_basis_est = knots.len().saturating_sub(effective_degree + 1);
                    Ok(
                        crate::solver::latent_cache::LatentBasisKind::TensorBspline {
                            knots: vec![knots.clone()],
                            degrees: vec![effective_degree],
                            chunk_size: crate::basis::auto_streaming_chunk_size_for_dense(
                                self.n_obs,
                                num_basis_est,
                            ),
                        },
                    )
                }
            }
            (
                SmoothBasisSpec::TensorBSpline { .. },
                BasisMetadata::TensorBSpline { knots, degrees, .. },
            ) => Ok(
                crate::solver::latent_cache::LatentBasisKind::TensorBspline {
                    knots: knots.clone(),
                    degrees: degrees.clone(),
                    chunk_size: None,
                },
            ),
            (
                SmoothBasisSpec::Pca { .. },
                BasisMetadata::Pca {
                    basis_matrix,
                    centered,
                    smooth_penalty,
                    center_mean,
                    pca_basis_path,
                    chunk_size,
                    ..
                },
            ) => {
                let center_mean_fingerprint = if *centered && pca_basis_path.is_none() {
                    let mean = center_mean.as_ref().ok_or_else(|| {
                        SmoothError::invalid_config(
                            "latent-coordinate Pca cache key requires center_mean when centered",
                        )
                    })?;
                    Some(crate::solver::latent_cache::pca_center_mean_fingerprint(
                        mean,
                    ))
                } else {
                    None
                };
                Ok(crate::solver::latent_cache::LatentBasisKind::Pca {
                    basis_matrix: basis_matrix.clone(),
                    centered: *centered,
                    center_mean_fingerprint,
                    smooth_penalty: *smooth_penalty,
                    pca_basis_path: pca_basis_path.clone(),
                    chunk_size: *chunk_size,
                })
            }
            _ => Err(SmoothError::invalid_config(
                "latent-coordinate design cache could not key the realized latent smooth basis"
                    .to_string(),
            )
            .into()),
        }
    }

    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
        {
            return Ok(());
        }
        let latent_flat_len = self.n_obs * self.latent_dim;
        let direct_hyper_count = latent_coord_direct_hyper_count(&self.id_mode, self.latent_dim);
        let expected =
            self.rho_dim + latent_flat_len + self.analytic_rho_count + direct_hyper_count;
        if theta.len() != expected {
            return Err(SmoothError::dimension_mismatch(format!(
                "latent-coordinate theta length mismatch: got {}, expected {} (rho_dim={}, n={}, d={}, analytic_rhos={}, direct_hypers={})",
                theta.len(),
                expected,
                self.rho_dim,
                self.n_obs,
                self.latent_dim,
                self.analytic_rho_count,
                direct_hyper_count
            ))
            .into());
        }
        let flat = theta
            .slice(s![self.rho_dim..self.rho_dim + latent_flat_len])
            .to_owned();
        let latent = std::sync::Arc::new(
            crate::terms::latent_coord::LatentCoordValues::from_flat_with_manifold_and_retraction_and_id(
                flat,
                self.n_obs,
                self.latent_dim,
                self.id_mode.clone(),
                self.manifold.clone(),
                self.retraction_registry.clone(),
                self.latent_id,
            ),
        );
        let latent_values_changed = self
            .current_latent
            .as_ref()
            .map(|cached| !latent_values_match(cached.as_flat(), latent.as_flat()))
            .unwrap_or(true);
        if latent_values_changed {
            self.latent_design_cache.invalidate_all();
            self.current_design_cache_id = None;
            self.design_revision = self.design_revision.wrapping_add(1);
        }
        for n in 0..self.n_obs {
            for axis in 0..self.latent_dim {
                let col = self.feature_cols[axis];
                self.data[[n, col]] = latent.as_flat()[n * self.latent_dim + axis];
            }
        }

        let basis_kind = self.latent_basis_kind()?;
        let rebuilt_width = self.design.design.ncols();
        let spec = self.spec.clone();
        let term_index = self.term_index;
        let analytic_rho_count = self.analytic_rho_count;
        let data = self.data.view();
        let design_context_digest =
            crate::solver::latent_cache::latent_design_context_cache_digest(
                data,
                &spec,
                term_index,
                analytic_rho_count,
                &self.feature_cols,
            )
            .map_err(|e| e.to_string())?;
        let lookup = self
            .latent_design_cache
            .lookup_or_compute(latent.clone(), basis_kind, design_context_digest, || {
                let rebuilt = build_term_collection_design(data, &spec).map_err(|e| {
                    EstimationError::InvalidInput(format!(
                        "failed to rebuild latent-coordinate design: {e}"
                    ))
                })?;
                if rebuilt.design.ncols() != rebuilt_width {
                    crate::bail_invalid_estim!(
                        "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
                        rebuilt.design.ncols(),
                        rebuilt_width
                    );
                }
                let hyper_dirs = try_build_latent_coord_hyper_dirs(
                    latent.clone(),
                    &spec,
                    &rebuilt,
                    &[term_index],
                    analytic_rho_count,
                )?
                .ok_or_else(|| {
                    EstimationError::InvalidInput(
                        "failed to build latent-coordinate hyper_dirs".to_string(),
                    )
                })?;
                Ok(crate::solver::latent_cache::ComputedLatentDesign {
                    design: rebuilt,
                    hyper_dirs,
                })
            })
            .map_err(|e| e.to_string())?;
        if lookup.cached.design.design.ncols() != self.design.design.ncols() {
            return Err(SmoothError::dimension_mismatch(format!(
                "latent-coordinate design topology changed: rebuilt p={}, cached p={}",
                lookup.cached.design.design.ncols(),
                self.design.design.ncols()
            ))
            .into());
        }
        self.design = lookup.cached.design.clone();
        self.current_hyper_dirs = Some(lookup.cached.hyper_dirs.clone());
        self.current_latent = Some(latent);
        self.current_theta = Some(theta.clone());
        self.last_cost = None;
        self.last_eval = None;
        self.last_outer_iter = None;
        if !latent_values_changed && self.current_design_cache_id != Some(lookup.entry_id) {
            self.design_revision = self.design_revision.wrapping_add(1);
        }
        self.current_design_cache_id = Some(lookup.entry_id);
        Ok(())
    }

    fn memoized_cost(&self, theta: &Array1<f64>) -> Option<f64> {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
            && self.last_outer_iter
                == Some(crate::solver::estimate::reml::runtime::current_outer_iter())
        {
            self.last_eval
                .as_ref()
                .map(|cached| cached.0)
                .or(self.last_cost)
        } else {
            None
        }
    }

    fn memoized_eval(
        &self,
        theta: &Array1<f64>,
    ) -> Option<(
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    )> {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
            && self.last_outer_iter
                == Some(crate::solver::estimate::reml::runtime::current_outer_iter())
        {
            self.last_eval.clone()
        } else {
            None
        }
    }

    fn store_eval(
        &mut self,
        eval: (
            f64,
            Array1<f64>,
            crate::solver::outer_strategy::HessianResult,
        ),
    ) {
        self.last_cost = Some(eval.0);
        self.last_eval = Some(eval);
        self.last_outer_iter = Some(crate::solver::estimate::reml::runtime::current_outer_iter());
    }

    fn store_cost(&mut self, cost: f64) {
        self.last_cost = Some(cost);
        self.last_outer_iter = Some(crate::solver::estimate::reml::runtime::current_outer_iter());
    }

    fn reset(&mut self) {
        self.current_theta = None;
        self.current_latent = None;
        self.current_hyper_dirs = None;
        self.current_design_cache_id = None;
        self.latent_design_cache.invalidate();
        self.last_cost = None;
        self.last_eval = None;
        self.last_outer_iter = None;
    }
}

fn try_exact_joint_spatial_length_scale_optimization(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    best: &FittedTermCollection,
    family: LikelihoodSpec,
    options: &FitOptions,
    kappa_options: &SpatialLengthScaleOptimizationOptions,
    spatial_terms: &[usize],
) -> Result<Option<FittedTermCollectionWithSpec>, EstimationError> {
    if spatial_terms.is_empty() {
        return Ok(None);
    }
    // Fail loud on nonsensical κ options rather than letting them propagate
    // silent NaNs (e.g. inverted min/max inverts the BFGS window, negative
    // scales produce NaN logs). This is the first function on every outer-κ
    // path; downstream paths assume validated options.
    kappa_options
        .validate()
        .map_err(EstimationError::InvalidInput)?;
    if try_build_spatial_log_kappa_hyper_dirs(data, resolvedspec, &best.design, spatial_terms)?
        .is_none()
    {
        return Ok(None);
    }

    const JOINT_RHO_BOUND: f64 = 12.0;
    let rho_dim = best.fit.lambdas.len();

    // Compute per-term dimensionality for anisotropic terms.
    let dims_per_term = spatial_dims_per_term(resolvedspec, spatial_terms);
    let use_aniso = has_aniso_terms(resolvedspec, spatial_terms);

    // Build initial ψ values and bounds, using aniso-aware constructors
    // when any term has d > 1 axes. Bounds are tied to each term's center
    // geometry (r_min, r_max) so κ cannot saturate at an upper bound that
    // has no relationship to the data's distance scale.
    let log_kappa0 = if use_aniso {
        SpatialLogKappaCoords::from_length_scales_aniso(resolvedspec, spatial_terms, kappa_options)
    } else {
        SpatialLogKappaCoords::from_length_scales(resolvedspec, spatial_terms, kappa_options)
    };
    // If the user/spec did not set a length_scale, re-seed ψ at the midpoint
    // of the data-derived window instead of the arbitrary options fallback.
    let log_kappa0 = log_kappa0.reseed_from_data(data, resolvedspec, spatial_terms, kappa_options);
    let log_kappa_lower = if use_aniso {
        SpatialLogKappaCoords::lower_bounds_aniso_from_data(
            data,
            resolvedspec,
            spatial_terms,
            &dims_per_term,
            kappa_options,
        )
    } else {
        SpatialLogKappaCoords::lower_bounds_from_data(
            data,
            resolvedspec,
            spatial_terms,
            kappa_options,
        )
    };
    let log_kappa_upper = if use_aniso {
        SpatialLogKappaCoords::upper_bounds_aniso_from_data(
            data,
            resolvedspec,
            spatial_terms,
            &dims_per_term,
            kappa_options,
        )
    } else {
        SpatialLogKappaCoords::upper_bounds_from_data(
            data,
            resolvedspec,
            spatial_terms,
            kappa_options,
        )
    };
    // Project seed onto data-derived bounds; spec.length_scale is a hint,
    // not a hard constraint. BFGS requires theta0 ∈ [lower, upper].
    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
    let setup = ExactJointHyperSetup::new(
        best.fit.lambdas.mapv(f64::ln),
        Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
        Array1::<f64>::from_elem(rho_dim, JOINT_RHO_BOUND),
        log_kappa0,
        log_kappa_lower,
        log_kappa_upper,
    );

    let theta0 = setup.theta0();
    let lower = setup.lower();
    let upper = setup.upper();

    // ───────────────────────────────────────────────────────────────────────
    //  Both coordinate kinds drive the SAME exact joint optimizer
    //  (`run_exact_joint_spatial_optimization`): the unified REML evaluator with
    //  ext_coords for joint [ρ, ψ] optimization, with analytic gradient +
    //  Hessian flowing through the
    //  AnisoBasisPsiDerivatives / SpatialPsiDerivative → DirectionalHyperParam →
    //  HyperCoord pipeline for Newton/BFGS quadratic convergence. The only
    //  difference is the coordinate kind: anisotropic carries one ψ per axis per
    //  term, isotropic one log-κ per term. `outer_strategy` handles the
    //  centralized degradation path when the analytic Hessian is unavailable.
    // ───────────────────────────────────────────────────────────────────────
    let kind = if use_aniso {
        SpatialHyperKind::Anisotropic
    } else {
        SpatialHyperKind::Isotropic
    };
    let (theta_star, joint_final_value) = run_exact_joint_spatial_optimization(
        kind,
        data,
        y,
        weights,
        offset,
        resolvedspec,
        &best.design,
        family.clone(),
        options,
        spatial_terms,
        &dims_per_term,
        &theta0,
        &lower,
        &upper,
        rho_dim,
        kappa_options,
    )?;

    let baseline_score = fit_score(&best.fit);
    let baseline_result = FittedTermCollectionWithSpec {
        fit: best.fit.clone(),
        design: best.design.clone(),
        resolvedspec: resolvedspec.clone(),
        adaptive_diagnostics: best.adaptive_diagnostics.clone(),
    };

    // Compare the joint optimizer's certified cost (final_value at theta*)
    // against the baseline. Tolerance ≥ options.tol because both endpoints
    // are outer-BFGS approximations accurate to options.tol; a tighter
    // gate would reject true improvements due to floating-point noise.
    let accept_tol = options.tol.max(1e-8 * baseline_score.abs()).max(1e-12);
    if joint_final_value > baseline_score + accept_tol {
        log::info!(
            "[spatial-kappa] exact joint spatial candidate worsened the profiled score (joint={:.6e}, baseline={:.6e}, tol={:.2e}); keeping the frozen baseline geometry",
            joint_final_value,
            baseline_score,
            accept_tol,
        );
        return Ok(Some(baseline_result));
    }

    let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
    let log_kappa_star =
        SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
    let resolvedspec = log_kappa_star.apply_tospec(resolvedspec, spatial_terms)?;
    let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
        data,
        y,
        weights,
        offset,
        &resolvedspec,
        rho_star.as_slice(),
        family,
        options,
    )?;

    // Stamp reml_score with joint_final_value so downstream consumers see a
    // score consistent with the gate decision; the refit serves as a
    // β/inference harvester at the certified (ρ*, ψ*).
    let mut fit = optimized.fit;
    fit.reml_score = joint_final_value;
    let optimized_result = FittedTermCollectionWithSpec {
        fit,
        design: optimized.design,
        resolvedspec,
        adaptive_diagnostics: optimized.adaptive_diagnostics,
    };

    Ok(Some(optimized_result))
}

/// Coordinate kind for the exact joint spatial hyperparameter optimizer.
///
/// Anisotropic and isotropic spatial terms drive the *same* joint `[ρ, ψ]`
/// optimizer: identical outer-Hessian policy, identical
/// `ExternalJointHyperEvaluator` wiring, identical multistart problem, identical
/// convergence processing, and an identical `eval_full / eval_efs / eval_cost`
/// inner loop that routes ψ through `try_build_spatial_log_kappa_hyper_dirs`.
/// The only difference is the coordinate *kind*: the anisotropic path carries
/// one log-scale coordinate per axis per term (ψ_a) while the isotropic path
/// carries one log-κ coordinate per term. The kind selects diagnostic labels
/// only — the numerics are shared verbatim.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum SpatialHyperKind {
    Anisotropic,
    Isotropic,
}

impl SpatialHyperKind {
    /// Stable diagnostic prefix used in every `log::*` line and as the
    /// `ExternalJointHyperEvaluator` / cost-only label root.
    fn label(self) -> &'static str {
        match self {
            SpatialHyperKind::Anisotropic => "spatial-aniso-joint",
            SpatialHyperKind::Isotropic => "spatial-iso-joint",
        }
    }

    /// Human-readable adjective for error strings ("anisotropic" / "isotropic").
    fn adjective(self) -> &'static str {
        match self {
            SpatialHyperKind::Anisotropic => "anisotropic",
            SpatialHyperKind::Isotropic => "isotropic",
        }
    }

    /// Name of the directional coordinate being optimized ("psi" / "kappa"),
    /// used only in hyper-direction construction error messages.
    fn coord_name(self) -> &'static str {
        match self {
            SpatialHyperKind::Anisotropic => "psi",
            SpatialHyperKind::Isotropic => "kappa",
        }
    }
}

/// Shared context for the exact joint spatial optimizer's closures. Holds the
/// realized-design cache and the joint REML evaluator, plus the coordinate
/// `kind` whose only effect is the diagnostic label routed into the cost-only
/// evaluation path. The `eval_full / eval_efs / eval_cost` methods are the
/// single source of truth for both anisotropic and isotropic spatial terms.
struct SpatialJointContext<'d> {
    data: ArrayView2<'d, f64>,
    rho_dim: usize,
    kind: SpatialHyperKind,
    cache: SingleBlockExactJointDesignCache<'d>,
    evaluator: crate::estimate::ExternalJointHyperEvaluator<'d>,
}

impl<'d> SpatialJointContext<'d> {
    /// Full evaluation on the current realized design + hyper_dirs.
    fn eval_full(
        &mut self,
        theta: &Array1<f64>,
        order: crate::solver::outer_strategy::OuterEvalOrder,
        analytic_outer_hessian_available: bool,
    ) -> Result<
        (
            f64,
            Array1<f64>,
            crate::solver::outer_strategy::HessianResult,
        ),
        EstimationError,
    > {
        use crate::solver::outer_strategy::OuterEvalOrder;
        let allow_second_order = matches!(order, OuterEvalOrder::ValueGradientHessian)
            && analytic_outer_hessian_available;
        if let Some(eval) = self.cache.memoized_eval(theta) {
            let cached_satisfies_order = !allow_second_order || eval.2.is_analytic();
            if cached_satisfies_order {
                return Ok(eval);
            }
        }
        self.cache
            .ensure_theta(theta)
            .map_err(EstimationError::InvalidInput)?;
        let kind = self.kind;
        let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
            self.data,
            self.cache.spec(),
            self.cache.design(),
            &self.cache.spatial_terms,
        )?
        .ok_or_else(|| {
            EstimationError::InvalidInput(format!(
                "failed to build {} hyper_dirs at current {}",
                kind.adjective(),
                kind.coord_name(),
            ))
        })?;

        let design_revision = Some(self.cache.design_revision());
        let eval = evaluate_joint_reml_outer_eval_at_theta(
            &mut self.evaluator,
            self.cache.design(),
            theta,
            self.rho_dim,
            hyper_dirs,
            None,
            if allow_second_order {
                order
            } else {
                OuterEvalOrder::ValueAndGradient
            },
            design_revision,
        );
        if let Ok(ref value) = eval {
            self.cache.store_eval(value.clone());
        }
        eval
    }

    fn eval_efs(
        &mut self,
        theta: &Array1<f64>,
    ) -> Result<crate::solver::outer_strategy::EfsEval, EstimationError> {
        self.cache
            .ensure_theta(theta)
            .map_err(EstimationError::InvalidInput)?;
        let kind = self.kind;
        let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
            self.data,
            self.cache.spec(),
            self.cache.design(),
            &self.cache.spatial_terms,
        )?
        .ok_or_else(|| {
            EstimationError::InvalidInput(format!(
                "failed to build {} hyper_dirs for exact-joint EFS",
                kind.adjective(),
            ))
        })?;
        let design_revision = Some(self.cache.design_revision());
        evaluate_joint_reml_efs_at_theta(
            &mut self.evaluator,
            self.cache.design(),
            theta,
            self.rho_dim,
            hyper_dirs,
            None,
            design_revision,
        )
    }

    /// Cost-only evaluation. BFGS line-search probes route through the
    /// evaluator's true value-only path so they neither construct
    /// `try_build_spatial_log_kappa_hyper_dirs` nor assemble a gradient that
    /// the line search will discard. Split-borrow on `self.cache` +
    /// `self.evaluator` matches the pattern already used by `eval_full`.
    fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
        if let Some(cost) = self.cache.memoized_cost(theta) {
            return cost;
        }
        if self.cache.ensure_theta(theta).is_err() {
            return f64::INFINITY;
        }
        let design_revision = Some(self.cache.design_revision());
        let cost_label = self.kind.label();
        let result = {
            let design = self.cache.design();
            self.evaluator.evaluate_cost_only(
                &design.design,
                &design.penalties,
                &design.nullspace_dims,
                design.linear_constraints.clone(),
                theta,
                self.rho_dim,
                None,
                cost_label,
                design_revision,
            )
        };
        match result {
            Ok(cost) => {
                self.cache.store_cost(cost);
                cost
            }
            Err(_) => f64::INFINITY,
        }
    }

    fn reset(&mut self) {
        self.cache.current_theta = None;
        self.cache.last_cost = None;
        self.cache.last_eval = None;
    }
}

/// Exact joint `[ρ, ψ]` optimization for spatial terms using analytic
/// derivatives through the unified REML evaluator. This is the single shared
/// engine for both the anisotropic and isotropic coordinate kinds (selected by
/// `kind`).
///
/// At each outer iteration, the frozen term topology is reused and only the
/// spatial realized blocks affected by the current ψ are refreshed before the
/// unified evaluator returns cost + gradient + Hessian for the full
/// θ = [ρ, ψ] vector. The ψ derivatives flow through:
///
///   `AnisoBasisPsiDerivatives` / `SpatialPsiDerivative` → `DirectionalHyperParam`
///     → `build_tau_unified_objects` → `HyperCoord` ext_coords → unified evaluator
///
/// This gives Newton/BFGS quadratic convergence on the length-scale /
/// anisotropy parameters while jointly optimizing the smoothing parameters.
///
/// The ψ coordinates are parameterized as unconstrained log-scales. For the
/// anisotropic kind the decomposition into isotropic scale (ψ̄ = mean(ψ_a)) and
/// anisotropy (η_a = ψ_a − ψ̄, with Ση_a = 0) happens only on writeback via
/// `SpatialLogKappaCoords::apply_tospec`; the all-ones direction in ψ-space is
/// NOT a gauge direction — it controls the identifiable isotropic scale
/// κ = exp(ψ̄). The isotropic kind carries one log-κ coordinate per term. In
/// neither case is a sum-to-zero constraint enforced during optimization.
fn run_exact_joint_spatial_optimization(
    kind: SpatialHyperKind,
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    baseline_design: &TermCollectionDesign,
    family: LikelihoodSpec,
    options: &FitOptions,
    spatial_terms: &[usize],
    dims_per_term: &[usize],
    theta0: &Array1<f64>,
    lower: &Array1<f64>,
    upper: &Array1<f64>,
    rho_dim: usize,
    kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<(Array1<f64>, f64), EstimationError> {
    let label = kind.label();
    // Use bounds and design metadata for validation.
    assert!(
        lower.len() == theta0.len() && upper.len() == theta0.len(),
        "spatial hyperparameter bounds must match theta length: lower_len={}, upper_len={}, theta_len={}",
        lower.len(),
        upper.len(),
        theta0.len()
    );
    assert!(
        baseline_design.smooth.terms.len() >= spatial_terms.len(),
        "baseline design must have at least one smooth term per spatial term: baseline_terms={}, spatial_terms={}",
        baseline_design.smooth.terms.len(),
        spatial_terms.len()
    );
    use crate::solver::outer_strategy::{
        DeclaredHessianForm, Derivative, OuterEval, OuterEvalOrder,
    };

    let theta_dim = theta0.len();
    // Directional-coordinate dimension: psi-per-axis (anisotropic) or
    // kappa-per-term (isotropic). The numerics below are identical either way.
    let coord_dim = theta_dim - rho_dim;
    // Capability is declared solely from derivative coverage, not from
    // problem size. The unified REML evaluator now exposes exact matrix-free
    // outer Hessian operators for the costly third/fourth-derivative
    // contractions used by spatial ψ coordinates; its internal
    // `(n, p, K)` work model chooses `HessianResult::Operator` at biobank
    // scale and the dense analytic matrix only below that crossover. Keeping
    // `Derivative::Analytic` here preserves ARC / trust-region-CG second-order
    // optimization for `n > 50_000` and `coord_dim > 30` instead of forcing the
    // obsolete HybridEFS compatibility path.
    let analytic_outer_hessian_available =
        exact_joint_spatial_outer_hessian_available(&family, baseline_design);
    if !analytic_outer_hessian_available {
        log::info!(
            "[{label}] analytic outer Hessian unavailable for family/design; routing without second-order geometry (coord_dim={coord_dim})"
        );
    }
    let prefer_gradient_only = false;

    log::trace!(
        "[{}] starting analytic optimization: rho_dim={}, coord_dim={}, dims_per_term={:?}",
        label,
        rho_dim,
        coord_dim,
        dims_per_term,
    );

    let mut ctx = SpatialJointContext {
        data,
        rho_dim,
        kind,
        cache: SingleBlockExactJointDesignCache::new(
            data,
            resolvedspec.clone(),
            baseline_design.clone(),
            spatial_terms.to_vec(),
            rho_dim,
            dims_per_term.to_vec(),
        )
        .map_err(EstimationError::InvalidInput)?,
        evaluator: crate::estimate::ExternalJointHyperEvaluator::new(
            y,
            weights,
            &baseline_design.design,
            offset,
            &baseline_design.penalties,
            &external_opts_for_design(&family, baseline_design, options),
            label,
        )?,
    };

    let problem = exact_joint_multistart_outer_problem(
        theta0,
        lower,
        upper,
        rho_dim,
        coord_dim,
        theta_dim,
        Derivative::Analytic,
        if analytic_outer_hessian_available {
            DeclaredHessianForm::Either
        } else {
            DeclaredHessianForm::Unavailable
        },
        prefer_gradient_only,
        // Single-block spatial path: penalty-like rho + spatial psi.
        // EFS/HybridEFS remain eligible; the Wood-Fasiolo PSD structure holds
        // for single-block families with β-independent joint H_L.
        false,
        seed_risk_profile_for_likelihood_family(&family),
        kappa_options.rel_tol.max(1e-6),
        kappa_options.max_outer_iter.max(1),
        // Rho-axis BFGS cap: log-λ's natural step is ≈ 5 per
        // `first_order_bfgs_loglambda_step_cap`. Anything tighter throttles
        // BFGS on flat REML valleys.
        Some(5.0),
        // Psi-axis BFGS cap: kappa / aniso-log-scale needs ~ln 2 per iter.
        Some(kappa_options.log_step.clamp(0.25, 1.0)),
        None,
    );

    let eval_outer = |ctx: &mut &mut SpatialJointContext<'_>,
                      theta: &Array1<f64>,
                      order: OuterEvalOrder|
     -> Result<OuterEval, EstimationError> {
        match ctx.eval_full(theta, order, analytic_outer_hessian_available) {
            Ok((cost, grad, hess)) => Ok(OuterEval {
                cost,
                gradient: grad,
                hessian: hess,
                inner_beta_hint: None,
            }),
            // A trial hyperparameter at which the spatial kernel design /
            // ψ-derivatives are non-constructible is an infeasible point, not
            // a fatal error: the gradient/Hessian path must retreat exactly as
            // the cost-only path (which already returns +∞) does. Returning
            // `OuterEval::infeasible` keeps the two paths symmetric so a single
            // bad probe — e.g. an anisotropy that overflows the Duchon radial
            // kernel — no longer aborts the whole REML optimization.
            Err(err) if is_recoverable_trial_point_error(&err) => {
                log::debug!(
                    "[{label}] trial point infeasible (kernel design \
                     not constructible at theta={theta:?}): {err}; retreating",
                );
                Ok(OuterEval::infeasible(theta_dim))
            }
            Err(err) => Err(err),
        }
    };

    let mut obj = problem.build_objective_with_eval_order(
        &mut ctx,
        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| {
            eval_outer(
                ctx,
                theta,
                if analytic_outer_hessian_available {
                    OuterEvalOrder::ValueGradientHessian
                } else {
                    OuterEvalOrder::ValueAndGradient
                },
            )
        },
        |ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
            eval_outer(ctx, theta, order)
        },
        Some(|ctx: &mut &mut SpatialJointContext<'_>| {
            ctx.reset();
        }),
        Some(|ctx: &mut &mut SpatialJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
    );

    let run_label = match kind {
        SpatialHyperKind::Anisotropic => "aniso-psi joint REML",
        SpatialHyperKind::Isotropic => "iso-kappa joint REML",
    };
    let result = problem.run(&mut obj, run_label).map_err(|e| {
        EstimationError::InvalidInput(format!(
            "{} analytic optimization failed after exhausting strategy fallbacks: {e}",
            kind.adjective(),
        ))
    })?;
    if !result.converged {
        // Mirror `fit_term_collectionwith_exact_spatial_adaptive_regularization`
        // (commit 0267d082): the strict absolute-floor gradient criterion is too
        // tight when the outer Hessian carries a near-null direction (η-anchor
        // drift, ill-conditioned operator-collocation Gram, etc.) — the iterate
        // settles into a flat valley with ‖g‖_proj at numerical-noise scale
        // (~1e-5 for cost ~1e1 in double precision) which is above the 1e-6
        // absolute floor but well below the textbook mgcv `magic` REML rule
        // ‖g‖_proj ≤ τ·(1 + |f|). Accept the iterate under the rel-to-cost
        // form when the absolute form has timed out; divergent runs (‖g‖
        // large relative to |f|) still surface as errors.
        let rel_to_cost_threshold = options.tol * (1.0_f64 + result.final_value.abs());
        if let Some(final_grad) = result
            .final_grad_norm
            .filter(|v| v.is_finite() && *v <= rel_to_cost_threshold)
        {
            log::info!(
                "[{}] outer optimization hit max_iter={} but \
                 projected gradient norm {:.3e} ≤ τ·(1+|f|) = {:.3e} \
                 (τ={:.3e}, |f|={:.3e}); accepting iterate under the mgcv-style \
                 relative-to-cost REML convergence criterion.",
                label,
                result.iterations,
                final_grad,
                rel_to_cost_threshold,
                options.tol,
                result.final_value.abs(),
            );
        } else {
            crate::bail_invalid_estim!(
                "{} analytic optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
                kind.adjective(),
                result.iterations,
                result.final_value,
                result.final_grad_norm_report(),
            );
        }
    }
    log::trace!(
        "[{}] converged in {} iterations, final_value={:.6e}, grad_norm={}",
        label,
        result.iterations,
        result.final_value,
        result.final_grad_norm_report(),
    );
    // No sum-to-zero enforcement needed: ψ coordinates are unconstrained during
    // optimization. For the anisotropic kind the decomposition into (ψ̄, η)
    // happens later in apply_tospec.
    let theta_star = result.rho;
    Ok((theta_star, result.final_value))
}

fn set_spatial_length_scale(
    spec: &mut TermCollectionSpec,
    term_idx: usize,
    length_scale: f64,
) -> Result<(), EstimationError> {
    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
        crate::bail_invalid_estim!("spatial length-scale term index {term_idx} out of range");
    };
    match &mut term.basis {
        SmoothBasisSpec::ThinPlate { spec, .. } => {
            spec.length_scale = length_scale;
            Ok(())
        }
        SmoothBasisSpec::Matern { spec, .. } => {
            spec.length_scale = length_scale;
            Ok(())
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            spec.length_scale = Some(length_scale);
            Ok(())
        }
        _ => Err(EstimationError::InvalidInput(format!(
            "term '{}' does not expose a spatial length scale",
            term.name
        ))),
    }
}

/// Apply a length scale to a single `SmoothTermSpec` (independent of any
/// outer `TermCollectionSpec`). Mirrors `set_spatial_length_scale` but on a
/// term in isolation; used by the incremental realizer's cached planned spec.
fn set_single_term_spatial_length_scale(
    term: &mut SmoothTermSpec,
    length_scale: f64,
) -> Result<(), EstimationError> {
    match &mut term.basis {
        SmoothBasisSpec::ThinPlate { spec, .. } => {
            spec.length_scale = length_scale;
            Ok(())
        }
        SmoothBasisSpec::Matern { spec, .. } => {
            spec.length_scale = length_scale;
            Ok(())
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            spec.length_scale = Some(length_scale);
            Ok(())
        }
        _ => Err(EstimationError::InvalidInput(format!(
            "term '{}' does not expose a spatial length scale",
            term.name
        ))),
    }
}

/// Apply anisotropy contrasts to a single `SmoothTermSpec`. Mirrors
/// `set_spatial_aniso_log_scales` but on a term in isolation; used by the
/// incremental realizer's cached planned spec.
fn set_single_term_spatial_aniso_log_scales(
    term: &mut SmoothTermSpec,
    eta: Vec<f64>,
) -> Result<(), EstimationError> {
    let eta = center_aniso_log_scales(&eta);
    match &mut term.basis {
        SmoothBasisSpec::Matern { spec, .. } => {
            spec.aniso_log_scales = Some(eta);
            Ok(())
        }
        SmoothBasisSpec::Duchon { spec, .. } => {
            spec.aniso_log_scales = Some(eta);
            Ok(())
        }
        _ => Err(EstimationError::InvalidInput(format!(
            "term '{}' does not support aniso_log_scales",
            term.name
        ))),
    }
}

pub fn get_spatial_length_scale(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
    spec.smooth_terms
        .get(term_idx)
        .and_then(|term| match &term.basis {
            SmoothBasisSpec::ThinPlate { spec, .. } => Some(spec.length_scale),
            SmoothBasisSpec::Matern { spec, .. } => Some(spec.length_scale),
            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
            _ => None,
        })
}

/// Freeze a `TermCollectionSpec` by baking in the concrete knots, centers,
/// identifiability transforms, and random-effect levels that were resolved
/// during design-matrix construction.  The result passes `validate_frozen`
/// and is safe to serialize for prediction.
///
/// This is the single canonical freezer — every model-save path should call
/// this rather than rolling ad-hoc freezing logic.
/// Freeze a smooth basis spec from its fit-time metadata so that predict-time
/// rebuilds reproduce the exact fitted geometry instead of recomputing any
/// data-dependent construction (knot selection, radial reparameterization,
/// eigen-truncation, identifiability constraint) on the prediction rows.
///
/// This is the SINGLE source of truth for freezing, shared by stand-alone
/// terms and by `by=`-wrapped / factor-sum-to-zero inner smooths. The wrapper
/// arms recurse into this same function, so every inner basis kind is frozen
/// with identical logic. A previous split implementation froze only B-spline
/// inners, leaving spatial inner bases (`bs='tp'`/`matern`/`duchon`/`sos`)
/// unfrozen and recomputed on the prediction grid (#704).
fn freeze_smooth_basis_from_metadata(
    basis: &mut SmoothBasisSpec,
    metadata: &BasisMetadata,
    term_name: &str,
) -> Result<(), EstimationError> {
    match (&mut *basis, metadata) {
        (SmoothBasisSpec::ByVariable { inner, .. }, meta)
        | (SmoothBasisSpec::FactorSumToZero { inner, .. }, meta) => {
            freeze_smooth_basis_from_metadata(inner, meta, term_name)?;
        }
        (
            SmoothBasisSpec::BSpline1D { spec: s, .. },
            BasisMetadata::BSpline1D {
                knots,
                identifiability_transform,
                periodic,
                degree: meta_degree,
                ..
            },
        ) => {
            // Issue #340: bake the fit-time effective degree into the
            // frozen spec so reload sees a self-consistent
            // (degree, knots) pair.
            if let Some(d) = meta_degree {
                s.degree = *d;
            }
            s.knotspec = periodic
                .map(
                    |(domain_start, period, num_basis)| BSplineKnotSpec::PeriodicUniform {
                        data_range: (domain_start, domain_start + period),
                        num_basis,
                    },
                )
                .unwrap_or_else(|| BSplineKnotSpec::Provided(knots.clone()));
            s.identifiability = match identifiability_transform {
                Some(z) => BSplineIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => BSplineIdentifiability::None,
            };
            // Boundary projections are folded into `identifiability_transform`
            // by `build_bspline_basis_1d`. A frozen prediction spec must
            // rebuild the same raw knot basis and apply the captured
            // transform exactly once; keeping the original boundary
            // conditions would project the raw basis a second time and
            // shrink its width before `FrozenTransform` is applied.
            s.boundary_conditions = Default::default();
        }
        (
            SmoothBasisSpec::ThinPlate {
                spec: s,
                input_scales,
                ..
            },
            BasisMetadata::ThinPlate {
                centers,
                length_scale,
                periodic: meta_periodic,
                identifiability_transform,
                input_scales: meta_scales,
                radial_reparam,
            },
        ) => {
            s.center_strategy = crate::basis::CenterStrategy::UserProvided(centers.clone());
            s.length_scale = *length_scale;
            s.identifiability = match identifiability_transform {
                Some(z) => SpatialIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => match &s.identifiability {
                    SpatialIdentifiability::FrozenTransform { .. } => s.identifiability.clone(),
                    _ => SpatialIdentifiability::None,
                },
            };
            s.radial_reparam = radial_reparam.clone();
            s.periodic = meta_periodic.clone();
            *input_scales = meta_scales.clone();
        }
        (
            SmoothBasisSpec::ThinPlate { feature_cols, .. },
            BasisMetadata::Duchon {
                centers,
                length_scale,
                periodic: meta_periodic,
                power,
                nullspace_order,
                identifiability_transform,
                input_scales: meta_scales,
                aniso_log_scales: meta_aniso,
                ..
            },
        ) => {
            // Auto-promotion path: the basis builder rewrote a canonical-TPS
            // request to a pure Duchon spline because k < polynomial-nullspace
            // size at this dimension. Bake the resolved Duchon parameters into
            // the spec so predict-time goes through the same Duchon code path.
            let identifiability = match identifiability_transform {
                Some(z) => SpatialIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => SpatialIdentifiability::None,
            };
            *basis = SmoothBasisSpec::Duchon {
                feature_cols: feature_cols.clone(),
                spec: DuchonBasisSpec {
                    periodic: meta_periodic.clone(),
                    center_strategy: crate::basis::CenterStrategy::UserProvided(centers.clone()),
                    length_scale: *length_scale,
                    power: *power,
                    nullspace_order: *nullspace_order,
                    identifiability,
                    aniso_log_scales: meta_aniso.clone(),
                    operator_penalties: Default::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: meta_scales.clone(),
            };
        }
        (
            SmoothBasisSpec::Sphere { spec: s, .. },
            BasisMetadata::Sphere {
                centers,
                penalty_order,
                method,
                max_degree,
                wahba_kernel,
                constraint_transform,
            },
        ) => {
            s.center_strategy = crate::basis::CenterStrategy::UserProvided(centers.clone());
            s.penalty_order = *penalty_order;
            s.method = *method;
            s.max_degree = *max_degree;
            s.wahba_kernel = *wahba_kernel;
            // #532: freeze the realized-design transform (the composed
            // `z · z_parametric` captured at fit time) so the predict-time
            // rebuild reuses it verbatim instead of recomputing the
            // center-space `z`, which would drop the parametric
            // orthogonalization and resurrect the intercept collision. The
            // Harmonic method never carries a constraint transform.
            s.identifiability = match constraint_transform {
                Some(z) => SphericalSplineIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => SphericalSplineIdentifiability::CenterSumToZero,
            };
        }
        (
            SmoothBasisSpec::Matern {
                spec: s,
                input_scales,
                ..
            },
            BasisMetadata::Matern {
                centers,
                length_scale,
                periodic: meta_periodic,
                nu,
                include_intercept,
                identifiability_transform,
                input_scales: meta_scales,
                aniso_log_scales: meta_aniso,
            },
        ) => {
            s.center_strategy = crate::basis::CenterStrategy::UserProvided(centers.clone());
            s.length_scale = *length_scale;
            s.nu = *nu;
            s.include_intercept = *include_intercept;
            s.identifiability = match identifiability_transform {
                Some(z) => MaternIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => MaternIdentifiability::None,
            };
            s.aniso_log_scales = meta_aniso.clone();
            s.periodic = meta_periodic.clone();
            *input_scales = meta_scales.clone();
        }
        (
            SmoothBasisSpec::Duchon {
                spec: s,
                input_scales,
                ..
            },
            BasisMetadata::Duchon {
                centers,
                length_scale,
                periodic: meta_periodic,
                power,
                nullspace_order,
                identifiability_transform,
                input_scales: meta_scales,
                aniso_log_scales: meta_aniso,
                ..
            },
        ) => {
            s.center_strategy = crate::basis::CenterStrategy::UserProvided(centers.clone());
            s.length_scale = *length_scale;
            s.power = *power;
            s.nullspace_order = *nullspace_order;
            s.identifiability = match identifiability_transform {
                Some(z) => SpatialIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => match &s.identifiability {
                    // If the spec already carries a frozen transform but the
                    // metadata lost it (e.g. raw rebuild stripped it), keep
                    // the existing frozen transform rather than downgrading.
                    SpatialIdentifiability::FrozenTransform { .. } => s.identifiability.clone(),
                    _ => SpatialIdentifiability::None,
                },
            };
            s.aniso_log_scales = meta_aniso.clone();
            s.periodic = meta_periodic.clone();
            *input_scales = meta_scales.clone();
        }
        (
            SmoothBasisSpec::Sphere { spec: s, .. },
            BasisMetadata::SphereHarmonics {
                max_degree,
                radians,
            },
        ) => {
            s.max_degree = Some(*max_degree);
            s.radians = *radians;
        }
        (
            SmoothBasisSpec::TensorBSpline {
                feature_cols,
                spec: s,
            },
            BasisMetadata::TensorBSpline {
                feature_cols: fitted_cols,
                knots,
                degrees,
                periods,
                identifiability_transform,
            },
        ) => {
            if s.marginalspecs.len() != knots.len() || s.marginalspecs.len() != degrees.len() {
                crate::bail_invalid_estim!(
                    "tensor freeze mismatch for '{}': marginalspecs={}, knots={}, degrees={}",
                    term_name,
                    s.marginalspecs.len(),
                    knots.len(),
                    degrees.len()
                );
            }
            *feature_cols = fitted_cols.clone();
            for i in 0..s.marginalspecs.len() {
                s.marginalspecs[i].degree = degrees[i];
                s.marginalspecs[i].knotspec = match (periods[i], knots[i].len()) {
                    (Some(period), num_basis) if num_basis >= 1 => {
                        // Periodic uniform reconstructs the open
                        // `[start, start + period)` data range from the
                        // first knot and the saved period. `knots[i].len()`
                        // is the periodic control-site count.
                        let domain_start = knots[i][0];
                        BSplineKnotSpec::PeriodicUniform {
                            data_range: (domain_start, domain_start + period),
                            num_basis,
                        }
                    }
                    _ => BSplineKnotSpec::Provided(knots[i].clone()),
                };
            }
            // Do NOT overwrite `s.periods` from `periods`: `frozen = spec.clone()`
            // already preserves the user's original `spec.periods`. The metadata
            // `periods` slot captures the effective per-margin period for
            // restoring `PeriodicUniform` knotspecs, but it may also reflect
            // periodicity implied by a `PeriodicUniform` knotspec on a margin
            // for which the user left `spec.periods` empty (intending the
            // periodic B-spline path, not the Fourier path). Overwriting
            // `s.periods` here would silently flip such a margin onto the
            // Fourier path at predict time and produce a basis distinct from
            // the one used at fit time.
            s.identifiability = match identifiability_transform {
                Some(z) => TensorBSplineIdentifiability::FrozenTransform {
                    transform: z.clone(),
                },
                None => TensorBSplineIdentifiability::None,
            };
        }
        (
            SmoothBasisSpec::FactorSmooth { spec: s },
            BasisMetadata::FactorSmooth {
                knots,
                degree,
                periodic,
                group_levels,
                ..
            },
        ) => {
            s.marginal.knotspec = periodic
                .map(
                    |(domain_start, period, num_basis)| BSplineKnotSpec::PeriodicUniform {
                        data_range: (domain_start, domain_start + period),
                        num_basis,
                    },
                )
                .unwrap_or_else(|| BSplineKnotSpec::Provided(knots.clone()));
            // Restore the FROZEN marginal degree (#555 predict-replay). With
            // a `Provided(knots)` knotspec the per-margin basis count is
            // `knots.len() - (degree + 1)`, so if fit-time auto-shrink
            // lowered the marginal degree (small per-group n: cubic →
            // quadratic/linear), rebuilding with the original spec degree
            // would yield a different per-level `p` and corrupt the
            // block-diagonal replay. Mirror the sibling BySmooth arm, which
            // restores `inner.degree = *degree` for exactly this reason.
            s.marginal.degree = *degree;
            s.group_frozen_levels = Some(group_levels.clone());
        }
        (
            SmoothBasisSpec::BySmooth { smooth, by_kind },
            BasisMetadata::FactorSmooth {
                knots,
                degree,
                periodic,
                group_levels,
                ..
            },
        ) => {
            if let ByVarKind::Factor { frozen_levels, .. } = by_kind {
                *frozen_levels = Some(group_levels.clone());
            }
            if let SmoothBasisSpec::BSpline1D { spec: inner, .. } = smooth.as_mut() {
                // Issue #340: FactorSmooth metadata records the per-axis
                // spline degree directly (not optional); reflect it on
                // the inner spec so reload sees a self-consistent
                // (degree, knots) pair after fit-time auto-shrink.
                inner.degree = *degree;
                inner.knotspec = periodic
                    .map(
                        |(domain_start, period, num_basis)| BSplineKnotSpec::PeriodicUniform {
                            data_range: (domain_start, domain_start + period),
                            num_basis,
                        },
                    )
                    .unwrap_or_else(|| BSplineKnotSpec::Provided(knots.clone()));
                inner.identifiability = BSplineIdentifiability::None;
            }
        }
        (
            SmoothBasisSpec::BySmooth { smooth, by_kind },
            BasisMetadata::BySmooth { inner, levels, .. },
        ) => {
            // A `by=` smooth whose metadata is wrapped in `BySmooth`:
            // restore the frozen grouping levels, then recurse so the inner
            // basis is frozen with EXACTLY the same logic used for a
            // stand-alone term.
            if let ByVarKind::Factor { frozen_levels, .. } = by_kind
                && let Some(levels) = levels
            {
                *frozen_levels = Some(levels.clone());
            }
            freeze_smooth_basis_from_metadata(smooth, inner, term_name)?;
        }
        (SmoothBasisSpec::BySmooth { smooth, .. }, metadata) => {
            // `by=` wrapper carrying the inner basis metadata directly
            // (numeric `by`, or a factor `by` lowered to one gated block).
            // Recurse so a spatial inner basis (thin-plate, Matern, Duchon,
            // sphere, tensor, …) is frozen identically to a stand-alone
            // term. Previously only a B-spline inner was frozen here, so a
            // `s(x, by=g, bs='tp')` smooth left its data-dependent kernel
            // and eigen-truncation to be recomputed on the prediction grid,
            // crashing the predict-time design rebuild (#704).
            freeze_smooth_basis_from_metadata(smooth, metadata, term_name)?;
        }
        _ => {
            crate::bail_invalid_estim!(
                "smooth metadata/spec type mismatch while freezing term '{}'",
                term_name
            );
        }
    }
    Ok(())
}

pub fn freeze_term_collection_from_design(
    spec: &TermCollectionSpec,
    design: &TermCollectionDesign,
) -> Result<TermCollectionSpec, EstimationError> {
    if spec.smooth_terms.len() != design.smooth.terms.len() {
        crate::bail_invalid_estim!(
            "freeze mismatch: smooth spec count {} != design smooth term count {}",
            spec.smooth_terms.len(),
            design.smooth.terms.len()
        );
    }
    if spec.random_effect_terms.len() != design.random_effect_levels.len() {
        crate::bail_invalid_estim!(
            "freeze mismatch: random-effect spec count {} != design random-effect term count {}",
            spec.random_effect_terms.len(),
            design.random_effect_levels.len()
        );
    }

    let mut frozen = spec.clone();

    // ── smooth terms ────────────────────────────────────────────────────
    for (term, fitted) in frozen
        .smooth_terms
        .iter_mut()
        .zip(design.smooth.terms.iter())
    {
        // Persist joint-null absorption rotation captured at fit time so
        // save → load → predict re-applies `X_new_raw · Q` identically to
        // in-memory prediction. `None` when the smooth had no joint null
        // space, or when rotation was suppressed (shape-constrained smooths
        // whose cone geometry would not survive arbitrary orthogonal
        // rotation). Without this propagation, models reloaded from disk
        // produce wrong η at predict-time for any smooth with `Some(Q)`.
        term.joint_null_rotation = fitted.joint_null_rotation.clone();
        freeze_smooth_basis_from_metadata(&mut term.basis, &fitted.metadata, &term.name)?;
    }

    // ── random-effect terms ─────────────────────────────────────────────
    for (idx, rt) in frozen.random_effect_terms.iter_mut().enumerate() {
        let (_, kept_levels) = &design.random_effect_levels[idx];
        rt.frozen_levels = Some(kept_levels.clone());
    }

    Ok(frozen)
}

#[derive(Debug, Clone)]
struct SingleSmoothTermRealization {
    design_local: DesignMatrix,
    term: SmoothTerm,
    dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
}

impl SingleSmoothTermRealization {
    fn active_penaltyinfo(&self) -> Vec<PenaltyInfo> {
        self.term
            .penaltyinfo_local
            .iter()
            .filter(|info| info.active)
            .cloned()
            .collect()
    }
}

fn build_single_smooth_term_realization(
    data: ArrayView2<'_, f64>,
    termspec: &SmoothTermSpec,
) -> Result<SingleSmoothTermRealization, BasisError> {
    let raw = build_smooth_design(data, std::slice::from_ref(termspec))?;
    finish_single_smooth_term_realization(raw)
}

fn finish_single_smooth_term_realization(
    raw: RawSmoothDesign,
) -> Result<SingleSmoothTermRealization, BasisError> {
    let RawSmoothDesign {
        term_designs,
        dropped_penaltyinfo,
        terms,
        ..
    } = raw;
    let term = terms.into_iter().next().ok_or_else(|| {
        BasisError::InvalidInput("single-term smooth build returned no term".to_string())
    })?;
    let design = term_designs.into_iter().next().ok_or_else(|| {
        BasisError::InvalidInput("single-term smooth build returned no term design".to_string())
    })?;

    Ok(SingleSmoothTermRealization {
        design_local: design,
        term,
        dropped_penaltyinfo,
    })
}

/// Wrap a fresh `LocalSmoothTermBuild` (produced by `build_single_local_smooth_term`)
/// into a `SingleSmoothTermRealization`. Mirrors the single-term portion of
/// `build_smooth_design_withworkspace_unvalidated`, but skips the joint center
/// planner and per-term workspace fork — the realizer drives κ-only rebuilds
/// directly with its persistent workspace so basis caches survive across BFGS
/// κ proposals.
fn wrap_local_build_as_realization(
    mut local: LocalSmoothTermBuild,
    termspec: &SmoothTermSpec,
) -> Result<SingleSmoothTermRealization, String> {
    let p_local = local.dim;
    let lb_local = if local.box_reparam {
        shape_lower_bounds_local(termspec.shape, p_local)
    } else {
        None
    };

    let active_count = local.penaltyinfo.iter().filter(|info| info.active).count();
    if active_count != local.penalties.len() {
        return Err(format!(
            "internal penalty info mismatch for term '{}': active_infos={}, penalties={}",
            termspec.name,
            active_count,
            local.penalties.len()
        ));
    }

    let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
    for info in local.penaltyinfo.iter().filter(|info| !info.active) {
        dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
            termname: Some(termspec.name.clone()),
            penalty: info.clone(),
        });
    }
    for info in &local.pre_dropped_penaltyinfo {
        dropped_penaltyinfo.push(DroppedPenaltyBlockInfo {
            termname: Some(termspec.name.clone()),
            penalty: info.clone(),
        });
    }

    // Stage-2 joint-null absorption rotation, same logic as the main
    // aggregation loop in `build_smooth_design_withworkspace_unvalidated`:
    // apply Q when Some AND the smooth has no shape constraints.
    let applied_rotation: Option<crate::terms::basis::JointNullRotation> = match (
        local.joint_null_rotation.take(),
        lb_local.is_some(),
        local.linear_constraints.is_some(),
    ) {
        (Some(rot), false, false) => {
            let q = &rot.rotation;
            let dense = local
                .design
                .try_to_dense_by_chunks("joint-null absorption rotation (single realization)")
                .map_err(|e| {
                    format!(
                        "joint-null absorption rotation: dense conversion failed for term '{}': {}",
                        termspec.name, e
                    )
                })?;
            let rotated = crate::linalg::faer_ndarray::fast_ab(&dense, q);
            local.design = DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(rotated));
            local.penalties = local
                .penalties
                .into_iter()
                .map(|s_local| {
                    let qt_s = crate::linalg::faer_ndarray::fast_atb(q, &s_local);
                    crate::linalg::faer_ndarray::fast_ab(&qt_s, q)
                })
                .collect();
            local.ops = vec![None; local.penalties.len()];
            local.kronecker_factored = None;
            Some(rot)
        }
        (Some(_), _, _) => None,
        (None, _, _) => None,
    };

    let smooth_term = SmoothTerm {
        name: termspec.name.clone(),
        coeff_range: 0..p_local,
        shape: termspec.shape,
        penalties_local: local.penalties.clone(),
        nullspace_dims: local.nullspaces.clone(),
        penaltyinfo_local: local.penaltyinfo.clone(),
        metadata: local.metadata.clone(),
        lower_bounds_local: lb_local,
        linear_constraints_local: local.linear_constraints.clone(),
        kronecker_factored: local.kronecker_factored.take(),
        joint_null_rotation: applied_rotation,
    };

    Ok(SingleSmoothTermRealization {
        design_local: local.design,
        term: smooth_term,
        dropped_penaltyinfo,
    })
}

/// Extract the κ-invariant pieces of a freshly-built spatial basis — center
/// cloud (in standardized coords) and `input_scales` — and bake them into a
/// `SmoothTermSpec` whose `center_strategy` becomes `UserProvided` and whose
/// `input_scales` is `Some`. Subsequent rebuilds driven from this cached spec
/// will short-circuit `select_centers_by_strategy` (KMeans / FarthestPoint /
/// EqualMass cluster searches over n×d data) and `compute_spatial_input_scales`
/// (per-axis variance over n rows), leaving only the κ-dependent kernel
/// values and basis assembly. Returns `None` for non-spatial families or when
/// the metadata does not yet expose the required pieces (for instance when a
/// ThinPlate request was auto-promoted to Duchon during the build).
fn freeze_geometry_from_metadata(
    termspec: &SmoothTermSpec,
    metadata: &BasisMetadata,
) -> Option<SmoothTermSpec> {
    let mut frozen = termspec.clone();
    match (&mut frozen.basis, metadata) {
        (
            SmoothBasisSpec::Matern {
                spec,
                input_scales: spec_scales,
                ..
            },
            BasisMetadata::Matern {
                centers,
                input_scales: meta_scales,
                ..
            },
        ) => {
            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
            if spec_scales.is_none()
                && let Some(s) = meta_scales.clone()
            {
                *spec_scales = Some(s);
            }
            Some(frozen)
        }
        (
            SmoothBasisSpec::Duchon {
                spec,
                input_scales: spec_scales,
                ..
            },
            BasisMetadata::Duchon {
                centers,
                input_scales: meta_scales,
                ..
            },
        ) => {
            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
            if spec_scales.is_none()
                && let Some(s) = meta_scales.clone()
            {
                *spec_scales = Some(s);
            }
            Some(frozen)
        }
        (
            SmoothBasisSpec::ThinPlate {
                spec,
                input_scales: spec_scales,
                ..
            },
            BasisMetadata::ThinPlate {
                centers,
                input_scales: meta_scales,
                ..
            },
        ) => {
            spec.center_strategy = CenterStrategy::UserProvided(centers.clone());
            if spec_scales.is_none()
                && let Some(s) = meta_scales.clone()
            {
                *spec_scales = Some(s);
            }
            Some(frozen)
        }
        // Family mismatch (e.g. ThinPlate auto-promotion to Duchon) leaves the
        // cache empty; we'll retry materialization on the next κ apply.
        _ => None,
    }
}

fn rebuild_smooth_auxiliary_state(
    smooth: &mut SmoothDesign,
    dropped_penaltyinfo_by_term: &[Vec<DroppedPenaltyBlockInfo>],
) -> Result<(), String> {
    if dropped_penaltyinfo_by_term.len() != smooth.terms.len() {
        return Err(SmoothError::dimension_mismatch(format!(
            "smooth dropped-penalty cache mismatch: terms={}, dropped_sets={}",
            smooth.terms.len(),
            dropped_penaltyinfo_by_term.len()
        ))
        .into());
    }

    let total_p = smooth.total_smooth_cols();
    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
    let mut any_bounds = false;
    let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
    let mut linear_constraint_b: Vec<f64> = Vec::new();

    for term in &smooth.terms {
        let range = term.coeff_range.clone();
        if let Some(lb_local) = term.lower_bounds_local.as_ref() {
            if lb_local.len() != range.len() {
                return Err(SmoothError::dimension_mismatch(format!(
                    "smooth lower-bound cache mismatch for term '{}': bounds={}, coeffs={}",
                    term.name,
                    lb_local.len(),
                    range.len()
                ))
                .into());
            }
            coefficient_lower_bounds
                .slice_mut(s![range.clone()])
                .assign(lb_local);
            any_bounds = true;
        }
        if let Some(lin_local) = term.linear_constraints_local.as_ref() {
            if lin_local.a.ncols() != range.len() {
                return Err(SmoothError::dimension_mismatch(format!(
                    "smooth linear-constraint cache mismatch for term '{}': cols={}, coeffs={}",
                    term.name,
                    lin_local.a.ncols(),
                    range.len()
                ))
                .into());
            }
            for r in 0..lin_local.a.nrows() {
                let mut row = Array1::<f64>::zeros(total_p);
                row.slice_mut(s![range.clone()]).assign(&lin_local.a.row(r));
                linear_constraintrows.push(row);
                linear_constraint_b.push(lin_local.b[r]);
            }
        }
    }

    smooth.coefficient_lower_bounds = if any_bounds {
        Some(coefficient_lower_bounds)
    } else {
        None
    };
    smooth.linear_constraints = if linear_constraintrows.is_empty() {
        None
    } else {
        let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), total_p));
        for (i, row) in linear_constraintrows.iter().enumerate() {
            a.row_mut(i).assign(row);
        }
        Some(LinearInequalityConstraints {
            a,
            b: Array1::from_vec(linear_constraint_b),
        })
    };
    smooth.dropped_penaltyinfo = dropped_penaltyinfo_by_term
        .iter()
        .flat_map(|infos| infos.iter().cloned())
        .collect();
    Ok(())
}

fn rebuild_term_collection_auxiliary_state(
    spec: &TermCollectionSpec,
    design: &mut TermCollectionDesign,
) -> Result<(), String> {
    if spec.linear_terms.len() != design.linear_ranges.len() {
        return Err(SmoothError::dimension_mismatch(format!(
            "term-collection linear bookkeeping mismatch: spec_terms={}, design_ranges={}",
            spec.linear_terms.len(),
            design.linear_ranges.len()
        ))
        .into());
    }

    let p_total = design.design.ncols();
    let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
    let mut any_bounds = false;
    let mut linear_constraintrows: Vec<Array1<f64>> = Vec::new();
    let mut linear_constraint_b: Vec<f64> = Vec::new();

    for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
        if range.len() != 1 {
            return Err(SmoothError::dimension_mismatch(format!(
                "linear term '{}' expected one coefficient column, found {}",
                linear.name,
                range.len()
            ))
            .into());
        }
        let col = range.start;
        if let Some(lb) = linear.coefficient_min {
            let mut row = Array1::<f64>::zeros(p_total);
            row[col] = 1.0;
            linear_constraintrows.push(row);
            linear_constraint_b.push(lb);
        }
        if let Some(ub) = linear.coefficient_max {
            let mut row = Array1::<f64>::zeros(p_total);
            row[col] = -1.0;
            linear_constraintrows.push(row);
            linear_constraint_b.push(-ub);
        }
    }

    if let Some(lb_smooth) = design.smooth.coefficient_lower_bounds.as_ref() {
        if lb_smooth.len() != design.smooth.total_smooth_cols() {
            return Err(SmoothError::dimension_mismatch(format!(
                "smooth lower-bound width mismatch: bounds={}, smooth_cols={}",
                lb_smooth.len(),
                design.smooth.total_smooth_cols()
            ))
            .into());
        }
        coefficient_lower_bounds
            .slice_mut(s![
                smooth_start..(smooth_start + design.smooth.total_smooth_cols())
            ])
            .assign(lb_smooth);
        any_bounds = true;
    }
    if let Some(lin_smooth) = design.smooth.linear_constraints.as_ref() {
        if lin_smooth.a.ncols() != design.smooth.total_smooth_cols() {
            return Err(SmoothError::dimension_mismatch(format!(
                "smooth linear-constraint width mismatch: cols={}, smooth_cols={}",
                lin_smooth.a.ncols(),
                design.smooth.total_smooth_cols()
            ))
            .into());
        }
        let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
        a_global
            .slice_mut(s![
                ..,
                smooth_start..(smooth_start + design.smooth.total_smooth_cols())
            ])
            .assign(&lin_smooth.a);
        for r in 0..a_global.nrows() {
            linear_constraintrows.push(a_global.row(r).to_owned());
            linear_constraint_b.push(lin_smooth.b[r]);
        }
    }

    let lower_bound_constraints = if any_bounds {
        linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
    } else {
        None
    };
    let explicit_linear_constraints = if linear_constraintrows.is_empty() {
        None
    } else {
        let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
        for (i, row) in linear_constraintrows.iter().enumerate() {
            a.row_mut(i).assign(row);
        }
        Some(LinearInequalityConstraints {
            a,
            b: Array1::from_vec(linear_constraint_b),
        })
    };

    design.coefficient_lower_bounds = if any_bounds {
        Some(coefficient_lower_bounds)
    } else {
        None
    };
    design.linear_constraints =
        merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints);
    design.dropped_penaltyinfo = design.smooth.dropped_penaltyinfo.clone();
    Ok(())
}

fn theta_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
    left.len() == right.len()
        && left
            .iter()
            .zip(right.iter())
            .all(|(&l, &r)| l.to_bits() == r.to_bits())
}

fn latent_values_match(left: &Array1<f64>, right: &Array1<f64>) -> bool {
    theta_values_match(left, right)
}

fn spatial_aniso_matches(left: Option<&[f64]>, right: Option<&[f64]>) -> bool {
    match (left, right) {
        (None, None) => true,
        (Some(a), Some(b)) => {
            a.len() == b.len()
                && a.iter()
                    .zip(b.iter())
                    .all(|(&x, &y)| x.to_bits() == y.to_bits())
        }
        _ => false,
    }
}

fn spatial_length_scale_matches(left: Option<f64>, right: Option<f64>) -> bool {
    match (left, right) {
        (None, None) => true,
        (Some(a), Some(b)) => a.to_bits() == b.to_bits(),
        _ => false,
    }
}

struct FrozenTermCollectionIncrementalRealizer<'d> {
    data: ArrayView2<'d, f64>,
    spec: TermCollectionSpec,
    design: TermCollectionDesign,
    fixed_blocks: Vec<DesignBlock>,
    dropped_penaltyinfo_by_term: Vec<Vec<DroppedPenaltyBlockInfo>>,
    smooth_penalty_ranges: Vec<Range<usize>>,
    full_penalty_ranges: Vec<Range<usize>>,
    /// Persistent workspace for basis cache reuse across κ proposals.
    /// Distance matrices are cached here so they're computed once and
    /// reused across repeated `apply_log_kappa_to_term` calls.
    basisworkspace: crate::basis::BasisWorkspace,
    /// Per-term cached realization geometry for incremental κ updates.
    ///
    /// On the first κ-driven rebuild of term `i`, this slot is populated with a
    /// `SmoothTermSpec` whose κ-invariant geometry — center cloud (as
    /// `CenterStrategy::UserProvided`) and `input_scales` — has been frozen
    /// out of the realized basis metadata. Subsequent
    /// `apply_log_kappa_to_term` calls reuse this spec, mutating only the
    /// κ / aniso fields. This short-circuits `select_centers_by_strategy`
    /// (KMeans / FarthestPoint / EqualMass cluster searches over the n×d data
    /// matrix) and `compute_spatial_input_scales` (per-axis variance pass
    /// over n rows) on every BFGS κ-eval, leaving the kernel-value pass and
    /// basis assembly as the only work.
    spatial_realization_geometry: Vec<Option<SmoothTermSpec>>,
    /// Monotonic counter incremented every time `apply_log_kappa` actually
    /// rebuilds the realized design / smooth penalties. Read by the
    /// design-revision-counter fast path in `ExternalJointHyperEvaluator`
    /// to skip redundant canonical-penalty rebuilds and cache wipes when
    /// the outer BFGS loop probes the same ψ twice in a row.
    design_revision: u64,
}

impl<'d> std::fmt::Debug for FrozenTermCollectionIncrementalRealizer<'d> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FrozenTermCollectionIncrementalRealizer")
            .field("data_shape", &(self.data.nrows(), self.data.ncols()))
            .field("fixed_blocks", &self.fixed_blocks.len())
            .finish_non_exhaustive()
    }
}

impl<'d> FrozenTermCollectionIncrementalRealizer<'d> {
    fn new(
        data: ArrayView2<'d, f64>,
        spec: TermCollectionSpec,
        design: TermCollectionDesign,
    ) -> Result<Self, String> {
        if spec.smooth_terms.len() != design.smooth.terms.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer smooth term mismatch: spec_terms={}, design_terms={}",
                spec.smooth_terms.len(),
                design.smooth.terms.len()
            ))
            .into());
        }

        let mut smooth_cursor = 0usize;
        let mut smooth_penalty_ranges = Vec::with_capacity(design.smooth.terms.len());
        for term in &design.smooth.terms {
            let next = smooth_cursor + term.penalties_local.len();
            smooth_penalty_ranges.push(smooth_cursor..next);
            smooth_cursor = next;
        }
        if smooth_cursor != design.smooth.penalties.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer smooth penalty mismatch: ranged={}, actual={}",
                smooth_cursor,
                design.smooth.penalties.len()
            ))
            .into());
        }

        let fixed_penalty_offset = design
            .penalties
            .len()
            .checked_sub(design.smooth.penalties.len())
            .ok_or_else(|| {
                "incremental realizer encountered invalid penalty bookkeeping".to_string()
            })?;
        let full_penalty_ranges = smooth_penalty_ranges
            .iter()
            .map(|range| (fixed_penalty_offset + range.start)..(fixed_penalty_offset + range.end))
            .collect::<Vec<_>>();
        let fixed_blocks = build_term_collection_fixed_blocks(data, &spec)
            .map_err(|e| format!("failed to cache fixed term-collection blocks: {e}"))?;

        let mut dropped_penaltyinfo_by_term = Vec::with_capacity(spec.smooth_terms.len());
        for (term_idx, termspec) in spec.smooth_terms.iter().enumerate() {
            let realization =
                build_single_smooth_term_realization(data, termspec).map_err(|e| {
                    format!(
                        "failed to build cached realization for smooth term '{}' (index {}): {e}",
                        termspec.name, term_idx
                    )
                })?;
            let expected_cols = design.smooth.terms[term_idx].coeff_range.len();
            if realization.design_local.ncols() != expected_cols {
                return Err(SmoothError::dimension_mismatch(format!(
                    "cached realization width mismatch for term '{}': cached_cols={}, design_cols={}",
                    termspec.name,
                    realization.design_local.ncols(),
                    expected_cols
                ))
                .into());
            }
            if realization.active_penaltyinfo().len()
                != design.smooth.terms[term_idx].penalties_local.len()
            {
                return Err(SmoothError::dimension_mismatch(format!(
                    "cached realization penalty mismatch for term '{}': cached_penalties={}, design_penalties={}",
                    termspec.name,
                    realization.active_penaltyinfo().len(),
                    design.smooth.terms[term_idx].penalties_local.len()
                ))
                .into());
            }
            dropped_penaltyinfo_by_term.push(realization.dropped_penaltyinfo);
        }

        let geometry_slots = spec.smooth_terms.len();
        Ok(Self {
            data,
            spec,
            design,
            fixed_blocks,
            dropped_penaltyinfo_by_term,
            smooth_penalty_ranges,
            full_penalty_ranges,
            basisworkspace: crate::basis::BasisWorkspace::new(),
            spatial_realization_geometry: vec![None; geometry_slots],
            design_revision: 0,
        })
    }

    fn design_revision(&self) -> u64 {
        self.design_revision
    }

    fn spec(&self) -> &TermCollectionSpec {
        &self.spec
    }

    fn design(&self) -> &TermCollectionDesign {
        &self.design
    }

    fn apply_log_kappa(
        &mut self,
        log_kappa: &SpatialLogKappaCoords,
        term_indices: &[usize],
    ) -> Result<(), String> {
        if term_indices.len() != log_kappa.dims_per_term().len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer log-kappa term mismatch: term_indices={}, dims_per_term={}",
                term_indices.len(),
                log_kappa.dims_per_term().len()
            ))
            .into());
        }

        let mut any_changed = false;
        for (slot, &term_idx) in term_indices.iter().enumerate() {
            any_changed |= self.apply_log_kappa_to_term(term_idx, log_kappa.term_slice(slot))?;
        }

        if any_changed {
            self.refresh_full_design_operator()?;
            rebuild_smooth_auxiliary_state(
                &mut self.design.smooth,
                &self.dropped_penaltyinfo_by_term,
            )?;
            rebuild_term_collection_auxiliary_state(&self.spec, &mut self.design)?;
            self.design_revision = self.design_revision.wrapping_add(1);
        }
        Ok(())
    }

    fn apply_log_kappa_to_term(&mut self, term_idx: usize, psi: &[f64]) -> Result<bool, String> {
        if !spatial_term_supports_hyper_optimization(&self.spec, term_idx) {
            return Err(SmoothError::invalid_config(format!(
                "incremental realizer term {term_idx} does not expose spatial hyperparameters"
            ))
            .into());
        }
        let current_length_scale = get_spatial_length_scale(&self.spec, term_idx);
        let current_aniso = get_spatial_aniso_log_scales(&self.spec, term_idx);
        let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
        let same_length = spatial_length_scale_matches(current_length_scale, next_length_scale);
        let same_aniso = spatial_aniso_matches(current_aniso.as_deref(), next_aniso.as_deref());
        if same_length && same_aniso {
            return Ok(false);
        }

        if let Some(length_scale) = next_length_scale {
            set_spatial_length_scale(&mut self.spec, term_idx, length_scale)
                .map_err(|e| e.to_string())?;
        }
        if let Some(eta) = next_aniso.clone() {
            set_spatial_aniso_log_scales(&mut self.spec, term_idx, eta)
                .map_err(|e| e.to_string())?;
        }

        // Pick the spec to drive the rebuild. If the per-term geometry cache
        // is populated, it carries already-resolved centers
        // (`CenterStrategy::UserProvided`) and frozen `input_scales`; reusing
        // it short-circuits `select_centers_by_strategy` (KMeans /
        // FarthestPoint / EqualMass cluster searches) and
        // `compute_spatial_input_scales` (per-axis variance over n rows) in
        // the family builders. Centers in the cached spec live in
        // standardized coordinates (matching the cached `input_scales`), so
        // the same standardization + kernel path runs without recomputation
        // of the geometry.
        let geometry_slot = self
            .spatial_realization_geometry
            .get(term_idx)
            .ok_or_else(|| format!("incremental realizer geometry slot {term_idx} out of range"))?;
        let mut build_spec = match geometry_slot {
            Some(cached) => cached.clone(),
            None => self
                .spec
                .smooth_terms
                .get(term_idx)
                .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
                .clone(),
        };
        if let Some(length_scale) = next_length_scale {
            set_single_term_spatial_length_scale(&mut build_spec, length_scale)
                .map_err(|e| e.to_string())?;
        }
        if let Some(eta) = next_aniso {
            set_single_term_spatial_aniso_log_scales(&mut build_spec, eta)
                .map_err(|e| e.to_string())?;
        }

        let termname = build_spec.name.clone();
        let local = build_single_local_smooth_term(
            self.data,
            &build_spec,
            &mut self.basisworkspace,
        )
        .map_err(|e| {
            format!(
                "failed to rebuild smooth term '{termname}' during incremental κ realization: {e}"
            )
        })?;

        // Populate the geometry cache from the realized metadata on first use.
        // Family auto-promotion (ThinPlate -> Duchon) is detected as a basis /
        // metadata mismatch in `freeze_geometry_from_metadata` and leaves the
        // cache empty so the next call re-tries with the (now stable) family.
        if self.spatial_realization_geometry[term_idx].is_none()
            && let Some(frozen) = freeze_geometry_from_metadata(&build_spec, &local.metadata)
        {
            self.spatial_realization_geometry[term_idx] = Some(frozen);
        }

        let realization = wrap_local_build_as_realization(local, &build_spec)?;
        self.replace_term_realization(term_idx, realization)?;
        Ok(true)
    }

    fn replace_term_realization(
        &mut self,
        term_idx: usize,
        realization: SingleSmoothTermRealization,
    ) -> Result<(), String> {
        let t_replace = std::time::Instant::now();
        let SingleSmoothTermRealization {
            design_local,
            term,
            dropped_penaltyinfo,
        } = realization;
        let SmoothTerm {
            name,
            penalties_local,
            nullspace_dims,
            penaltyinfo_local,
            metadata,
            lower_bounds_local,
            linear_constraints_local,
            joint_null_rotation,
            ..
        } = term;
        let coeff_range = self
            .design
            .smooth
            .terms
            .get(term_idx)
            .ok_or_else(|| format!("incremental realizer smooth term {term_idx} out of range"))?
            .coeff_range
            .clone();
        if design_local.ncols() != coeff_range.len() {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer width mismatch for term {}: rebuilt_cols={}, cached_cols={}",
                term_idx,
                design_local.ncols(),
                coeff_range.len()
            ))
            .into());
        }
        if design_local.nrows() != self.design.design.nrows() {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer row mismatch for term {}: rebuilt_rows={}, design_rows={}",
                term_idx,
                design_local.nrows(),
                self.design.design.nrows()
            ))
            .into());
        }

        let active_penaltyinfo = penaltyinfo_local
            .iter()
            .filter(|info| info.active)
            .cloned()
            .collect::<Vec<_>>();
        let smooth_penalty_range = self
            .smooth_penalty_ranges
            .get(term_idx)
            .ok_or_else(|| {
                format!("incremental realizer missing smooth penalty range for term {term_idx}")
            })?
            .clone();
        let full_penalty_range = self
            .full_penalty_ranges
            .get(term_idx)
            .ok_or_else(|| {
                format!("incremental realizer missing full penalty range for term {term_idx}")
            })?
            .clone();
        if active_penaltyinfo.len() != smooth_penalty_range.len()
            || penalties_local.len() != smooth_penalty_range.len()
            || nullspace_dims.len() != smooth_penalty_range.len()
        {
            return Err(SmoothError::dimension_mismatch(format!(
                "incremental realizer topology changed for term '{}': penalties={}, infos={}, nullspaces={}, cached_penalties={}",
                name,
                penalties_local.len(),
                active_penaltyinfo.len(),
                nullspace_dims.len(),
                smooth_penalty_range.len()
            ))
            .into());
        }

        self.design.smooth.term_designs[term_idx] = design_local;

        for (offset, penalty_local) in penalties_local.iter().enumerate() {
            let smooth_penalty_idx = smooth_penalty_range.start + offset;
            let full_penalty_idx = full_penalty_range.start + offset;
            let nullspace_dim = nullspace_dims[offset];
            let penalty_info = active_penaltyinfo[offset].clone();

            if penalty_local.nrows() != coeff_range.len()
                || penalty_local.ncols() != coeff_range.len()
            {
                return Err(SmoothError::dimension_mismatch(format!(
                    "incremental realizer penalty shape mismatch for term '{}' penalty {}: \
                     penalty is {}x{} but coeff_range has {} columns",
                    name,
                    offset,
                    penalty_local.nrows(),
                    penalty_local.ncols(),
                    coeff_range.len()
                ))
                .into());
            }

            let smooth_penalty = self
                .design
                .smooth
                .penalties
                .get_mut(smooth_penalty_idx)
                .ok_or_else(|| {
                    format!(
                        "incremental realizer smooth penalty {} out of range for term {}",
                        smooth_penalty_idx, term_idx
                    )
                })?;
            // With per-term block-local penalties, col_range already targets
            // this specific term, so .local is p_k × p_k.
            smooth_penalty.local.assign(penalty_local);

            let full_bp = self
                .design
                .penalties
                .get_mut(full_penalty_idx)
                .ok_or_else(|| {
                    format!(
                        "incremental realizer full penalty {} out of range for term {}",
                        full_penalty_idx, term_idx
                    )
                })?;
            // With per-term block-local penalties, col_range already targets
            // this specific term, so .local is p_k × p_k.
            full_bp.local.assign(penalty_local);

            self.design.smooth.nullspace_dims[smooth_penalty_idx] = nullspace_dim;
            self.design.nullspace_dims[full_penalty_idx] = nullspace_dim;

            self.design.smooth.penaltyinfo[smooth_penalty_idx].global_index = smooth_penalty_idx;
            self.design.smooth.penaltyinfo[smooth_penalty_idx].termname = Some(name.clone());
            self.design.smooth.penaltyinfo[smooth_penalty_idx].penalty = penalty_info.clone();

            self.design.penaltyinfo[full_penalty_idx].global_index = full_penalty_idx;
            self.design.penaltyinfo[full_penalty_idx].termname = Some(name.clone());
            self.design.penaltyinfo[full_penalty_idx].penalty = penalty_info;
        }

        let target_term = self.design.smooth.terms.get_mut(term_idx).ok_or_else(|| {
            format!("incremental realizer smooth term {term_idx} disappeared during replacement")
        })?;
        target_term.penalties_local = penalties_local;
        target_term.nullspace_dims = nullspace_dims;
        target_term.penaltyinfo_local = penaltyinfo_local;
        target_term.metadata = metadata;
        target_term.lower_bounds_local = lower_bounds_local;
        target_term.linear_constraints_local = linear_constraints_local;
        target_term.joint_null_rotation = joint_null_rotation;
        self.dropped_penaltyinfo_by_term[term_idx] = dropped_penaltyinfo;
        log::info!(
            "[STAGE] smooth basis rebuild (term {}, '{}', cols={}): {:.3}s",
            term_idx,
            target_term.name,
            coeff_range.len(),
            t_replace.elapsed().as_secs_f64(),
        );
        Ok(())
    }

    fn refresh_full_design_operator(&mut self) -> Result<(), String> {
        let mut blocks = Vec::<DesignBlock>::with_capacity(
            self.fixed_blocks.len() + self.design.smooth.term_designs.len(),
        );
        blocks.extend(self.fixed_blocks.iter().cloned());
        for term_design in &self.design.smooth.term_designs {
            blocks.push(DesignBlock::from(term_design));
        }
        self.design.design = assemble_term_collection_design_matrix(blocks)
            .map_err(|e| format!("failed to refresh term-collection design: {e}"))?;
        Ok(())
    }
}

fn build_term_collection_fixed_blocks(
    data: ArrayView2<'_, f64>,
    spec: &TermCollectionSpec,
) -> Result<Vec<DesignBlock>, BasisError> {
    let mut blocks = Vec::<DesignBlock>::new();
    blocks.push(DesignBlock::Intercept(data.nrows()));

    if !spec.linear_terms.is_empty() {
        for linear in &spec.linear_terms {
            for &col in &linear.effective_feature_cols() {
                if col >= data.ncols() {
                    crate::bail_dim_basis!(
                        "linear term '{}' feature column {} out of bounds for {} columns",
                        linear.name,
                        col,
                        data.ncols()
                    );
                }
            }
        }
        let mut linear_block = Array2::<f64>::zeros((data.nrows(), spec.linear_terms.len()));
        for (j, linear) in spec.linear_terms.iter().enumerate() {
            let cols = linear.effective_feature_cols();
            let mut col_view = linear_block.column_mut(j);
            col_view.assign(&data.column(cols[0]));
            for &c in cols.iter().skip(1) {
                let factor = data.column(c);
                for (out, &x) in col_view.iter_mut().zip(factor.iter()) {
                    *out *= x;
                }
            }
        }
        blocks.push(DesignBlock::Dense(crate::matrix::DenseDesignMatrix::from(
            linear_block,
        )));
    }

    for term in &spec.random_effect_terms {
        let block = build_random_effect_block(data, term)?;
        let re_op = RandomEffectOperator::new(block.group_ids, block.num_groups);
        blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
    }

    Ok(blocks)
}

// ---------------------------------------------------------------------------
// N-block spatial length-scale optimizer.
// ---------------------------------------------------------------------------

pub struct SpatialLengthScaleOptimizationResult<FitOut> {
    pub resolved_specs: Vec<TermCollectionSpec>,
    pub designs: Vec<TermCollectionDesign>,
    pub fit: FitOut,
}

/// Exact-joint hyper-parameter setup for N-block spatial length-scale optimization.
#[derive(Debug, Clone)]
pub struct ExactJointHyperSetup {
    rho0: Array1<f64>,
    rho_lower: Array1<f64>,
    rho_upper: Array1<f64>,
    log_kappa0: SpatialLogKappaCoords,
    log_kappa_lower: SpatialLogKappaCoords,
    log_kappa_upper: SpatialLogKappaCoords,
    auxiliary0: Array1<f64>,
    auxiliary_lower: Array1<f64>,
    auxiliary_upper: Array1<f64>,
}

impl ExactJointHyperSetup {
    fn sanitize_rho_seed(
        rho0: Array1<f64>,
        rho_lower: &Array1<f64>,
        rho_upper: &Array1<f64>,
    ) -> Array1<f64> {
        Array1::from_iter(rho0.iter().enumerate().map(|(idx, &value)| {
            let lo = rho_lower[idx];
            let hi = rho_upper[idx];
            let fallback = 0.0_f64.clamp(lo, hi);
            if value.is_finite() {
                value.clamp(lo, hi)
            } else {
                fallback
            }
        }))
    }

    pub(crate) fn new(
        rho0: Array1<f64>,
        rho_lower: Array1<f64>,
        rho_upper: Array1<f64>,
        log_kappa0: SpatialLogKappaCoords,
        log_kappa_lower: SpatialLogKappaCoords,
        log_kappa_upper: SpatialLogKappaCoords,
    ) -> Self {
        let rho0 = Self::sanitize_rho_seed(rho0, &rho_lower, &rho_upper);
        Self {
            rho0,
            rho_lower,
            rho_upper,
            log_kappa0,
            log_kappa_lower,
            log_kappa_upper,
            auxiliary0: Array1::zeros(0),
            auxiliary_lower: Array1::zeros(0),
            auxiliary_upper: Array1::zeros(0),
        }
    }

    pub(crate) fn with_auxiliary(
        mut self,
        auxiliary0: Array1<f64>,
        auxiliary_lower: Array1<f64>,
        auxiliary_upper: Array1<f64>,
    ) -> Self {
        assert_eq!(
            auxiliary0.len(),
            auxiliary_lower.len(),
            "auxiliary lower bound length mismatch"
        );
        assert_eq!(
            auxiliary0.len(),
            auxiliary_upper.len(),
            "auxiliary upper bound length mismatch"
        );
        self.auxiliary0 = Self::sanitize_rho_seed(auxiliary0, &auxiliary_lower, &auxiliary_upper);
        self.auxiliary_lower = auxiliary_lower;
        self.auxiliary_upper = auxiliary_upper;
        self
    }

    pub(crate) fn rho_dim(&self) -> usize {
        self.rho0.len()
    }

    pub(crate) fn log_kappa_dim(&self) -> usize {
        self.log_kappa0.len()
    }

    pub(crate) fn auxiliary_dim(&self) -> usize {
        self.auxiliary0.len()
    }

    pub(crate) fn theta0(&self) -> Array1<f64> {
        let mut out =
            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho0);
        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
            .assign(self.log_kappa0.as_array());
        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
            .assign(&self.auxiliary0);
        out
    }

    pub(crate) fn lower(&self) -> Array1<f64> {
        let mut out =
            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_lower);
        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
            .assign(self.log_kappa_lower.as_array());
        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
            .assign(&self.auxiliary_lower);
        out
    }

    pub(crate) fn upper(&self) -> Array1<f64> {
        let mut out =
            Array1::<f64>::zeros(self.rho_dim() + self.log_kappa_dim() + self.auxiliary_dim());
        out.slice_mut(s![..self.rho_dim()]).assign(&self.rho_upper);
        out.slice_mut(s![self.rho_dim()..self.rho_dim() + self.log_kappa_dim()])
            .assign(self.log_kappa_upper.as_array());
        out.slice_mut(s![self.rho_dim() + self.log_kappa_dim()..])
            .assign(&self.auxiliary_upper);
        out
    }

    /// Per-term dimensionality layout for the psi block.
    pub(crate) fn log_kappa_dims_per_term(&self) -> Vec<usize> {
        self.log_kappa0.dims_per_term().to_vec()
    }
}

/// N-block design cache for exact-joint spatial length-scale optimization.
///
/// Each block owns a `FrozenTermCollectionIncrementalRealizer` and a list of
/// spatial term indices within that block's spec. The cache splits the
/// combined psi vector into per-block slices using precomputed offsets.
struct ExactJointDesignCache<'d> {
    realizers: Vec<FrozenTermCollectionIncrementalRealizer<'d>>,
    block_term_indices: Vec<Vec<usize>>,
    current_theta: Option<Array1<f64>>,
    last_cost: Option<f64>,
    last_eval: Option<(
        f64,
        Array1<f64>,
        crate::solver::outer_strategy::HessianResult,
    )>,
    rho_dim: usize,
    all_dims: Vec<usize>,
    log_kappa_dim: usize,
    block_term_counts: Vec<usize>,
}

impl<'d> ExactJointDesignCache<'d> {
    fn new(
        data: ArrayView2<'d, f64>,
        blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)>,
        rho_dim: usize,
        all_dims: Vec<usize>,
    ) -> Result<Self, String> {
        let n_blocks = blocks.len();
        let mut realizers = Vec::with_capacity(n_blocks);
        let mut block_term_indices = Vec::with_capacity(n_blocks);
        let mut block_term_counts = Vec::with_capacity(n_blocks);

        for (spec, design, terms) in blocks {
            block_term_counts.push(terms.len());
            block_term_indices.push(terms);
            realizers.push(FrozenTermCollectionIncrementalRealizer::new(
                data, spec, design,
            )?);
        }

        Ok(Self {
            realizers,
            block_term_indices,
            current_theta: None,
            last_cost: None,
            last_eval: None,
            rho_dim,
            log_kappa_dim: all_dims.iter().sum(),
            all_dims,
            block_term_counts,
        })
    }

    fn ensure_theta(&mut self, theta: &Array1<f64>) -> Result<(), String> {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
        {
            return Ok(());
        }

        let t_ensure = std::time::Instant::now();
        let kappa_theta_len = self.rho_dim + self.log_kappa_dim;
        if theta.len() < kappa_theta_len {
            return Err(SmoothError::dimension_mismatch(format!(
                "exact-joint theta length mismatch: got {}, expected at least {} (rho_dim={}, log_kappa_dim={})",
                theta.len(),
                kappa_theta_len,
                self.rho_dim,
                self.log_kappa_dim
            ))
            .into());
        }
        let theta_kappa = theta.slice(s![..kappa_theta_len]).to_owned();
        let full_log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
            &theta_kappa,
            self.rho_dim,
            self.all_dims.clone(),
        );

        // Split the full log_kappa into per-block sub-coords using split_at.
        // We split from the front iteratively: after extracting block 0..N-2,
        // the remainder is the last block.
        let n = self.realizers.len();
        let mut remaining = full_log_kappa;
        for block_idx in 0..n {
            let count = self.block_term_counts[block_idx];
            if block_idx < n - 1 {
                let (block_lk, rest) = remaining.split_at(count);
                self.realizers[block_idx]
                    .apply_log_kappa(&block_lk, &self.block_term_indices[block_idx])?;
                remaining = rest;
            } else {
                // Last block gets the remainder.
                self.realizers[block_idx]
                    .apply_log_kappa(&remaining, &self.block_term_indices[block_idx])?;
            }
        }

        log::info!(
            "[STAGE] ensure_theta (n-block, {} blocks, {} realizers): {:.3}s",
            n,
            self.realizers.len(),
            t_ensure.elapsed().as_secs_f64(),
        );
        self.current_theta = Some(theta.clone());
        self.last_cost = None;
        self.last_eval = None;
        Ok(())
    }

    impl_exact_joint_theta_memo!();

    /// Cache a cost-only result. Called after `ensure_theta(theta)` for
    /// line-search probes that pay only for the cost evaluation. We
    /// intentionally do not populate `last_eval` because no gradient was
    /// computed; the next outer evaluation at this θ will recompute
    /// (V, ∇V) via `evaluate_with_order` if the optimizer asks for it.
    fn store_cost_only(&mut self, theta: &Array1<f64>, cost: f64) {
        if self
            .current_theta
            .as_ref()
            .is_some_and(|cached| theta_values_match(cached, theta))
        {
            self.last_cost = Some(cost);
        }
    }

    fn specs(&self) -> Vec<&TermCollectionSpec> {
        self.realizers.iter().map(|r| r.spec()).collect()
    }

    fn designs(&self) -> Vec<&TermCollectionDesign> {
        self.realizers.iter().map(|r| r.design()).collect()
    }

    /// Combined monotonic design revision across all per-block realizers.
    ///
    /// Mirrors `SingleBlockExactJointDesignCache::design_revision` for the
    /// n-block exact-joint path. Each realizer's `design_revision` counter
    /// advances iff `apply_log_kappa` actually rebuilt that block's realized
    /// design / smooth penalties; the wrapping sum therefore changes iff
    /// *any* block rebuilt. Equal values across two calls imply no realizer
    /// has been rebuilt in between, which is the invariant the
    /// `ExternalJointHyperEvaluator` canonical-penalty fast path needs.
    fn design_revision(&self) -> u64 {
        self.realizers
            .iter()
            .fold(0u64, |acc, r| acc.wrapping_add(r.design_revision()))
    }
}

pub(crate) fn seed_risk_profile_for_likelihood_family(
    family: &LikelihoodSpec,
) -> crate::seeding::SeedRiskProfile {
    match &family.response {
        ResponseFamily::Gaussian => crate::seeding::SeedRiskProfile::Gaussian,
        ResponseFamily::RoystonParmar => crate::seeding::SeedRiskProfile::Survival,
        ResponseFamily::Binomial
        | ResponseFamily::Poisson
        | ResponseFamily::Tweedie { .. }
        | ResponseFamily::NegativeBinomial { .. }
        | ResponseFamily::Beta { .. }
        | ResponseFamily::Gamma => crate::seeding::SeedRiskProfile::GeneralizedLinear,
    }
}

pub(crate) fn exact_joint_multistart_outer_problem(
    theta0: &Array1<f64>,
    lower: &Array1<f64>,
    upper: &Array1<f64>,
    rho_dim: usize,
    auxiliary_dim: usize,
    n_params: usize,
    gradient: crate::solver::outer_strategy::Derivative,
    hessian: crate::solver::outer_strategy::DeclaredHessianForm,
    prefer_gradient_only: bool,
    disable_fixed_point: bool,
    risk_profile: crate::seeding::SeedRiskProfile,
    tolerance: f64,
    max_iter: usize,
    // BFGS step caps split by parameter type. `bfgs_step_cap` (rho-axis cap)
    // bounds first-trial moves on log-λ; documented natural step is ≈ 5.
    // `bfgs_step_cap_psi` bounds moves on the trailing `auxiliary_dim`
    // psi-axes (kappa / aniso-log-scales), where ≈ ln 2 keeps the kernel
    // scale from oscillating across orders of magnitude per iter. Using a
    // single uniform cap (the old API) starved rho on the survival-marg-slope
    // joint solver because the psi-calibrated value (`ln 2 ≈ 0.69`) was
    // applied to log-λ, where |d|≈5 is the natural quasi-Newton magnitude.
    bfgs_step_cap: Option<f64>,
    bfgs_step_cap_psi: Option<f64>,
    screening_cap: Option<Arc<AtomicUsize>>,
) -> crate::solver::outer_strategy::OuterProblem {
    let mut seed_heuristic = theta0.to_vec();
    for value in &mut seed_heuristic[..rho_dim] {
        *value = value.exp();
    }
    let mut problem = crate::solver::outer_strategy::OuterProblem::new(n_params)
        .with_gradient(gradient)
        .with_hessian(hessian)
        .with_prefer_gradient_only(prefer_gradient_only)
        .with_disable_fixed_point(disable_fixed_point)
        // Re-enable the automatic fallback ladder for exact joint spatial
        // problems. It was previously `Disabled` to suppress a geo-bench
        // fallback bug where HybridEFS ψ stagnation degraded silently to
        // BfgsApprox on a Charbonnier surface. With the ψ-stagnation guard
        // in OuterFixedPointBridge (`MAX_CONSECUTIVE_PSI_STAGNATION`) the
        // bridge now surfaces `EFS_FIRST_ORDER_FALLBACK_MARKER` when ψ
        // stationarity cannot be enforced, so the ladder routes correctly
        // to a joint gradient-based solver instead of grinding HybridEFS
        // for thousands of iterations.
        .with_fallback_policy(crate::solver::outer_strategy::FallbackPolicy::Automatic)
        .with_psi_dim(auxiliary_dim)
        .with_tolerance(tolerance)
        .with_max_iter(max_iter)
        .with_bounds(lower.clone(), upper.clone())
        .with_initial_rho(theta0.clone())
        .with_bfgs_step_cap(bfgs_step_cap)
        .with_bfgs_step_cap_psi(bfgs_step_cap_psi)
        .with_seed_config(crate::seeding::SeedConfig {
            max_seeds: 4,
            seed_budget: 2,
            risk_profile,
            num_auxiliary_trailing: auxiliary_dim,
            ..Default::default()
        })
        .with_rho_bound(12.0)
        .with_heuristic_lambdas(seed_heuristic);
    if let Some(screening_cap) = screening_cap {
        problem = problem
            .with_screening_cap(screening_cap)
            .with_screen_initial_rho(true);
    }
    problem
}

pub fn optimize_spatial_length_scale_exact_joint<FitOut, FitFn, ExactFn, ExactEfsFn, SeedFn>(
    data: ArrayView2<'_, f64>,
    block_specs: &[TermCollectionSpec],
    block_term_indices: &[Vec<usize>],
    kappa_options: &SpatialLengthScaleOptimizationOptions,
    joint_setup: &ExactJointHyperSetup,
    seed_risk_profile: crate::seeding::SeedRiskProfile,
    analytic_joint_gradient_available: bool,
    analytic_joint_hessian_available: bool,
    disable_fixed_point: bool,
    screening_cap: Option<Arc<AtomicUsize>>,
    outer_derivative_policy: crate::families::custom_family::OuterDerivativePolicy,
    mut fit_fn: FitFn,
    mut exact_fn: ExactFn,
    mut exact_efs_fn: ExactEfsFn,
    mut seed_inner_beta_fn: SeedFn,
) -> Result<SpatialLengthScaleOptimizationResult<FitOut>, String>
where
    FitOut: Clone,
    FitFn: FnMut(
        &Array1<f64>,
        &[TermCollectionSpec],
        &[TermCollectionDesign],
    ) -> Result<FitOut, String>,
    ExactFn: FnMut(
        &Array1<f64>,
        &[TermCollectionSpec],
        &[TermCollectionDesign],
        crate::solver::estimate::reml::unified::EvalMode,
        &crate::families::row_kernel::RowSet,
    ) -> Result<
        (
            f64,
            Array1<f64>,
            crate::solver::outer_strategy::HessianResult,
        ),
        String,
    >,
    ExactEfsFn: FnMut(
        &Array1<f64>,
        &[TermCollectionSpec],
        &[TermCollectionDesign],
    ) -> Result<crate::solver::outer_strategy::EfsEval, String>,
    SeedFn: FnMut(&Array1<f64>) -> Result<(), EstimationError>,
{
    let n_blocks = block_specs.len();
    if block_term_indices.len() != n_blocks {
        return Err(SmoothError::dimension_mismatch(format!(
            "block_specs ({}) and block_term_indices ({}) length mismatch",
            n_blocks,
            block_term_indices.len()
        ))
        .into());
    }

    let log_kappa_dim = joint_setup.log_kappa_dim();

    // -----------------------------------------------------------------------
    // Fast path: kappa disabled or no spatial terms — build designs once.
    // -----------------------------------------------------------------------
    if joint_setup.auxiliary_dim() == 0 && (!kappa_options.enabled || log_kappa_dim == 0) {
        let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
            data, block_specs,
        )
        .map_err(|e| {
            format!("failed to build and freeze joint block designs during exact joint kappa optimization: {e}")
        })?;
        let theta0 = joint_setup.theta0();

        // Build temporary owned slices for the closure call.
        let spec_refs: Vec<TermCollectionSpec> = resolved_specs.clone();
        let design_refs: Vec<TermCollectionDesign> = designs.clone();
        let fit = fit_fn(&theta0, &spec_refs, &design_refs)?;
        return Ok(SpatialLengthScaleOptimizationResult {
            resolved_specs,
            designs,
            fit,
        });
    }

    // -----------------------------------------------------------------------
    // Full optimization path.
    // -----------------------------------------------------------------------
    let theta0 = joint_setup.theta0();
    let lower = joint_setup.lower();
    let upper = joint_setup.upper();
    if theta0.len() < log_kappa_dim || lower.len() != theta0.len() || upper.len() != theta0.len() {
        return Err(SmoothError::dimension_mismatch(format!(
            "invalid exact joint theta setup: theta0={}, lower={}, upper={}, required_log_kappa_dim={}",
            theta0.len(),
            lower.len(),
            upper.len(),
            log_kappa_dim
        ))
        .into());
    }
    let rho_dim = joint_setup.rho_dim();
    let all_dims = joint_setup.log_kappa_dims_per_term();

    // Build bootstrap designs and frozen specs for each block.
    let (boot_designs, best_specs) = build_term_collection_designs_and_freeze_joint(
        data,
        block_specs,
    )
    .map_err(|e| {
        format!(
            "failed to build and freeze joint block designs during exact joint kappa bootstrap: {e}"
        )
    })?;
    // Capability vs realized policy: the family may *advertise* an exact
    // analytic outer Hessian, but at this realized (n, psi_dim, rho_dim,
    // p_total) the predicted per-eval cost can still exceed the universal
    // outer-Hessian work budget. In that regime we route the outer optimizer
    // through gradient-only BFGS / L-BFGS, which is **convergent** to the
    // exact MLE — it just takes more line-search iterations. This is **not**
    // a feature drop: quasi-Newton picks up curvature from successive
    // analytic gradients, and the per-eval cost saving (`O(p)` instead of
    // `O(p²)`) more than pays for the iteration overhead at biobank scale.
    let policy_hessian_form = outer_derivative_policy.declared_hessian_form();
    let analytic_outer_hessian_available = analytic_joint_hessian_available
        && matches!(
            policy_hessian_form,
            crate::solver::outer_strategy::DeclaredHessianForm::Either
                | crate::solver::outer_strategy::DeclaredHessianForm::Dense
                | crate::solver::outer_strategy::DeclaredHessianForm::Operator { .. }
        );
    let prefer_gradient_only = !analytic_outer_hessian_available;

    let theta_dim = theta0.len();
    let psi_dim = theta_dim - rho_dim;

    // Build the cache with one realizer per block.
    let cache_blocks: Vec<(TermCollectionSpec, TermCollectionDesign, Vec<usize>)> = best_specs
        .iter()
        .zip(boot_designs.iter())
        .zip(block_term_indices.iter())
        .map(|((spec, design), terms)| (spec.clone(), design.clone(), terms.clone()))
        .collect();

    struct NBlockExactJointState<'d> {
        cache: ExactJointDesignCache<'d>,
    }

    let mut state = NBlockExactJointState {
        cache: ExactJointDesignCache::new(data, cache_blocks, rho_dim, all_dims.clone())?,
    };

    // ── P7: staged-κ schedule ────────────────────────────────────────────
    //
    // The κ MLE for a stationary spatial process is asymptotically
    // *invariant* in `n` once `n` is past the Monte-Carlo resolution of
    // the cell-moment kernel. At biobank scale (`n ≥ STAGED_KAPPA_*`) the
    // Monte-Carlo error of a `K = 5_000`-row pilot is ≪ the κ posterior
    // width, so estimating θ on a stratified `K`-row pilot returns
    // statistically the *same* estimate as the full-data fit at a
    // fraction of the wall-clock cost. We then do one Gauss-Newton-style
    // polish at `K_polish` to absorb residual Monte-Carlo error before
    // the final coefficient fit at the polished θ on the full data.
    //
    // This is **not a heuristic shortcut**. It is the textbook
    // pilot-then-refine schedule for stationary-process likelihoods,
    // chosen here because the per-eval cost of the κ gradient grows
    // linearly in `n` and the pilot subsample reduces that cost by a
    // factor of `n / K`. The final coefficient fit at θ̂_polished on the
    // full data preserves estimation accuracy for β.
    //
    // At `n < STAGED_KAPPA_TRIGGER_N` the schedule collapses to one
    // full-data stage — identical to the pre-P7 behaviour.
    // Note: the n≥30_000 pilot trigger lives in
    // `outer_derivative_policy.should_use_staged_kappa(n_total)`; this fn
    // only carries the constants it consumes directly.
    const KAPPA_PILOT_K: usize = 5_000;
    const KAPPA_POLISH_K: usize = 25_000;
    const KAPPA_POLISH_TRIGGER_N: usize = 100_000;

    let n_total = data.nrows();
    let use_staged_kappa = outer_derivative_policy.should_use_staged_kappa(n_total);
    if use_staged_kappa {
        log::info!(
            "[KAPPA-STAGED] auto-engaging pilot+polish schedule: n={} pilot_k={} polish_k={}",
            n_total,
            KAPPA_PILOT_K,
            KAPPA_POLISH_K,
        );
    }

    // Build the initial row mask for the κ optimization.
    //
    // * `use_staged_kappa = false`: full data (`RowSet::All`). The
    //   schedule collapses to the historical single-stage path.
    // * `use_staged_kappa = true`: deterministic uniform pilot of size
    //   `min(KAPPA_PILOT_K, n_total)`, wrapped as a `RowSet::Subsample`
    //   with per-row HT weight `n_total / k_pilot`. The uniform pick is
    //   a valid unbiased estimator on its own; the stratified
    //   per-decile picker
    //   (`marginal_slope_shared::auto_outer_score_subsample`) requires
    //   the response vector `z`, which only the family evaluator can
    //   produce. **Agent C replaces this with the stratified pick once
    //   `exact_fn` exposes the per-row score.**
    //
    // Sampling RNG is seeded from `n_total` so the pilot is
    // deterministic across reruns at fixed `n`.
    fn build_uniform_pilot_subsample(
        n_total: usize,
        k_target: usize,
        seed: u64,
    ) -> crate::families::marginal_slope_shared::OuterScoreSubsample {
        use crate::families::marginal_slope_shared::OuterScoreSubsample;
        let k = k_target.min(n_total);
        if k == 0 || n_total == 0 {
            return OuterScoreSubsample::new(Vec::new(), n_total, seed);
        }
        // Reservoir-free deterministic pick: linear congruential walk
        // over a shuffled index set; for the pilot, a fast Floyd-style
        // sample is sufficient.
        let mut mask: Vec<usize> = Vec::with_capacity(k);
        // Splitmix64-driven Floyd's sampler.
        let mut state = seed.wrapping_add(0x9E3779B97F4A7C15);
        let splitmix = |s: &mut u64| -> u64 {
            *s = s.wrapping_add(0x9E3779B97F4A7C15);
            let mut z = *s;
            z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
            z ^ (z >> 31)
        };
        let mut taken = std::collections::HashSet::with_capacity(k);
        for j in (n_total - k)..n_total {
            let r = (splitmix(&mut state) % (j as u64 + 1)) as usize;
            if !taken.insert(r) {
                taken.insert(j);
                mask.push(j);
            } else {
                mask.push(r);
            }
        }
        mask.sort_unstable();
        mask.dedup();
        OuterScoreSubsample::new(mask, n_total, seed)
    }

    let current_row_set: std::cell::RefCell<crate::families::row_kernel::RowSet> =
        if use_staged_kappa {
            let pilot = build_uniform_pilot_subsample(n_total, KAPPA_PILOT_K, n_total as u64);
            std::cell::RefCell::new(crate::families::row_kernel::RowSet::Subsample {
                rows: std::sync::Arc::clone(&pilot.rows),
                n_full: n_total,
            })
        } else {
            std::cell::RefCell::new(crate::families::row_kernel::RowSet::All)
        };

    let exact_fn_cell = std::cell::RefCell::new(&mut exact_fn);
    let exact_efs_fn_cell = std::cell::RefCell::new(&mut exact_efs_fn);

    // ── κ-optimization scaling instrumentation ──
    //
    // Per-phase wall-clock counters for the three kinds of evaluator
    // invocation the κ outer drives: cost-only line-search probes,
    // value-and-gradient(/Hessian) evaluations at accepted iterates, and
    // EFS fixed-point evaluations. Each invocation emits one
    // `[KAPPA-PHASE]` log line with a per-call elapsed time, plus the
    // running call counter and a summary `theta_norm` /
    // `log_kappa_norm` so the bench runner can attribute cost to
    // particular trajectory regions. A single `[KAPPA-PHASE-SUMMARY]`
    // line is emitted on optimization exit. Grepping these is the
    // production-fit κ-scaling probe (task #32) — measurement happens
    // in real biobank fits rather than a synthetic harness, so the
    // scaling law reflects the actual workload.
    use std::cell::Cell;
    let kphase_cost_calls: Cell<usize> = Cell::new(0);
    let kphase_cost_total_s: Cell<f64> = Cell::new(0.0);
    let kphase_eval_calls: Cell<usize> = Cell::new(0);
    let kphase_eval_total_s: Cell<f64> = Cell::new(0.0);
    let kphase_efs_calls: Cell<usize> = Cell::new(0);
    let kphase_efs_total_s: Cell<f64> = Cell::new(0.0);
    let kphase_optim_start = std::time::Instant::now();
    let kphase_log_kappa_dim = log_kappa_dim;
    let kphase_log_norms = |theta: &Array1<f64>| -> (f64, f64) {
        let theta_norm = theta.iter().map(|v| v * v).sum::<f64>().sqrt();
        let log_kappa_norm = if kphase_log_kappa_dim > 0 && theta.len() >= kphase_log_kappa_dim {
            let start = theta.len() - kphase_log_kappa_dim;
            theta.iter().skip(start).map(|v| v * v).sum::<f64>().sqrt()
        } else {
            0.0
        };
        (theta_norm, log_kappa_norm)
    };

    use crate::solver::outer_strategy::{
        DeclaredHessianForm, Derivative, OuterEval, OuterEvalOrder,
    };

    let problem = exact_joint_multistart_outer_problem(
        &theta0,
        &lower,
        &upper,
        rho_dim,
        psi_dim,
        theta_dim,
        if analytic_joint_gradient_available {
            Derivative::Analytic
        } else {
            Derivative::Unavailable
        },
        if analytic_outer_hessian_available {
            DeclaredHessianForm::Either
        } else {
            DeclaredHessianForm::Unavailable
        },
        prefer_gradient_only,
        disable_fixed_point,
        seed_risk_profile,
        kappa_options.rel_tol.max(1e-6),
        kappa_options.max_outer_iter.max(1),
        // Rho-axis cap: log-λ natural step ≈ 5.
        Some(5.0),
        // Psi-axis cap: kappa scale needs ~ln 2 per iter.
        Some(kappa_options.log_step.clamp(0.25, 1.0)),
        screening_cap.clone(),
    );

    // Helper: collect specs and designs from cache into owned Vecs for closure calls.
    fn collect_specs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionSpec> {
        cache.specs().into_iter().cloned().collect()
    }
    fn collect_designs(cache: &ExactJointDesignCache<'_>) -> Vec<TermCollectionDesign> {
        cache.designs().into_iter().cloned().collect()
    }

    let result = {
        let eval_outer = |ctx: &mut &mut NBlockExactJointState<'_>,
                          theta: &Array1<f64>,
                          order: OuterEvalOrder|
         -> Result<OuterEval, EstimationError> {
            if let Some((cost, grad, hess)) = ctx.cache.memoized_eval(theta) {
                let cached_satisfies_order = match order {
                    OuterEvalOrder::ValueAndGradient => true,
                    OuterEvalOrder::ValueGradientHessian => hess.is_analytic(),
                };
                if cached_satisfies_order {
                    if !cost.is_finite() {
                        return Ok(OuterEval::infeasible(theta.len()));
                    }
                    if grad.iter().any(|v| !v.is_finite()) {
                        return Err(EstimationError::RemlOptimizationFailed(
                            "n-block exact-joint gradient contained non-finite values".to_string(),
                        ));
                    }
                    return Ok(OuterEval {
                        cost,
                        gradient: grad,
                        hessian: hess,
                        inner_beta_hint: None,
                    });
                }
            }
            if let Err(err) = ctx.cache.ensure_theta(theta) {
                log::warn!(
                    "[OUTER] n-block exact-joint spatial: ensure_theta failed during gradient evaluation: {err}"
                );
                return Ok(OuterEval::infeasible(theta.len()));
            }
            let design_revision = Some(ctx.cache.design_revision());
            let specs = collect_specs(&ctx.cache);
            let designs = collect_designs(&ctx.cache);
            // Clamp the requested order against the realized outer
            // derivative policy. The capability-aware
            // `analytic_outer_hessian_available` already encodes the
            // policy gate; re-checking through `order_for_evaluation`
            // here keeps the per-eval branch in lockstep with the
            // top-of-function declaration so the optimizer and the
            // evaluator never disagree on what was requested.
            let clamped = outer_derivative_policy.order_for_evaluation(order);
            let need_hessian = matches!(clamped, OuterEvalOrder::ValueGradientHessian)
                && analytic_outer_hessian_available;
            let eval_mode = if need_hessian {
                crate::solver::estimate::reml::unified::EvalMode::ValueGradientHessian
            } else {
                crate::solver::estimate::reml::unified::EvalMode::ValueAndGradient
            };
            let t0 = std::time::Instant::now();
            let result = {
                let row_set_borrow = current_row_set.borrow();
                (*exact_fn_cell.borrow_mut())(theta, &specs, &designs, eval_mode, &row_set_borrow)
            };
            let elapsed_s = t0.elapsed().as_secs_f64();
            kphase_eval_calls.set(kphase_eval_calls.get() + 1);
            kphase_eval_total_s.set(kphase_eval_total_s.get() + elapsed_s);
            let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
            log::info!(
                "[KAPPA-PHASE] phase=eval_outer call={} order={:?} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
                kphase_eval_calls.get(),
                order,
                design_revision,
                theta_norm,
                log_kappa_norm,
                elapsed_s,
            );
            match result {
                Ok((cost, grad, hess)) => {
                    ctx.cache.store_eval((cost, grad.clone(), hess.clone()));
                    if !cost.is_finite() {
                        return Ok(OuterEval::infeasible(theta.len()));
                    }
                    if grad.iter().any(|v| !v.is_finite()) {
                        return Err(EstimationError::RemlOptimizationFailed(
                            "n-block exact-joint gradient contained non-finite values".to_string(),
                        ));
                    }
                    Ok(OuterEval {
                        cost,
                        gradient: grad,
                        hessian: hess,
                        inner_beta_hint: None,
                    })
                }
                Err(err) => {
                    log::warn!(
                        "[OUTER] n-block exact-joint spatial: exact evaluation failed: {err}"
                    );
                    Ok(OuterEval::infeasible(theta.len()))
                }
            }
        };

        let obj = problem.build_objective_with_eval_order(
            &mut state,
            |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
                if let Some(cost) = ctx.cache.memoized_cost(theta) {
                    return Ok(cost);
                }
                if let Err(err) = ctx.cache.ensure_theta(theta) {
                    log::warn!(
                        "[OUTER] n-block exact-joint spatial: ensure_theta failed during cost evaluation: {err}"
                    );
                    return Ok(f64::INFINITY);
                }
                let design_revision = Some(ctx.cache.design_revision());
                let specs = collect_specs(&ctx.cache);
                let designs = collect_designs(&ctx.cache);
                // Cost-only line-search probe: pass `ValueOnly` so the closure
                // skips gradient and Hessian assembly. This is the principled
                // fix for the N-block joint optimization V+G-per-probe waste —
                // gradient construction (≈ 6.5·10⁹ FLOPs per CTN step at
                // n=320 000, n_grid=293, p_resp=32, p_cov=23) is now paid only
                // when the outer evaluator actually requests it.
                let t0 = std::time::Instant::now();
                let result = {
                    let row_set_borrow = current_row_set.borrow();
                    (*exact_fn_cell.borrow_mut())(
                        theta,
                        &specs,
                        &designs,
                        crate::solver::estimate::reml::unified::EvalMode::ValueOnly,
                        &row_set_borrow,
                    )
                };
                let elapsed_s = t0.elapsed().as_secs_f64();
                kphase_cost_calls.set(kphase_cost_calls.get() + 1);
                kphase_cost_total_s.set(kphase_cost_total_s.get() + elapsed_s);
                let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
                log::info!(
                    "[KAPPA-PHASE] phase=cost call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
                    kphase_cost_calls.get(),
                    design_revision,
                    theta_norm,
                    log_kappa_norm,
                    elapsed_s,
                );
                match result {
                    Ok((cost, _grad, _hess)) => {
                        // Don't `store_eval`: that path is only valid when the
                        // closure produced a real gradient. The next outer-eval
                        // call will recompute (V, ∇V) at this θ if needed; the
                        // memoized_cost path covers the common case where the
                        // line search returns to an accepted iterate.
                        ctx.cache.store_cost_only(theta, cost);
                        Ok(cost)
                    }
                    Err(err) => {
                        log::warn!(
                            "[OUTER] n-block exact-joint spatial: exact cost evaluation failed: {err}"
                        );
                        Ok(f64::INFINITY)
                    }
                }
            },
            |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
                eval_outer(
                    ctx,
                    theta,
                    if analytic_outer_hessian_available {
                        OuterEvalOrder::ValueGradientHessian
                    } else {
                        OuterEvalOrder::ValueAndGradient
                    },
                )
            },
            |ctx: &mut &mut NBlockExactJointState<'_>,
             theta: &Array1<f64>,
             order: OuterEvalOrder| { eval_outer(ctx, theta, order) },
            None::<fn(&mut &mut NBlockExactJointState<'_>)>,
            Some(
                |ctx: &mut &mut NBlockExactJointState<'_>, theta: &Array1<f64>| {
                    ctx.cache
                        .ensure_theta(theta)
                        .map_err(EstimationError::InvalidInput)?;
                    let design_revision = Some(ctx.cache.design_revision());
                    let specs = collect_specs(&ctx.cache);
                    let designs = collect_designs(&ctx.cache);
                    let t0 = std::time::Instant::now();
                    let eval_result = (*exact_efs_fn_cell.borrow_mut())(
                        theta,
                        &specs,
                        &designs,
                    );
                    let elapsed_s = t0.elapsed().as_secs_f64();
                    kphase_efs_calls.set(kphase_efs_calls.get() + 1);
                    kphase_efs_total_s.set(kphase_efs_total_s.get() + elapsed_s);
                    let (theta_norm, log_kappa_norm) = kphase_log_norms(theta);
                    log::info!(
                        "[KAPPA-PHASE] phase=efs call={} design_revision={:?} theta_norm={:.4e} log_kappa_norm={:.4e} elapsed_s={:.4}",
                        kphase_efs_calls.get(),
                        design_revision,
                        theta_norm,
                        log_kappa_norm,
                        elapsed_s,
                    );
                    let eval = eval_result.map_err(EstimationError::RemlOptimizationFailed)?;
                    Ok(eval)
                },
            ),
        );
        let mut obj = obj.with_seed_inner_state(
            move |_ctx: &mut &mut NBlockExactJointState<'_>, beta: &Array1<f64>| {
                (seed_inner_beta_fn)(beta)
            },
        );

        problem
            .run(&mut obj, "n-block exact-joint spatial")
            .map_err(|e| e.to_string())?
    }; // obj dropped here, releasing mutable borrow on state

    // ── κ-optimization scaling summary ──
    //
    // Single line summarizing all per-call wall-clock counters
    // accumulated above. The bench runner / scaling-law analyzer
    // can pivot on this directly without parsing the per-call
    // [KAPPA-PHASE] markers (which remain available for
    // attribution).
    let kphase_total_s = kphase_optim_start.elapsed().as_secs_f64();
    log::info!(
        "[KAPPA-PHASE-SUMMARY] log_kappa_dim={} n_cost={} cost_total_s={:.4} n_eval={} eval_total_s={:.4} n_efs={} efs_total_s={:.4} optim_total_s={:.4}",
        kphase_log_kappa_dim,
        kphase_cost_calls.get(),
        kphase_cost_total_s.get(),
        kphase_eval_calls.get(),
        kphase_eval_total_s.get(),
        kphase_efs_calls.get(),
        kphase_efs_total_s.get(),
        kphase_total_s,
    );

    let theta_star = result.rho;

    // ── P7 stage rotation ────────────────────────────────────────────────
    //
    // The optimization above ran against `current_row_set` — the pilot
    // subsample under `use_staged_kappa`, otherwise the full data. We
    // now:
    //
    // 1. If `n_total ≥ KAPPA_POLISH_TRIGGER_N`, rotate to a larger
    //    polish subsample and request a single value+gradient evaluation
    //    at `theta_star` so the family caches its polished score. This
    //    is the Gauss-Newton-style polish in the schedule — one step
    //    rather than a full re-run because the pilot has already
    //    consumed most of the curvature information.
    //
    // 2. Always rotate back to `RowSet::All` before the final
    //    coefficient fit `fit_fn(theta_star)`. The final β estimate at
    //    θ̂ uses the full data so no estimation accuracy is lost.
    if use_staged_kappa && n_total >= KAPPA_POLISH_TRIGGER_N {
        let polish = build_uniform_pilot_subsample(
            n_total,
            KAPPA_POLISH_K,
            (n_total as u64).wrapping_add(0xA5A5A5A5),
        );
        *current_row_set.borrow_mut() = crate::families::row_kernel::RowSet::Subsample {
            rows: std::sync::Arc::clone(&polish.rows),
            n_full: n_total,
        };
        log::info!(
            "[KAPPA-STAGED] rotating to polish subsample: k={} at theta_star",
            polish.rows.len(),
        );
        // One V+G evaluation at theta_star on the polish subsample. The
        // returned objective pieces must be usable; the family-side cache
        // update inside `exact_fn` is consumed by the final fit.
        state.cache.ensure_theta(&theta_star)?;
        let (polish_cost, polish_grad, _) = {
            let specs = collect_specs(&state.cache);
            let designs = collect_designs(&state.cache);
            let row_set_borrow = current_row_set.borrow();
            exact_fn(
                &theta_star,
                &specs,
                &designs,
                crate::solver::estimate::reml::unified::EvalMode::ValueAndGradient,
                &row_set_borrow,
            )?
        };
        if !polish_cost.is_finite() || polish_grad.iter().any(|value| !value.is_finite()) {
            return Err(
                "polish subsample exact-joint evaluation produced non-finite objective pieces"
                    .to_string(),
            );
        }
    }
    *current_row_set.borrow_mut() = crate::families::row_kernel::RowSet::All;
    if use_staged_kappa {
        log::info!(
            "[KAPPA-STAGED] rotating to full data for final coefficient fit (n={})",
            n_total,
        );
    }

    state.cache.ensure_theta(&theta_star)?;

    let resolved_specs: Vec<TermCollectionSpec> = collect_specs(&state.cache);
    let designs: Vec<TermCollectionDesign> = collect_designs(&state.cache);

    let fit = fit_fn(&theta_star, &resolved_specs, &designs)?;

    for spec in &resolved_specs {
        log_spatial_aniso_scales(spec);
    }

    Ok(SpatialLengthScaleOptimizationResult {
        resolved_specs,
        designs,
        fit,
    })
}

fn try_exact_joint_latent_coord_optimization(
    data: ArrayView2<'_, f64>,
    y: ArrayView1<'_, f64>,
    weights: ArrayView1<'_, f64>,
    offset: ArrayView1<'_, f64>,
    resolvedspec: &TermCollectionSpec,
    best: &FittedTermCollection,
    family: LikelihoodSpec,
    options: &FitOptions,
    latent: &StandardLatentCoordConfig,
) -> Result<FittedTermCollectionWithSpec, EstimationError> {
    use crate::solver::outer_strategy::{
        DeclaredHessianForm, Derivative, OuterEval, OuterEvalOrder,
    };

    let rho_dim = best.fit.lambdas.len();
    let latent_flat_dim = latent.values.len();
    if latent_flat_dim == 0 {
        crate::bail_invalid_estim!(
            "latent-coordinate optimization requires a non-empty latent block"
        );
    }
    let direct_hypers =
        latent_coord_initial_direct_hypers(latent.values.id_mode(), latent.values.latent_dim())?;
    let analytic_rho_count = latent
        .analytic_penalties
        .as_ref()
        .map_or(0, |registry| registry.total_rho_count());
    let latent_coord_ext_dim = latent_flat_dim + analytic_rho_count + direct_hypers.len();

    let mut theta0 = Array1::<f64>::zeros(rho_dim + latent_coord_ext_dim);
    theta0
        .slice_mut(s![..rho_dim])
        .assign(&best.fit.lambdas.mapv(f64::ln));
    theta0
        .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
        .assign(latent.values.as_flat());
    if !direct_hypers.is_empty() {
        let direct_start = rho_dim + latent_flat_dim + analytic_rho_count;
        theta0
            .slice_mut(s![direct_start..direct_start + direct_hypers.len()])
            .assign(&direct_hypers);
    }

    let mut lower = Array1::<f64>::from_elem(theta0.len(), -12.0);
    let mut upper = Array1::<f64>::from_elem(theta0.len(), 12.0);
    let latent_bound = latent
        .values
        .as_flat()
        .iter()
        .fold(1.0_f64, |acc, &v| acc.max(v.abs()))
        + 10.0;
    for axis in rho_dim..rho_dim + latent_flat_dim {
        lower[axis] = -latent_bound;
        upper[axis] = latent_bound;
    }

    struct LatentJointContext<'d> {
        rho_dim: usize,
        cache: SingleBlockLatentCoordDesignCache,
        evaluator: crate::estimate::ExternalJointHyperEvaluator<'d>,
    }

    impl<'d> LatentJointContext<'d> {
        fn eval_full(
            &mut self,
            theta: &Array1<f64>,
            order: OuterEvalOrder,
        ) -> Result<
            (
                f64,
                Array1<f64>,
                crate::solver::outer_strategy::HessianResult,
            ),
            EstimationError,
        > {
            if let Some(eval) = self.cache.memoized_eval(theta) {
                return Ok(eval);
            }
            self.cache
                .ensure_theta(theta)
                .map_err(EstimationError::InvalidInput)?;
            let hyper_dirs = self
                .cache
                .hyper_dirs()
                .map_err(EstimationError::InvalidInput)?;
            let design_revision = Some(self.cache.design_revision());
            let registry_for_key = self.cache.analytic_penalties();
            self.evaluator
                .set_analytic_penalty_registry(registry_for_key.as_deref());
            let mut eval = evaluate_joint_reml_outer_eval_at_theta(
                &mut self.evaluator,
                self.cache.design(),
                theta,
                self.rho_dim,
                hyper_dirs,
                None,
                order,
                design_revision,
            )?;
            let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
            if let Some(registry) = registry_for_key {
                let mut registry = registry.as_ref().clone();
                registry.apply_weight_schedules(
                    crate::solver::estimate::reml::runtime::current_outer_iter() as usize,
                );
                add_analytic_penalty_objective_to_eval(
                    theta,
                    self.rho_dim,
                    latent.as_ref(),
                    &registry,
                    &mut eval,
                )?;
            }
            add_latent_id_objective_to_eval(
                theta,
                self.rho_dim,
                self.cache.analytic_penalty_rho_count(),
                latent.as_ref(),
                &mut eval,
            )?;
            self.cache.store_eval(eval.clone());
            Ok(eval)
        }

        fn eval_efs(
            &mut self,
            theta: &Array1<f64>,
        ) -> Result<crate::solver::outer_strategy::EfsEval, EstimationError> {
            self.cache
                .ensure_theta(theta)
                .map_err(EstimationError::InvalidInput)?;
            let hyper_dirs = self
                .cache
                .hyper_dirs()
                .map_err(EstimationError::InvalidInput)?;
            let registry_for_key = self.cache.analytic_penalties();
            self.evaluator
                .set_analytic_penalty_registry(registry_for_key.as_deref());
            let mut efs = evaluate_joint_reml_efs_at_theta(
                &mut self.evaluator,
                self.cache.design(),
                theta,
                self.rho_dim,
                hyper_dirs,
                None,
                Some(self.cache.design_revision()),
            )?;
            if let Some(registry) = registry_for_key {
                let mut registry = registry.as_ref().clone();
                registry.apply_weight_schedules(
                    crate::solver::estimate::reml::runtime::current_outer_iter() as usize,
                );
                let latent = self.cache.latent().map_err(EstimationError::InvalidInput)?;
                let contribution = analytic_penalty_objective_contribution(
                    theta,
                    self.rho_dim,
                    latent.as_ref(),
                    &registry,
                )?;
                efs.cost += contribution.cost;
                if let (Some(psi_gradient), Some(psi_indices)) =
                    (efs.psi_gradient.as_mut(), efs.psi_indices.as_ref())
                {
                    if psi_gradient.len() != psi_indices.len() {
                        crate::bail_invalid_estim!(
                            "latent-coordinate analytic penalty EFS psi gradient length mismatch: gradient={}, indices={}",
                            psi_gradient.len(),
                            psi_indices.len()
                        );
                    }
                    for (local_idx, &theta_idx) in psi_indices.iter().enumerate() {
                        psi_gradient[local_idx] += contribution.gradient[theta_idx];
                    }
                }
            }
            Ok(efs)
        }

        fn eval_cost(&mut self, theta: &Array1<f64>) -> f64 {
            if let Some(cost) = self.cache.memoized_cost(theta) {
                return cost;
            }
            if self.cache.ensure_theta(theta).is_err() {
                return f64::INFINITY;
            }
            let design_revision = Some(self.cache.design_revision());
            let registry_for_key = self.cache.analytic_penalties();
            self.evaluator
                .set_analytic_penalty_registry(registry_for_key.as_deref());
            let result = {
                let design = self.cache.design();
                self.evaluator.evaluate_cost_only(
                    &design.design,
                    &design.penalties,
                    &design.nullspace_dims,
                    design.linear_constraints.clone(),
                    theta,
                    self.rho_dim,
                    None,
                    "latent-coordinate-joint cost-only",
                    design_revision,
                )
            };
            match result {
                Ok(cost) => {
                    let latent = match self.cache.latent() {
                        Ok(latent) => latent,
                        Err(_) => return f64::INFINITY,
                    };
                    let contribution = match latent_id_objective_contribution(
                        theta,
                        self.rho_dim,
                        self.cache.analytic_penalty_rho_count(),
                        latent.as_ref(),
                    ) {
                        Ok(contribution) => contribution,
                        Err(_) => return f64::INFINITY,
                    };
                    let cost = cost + contribution.cost;
                    let cost = if let Some(registry) = registry_for_key {
                        let mut registry = registry.as_ref().clone();
                        registry.apply_weight_schedules(
                            crate::solver::estimate::reml::runtime::current_outer_iter() as usize,
                        );
                        match analytic_penalty_objective_contribution(
                            theta,
                            self.rho_dim,
                            latent.as_ref(),
                            &registry,
                        ) {
                            Ok(contribution) => cost + contribution.cost,
                            Err(_) => return f64::INFINITY,
                        }
                    } else {
                        cost
                    };
                    self.cache.store_cost(cost);
                    cost
                }
                Err(_) => f64::INFINITY,
            }
        }
    }

    let mut ctx = LatentJointContext {
        rho_dim,
        cache: SingleBlockLatentCoordDesignCache::new(
            data.to_owned(),
            resolvedspec.clone(),
            best.design.clone(),
            latent,
            rho_dim,
        )
        .map_err(EstimationError::InvalidInput)?,
        evaluator: crate::estimate::ExternalJointHyperEvaluator::new(
            y,
            weights,
            &best.design.design,
            offset,
            &best.design.penalties,
            &external_opts_for_design(&family, &best.design, options),
            "latent-coordinate-joint",
        )?,
    };
    let registry_for_key = ctx.cache.analytic_penalties();
    ctx.evaluator
        .set_analytic_penalty_registry(registry_for_key.as_deref());
    ctx.evaluator
        .set_persistent_latent_values_fingerprint(latent.values.id_mode());
    if let Some(cached_t) = ctx
        .evaluator
        .load_persistent_latent_values(latent.values.n_obs(), latent.values.latent_dim())
    {
        for (dst, src) in theta0
            .slice_mut(s![rho_dim..rho_dim + latent_flat_dim])
            .iter_mut()
            .zip(cached_t.iter())
        {
            *dst = *src;
        }
    }

    let problem = exact_joint_multistart_outer_problem(
        &theta0,
        &lower,
        &upper,
        rho_dim,
        latent_coord_ext_dim,
        theta0.len(),
        Derivative::Analytic,
        DeclaredHessianForm::Unavailable,
        false,
        false,
        seed_risk_profile_for_likelihood_family(&family),
        options.tol,
        options.max_iter.max(1),
        Some(5.0),
        Some(0.5),
        None,
    );

    let eval_outer = |ctx: &mut &mut LatentJointContext<'_>,
                      theta: &Array1<f64>,
                      order: OuterEvalOrder|
     -> Result<OuterEval, EstimationError> {
        let (cost, gradient, hessian) = ctx.eval_full(theta, order)?;
        Ok(OuterEval {
            cost,
            gradient,
            hessian,
            inner_beta_hint: None,
        })
    };

    let result = {
        let mut obj = problem.build_objective_with_eval_order(
            &mut ctx,
            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| Ok(ctx.eval_cost(theta)),
            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| {
                eval_outer(ctx, theta, OuterEvalOrder::ValueAndGradient)
            },
            |ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>, order: OuterEvalOrder| {
                eval_outer(ctx, theta, order)
            },
            Some(|ctx: &mut &mut LatentJointContext<'_>| {
                ctx.cache.reset();
            }),
            Some(|ctx: &mut &mut LatentJointContext<'_>, theta: &Array1<f64>| ctx.eval_efs(theta)),
        );

        problem
            .run(&mut obj, "latent-coordinate joint REML")
            .map_err(|e| {
                EstimationError::InvalidInput(format!(
                    "latent-coordinate joint optimization failed after exhausting strategy fallbacks: {e}"
                ))
            })?
    };
    if !result.converged {
        crate::bail_invalid_estim!(
            "latent-coordinate joint optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
            result.iterations,
            result.final_value,
            result.final_grad_norm_report(),
        );
    }

    let theta_star = result.rho;
    let rho_star = theta_star.slice(s![..rho_dim]).mapv(f64::exp);
    let mut final_data = data.to_owned();
    let flat_t = theta_star
        .slice(s![rho_dim..rho_dim + latent_flat_dim])
        .to_owned();
    let mut fitted_latent_values =
        Array2::<f64>::zeros((latent.values.n_obs(), latent.values.latent_dim()));
    for n in 0..latent.values.n_obs() {
        for axis in 0..latent.values.latent_dim() {
            let value = flat_t[n * latent.values.latent_dim() + axis];
            fitted_latent_values[[n, axis]] = value;
            final_data[[n, latent.feature_cols[axis]]] = value;
        }
    }
    let optimized = fit_term_collection_forspecwith_heuristic_lambdas(
        final_data.view(),
        y,
        weights,
        offset,
        resolvedspec,
        rho_star.as_slice(),
        family,
        options,
    )?;
    ctx.evaluator
        .store_persistent_latent_values(&fitted_latent_values);
    let mut fit = optimized.fit;
    fit.reml_score = result.final_value;
    fit.penalized_objective = result.final_value;
    Ok(FittedTermCollectionWithSpec {
        fit,
        design: optimized.design,
        resolvedspec: resolvedspec.clone(),
        adaptive_diagnostics: optimized.adaptive_diagnostics,
    })
}

pub fn fit_term_collectionwith_latent_coord_optimization(
    data: ArrayView2<'_, f64>,
    y: Array1<f64>,
    weights: Array1<f64>,
    offset: Array1<f64>,
    spec: &TermCollectionSpec,
    latent: &StandardLatentCoordConfig,
    family: LikelihoodSpec,
    options: &FitOptions,
) -> Result<FittedTermCollectionWithSpec, EstimationError> {
    let n = data.nrows();
    if !(y.len() == n && weights.len() == n && offset.len() == n) {
        crate::bail_invalid_estim!(
            "fit_term_collectionwith_latent_coord_optimization row mismatch: n={}, y={}, weights={}, offset={}",
            n,
            y.len(),
            weights.len(),
            offset.len()
        );
    }
    let best = fit_term_collection_forspec(
        data,
        y.view(),
        weights.view(),
        offset.view(),
        spec,
        family.clone(),
        options,
    )?;
    let resolvedspec = freeze_term_collection_from_design(spec, &best.design)?;
    try_exact_joint_latent_coord_optimization(
        data,
        y.view(),
        weights.view(),
        offset.view(),
        &resolvedspec,
        &best,
        family,
        options,
        latent,
    )
}

pub fn fit_term_collectionwith_spatial_length_scale_optimization(
    data: ArrayView2<'_, f64>,
    y: Array1<f64>,
    weights: Array1<f64>,
    offset: Array1<f64>,
    spec: &TermCollectionSpec,
    family: LikelihoodSpec,
    options: &FitOptions,
    kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<FittedTermCollectionWithSpec, EstimationError> {
    // Spatial hyperparameters change kernel geometry nonlinearly, so each
    // proposal rebuilds the spatial basis. Hybrid/isotropic terms expose a
    // scalar κ (= 1/length_scale); pure Duchon anisotropy exposes only
    // per-axis shape coordinates.
    //
    // When exact derivative information is available for the rebuilt basis and
    // penalty, kappa is promoted to a first-class outer hyperparameter beside
    // rho = log(lambda). In that mode this routine runs a joint outer solve in
    // theta = [rho, psi], where psi = log(kappa) = -log(length_scale), and the
    // optimizer is expected to consume a real joint Hessian. ARC is not meant
    // to run on a gradient-only surrogate here.
    //
    // Any eligible spatial smooth participates in this outer solve. If an
    // eligible spatial basis does not expose derivative information, that is
    // now a hard error.
    let mut resolvedspec = spec.clone();
    let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
    let n = data.nrows();
    if !(y.len() == n && weights.len() == n && offset.len() == n) {
        crate::bail_invalid_estim!(
            "fit_term_collectionwith_spatial_length_scale_optimization row mismatch: n={}, y={}, weights={}, offset={}",
            n,
            y.len(),
            weights.len(),
            offset.len()
        );
    }
    if !kappa_options.enabled || spatial_terms.is_empty() {
        let out = fit_term_collection_forspec(
            data,
            y.view(),
            weights.view(),
            offset.view(),
            &resolvedspec,
            family,
            options,
        )?;
        let resolvedspec = freeze_term_collection_from_design(&resolvedspec, &out.design)?;
        return Ok(FittedTermCollectionWithSpec {
            fit: out.fit,
            design: out.design,
            resolvedspec,
            adaptive_diagnostics: out.adaptive_diagnostics,
        });
    }
    if kappa_options.max_outer_iter == 0 {
        crate::bail_invalid_estim!("spatial kappa optimization requires max_outer_iter >= 1");
    }
    if !(kappa_options.log_step.is_finite() && kappa_options.log_step > 0.0) {
        crate::bail_invalid_estim!("spatial kappa optimization requires log_step > 0");
    }
    if !(kappa_options.min_length_scale.is_finite()
        && kappa_options.max_length_scale.is_finite()
        && kappa_options.min_length_scale > 0.0
        && kappa_options.max_length_scale >= kappa_options.min_length_scale)
    {
        crate::bail_invalid_estim!(
            "spatial kappa optimization requires valid positive length_scale bounds"
        );
    }

    let pilot_threshold = kappa_options.pilot_subsample_threshold;
    if pilot_threshold > 0 && n > pilot_threshold * 2 {
        log::info!(
            "[spatial-kappa] n={n} exceeds pilot threshold {}; using pilot geometry only for deterministic anisotropy initialization",
            pilot_threshold * 2,
        );
        apply_spatial_anisotropy_pilot_initializer(
            data,
            &mut resolvedspec,
            &spatial_terms,
            pilot_threshold,
            kappa_options,
        );
    }

    let best = fit_term_collection_forspec(
        data,
        y.view(),
        weights.view(),
        offset.view(),
        &resolvedspec,
        family.clone(),
        options,
    )?;
    resolvedspec = freeze_term_collection_from_design(&resolvedspec, &best.design)?;
    // The freeze step can rewrite a term's basis variant — most notably when
    // `build_thin_plate_basis_with_workspace` auto-promotes an infeasible
    // canonical-TPS request to a pure Duchon spline (length_scale = None,
    // no anisotropy). The pre-fit eligibility list was computed against the
    // ThinPlate spec, which has length_scale set, so it included that term.
    // After the rewrite the same term is a *pure* Duchon basis with no free
    // length-scale parameter to optimize, and the downstream kappa solver
    // (which assumes hybrid Duchon for log-κ derivatives) errors out. Refresh
    // the index list so it reflects the post-freeze spec.
    let spatial_terms = spatial_length_scale_term_indices(&resolvedspec);
    // Sync knot-cloud-derived aniso contrasts from the basis metadata back
    // into the spec so the optimizer starts from the geometry-informed η values
    // rather than the zero sentinel from --scale-dimensions.
    sync_aniso_contrasts_from_metadata(&mut resolvedspec, &best.design.smooth);
    let initial_score = fit_score(&best.fit);
    if !initial_score.is_finite() {
        log::debug!("[spatial-kappa] initial profiled score is non-finite");
    }
    if !spatial_terms.is_empty() {
        let exact_joint = require_successful_spatial_optimization_result(
            initial_score,
            try_exact_joint_spatial_length_scale_optimization(
                data,
                y.view(),
                weights.view(),
                offset.view(),
                &resolvedspec,
                &best,
                family,
                options,
                kappa_options,
                &spatial_terms,
            )
            .map(|opt| {
                opt.map(|fit| {
                    let score = fit_score(&fit.fit);
                    (fit, score)
                })
            }),
        )?;
        log_spatial_aniso_scales(&exact_joint.resolvedspec);
        return Ok(exact_joint);
    }

    Ok(FittedTermCollectionWithSpec {
        fit: best.fit,
        design: best.design,
        resolvedspec,
        adaptive_diagnostics: best.adaptive_diagnostics,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::basis::{
        BSplineBasisSpec, BSplineIdentifiability, BSplineKnotSpec, CenterStrategy, DuchonBasisSpec,
        DuchonNullspaceOrder, DuchonOperatorPenaltySpec, MaternBasisSpec, MaternIdentifiability,
        MaternNu, SpatialIdentifiability, ThinPlateBasisSpec,
    };
    use crate::estimate::AdaptiveRegularizationOptions;
    use crate::faer_ndarray::{FaerEigh, FaerSvd};
    use ndarray::{Axis, array};
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    /// A minimal frozen 1-D B-spline basis at `feature_col`, used to exercise
    /// the column-remap walk without standing up a full fit.
    fn remap_test_bspline(feature_col: usize) -> SmoothBasisSpec {
        SmoothBasisSpec::BSpline1D {
            feature_col,
            spec: BSplineBasisSpec {
                degree: 3,
                penalty_order: 2,
                knotspec: BSplineKnotSpec::Generate {
                    data_range: (0.0, 1.0),
                    num_internal_knots: 4,
                },
                double_penalty: false,
                identifiability: BSplineIdentifiability::None,
                boundary_conditions: BSplineBoundaryConditions::default(),
                boundary: OneDimensionalBoundary::Open,
            },
        }
    }

    #[test]
    fn remap_feature_columns_rewrites_every_index_bearing_field() {
        // Exhaustively verify that TermCollectionSpec::remap_feature_columns
        // re-resolves *every* stored column index across every basis variant —
        // including the two that the old survival-only walk silently skipped
        // (BySmooth's by_kind.feature_col and FactorSmooth's
        // continuous_cols/group_col). This is the predict-time realignment
        // contract (#803): a stale training index that survives the walk would
        // dereference the wrong predict column.
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin".to_string(),
                feature_col: 1,
                feature_cols: vec![1],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![RandomEffectTermSpec {
                name: "re".to_string(),
                feature_col: 2,
                drop_first_level: false,
                penalized: true,
                frozen_levels: Some(vec![0, 1]),
            }],
            smooth_terms: vec![
                SmoothTermSpec {
                    name: "bspline".to_string(),
                    basis: remap_test_bspline(3),
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "by_variable".to_string(),
                    basis: SmoothBasisSpec::ByVariable {
                        inner: Box::new(remap_test_bspline(4)),
                        by_col: 5,
                        kind: BySmoothKind::Numeric,
                        by: ByVariableSpec::Numeric,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "by_smooth".to_string(),
                    basis: SmoothBasisSpec::BySmooth {
                        smooth: Box::new(remap_test_bspline(6)),
                        by_kind: ByVarKind::Numeric { feature_col: 7 },
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "factor_smooth".to_string(),
                    basis: SmoothBasisSpec::FactorSmooth {
                        spec: FactorSmoothSpec {
                            continuous_cols: vec![8],
                            group_col: 9,
                            marginal: match remap_test_bspline(0) {
                                SmoothBasisSpec::BSpline1D { spec, .. } => spec,
                                _ => unreachable!(),
                            },
                            flavour: FactorSmoothFlavour::Sz,
                            group_frozen_levels: Some(vec![0, 1]),
                        },
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "thin_plate".to_string(),
                    basis: SmoothBasisSpec::ThinPlate {
                        feature_cols: vec![10, 11],
                        spec: ThinPlateBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                            length_scale: 1.0,
                            double_penalty: true,
                            identifiability: SpatialIdentifiability::default(),
                            radial_reparam: None,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
            ],
        };

        // Remap every index by +100 (injective, so any missed field stays < 100
        // and is caught below).
        let remapped: TermCollectionSpec = spec
            .remap_feature_columns(|i| Ok::<usize, String>(i + 100))
            .expect("remap must succeed");

        assert_eq!(remapped.linear_terms[0].feature_col, 101);
        assert_eq!(remapped.random_effect_terms[0].feature_col, 102);

        let collected = collect_feature_columns(&remapped);
        let mut expected: Vec<usize> = (3..=11).map(|i| i + 100).collect();
        expected.sort_unstable();
        assert_eq!(
            collected, expected,
            "every smooth-basis column index must be remapped exactly once"
        );

        // The remap closure's error must short-circuit the whole walk.
        let err = spec.remap_feature_columns(|_| Err::<usize, String>("boom".to_string()));
        assert_eq!(err.unwrap_err(), "boom");
    }

    /// Gather every column index referenced by the smooth bases of a spec, sorted.
    fn collect_feature_columns(spec: &TermCollectionSpec) -> Vec<usize> {
        fn walk(basis: &SmoothBasisSpec, out: &mut Vec<usize>) {
            match basis {
                SmoothBasisSpec::ByVariable { inner, by_col, .. }
                | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
                    out.push(*by_col);
                    walk(inner, out);
                }
                SmoothBasisSpec::BSpline1D { feature_col, .. } => out.push(*feature_col),
                SmoothBasisSpec::BySmooth { smooth, by_kind } => {
                    match by_kind {
                        ByVarKind::Numeric { feature_col }
                        | ByVarKind::Factor { feature_col, .. } => out.push(*feature_col),
                    }
                    walk(smooth, out);
                }
                SmoothBasisSpec::FactorSmooth { spec } => {
                    out.extend(spec.continuous_cols.iter().copied());
                    out.push(spec.group_col);
                }
                SmoothBasisSpec::ThinPlate { feature_cols, .. }
                | SmoothBasisSpec::Sphere { feature_cols, .. }
                | SmoothBasisSpec::Matern { feature_cols, .. }
                | SmoothBasisSpec::Duchon { feature_cols, .. }
                | SmoothBasisSpec::Pca { feature_cols, .. }
                | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
                    out.extend(feature_cols.iter().copied())
                }
            }
        }
        let mut out = Vec::new();
        for st in &spec.smooth_terms {
            walk(&st.basis, &mut out);
        }
        out.sort_unstable();
        out
    }

    #[test]
    fn bspline_boundary_conditions_emit_paired_equality_constraints() {
        let x = Array1::linspace(0.0, 1.0, 25);
        let data = x.clone().insert_axis(Axis(1));
        let spec = SmoothTermSpec {
            name: "s(x, bc_left=anchored, anchor_left=2, bc_right=clamped)".to_string(),
            basis: SmoothBasisSpec::BSpline1D {
                feature_col: 0,
                spec: BSplineBasisSpec {
                    degree: 3,
                    penalty_order: 2,
                    knotspec: BSplineKnotSpec::Generate {
                        data_range: (0.0, 1.0),
                        num_internal_knots: 4,
                    },
                    double_penalty: false,
                    identifiability: BSplineIdentifiability::None,
                    boundary_conditions: BSplineBoundaryConditions {
                        left: BSplineEndpointBoundaryCondition::Anchored { value: 2.0 },
                        right: BSplineEndpointBoundaryCondition::Clamped,
                    },
                    boundary: OneDimensionalBoundary::Open,
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let design = build_smooth_design(data.view(), &[spec]).expect("boundary design");
        let constraints = design
            .linear_constraints
            .as_ref()
            .expect("boundary constraints");
        assert_eq!(constraints.a.nrows(), 4);
        assert_eq!(constraints.a.ncols(), design.total_smooth_cols());
        assert_eq!(constraints.b.to_vec(), vec![2.0, -2.0, 0.0, -0.0]);

        let metadata = &design.terms[0].metadata;
        let BasisMetadata::BSpline1D { knots, .. } = metadata else {
            panic!("expected B-spline metadata");
        };
        let (left_value, _) = crate::basis::create_basis::<Dense>(
            array![0.0].view(),
            KnotSource::Provided(knots.view()),
            3,
            BasisOptions::value(),
        )
        .expect("left endpoint basis");
        let (right_slope, _) = crate::basis::create_basis::<Dense>(
            array![1.0].view(),
            KnotSource::Provided(knots.view()),
            3,
            BasisOptions::first_derivative(),
        )
        .expect("right endpoint derivative");
        for j in 0..constraints.a.ncols() {
            assert!((constraints.a[[0, j]] - left_value[[0, j]]).abs() < 1e-12);
            assert!((constraints.a[[1, j]] + left_value[[0, j]]).abs() < 1e-12);
            assert!((constraints.a[[2, j]] - right_slope[[0, j]]).abs() < 1e-12);
            assert!((constraints.a[[3, j]] + right_slope[[0, j]]).abs() < 1e-12);
        }
    }

    #[test]
    fn bspline_boundary_conditions_follow_frozen_identifiability_transform() {
        let x = Array1::linspace(0.0, 1.0, 20);
        let data = x.insert_axis(Axis(1));
        let raw_cols = 8;
        let mut z = Array2::<f64>::zeros((raw_cols, raw_cols - 1));
        for j in 0..(raw_cols - 1) {
            z[[j, j]] = 1.0;
            z[[raw_cols - 1, j]] = -1.0;
        }
        let spec = SmoothTermSpec {
            name: "half-open anchored smooth".to_string(),
            basis: SmoothBasisSpec::BSpline1D {
                feature_col: 0,
                spec: BSplineBasisSpec {
                    degree: 3,
                    penalty_order: 2,
                    knotspec: BSplineKnotSpec::Generate {
                        data_range: (0.0, 1.0),
                        num_internal_knots: 4,
                    },
                    double_penalty: false,
                    identifiability: BSplineIdentifiability::FrozenTransform {
                        transform: z.clone(),
                    },
                    boundary_conditions: BSplineBoundaryConditions {
                        left: BSplineEndpointBoundaryCondition::Anchored { value: 0.0 },
                        right: BSplineEndpointBoundaryCondition::Free,
                    },
                    boundary: OneDimensionalBoundary::Open,
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let design = build_smooth_design(data.view(), &[spec]).expect("boundary design");
        let constraints = design
            .linear_constraints
            .as_ref()
            .expect("boundary constraints");
        assert_eq!(constraints.a.nrows(), 2);
        assert_eq!(constraints.a.ncols(), raw_cols - 1);
        let BasisMetadata::BSpline1D { knots, .. } = &design.terms[0].metadata else {
            panic!("expected B-spline metadata");
        };
        let (left_value, _) = crate::basis::create_basis::<Dense>(
            array![0.0].view(),
            KnotSource::Provided(knots.view()),
            3,
            BasisOptions::value(),
        )
        .expect("left endpoint basis");
        let expected = left_value.row(0).to_owned().dot(&z);
        for j in 0..expected.len() {
            assert!((constraints.a[[0, j]] - expected[j]).abs() < 1e-12);
            assert!((constraints.a[[1, j]] + expected[j]).abs() < 1e-12);
        }
    }

    fn assert_spatial_derivative_width(
        label: &str,
        dense: &Array2<f64>,
        implicit: Option<&crate::terms::basis::ImplicitDesignPsiDerivative>,
        expected: usize,
    ) {
        if let Some(op) = implicit {
            assert_eq!(
                op.p_out(),
                expected,
                "{label} implicit derivative width should match term coefficient width"
            );
        } else {
            assert_eq!(
                dense.ncols(),
                expected,
                "{label} dense derivative width should match term coefficient width"
            );
        }
    }

    fn numerical_rank(x: &Array2<f64>) -> usize {
        let (_, s, _) = x
            .svd(false, false)
            .expect("SVD should succeed in rank test");
        let sigma_max = s.iter().copied().fold(0.0_f64, f64::max);
        let tol = (x.nrows().max(x.ncols()).max(1) as f64) * f64::EPSILON * sigma_max.max(1.0);
        s.iter().filter(|&&sv| sv > tol).count()
    }

    fn residual_norm_to_column_space(x: &Array2<f64>, y: &Array1<f64>) -> f64 {
        let (u_opt, _, _) = x
            .svd(true, false)
            .expect("SVD should succeed in projection residual test");
        let u = u_opt.expect("left singular vectors should be present");
        let rank = numerical_rank(x);
        let mut proj = Array1::<f64>::zeros(y.len());
        for j in 0..rank.min(u.ncols()) {
            let uj = u.column(j);
            let coeff = uj.dot(y);
            proj.scaled_add(coeff, &uj);
        }
        let resid = y - &proj;
        resid.dot(&resid).sqrt()
    }

    fn spatial_log_kappa_bounds_from_options(
        dims_per_term: &[usize],
        options: &SpatialLengthScaleOptimizationOptions,
        lower: bool,
    ) -> SpatialLogKappaCoords {
        let total: usize = dims_per_term.iter().sum();
        let value = if lower {
            -options.max_length_scale.ln()
        } else {
            -options.min_length_scale.ln()
        };
        SpatialLogKappaCoords::new_with_dims(
            Array1::<f64>::from_elem(total, value),
            dims_per_term.to_vec(),
        )
    }

    fn two_block_exact_joint_hyper_setup(
        meanspec: &TermCollectionSpec,
        noisespec: &TermCollectionSpec,
        kappa_options: &SpatialLengthScaleOptimizationOptions,
    ) -> ExactJointHyperSetup {
        let mean_terms = spatial_length_scale_term_indices(meanspec);
        let noise_terms = spatial_length_scale_term_indices(noisespec);
        let mean_dims_per_term = spatial_dims_per_term(meanspec, &mean_terms);
        let noise_dims_per_term = spatial_dims_per_term(noisespec, &noise_terms);
        let mean_use_aniso = has_aniso_terms(meanspec, &mean_terms);
        let noise_use_aniso = has_aniso_terms(noisespec, &noise_terms);
        let mean_log_kappa = if mean_use_aniso {
            SpatialLogKappaCoords::from_length_scales_aniso(meanspec, &mean_terms, kappa_options)
        } else {
            SpatialLogKappaCoords::from_length_scales(meanspec, &mean_terms, kappa_options)
        };
        let noise_log_kappa = if noise_use_aniso {
            SpatialLogKappaCoords::from_length_scales_aniso(noisespec, &noise_terms, kappa_options)
        } else {
            SpatialLogKappaCoords::from_length_scales(noisespec, &noise_terms, kappa_options)
        };
        let dims_per_term = mean_log_kappa
            .dims_per_term()
            .iter()
            .copied()
            .chain(noise_log_kappa.dims_per_term().iter().copied())
            .collect::<Vec<_>>();
        assert_eq!(
            dims_per_term,
            mean_dims_per_term
                .iter()
                .copied()
                .chain(noise_dims_per_term.iter().copied())
                .collect::<Vec<_>>()
        );
        let log_kappa0 = SpatialLogKappaCoords::new_with_dims(
            Array1::from_iter(
                mean_log_kappa
                    .as_array()
                    .iter()
                    .chain(noise_log_kappa.as_array().iter())
                    .copied(),
            ),
            dims_per_term.clone(),
        );
        ExactJointHyperSetup::new(
            Array1::zeros(0),
            Array1::zeros(0),
            Array1::zeros(0),
            log_kappa0,
            spatial_log_kappa_bounds_from_options(&dims_per_term, kappa_options, true),
            spatial_log_kappa_bounds_from_options(&dims_per_term, kappa_options, false),
        )
    }

    fn max_abs_diff_matrix(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
        assert_eq!(a.dim(), b.dim());
        a.iter()
            .zip(b.iter())
            .map(|(&x, &y)| (x - y).abs())
            .fold(0.0_f64, f64::max)
    }

    fn assert_frozen_replay_matches_fit(
        data: ArrayView2<'_, f64>,
        spec: &TermCollectionSpec,
        label: &str,
    ) {
        let fit_design = build_term_collection_design(data, spec).expect("fit-time design");
        let frozen =
            freeze_term_collection_from_design(spec, &fit_design).expect("freeze term collection");
        let replay_design = build_term_collection_design(data, &frozen).expect("replay design");
        let max_abs = max_abs_diff_matrix(
            &fit_design.design.to_dense(),
            &replay_design.design.to_dense(),
        );
        assert!(
            max_abs <= 1e-10,
            "{label} frozen replay changed realized design: max_abs={max_abs}"
        );
    }

    fn dense_kronecker_pseudo_logdet_reference(
        marginal_penalties: &[Array2<f64>],
        lambdas: &[f64],
        ridge: f64,
    ) -> (f64, Array1<f64>, Array2<f64>) {
        let p_total: usize = marginal_penalties
            .iter()
            .map(|penalty| penalty.nrows())
            .product();
        let mut s_dense = Array2::<f64>::zeros((p_total, p_total));
        for (axis, penalty) in marginal_penalties.iter().enumerate() {
            let mut kron_term = Array2::<f64>::eye(1);
            for (other_axis, other_penalty) in marginal_penalties.iter().enumerate() {
                let factor = if axis == other_axis {
                    penalty.clone()
                } else {
                    Array2::<f64>::eye(other_penalty.nrows())
                };
                kron_term = crate::construction::kronecker_product(&kron_term, &factor);
            }
            s_dense.scaled_add(lambdas[axis], &kron_term);
        }
        if ridge > 0.0 {
            for idx in 0..p_total {
                s_dense[[idx, idx]] += ridge;
            }
        }
        let (evals_dense, evecs_dense): (Array1<f64>, Array2<f64>) = s_dense
            .eigh(faer::Side::Lower)
            .expect("dense Kronecker eigh");
        let tol = 1e-12;
        let positive_indices: Vec<usize> = evals_dense
            .iter()
            .enumerate()
            .filter_map(|(idx, &value)| (value > tol).then_some(idx))
            .collect();
        let logdet = positive_indices
            .iter()
            .map(|&idx| evals_dense[idx].ln())
            .sum();
        let mut grad = Array1::<f64>::zeros(lambdas.len());
        let mut hess = Array2::<f64>::zeros((lambdas.len(), lambdas.len()));
        for (axis, penalty) in marginal_penalties.iter().enumerate() {
            let mut kron_term = Array2::<f64>::eye(1);
            for (other_axis, other_penalty) in marginal_penalties.iter().enumerate() {
                let factor = if axis == other_axis {
                    penalty.clone()
                } else {
                    Array2::<f64>::eye(other_penalty.nrows())
                };
                kron_term = crate::construction::kronecker_product(&kron_term, &factor);
            }
            for &eig_idx in &positive_indices {
                let eigval = evals_dense[eig_idx];
                let eigvec = evecs_dense.column(eig_idx).to_owned();
                let projected = kron_term.dot(&eigvec);
                let ck = lambdas[axis] * eigvec.dot(&projected);
                grad[axis] += ck / eigval;
                hess[[axis, axis]] += ck / eigval - (ck * ck) / (eigval * eigval);
                for other_axis in (axis + 1)..lambdas.len() {
                    let mut other_kron = Array2::<f64>::eye(1);
                    for (inner_axis, inner_penalty) in marginal_penalties.iter().enumerate() {
                        let factor = if other_axis == inner_axis {
                            inner_penalty.clone()
                        } else {
                            Array2::<f64>::eye(inner_penalty.nrows())
                        };
                        other_kron = crate::construction::kronecker_product(&other_kron, &factor);
                    }
                    let other_projected = other_kron.dot(&eigvec);
                    let cl = lambdas[other_axis] * eigvec.dot(&other_projected);
                    let off = -(ck * cl) / (eigval * eigval);
                    hess[[axis, other_axis]] += off;
                    hess[[other_axis, axis]] += off;
                }
            }
        }
        (logdet, grad, hess)
    }

    fn max_abs_diff_vector(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
        assert_eq!(a.len(), b.len());
        a.iter()
            .zip(b.iter())
            .map(|(&x, &y)| (x - y).abs())
            .fold(0.0_f64, f64::max)
    }

    #[test]
    fn kronecker_penalty_system_logdet_matches_dense_reference() {
        let q1 = 3usize;
        let q2 = 4usize;
        let s1 = array![[1.0, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 1.0]];
        let s2 = array![
            [1.0, -1.0, 0.0, 0.0],
            [-1.0, 2.0, -1.0, 0.0],
            [0.0, -1.0, 2.0, -1.0],
            [0.0, 0.0, -1.0, 1.0]
        ];
        let marginal_penalties = vec![s1, s2];
        let lambdas = vec![2.5, 1.3];
        let ridge = 0.0;

        let system = KroneckerPenaltySystem::new(marginal_penalties.clone(), vec![q1, q2], false)
            .expect("KroneckerPenaltySystem");
        let (logdet, grad, hess) = system.logdet_and_derivatives(&lambdas, ridge);
        let (dense_logdet, dense_grad, dense_hess) =
            dense_kronecker_pseudo_logdet_reference(&marginal_penalties, &lambdas, ridge);

        assert!(
            (logdet - dense_logdet).abs() < 1e-8,
            "KroneckerPenaltySystem logdet mismatch: factored={} dense={}",
            logdet,
            dense_logdet
        );
        let grad_diff = max_abs_diff_vector(&grad, &dense_grad);
        assert!(
            grad_diff < 1e-8,
            "KroneckerPenaltySystem gradient mismatch: max diff={grad_diff}"
        );
        let hess_diff = max_abs_diff_matrix(&hess, &dense_hess);
        assert!(
            hess_diff < 1e-8,
            "KroneckerPenaltySystem Hessian mismatch: max diff={hess_diff}"
        );
    }

    fn assert_term_collection_designs_match(
        left: &TermCollectionDesign,
        right: &TermCollectionDesign,
        label: &str,
    ) {
        let left_design = left.design.to_dense();
        let right_design = right.design.to_dense();
        let design_diff = max_abs_diff_matrix(&left_design, &right_design);
        assert!(
            design_diff <= 1e-10,
            "{label} design mismatch max_abs={design_diff}"
        );
        assert_eq!(
            left.penalties.len(),
            right.penalties.len(),
            "{label} penalty count mismatch"
        );
        for (idx, (lp, rp)) in left
            .penalties
            .iter()
            .zip(right.penalties.iter())
            .enumerate()
        {
            assert_eq!(
                lp.col_range, rp.col_range,
                "{label} penalty {idx} col_range mismatch"
            );
            let penalty_diff = max_abs_diff_matrix(&lp.local, &rp.local);
            assert!(
                penalty_diff <= 1e-10,
                "{label} penalty {idx} mismatch max_abs={penalty_diff}"
            );
        }
        assert_eq!(
            left.nullspace_dims, right.nullspace_dims,
            "{label} nullspace dims mismatch"
        );
        assert_eq!(
            left.penaltyinfo.len(),
            right.penaltyinfo.len(),
            "{label} penaltyinfo length mismatch"
        );
        for (idx, (linfo, rinfo)) in left
            .penaltyinfo
            .iter()
            .zip(right.penaltyinfo.iter())
            .enumerate()
        {
            assert_eq!(
                linfo.termname, rinfo.termname,
                "{label} penaltyinfo termname mismatch at {idx}"
            );
            assert_eq!(
                linfo.penalty.source, rinfo.penalty.source,
                "{label} penalty source mismatch at {idx}"
            );
            assert_eq!(
                linfo.penalty.active, rinfo.penalty.active,
                "{label} penalty active mismatch at {idx}"
            );
            assert_eq!(
                linfo.penalty.effective_rank, rinfo.penalty.effective_rank,
                "{label} penalty rank mismatch at {idx}"
            );
            assert_eq!(
                linfo.penalty.nullspace_dim_hint, rinfo.penalty.nullspace_dim_hint,
                "{label} penalty nullspace hint mismatch at {idx}"
            );
            assert!(
                (linfo.penalty.normalization_scale - rinfo.penalty.normalization_scale).abs()
                    <= 1e-10,
                "{label} penalty normalization mismatch at {idx}"
            );
        }
        match (
            left.coefficient_lower_bounds.as_ref(),
            right.coefficient_lower_bounds.as_ref(),
        ) {
            (Some(lb_left), Some(lb_right)) => {
                let diff = max_abs_diff_vector(lb_left, lb_right);
                assert!(diff <= 1e-10, "{label} lower-bound mismatch max_abs={diff}");
            }
            (None, None) => {}
            _ => panic!("{label} lower-bound presence mismatch"),
        }
        match (
            left.linear_constraints.as_ref(),
            right.linear_constraints.as_ref(),
        ) {
            (Some(c_left), Some(c_right)) => {
                let a_diff = max_abs_diff_matrix(&c_left.a, &c_right.a);
                let b_diff = max_abs_diff_vector(&c_left.b, &c_right.b);
                assert!(
                    a_diff <= 1e-10,
                    "{label} linear-constraint A mismatch max_abs={a_diff}"
                );
                assert!(
                    b_diff <= 1e-10,
                    "{label} linear-constraint b mismatch max_abs={b_diff}"
                );
            }
            (None, None) => {}
            _ => panic!("{label} linear-constraint presence mismatch"),
        }
    }

    #[test]
    fn smooth_design_assembles_terms_and_penalties() {
        let data = array![
            [0.0, 0.0, 0.2],
            [0.2, 0.1, 0.4],
            [0.4, 0.2, 0.6],
            [0.6, 0.4, 0.7],
            [0.8, 0.7, 0.9],
            [1.0, 1.0, 1.1]
        ];

        let terms = vec![
            SmoothTermSpec {
                name: "s_x0".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 4,
                        },
                        double_penalty: true,
                        identifiability: BSplineIdentifiability::default(),
                        boundary: OneDimensionalBoundary::Open,
                        boundary_conditions: BSplineBoundaryConditions::default(),
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            },
            SmoothTermSpec {
                name: "tps_x1x2".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![1, 2],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: true,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            },
        ];

        let sd = build_smooth_design(data.view(), &terms).unwrap();
        assert_eq!(sd.nrows(), data.nrows());
        assert_eq!(sd.terms.len(), 2);
        // bspline double-penalty contributes two blocks; tps double-penalty also
        // contributes two blocks (bending + nullspace ridge).
        assert_eq!(sd.penalties.len(), 4);
        assert_eq!(sd.nullspace_dims.len(), 4);
        for bp in &sd.penalties {
            assert_eq!(bp.local.nrows(), bp.block_size());
            assert_eq!(bp.local.ncols(), bp.block_size());
            assert!(bp.col_range.end <= sd.total_smooth_cols());
        }
    }

    #[test]
    fn shape_mapping_monotone_increasing_is_non_decreasing() {
        let theta = array![-1.0, 0.5, -0.2, 0.3];
        let beta = SmoothDesign::map_term_coefficients(&theta, ShapeConstraint::MonotoneIncreasing)
            .unwrap();
        for i in 1..beta.len() {
            assert!(beta[i] >= beta[i - 1]);
        }
    }

    #[test]
    fn build_smooth_design_rejectsmultiaxis_spatial_shape_constraints() {
        let data = array![[0.0, 0.0], [0.5, 0.2], [1.0, 0.4], [1.5, 0.6],];
        let terms = vec![SmoothTermSpec {
            name: "tps_shape".to_string(),
            basis: SmoothBasisSpec::ThinPlate {
                feature_cols: vec![0, 1],
                spec: ThinPlateBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 3 },
                    length_scale: 1.0,
                    double_penalty: false,
                    identifiability: SpatialIdentifiability::default(),
                    radial_reparam: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::MonotoneIncreasing,
            joint_null_rotation: None,
        }];

        let err = build_smooth_design(data.view(), &terms).expect_err("shape should be rejected");
        match err {
            BasisError::InvalidInput(msg) => {
                assert!(msg.contains("requires exactly 1 feature axis"));
            }
            other => panic!("unexpected error: {other:?}"),
        }
    }

    #[test]
    fn build_smooth_design_accepts_monotone_thin_plate_1dwith_linear_constraints() {
        // 1D TPS with `num_centers = 4` requests exactly 4 centers. The
        // polynomial nullspace is represented by separate columns, not by
        // hidden extra knots.
        let data = array![[0.0], [0.15], [0.35], [0.5], [0.65], [0.85], [1.0]];
        let terms = vec![SmoothTermSpec {
            name: "mono_tps".to_string(),
            basis: SmoothBasisSpec::ThinPlate {
                feature_cols: vec![0],
                spec: ThinPlateBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                    length_scale: 1.0,
                    double_penalty: false,
                    identifiability: SpatialIdentifiability::default(),
                    radial_reparam: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::MonotoneIncreasing,
            joint_null_rotation: None,
        }];
        let sd = build_smooth_design(data.view(), &terms).expect("shape-constrained thin-plate");
        assert!(sd.coefficient_lower_bounds.is_none());
        let lin = sd
            .linear_constraints
            .as_ref()
            .expect("linear constraints should be generated");
        assert!(lin.a.nrows() > 0);
        assert_eq!(lin.a.ncols(), sd.total_smooth_cols());
        assert_eq!(lin.b.len(), lin.a.nrows());
    }

    #[test]
    fn build_smooth_design_auto_promotes_thin_plate_below_canonical_polynomial_dimension() {
        // d=3 with k=3 centers is below the canonical TPS polynomial-nullspace
        // size M(3, m=2) = 4. The basis builder auto-promotes to a pure Duchon
        // spline (Riesz-fractional generalization) — the principled fix for
        // canonical-TPS infeasibility — instead of rejecting. The resulting
        // metadata is Duchon, confirming the route fired.
        let data = array![
            [0.0, 0.0, 0.0],
            [0.2, 0.1, 0.3],
            [0.4, 0.3, 0.5],
            [0.7, 0.6, 0.8],
        ];
        let terms = vec![SmoothTermSpec {
            name: "thinplate(pc1, pc2, pc3)".to_string(),
            basis: SmoothBasisSpec::ThinPlate {
                feature_cols: vec![0, 1, 2],
                spec: ThinPlateBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 3 },
                    length_scale: 1.0,
                    double_penalty: false,
                    identifiability: SpatialIdentifiability::default(),
                    radial_reparam: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let sd = build_smooth_design(data.view(), &terms)
            .expect("auto-promotion to Duchon should succeed at infeasible canonical (d, k)");
        let metadata = &sd.terms.first().expect("at least one smooth term").metadata;
        assert!(
            matches!(metadata, BasisMetadata::Duchon { .. }),
            "expected Duchon metadata after auto-promotion, got {metadata:?}"
        );
    }

    #[test]
    fn freeze_term_collection_handles_thin_plate_auto_promotion_to_duchon() {
        // Reproducer for the freezer falling into its catch-all "smooth
        // metadata/spec type mismatch" arm whenever `build_thin_plate_basis`
        // delegates to `build_duchon_basis` (the auto-promotion path that
        // fires whenever canonical TPS is mathematically infeasible at the
        // requested d, k).  Without the rewrite step in
        // `freeze_term_collection_from_design`, the (ThinPlate spec, Duchon
        // metadata) pairing aborts the entire fit at serialization time even
        // though the fit itself succeeded against the promoted Duchon basis.
        //
        // d=5, k=10 hits the auto-promotion branch (canonical TPS at d=5 needs
        // M(5, m=3)=21 polynomial columns, above k=10) AND the Duchon fallback
        // is admissible (Linear nullspace at p=2 needs m_poly=6 centers, so
        // k=10 ≥ 6, with the smallest s satisfying both 2(p+s) > d and
        // 2s < d giving s=1).
        let mut rng = StdRng::seed_from_u64(20260504);
        let n = 200usize;
        let mut data = Array2::<f64>::zeros((n, 5));
        for i in 0..n {
            for j in 0..5 {
                data[[i, j]] = rng.random_range(-1.0..1.0);
            }
        }
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "thinplate(pc1, pc2, pc3, pc4, pc5)".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1, 2, 3, 4],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 10 },
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let fit_design = build_term_collection_design(data.view(), &spec).expect("fit-time design");
        // Confirm we actually exercised the auto-promotion branch.
        let metadata = &fit_design
            .smooth
            .terms
            .first()
            .expect("at least one smooth term")
            .metadata;
        assert!(
            matches!(metadata, BasisMetadata::Duchon { .. }),
            "expected auto-promotion to Duchon, got {metadata:?}"
        );

        let frozen = freeze_term_collection_from_design(&spec, &fit_design).expect(
            "freeze must succeed across the auto-promoted (ThinPlate spec, Duchon metadata) pair",
        );
        assert!(
            matches!(frozen.smooth_terms[0].basis, SmoothBasisSpec::Duchon { .. }),
            "frozen spec should reflect the auto-promotion as a Duchon variant"
        );

        // Predict-time replay must reproduce the fit-time design bit-for-bit:
        // the frozen Duchon spec carries the exact centers, power, and
        // nullspace_order that the basis builder selected during the fit.
        let replay_design =
            build_term_collection_design(data.view(), &frozen).expect("replay design");
        let max_abs = fit_design
            .design
            .to_dense()
            .iter()
            .zip(replay_design.design.to_dense().iter())
            .map(|(&a, &b)| (a - b).abs())
            .fold(0.0_f64, f64::max);
        assert!(
            max_abs <= 1e-10,
            "auto-promoted frozen replay changed realized design: max_abs={max_abs}"
        );
    }

    #[test]
    fn build_smooth_design_accepts_monotone_matern_1dwith_linear_constraints() {
        let data = array![[0.0], [0.2], [0.4], [0.6], [0.8], [1.0]];
        let terms = vec![SmoothTermSpec {
            name: "mono_matern".to_string(),
            basis: SmoothBasisSpec::Matern {
                feature_cols: vec![0],
                spec: MaternBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                    length_scale: 0.7,
                    nu: MaternNu::FiveHalves,
                    include_intercept: false,
                    double_penalty: false,
                    identifiability: MaternIdentifiability::CenterSumToZero,
                    aniso_log_scales: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::MonotoneIncreasing,
            joint_null_rotation: None,
        }];
        let sd = build_smooth_design(data.view(), &terms).expect("shape-constrained Matérn");
        assert!(sd.coefficient_lower_bounds.is_none());
        let lin = sd
            .linear_constraints
            .as_ref()
            .expect("linear constraints should be generated");
        assert!(lin.a.nrows() > 0);
        assert_eq!(lin.a.ncols(), sd.total_smooth_cols());
        assert_eq!(lin.b.len(), lin.a.nrows());
    }

    #[test]
    fn build_smooth_design_accepts_monotone_duchon_1dwith_linear_constraints() {
        let data = array![[0.0], [0.2], [0.4], [0.6], [0.8], [1.0]];
        let terms = vec![SmoothTermSpec {
            name: "mono_duchon".to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: vec![0],
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                    length_scale: Some(0.9),
                    power: 5.0,
                    nullspace_order: DuchonNullspaceOrder::Linear,
                    identifiability: SpatialIdentifiability::OrthogonalToParametric,
                    aniso_log_scales: None,
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::MonotoneIncreasing,
            joint_null_rotation: None,
        }];
        let sd = build_smooth_design(data.view(), &terms).expect("shape-constrained Duchon");
        assert!(sd.coefficient_lower_bounds.is_none());
        let lin = sd
            .linear_constraints
            .as_ref()
            .expect("linear constraints should be generated");
        assert!(lin.a.nrows() > 0);
        assert_eq!(lin.a.ncols(), sd.total_smooth_cols());
        assert_eq!(lin.b.len(), lin.a.nrows());
    }

    #[test]
    fn build_smooth_design_accepts_monotone_bsplinewith_bounds() {
        let data = array![[0.0], [0.25], [0.5], [0.75], [1.0]];
        let terms = vec![SmoothTermSpec {
            name: "mono_bs".to_string(),
            basis: SmoothBasisSpec::BSpline1D {
                feature_col: 0,
                spec: BSplineBasisSpec {
                    degree: 3,
                    penalty_order: 2,
                    knotspec: BSplineKnotSpec::Generate {
                        data_range: (0.0, 1.0),
                        num_internal_knots: 3,
                    },
                    double_penalty: false,
                    identifiability: BSplineIdentifiability::default(),
                    boundary: OneDimensionalBoundary::Open,
                    boundary_conditions: BSplineBoundaryConditions::default(),
                },
            },
            shape: ShapeConstraint::MonotoneIncreasing,
            joint_null_rotation: None,
        }];
        let sd = build_smooth_design(data.view(), &terms).expect("shape-constrained bspline");
        let lb = sd
            .coefficient_lower_bounds
            .as_ref()
            .expect("lower bounds should be generated");
        assert_eq!(lb.len(), sd.total_smooth_cols());
        assert!(lb[0].is_infinite() && lb[0].is_sign_negative());
        for j in 1..lb.len() {
            assert_eq!(lb[j], 0.0);
        }
    }

    #[test]
    fn term_collection_design_combines_linear_and_smooth() {
        let data = array![
            [0.0, 0.0, 0.2],
            [0.2, 0.1, 0.4],
            [0.4, 0.2, 0.6],
            [0.6, 0.4, 0.7],
            [0.8, 0.7, 0.9],
            [1.0, 1.0, 1.1]
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin_x0".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: true,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_x1x2".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![1, 2],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: true,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec).unwrap();
        let design_dense = design.design.to_dense();
        assert_eq!(design.design.nrows(), data.nrows());
        assert_eq!(design.intercept_range, 0..1);
        assert!(
            design_dense
                .column(design.intercept_range.start)
                .iter()
                .all(|&v: &f64| (v - 1.0).abs() < 1e-12)
        );
        assert!(design.design.ncols() >= 2);
        assert_eq!(design.linear_ranges.len(), 1);
        assert_eq!(design.random_effect_ranges.len(), 0);
        assert_eq!(design.penalties.len(), 3); // linear ridge + 2 smooth penalties (bending + nullspace)
        assert_eq!(design.nullspace_dims.len(), 3);
    }

    #[test]
    fn term_collection_design_keeps_intercept_plus_bspline_sparse() {
        // Use `BSplineIdentifiability::None` so the smooth block stays sparse
        // through `build_bspline_basis_1d`: the default `WeightedSumToZero`
        // policy is deliberately densified by `apply_sum_to_zero_constraint_sparse`
        // (orthonormal Z, so ZZᵀ projects), which would prevent the full
        // assembled design from ever landing in `DesignMatrix::Sparse`
        // — that is, the pin only makes sense for the None branch.
        let n = 96usize;
        let x = Array1::linspace(0.0, 1.0, n);
        let mut data = Array2::<f64>::zeros((n, 1));
        data.column_mut(0).assign(&x);
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "s_x".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 32,
                        },
                        double_penalty: false,
                        identifiability: BSplineIdentifiability::None,
                        boundary: OneDimensionalBoundary::Open,
                        boundary_conditions: BSplineBoundaryConditions::default(),
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design =
            build_term_collection_design(data.view(), &spec).expect("term collection design");
        assert!(
            matches!(design.design, DesignMatrix::Sparse(_)),
            "expected sparse full design, got {:?}",
            design.design
        );
    }

    #[test]
    fn spatial_smooth_columns_do_not_duplicate_global_intercept() {
        let data = array![
            [0.0, 0.0],
            [0.2, 0.1],
            [0.4, 0.3],
            [0.6, 0.6],
            [0.8, 0.7],
            [1.0, 1.0],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).unwrap();
        let design_dense = design.design.to_dense();
        let smooth_start = 1usize;
        let smooth_end = smooth_start + design.smooth.total_smooth_cols();
        for col in smooth_start..smooth_end {
            let is_all_ones = design_dense
                .column(col)
                .iter()
                .all(|&v: &f64| (v - 1.0).abs() < 1e-12);
            assert!(
                !is_all_ones,
                "smooth column {col} unexpectedly duplicated intercept"
            );
        }
    }

    #[test]
    fn spatial_smooth_drops_matching_linear_trend_columns() {
        let data = array![
            [0.0, 0.1],
            [0.2, 0.0],
            [0.3, 0.4],
            [0.5, 0.2],
            [0.7, 0.9],
            [1.0, 0.8],
            [1.2, 1.1],
            [1.4, 1.3],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin_x0".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).unwrap();

        // Raw TPS width for k=4,d=2 is 4; we drop intercept + matching x0 linear component.
        assert_eq!(design.smooth.total_smooth_cols(), 2);

        let dense = design.design.to_dense_cow();
        let lin_col = design.linear_ranges[0].1.start;
        let linvalues = dense.column(lin_col).to_owned();
        let smooth_start = 1 + spec.linear_terms.len();
        let smooth_end = smooth_start + design.smooth.total_smooth_cols();
        for col in smooth_start..smooth_end {
            let same_as_linear = dense
                .column(col)
                .iter()
                .zip(linvalues.iter())
                .all(|(&a, &b)| (a - b).abs() < 1e-12);
            assert!(
                !same_as_linear,
                "smooth column {col} unexpectedly duplicated linear term column"
            );
        }
    }

    #[test]
    fn spatial_option5_is_orthogonal_to_parametric_block() {
        let data = array![
            [0.0, 0.1],
            [0.2, 0.0],
            [0.3, 0.4],
            [0.5, 0.2],
            [0.7, 0.9],
            [1.0, 0.8],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin_x0".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::OrthogonalToParametric,
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).unwrap();
        let design_dense = design.design.to_dense();
        let n = data.nrows();
        let mut c = Array2::<f64>::zeros((n, 2));
        c.column_mut(0).fill(1.0);
        c.column_mut(1).assign(&data.column(0));
        let smooth_start = 1 + spec.linear_terms.len();
        let b = design_dense
            .slice(s![
                ..,
                smooth_start..(smooth_start + design.smooth.total_smooth_cols())
            ])
            .to_owned();
        let cross = b.t().dot(&c);
        let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
        let b_norm = b.iter().map(|v| v * v).sum::<f64>().sqrt();
        let c_norm = c.iter().map(|v| v * v).sum::<f64>().sqrt();
        let rel = num / (b_norm * c_norm).max(1e-300);
        assert!(
            rel <= 1e-10,
            "smooth residual against model-owned parametric block too large: {rel}"
        );
    }

    #[test]
    fn thin_plate_default_identifiability_centers_against_intercept_only_without_linear_terms() {
        // Without any explicit linear term in the formula, the parametric
        // constraint block built by `build_parametric_constraint_block_for_term`
        // contains only the intercept column — see the
        // `SpatialIdentifiability` docs:
        //
        //   "The term-collection builder augments `C` with explicit linear
        //    terms when those terms are present in the formula."
        //
        // So a standalone TPS smooth marked `OrthogonalToParametric` is
        // orthogonalized only against `[1]`; its full polynomial nullspace
        // (the linear axes that thin-plate splines own as part of their
        // canonical model surface) stays in the smooth's column span.
        // Companions: `standalone_tps_keeps_centered_linear_nullspace` and
        // `term_collection_joint_duchon_carries_frozen_transform_into_metadata`
        // assert the dimension count from the same contract.
        let data = array![
            [-1.9, -1.2],
            [-1.3, -0.7],
            [-0.8, -0.4],
            [-0.2, 0.1],
            [0.0, 0.3],
            [0.4, 0.5],
            [0.9, 0.8],
            [1.4, 1.1],
            [1.9, 1.5],
            [2.3, 1.8],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: (0..2)
                .map(|feature| SmoothTermSpec {
                    name: format!("tps_x{feature}"),
                    basis: SmoothBasisSpec::ThinPlate {
                        feature_cols: vec![feature],
                        spec: ThinPlateBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::EqualMass { num_centers: 4 },
                            length_scale: 1.0,
                            double_penalty: false,
                            identifiability: SpatialIdentifiability::OrthogonalToParametric,
                            radial_reparam: None,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                })
                .collect(),
        };

        let design = build_term_collection_design(data.view(), &spec).unwrap();
        let design_dense = design.design.to_dense();
        let smooth_start = 1 + spec.linear_terms.len();
        let intercept = Array2::<f64>::ones((data.nrows(), 1));
        for (term_idx, term) in design.smooth.terms.iter().enumerate() {
            let block = design_dense
                .slice(s![
                    ..,
                    (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end)
                ])
                .to_owned();
            let cross = block.t().dot(&intercept);
            let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
            let block_norm = block.iter().map(|v| v * v).sum::<f64>().sqrt();
            let intercept_norm = intercept.iter().map(|v| v * v).sum::<f64>().sqrt();
            let rel = num / (block_norm * intercept_norm).max(1e-300);
            assert!(
                rel <= 1e-10,
                "ThinPlate term {term_idx} should be centered against the intercept (no linear terms in formula); got rel={rel:.3e}"
            );
        }
    }

    #[test]
    fn spatial_option5_does_not_overconstrain_on_nonoverlapping_linear_terms() {
        let n = 40usize;
        let p = 16usize;
        let mut data = Array2::<f64>::zeros((n, p));
        for i in 0..n {
            for j in 0..p {
                // Deterministic, non-collinear synthetic PCs.
                data[[i, j]] =
                    (i as f64) * 0.03 + (j as f64) * 0.11 + ((i * (j + 1)) as f64) * 1e-3;
            }
        }

        let spec = TermCollectionSpec {
            linear_terms: (5..16)
                .map(|j| LinearTermSpec {
                    name: format!("pc{j}"),
                    feature_col: j,
                    feature_cols: vec![j],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                })
                .collect(),
            random_effect_terms: vec![],
            smooth_terms: vec![
                SmoothTermSpec {
                    name: "tps_pc1".to_string(),
                    basis: SmoothBasisSpec::ThinPlate {
                        feature_cols: vec![1],
                        spec: ThinPlateBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::FarthestPoint { num_centers: 12 },
                            length_scale: 1.0,
                            double_penalty: true,
                            identifiability: SpatialIdentifiability::OrthogonalToParametric,
                            radial_reparam: None,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "tps_pc2".to_string(),
                    basis: SmoothBasisSpec::ThinPlate {
                        feature_cols: vec![2],
                        spec: ThinPlateBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::FarthestPoint { num_centers: 12 },
                            length_scale: 1.0,
                            double_penalty: true,
                            identifiability: SpatialIdentifiability::OrthogonalToParametric,
                            radial_reparam: None,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
            ],
        };

        let out = build_term_collection_design(data.view(), &spec);
        assert!(
            out.is_ok(),
            "term-local Option 5 should not over-constrain non-overlapping smooth/linear terms: {:?}",
            out.err()
        );
    }

    #[test]
    fn overlapping_linear_term_residualizes_bspline_smooth() {
        let data = array![
            [0.0],
            [0.1],
            [0.2],
            [0.3],
            [0.4],
            [0.5],
            [0.6],
            [0.7],
            [0.8],
            [0.9],
            [1.0],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "x".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "s_x".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 4,
                        },
                        double_penalty: false,
                        identifiability: BSplineIdentifiability::default(),
                        boundary: OneDimensionalBoundary::Open,
                        boundary_conditions: BSplineBoundaryConditions::default(),
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("bspline design");
        let mut c = Array2::<f64>::zeros((data.nrows(), 2));
        c.column_mut(0).fill(1.0);
        c.column_mut(1).assign(&data.column(0));
        let rel =
            orthogonality_relative_residual_for_design(&design.smooth.term_designs[0], c.view())
                .expect("orthogonality residual");
        assert!(
            rel <= 1e-10,
            "B-spline smooth should be orthogonal to [1, x] when linear(x) is present; rel={rel}"
        );
    }

    #[test]
    fn standalone_tps_keeps_centered_linear_nullspace() {
        let data = array![[-1.5], [-0.7], [0.2], [0.8], [1.6]];
        let centers = array![[-1.5], [0.2], [1.6]];
        let smooth = SmoothTermSpec {
            name: "s_x".to_string(),
            basis: SmoothBasisSpec::ThinPlate {
                feature_cols: vec![0],
                spec: ThinPlateBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::UserProvided(centers),
                    length_scale: 1.0,
                    double_penalty: false,
                    identifiability: SpatialIdentifiability::OrthogonalToParametric,
                    radial_reparam: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![smooth],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("tps design");

        assert_eq!(design.smooth.term_designs[0].ncols(), 2);
        assert_eq!(design.smooth.nullspace_dims, vec![1]);
        let intercept = Array2::<f64>::ones((data.nrows(), 1));
        let rel = orthogonality_relative_residual_for_design(
            &design.smooth.term_designs[0],
            intercept.view(),
        )
        .expect("intercept residual");
        assert!(
            rel <= 1e-10,
            "standalone TPS should be centered against the intercept while retaining its linear nullspace; rel={rel}"
        );
    }

    #[test]
    fn spatial_parametric_ownership_projects_only_explicit_linear_axes() {
        let term = SmoothTermSpec {
            name: "s_xy".to_string(),
            basis: SmoothBasisSpec::ThinPlate {
                feature_cols: vec![0, 1],
                spec: ThinPlateBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::EqualMass { num_centers: 4 },
                    length_scale: 1.0,
                    double_penalty: false,
                    identifiability: SpatialIdentifiability::OrthogonalToParametric,
                    radial_reparam: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let linear_terms = vec![LinearTermSpec {
            name: "x0".to_string(),
            feature_col: 0,
            feature_cols: vec![0],
            double_penalty: false,
            coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
            coefficient_min: None,
            coefficient_max: None,
        }];

        assert_eq!(
            smooth_intrinsic_parametric_feature_cols(&linear_terms, &term),
            vec![0],
            "a linear term on x0 should not claim the smooth's x1 nullspace"
        );
    }

    #[test]
    fn hierarchical_smooth_ownership_is_order_independent_for_bspline_and_duchon() {
        let data = array![
            [0.00, 0.00],
            [0.10, 0.15],
            [0.18, 0.30],
            [0.27, 0.10],
            [0.35, 0.55],
            [0.46, 0.25],
            [0.54, 0.70],
            [0.63, 0.40],
            [0.72, 0.85],
            [0.81, 0.60],
            [0.90, 0.95],
            [1.00, 0.75],
        ];

        let bspline_term = SmoothTermSpec {
            name: "s_x".to_string(),
            basis: SmoothBasisSpec::BSpline1D {
                feature_col: 0,
                spec: BSplineBasisSpec {
                    degree: 3,
                    penalty_order: 2,
                    knotspec: BSplineKnotSpec::Generate {
                        data_range: (0.0, 1.0),
                        num_internal_knots: 5,
                    },
                    double_penalty: false,
                    identifiability: BSplineIdentifiability::default(),
                    boundary: OneDimensionalBoundary::Open,
                    boundary_conditions: BSplineBoundaryConditions::default(),
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let duchon_term = SmoothTermSpec {
            name: "duchon_xy".to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: vec![0, 1],
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                    length_scale: Some(1.0),
                    power: 5.0,
                    nullspace_order: DuchonNullspaceOrder::Linear,
                    identifiability: SpatialIdentifiability::default(),
                    aniso_log_scales: None,
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let spec_a = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![duchon_term.clone(), bspline_term.clone()],
        };
        let spec_b = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![bspline_term, duchon_term],
        };

        let design_a = build_term_collection_design(data.view(), &spec_a).expect("design a");
        let design_b = build_term_collection_design(data.view(), &spec_b).expect("design b");

        for design in [&design_a, &design_b] {
            let owner_idx = design
                .smooth
                .terms
                .iter()
                .position(|term| term.name == "s_x")
                .expect("owner term");
            let target_idx = design
                .smooth
                .terms
                .iter()
                .position(|term| term.name == "duchon_xy")
                .expect("target term");
            let owner_dense = design.smooth.term_designs[owner_idx].to_dense();
            let rel = orthogonality_relative_residual_for_design(
                &design.smooth.term_designs[target_idx],
                owner_dense.view(),
            )
            .expect("orthogonality residual");
            assert!(
                rel <= 1e-10,
                "multivariate Duchon term should be residualized against owned 1D spline space; rel={rel}"
            );
        }

        let duchon_a_idx = design_a
            .smooth
            .terms
            .iter()
            .position(|term| term.name == "duchon_xy")
            .expect("duchon in design a");
        let duchon_b_idx = design_b
            .smooth
            .terms
            .iter()
            .position(|term| term.name == "duchon_xy")
            .expect("duchon in design b");
        let duchon_a = design_a.smooth.term_designs[duchon_a_idx].to_dense();
        let duchon_b = design_b.smooth.term_designs[duchon_b_idx].to_dense();
        assert_eq!(duchon_a.dim(), duchon_b.dim());
        let max_abs = duchon_a
            .iter()
            .zip(duchon_b.iter())
            .map(|(lhs, rhs)| (lhs - rhs).abs())
            .fold(0.0_f64, f64::max);
        assert!(
            max_abs <= 1e-10,
            "hierarchical ownership should not depend on user term order; max_abs={max_abs}"
        );
    }

    #[test]
    fn freeze_roundtrip_preserves_hierarchical_smooth_transforms() {
        let data = array![
            [0.00, 0.00],
            [0.10, 0.15],
            [0.18, 0.30],
            [0.27, 0.10],
            [0.35, 0.55],
            [0.46, 0.25],
            [0.54, 0.70],
            [0.63, 0.40],
            [0.72, 0.85],
            [0.81, 0.60],
            [0.90, 0.95],
            [1.00, 0.75],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "x".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![
                SmoothTermSpec {
                    name: "duchon_xy".to_string(),
                    basis: SmoothBasisSpec::Duchon {
                        feature_cols: vec![0, 1],
                        spec: DuchonBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                            length_scale: Some(1.0),
                            power: 1.0,
                            nullspace_order: DuchonNullspaceOrder::Linear,
                            identifiability: SpatialIdentifiability::default(),
                            aniso_log_scales: None,
                            operator_penalties: DuchonOperatorPenaltySpec::default(),
                            boundary: OneDimensionalBoundary::Open,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "s_x".to_string(),
                    basis: SmoothBasisSpec::BSpline1D {
                        feature_col: 0,
                        spec: BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (0.0, 1.0),
                                num_internal_knots: 5,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::default(),
                            boundary: OneDimensionalBoundary::Open,
                            boundary_conditions: BSplineBoundaryConditions::default(),
                        },
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
            ],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("fit-time design");
        let frozen =
            freeze_term_collection_from_design(&spec, &design).expect("freeze hierarchical design");
        let replay = build_term_collection_design(data.view(), &frozen).expect("replay design");

        let dense_fit = design.design.to_dense();
        let dense_replay = replay.design.to_dense();
        assert_eq!(dense_fit.dim(), dense_replay.dim());
        let max_abs = dense_fit
            .iter()
            .zip(dense_replay.iter())
            .map(|(lhs, rhs)| (lhs - rhs).abs())
            .fold(0.0_f64, f64::max);
        assert!(
            max_abs <= 1e-10,
            "frozen hierarchical transforms should replay exactly on the training data; max_abs={max_abs}"
        );
    }

    #[test]
    fn spatial_option5_preserves_lazy_thin_plate_terms_at_large_scale() {
        let n = 17_000usize;
        let k = 2_000usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut centers = Array2::<f64>::zeros((k, 1));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n - 1) as f64;
        }
        for j in 0..k {
            centers[[j, 0]] = j as f64 / (k - 1) as f64;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "x".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_x".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(centers),
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::OrthogonalToParametric,
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design =
            build_term_collection_design(data.view(), &spec).expect("large option-5 design");
        assert!(matches!(
            &design.smooth.term_designs[0],
            DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::Lazy(_))
        ));
        let mut c = Array2::<f64>::zeros((n, 2));
        c.column_mut(0).fill(1.0);
        c.column_mut(1).assign(&data.column(0));
        let rel =
            orthogonality_relative_residual_for_design(&design.smooth.term_designs[0], c.view())
                .expect("orthogonality residual");
        assert!(rel <= 1e-8, "lazy option-5 residual too large: {rel}");
    }

    #[test]
    fn spatial_frozen_transform_rebuild_is_exact_on_trainingrows() {
        let data = array![
            [0.0, 0.1],
            [0.2, 0.0],
            [0.3, 0.4],
            [0.5, 0.2],
            [0.7, 0.9],
            [1.0, 0.8],
        ];
        let fitspec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin_x0".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: 1.0,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::OrthogonalToParametric,
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_design = build_term_collection_design(data.view(), &fitspec).unwrap();
        let term_meta = &fit_design.smooth.terms[0].metadata;
        let (centers, length_scale, z) = match term_meta {
            BasisMetadata::ThinPlate {
                centers,
                length_scale,
                identifiability_transform,
                ..
            } => (
                centers.clone(),
                *length_scale,
                identifiability_transform
                    .clone()
                    .expect("fit-time Option 5 should store transform"),
            ),
            other => panic!("unexpected metadata variant: {other:?}"),
        };

        let frozenspec = TermCollectionSpec {
            linear_terms: fitspec.linear_terms.clone(),
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(centers),
                        length_scale,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::FrozenTransform { transform: z },
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let frozen_design = build_term_collection_design(data.view(), &frozenspec).unwrap();

        assert_eq!(
            fit_design.smooth.term_designs.len(),
            frozen_design.smooth.term_designs.len(),
            "frozen transform rebuild term count mismatch"
        );
        let max_abs = fit_design
            .smooth
            .term_designs
            .iter()
            .zip(frozen_design.smooth.term_designs.iter())
            .flat_map(|(a, b)| {
                let a_dense = a.to_dense();
                let b_dense = b.to_dense();
                assert_eq!(a_dense.dim(), b_dense.dim());
                a_dense
                    .iter()
                    .zip(b_dense.iter())
                    .map(|(&x, &y)| (x - y).abs())
                    .collect::<Vec<_>>()
            })
            .fold(0.0_f64, f64::max);
        assert!(
            max_abs <= 1e-12,
            "frozen transform rebuild mismatch max_abs={max_abs}"
        );
    }

    #[test]
    fn frozen_spatial_replay_preserves_standardized_length_scale_compensation() {
        assert!(file!().ends_with(".rs"));
        let n = 16usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = 0.07 * i as f64 + 0.02 * (3.0 * t).sin();
            data[[i, 1]] = 4.0 * t + 0.35 * (5.0 * t).cos();
        }

        let tps_spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "tps_xy".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 1.3,
                        double_penalty: true,
                        identifiability: SpatialIdentifiability::OrthogonalToParametric,
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        assert_frozen_replay_matches_fit(data.view(), &tps_spec, "thin-plate");

        let matern_spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern_xy".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 1.1,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        assert_frozen_replay_matches_fit(data.view(), &matern_spec, "matern");

        let duchon_spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_xy".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: Some(1.4),
                        power: 5.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::OrthogonalToParametric,
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        assert_frozen_replay_matches_fit(data.view(), &duchon_spec, "duchon");
    }

    #[test]
    fn term_collection_design_adds_random_effect_dummy_blockwithridge() {
        let data = array![
            [0.1, 0.0],
            [0.2, 1.0],
            [0.3, 0.0],
            [0.4, 2.0],
            [0.5, 1.0],
            [0.6, 2.0],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![RandomEffectTermSpec {
                name: "id".to_string(),
                feature_col: 1,
                drop_first_level: false,
                penalized: true,
                frozen_levels: None,
            }],
            smooth_terms: vec![],
        };
        let design = build_term_collection_design(data.view(), &spec).unwrap();
        assert_eq!(design.intercept_range, 0..1);
        // 3 observed levels -> 3 dummy columns
        assert_eq!(design.design.ncols(), 4);
        assert_eq!(design.random_effect_ranges.len(), 1);
        assert_eq!(design.penalties.len(), 1);
        assert_eq!(design.nullspace_dims, vec![0]);
        let (_, range) = &design.random_effect_ranges[0];
        let dense = design.design.to_dense_cow();
        for i in 0..dense.nrows() {
            let row_sum: f64 = dense.slice(s![i, range.clone()]).sum();
            assert!((row_sum - 1.0).abs() < 1e-12);
        }
    }

    #[test]
    fn matern_smooth_buildswith_double_penalty_in_high_dim() {
        let n = 12usize;
        let d = 10usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            for j in 0..d {
                data[[i, j]] = (i as f64) * 0.1 + (j as f64) * 0.03;
            }
        }

        let terms = vec![SmoothTermSpec {
            name: "matern_x".to_string(),
            basis: SmoothBasisSpec::Matern {
                feature_cols: (0..d).collect(),
                spec: MaternBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 5 },
                    length_scale: 0.75,
                    nu: MaternNu::FiveHalves,
                    include_intercept: false,
                    double_penalty: true,
                    identifiability: MaternIdentifiability::CenterSumToZero,
                    aniso_log_scales: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let sd = build_smooth_design(data.view(), &terms).unwrap();
        assert_eq!(sd.nrows(), n);
        assert_eq!(sd.terms.len(), 1);
        // Spatial smooths use the canonical operator penalty triplet:
        // mass + tension + stiffness.
        assert_eq!(sd.penalties.len(), 3);
        assert_eq!(sd.nullspace_dims.len(), 3);
    }

    #[test]
    fn duchon_linear_nullspace_builds_and_reports_nullspace_dim() {
        // DuchonNullspaceOrder::Linear in dimension d needs at least d+1
        // affinely independent centers to span [1, x_1, ..., x_d]. Data must
        // also genuinely vary along all d axes (not collapse to a 1-D
        // manifold) so FarthestPoint sampling can find ≥ d+1 affinely
        // independent centers; otherwise the polynomial-block rank drops
        // below d+1 at the centers and the radial kernel block over-
        // parameterizes the basis.
        let n = 20usize;
        let d = 10usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            for j in 0..d {
                // Independent per-(i,j) values via splitmix64-like mixing —
                // gives full d-rank data without the additive (i*c1 + j*c2)
                // collapse that puts rows on a 1-D affine manifold.
                let mut key = (i as u64).wrapping_mul(0x9E3779B97F4A7C15);
                key ^= (j as u64).wrapping_mul(0xBF58476D1CE4E5B9);
                key = (key ^ (key >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
                key = (key ^ (key >> 27)).wrapping_mul(0x94D049BB133111EB);
                let v = ((key ^ (key >> 31)) as f64) / (u64::MAX as f64);
                data[[i, j]] = v;
            }
        }

        let terms = vec![SmoothTermSpec {
            name: "duchon_x".to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: (0..d).collect(),
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 12 },
                    length_scale: Some(0.9),
                    power: 5.0,
                    nullspace_order: DuchonNullspaceOrder::Linear,
                    identifiability: SpatialIdentifiability::default(),
                    aniso_log_scales: None,
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let sd = build_smooth_design(data.view(), &terms).unwrap();
        assert_eq!(sd.nrows(), n);
        assert_eq!(sd.terms.len(), 1);
        assert_eq!(sd.penalties.len(), 4);
        assert_eq!(sd.nullspace_dims.len(), 4);
    }

    #[test]
    fn joint_duchon_orderzero_raw_smooth_build_preserves_unconstrained_basis() {
        let n = 12usize;
        let d = 4usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            for j in 0..d {
                data[[i, j]] = (i as f64) * 0.13 + (j as f64) * 0.17;
            }
        }

        let terms = vec![SmoothTermSpec {
            name: "duchon_joint".to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: (0..d).collect(),
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                    length_scale: Some(1.0),
                    power: 3.0,
                    nullspace_order: DuchonNullspaceOrder::Zero,
                    identifiability: SpatialIdentifiability::default(),
                    aniso_log_scales: None,
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let sd = build_smooth_design(data.view(), &terms).expect("joint duchon build");
        assert_eq!(sd.total_smooth_cols(), 4);
        match &sd.terms[0].metadata {
            BasisMetadata::Duchon {
                identifiability_transform,
                ..
            } => {
                assert!(
                    identifiability_transform.is_some(),
                    "raw smooth build should freeze Duchon orthogonality once the basis is built"
                );
            }
            other => panic!("expected Duchon metadata, got {other:?}"),
        }
    }

    #[test]
    fn term_collection_joint_duchon_carries_frozen_transform_into_metadata() {
        let n = 12usize;
        let d = 4usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            for j in 0..d {
                data[[i, j]] = (i as f64) * 0.13 + (j as f64) * 0.17;
            }
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_joint".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: (0..d).collect(),
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: Some(1.0),
                        power: 3.0,
                        nullspace_order: DuchonNullspaceOrder::Zero,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design =
            build_term_collection_design(data.view(), &spec).expect("term collection design");
        let term = &design.smooth.terms[0];
        assert_eq!(term.coeff_range.len(), 3);
        match &term.metadata {
            BasisMetadata::Duchon {
                identifiability_transform,
                ..
            } => {
                let z = identifiability_transform
                    .as_ref()
                    .expect("term collection should store frozen Duchon transform");
                assert_eq!(z.nrows(), 4);
                assert_eq!(z.ncols(), 3);
            }
            other => panic!("expected Duchon metadata, got {other:?}"),
        }
    }

    #[test]
    fn frozen_joint_maternspec_rebuild_keeps_adaptive_cache_in_sync() {
        let n = 12usize;
        let d = 2usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            data[[i, 0]] = i as f64 * 0.13;
            data[[i, 1]] = (i as f64 * 0.17).sin();
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern_joint".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: (0..d).collect(),
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 1.0,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("base design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze spec");
        let rebuilt = build_term_collection_design(data.view(), &frozen).expect("rebuilt design");
        let caches =
            extract_spatial_operator_runtime_caches(&frozen, &rebuilt).expect("adaptive caches");
        assert_eq!(caches.len(), 1);
        assert_eq!(caches[0].termname, "matern_joint");
        assert_eq!(rebuilt.smooth.terms.len(), 1);
        assert!(!rebuilt.smooth.terms[0].coeff_range.is_empty());
    }

    #[test]
    fn tensor_bspline_term_builds_te_style_design_and_penalties() {
        let n = 10usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
            data[[i, 1]] = (i as f64 / (n as f64 - 1.0)).powi(2);
        }

        let spec_x = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 3,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::default(),
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };
        let spec_y = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 2,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::default(),
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };

        let terms = vec![SmoothTermSpec {
            name: "te_xy".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    periods: Vec::new(),
                    marginalspecs: vec![spec_x, spec_y],
                    double_penalty: true,
                    identifiability: TensorBSplineIdentifiability::default(),
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let sd = build_smooth_design(data.view(), &terms).unwrap();
        assert_eq!(sd.nrows(), n);
        assert_eq!(sd.terms.len(), 1);
        // one Kronecker penalty per marginal + optional ridge
        assert_eq!(sd.penalties.len(), 3);
        assert_eq!(sd.nullspace_dims.len(), 3);
        assert!(
            sd.penalties
                .iter()
                .all(|bp| bp.local.nrows() == bp.block_size())
        );
        assert!(
            sd.penalties
                .iter()
                .all(|bp| bp.col_range.end <= sd.total_smooth_cols())
        );
    }

    #[test]
    fn tensor_binary_margin_is_penalized_factor_smooth_not_unidentified_raw_tensor() {
        let n = 18usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
            data[[i, 1]] = if i % 2 == 0 { 0.0 } else { 1.0 };
        }

        let age_margin = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 1,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };
        let binary_margin = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 1,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };
        let terms = vec![SmoothTermSpec {
            name: "te_age_binary".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    periods: Vec::new(),
                    marginalspecs: vec![age_margin, binary_margin],
                    double_penalty: false,
                    identifiability: TensorBSplineIdentifiability::None,
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        }];

        let design = build_smooth_design(data.view(), &terms).expect("binary-margin tensor");
        let kron = design.terms[0]
            .kronecker_factored
            .as_ref()
            .expect("tensor term should preserve Kronecker marginal metadata");

        assert_eq!(kron.marginal_dims, vec![5, 5]);
        assert_eq!(
            numerical_rank(&kron.marginal_designs[1]),
            2,
            "a binary tensor margin has two data-supported levels even when its raw spline margin has five columns"
        );
        assert_eq!(
            numerical_rank(&kron.marginal_penalties[1]),
            3,
            "the binary margin's second-difference roughness penalty must shrink the three unsupported spline-range directions"
        );
        assert_eq!(design.penalties.len(), 2);
        assert!(design.penalties.iter().all(|penalty| {
            penalty.local.nrows() == 25
                && penalty.local.ncols() == 25
                && penalty.local.iter().all(|value| value.is_finite())
        }));
    }

    #[test]
    fn centered_tensor_penalties_canonicalize_in_transformed_basis_width() {
        let n = 16usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = 0.5 + 0.25 * (2.0 * std::f64::consts::PI * t).sin();
        }

        let tensor_term = SmoothTermSpec {
            name: "te_centered".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    marginalspecs: vec![
                        BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (0.0, 1.0),
                                num_internal_knots: 3,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::default(),
                            boundary_conditions: Default::default(),
                            boundary: OneDimensionalBoundary::Open,
                        },
                        BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (0.0, 1.0),
                                num_internal_knots: 2,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::default(),
                            boundary_conditions: Default::default(),
                            boundary: OneDimensionalBoundary::Open,
                        },
                    ],
                    periods: Vec::new(),
                    double_penalty: false,
                    identifiability: TensorBSplineIdentifiability::default(),
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![tensor_term],
        };
        let design = build_term_collection_design(data.view(), &spec).unwrap();
        let penalty_specs = design
            .penalties
            .iter()
            .map(crate::estimate::PenaltySpec::from_blockwise_ref)
            .collect::<Vec<_>>();
        let (canonical, _) = crate::terms::construction::canonicalize_penalty_specs(
            &penalty_specs,
            &design.nullspace_dims,
            design.design.ncols(),
            "centered tensor penalty regression",
        )
        .unwrap();
        for cp in canonical {
            assert_eq!(cp.root.ncols(), cp.col_range.len());
            assert_eq!(cp.local.nrows(), cp.col_range.len());
            assert_eq!(cp.local.ncols(), cp.col_range.len());
        }
    }

    #[test]
    fn periodic_bspline_margin_wraps_exactly_at_period() {
        let x = array![0.0, 1.25, 2.5, 3.75, 7.0, 8.25];
        let spec = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::PeriodicUniform {
                data_range: (0.0, 1.0),
                num_basis: 8,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary_conditions: Default::default(),
            boundary: OneDimensionalBoundary::Open,
        };
        let built = build_bspline_basis_1d(x.view(), &spec).expect("periodic bspline");
        let dense = built.design.to_dense();
        assert_eq!(dense.ncols(), 8);
        for j in 0..dense.ncols() {
            assert!(
                (dense[[0, j]] - dense[[4, j]]).abs() < 1e-12,
                "seam row differs at column {j}: {} vs {}",
                dense[[0, j]],
                dense[[4, j]]
            );
            assert!(
                (dense[[1, j]] - dense[[5, j]]).abs() < 1e-12,
                "wrapped row differs at column {j}: {} vs {}",
                dense[[1, j]],
                dense[[5, j]]
            );
        }
        for row in dense.rows() {
            assert!((row.sum() - 1.0).abs() < 1e-12);
        }
        assert_eq!(built.nullspace_dims[0], 1);
    }

    #[test]
    fn tensor_bspline_supports_two_periodic_margins_as_torus() {
        let data = array![[0.0, 0.0], [7.0, 0.0], [0.0, 24.0], [7.0, 24.0], [1.5, 6.0]];
        let spec_day = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::PeriodicUniform {
                data_range: (0.0, 7.0),
                num_basis: 7,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary_conditions: Default::default(),
            boundary: OneDimensionalBoundary::Open,
        };
        let spec_hour = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::PeriodicUniform {
                data_range: (0.0, 24.0),
                num_basis: 8,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary_conditions: Default::default(),
            boundary: OneDimensionalBoundary::Open,
        };
        let spec_collection = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "te_day_hour".to_string(),
                basis: SmoothBasisSpec::TensorBSpline {
                    feature_cols: vec![0, 1],
                    spec: TensorBSplineSpec {
                        marginalspecs: vec![spec_day, spec_hour],
                        periods: Vec::new(),
                        double_penalty: false,
                        identifiability: TensorBSplineIdentifiability::None,
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec_collection)
            .expect("periodic tensor design");
        let sd = &design.smooth;
        let dense = sd.term_designs[0].to_dense();
        assert_eq!(dense.ncols(), 56);
        for j in 0..dense.ncols() {
            assert!((dense[[0, j]] - dense[[1, j]]).abs() < 1e-12);
            assert!((dense[[0, j]] - dense[[2, j]]).abs() < 1e-12);
            assert!((dense[[0, j]] - dense[[3, j]]).abs() < 1e-12);
        }
        assert_eq!(sd.penalties.len(), 2);
        assert!(sd.penalties.iter().all(|p| p.local.nrows() == 56));

        let frozen = freeze_term_collection_from_design(&spec_collection, &design)
            .expect("freeze periodic tensor");
        match &frozen.smooth_terms[0].basis {
            SmoothBasisSpec::TensorBSpline { spec, .. } => {
                assert!(matches!(
                    spec.marginalspecs[0].knotspec,
                    BSplineKnotSpec::PeriodicUniform { data_range, .. }
                        if (data_range.1 - data_range.0 - 7.0).abs() < 1e-9
                ));
                assert!(matches!(
                    spec.marginalspecs[1].knotspec,
                    BSplineKnotSpec::PeriodicUniform { data_range, .. }
                        if (data_range.1 - data_range.0 - 24.0).abs() < 1e-9
                ));
            }
            _ => panic!("expected tensor spec"),
        }
    }

    #[test]
    fn tensor_bspline_design_matches_extended_marginal_kronecker_product() {
        let data = array![[-0.2, 0.1], [0.2, 0.4], [0.7, 0.8], [1.2, 1.1],];
        let spec_x = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 3,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };
        let spec_y = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 1.0),
                num_internal_knots: 2,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary: OneDimensionalBoundary::Open,
            boundary_conditions: BSplineBoundaryConditions::default(),
        };
        let mx = build_bspline_basis_1d(data.column(0), &spec_x)
            .unwrap()
            .design
            .to_dense();
        let my = build_bspline_basis_1d(data.column(1), &spec_y)
            .unwrap()
            .design
            .to_dense();
        let expected = tensor_product_design_from_marginals(&[mx.clone(), my.clone()]).unwrap();

        let term = SmoothTermSpec {
            name: "te_xy".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    periods: Vec::new(),
                    marginalspecs: vec![spec_x, spec_y],
                    double_penalty: false,
                    identifiability: TensorBSplineIdentifiability::None,
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let got = build_smooth_design(data.view(), &[term])
            .unwrap()
            .term_designs
            .into_iter()
            .next()
            .unwrap()
            .to_dense();
        assert_eq!(got.dim(), expected.dim());
        for i in 0..got.nrows() {
            for j in 0..got.ncols() {
                assert!((got[[i, j]] - expected[[i, j]]).abs() < 1e-10);
            }
        }
    }

    #[test]
    fn tensor_bspline_periodic_margins_wrap_as_torus() {
        let data = array![[1.25, 3.5], [8.25, 27.5], [-5.75, -20.5], [1.25, 27.5]];
        let periodic_day = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 7.0),
                num_internal_knots: 4,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary_conditions: BSplineBoundaryConditions::default(),
            boundary: OneDimensionalBoundary::Cyclic {
                start: 0.0,
                end: 7.0,
            },
        };
        let periodic_hour = BSplineBasisSpec {
            degree: 3,
            penalty_order: 2,
            knotspec: BSplineKnotSpec::Generate {
                data_range: (0.0, 24.0),
                num_internal_knots: 4,
            },
            double_penalty: false,
            identifiability: BSplineIdentifiability::None,
            boundary_conditions: BSplineBoundaryConditions::default(),
            boundary: OneDimensionalBoundary::Cyclic {
                start: 0.0,
                end: 24.0,
            },
        };
        let term = SmoothTermSpec {
            name: "te_day_hour".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    marginalspecs: vec![periodic_day, periodic_hour],
                    periods: Vec::new(),
                    double_penalty: false,
                    identifiability: TensorBSplineIdentifiability::None,
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let design = build_smooth_design(data.view(), &[term])
            .unwrap()
            .term_designs
            .into_iter()
            .next()
            .unwrap()
            .to_dense();
        for j in 0..design.ncols() {
            assert!((design[[0, j]] - design[[1, j]]).abs() < 1e-12);
            assert!((design[[0, j]] - design[[2, j]]).abs() < 1e-12);
        }
    }

    #[test]
    fn tensor_bspline_design_is_identifiable_against_global_intercept() {
        let n = 120usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (3.0 * t).sin();
        }

        let tensor_term = SmoothTermSpec {
            name: "te_xy".to_string(),
            basis: SmoothBasisSpec::TensorBSpline {
                feature_cols: vec![0, 1],
                spec: TensorBSplineSpec {
                    periods: Vec::new(),
                    marginalspecs: vec![
                        BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (0.0, 1.0),
                                num_internal_knots: 6,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::default(),
                            boundary: OneDimensionalBoundary::Open,
                            boundary_conditions: BSplineBoundaryConditions::default(),
                        },
                        BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (-1.0, 1.0),
                                num_internal_knots: 6,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::default(),
                            boundary: OneDimensionalBoundary::Open,
                            boundary_conditions: BSplineBoundaryConditions::default(),
                        },
                    ],
                    double_penalty: false,
                    identifiability: TensorBSplineIdentifiability::default(),
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let sd = build_smooth_design(data.view(), &[tensor_term.clone()]).unwrap();
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![tensor_term],
        };
        let full = build_term_collection_design(data.view(), &spec).unwrap();
        let ones = Array1::<f64>::ones(n);
        let sd_dense_terms = sd
            .term_designs
            .iter()
            .map(|d| d.to_dense())
            .collect::<Vec<_>>();
        let sd_assembled = ndarray::concatenate(
            ndarray::Axis(1),
            &sd_dense_terms.iter().map(|d| d.view()).collect::<Vec<_>>(),
        )
        .unwrap();
        let residualvs_tensor = residual_norm_to_column_space(&sd_assembled, &ones);
        let full_design_dense = full.design.to_dense();
        let residualvs_full = residual_norm_to_column_space(&full_design_dense, &ones);

        // Tensor block alone must not be able to represent the constant surface.
        assert!(residualvs_tensor > 1e-6);
        // With explicit intercept, constants should be represented (near) exactly.
        assert!(residualvs_full < 1e-8);
    }

    /// Runs a Gaussian baseline fit and the two-outer-iteration spatial
    /// length-scale optimization for a single Matérn term, then asserts the
    /// optimized score is monotone-non-worse and that the resolved term froze
    /// its centers / identifiability transform with a finite in-range length
    /// scale. Shared verbatim between the 2- and 3-feature Matérn monotone
    /// pins, which differ only in their data generation, term dimensionality,
    /// and seed length scale.
    fn assert_matern_spatial_length_scale_optimization_monotone(
        data: ArrayView2<'_, f64>,
        y: &Array1<f64>,
        weights: &Array1<f64>,
        offset: &Array1<f64>,
        spec: &TermCollectionSpec,
        fit_opts: &FitOptions,
    ) {
        let baseline = fit_term_collection_forspec(
            data,
            y.view(),
            weights.view(),
            offset.view(),
            spec,
            LikelihoodSpec::gaussian_identity(),
            fit_opts,
        )
        .expect("baseline fit should succeed");
        let baseline_score = fit_score(&baseline.fit);

        let optimized = fit_term_collectionwith_spatial_length_scale_optimization(
            data,
            y.clone(),
            weights.clone(),
            offset.clone(),
            spec,
            LikelihoodSpec::gaussian_identity(),
            fit_opts,
            &SpatialLengthScaleOptimizationOptions {
                max_outer_iter: 2,
                rel_tol: 1e-5,
                pilot_subsample_threshold: 0,
                ..SpatialLengthScaleOptimizationOptions::default()
            },
        )
        .expect("optimized fit should succeed");
        let optimized_score = fit_score(&optimized.fit);
        assert!(optimized_score <= baseline_score + 1e-10);

        let ls = match &optimized.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => spec.length_scale,
            _ => panic!("expected Matérn term"),
        };
        assert!(ls.is_finite() && (1e-3..=1e3).contains(&ls));

        match &optimized.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => {
                assert!(matches!(
                    spec.center_strategy,
                    CenterStrategy::UserProvided(_)
                ));
                assert!(matches!(
                    spec.identifiability,
                    MaternIdentifiability::FrozenTransform { .. }
                ));
            }
            _ => panic!("expected Matérn term"),
        }
    }

    #[test]
    fn spatial_length_scale_optimization_monotone_improves_or_keeps_score_for_matern_two_feature() {
        let n = 60usize;
        let d = 3usize;
        let mut data = Array2::<f64>::zeros((n, d));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (i as f64 * 0.13).sin();
            let x2 = (i as f64 * 0.07).cos();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            data[[i, 2]] = x2;
            y[i] = (2.5 * x0).sin() + 0.4 * x1 - 0.2 * x2;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1, 2],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 12 },
                        length_scale: 20.0,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            max_iter: 40,
            ..FitOptions::default()
        };
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);

        assert_matern_spatial_length_scale_optimization_monotone(
            data.view(),
            &y,
            &weights,
            &offset,
            &spec,
            &fit_opts,
        );
    }

    /// Drives a two-block exact-joint κ optimization with the canonical
    /// zero-work test closures (cost = total design ncols + penalty count;
    /// flat gradient/Hessian; trivial EFS) and returns the resolved result.
    /// Shared verbatim across the Matérn- and Duchon-freezing pins; only the
    /// final `.expect` diagnostic differs, passed via `expect_msg`.
    fn run_two_block_exact_joint_optimize(
        data: ArrayView2<'_, f64>,
        meanspec: &TermCollectionSpec,
        noisespec: &TermCollectionSpec,
        expect_msg: &str,
    ) -> SpatialLengthScaleOptimizationResult<f64> {
        let kappa_options = SpatialLengthScaleOptimizationOptions {
            max_outer_iter: 1,
            rel_tol: 1e-6,
            pilot_subsample_threshold: 0,
            ..SpatialLengthScaleOptimizationOptions::default()
        };
        let joint_setup = two_block_exact_joint_hyper_setup(meanspec, noisespec, &kappa_options);
        let theta_dim = joint_setup.theta0().len();

        let mean_terms = spatial_length_scale_term_indices(meanspec);
        let noise_terms = spatial_length_scale_term_indices(noisespec);
        let policy = crate::families::custom_family::OuterDerivativePolicy {
            capability: crate::families::custom_family::ExactOuterDerivativeOrder::Second,
            predicted_hessian_work: 0,
            predicted_gradient_work: 0,
            // Test-style construction with zero predicted work — these
            // paths never engage staged-κ, so the capability bit is
            // moot. Keep `false` as the safe default.
            subsample_capable: false,
        };
        optimize_spatial_length_scale_exact_joint(
            data,
            &[meanspec.clone(), noisespec.clone()],
            &[mean_terms, noise_terms],
            &kappa_options,
            &joint_setup,
            crate::seeding::SeedRiskProfile::Gaussian,
            true,
            true,
            false,
            None,
            policy,
            |theta, specs, designs| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                Ok(designs[0].design.ncols() as f64
                    + designs[1].design.ncols() as f64
                    + designs[0].penalties.len() as f64
                    + designs[1].penalties.len() as f64)
            },
            |theta, specs, designs, eval_mode, _| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                assert!(!designs.is_empty());
                Ok((
                    0.0,
                    Array1::zeros(theta_dim),
                    if matches!(
                        eval_mode,
                        crate::solver::estimate::reml::unified::EvalMode::ValueGradientHessian
                    ) {
                        crate::solver::outer_strategy::HessianResult::Analytic(Array2::zeros((
                            theta_dim, theta_dim,
                        )))
                    } else {
                        crate::solver::outer_strategy::HessianResult::Unavailable
                    },
                ))
            },
            |theta, specs, designs| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                assert!(!designs.is_empty());
                Ok(crate::solver::outer_strategy::EfsEval {
                    cost: 0.0,
                    steps: vec![0.0; theta_dim],
                    beta: None,
                    psi_gradient: None,
                    psi_indices: None,
                    inner_hessian_scale: None,
                })
            },
            |_beta: &Array1<f64>| Ok(()),
        )
        .expect(expect_msg)
    }

    #[test]
    fn exact_joint_two_block_spatial_length_scale_freezes_matern_centers() {
        let n = 40usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (i as f64 * 0.21).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
        }

        // ANISOTROPIC Matérn (`aniso_log_scales = Some`): the joint κ/η outer
        // optimizer only engages for anisotropic spatial terms (#519 —
        // isotropic Matérn anchors its data-seeded κ and learns smoothness
        // through ρ alone, so it contributes no κ axis). This test exercises
        // the joint-optimizer center-freezing path, so it must carry per-axis
        // anisotropy scales to produce the κ/η hyper axes it is asserting on.
        let matern_term = |name: &str, length_scale: f64| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::Matern {
                feature_cols: vec![0, 1],
                spec: MaternBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                    length_scale,
                    nu: MaternNu::FiveHalves,
                    include_intercept: false,
                    double_penalty: true,
                    identifiability: MaternIdentifiability::CenterSumToZero,
                    aniso_log_scales: Some(vec![0.0, 0.0]),
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let meanspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("mean_matern", 0.8)],
        };
        let noisespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("noise_matern", 1.1)],
        };

        let solved = run_two_block_exact_joint_optimize(
            data.view(),
            &meanspec,
            &noisespec,
            "exact joint two-block κ optimization should succeed",
        );

        for resolved in [&solved.resolved_specs[0], &solved.resolved_specs[1]] {
            match &resolved.smooth_terms[0].basis {
                SmoothBasisSpec::Matern { spec, .. } => {
                    assert!(matches!(
                        spec.center_strategy,
                        CenterStrategy::UserProvided(_)
                    ));
                    assert!(matches!(
                        spec.identifiability,
                        MaternIdentifiability::FrozenTransform { .. }
                    ));
                }
                _ => panic!("expected Matérn term"),
            }
        }
    }

    #[test]
    fn exact_joint_spatial_outer_hessian_available_for_dense_non_gaussian_designs() {
        let n = 5_001usize;
        let p = 100usize;
        let data = Array2::from_shape_fn((n, p), |(i, j)| ((i + j + 1) as f64).sin());
        let spec = TermCollectionSpec {
            linear_terms: (0..p)
                .map(|j| LinearTermSpec {
                    name: format!("x{j}"),
                    feature_col: j,
                    feature_cols: vec![j],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                })
                .collect(),
            random_effect_terms: vec![],
            smooth_terms: vec![],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");

        // Both Gaussian and BinomialLogit on a dense design now report
        // analytic Hessian available; the unified evaluator routes
        // non-Gaussian dense-lazy designs through `build_outer_hessian_operator`
        // at biobank scale and through `compute_outer_hessian` otherwise.
        assert!(exact_joint_spatial_outer_hessian_available(
            &LikelihoodSpec::binomial_logit(),
            &design,
        ));
        assert!(exact_joint_spatial_outer_hessian_available(
            &LikelihoodSpec::gaussian_identity(),
            &design,
        ));
    }

    #[test]
    fn spatial_aniso_joint_exact_hessian_materializes_small_case() {
        let n = 18usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (0.41 * i as f64).sin();
        }
        let y = Array1::from_iter((0..n).map(|i| {
            let t = i as f64 / (n as f64 - 1.0);
            0.4 + (2.0 * std::f64::consts::PI * t).sin()
        }));
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern_aniso".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 5 },
                        length_scale: 0.85,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: Some(vec![0.2, -0.2]),
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 120,
            tol: 1e-10,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        assert_eq!(dims_per_term, vec![2]);
        let rho_dim = design.penalties.len();
        let log_kappa0 = SpatialLogKappaCoords::from_length_scales_aniso(
            &frozen,
            &spatial_terms,
            &SpatialLengthScaleOptimizationOptions::default(),
        );
        let mut theta = Array1::<f64>::zeros(rho_dim + log_kappa0.as_array().len());
        for j in 0..rho_dim {
            theta[j] = -0.15 + 0.07 * j as f64;
        }
        theta.slice_mut(s![rho_dim..]).assign(log_kappa0.as_array());

        let external_opts =
            external_opts_for_design(&LikelihoodSpec::gaussian_identity(), &design, &fit_opts);
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen,
            design.clone(),
            spatial_terms,
            rho_dim,
            dims_per_term,
        )
        .expect("single-block cache");
        let mut evaluator = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &design.design,
            offset.view(),
            &design.penalties,
            &external_opts,
            "small aniso Hessian finite-difference evaluator",
        )
        .expect("evaluator");

        let eval_at = |theta: &Array1<f64>,
                       cache: &mut SingleBlockExactJointDesignCache<'_>,
                       evaluator: &mut crate::estimate::ExternalJointHyperEvaluator<'_>,
                       order: crate::solver::outer_strategy::OuterEvalOrder| {
            cache.ensure_theta(theta).expect("theta applied");
            let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
                data.view(),
                cache.spec(),
                cache.design(),
                &cache.spatial_terms,
            )
            .expect("hyper dirs build")
            .expect("hyper dirs present");
            evaluate_joint_reml_outer_eval_at_theta(
                evaluator,
                cache.design(),
                theta,
                rho_dim,
                hyper_dirs,
                None,
                order,
                None,
            )
            .expect("outer eval")
        };

        let (_, gradient, hessian_result) = eval_at(
            &theta,
            &mut cache,
            &mut evaluator,
            crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian,
        );
        let hessian = hessian_result
            .materialize_dense()
            .expect("hessian materializes")
            .expect("hessian present");
        assert_eq!(hessian.nrows(), theta.len());
        assert_eq!(hessian.ncols(), theta.len());
        assert!(hessian.iter().all(|value| value.is_finite()));
        assert!(gradient.iter().all(|value| value.is_finite()));

        let symmetry_diff = max_abs_diff_matrix(&hessian, &hessian.t().to_owned());
        assert!(
            symmetry_diff <= 1e-10,
            "small aniso exact Hessian should be symmetric, max diff={symmetry_diff}"
        );
        let psi_block = hessian.slice(s![rho_dim.., rho_dim..]).to_owned();
        assert!(
            psi_block.iter().any(|value| value.abs() > 1e-10),
            "small aniso exact Hessian should carry non-zero ψ curvature"
        );
    }

    /// Finite-difference verification of the joint REML gradient on a Duchon
    /// BinomialProbit configuration that reproduces the iso-kappa joint REML
    /// `|g|≈3.5e7` blow-up. The 1D Duchon term uses `length_scale=Some(1.0)`
    /// and a moderate number of centers so the active penalty count produces a
    /// non-trivial ρ block alongside the single log-κ axis. The analytic
    /// gradient assembled via `evaluate_joint_reml_outer_eval_at_theta` is
    /// compared component-wise against a centered finite difference of the
    /// cost via `evaluator.evaluate_cost_only` (the same cost path the joint
    /// outer optimizer uses). Disagreement on the ext-coord component
    /// isolates a wrong derivative in the log-κ gradient path.
    #[test]
    fn iso_kappa_duchon_binomial_probit_joint_gradient_matches_finite_difference() {
        let (pass, worst, violations) = iso_kappa_fd_variant_driver(
            "duchon_probit_n80",
            80,
            LikelihoodSpec::binomial_probit(),
            false,
        );
        assert!(
            pass,
            "Duchon BinomialProbit n=80 FD failed; worst_psi_rel={worst:.3e}\n  {}",
            violations.join("\n  ")
        );
    }

    /// Shared driver for iso-κ joint REML gradient FD variants. Returns the
    /// worst psi rel_err across the four theta probes (zero / psi_only / base /
    /// alt) and panics with full violations only if `assert_pass` is true.
    /// Knobs let one-at-a-time variants of the original BinomialProbit Duchon
    /// failure isolate which dimension triggers the analytic-vs-FD blow-up.
    fn iso_kappa_fd_variant_driver(
        label: &str,
        n: usize,
        family: LikelihoodSpec,
        skip_psi: bool,
    ) -> (bool, f64, Vec<String>) {
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            let eta = 1.4 * (2.0 * std::f64::consts::PI * t).sin() + 0.5 * (t - 0.5);
            let raw = eta + 0.7 * (3.7 * (i as f64) + 1.0).sin();
            y[i] = if family.is_gaussian_identity() {
                raw
            } else if raw > 0.0 {
                1.0
            } else {
                0.0
            };
        }
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);
        // Every active caller is a Duchon-iso-κ FD probe. Thin-plate is
        // deliberately excluded from the spatial κ-axis enrollment (see
        // `spatial_term_supports_hyper_optimization`), so the driver only
        // needs to build the Duchon basis.
        let basis = SmoothBasisSpec::Duchon {
            feature_cols: vec![0],
            spec: DuchonBasisSpec {
                periodic: None,
                center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                length_scale: Some(1.0),
                power: 1.0,
                nullspace_order: DuchonNullspaceOrder::Linear,
                identifiability: SpatialIdentifiability::default(),
                aniso_log_scales: None,
                operator_penalties: DuchonOperatorPenaltySpec::default(),
                boundary: OneDimensionalBoundary::Open,
            },
            input_scales: None,
        };
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "variant_1d".to_string(),
                basis,
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 200,
            tol: 1e-12,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze");
        let frozen_design =
            build_term_collection_design(data.view(), &frozen).expect("frozen design");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        assert_eq!(dims_per_term, vec![1], "{label}: expect one log-κ axis");
        let rho_dim = frozen_design.penalties.len();
        let psi_dim: usize = dims_per_term.iter().sum();
        assert!(psi_dim >= 1);

        let external_opts = external_opts_for_design(&family, &frozen_design, &fit_opts);
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
            spatial_terms.clone(),
            rho_dim,
            dims_per_term.clone(),
        )
        .expect("single-block cache");
        let mut evaluator = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &frozen_design.design,
            offset.view(),
            &frozen_design.penalties,
            &external_opts,
            "iso-κ variant FD evaluator",
        )
        .expect("evaluator");

        let cost_at = |theta: &Array1<f64>,
                       cache: &mut SingleBlockExactJointDesignCache<'_>,
                       evaluator: &mut crate::estimate::ExternalJointHyperEvaluator<'_>|
         -> f64 {
            cache.ensure_theta(theta).expect("ensure_theta");
            let design = cache.design();
            evaluator
                .evaluate_cost_only(
                    &design.design,
                    &design.penalties,
                    &design.nullspace_dims,
                    design.linear_constraints.clone(),
                    theta,
                    rho_dim,
                    None,
                    "iso-κ variant FD cost-only",
                    None,
                )
                .expect("cost-only eval")
        };

        let analytic_at = |theta: &Array1<f64>,
                           cache: &mut SingleBlockExactJointDesignCache<'_>,
                           evaluator: &mut crate::estimate::ExternalJointHyperEvaluator<'_>|
         -> (f64, Array1<f64>) {
            cache.ensure_theta(theta).expect("ensure_theta");
            let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
                data.view(),
                cache.spec(),
                cache.design(),
                &cache.spatial_terms,
            )
            .expect("hyper dirs build")
            .expect("hyper dirs present");
            let (cost, grad, _hess) = evaluate_joint_reml_outer_eval_at_theta(
                evaluator,
                cache.design(),
                theta,
                rho_dim,
                hyper_dirs,
                None,
                crate::solver::outer_strategy::OuterEvalOrder::ValueAndGradient,
                None,
            )
            .expect("outer eval");
            (cost, grad)
        };

        let theta_dim = rho_dim + psi_dim;
        let theta_zero = Array1::<f64>::zeros(theta_dim);
        let mut theta_base = Array1::<f64>::zeros(theta_dim);
        for j in 0..rho_dim {
            theta_base[j] = 0.2 - 0.1 * j as f64;
        }
        let mut theta_psi_only = Array1::<f64>::zeros(theta_dim);
        for k in 0..psi_dim {
            theta_psi_only[rho_dim + k] = 0.4;
        }
        let mut theta_alt = theta_base.clone();
        for j in 0..rho_dim {
            theta_alt[j] = 1.0 + 0.05 * j as f64;
        }
        for k in 0..psi_dim {
            theta_alt[rho_dim + k] = 0.4;
        }

        let h = 1e-5_f64;
        let rel_tol = 5e-3_f64;
        let mut violations: Vec<String> = Vec::new();
        let mut worst_psi_rel = 0.0_f64;
        for (probe, theta) in [
            ("zero", &theta_zero),
            ("psi_only", &theta_psi_only),
            ("base", &theta_base),
            ("alt", &theta_alt),
        ] {
            let (cost_an, grad_an) = analytic_at(theta, &mut cache, &mut evaluator);
            assert!(cost_an.is_finite(), "{label} {probe}: cost not finite");
            for j in 0..theta_dim {
                let is_psi = j >= rho_dim;
                if skip_psi && is_psi {
                    continue;
                }
                let mut plus = theta.clone();
                plus[j] += h;
                let mut minus = theta.clone();
                minus[j] -= h;
                let cp = cost_at(&plus, &mut cache, &mut evaluator);
                let cm = cost_at(&minus, &mut cache, &mut evaluator);
                let fd = (cp - cm) / (2.0 * h);
                let denom = fd.abs().max(grad_an[j].abs()).max(1e-3);
                let rel = (grad_an[j] - fd).abs() / denom;
                let kind = if is_psi { "psi" } else { "rho" };
                eprintln!(
                    "[{label} {probe}] {kind} j={j} an={:+.4e} fd={:+.4e} rel={:.3e}",
                    grad_an[j], fd, rel
                );
                if is_psi && rel > worst_psi_rel {
                    worst_psi_rel = rel;
                }
                if rel >= rel_tol {
                    violations.push(format!(
                        "{probe} {kind} j={j}: analytic={:+.6e} fd={:+.6e} rel={:.3e}",
                        grad_an[j], fd, rel
                    ));
                }
            }
        }
        let pass = violations.is_empty();
        eprintln!(
            "[{label} SUMMARY] pass={pass} worst_psi_rel={worst_psi_rel:.3e} \
             violations={}",
            violations.len()
        );
        (pass, worst_psi_rel, violations)
    }

    #[test]
    fn iso_kappa_duchon_gaussian_identity_fd() {
        let (pass, worst, violations) = iso_kappa_fd_variant_driver(
            "duchon_gaussian",
            80,
            LikelihoodSpec::gaussian_identity(),
            false,
        );
        assert!(
            pass,
            "Gaussian Identity FD failed; worst_psi_rel={worst:.3e}\n  {}",
            violations.join("\n  ")
        );
    }

    /// Parity test for the unified exact-spatial joint optimizer (issue #427).
    ///
    /// Before unification, anisotropic and isotropic spatial joint optimization
    /// were two near-identical functions that differed only in diagnostic
    /// labels. The shared engine `run_exact_joint_spatial_optimization` now
    /// drives both, selected by `SpatialHyperKind`. For a 1-D spatial term the
    /// two coordinate kinds are mathematically identical — `dims_per_term ==
    /// [1]`, so each carries exactly one log-scale coordinate per term and both
    /// route the same θ through `try_build_spatial_log_kappa_hyper_dirs`. The
    /// converged hyperparameters and certified REML cost must therefore agree to
    /// numerical round-off when the engine is invoked under either kind with
    /// identical inputs. Any divergence would mean the kind discriminator leaked
    /// into the numerics rather than staying confined to labels.
    #[test]
    fn exact_spatial_joint_engine_aniso_iso_parity_1d() {
        let n = 80usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            let eta = 1.4 * (2.0 * std::f64::consts::PI * t).sin() + 0.5 * (t - 0.5);
            y[i] = eta + 0.7 * (3.7 * (i as f64) + 1.0).sin();
        }
        let weights = Array1::<f64>::ones(n);
        let offset = Array1::<f64>::zeros(n);
        let family = LikelihoodSpec::gaussian_identity();

        // 1-D Duchon term: a single log-scale axis (dims_per_term == [1]), the
        // shared geometry across both coordinate kinds.
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "parity_1d".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                        length_scale: Some(1.0),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::all_active(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 200,
            tol: 1e-12,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze");
        let frozen_design =
            build_term_collection_design(data.view(), &frozen).expect("frozen design");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        assert_eq!(spatial_terms.len(), 1, "expect a single spatial term");
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        assert_eq!(dims_per_term, vec![1], "expect one log-scale axis");
        let rho_dim = frozen_design.penalties.len();
        assert!(rho_dim >= 1, "expect at least one penalty block");

        // Construct the joint setup exactly as the production caller does,
        // shared verbatim between the two engine invocations so that any
        // difference in the result can only come from the coordinate kind.
        const JOINT_RHO_BOUND: f64 = 12.0;
        let kappa_options = SpatialLengthScaleOptimizationOptions::default();
        let log_kappa0 =
            SpatialLogKappaCoords::from_length_scales(&frozen, &spatial_terms, &kappa_options);
        let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_from_data(
            data.view(),
            &frozen,
            &spatial_terms,
            &kappa_options,
        );
        let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_from_data(
            data.view(),
            &frozen,
            &spatial_terms,
            &kappa_options,
        );
        let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
        let setup = ExactJointHyperSetup::new(
            Array1::<f64>::zeros(rho_dim), // log λ seed (λ = 1)
            Array1::<f64>::from_elem(rho_dim, -JOINT_RHO_BOUND),
            Array1::<f64>::from_elem(rho_dim, JOINT_RHO_BOUND),
            log_kappa0,
            log_kappa_lower,
            log_kappa_upper,
        );
        let theta0 = setup.theta0();
        let lower = setup.lower();
        let upper = setup.upper();

        let run = |kind: SpatialHyperKind| -> (Array1<f64>, f64) {
            run_exact_joint_spatial_optimization(
                kind,
                data.view(),
                y.view(),
                weights.view(),
                offset.view(),
                &frozen,
                &frozen_design,
                family.clone(),
                &fit_opts,
                &spatial_terms,
                &dims_per_term,
                &theta0,
                &lower,
                &upper,
                rho_dim,
                &kappa_options,
            )
            .expect("exact joint spatial optimization")
        };

        let (theta_aniso, value_aniso) = run(SpatialHyperKind::Anisotropic);
        let (theta_iso, value_iso) = run(SpatialHyperKind::Isotropic);

        assert_eq!(
            theta_aniso.len(),
            theta_iso.len(),
            "converged θ dimension must match across coordinate kinds"
        );
        // In 1-D the two kinds are numerically identical: the only difference is
        // diagnostic labels, so converged hyperparameters and the certified REML
        // cost must agree to round-off. No tolerance weakening — this is an
        // equality, not an approximation.
        for j in 0..theta_aniso.len() {
            let diff = (theta_aniso[j] - theta_iso[j]).abs();
            assert!(
                diff <= 1e-9 * (1.0 + theta_aniso[j].abs()),
                "θ[{j}] differs across kinds: aniso={:+.12e} iso={:+.12e} diff={:.3e}",
                theta_aniso[j],
                theta_iso[j],
                diff,
            );
        }
        let value_diff = (value_aniso - value_iso).abs();
        assert!(
            value_diff <= 1e-9 * (1.0 + value_aniso.abs()),
            "final REML value differs across kinds: aniso={:+.12e} iso={:+.12e} diff={:.3e}",
            value_aniso,
            value_iso,
            value_diff,
        );
        assert!(
            value_aniso.is_finite() && value_iso.is_finite(),
            "both kinds must produce a finite certified REML cost"
        );
    }

    #[test]
    fn iso_kappa_duchon_binomial_logit_fd() {
        let (pass, worst, violations) = iso_kappa_fd_variant_driver(
            "duchon_logit",
            80,
            LikelihoodSpec::binomial_logit(),
            false,
        );
        assert!(
            pass,
            "BinomialLogit FD failed; worst_psi_rel={worst:.3e}\n  {}",
            violations.join("\n  ")
        );
    }

    // No `iso_kappa_thinplate_*_fd` companion to the Duchon FD tests above:
    // thin-plate is deliberately excluded from the spatial κ-axis enrollment
    // by `spatial_term_supports_hyper_optimization` (a scalar TPS κ creates
    // the flat ρ/κ valleys tracked in #718 / #721 / #731 / #732), so there
    // is no analytic κ-gradient on which an FD comparison could land.

    #[test]
    fn iso_kappa_duchon_n_smaller_fd() {
        let (pass, worst, violations) = iso_kappa_fd_variant_driver(
            "duchon_probit_n20",
            20,
            LikelihoodSpec::binomial_probit(),
            false,
        );
        assert!(
            pass,
            "Duchon Probit n=20 FD failed; worst_psi_rel={worst:.3e}\n  {}",
            violations.join("\n  ")
        );
    }

    #[test]
    fn iso_kappa_duchon_no_psi_fd() {
        let (pass, _worst, violations) = iso_kappa_fd_variant_driver(
            "duchon_probit_rho_only",
            80,
            LikelihoodSpec::binomial_probit(),
            true,
        );
        assert!(
            pass,
            "Duchon Probit ρ-only FD failed:\n  {}",
            violations.join("\n  ")
        );
    }

    /// Owned 1-D Duchon BinomialProbit setup shared verbatim across the
    /// `duchon_probit_*` mechanism pins. Holds only non-self-referential
    /// owners; each test constructs its own `external_opts` / cache /
    /// evaluator inline (the borrow-entangled, per-test-labelled parts).
    struct DuchonProbitSetup {
        data: Array2<f64>,
        y: Array1<f64>,
        weights: Array1<f64>,
        offset: Array1<f64>,
        frozen: TermCollectionSpec,
        frozen_design: TermCollectionDesign,
        spatial_terms: Vec<usize>,
        dims_per_term: Vec<usize>,
        rho_dim: usize,
        psi_dim: usize,
    }

    /// Builds the verbatim 1-D Duchon BinomialProbit data + frozen design used
    /// by the ψ-trace / per-row / PIRLS-determinism mechanism pins.
    fn build_duchon_probit_setup() -> DuchonProbitSetup {
        let n = 80usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            let eta = 1.4 * (2.0 * std::f64::consts::PI * t).sin() + 0.5 * (t - 0.5);
            y[i] = if eta + 0.7 * (3.7 * (i as f64) + 1.0).sin() > 0.0 {
                1.0
            } else {
                0.0
            };
        }
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_1d".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                        length_scale: Some(1.0),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::all_active(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze");
        let frozen_design =
            build_term_collection_design(data.view(), &frozen).expect("frozen design");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        let rho_dim = frozen_design.penalties.len();
        let psi_dim: usize = dims_per_term.iter().sum();
        DuchonProbitSetup {
            data,
            y,
            weights,
            offset,
            frozen,
            frozen_design,
            spatial_terms,
            dims_per_term,
            rho_dim,
            psi_dim,
        }
    }

    /// Architectural pin: for the iso-κ Duchon ψ-axis under BinomialProbit,
    /// the penalty-subspace projection (`solution.penalty_subspace_trace =
    /// Some(...)`) must be *active* AND must change the trace by an order of
    /// magnitude relative to the full-space `G_ε(H)` kernel.
    ///
    /// Specifically:
    ///   * `unprojected_tr ≈ tr(G_ε(H) · op_total)`
    ///     ≫ `production_tr ≈ tr_projected(...)` in magnitude.
    ///   * `production_tr ≈ ∂log|H|/∂ψ` (FD truth).
    ///
    /// If either invariant breaks, the cost identity for Duchon ψ silently
    /// drifts and the joint-gradient test (and any downstream optimizer) is
    /// solving the wrong problem. This test exists to make the projection's
    /// role legible: it captures the mechanism, not just the outcome.
    #[test]
    fn iso_kappa_duchon_penalty_subspace_projection_pins_trace() {
        let DuchonProbitSetup {
            data,
            y,
            weights,
            offset,
            frozen,
            frozen_design,
            spatial_terms,
            dims_per_term,
            rho_dim,
            psi_dim,
        } = build_duchon_probit_setup();
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 200,
            tol: 1e-12,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let external_opts = external_opts_for_design(
            &LikelihoodSpec::binomial_probit(),
            &frozen_design,
            &fit_opts,
        );
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
            spatial_terms.clone(),
            rho_dim,
            dims_per_term.clone(),
        )
        .expect("cache");
        let mut evaluator = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &frozen_design.design,
            offset.view(),
            &frozen_design.penalties,
            &external_opts,
            "penalty-subspace projection pin",
        )
        .expect("evaluator");

        let theta_dim = rho_dim + psi_dim;
        let theta_zero = Array1::<f64>::zeros(theta_dim);

        cache.ensure_theta(&theta_zero).expect("ensure_theta");
        let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
            data.view(),
            cache.spec(),
            cache.design(),
            &cache.spatial_terms,
        )
        .expect("hyper dirs build")
        .expect("hyper dirs present");
        let (cost_at_zero, grad_at_zero, _hess) = evaluate_joint_reml_outer_eval_at_theta(
            &mut evaluator,
            cache.design(),
            &theta_zero,
            rho_dim,
            hyper_dirs,
            None,
            crate::solver::outer_strategy::OuterEvalOrder::ValueAndGradient,
            None,
        )
        .expect("analytic outer eval");
        let stash = crate::test_support::debug_stash::take_terms();

        let unprojected_tr = stash
            .unprojected_tr
            .expect("unprojected_tr must be captured by the EIG-DECOMP block at small p");
        let production_tr = stash
            .production_tr
            .expect("production_tr must be captured by the EIG-DECOMP block at small p");
        let projection_active = stash
            .projection_active
            .expect("projection_active must be captured");

        // INVARIANT 1: The projection must be active for the Duchon ψ axis
        // at non-canonical links. If this regresses (someone disables
        // penalty_subspace_trace for the ψ path), the Duchon ψ-gradient
        // silently breaks at the order-of-magnitude level.
        assert!(
            projection_active,
            "Duchon ψ axis must route the trace through penalty_subspace_trace"
        );

        // INVARIANT 2: Production trace agrees with FD ∂log|H_proj|/∂ψ —
        // i.e. the centered derivative of the **projected** Hessian
        // log-determinant `log|U_Sᵀ H U_S|_+`, evaluated through the same
        // assembly path the outer cost identity uses.
        //
        // Why the projected logdet, not the full one: under the
        // rank-deficient LAML fix the REML/LAML cost reads
        // `log|H_proj|`, not `log|H_full|`.  These two scalars differ by
        // the `null(S)` directions' contribution to `log|X'WX|` — visible
        // here as `hop_logdet ≠ log_det_h_proj` whenever
        // `range(S_+)` is a strict subspace (for this 1-D Duchon test
        // `S_λ` has rank 5 in a 7-dim transformed basis, so two
        // directions of H sit outside `range(S)` and contribute their
        // own ψ-dependence to `log|H_full|`).  Finite-differencing the
        // full-space logdet picks up that extra contribution; finite-
        // differencing the projected logdet exposes exactly what the
        // analytic trace formula `kernel.trace_projected_logdet(op_total)`
        // computes.  Pinning against the wrong finite-difference reference
        // made this invariant fire even on mathematically-correct
        // production output, so the test pinned an impossible identity.
        // The corrected finite-difference reference is the only one the
        // production code actually claims to satisfy.
        let h = 1e-5_f64;
        let psi_idx = rho_dim;
        let mut theta_p = theta_zero.clone();
        theta_p[psi_idx] += h;
        let mut theta_m = theta_zero.clone();
        theta_m[psi_idx] -= h;
        cache.ensure_theta(&theta_p).expect("ensure +h");
        let dp = cache.design().clone();
        let log_det_p = evaluator
            .debug_logdet_h_proj(
                &dp.design,
                &dp.penalties,
                &dp.nullspace_dims,
                dp.linear_constraints.clone(),
                &theta_p,
                rho_dim,
                "subspace-pin +h",
            )
            .expect("debug_logdet_h_proj +h");
        cache.ensure_theta(&theta_m).expect("ensure -h");
        let dm = cache.design().clone();
        let log_det_m = evaluator
            .debug_logdet_h_proj(
                &dm.design,
                &dm.penalties,
                &dm.nullspace_dims,
                dm.linear_constraints.clone(),
                &theta_m,
                rho_dim,
                "subspace-pin -h",
            )
            .expect("debug_logdet_h_proj -h");
        let fd_d_logdet_h_proj = (log_det_p - log_det_m) / (2.0 * h);
        let rel_to_fd =
            (production_tr - fd_d_logdet_h_proj).abs() / fd_d_logdet_h_proj.abs().max(1e-30);
        assert!(
            rel_to_fd < 1e-2,
            "production trace must match FD ∂log|H_proj|/∂ψ: production_tr={:+.4e} \
             fd_d_logdet_h_proj={:+.4e} rel={:+.3e}",
            production_tr,
            fd_d_logdet_h_proj,
            rel_to_fd
        );

        // INVARIANT 3: The penalty-subspace projection must change the trace
        // substantially — otherwise it's a no-op. For this Duchon ψ test
        // unprojected ≈ +22.5 while production ≈ −1.05.  A gap collapse to
        // near zero indicates the projection is broken or disabled.
        let gap = (unprojected_tr - production_tr).abs();
        assert!(
            gap > 5.0,
            "penalty-subspace projection should eliminate a substantial \
             null-space contribution (unprojected={:+.4e}, production={:+.4e}, \
             gap={:+.4e})",
            unprojected_tr,
            production_tr,
            gap
        );

        // Sanity: cost & gradient finite at θ=0.
        assert!(
            cost_at_zero.is_finite() && grad_at_zero.iter().all(|v| v.is_finite()),
            "ψ-gradient and cost must be finite at θ=0"
        );
    }

    /// FD `c · dη/dψ_total` vs analytic `c·X_τβ̂ - c·X·v_ψ` per-row.
    /// If they match: the M = term4 + correction matrix has the right diagonal,
    /// and the trace bug must lie elsewhere (e.g., basis or weights).
    /// If they don't match: the IFT formula or v_ψ is wrong.
    #[test]
    fn duchon_probit_per_row_dnu_dpsi_fd_vs_analytic() {
        let DuchonProbitSetup {
            data,
            y,
            weights,
            offset,
            frozen,
            frozen_design,
            spatial_terms,
            dims_per_term,
            rho_dim,
            psi_dim,
        } = build_duchon_probit_setup();
        let n = data.nrows();
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 200,
            tol: 1e-12,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let external_opts = external_opts_for_design(
            &LikelihoodSpec::binomial_probit(),
            &frozen_design,
            &fit_opts,
        );
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
            spatial_terms.clone(),
            rho_dim,
            dims_per_term.clone(),
        )
        .expect("cache");
        let mut evaluator = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &frozen_design.design,
            offset.view(),
            &frozen_design.penalties,
            &external_opts,
            "per-row dη/dψ FD vs analytic",
        )
        .expect("evaluator");

        let theta_dim = rho_dim + psi_dim;
        let theta_zero = Array1::<f64>::zeros(theta_dim);
        let psi_idx = rho_dim;
        let h = 1e-5_f64;

        // ── η, W, c at ψ=0
        cache.ensure_theta(&theta_zero).expect("ensure zero");
        let dz = cache.design().clone();
        let (eta_0, w_0, c_0) = evaluator
            .debug_full_eta_w_c(
                &dz.design,
                &dz.penalties,
                &dz.nullspace_dims,
                dz.linear_constraints.clone(),
                &theta_zero,
                rho_dim,
                "ψ=0",
            )
            .expect("debug at 0");

        // η at ψ=±h
        let mut theta_p = theta_zero.clone();
        theta_p[psi_idx] += h;
        cache.ensure_theta(&theta_p).expect("ensure +h");
        let dp = cache.design().clone();
        let (eta_p, w_p, c_p) = evaluator
            .debug_full_eta_w_c(
                &dp.design,
                &dp.penalties,
                &dp.nullspace_dims,
                dp.linear_constraints.clone(),
                &theta_p,
                rho_dim,
                "ψ=+h",
            )
            .expect("debug at +h");

        let mut theta_m = theta_zero.clone();
        theta_m[psi_idx] -= h;
        cache.ensure_theta(&theta_m).expect("ensure -h");
        let dm = cache.design().clone();
        let (eta_m, _w_m, _c_m) = evaluator
            .debug_full_eta_w_c(
                &dm.design,
                &dm.penalties,
                &dm.nullspace_dims,
                dm.linear_constraints.clone(),
                &theta_m,
                rho_dim,
                "ψ=-h",
            )
            .expect("debug at -h");

        // dη/dψ_FD per row
        let mut dnu_dpsi_fd = Array1::<f64>::zeros(n);
        for i in 0..n {
            dnu_dpsi_fd[i] = (eta_p[i] - eta_m[i]) / (2.0 * h);
        }
        let l2_dnu = (dnu_dpsi_fd.iter().map(|v| v * v).sum::<f64>()).sqrt();
        eprintln!(
            "[ROW] ||dη/dψ_FD||_L2 = {:+.6e}  per-row max|·| = {:+.4e}",
            l2_dnu,
            dnu_dpsi_fd.iter().map(|v| v.abs()).fold(0.0_f64, f64::max)
        );

        // c · dη/dψ_FD per row
        let c_dnu_fd: Array1<f64> = &c_0 * &dnu_dpsi_fd;
        let l2_c_dnu = (c_dnu_fd.iter().map(|v| v * v).sum::<f64>()).sqrt();
        eprintln!(
            "[ROW] ||c·dη/dψ_FD||_L2 = {:+.6e}  per-row max|·| = {:+.4e}",
            l2_c_dnu,
            c_dnu_fd.iter().map(|v| v.abs()).fold(0.0_f64, f64::max)
        );

        // ── Get ANALYTIC per-row diagonals via debug_stash.
        cache.ensure_theta(&theta_zero).expect("ensure_theta zero");
        let hyper_dirs = try_build_spatial_log_kappa_hyper_dirs(
            data.view(),
            cache.spec(),
            cache.design(),
            &cache.spatial_terms,
        )
        .expect("hyper dirs build")
        .expect("hyper dirs present");
        let (analytic_cost, analytic_gradient, _) = evaluate_joint_reml_outer_eval_at_theta(
            &mut evaluator,
            cache.design(),
            &theta_zero,
            rho_dim,
            hyper_dirs,
            None,
            crate::solver::outer_strategy::OuterEvalOrder::ValueAndGradient,
            None,
        )
        .expect("analytic outer eval");
        assert!(analytic_cost.is_finite());
        assert!(analytic_gradient.iter().all(|value| value.is_finite()));
        let stash = crate::test_support::debug_stash::take_terms();
        let c_x_tau_beta = stash.c_x_tau_beta_diag.clone().expect("term4 diag stashed");
        let x_v_psi = stash.c_x_v_psi_diag.clone().expect("X·v_ψ stashed");
        assert_eq!(c_x_tau_beta.len(), n);
        assert_eq!(x_v_psi.len(), n);

        // Analytic per-row IFT diagonal:
        //   d_i = (c·X_τβ̂)_i − (c·X·v_ψ)_i
        //       = c_i · (X_τβ̂ + X·∂β̂/∂ψ)_i = c_i · (dη/dψ_total)_i.
        // FD reference: c_i · (η_+ − η_−)_i / 2h.
        let c_x_v: Array1<f64> = &c_0 * &x_v_psi;
        let analytic_diag: Array1<f64> = &c_x_tau_beta - &c_x_v;

        // Each half is individually large (Frobenius ~ 88 for our test); the
        // physics is in their row-wise cancellation. Pin the cancellation:
        // |analytic - FD| / |FD|  must stay below 1e-3.  Our current
        // measurement is rel ≈ 2.9e-5 with both halves at ||·||_L2 ≈ 88.5.
        let l2_term4_only = c_x_tau_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
        let l2_corr_only = c_x_v.iter().map(|v| v * v).sum::<f64>().sqrt();
        assert!(
            l2_term4_only > 10.0 && l2_corr_only > 10.0,
            "expected both IFT halves to be large in isolation (only their \
             row-wise sum should be small), got ||c·X_τβ̂||_L2={:+.3e}, \
             ||c·X·v_ψ||_L2={:+.3e}",
            l2_term4_only,
            l2_corr_only
        );

        let diff: Array1<f64> = &analytic_diag - &c_dnu_fd;
        let l2_diff = diff.iter().map(|v| v * v).sum::<f64>().sqrt();
        let rel = l2_diff / l2_c_dnu.max(1e-30);
        assert!(
            rel < 1e-3,
            "analytic IFT row-diagonal must match FD to 1e-3 relative \
             (got rel={:+.4e}, ||Δ||_L2={:+.4e}, ||FD||_L2={:+.4e}). \
             Regression in IFT chain rule (v_ψ formula or W-surface mix).",
            rel,
            l2_diff,
            l2_c_dnu
        );

        assert_eq!(eta_0.len(), n);
        assert_eq!(w_0.len(), n);
        assert_eq!(c_p.len(), n);
        assert_eq!(w_p.len(), n);
        assert!(eta_0.iter().all(|value| value.is_finite()));
        assert!(w_0.iter().all(|value| value.is_finite()));
        assert!(c_p.iter().all(|value| value.is_finite()));
        assert!(w_p.iter().all(|value| value.is_finite()));
    }

    /// Test PIRLS structural determinism: call debug_full_h three times at
    /// the SAME theta=0 and check the returned H matrices agree to a tight
    /// relative tolerance. We deliberately do NOT require bit-identical
    /// matrices: PIRLS feeds intermediate state through rayon `.reduce`
    /// fold/combine on f64 (deviance, weighted sums, X'WX accumulation),
    /// and floating-point addition is non-associative, so the bit pattern
    /// of those reductions varies with thread scheduling. What we *do*
    /// require is structural agreement — the same fixed point in the same
    /// Qs frame. A non-deterministic Qs reparametrization (e.g. an
    /// eigenvector sign flip) shows up as O(‖H‖) entry-wise drift, two
    /// to ten orders of magnitude above the rayon summation floor; the
    /// 1e-5 relative band catches that while tolerating the latter. The
    /// band is 1e-5 rather than 1e-6 because at θ=0 the BinomialProbit
    /// IRLS weights on near-separable data drive ‖H‖_∞ to ~2e9, and the
    /// measured rayon-reduction floor on cancellation-heavy X'WX sums at
    /// that scale is ~2.5e-6 relative — already above 1e-6. 1e-5 still
    /// sits ~5 orders below an O(1)-relative Qs sign flip.
    #[test]
    fn duchon_probit_pirls_determinism_at_zero() {
        let DuchonProbitSetup {
            data,
            y,
            weights,
            offset,
            frozen,
            frozen_design,
            spatial_terms,
            dims_per_term,
            rho_dim,
            psi_dim,
        } = build_duchon_probit_setup();
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 200,
            // This test certifies that `debug_full_h` is deterministic at a
            // fixed θ; it needs the inner PIRLS to converge so a Hessian is
            // returned. The binomial-probit Duchon fit's stationarity residual
            // floors at ‖g‖≈2e-6 (the LM-ridge noise floor), so a sub-floor
            // request such as 1e-12 can never be certified and surfaces as
            // `PirlsDidNotConverge`. 1e-6 is the standard GLM convergence
            // tolerance and clears the floor with margin (the scale-invariant
            // KKT bound certifies at ‖g‖ < 1e-6·√n·√p ≈ 2.8e-5). 1e-12 only
            // ever "passed" because the near-stationary band silently carried a
            // 1e-6 floor, since removed.
            tol: 1e-6,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let external_opts = external_opts_for_design(
            &LikelihoodSpec::binomial_probit(),
            &frozen_design,
            &fit_opts,
        );
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
            spatial_terms.clone(),
            rho_dim,
            dims_per_term.clone(),
        )
        .expect("cache");
        let mut evaluator = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &frozen_design.design,
            offset.view(),
            &frozen_design.penalties,
            &external_opts,
            "PIRLS-determinism",
        )
        .expect("evaluator");

        let theta_dim = rho_dim + psi_dim;
        let theta_zero = Array1::<f64>::zeros(theta_dim);

        let mut h_calls = Vec::new();
        for trial in 0..3 {
            cache.ensure_theta(&theta_zero).expect("ensure_theta");
            let d = cache.design().clone();
            let h_i = evaluator
                .debug_full_h(
                    &d.design,
                    &d.penalties,
                    &d.nullspace_dims,
                    d.linear_constraints.clone(),
                    &theta_zero,
                    rho_dim,
                    &format!("determinism trial {}", trial),
                )
                .expect("debug_full_h");
            h_calls.push(h_i);
        }

        let h0 = &h_calls[0];
        let norm_h0 = h0.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
        let abs_tol = (1e-5_f64) * norm_h0 + 1e-12_f64;
        for trial in 1..3 {
            let diff = &h_calls[trial] - h0;
            let max_abs = diff.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
            assert!(
                max_abs <= abs_tol,
                "PIRLS non-deterministic at fixed θ (trial {trial} vs 0): \
                 max|Δ|={max_abs:+.6e} > tol={abs_tol:+.6e} \
                 (‖H‖_∞={norm_h0:+.6e})"
            );
        }
    }

    #[test]
    fn spatial_aniso_joint_large_psi_dim_keeps_second_order_route() {
        let cap = crate::solver::outer_strategy::OuterCapability {
            gradient: crate::solver::outer_strategy::Derivative::Analytic,
            hessian: crate::solver::outer_strategy::DeclaredHessianForm::Either,
            n_params: 40,
            psi_dim: 31,
            fixed_point_available: true,
            barrier_config: None,
            prefer_gradient_only: false,
            disable_fixed_point: false,
        };
        let route = crate::solver::outer_strategy::plan(&cap);
        assert_eq!(route.solver, crate::solver::outer_strategy::Solver::Arc);
        assert_eq!(
            route.hessian_source,
            crate::solver::outer_strategy::HessianSource::Analytic
        );
        assert!(route.routing_log_line().contains("matrix-free=false"));
    }

    #[test]
    fn exact_joint_spatial_outer_hessian_available_for_sparse_designs() {
        let n = 96usize;
        let x = Array1::linspace(0.0, 1.0, n);
        let mut data = Array2::<f64>::zeros((n, 1));
        data.column_mut(0).assign(&x);
        // `BSplineIdentifiability::None` keeps the smooth block sparse through
        // `build_bspline_basis_1d`. The default `WeightedSumToZero` policy is
        // deliberately densified by `apply_sum_to_zero_constraint_sparse`
        // (orthonormal Z, so ZZᵀ is a true projector — `B·Z` is mathematically
        // dense), which would prevent the assembled design from landing in
        // `DesignMatrix::Sparse`; the sparse-design pin only holds on the None
        // branch (see `test_build_bspline_basis_1d_default_identifiability_densifies_via_orthonormal_centering`).
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "s_x".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 32,
                        },
                        double_penalty: false,
                        identifiability: BSplineIdentifiability::None,
                        boundary: OneDimensionalBoundary::Open,
                        boundary_conditions: BSplineBoundaryConditions::default(),
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");

        assert!(matches!(design.design, DesignMatrix::Sparse(_)));
        assert!(exact_joint_spatial_outer_hessian_available(
            &LikelihoodSpec::binomial_logit(),
            &design,
        ));
    }

    #[test]
    fn iso_kappa_duchon_dx_dpsi_matches_fd() {
        // Compare the production frozen-spec dX/dψ path against centered FD
        // of X(ψ+h) - X(ψ-h). This intentionally goes through
        // `try_build_spatial_term_log_kappa_derivative`: the formula layer owns
        // the frozen centers, length-scale compensation, and composed
        // identifiability transform.
        let n = 80usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
        }
        let spec_orig = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_1d".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                        length_scale: Some(1.0),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec_orig).expect("design");
        let frozen = freeze_term_collection_from_design(&spec_orig, &design).expect("freeze");

        let build_design_at = |psi: f64| -> Array2<f64> {
            // Rebuild design at psi via direct kernel build using frozen spec.
            let mut s = frozen.clone();
            if let SmoothBasisSpec::Duchon {
                spec: ref mut duchon,
                ..
            } = s.smooth_terms[0].basis
            {
                duchon.length_scale = Some((-psi).exp());
            }
            let d = build_term_collection_design(data.view(), &s).expect("rebuild");
            d.design.to_dense()
        };

        // Build derivative at psi=0.
        let psi_eval = 0.0_f64;
        let derivative_bundle =
            try_build_spatial_term_log_kappa_derivative(data.view(), &frozen, &design, 0)
                .expect("formula Duchon derivative should build")
                .expect("Duchon derivative should be available");
        let global_range = derivative_bundle.0;
        let p_total = derivative_bundle.1;
        let implicit_operator = derivative_bundle.8;
        let op = implicit_operator.expect("Duchon derivative should expose implicit operator");
        let p = op.p_out();
        assert_eq!(p_total, design.design.ncols());
        assert_eq!(global_range.end - global_range.start, p);

        // FD reference.
        let h = 1e-4_f64;
        let x_plus = build_design_at(psi_eval + h);
        let x_minus = build_design_at(psi_eval - h);
        eprintln!(
            "[DXDPSI_FD] X(+h)[0,0..3]={:?} X(-h)[0,0..3]={:?}",
            x_plus.row(0).iter().take(3).copied().collect::<Vec<_>>(),
            x_minus.row(0).iter().take(3).copied().collect::<Vec<_>>(),
        );
        eprintln!(
            "[DXDPSI_FD] X(+h) shape={:?} X(-h) shape={:?} p_out={}",
            x_plus.shape(),
            x_minus.shape(),
            p,
        );
        // Also build at psi_eval to compare cols.
        let x_at = build_design_at(psi_eval);
        let orig_design =
            build_term_collection_design(data.view(), &spec_orig).expect("rebuild orig");
        eprintln!(
            "[DXDPSI_FD] X(psi_eval) shape={:?} orig_design.ncols={}",
            x_at.shape(),
            orig_design.design.ncols(),
        );

        // Multiply analytic operator by unit basis vectors.
        let mut analytic = Array2::<f64>::zeros((n, p));
        let mut basisv = Array1::<f64>::zeros(p);
        for j in 0..p {
            basisv[j] = 1.0;
            let col = op.forward_mul(0, &basisv.view()).expect("forward_mul");
            analytic.column_mut(j).assign(&col);
            basisv[j] = 0.0;
        }

        // Also check transpose_mul: X_tau^T v for v of length n.
        // FD reference: X_tau^T v should be (X(+h)^T - X(-h)^T)/(2h) · v.
        let smooth_start = global_range.start;
        let v_test = Array1::<f64>::from_shape_fn(n, |i| (i as f64 * 0.07).sin());
        let analytic_tv = op.transpose_mul(0, &v_test.view()).expect("transpose_mul");
        let fd_tv_full = (&x_plus.t() - &x_minus.t()) / (2.0 * h);
        let fd_tv = fd_tv_full.dot(&v_test);
        // Extract smooth portion only
        let fd_tv_smooth = fd_tv.slice(s![smooth_start..(smooth_start + p)]).to_owned();
        let max_tv_diff = analytic_tv
            .iter()
            .zip(fd_tv_smooth.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f64, f64::max);
        let max_tv_abs = analytic_tv.iter().map(|v| v.abs()).fold(0.0f64, f64::max);
        eprintln!(
            "[DXDPSI_TV] max|analytic_tv - fd_tv|={:.3e}  max|analytic_tv|={:.3e}",
            max_tv_diff, max_tv_abs
        );
        eprintln!(
            "[DXDPSI_TV] analytic_tv={:?}",
            analytic_tv.iter().take(p).copied().collect::<Vec<_>>()
        );
        eprintln!(
            "[DXDPSI_TV] fd_tv_smooth={:?}",
            fd_tv_smooth.iter().take(p).copied().collect::<Vec<_>>()
        );
        let fd_full = (&x_plus - &x_minus) / (2.0 * h);
        let fd = fd_full
            .slice(s![.., smooth_start..(smooth_start + p)])
            .to_owned();
        let mut max_diff = 0.0_f64;
        let mut max_abs = 0.0_f64;
        for i in 0..n {
            for j in 0..p {
                let d = (analytic[[i, j]] - fd[[i, j]]).abs();
                if d > max_diff {
                    max_diff = d;
                }
                if analytic[[i, j]].abs() > max_abs {
                    max_abs = analytic[[i, j]].abs();
                }
            }
        }
        eprintln!(
            "[DXDPSI_FD] max|analytic - fd|={:.3e}  max|analytic|={:.3e}",
            max_diff, max_abs
        );
        eprintln!(
            "[DXDPSI_FD] analytic[0,..]={:?}",
            analytic.row(0).iter().take(p).copied().collect::<Vec<_>>(),
        );
        eprintln!(
            "[DXDPSI_FD] fd[0,..]={:?}",
            fd.row(0).iter().take(p).copied().collect::<Vec<_>>(),
        );
        assert!(max_diff < 5e-3 * max_abs.max(1e-3), "dX/dψ mismatch");
    }

    #[test]
    fn joint_build_and_freeze_shares_auto_spatial_centers_across_blocks() {
        let n = 400usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
            data[[i, 1]] = (i as f64 * 0.19).sin();
        }

        let matern_term = |name: &str| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::Matern {
                feature_cols: vec![0, 1],
                spec: MaternBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::Auto(Box::new(
                        CenterStrategy::FarthestPoint { num_centers: 8 },
                    )),
                    length_scale: 0.8,
                    nu: MaternNu::FiveHalves,
                    include_intercept: false,
                    double_penalty: true,
                    identifiability: MaternIdentifiability::CenterSumToZero,
                    aniso_log_scales: None,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let marginalspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("marginal")],
        };
        let logslopespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("logslope")],
        };

        let (designs, resolved_specs) = build_term_collection_designs_and_freeze_joint(
            data.view(),
            &[marginalspec.clone(), logslopespec.clone()],
        )
        .expect("joint build and freeze should succeed");

        assert_eq!(designs.len(), 2);
        assert_eq!(resolved_specs.len(), 2);

        let extract_centers = |spec: &TermCollectionSpec| match &spec.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => match &spec.center_strategy {
                CenterStrategy::UserProvided(centers) => centers.clone(),
                other => panic!("expected frozen user-provided centers, got {other:?}"),
            },
            other => panic!("expected Matérn term, got {other:?}"),
        };

        let marginal_centers = extract_centers(&resolved_specs[0]);
        let logslope_centers = extract_centers(&resolved_specs[1]);
        let separate_marginal_design =
            build_term_collection_design(data.view(), &marginalspec).expect("separate marginal");
        let separate_marginal =
            freeze_term_collection_from_design(&marginalspec, &separate_marginal_design)
                .expect("freeze separate marginal");
        let separate_marginal_centers = extract_centers(&separate_marginal);

        assert_eq!(marginal_centers, logslope_centers);
        assert_eq!(marginal_centers.ncols(), 2);
        assert_eq!(marginal_centers.nrows(), separate_marginal_centers.nrows());
    }

    #[test]
    fn exact_joint_two_block_no_spatial_fast_path_returns_fully_frozen_specs() {
        let n = 24usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (i % 3) as f64;
        }

        let pspline_term = |name: &str| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::BSpline1D {
                feature_col: 0,
                spec: BSplineBasisSpec {
                    degree: 3,
                    penalty_order: 2,
                    knotspec: BSplineKnotSpec::Generate {
                        data_range: (0.0, 1.0),
                        num_internal_knots: 3,
                    },
                    double_penalty: true,
                    identifiability: BSplineIdentifiability::None,
                    boundary: OneDimensionalBoundary::Open,
                    boundary_conditions: BSplineBoundaryConditions::default(),
                },
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };
        let random_effect = RandomEffectTermSpec {
            name: "grp".to_string(),
            feature_col: 1,
            drop_first_level: false,
            penalized: true,
            frozen_levels: None,
        };

        let meanspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![random_effect.clone()],
            smooth_terms: vec![pspline_term("mean_ps")],
        };
        let noisespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![random_effect],
            smooth_terms: vec![pspline_term("noise_ps")],
        };

        let kappa_options = SpatialLengthScaleOptimizationOptions {
            max_outer_iter: 1,
            rel_tol: 1e-6,
            pilot_subsample_threshold: 0,
            ..SpatialLengthScaleOptimizationOptions::default()
        };
        let joint_setup = two_block_exact_joint_hyper_setup(&meanspec, &noisespec, &kappa_options);
        let theta_dim = joint_setup.theta0().len();

        let mean_terms = spatial_length_scale_term_indices(&meanspec);
        let noise_terms = spatial_length_scale_term_indices(&noisespec);
        assert!(mean_terms.is_empty());
        assert!(noise_terms.is_empty());

        let policy = crate::families::custom_family::OuterDerivativePolicy {
            capability: crate::families::custom_family::ExactOuterDerivativeOrder::Second,
            predicted_hessian_work: 0,
            predicted_gradient_work: 0,
            // Test-style construction with zero predicted work — these
            // paths never engage staged-κ, so the capability bit is
            // moot. Keep `false` as the safe default.
            subsample_capable: false,
        };
        let solved = optimize_spatial_length_scale_exact_joint(
            data.view(),
            &[meanspec.clone(), noisespec.clone()],
            &[mean_terms, noise_terms],
            &kappa_options,
            &joint_setup,
            crate::seeding::SeedRiskProfile::Gaussian,
            true,
            true,
            false,
            None,
            policy,
            |theta, specs, designs| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                assert_eq!(designs.len(), 2);
                Ok(designs[0].design.ncols() as f64 + designs[1].design.ncols() as f64)
            },
            |theta, specs, designs, eval_mode, _| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                assert_eq!(designs.len(), 2);
                Ok((
                    0.0,
                    Array1::zeros(theta_dim),
                    if matches!(
                        eval_mode,
                        crate::solver::estimate::reml::unified::EvalMode::ValueGradientHessian
                    ) {
                        crate::solver::outer_strategy::HessianResult::Analytic(Array2::zeros((
                            theta_dim, theta_dim,
                        )))
                    } else {
                        crate::solver::outer_strategy::HessianResult::Unavailable
                    },
                ))
            },
            |theta, specs, designs| {
                assert_eq!(theta.len(), theta_dim);
                assert_eq!(specs.len(), 2);
                assert_eq!(designs.len(), 2);
                Ok(crate::solver::outer_strategy::EfsEval {
                    cost: 0.0,
                    steps: vec![0.0; theta_dim],
                    beta: None,
                    psi_gradient: None,
                    psi_indices: None,
                    inner_hessian_scale: None,
                })
            },
            |_beta: &Array1<f64>| Ok(()),
        )
        .expect("exact joint no-spatial fast path should succeed");

        for resolved in [&solved.resolved_specs[0], &solved.resolved_specs[1]] {
            resolved
                .validate_frozen("resolvedspec")
                .expect("exact joint no-spatial fast path should fully freeze specs");
            match &resolved.smooth_terms[0].basis {
                SmoothBasisSpec::BSpline1D { spec, .. } => {
                    assert!(matches!(spec.knotspec, BSplineKnotSpec::Provided(_)));
                }
                _ => panic!("expected P-spline term"),
            }
            assert!(
                resolved.random_effect_terms[0].frozen_levels.is_some(),
                "random-effect levels should be frozen in exact joint no-spatial fast path"
            );
        }
    }

    #[test]
    fn incremental_frozen_realizer_matches_unified_full_rebuild() {
        let n = 24usize;
        let mut data = Array2::<f64>::zeros((n, 4));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (0.35 * i as f64).sin();
            data[[i, 2]] = (i % 3) as f64;
            data[[i, 3]] = t * t;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "lin".to_string(),
                feature_col: 1,
                feature_cols: vec![1],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: Some(-0.5),
                coefficient_max: None,
            }],
            random_effect_terms: vec![RandomEffectTermSpec {
                name: "grp".to_string(),
                feature_col: 2,
                drop_first_level: false,
                penalized: true,
                frozen_levels: None,
            }],
            smooth_terms: vec![
                SmoothTermSpec {
                    name: "spatial".to_string(),
                    basis: SmoothBasisSpec::Matern {
                        feature_cols: vec![0, 1],
                        spec: MaternBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                            length_scale: 0.8,
                            nu: MaternNu::FiveHalves,
                            include_intercept: false,
                            double_penalty: true,
                            identifiability: MaternIdentifiability::CenterSumToZero,
                            aniso_log_scales: Some(vec![0.15, -0.15]),
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "mono".to_string(),
                    basis: SmoothBasisSpec::BSpline1D {
                        feature_col: 3,
                        spec: BSplineBasisSpec {
                            degree: 3,
                            penalty_order: 2,
                            knotspec: BSplineKnotSpec::Generate {
                                data_range: (0.0, 1.0),
                                num_internal_knots: 3,
                            },
                            double_penalty: false,
                            identifiability: BSplineIdentifiability::None,
                            boundary: OneDimensionalBoundary::Open,
                            boundary_conditions: BSplineBoundaryConditions::default(),
                        },
                    },
                    shape: ShapeConstraint::MonotoneIncreasing,
                    joint_null_rotation: None,
                },
            ],
        };

        let base_design = build_term_collection_design(data.view(), &spec).expect("base design");
        let frozen = freeze_term_collection_from_design(&spec, &base_design).expect("freeze");
        let frozen_design =
            build_term_collection_design(data.view(), &frozen).expect("frozen design");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        assert_eq!(spatial_terms, vec![0]);

        let smooth_start = frozen_design.design.ncols() - frozen_design.smooth.total_smooth_cols();
        let fixed_before = frozen_design.design.clone();
        let nonspatial_range = frozen_design.smooth.terms[1].coeff_range.clone();
        let full_nonspatial_range =
            (smooth_start + nonspatial_range.start)..(smooth_start + nonspatial_range.end);
        let mut realizer = FrozenTermCollectionIncrementalRealizer::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
        )
        .expect("incremental realizer");

        let updated_log_kappa = SpatialLogKappaCoords::new_with_dims(array![0.30, -0.20], vec![2]);
        let updated_spec = updated_log_kappa
            .apply_tospec(&frozen, &spatial_terms)
            .expect("updated spec");
        realizer
            .apply_log_kappa(&updated_log_kappa, &spatial_terms)
            .expect("incremental update");
        let rebuilt =
            build_term_collection_design(data.view(), &updated_spec).expect("rebuilt design");

        assert_term_collection_designs_match(realizer.design(), &rebuilt, "incremental realizer");

        let linear_range = frozen_design.linear_ranges[0].1.clone();
        let random_range = frozen_design.random_effect_ranges[0].1.clone();
        let fixed_before_dense = fixed_before.to_dense();
        let updated_full_dense = realizer.design().design.to_dense();
        let linear_diff = max_abs_diff_matrix(
            &fixed_before_dense
                .slice(s![.., linear_range.clone()])
                .to_owned(),
            &updated_full_dense.slice(s![.., linear_range]).to_owned(),
        );
        let random_diff = max_abs_diff_matrix(
            &fixed_before_dense
                .slice(s![.., random_range.clone()])
                .to_owned(),
            &updated_full_dense.slice(s![.., random_range]).to_owned(),
        );
        let nonspatial_diff = max_abs_diff_matrix(
            &fixed_before_dense
                .slice(s![.., full_nonspatial_range.clone()])
                .to_owned(),
            &updated_full_dense
                .slice(s![.., full_nonspatial_range.clone()])
                .to_owned(),
        );
        let spatial_range = frozen_design.smooth.terms[0].coeff_range.clone();
        let full_spatial_range =
            (smooth_start + spatial_range.start)..(smooth_start + spatial_range.end);
        let spatial_change = max_abs_diff_matrix(
            &fixed_before_dense
                .slice(s![.., full_spatial_range.clone()])
                .to_owned(),
            &updated_full_dense
                .slice(s![.., full_spatial_range])
                .to_owned(),
        );
        assert!(
            linear_diff <= 1e-12,
            "linear block changed max_abs={linear_diff}"
        );
        assert!(
            random_diff <= 1e-12,
            "random-effect block changed max_abs={random_diff}"
        );
        assert!(
            nonspatial_diff <= 1e-12,
            "unchanged smooth block changed max_abs={nonspatial_diff}"
        );
        assert!(
            spatial_change > 1e-8,
            "spatial block did not update max_abs={spatial_change}"
        );
    }

    #[test]
    fn two_block_exact_joint_design_cache_clears_memo_on_theta_change() {
        let n = 20usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.19 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
        }

        // ANISOTROPIC Matérn (`aniso_log_scales = Some`): the two-block
        // exact-joint design cache memoizes per-block κ/η axes, which (#519)
        // only exist for anisotropic spatial terms — isotropic Matérn anchors
        // its data-seeded κ and contributes no κ axis. Per-axis scales give
        // each block the log-κ/η hyper axes this cache test drives.
        let matern_term = |name: &str, length_scale: f64| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::Matern {
                feature_cols: vec![0, 1],
                spec: MaternBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 5 },
                    length_scale,
                    nu: MaternNu::FiveHalves,
                    include_intercept: false,
                    double_penalty: true,
                    identifiability: MaternIdentifiability::CenterSumToZero,
                    aniso_log_scales: Some(vec![0.0, 0.0]),
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let meanspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("mean", 0.7)],
        };
        let noisespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![matern_term("noise", 1.1)],
        };
        let kappa_options = SpatialLengthScaleOptimizationOptions {
            max_outer_iter: 1,
            rel_tol: 1e-6,
            pilot_subsample_threshold: 0,
            ..SpatialLengthScaleOptimizationOptions::default()
        };
        let joint_setup = two_block_exact_joint_hyper_setup(&meanspec, &noisespec, &kappa_options);
        let theta0 = joint_setup.theta0();

        let mean_design = build_term_collection_design(data.view(), &meanspec).expect("mean");
        let noise_design = build_term_collection_design(data.view(), &noisespec).expect("noise");
        let mean_frozen =
            freeze_term_collection_from_design(&meanspec, &mean_design).expect("freeze mean");
        let noise_frozen =
            freeze_term_collection_from_design(&noisespec, &noise_design).expect("freeze noise");

        let mean_term_indices = spatial_length_scale_term_indices(&mean_frozen);
        let noise_term_indices = spatial_length_scale_term_indices(&noise_frozen);
        let mut cache = ExactJointDesignCache::new(
            data.view(),
            vec![
                (
                    mean_frozen.clone(),
                    mean_design.clone(),
                    mean_term_indices.clone(),
                ),
                (
                    noise_frozen.clone(),
                    noise_design.clone(),
                    noise_term_indices.clone(),
                ),
            ],
            joint_setup.rho_dim(),
            joint_setup.log_kappa_dims_per_term(),
        )
        .expect("n-block cache");

        cache.ensure_theta(&theta0).expect("initial theta");
        assert!(cache.memoized_cost(&theta0).is_none());
        assert!(cache.memoized_eval(&theta0).is_none());

        let eval = (
            2.25,
            Array1::<f64>::ones(theta0.len()),
            crate::solver::outer_strategy::HessianResult::Analytic(Array2::<f64>::eye(
                theta0.len(),
            )),
        );
        cache.store_eval(eval.clone());
        let cached_eval = cache.memoized_eval(&theta0).expect("cached eval");
        assert!((cached_eval.0 - eval.0).abs() <= 1e-12);
        assert_eq!(cached_eval.1, eval.1);
        assert_eq!(
            cached_eval
                .2
                .materialize_dense()
                .expect("materialize cached hessian"),
            eval.2
                .materialize_dense()
                .expect("materialize eval hessian"),
        );

        let mut theta1 = theta0.clone();
        theta1[joint_setup.rho_dim()] += 0.25;
        cache.ensure_theta(&theta1).expect("updated theta");
        assert!(cache.memoized_cost(&theta1).is_none());
        assert!(cache.memoized_eval(&theta1).is_none());

        let log_kappa = SpatialLogKappaCoords::from_theta_tail_with_dims(
            &theta1,
            joint_setup.rho_dim(),
            joint_setup.log_kappa_dims_per_term(),
        );
        let mean_terms = spatial_length_scale_term_indices(&mean_frozen);
        let noise_terms = spatial_length_scale_term_indices(&noise_frozen);
        let (mean_lk, noise_lk) = log_kappa.split_at(mean_terms.len());
        let mean_updated = mean_lk
            .apply_tospec(&mean_frozen, &mean_terms)
            .expect("mean updated spec");
        let noise_updated = noise_lk
            .apply_tospec(&noise_frozen, &noise_terms)
            .expect("noise updated spec");
        let mean_rebuilt =
            build_term_collection_design(data.view(), &mean_updated).expect("mean rebuilt");
        let noise_rebuilt =
            build_term_collection_design(data.view(), &noise_updated).expect("noise rebuilt");
        let cache_designs = cache.designs();
        assert_term_collection_designs_match(cache_designs[0], &mean_rebuilt, "mean cache");
        assert_term_collection_designs_match(cache_designs[1], &noise_rebuilt, "noise cache");
    }

    #[test]
    fn single_block_exact_joint_design_cache_clears_memo_on_theta_change() {
        let n = 22usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.23 * i as f64).cos();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
        }

        // Hybrid Duchon term with an explicit scalar `length_scale`: this is
        // the canonical single-log-κ-axis spatial term (`dims_per_term == [1]`)
        // that the single-block exact-joint design cache is built to memoize.
        // (#519 — isotropic Matérn no longer contributes a κ axis; it anchors
        // its data-seeded κ and learns smoothness through ρ alone, so it is the
        // wrong fixture for a single-κ-axis cache test. Hybrid Duchon keeps the
        // scalar κ axis without any of the brittle isotropic-Matérn κ-search.)
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_hybrid".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: Some(0.9),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze spec");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        let rho_dim = design.penalties.len();
        let dims_per_term = vec![1];
        let mut theta0 = Array1::<f64>::zeros(rho_dim + 1);
        theta0[rho_dim] = -get_spatial_length_scale(&frozen, spatial_terms[0])
            .expect("length scale")
            .ln();

        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            design.clone(),
            spatial_terms.clone(),
            rho_dim,
            dims_per_term.clone(),
        )
        .expect("single-block cache");

        cache.ensure_theta(&theta0).expect("initial theta");
        assert!(cache.memoized_cost(&theta0).is_none());
        assert!(cache.memoized_eval(&theta0).is_none());

        let eval = (
            0.5,
            Array1::<f64>::ones(theta0.len()),
            crate::solver::outer_strategy::HessianResult::Analytic(Array2::<f64>::eye(
                theta0.len(),
            )),
        );
        cache.store_eval(eval.clone());
        let cached_eval = cache.memoized_eval(&theta0).expect("cached eval");
        assert!((cached_eval.0 - eval.0).abs() <= 1e-12);
        assert_eq!(cached_eval.1, eval.1);
        assert_eq!(
            cached_eval
                .2
                .materialize_dense()
                .expect("materialize cached hessian"),
            eval.2
                .materialize_dense()
                .expect("materialize eval hessian"),
        );

        let mut theta1 = theta0.clone();
        theta1[rho_dim] += 0.35;
        cache.ensure_theta(&theta1).expect("updated theta");
        assert!(cache.memoized_cost(&theta1).is_none());
        assert!(cache.memoized_eval(&theta1).is_none());

        let updated_log_kappa =
            SpatialLogKappaCoords::from_theta_tail_with_dims(&theta1, rho_dim, dims_per_term);
        let updated_spec = updated_log_kappa
            .apply_tospec(&frozen, &spatial_terms)
            .expect("updated spec");
        let rebuilt =
            build_term_collection_design(data.view(), &updated_spec).expect("rebuilt design");
        assert_term_collection_designs_match(cache.design(), &rebuilt, "single-block cache");
    }

    #[test]
    fn single_block_latent_coord_design_cache_invalidates_memo_on_outer_iter_advance() {
        // Pins θ and re-evaluates after current_outer_iter() advances; the
        // scheduled penalty weight at that θ has changed, so the memo must miss.
        use crate::solver::estimate::reml::runtime::{
            current_outer_iter, record_current_outer_iter_for_ift,
        };
        use crate::terms::latent_coord::{LatentCoordValues, LatentIdMode, LatentManifold};

        let n = 16usize;
        let latent_dim = 1usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "s(x0)".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 3,
                        },
                        double_penalty: false,
                        identifiability: BSplineIdentifiability::WeightedSumToZero {
                            weights: None,
                        },
                        boundary_conditions: BSplineBoundaryConditions::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let rho_dim = design.penalties.len();

        let flat = Array1::<f64>::from_iter((0..n).map(|i| i as f64 / (n as f64 - 1.0)));
        let latent_values = std::sync::Arc::new(LatentCoordValues::from_flat_with_manifold(
            flat,
            n,
            latent_dim,
            LatentIdMode::None,
            LatentManifold::Euclidean,
        ));
        let latent = StandardLatentCoordConfig {
            values: latent_values,
            term_index: crate::types::SmoothTermIdx::new(0),
            feature_cols: vec![0],
            manifold: LatentManifold::Euclidean,
            manifold_auto: false,
            retraction_registry: crate::solver::latent_cache::LatentRetractionRegistry::default(),
            analytic_penalties: None,
        };

        let mut cache = SingleBlockLatentCoordDesignCache::new(
            data.clone(),
            spec.clone(),
            design.clone(),
            &latent,
            rho_dim,
        )
        .expect("latent-coord cache");

        // Seed θ + cached eval directly (skip `ensure_theta`'s latent rebuild;
        // we only exercise the memo invalidation property). Same-module access.
        let theta = Array1::<f64>::zeros(rho_dim + n * latent_dim);
        cache.current_theta = Some(theta.clone());

        let initial_iter = current_outer_iter();
        let eval = (
            1.25_f64,
            Array1::<f64>::from_elem(theta.len(), 0.5),
            crate::solver::outer_strategy::HessianResult::Analytic(Array2::<f64>::eye(theta.len())),
        );
        cache.store_eval(eval.clone());

        let hit = cache
            .memoized_eval(&theta)
            .expect("memo should hit at the same outer iter");
        assert!((hit.0 - eval.0).abs() <= 1e-12);
        assert_eq!(cache.memoized_cost(&theta), Some(eval.0));

        record_current_outer_iter_for_ift(initial_iter.wrapping_add(1));
        assert!(
            cache.memoized_eval(&theta).is_none(),
            "memoized_eval returned a stale cached eval after current_outer_iter advanced"
        );
        assert!(
            cache.memoized_cost(&theta).is_none(),
            "memoized_cost returned a stale cached cost after current_outer_iter advanced"
        );

        record_current_outer_iter_for_ift(initial_iter);
    }

    #[test]
    fn external_joint_evaluator_reuse_matches_fresh_state_after_theta_update() {
        let n = 26usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.21 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            y[i] = (2.0 * std::f64::consts::PI * x0).sin() + 0.35 * x1;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "x0".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: None,
                coefficient_max: None,
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 0.85,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);
        let fit_opts = FitOptions {
            compute_inference: false,
            max_iter: 40,
            tol: 1e-7,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let frozen = freeze_term_collection_from_design(&spec, &design).expect("freeze");
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        let rho_dim = design.penalties.len();
        let mut theta0 = Array1::<f64>::zeros(rho_dim + dims_per_term.iter().sum::<usize>());
        for j in 0..rho_dim {
            theta0[j] = 0.2 - 0.1 * j as f64;
        }
        theta0[rho_dim] = -get_spatial_length_scale(&frozen, spatial_terms[0])
            .expect("length scale")
            .ln();
        let mut theta1 = theta0.clone();
        theta1[rho_dim] += 0.3;

        let external_opts =
            external_opts_for_design(&LikelihoodSpec::gaussian_identity(), &design, &fit_opts);
        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen,
            design.clone(),
            spatial_terms,
            rho_dim,
            dims_per_term,
        )
        .expect("single-block cache");
        let mut reused = crate::estimate::ExternalJointHyperEvaluator::new(
            y.view(),
            weights.view(),
            &design.design,
            offset.view(),
            &design.penalties,
            &external_opts,
            "reused evaluator",
        )
        .expect("reused evaluator");

        let compare_eval =
            |theta: &Array1<f64>,
             cache: &mut SingleBlockExactJointDesignCache<'_>,
             reused: &mut crate::estimate::ExternalJointHyperEvaluator<'_>| {
                cache.ensure_theta(theta).expect("theta applied");

                let build_hyper_dirs = || {
                    try_build_spatial_log_kappa_hyper_dirs(
                        data.view(),
                        cache.spec(),
                        cache.design(),
                        &cache.spatial_terms,
                    )
                    .expect("hyper dirs build")
                    .expect("hyper dirs present")
                };

                let reused_eval = evaluate_joint_reml_outer_eval_at_theta(
                    reused,
                    cache.design(),
                    theta,
                    rho_dim,
                    build_hyper_dirs(),
                    None,
                    crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian,
                    None,
                )
                .expect("reused eval");

                let fresh_opts = external_opts_for_design(
                    &LikelihoodSpec::gaussian_identity(),
                    cache.design(),
                    &fit_opts,
                );
                let mut fresh = crate::estimate::ExternalJointHyperEvaluator::new(
                    y.view(),
                    weights.view(),
                    &cache.design().design,
                    offset.view(),
                    &cache.design().penalties,
                    &fresh_opts,
                    "fresh evaluator",
                )
                .expect("fresh evaluator");
                let fresh_eval = evaluate_joint_reml_outer_eval_at_theta(
                    &mut fresh,
                    cache.design(),
                    theta,
                    rho_dim,
                    build_hyper_dirs(),
                    None,
                    crate::solver::outer_strategy::OuterEvalOrder::ValueGradientHessian,
                    None,
                )
                .expect("fresh eval");

                let cost_diff = (reused_eval.0 - fresh_eval.0).abs();
                assert!(cost_diff <= 1e-10, "cost mismatch: {cost_diff}");

                let grad_diff = reused_eval
                    .1
                    .iter()
                    .zip(fresh_eval.1.iter())
                    .map(|(left, right)| (left - right).abs())
                    .fold(0.0_f64, f64::max);
                assert!(grad_diff <= 1e-9, "gradient mismatch: {grad_diff}");

                let reused_hess = reused_eval
                    .2
                    .materialize_dense()
                    .expect("reused hessian materializes")
                    .expect("reused hessian present");
                let fresh_hess = fresh_eval
                    .2
                    .materialize_dense()
                    .expect("fresh hessian materializes")
                    .expect("fresh hessian present");
                let hess_diff = max_abs_diff_matrix(&reused_hess, &fresh_hess);
                assert!(hess_diff <= 1e-9, "hessian mismatch: {hess_diff}");

                let reused_efs = evaluate_joint_reml_efs_at_theta(
                    reused,
                    cache.design(),
                    theta,
                    rho_dim,
                    build_hyper_dirs(),
                    None,
                    None,
                )
                .expect("reused EFS eval");

                let mut fresh_efs_eval = crate::estimate::ExternalJointHyperEvaluator::new(
                    y.view(),
                    weights.view(),
                    &cache.design().design,
                    offset.view(),
                    &cache.design().penalties,
                    &fresh_opts,
                    "fresh EFS evaluator",
                )
                .expect("fresh EFS evaluator");
                let fresh_efs = evaluate_joint_reml_efs_at_theta(
                    &mut fresh_efs_eval,
                    cache.design(),
                    theta,
                    rho_dim,
                    build_hyper_dirs(),
                    None,
                    None,
                )
                .expect("fresh EFS eval");

                let efs_cost_diff = (reused_efs.cost - fresh_efs.cost).abs();
                assert!(efs_cost_diff <= 1e-10, "EFS cost mismatch: {efs_cost_diff}");
                assert_eq!(reused_efs.steps.len(), fresh_efs.steps.len());
                let efs_step_diff = reused_efs
                    .steps
                    .iter()
                    .zip(fresh_efs.steps.iter())
                    .map(|(left, right)| (left - right).abs())
                    .fold(0.0_f64, f64::max);
                assert!(efs_step_diff <= 1e-9, "EFS step mismatch: {efs_step_diff}");
            };

        compare_eval(&theta0, &mut cache, &mut reused);
        compare_eval(&theta1, &mut cache, &mut reused);
    }

    #[test]
    fn exact_matern_log_kappa_derivative_uses_feature_columns_only() {
        let n = 24usize;
        let p = 17usize;
        let mut data = Array2::<f64>::zeros((n, p));
        for i in 0..n {
            let x = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = x;
            for j in 1..p {
                data[[i, j]] = ((i + j) as f64 * 0.13).sin();
            }
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 0.4,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec)
            .expect("baseline Matérn design should build");
        let frozenspec = freeze_term_collection_from_design(&spec, &design)
            .expect("freezing Matérn centers from design should succeed");

        match &frozenspec.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => match &spec.center_strategy {
                CenterStrategy::UserProvided(centers) => {
                    assert_eq!(centers.ncols(), 1, "frozen centers should stay term-local");
                }
                _ => panic!("expected frozen user-provided centers"),
            },
            _ => panic!("expected Matérn term"),
        }

        let derivative =
            try_build_spatial_term_log_kappa_derivative(data.view(), &frozenspec, &design, 0);
        assert!(
            derivative.is_ok(),
            "exact Matérn log-kappa derivative should use only feature_cols; got {derivative:?}"
        );
        assert!(
            derivative
                .expect("derivative call should succeed")
                .is_some(),
            "Matérn term should expose an exact derivative"
        );
    }

    #[test]
    fn exact_thin_plate_log_kappa_derivative_uses_feature_columns_only() {
        let n = 28usize;
        let p = 15usize;
        let mut data = Array2::<f64>::zeros((n, p));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.17 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            for j in 2..p {
                data[[i, j]] = ((i + 3 * j) as f64 * 0.07).cos();
            }
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "thinplate".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 7 },
                        length_scale: 0.7,
                        double_penalty: true,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec)
            .expect("baseline ThinPlate design should build");
        let frozenspec = freeze_term_collection_from_design(&spec, &design)
            .expect("freezing ThinPlate centers from design should succeed");

        match &frozenspec.smooth_terms[0].basis {
            SmoothBasisSpec::ThinPlate { spec, .. } => match &spec.center_strategy {
                CenterStrategy::UserProvided(centers) => {
                    assert_eq!(centers.ncols(), 2, "frozen centers should stay term-local");
                }
                _ => panic!("expected frozen user-provided centers"),
            },
            _ => panic!("expected ThinPlate term"),
        }

        let smooth_term = &design.smooth.terms[0];
        let termspec = &frozenspec.smooth_terms[0];
        let BasisPsiDerivativeResult {
            design_derivative: local_x_psi,
            penalties_derivative: local_s_psi,
            ..
        } = match &termspec.basis {
            SmoothBasisSpec::ThinPlate {
                feature_cols, spec, ..
            } => {
                let x = select_columns(data.view(), feature_cols)
                    .expect("select ThinPlate feature cols");
                crate::basis::build_thin_plate_basis_log_kappa_derivative(x.view(), spec)
                    .expect("direct ThinPlate derivative should build")
            }
            _ => panic!("expected ThinPlate term"),
        };
        let BasisPsiSecondDerivativeResult {
            designsecond_derivative: local_x_psi_psi,
            penaltiessecond_derivative: local_s_psi_psi,
            ..
        } = match &termspec.basis {
            SmoothBasisSpec::ThinPlate {
                feature_cols, spec, ..
            } => {
                let x = select_columns(data.view(), feature_cols)
                    .expect("select ThinPlate feature cols");
                crate::basis::build_thin_plate_basis_log_kappasecond_derivative(x.view(), spec)
                    .expect("direct ThinPlate second derivative should build")
            }
            _ => panic!("expected ThinPlate term"),
        };
        assert_eq!(local_x_psi.ncols(), smooth_term.coeff_range.len());
        assert_eq!(local_x_psi_psi.ncols(), smooth_term.coeff_range.len());
        assert!(!local_s_psi.is_empty());
        assert_eq!(local_s_psi.len(), local_s_psi_psi.len());
        assert!(local_s_psi.iter().all(|s| {
            s.nrows() == smooth_term.coeff_range.len() && s.ncols() == smooth_term.coeff_range.len()
        }));
        assert!(local_s_psi_psi.iter().all(|s| {
            s.nrows() == smooth_term.coeff_range.len() && s.ncols() == smooth_term.coeff_range.len()
        }));

        let derivative =
            try_build_spatial_term_log_kappa_derivative(data.view(), &frozenspec, &design, 0);
        assert!(
            derivative.is_ok(),
            "exact ThinPlate log-kappa derivative should use only feature_cols; got {derivative:?}"
        );
        let derivative = derivative.expect("derivative call should succeed");
        assert!(
            derivative.is_some(),
            "ThinPlate term should expose an exact derivative"
        );
    }

    #[test]
    fn exact_duchon_log_kappa_derivative_uses_feature_columns_only() {
        let n = 28usize;
        let p = 15usize;
        let mut data = Array2::<f64>::zeros((n, p));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.21 * i as f64).cos();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            for j in 2..p {
                data[[i, j]] = ((i + 2 * j) as f64 * 0.09).sin();
            }
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 7 },
                        length_scale: Some(0.7),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec)
            .expect("baseline Duchon design should build");
        let frozenspec = freeze_term_collection_from_design(&spec, &design)
            .expect("freezing Duchon centers from design should succeed");

        match &frozenspec.smooth_terms[0].basis {
            SmoothBasisSpec::Duchon { spec, .. } => match &spec.center_strategy {
                CenterStrategy::UserProvided(centers) => {
                    assert_eq!(centers.ncols(), 2, "frozen centers should stay term-local");
                }
                _ => panic!("expected frozen user-provided centers"),
            },
            _ => panic!("expected Duchon term"),
        }

        let smooth_term = &design.smooth.terms[0];
        let termspec = &frozenspec.smooth_terms[0];
        let derivative_bundle = match &termspec.basis {
            SmoothBasisSpec::Duchon {
                feature_cols, spec, ..
            } => {
                let x =
                    select_columns(data.view(), feature_cols).expect("select Duchon feature cols");
                crate::basis::build_duchon_basis_log_kappa_derivatives(x.view(), spec)
                    .expect("direct Duchon derivative bundle should build")
            }
            _ => panic!("expected Duchon term"),
        };
        let local_implicit = derivative_bundle.implicit_operator;
        let BasisPsiDerivativeResult {
            design_derivative: local_x_psi,
            penalties_derivative: local_s_psi,
            implicit_operator: local_implicit_psi_unused,
        } = derivative_bundle.first;
        let BasisPsiSecondDerivativeResult {
            designsecond_derivative: local_x_psi_psi,
            penaltiessecond_derivative: local_s_psi_psi,
            implicit_operator: local_implicit_psi_psi_unused,
        } = derivative_bundle.second;
        assert!(local_implicit_psi_unused.is_none());
        assert!(local_implicit_psi_psi_unused.is_none());
        assert_spatial_derivative_width(
            "Duchon first log-kappa",
            &local_x_psi,
            local_implicit.as_ref(),
            smooth_term.coeff_range.len(),
        );
        assert_spatial_derivative_width(
            "Duchon second log-kappa",
            &local_x_psi_psi,
            local_implicit.as_ref(),
            smooth_term.coeff_range.len(),
        );
        assert!(!local_s_psi.is_empty());
        assert_eq!(local_s_psi.len(), local_s_psi_psi.len());
        assert!(local_s_psi.iter().all(|s| {
            s.nrows() == smooth_term.coeff_range.len() && s.ncols() == smooth_term.coeff_range.len()
        }));
        assert!(local_s_psi_psi.iter().all(|s| {
            s.nrows() == smooth_term.coeff_range.len() && s.ncols() == smooth_term.coeff_range.len()
        }));

        let derivative =
            try_build_spatial_term_log_kappa_derivative(data.view(), &frozenspec, &design, 0);
        assert!(
            derivative.is_ok(),
            "exact Duchon log-kappa derivative should use only feature_cols; got {derivative:?}"
        );
        let derivative = derivative.expect("derivative call should succeed");
        assert!(
            derivative.is_some(),
            "Duchon term should expose an exact derivative"
        );
    }

    #[test]
    fn spatial_length_scale_optimization_monotone_improves_or_keeps_score_for_matern() {
        let n = 60usize;
        let d = 2usize;
        let mut data = Array2::<f64>::zeros((n, d));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (i as f64 * 0.17).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            y[i] = (3.0 * x0).cos() + 0.35 * x1;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 12 },
                        length_scale: 12.0,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            max_iter: 40,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);

        assert_matern_spatial_length_scale_optimization_monotone(
            data.view(),
            &y,
            &weights,
            &offset,
            &spec,
            &fit_opts,
        );
    }

    #[test]
    fn spatial_length_scale_optimization_runs_binomial_logit_matern_with_exact_laml_derivatives() {
        assert!(file!().ends_with(".rs"));
        let n = 80usize;
        let d = 2usize;
        let mut data = Array2::<f64>::zeros((n, d));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (i as f64 * 0.19).cos();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            let eta = -0.15 + 0.45 * x0 - 0.25 * x1 + 0.10 * (6.0 * x0).sin();
            let mu = 1.0 / (1.0 + (-eta).exp());
            let u = (((i * 37 + 17) % 101) as f64 + 0.5) / 101.0;
            y[i] = if u < mu { 1.0 } else { 0.0 };
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 10 },
                        length_scale: 1.8,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            max_iter: 60,
            ..FitOptions::default()
        };
        let weights = Array1::ones(n);
        let offset = Array1::zeros(n);

        fit_term_collectionwith_spatial_length_scale_optimization(
            data.view(),
            y,
            weights,
            offset,
            &spec,
            LikelihoodSpec::binomial_logit(),
            &fit_opts,
            &SpatialLengthScaleOptimizationOptions {
                max_outer_iter: 2,
                rel_tol: 1e-5,
                pilot_subsample_threshold: 0,
                ..SpatialLengthScaleOptimizationOptions::default()
            },
        )
        .expect("standard binomial-logit spatial kappa optimization should use exact non-TK LAML derivatives");
    }

    #[test]
    fn spatial_kappa_result_requires_exact_availability() {
        let err = require_successful_spatial_optimization_result::<()>(0.0, Ok(None))
            .expect_err("missing exact spatial result must be surfaced");
        let msg = err.to_string();
        assert!(msg.contains("unavailable"), "unexpected error: {msg}");
    }

    #[test]
    fn spatial_kappa_result_rejects_worse_exact_score() {
        let err = require_successful_spatial_optimization_result(1.0, Ok(Some(((), 1.5))))
            .expect_err("worse exact spatial result must be rejected");
        let msg = err.to_string();
        assert!(
            msg.contains("made REML score worse"),
            "unexpected error: {msg}"
        );
        assert!(msg.contains("1.000000e0"), "unexpected error: {msg}");
        assert!(msg.contains("1.500000e0"), "unexpected error: {msg}");
    }

    #[test]
    fn spatial_kappa_result_surfaces_optimizer_failure() {
        let err = require_successful_spatial_optimization_result::<()>(
            0.0,
            Err(EstimationError::InvalidInput("boom".to_string())),
        )
        .expect_err("exact spatial optimizer failure must be surfaced");
        let msg = err.to_string();
        assert!(
            msg.contains("spatial kappa optimization failed"),
            "unexpected error: {msg}"
        );
        assert!(msg.contains("boom"), "unexpected error: {msg}");
    }

    #[test]
    fn duchon_terms_participate_in_kappa_optimization() {
        let data = array![
            [0.0, 0.1, 0.2],
            [0.2, 0.0, 0.4],
            [0.4, 0.3, 0.1],
            [0.6, 0.5, 0.7],
            [0.8, 0.7, 0.3],
            [1.0, 0.9, 0.8],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: Some(0.9),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        assert_eq!(spatial_length_scale_term_indices(&spec), vec![0]);

        let fit_opts = FitOptions {
            max_iter: 40,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };
        let y = Array1::linspace(0.0, 1.0, data.nrows());
        let weights = Array1::ones(data.nrows());
        let offset = Array1::zeros(data.nrows());

        let design = build_term_collection_design(data.view(), &spec)
            .expect("baseline Duchon design should build");
        let frozenspec = freeze_term_collection_from_design(&spec, &design)
            .expect("freezing Duchon centers from design should succeed");
        let derivative =
            try_build_spatial_term_log_kappa_derivative(data.view(), &frozenspec, &design, 0);
        assert!(
            derivative
                .expect("Duchon exact derivative call should succeed")
                .is_some(),
            "Duchon term should expose an exact derivative"
        );

        let optimized = fit_term_collectionwith_spatial_length_scale_optimization(
            data.view(),
            y,
            weights,
            offset,
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &fit_opts,
            &SpatialLengthScaleOptimizationOptions::default(),
        )
        .expect("Duchon fit should use exact κ optimization");

        let optimized_ls = match &optimized.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
            _ => panic!("expected Duchon term"),
        };
        assert!(optimized_ls.is_some());
        match &optimized.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::Duchon { spec, .. } => {
                assert!(matches!(
                    spec.center_strategy,
                    CenterStrategy::UserProvided(_)
                ));
                assert!(matches!(
                    spec.identifiability,
                    SpatialIdentifiability::FrozenTransform { .. }
                ));
            }
            _ => panic!("expected Duchon term"),
        }
    }

    #[test]
    fn pure_duchon_scale_dimensions_seed_geometry_but_enroll_no_hyper_axis() {
        let mut spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pure_duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 4 },
                        length_scale: None,
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        crate::term_builder::enable_scale_dimensions(&mut spec);
        // Duchon anisotropy η is a fixed, geometry-derived basis parameter,
        // never a REML hyper axis (see `spatial_term_supports_hyper_optimization`).
        // `scale_dims` seeds the per-axis metric on the spec, but a pure Duchon
        // (no explicit κ) still contributes no outer length-scale/ψ optimization
        // axis — "standardize the geometry, then learn the smoothness."
        assert!(
            spatial_length_scale_term_indices(&spec).is_empty(),
            "pure Duchon must enroll no outer hyper axis even with scale_dims on"
        );
        match &spec.smooth_terms[0].basis {
            SmoothBasisSpec::Duchon { spec, .. } => {
                assert_eq!(spec.length_scale, None);
                assert_eq!(spec.aniso_log_scales.as_deref(), Some(&[0.0, 0.0][..]));
            }
            _ => panic!("expected Duchon term"),
        }
    }

    #[test]
    fn thin_plate_terms_anchor_length_scale_and_enroll_no_kappa_axis() {
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "thin_plate".to_string(),
                basis: SmoothBasisSpec::ThinPlate {
                    feature_cols: vec![0, 1],
                    spec: ThinPlateBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 6 },
                        length_scale: 0.75,
                        double_penalty: false,
                        identifiability: SpatialIdentifiability::default(),
                        radial_reparam: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        assert!(
            spatial_length_scale_term_indices(&spec).is_empty(),
            "penalized thin-plate regression splines must not contribute a redundant isotropic kappa axis"
        );
        assert!(
            all_spatial_terms_kappa_fixed(&spec),
            "with no TPS kappa axis, all spatial terms are effectively fixed-geometry"
        );
    }

    #[test]
    fn pure_duchon_from_length_scales_aniso_is_isotropic_single_psi() {
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pure_duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1, 2],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(array![
                            [0.0, 0.0, 0.0],
                            [1.0, 0.0, 0.0],
                            [0.0, 1.0, 0.0],
                            [0.0, 0.0, 1.0],
                        ]),
                        length_scale: None,
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::None,
                        aniso_log_scales: Some(vec![0.7, 0.2, 0.1]),
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let opts = SpatialLengthScaleOptimizationOptions::default();
        let coords = SpatialLogKappaCoords::from_length_scales_aniso(&spec, &[0], &opts);

        // Duchon anisotropy η is a fixed, geometry-derived basis parameter, not
        // a REML hyper axis. Even with multi-axis `aniso_log_scales`,
        // `from_length_scales_aniso` enrolls a Duchon term as a single isotropic
        // ψ̄ slot — matching the lone `SpatialPsiDerivative` the hyper_dirs
        // builder emits — via the `spatial_term_uses_per_axis_psi` single source
        // of truth. A pure Duchon carries no explicit κ, so ψ̄ defaults to
        // −ln(min_length_scale).
        assert_eq!(coords.dims_per_term(), &[1]);
        assert_eq!(coords.as_array().len(), 1);
        let expected_psi = -opts.min_length_scale.ln();
        assert!((coords.as_array()[0] - expected_psi).abs() <= 1e-12);
    }

    #[test]
    fn explicit_duchon_aniso_length_scale_is_locked_kappa() {
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon_fixed_geometry".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1, 2],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(array![
                            [0.0, 0.0, 0.0],
                            [1.0, 0.0, 0.0],
                            [0.0, 1.0, 0.0],
                            [0.0, 0.0, 1.0],
                        ]),
                        length_scale: Some(1.0),
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::None,
                        aniso_log_scales: Some(vec![0.7, 0.2, 0.1]),
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        assert!(
            spatial_term_has_locked_kappa(&spec, 0),
            "Duchon anisotropy is fixed geometry and must not force ψ optimization"
        );
        assert!(
            all_spatial_terms_kappa_fixed(&spec),
            "a Duchon term with explicit length_scale and fixed anisotropy has no REML κ/ψ axis"
        );
    }

    #[test]
    fn from_length_scales_aniso_keeps_nonaniso_spatial_terms_scalar() {
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![
                SmoothTermSpec {
                    name: "matern_aniso".to_string(),
                    basis: SmoothBasisSpec::Matern {
                        feature_cols: vec![0, 1],
                        spec: MaternBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::UserProvided(array![
                                [0.0, 0.0],
                                [1.0, 0.0],
                                [0.0, 1.0],
                            ]),
                            length_scale: 0.5,
                            nu: MaternNu::FiveHalves,
                            include_intercept: false,
                            double_penalty: false,
                            identifiability: MaternIdentifiability::None,
                            aniso_log_scales: Some(vec![0.3, -0.3]),
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
                SmoothTermSpec {
                    name: "matern_iso".to_string(),
                    basis: SmoothBasisSpec::Matern {
                        feature_cols: vec![0, 1],
                        spec: MaternBasisSpec {
                            periodic: None,
                            center_strategy: CenterStrategy::UserProvided(array![
                                [0.0, 0.0],
                                [1.0, 0.0],
                                [0.0, 1.0],
                            ]),
                            length_scale: 0.25,
                            nu: MaternNu::ThreeHalves,
                            include_intercept: false,
                            double_penalty: false,
                            identifiability: MaternIdentifiability::None,
                            aniso_log_scales: None,
                        },
                        input_scales: None,
                    },
                    shape: ShapeConstraint::None,
                    joint_null_rotation: None,
                },
            ],
        };

        let term_indices = [0usize, 1usize];
        let coords = SpatialLogKappaCoords::from_length_scales_aniso(
            &spec,
            &term_indices,
            &SpatialLengthScaleOptimizationOptions::default(),
        );

        assert_eq!(spatial_dims_per_term(&spec, &term_indices), vec![2, 1]);
        assert_eq!(coords.dims_per_term(), &[2, 1]);
        let expected = [-0.5_f64.ln() + 0.3, -0.5_f64.ln() - 0.3, -0.25_f64.ln()];
        for (got, want) in coords.as_array().iter().zip(expected.iter()) {
            assert!((got - want).abs() <= 1e-12);
        }
    }

    #[test]
    fn aniso_bounds_clamp_preserves_in_range_global_length_scale_and_eta() {
        let data = array![[0.0, 0.0], [1.0, 0.2], [0.1, 1.0], [1.1, 1.2]];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern_aniso".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(array![
                            [0.0, 0.0],
                            [1.0, 0.0],
                            [0.0, 1.0],
                            [1.0, 1.0],
                        ]),
                        length_scale: 1.0,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::None,
                        aniso_log_scales: Some(vec![3.0, -3.0]),
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let options = SpatialLengthScaleOptimizationOptions {
            max_outer_iter: 1,
            rel_tol: 1e-6,
            min_length_scale: (-2.0_f64).exp(),
            max_length_scale: 1.0_f64.exp(),
            pilot_subsample_threshold: 0,
            ..SpatialLengthScaleOptimizationOptions::default()
        };
        let spatial_terms = vec![0];
        let dims_per_term = spatial_dims_per_term(&spec, &spatial_terms);
        let seed = SpatialLogKappaCoords::from_length_scales_aniso(&spec, &spatial_terms, &options);
        let lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
            data.view(),
            &spec,
            &spatial_terms,
            &dims_per_term,
            &options,
        );
        let upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
            data.view(),
            &spec,
            &spatial_terms,
            &dims_per_term,
            &options,
        );

        let projected = seed.clone().clamp_to_bounds(&lower, &upper);
        assert_eq!(projected.as_array(), seed.as_array());

        let updated = projected
            .apply_tospec(&spec, &spatial_terms)
            .expect("aniso projection should decode");
        match &updated.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => {
                assert!((spec.length_scale - 1.0).abs() <= 1e-12);
                let eta = spec
                    .aniso_log_scales
                    .as_ref()
                    .expect("anisotropy should be preserved");
                assert!((eta[0] - 3.0).abs() <= 1e-12);
                assert!((eta[1] + 3.0).abs() <= 1e-12);
            }
            _ => panic!("expected Matérn term"),
        }
    }

    #[test]
    fn pure_duchon_aniso_fit_optimizes_without_introducing_hybrid_scale() {
        let data = array![
            [0.0, 0.1, 0.2],
            [0.2, 0.0, 0.4],
            [0.4, 0.3, 0.1],
            [0.6, 0.5, 0.7],
            [0.8, 0.7, 0.3],
            [1.0, 0.9, 0.8],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pure_duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1, 2],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 5 },
                        length_scale: None,
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: Some(vec![0.0, 0.0, 0.0]),
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let fit_opts = FitOptions {
            max_iter: 40,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let optimized = fit_term_collectionwith_spatial_length_scale_optimization(
            data.view(),
            Array1::linspace(0.0, 1.0, data.nrows()),
            Array1::ones(data.nrows()),
            Array1::zeros(data.nrows()),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &fit_opts,
            &SpatialLengthScaleOptimizationOptions::default(),
        )
        .expect("pure Duchon anisotropic fit should optimize");

        match &optimized.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::Duchon { spec, .. } => {
                assert_eq!(spec.length_scale, None);
                assert!(
                    spec.aniso_log_scales.is_some(),
                    "pure Duchon anisotropy should remain enabled"
                );
            }
            _ => panic!("expected Duchon term"),
        }
    }

    #[test]
    fn spatial_anisotropy_pilot_initializer_seeds_geometry_without_fit() {
        let data = Array2::from_shape_fn((32, 2), |(i, j)| {
            if j == 0 {
                i as f64 / 31.0
            } else {
                ((i % 8) as f64) * 0.03
            }
        });
        let mut spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pc_matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::UserProvided(array![
                            [0.0, 0.0],
                            [1.0, 0.0],
                            [0.0, 0.05],
                            [1.0, 0.05],
                        ]),
                        length_scale: 1.0,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::None,
                        aniso_log_scales: Some(vec![0.0, 0.0]),
                    },
                    input_scales: Some(vec![1.0, 1.0]),
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let spatial_terms = spatial_length_scale_term_indices(&spec);
        let updated = apply_spatial_anisotropy_pilot_initializer(
            data.view(),
            &mut spec,
            &spatial_terms,
            8,
            &SpatialLengthScaleOptimizationOptions::default(),
        );

        assert_eq!(updated, 1);
        match &spec.smooth_terms[0].basis {
            SmoothBasisSpec::Matern { spec, .. } => {
                let eta = spec
                    .aniso_log_scales
                    .as_ref()
                    .expect("pilot initializer should preserve anisotropy");
                assert_eq!(eta.len(), 2);
                assert!((eta[0] + eta[1]).abs() <= 1e-12);
                assert!(
                    eta.iter().any(|value| value.abs() > 1e-6),
                    "pilot geometry should seed nonzero axis contrast"
                );
                assert!(spec.length_scale.is_finite() && spec.length_scale > 0.0);
            }
            _ => panic!("expected Matern term"),
        }
    }

    #[test]
    fn pure_duchon_aniso_penalties_stay_symmetric_through_freeze_and_cache() {
        fn max_asymmetry(matrix: &Array2<f64>) -> f64 {
            let n = matrix.nrows().min(matrix.ncols());
            let mut max_asym = 0.0_f64;
            for i in 0..n {
                for j in 0..i {
                    max_asym = max_asym.max((matrix[[i, j]] - matrix[[j, i]]).abs());
                }
            }
            max_asym
        }

        fn assert_design_penalties_symmetric(label: &str, design: &TermCollectionDesign) {
            for (penalty_idx, penalty) in design.penalties.iter().enumerate() {
                let max_asym = max_asymmetry(&penalty.local);
                assert!(
                    max_asym <= 1e-10,
                    "{label} penalty {penalty_idx} asymmetry too large: {max_asym:.3e}"
                );
            }
        }

        fn assert_reparam_penalty_symmetric(label: &str, design: &TermCollectionDesign) {
            let p_total = design.design.ncols();
            let penalty_specs = design
                .penalties
                .iter()
                .map(|penalty| crate::estimate::PenaltySpec::Dense(penalty.to_global(p_total)))
                .collect::<Vec<_>>();
            let (canonical_penalties, _) = crate::construction::canonicalize_penalty_specs(
                &penalty_specs,
                &design.nullspace_dims,
                p_total,
                label,
            )
            .expect("canonicalize penalties");
            let invariant = crate::construction::precompute_reparam_invariant_from_canonical(
                &canonical_penalties,
                p_total,
            )
            .expect("reparam invariant");
            let lambdas = vec![1.0; canonical_penalties.len()];
            let reparam = crate::construction::stable_reparameterizationwith_invariant(
                &canonical_penalties,
                &lambdas,
                p_total,
                &invariant,
                None,
            )
            .expect("stable reparameterization");
            let max_asym = max_asymmetry(&reparam.s_transformed);
            assert!(
                max_asym <= 1e-10,
                "{label} transformed penalty asymmetry too large: {max_asym:.3e}"
            );
        }

        let data = array![
            [0.0, 0.1, 0.2],
            [0.2, 0.0, 0.4],
            [0.4, 0.3, 0.1],
            [0.6, 0.5, 0.7],
            [0.8, 0.7, 0.3],
            [1.0, 0.9, 0.8],
        ];
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pure_duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1, 2],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 5 },
                        length_scale: None,
                        power: 1.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: Some(vec![0.0, 0.0, 0.0]),
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let base_design = build_term_collection_design(data.view(), &spec).expect("base design");
        assert_design_penalties_symmetric("base", &base_design);
        assert_reparam_penalty_symmetric("base", &base_design);

        let frozen = freeze_term_collection_from_design(&spec, &base_design).expect("freeze spec");
        let frozen_design =
            build_term_collection_design(data.view(), &frozen).expect("frozen rebuild");
        assert_design_penalties_symmetric("frozen", &frozen_design);
        assert_reparam_penalty_symmetric("frozen", &frozen_design);

        // Design B: a pure Duchon enrolls no outer ψ axis (η is a fixed,
        // geometry-derived basis parameter), so the single-block exact-joint
        // cache for this term is ρ-only. The penalties must stay symmetric
        // through that cache exactly as they do through the base build and the
        // freeze/rebuild above.
        let spatial_terms = spatial_length_scale_term_indices(&frozen);
        assert!(
            spatial_terms.is_empty(),
            "pure Duchon enrolls no outer κ/ψ axis"
        );
        let rho_dim = frozen_design.penalties.len();
        let dims_per_term = spatial_dims_per_term(&frozen, &spatial_terms);
        let theta = Array1::<f64>::zeros(rho_dim);

        let mut cache = SingleBlockExactJointDesignCache::new(
            data.view(),
            frozen.clone(),
            frozen_design.clone(),
            spatial_terms,
            rho_dim,
            dims_per_term,
        )
        .expect("single-block cache");
        cache.ensure_theta(&theta).expect("updated theta");
        assert_design_penalties_symmetric("cache", cache.design());
        assert_reparam_penalty_symmetric("cache", cache.design());
    }

    #[test]
    fn single_block_no_spatial_fast_path_returns_fully_frozen_spec() {
        let n = 48usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (i % 4) as f64;
            y[i] = 0.5 + 1.5 * t;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![RandomEffectTermSpec {
                name: "grp".to_string(),
                feature_col: 1,
                drop_first_level: false,
                penalized: true,
                frozen_levels: None,
            }],
            smooth_terms: vec![SmoothTermSpec {
                name: "ps".to_string(),
                basis: SmoothBasisSpec::BSpline1D {
                    feature_col: 0,
                    spec: BSplineBasisSpec {
                        degree: 3,
                        penalty_order: 2,
                        knotspec: BSplineKnotSpec::Generate {
                            data_range: (0.0, 1.0),
                            num_internal_knots: 4,
                        },
                        double_penalty: true,
                        identifiability: BSplineIdentifiability::None,
                        boundary: OneDimensionalBoundary::Open,
                        boundary_conditions: BSplineBoundaryConditions::default(),
                    },
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit_opts = FitOptions {
            max_iter: 40,
            penalty_shrinkage_floor: None,
            ..FitOptions::default()
        };

        let fitted = fit_term_collectionwith_spatial_length_scale_optimization(
            data.view(),
            y,
            Array1::ones(n),
            Array1::zeros(n),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &fit_opts,
            &SpatialLengthScaleOptimizationOptions::default(),
        )
        .expect("single-block no-spatial fit should succeed");

        fitted
            .resolvedspec
            .validate_frozen("resolvedspec")
            .expect("single-block no-spatial fast path should fully freeze specs");
        match &fitted.resolvedspec.smooth_terms[0].basis {
            SmoothBasisSpec::BSpline1D { spec, .. } => {
                assert!(matches!(spec.knotspec, BSplineKnotSpec::Provided(_)));
            }
            _ => panic!("expected P-spline term"),
        }
        assert!(
            fitted.resolvedspec.random_effect_terms[0]
                .frozen_levels
                .is_some(),
            "random-effect levels should be frozen in single-block no-spatial fast path"
        );
    }

    #[test]
    fn exact_joint_two_block_spatial_length_scale_freezes_duchon_centers() {
        let n = 40usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (i as f64 * 0.19).cos();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
        }

        let duchon_term = |name: &str, length_scale: f64| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: vec![0, 1],
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                    length_scale: Some(length_scale),
                    power: 3.0,
                    nullspace_order: DuchonNullspaceOrder::Linear,
                    identifiability: SpatialIdentifiability::default(),
                    aniso_log_scales: None,
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let meanspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![duchon_term("mean_duchon", 0.8)],
        };
        let noisespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![duchon_term("noise_duchon", 1.1)],
        };

        let solved = run_two_block_exact_joint_optimize(
            data.view(),
            &meanspec,
            &noisespec,
            "exact joint two-block spatial length-scale optimization should succeed",
        );

        for resolved in [&solved.resolved_specs[0], &solved.resolved_specs[1]] {
            match &resolved.smooth_terms[0].basis {
                SmoothBasisSpec::Duchon { spec, .. } => {
                    assert!(matches!(
                        spec.center_strategy,
                        CenterStrategy::UserProvided(_)
                    ));
                    assert!(matches!(
                        spec.identifiability,
                        SpatialIdentifiability::FrozenTransform { .. }
                    ));
                }
                _ => panic!("expected Duchon term"),
            }
        }
    }

    #[test]
    fn joint_build_and_cache_rebuild_frozen_pure_duchon_blocks() {
        let n = 72usize;
        let d = 5usize;
        let mut data = Array2::<f64>::zeros((n, d));
        for i in 0..n {
            let t = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = t;
            data[[i, 1]] = (0.17 * i as f64).sin();
            data[[i, 2]] = (0.11 * i as f64).cos();
            data[[i, 3]] = ((i % 7) as f64) / 6.0;
            data[[i, 4]] = t * (1.0 - t);
        }

        let pure_duchon_term = |name: &str| SmoothTermSpec {
            name: name.to_string(),
            basis: SmoothBasisSpec::Duchon {
                feature_cols: (0..d).collect(),
                spec: DuchonBasisSpec {
                    periodic: None,
                    center_strategy: CenterStrategy::FarthestPoint { num_centers: 24 },
                    length_scale: None,
                    power: 2.0,
                    nullspace_order: DuchonNullspaceOrder::Linear,
                    identifiability: SpatialIdentifiability::default(),
                    aniso_log_scales: Some(vec![0.0; d]),
                    operator_penalties: DuchonOperatorPenaltySpec::default(),
                    boundary: OneDimensionalBoundary::Open,
                },
                input_scales: None,
            },
            shape: ShapeConstraint::None,
            joint_null_rotation: None,
        };

        let meanspec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![pure_duchon_term("mean_pure_duchon")],
        };
        let noisespec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![pure_duchon_term("noise_pure_duchon")],
        };

        let (boot_designs, frozen_specs) = build_term_collection_designs_and_freeze_joint(
            data.view(),
            &[meanspec.clone(), noisespec.clone()],
        )
        .expect("initial joint pure Duchon build");
        assert_eq!(boot_designs.len(), 2);
        assert_eq!(frozen_specs.len(), 2);
        assert_eq!(boot_designs[0].smooth.terms[0].coeff_range.len(), 23);
        assert_eq!(boot_designs[1].smooth.terms[0].coeff_range.len(), 23);

        let (rebuilt_designs, refrozen_specs) =
            build_term_collection_designs_and_freeze_joint(data.view(), &frozen_specs)
                .expect("rebuilding frozen joint pure Duchon specs should succeed");
        assert_eq!(rebuilt_designs.len(), 2);
        assert_eq!(refrozen_specs.len(), 2);

        for idx in 0..2 {
            let direct = build_term_collection_design(data.view(), &frozen_specs[idx])
                .expect("direct frozen pure Duchon rebuild");
            assert_term_collection_designs_match(
                &rebuilt_designs[idx],
                &direct,
                if idx == 0 {
                    "mean pure Duchon frozen rebuild"
                } else {
                    "noise pure Duchon frozen rebuild"
                },
            );
            assert_eq!(rebuilt_designs[idx].smooth.terms[0].coeff_range.len(), 23);
            match &refrozen_specs[idx].smooth_terms[0].basis {
                SmoothBasisSpec::Duchon { spec, .. } => {
                    assert!(matches!(
                        spec.identifiability,
                        SpatialIdentifiability::FrozenTransform { .. }
                    ));
                }
                _ => panic!("expected Duchon term"),
            }
        }

        let kappa_options = SpatialLengthScaleOptimizationOptions {
            max_outer_iter: 1,
            rel_tol: 1e-6,
            pilot_subsample_threshold: 0,
            ..SpatialLengthScaleOptimizationOptions::default()
        };
        let joint_setup =
            two_block_exact_joint_hyper_setup(&frozen_specs[0], &frozen_specs[1], &kappa_options);
        // Design B: Duchon anisotropy η is a fixed, geometry-derived basis
        // parameter, never a REML axis, so two pure-Duchon blocks contribute no
        // outer log-κ axis — the joint outer vector is ρ-only.
        assert_eq!(joint_setup.log_kappa_dim(), 0);

        let mean_term_indices = spatial_length_scale_term_indices(&frozen_specs[0]);
        let noise_term_indices = spatial_length_scale_term_indices(&frozen_specs[1]);
        assert!(
            mean_term_indices.is_empty() && noise_term_indices.is_empty(),
            "pure Duchon blocks enroll no outer κ/ψ axis"
        );
        let mut cache = ExactJointDesignCache::new(
            data.view(),
            vec![
                (
                    frozen_specs[0].clone(),
                    rebuilt_designs[0].clone(),
                    mean_term_indices.clone(),
                ),
                (
                    frozen_specs[1].clone(),
                    rebuilt_designs[1].clone(),
                    noise_term_indices.clone(),
                ),
            ],
            joint_setup.rho_dim(),
            joint_setup.log_kappa_dims_per_term(),
        )
        .expect("pure Duchon exact-joint cache");

        // With no κ axis the joint outer vector is ρ-only; realizing the cache
        // at θ₀ must reproduce the directly-rebuilt frozen designs, since there
        // is no per-axis log-κ update to apply.
        let theta0 = joint_setup.theta0();
        assert_eq!(theta0.len(), joint_setup.rho_dim());
        cache
            .ensure_theta(&theta0)
            .expect("pure Duchon cache theta update");
        let cache_designs = cache.designs();
        assert_term_collection_designs_match(
            cache_designs[0],
            &rebuilt_designs[0],
            "mean pure Duchon cache",
        );
        assert_term_collection_designs_match(
            cache_designs[1],
            &rebuilt_designs[1],
            "noise pure Duchon cache",
        );
    }

    #[test]
    fn bounded_linear_gaussian_fit_respects_interval() {
        let n = 64usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = (i as f64) / ((n - 1) as f64);
            let x = -1.0 + 2.0 * t;
            // z must be linearly independent of {1, x}: a ramp z = (x+1)/2
            // is exactly collinear with the intercept and x, so the bounded
            // column's true ∂η/∂θ Jacobian is genuinely rank-deficient and the
            // identifiability audit (correctly) refuses the fit. A 2-cycle
            // sinusoid is orthogonal to the constant and the linear ramp.
            let z = (2.0 * std::f64::consts::PI * 2.0 * t).sin();
            data[[i, 0]] = x;
            data[[i, 1]] = z;
            y[i] = 0.25 + 0.8 * x + 0.05 * z;
        }
        let spec = TermCollectionSpec {
            linear_terms: vec![
                LinearTermSpec {
                    name: "x".to_string(),
                    feature_col: 0,
                    feature_cols: vec![0],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Bounded {
                        min: 0.0,
                        max: 0.5,
                        prior: BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 },
                    },
                    coefficient_min: None,
                    coefficient_max: None,
                },
                LinearTermSpec {
                    name: "z".to_string(),
                    feature_col: 1,
                    feature_cols: vec![1],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                },
            ],
            random_effect_terms: vec![],
            smooth_terms: vec![],
        };

        let fitted = fit_term_collectionwith_spatial_length_scale_optimization(
            data.view(),
            y,
            Array1::ones(n),
            Array1::zeros(n),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 40,
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
            &SpatialLengthScaleOptimizationOptions {
                enabled: false,
                ..SpatialLengthScaleOptimizationOptions::default()
            },
        )
        .expect("bounded gaussian fit");

        let bounded_idx = fitted.design.linear_ranges[0].1.start;
        let estimate = fitted.fit.beta[bounded_idx];
        assert!(
            (0.0..=0.5).contains(&estimate),
            "bounded coefficient escaped interval: {estimate}"
        );
        assert!(
            estimate > 0.1,
            "bounded coefficient should move into the positive interior, got {estimate}"
        );
    }

    #[test]
    fn bounded_fit_geometry_precision_is_on_user_scale() {
        use crate::faer_ndarray::FaerCholesky;

        let n = 72usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = (i as f64) / ((n - 1) as f64);
            let x = -1.0 + 2.0 * t;
            let z = (4.0 * std::f64::consts::PI * t).sin();
            data[[i, 0]] = x;
            data[[i, 1]] = z;
            y[i] = 0.2 + 0.35 * x - 0.15 * z;
        }
        let spec = TermCollectionSpec {
            linear_terms: vec![
                LinearTermSpec {
                    name: "x".to_string(),
                    feature_col: 0,
                    feature_cols: vec![0],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Bounded {
                        min: -0.5,
                        max: 0.5,
                        prior: BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 },
                    },
                    coefficient_min: None,
                    coefficient_max: None,
                },
                LinearTermSpec {
                    name: "z".to_string(),
                    feature_col: 1,
                    feature_cols: vec![1],
                    double_penalty: false,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                },
            ],
            random_effect_terms: vec![],
            smooth_terms: vec![],
        };

        let fitted = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 40,
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("bounded gaussian fit");
        let precision = &fitted
            .fit
            .geometry
            .as_ref()
            .expect("bounded fit geometry")
            .penalized_hessian;
        let covariance = fitted
            .fit
            .beta_covariance()
            .expect("bounded user covariance");
        let chol = precision
            .cholesky(faer::Side::Lower)
            .expect("bounded user precision cholesky");
        let solved = chol.solve_mat(&Array2::eye(covariance.nrows()));
        for i in 0..solved.nrows() {
            for j in 0..solved.ncols() {
                assert!(
                    (solved[[i, j]] - covariance[[i, j]]).abs() < 1e-5,
                    "user-scale precision/covariance mismatch at ({i},{j}): inverse {}, covariance {}",
                    solved[[i, j]],
                    covariance[[i, j]]
                );
            }
        }
    }

    #[test]
    fn term_collection_design_emits_linear_coefficient_constraints() {
        let data = array![[0.0], [1.0], [2.0], [3.0]];
        let spec = TermCollectionSpec {
            linear_terms: vec![LinearTermSpec {
                name: "x".to_string(),
                feature_col: 0,
                feature_cols: vec![0],
                double_penalty: false,
                coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                coefficient_min: Some(0.0),
                coefficient_max: Some(1.0),
            }],
            random_effect_terms: vec![],
            smooth_terms: vec![],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let constraints = design.linear_constraints.expect("constraints");
        assert_eq!(constraints.a.ncols(), design.design.ncols());
        assert_eq!(constraints.a.nrows(), 2);
        let linear_idx = design.linear_ranges[0].1.start;
        assert_eq!(constraints.a[[0, linear_idx]], 1.0);
        assert_eq!(constraints.b[0], 0.0);
        assert_eq!(constraints.a[[1, linear_idx]], -1.0);
        assert_eq!(constraints.b[1], -1.0);
    }

    #[test]
    fn linear_termspec_defaults_to_unpenalizedwhen_field_is_omitted() {
        // Parametric/linear terms are unpenalized by default — mature tools
        // (mgcv/glm/survreg/VGAM) leave parametric terms unpenalized and gam
        // matches that, reporting the MLE rather than a ridge-shrunk estimate.
        // See `default_linear_term_double_penalty`.
        let json = r#"{"name":"x","feature_col":0}"#;
        let term: LinearTermSpec = serde_json::from_str(json).expect("deserialize linear term");
        assert!(!term.double_penalty);
        assert!(matches!(
            term.coefficient_geometry,
            LinearCoefficientGeometry::Unconstrained
        ));
    }

    #[test]
    fn linear_double_penalties_share_one_globalridge_block() {
        let data = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
        let spec = TermCollectionSpec {
            linear_terms: vec![
                LinearTermSpec {
                    name: "x1".to_string(),
                    feature_col: 0,
                    feature_cols: vec![0],
                    double_penalty: true,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                },
                LinearTermSpec {
                    name: "x2".to_string(),
                    feature_col: 1,
                    feature_cols: vec![1],
                    double_penalty: true,
                    coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
                    coefficient_min: None,
                    coefficient_max: None,
                },
            ],
            random_effect_terms: vec![],
            smooth_terms: vec![],
        };
        let design = build_term_collection_design(data.view(), &spec).expect("design");
        assert_eq!(design.penalties.len(), 1);
        assert_eq!(design.penaltyinfo.len(), 1);
        assert_eq!(design.penaltyinfo[0].termname.as_deref(), Some("linear"));
        assert_eq!(design.penaltyinfo[0].penalty.effective_rank, 2);
        let x1 = design.linear_ranges[0].1.start;
        let x2 = design.linear_ranges[1].1.start;
        let bp = &design.penalties[0];
        let x1_local = x1 - bp.col_range.start;
        let x2_local = x2 - bp.col_range.start;
        assert_eq!(bp.local[[x1_local, x1_local]], 1.0);
        assert_eq!(bp.local[[x2_local, x2_local]], 1.0);
    }

    #[test]
    fn bounded_uniform_prior_matches_beta_one_one_terms() {
        let theta = 0.7;
        let uniform = bounded_prior_terms(theta, &BoundedCoefficientPriorSpec::Uniform);
        let beta11 =
            bounded_prior_terms(theta, &BoundedCoefficientPriorSpec::Beta { a: 1.0, b: 1.0 });
        assert!((uniform.0 - beta11.0).abs() < 1e-12);
        assert!((uniform.1 - beta11.1).abs() < 1e-12);
        assert!((uniform.2 - beta11.2).abs() < 1e-12);
        assert!((uniform.3 - beta11.3).abs() < 1e-12);
    }

    #[test]
    fn boundednone_prior_has_no_extra_latentobjective_terms() {
        let theta = 0.7;
        let none = bounded_prior_terms(theta, &BoundedCoefficientPriorSpec::None);
        assert_eq!(none, (0.0, 0.0, 0.0, 0.0));

        let uniform = bounded_prior_terms(theta, &BoundedCoefficientPriorSpec::Uniform);
        assert!(uniform.0.is_finite());
        assert!(uniform.0 < 0.0);
        assert!(uniform.1.abs() > 1e-6);
        assert!(uniform.2 > 0.0);
        assert!(uniform.3.is_finite());
    }

    #[test]
    fn exact_bounded_edf_matches_trace_formula_for_simple_penalty() {
        let penalties = vec![PenaltySpec::Dense(Array2::eye(1))];
        let lambdas = array![0.25];
        let cov = array![[2.0]];
        let (edf_by_block, edf_total) =
            exact_bounded_edf(&penalties, &lambdas, &cov).expect("exact bounded edf");
        assert_eq!(edf_by_block.len(), 1);
        assert!((edf_by_block[0] - 0.5).abs() < 1e-12);
        assert!((edf_total - 0.5).abs() < 1e-12);
    }

    #[test]
    fn bounded_joint_hessian_directional_derivative_matches_finite_difference() {
        let x = array![[0.2, -1.0], [0.8, 0.5], [1.1, 1.2], [1.7, -0.3]];
        let y = array![0.4, 1.0, 1.7, 2.2];
        let weights = Array1::ones(y.len());
        let family = BoundedLinearFamily {
            family: LikelihoodSpec::gaussian_identity(),
            latent_cloglog_state: None,
            mixture_link_state: None,
            sas_link_state: None,
            y: y.clone(),
            weights: weights.clone(),
            design: x.clone(),
            designzeroed: {
                let mut dz = x.clone();
                dz.column_mut(0).fill(0.0);
                dz
            },
            offset: Array1::zeros(y.len()),
            bounded_terms: vec![BoundedLinearTermMeta {
                col_idx: 0,
                min: 0.0,
                max: 1.0,
                prior: BoundedCoefficientPriorSpec::Uniform,
            }],
        };
        let state = vec![ParameterBlockState {
            beta: array![0.4, -0.2],
            eta: Array1::zeros(y.len()),
        }];
        let direction = array![0.3, -0.4];

        let analytic = family
            .exact_newton_joint_hessian_directional_derivative(&state, &direction)
            .expect("analytic derivative")
            .expect("joint derivative");

        let h = 1e-6;
        let plus_state = vec![ParameterBlockState {
            beta: &state[0].beta + &(direction.clone() * h),
            eta: Array1::zeros(y.len()),
        }];
        let minus_state = vec![ParameterBlockState {
            beta: &state[0].beta - &(direction.clone() * h),
            eta: Array1::zeros(y.len()),
        }];
        let plus = family
            .exact_newton_joint_hessian(&plus_state)
            .expect("plus hessian")
            .expect("plus exact hessian");
        let minus = family
            .exact_newton_joint_hessian(&minus_state)
            .expect("minus hessian")
            .expect("minus exact hessian");
        let fd = (plus - minus) / (2.0 * h);

        for i in 0..analytic.nrows() {
            for j in 0..analytic.ncols() {
                assert_eq!(
                    analytic[[i, j]].signum(),
                    fd[[i, j]].signum(),
                    "directional derivative sign mismatch at ({i},{j}): analytic={}, fd={}",
                    analytic[[i, j]],
                    fd[[i, j]]
                );
                assert!(
                    (analytic[[i, j]] - fd[[i, j]]).abs() < 1e-5,
                    "directional derivative mismatch at ({i},{j}): analytic={}, fd={}",
                    analytic[[i, j]],
                    fd[[i, j]]
                );
            }
        }
    }

    #[test]
    fn adaptive_initial_epsilons_use_mean_fallbackwhen_median_is_tiny() {
        let cache = SpatialOperatorRuntimeCache {
            termname: "matern".to_string(),
            feature_cols: vec![0, 1],
            coeff_global_range: 0..2,
            mass_penalty_global_idx: 0,
            tension_penalty_global_idx: 1,
            stiffness_penalty_global_idx: 2,
            d0: array![[5e-10, 0.0], [6e-10, 0.0]],
            d1: array![[1e-10, 0.0], [0.0, 1e-10], [2e-10, 0.0], [0.0, 2e-10]],
            // d2 layout: rows = (point k, axis a, axis b) with row = (k*d + a)*d + b.
            // For P=2 points and d=2 axes that is 2*2*2 = 8 rows.
            d2: array![
                [3e-10, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
                [3e-10, 0.0],
                [4e-10, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
                [4e-10, 0.0],
            ],
            collocation_points: array![[0.0, 0.0], [1.0, 1.0]],
            dimension: 2,
        };
        let beta = array![1.0, 1.0];
        let (eps_0, eps_g, eps_c) =
            compute_initial_epsilons(&beta, &[cache], 1e-8).expect("initial epsilons");
        assert!(eps_0 >= 1e-8);
        assert!(eps_g >= 1e-8);
        assert!(eps_c >= 1e-8);
    }

    #[test]
    fn adaptive_exact_psigradient_symmetrizes_nearly_symmetrichessian() {
        let family = SpatialAdaptiveExactFamily {
            family: LikelihoodSpec::gaussian_identity(),
            latent_cloglog_state: None,
            mixture_link_state: None,
            sas_link_state: None,
            y: Arc::new(array![0.0, 0.0]),
            weights: Arc::new(array![1.0, 1.0]),
            design: Arc::new(array![[1.0, 0.0], [0.0, 1.0]]),
            offset: Arc::new(array![0.0, 0.0]),
            linear_constraints: None,
            runtime_caches: Arc::new(vec![SpatialOperatorRuntimeCache {
                termname: "toy".to_string(),
                feature_cols: vec![0],
                coeff_global_range: 0..2,
                mass_penalty_global_idx: 0,
                tension_penalty_global_idx: 1,
                stiffness_penalty_global_idx: 2,
                d0: array![[1.0, 0.0], [0.0, 1.0]],
                d1: array![[1.0, 0.0], [0.0, 1.0]],
                d2: array![[1.0, 0.0], [0.0, 1.0]],
                collocation_points: array![[0.0], [1.0]],
                dimension: 1,
            }]),
            adaptive_params: vec![SpatialAdaptiveTermHyperParams {
                lambda: [1.0, 1.0, 1.0],
                epsilon: [1.0, 1.0, 1.0],
            }],
            fixed_quadratichessian: Arc::new(array![[0.0, 0.1], [3.0, 0.0]]),
            hyperspecs: Arc::new(build_spatial_adaptive_hyperspecs(1)),
            exact_eval_cache: Arc::new(Mutex::new(None)),
        };
        let spec = ParameterBlockSpec {
            name: "toy".to_string(),
            design: DesignMatrix::Dense(crate::matrix::DenseDesignMatrix::from(array![
                [1.0, 0.0],
                [0.0, 1.0]
            ])),
            offset: array![0.0, 0.0],
            penalties: vec![],
            nullspace_dims: vec![],
            initial_log_lambdas: Array1::zeros(0),
            initial_beta: Some(array![0.0, 0.0]),
            gauge_priority: 100,
            jacobian_callback: None,
            stacked_design: None,
            stacked_offset: None,
        };
        let deriv = CustomFamilyBlockPsiDerivative {
            penalty_index: None,
            x_psi: Array2::zeros((2, 2)),
            s_psi: Array2::zeros((2, 2)),
            s_psi_components: None,
            s_psi_penalty_components: None,
            x_psi_psi: None,
            s_psi_psi: None,
            s_psi_psi_components: None,
            s_psi_psi_penalty_components: None,
            implicit_operator: None,
            implicit_axis: 0,
            implicit_group_id: None,
        };
        let state = vec![ParameterBlockState {
            beta: array![0.0, 0.0],
            eta: array![0.0, 0.0],
        }];

        let gradient = family
            .exact_newton_joint_psi_terms(&state, std::slice::from_ref(&spec), &[vec![deriv]], 0)
            .expect("adaptive joint psi terms should tolerate nearly symmetric Hessian")
            .expect("adaptive joint psi terms should be present")
            .objective_psi;

        assert!(
            gradient.is_finite(),
            "expected finite adaptive joint psi objective term after symmetrization, got {gradient}"
        );
    }

    #[test]
    fn adaptiveweight_clamp_is_applied_in_u_space() {
        let cache = SpatialOperatorRuntimeCache {
            termname: "matern".to_string(),
            feature_cols: vec![0, 1],
            coeff_global_range: 0..2,
            mass_penalty_global_idx: 0,
            tension_penalty_global_idx: 1,
            stiffness_penalty_global_idx: 2,
            d0: array![[0.0, 0.0]],
            d1: array![[0.0, 0.0], [0.0, 0.0]],
            // d2 layout: P=1 collocation point and d=2 axes give 1*2*2 = 4 rows
            // ordered (axis_a, axis_b) = (0,0), (0,1), (1,0), (1,1).
            d2: array![[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
            collocation_points: array![[0.0, 0.0]],
            dimension: 2,
        };
        let beta = array![0.0, 0.0];
        let out = compute_spatial_adaptiveweights_for_beta(
            &beta,
            &[cache],
            1e-8,
            1e-8,
            1e-8,
            1e-8,
            1e2,
            None,
        )
        .expect("adaptive weights");
        assert_eq!(out.len(), 1);
        // Raw u would be 1/eps = 1e8, so clamping to 1e2 yields diagnostics 1/u.
        assert!((out[0].inv_magweight[0] - 1e-2).abs() < 1e-12);
        assert!((out[0].invgradweight[0] - 1e-2).abs() < 1e-12);
        assert!((out[0].inv_lapweight[0] - 1e-2).abs() < 1e-12);
    }

    #[test]
    fn adaptiveweight_inverse_consistencywithout_clamp() {
        let cache = SpatialOperatorRuntimeCache {
            termname: "matern".to_string(),
            feature_cols: vec![0, 1],
            coeff_global_range: 0..2,
            mass_penalty_global_idx: 0,
            tension_penalty_global_idx: 1,
            stiffness_penalty_global_idx: 2,
            d0: array![[1.0, 0.0], [2.0, 0.0]],
            d1: array![[1.0, 0.0], [0.0, 1.0], [2.0, 0.0], [0.0, 2.0]],
            // d2 layout: rows = (point k, axis a, axis b) with row = (k*d + a)*d + b.
            // P=2 and d=2 gives 8 rows; row 0 (point 0, a=b=0) and row 4 (point 1, a=b=0)
            // carry the original diagonal-curvature signals.
            d2: array![
                [1.0, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
                [2.0, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
            ],
            collocation_points: array![[0.0, 0.0], [1.0, 1.0]],
            dimension: 2,
        };
        let beta = array![1.0, 1.0];
        let out = compute_spatial_adaptiveweights_for_beta(
            &beta,
            &[cache],
            1e-6,
            1e-6,
            1e-6,
            1e-12,
            1e12,
            None,
        )
        .expect("adaptive weights");
        assert_eq!(out.len(), 1);
        assert!((out[0].inv_magweight[0] - 1.0).abs() < 1e-10);
        assert!((out[0].inv_magweight[1] - 2.0).abs() < 1e-10);
        assert!((out[0].invgradweight[0] - 2.0_f64.sqrt()).abs() < 1e-10);
        assert!((out[0].invgradweight[1] - 8.0_f64.sqrt()).abs() < 1e-10);
        assert!((out[0].inv_lapweight[0] - 1.0).abs() < 1e-10);
        assert!((out[0].inv_lapweight[1] - 2.0).abs() < 1e-10);
    }

    #[test]
    fn adaptiveweight_is_monotone_in_signal_magnitude() {
        let cache = SpatialOperatorRuntimeCache {
            termname: "matern".to_string(),
            feature_cols: vec![0, 1],
            coeff_global_range: 0..2,
            mass_penalty_global_idx: 0,
            tension_penalty_global_idx: 1,
            stiffness_penalty_global_idx: 2,
            d0: array![[1.0, 0.0]],
            d1: array![[1.0, 0.0], [0.0, 1.0]],
            // d2 layout: P=1 collocation point and d=2 axes give 4 rows ordered
            // (axis_a, axis_b) = (0,0), (0,1), (1,0), (1,1). Row 0 carries the
            // diagonal curvature signal so monotonicity in beta is preserved.
            d2: array![[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
            collocation_points: array![[0.0, 0.0]],
            dimension: 2,
        };
        let beta_small = array![0.25, 0.25];
        let beta_large = array![2.0, 2.0];
        let small = compute_spatial_adaptiveweights_for_beta(
            &beta_small,
            std::slice::from_ref(&cache),
            1e-8,
            1e-8,
            1e-8,
            1e-12,
            1e12,
            None,
        )
        .expect("small adaptive weights");
        let large = compute_spatial_adaptiveweights_for_beta(
            &beta_large,
            &[cache],
            1e-8,
            1e-8,
            1e-8,
            1e-12,
            1e12,
            None,
        )
        .expect("large adaptive weights");
        assert!(small[0].inv_magweight[0] < large[0].inv_magweight[0]);
        assert!(small[0].invgradweight[0] < large[0].invgradweight[0]);
        assert!(small[0].inv_lapweight[0] < large[0].inv_lapweight[0]);
    }

    // ------------------------------------------------------------------
    // Posterior-SNR adaptive weighting: end-to-end objective-quality test.
    //
    // Truth: a 1D function on a uniform grid with a genuine sharp edge in the
    // middle (high curvature, *credibly* determined) and two flat regions. The
    // LEFT flat region is a low-information region whose grid coefficients are
    // poorly determined: its posterior covariance Sigma_beta = H^{-1} carries
    // large variance there, and the noisy point-estimate beta_hat shows
    // spurious curvature there. The RIGHT flat region is well determined.
    //
    // We drive the *real* adaptive-weight machinery
    // (`compute_spatial_adaptiveweights_for_beta`) two ways:
    //   * magnitude-only baseline  -> covariance `None`,
    //   * posterior-SNR (default)  -> covariance `Some(Sigma_beta)`,
    // then build the curvature surrogate penalty K = D2^T diag(w_c) D2 each
    // weighting implies, solve the penalized least-squares fit
    // beta = (X^T X + lambda K)^{-1} X^T y on the identity design (X = I, so
    // the coefficients ARE the fitted function values), and compare MSE-to-truth
    //
    //   * in the noisy LEFT flat region  -> must be STRICTLY LOWER for SNR
    //     (it does not chase noise), and
    //   * at the EDGE                     -> must be NO WORSE for SNR
    //     (the credible edge is preserved).
    // ------------------------------------------------------------------
    fn posterior_snr_finite_difference_d2(m: usize, h: f64) -> Array2<f64> {
        // Second-difference operator on a uniform 1D grid of `m` points: one
        // collocation row per interior point, row layout matching the grouped
        // curvature operator (block_dim = dimension^2 = 1 here). Rows for the
        // two endpoints are zero (no curvature defined there).
        let mut d2 = Array2::<f64>::zeros((m, m));
        for k in 1..m - 1 {
            d2[[k, k - 1]] = 1.0 / (h * h);
            d2[[k, k]] = -2.0 / (h * h);
            d2[[k, k + 1]] = 1.0 / (h * h);
        }
        d2
    }

    #[test]
    fn posterior_snr_weighting_suppresses_noise_and_preserves_edge() {
        use crate::faer_ndarray::FaerCholesky;
        use faer::Side;

        let m = 41usize;
        let h = 1.0 / (m as f64 - 1.0);
        let xs: Vec<f64> = (0..m).map(|j| j as f64 * h).collect();

        // True function: flat-low on the left, a sharp tanh edge centered at
        // x = 0.5, flat-high on the right.
        let edge_center = 0.5;
        let edge_sharpness = 28.0;
        let amplitude = 2.0;
        let truth = Array1::from_iter(
            xs.iter()
                .map(|&x| 0.5 * amplitude * (1.0 + (edge_sharpness * (x - edge_center)).tanh())),
        );

        // Region indices.
        let left_flat: Vec<usize> = (0..m).filter(|&j| xs[j] <= 0.30).collect();
        let edge_band: Vec<usize> = (0..m)
            .filter(|&j| (xs[j] - edge_center).abs() <= 0.07)
            .collect();
        assert!(!left_flat.is_empty() && !edge_band.is_empty());

        // Deterministic, reproducible "noise" pattern. The LEFT flat region is
        // a low-information region: large noise. Elsewhere noise is tiny.
        let noise = |j: usize| -> f64 {
            let s = ((j as f64) * 12.9898).sin() * 43758.5453;
            let frac = s - s.floor(); // pseudo-uniform in [0,1)
            2.0 * frac - 1.0 // in [-1,1)
        };
        let mut y = truth.clone();
        let mut beta_hat = truth.clone();
        for &j in &left_flat {
            let nz = 0.85 * noise(j);
            y[j] += nz;
            beta_hat[j] += nz; // noisy point estimate drives the weights
        }
        for j in 0..m {
            if !left_flat.contains(&j) {
                let nz = 0.02 * noise(j + 7);
                y[j] += nz;
                beta_hat[j] += nz;
            }
        }

        // Working-Laplace conditional covariance proxy Sigma_beta = H^{-1}.
        // Diagonal posterior variances: large in the poorly-determined LEFT
        // flat region, small elsewhere. (Diagonal is sufficient and is the
        // honest leading-order structure of a per-coefficient variance.)
        let mut sigma = Array2::<f64>::zeros((m, m));
        for j in 0..m {
            let var = if left_flat.contains(&j) { 0.55 } else { 1e-4 };
            sigma[[j, j]] = var;
        }

        let d2 = posterior_snr_finite_difference_d2(m, h);
        // Magnitude-only and SNR machinery share the identical cache, betas and
        // epsilons; only the covariance argument differs.
        let cache = SpatialOperatorRuntimeCache {
            termname: "snr_1d".to_string(),
            feature_cols: vec![0],
            coeff_global_range: 0..m,
            mass_penalty_global_idx: 0,
            tension_penalty_global_idx: 1,
            stiffness_penalty_global_idx: 2,
            d0: Array2::<f64>::zeros((m, m)),
            d1: Array2::<f64>::zeros((m, m)),
            d2: d2.clone(),
            collocation_points: {
                let mut cp = Array2::<f64>::zeros((m, 1));
                for j in 0..m {
                    cp[[j, 0]] = xs[j];
                }
                cp
            },
            dimension: 1,
        };

        let eps = 0.05; // shared Charbonnier transition scale
        let weight_floor = 1e-12;
        let weight_ceiling = 1e12;

        let mag = compute_spatial_adaptiveweights_for_beta(
            &beta_hat,
            std::slice::from_ref(&cache),
            eps,
            eps,
            eps,
            weight_floor,
            weight_ceiling,
            None,
        )
        .expect("magnitude-only weights");
        let snr = compute_spatial_adaptiveweights_for_beta(
            &beta_hat,
            std::slice::from_ref(&cache),
            eps,
            eps,
            eps,
            weight_floor,
            weight_ceiling,
            Some(&sigma),
        )
        .expect("posterior-SNR weights");

        // inv_lapweight = 1 / w_c; recover the curvature weights w_c.
        let w_mag = mag[0].inv_lapweight.mapv(|iv| 1.0 / iv);
        let w_snr = snr[0].inv_lapweight.mapv(|iv| 1.0 / iv);

        // Sanity on the mechanism itself: in the noisy LEFT flat region the SNR
        // curvature weight must be substantially LARGER (more smoothing) than
        // the magnitude-only weight, because the spurious point-estimate
        // curvature there is swamped by the posterior variance. (The real
        // quality bar is the MSE assertions below; this only confirms the
        // covariance path is genuinely engaged.)
        let mean = |idx: &[usize], w: &Array1<f64>| -> f64 {
            idx.iter().map(|&j| w[j]).sum::<f64>() / idx.len() as f64
        };
        assert!(
            mean(&left_flat, &w_snr) > mean(&left_flat, &w_mag) * 1.5,
            "SNR must penalize the noisy flat region more: w_snr={:.4e} vs w_mag={:.4e}",
            mean(&left_flat, &w_snr),
            mean(&left_flat, &w_mag),
        );

        // Penalized least-squares fit on the identity design X = I:
        //   beta = (I + lambda * D2^T diag(w_c) D2)^{-1} y.
        let lambda = 0.02;
        let fit = |w: &Array1<f64>| -> Array1<f64> {
            let k = scalar_operatorhessian(&d2, w); // D2^T diag(w) D2 (symmetric)
            let mut a = Array2::<f64>::eye(m);
            a.scaled_add(lambda, &k);
            let factor = a
                .cholesky(Side::Lower)
                .expect("penalized normal matrix is SPD");
            factor.solvevec(&y)
        };
        let fit_mag = fit(&w_mag);
        let fit_snr = fit(&w_snr);

        let region_mse = |idx: &[usize], f: &Array1<f64>| -> f64 {
            idx.iter().map(|&j| (f[j] - truth[j]).powi(2)).sum::<f64>() / idx.len() as f64
        };

        let mse_flat_mag = region_mse(&left_flat, &fit_mag);
        let mse_flat_snr = region_mse(&left_flat, &fit_snr);
        let mse_edge_mag = region_mse(&edge_band, &fit_mag);
        let mse_edge_snr = region_mse(&edge_band, &fit_snr);

        // Objective quality assertions.
        // 1. Noisy flat region: SNR fit is STRICTLY closer to truth (does not
        //    chase noise). Require a clear margin, not a hairline win.
        assert!(
            mse_flat_snr < mse_flat_mag * 0.9,
            "posterior-SNR should be strictly smoother in the noisy flat region: \
             mse_flat_snr={mse_flat_snr:.6e} vs mse_flat_mag={mse_flat_mag:.6e}"
        );
        // 2. Edge region: SNR recovers the edge at least as sharply (MSE no
        //    worse, with a small tolerance for numerical wiggle).
        assert!(
            mse_edge_snr <= mse_edge_mag * 1.05,
            "posterior-SNR must recover the edge at least as sharply: \
             mse_edge_snr={mse_edge_snr:.6e} vs mse_edge_mag={mse_edge_mag:.6e}"
        );

        // Guard against the degenerate pass where both fits are identical (the
        // covariance path must actually change the weights).
        assert!(
            (mse_flat_mag - mse_flat_snr).abs() > 1e-9,
            "the two weightings must produce materially different flat-region fits"
        );
    }

    #[test]
    fn exact_spatial_adaptive_regularization_fit_runswithout_mm() {
        let n = 48usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.19 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            y[i] = (3.0 * x0).sin() + 0.25 * x1;
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 10 },
                        length_scale: 0.7,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 25,
                tol: 1e-5,
                adaptive_regularization: Some(AdaptiveRegularizationOptions {
                    enabled: true,
                    max_mm_iter: 4,
                    beta_rel_tol: 1e-4,
                    max_epsilon_outer_iter: 2,
                    min_epsilon: 1e-6,
                    ..AdaptiveRegularizationOptions::default()
                }),
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("exact adaptive spatial fit should succeed");

        let diag = fit
            .adaptive_diagnostics
            .as_ref()
            .expect("adaptive diagnostics should be present");
        assert_eq!(diag.mm_iterations, 0);
        assert!(diag.epsilon_0.is_finite() && diag.epsilon_0 > 0.0);
        assert!(diag.epsilon_g.is_finite() && diag.epsilon_g > 0.0);
        assert!(diag.epsilon_c.is_finite() && diag.epsilon_c > 0.0);
        assert_eq!(diag.maps.len(), 1);
        assert!(fit.fit.beta.iter().all(|v| v.is_finite()));
        assert!(fit.fit.reml_score.is_finite());
    }

    #[test]
    fn pure_duchon_skips_operator_triplet_adaptive_overlay() {
        let n = 56usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let t = 2.0 * std::f64::consts::PI * i as f64 / n as f64;
            data[[i, 0]] = t.cos();
            data[[i, 1]] = t.sin();
            y[i] = 0.55 * t.sin() + 0.18 * (3.0 * t).cos();
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "pure_duchon_circle".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0, 1],
                    spec: DuchonBasisSpec {
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 16 },
                        length_scale: None,
                        power: 0.0,
                        nullspace_order: DuchonNullspaceOrder::Degree(2),
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        periodic: None,
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let fit = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                compute_inference: false,
                // The scale-free Duchon mass candidate is now the centered
                // design Gram at data points (commit 718810d1), which has
                // an exact zero eigenvalue in the constant direction
                // instead of the rank-full collocation Gram the test was
                // first calibrated against. The outer REML iteration
                // sees a rank-deficient `S_mass` and converges in a few
                // more iterations than under the old construction; the
                // *fit quality* matches (final grad norm ≈ 6e-5 at iter
                // 18). Budget bumped to 28 so the convergence check
                // stays a real assertion rather than an iter-budget
                // artefact.
                max_iter: 28,
                tol: 1e-5,
                adaptive_regularization: Some(AdaptiveRegularizationOptions {
                    enabled: true,
                    max_mm_iter: 4,
                    beta_rel_tol: 1e-4,
                    max_epsilon_outer_iter: 2,
                    min_epsilon: 1e-6,
                    ..AdaptiveRegularizationOptions::default()
                }),
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("pure Duchon exact adaptive fit should succeed");

        // The Charbonnier spatial-adaptive overlay reweights the full
        // mass/tension/stiffness operator triplet. A redesigned Duchon basis
        // deliberately ships its exact RKHS `Primary` curvature plus the lower
        // mass/tension dials — never the stiffness order — so it does not carry
        // the triplet the overlay consumes (see `extract_spatial_operator_runtime_caches`,
        // which has no Duchon arm). Requesting adaptive regularization on a pure
        // Duchon must therefore still fit finitely, but produces no operator-triplet
        // overlay: only Matérn collocation bases feed the adaptive path.
        assert!(
            fit.adaptive_diagnostics.is_none(),
            "pure Duchon carries no operator triplet, so the Charbonnier overlay must not run"
        );
        assert!(fit.fit.beta.iter().all(|v| v.is_finite()));
        assert!(fit.fit.reml_score.is_finite());
    }

    #[test]
    fn exact_spatial_adaptive_binomial_sas_fit_preserves_link_state() {
        let n = 36usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = -1.0 + 2.0 * (i as f64 / (n as f64 - 1.0));
            let x1 = (0.23 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            let eta = 0.55 * x0 - 0.2 * x1 + 0.1 * x0 * x1;
            let p = 1.0 / (1.0 + (-eta).exp());
            let u = ((i * 37 + 13) % 100) as f64 / 100.0;
            y[i] = if u < p { 1.0 } else { 0.0 };
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 10 },
                        length_scale: 0.7,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let fit = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::binomial_sas(
                SasLinkState::new(0.1, -0.2).expect("valid SAS link state"),
            ),
            &FitOptions {
                sas_link: Some(crate::types::SasLinkSpec {
                    initial_epsilon: 0.1,
                    initial_log_delta: -0.2,
                }),
                max_iter: 15,
                tol: 1e-5,
                adaptive_regularization: Some(AdaptiveRegularizationOptions {
                    enabled: true,
                    max_mm_iter: 4,
                    beta_rel_tol: 1e-4,
                    max_epsilon_outer_iter: 2,
                    min_epsilon: 1e-6,
                    ..AdaptiveRegularizationOptions::default()
                }),
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("exact adaptive SAS fit should succeed");

        match fit.fit.fitted_link {
            FittedLinkState::Sas { state, covariance } => {
                assert!(state.epsilon.is_finite());
                assert!(state.log_delta.is_finite());
                assert!(state.delta.is_finite() && state.delta > 0.0);
                assert!(covariance.is_none());
            }
            other => panic!("expected SAS link parameters, got {other:?}"),
        }
    }

    /// Builds the shared spatial-adaptive joint-hyper evaluation scaffolding
    /// (hyperspecs, the zero-ψ derivative blocks, the `base_family` with empty
    /// adaptive params, and the single `eta` block spec) from a Gaussian
    /// baseline fit. Shared verbatim across the FD-gradient and
    /// gradient-lambda-profile pins; the per-test outer-loop cycle counts and
    /// θ probing stay at each call site.
    fn build_spatial_adaptive_joint_hyper_scaffold(
        baseline: &FittedTermCollection,
        runtime_caches: &[SpatialOperatorRuntimeCache],
        y: &Array1<f64>,
        n: usize,
    ) -> (
        SpatialAdaptiveExactFamily,
        ParameterBlockSpec,
        Vec<Vec<CustomFamilyBlockPsiDerivative>>,
    ) {
        let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
        let zero_psi_op: std::sync::Arc<
            dyn crate::custom_family::CustomFamilyPsiDerivativeOperator,
        > = std::sync::Arc::new(crate::custom_family::ZeroPsiDerivativeOperator::new(
            baseline.design.design.nrows(),
            baseline.design.design.ncols(),
        ));
        let derivative_blocks = vec![
            hyperspecs
                .iter()
                .map(|_| CustomFamilyBlockPsiDerivative {
                    penalty_index: None,
                    x_psi: Array2::<f64>::zeros((0, 0)),
                    s_psi: Array2::<f64>::zeros((0, 0)),
                    s_psi_components: None,
                    s_psi_penalty_components: None,
                    x_psi_psi: None,
                    s_psi_psi: None,
                    s_psi_psi_components: None,
                    s_psi_psi_penalty_components: None,
                    implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
                    implicit_axis: 0,
                    implicit_group_id: None,
                })
                .collect::<Vec<_>>(),
        ];
        let base_family = SpatialAdaptiveExactFamily {
            family: LikelihoodSpec::gaussian_identity(),
            latent_cloglog_state: None,
            mixture_link_state: None,
            sas_link_state: None,
            y: Arc::new(y.clone()),
            weights: Arc::new(Array1::ones(n)),
            design: baseline.design.design.to_dense_arc(),
            offset: Arc::new(Array1::zeros(n)),
            linear_constraints: baseline.design.linear_constraints.clone(),
            runtime_caches: Arc::new(runtime_caches.to_vec()),
            adaptive_params: Vec::new(),
            fixed_quadratichessian: Arc::new(Array2::<f64>::zeros((
                baseline.design.design.ncols(),
                baseline.design.design.ncols(),
            ))),
            hyperspecs: Arc::new(hyperspecs),
            exact_eval_cache: Arc::new(Mutex::new(None)),
        };
        let blockspec = ParameterBlockSpec {
            name: "eta".to_string(),
            design: baseline.design.design.clone(),
            offset: Array1::zeros(n),
            penalties: vec![],
            nullspace_dims: vec![],
            initial_log_lambdas: Array1::zeros(0),
            initial_beta: Some(baseline.fit.beta.clone()),
            gauge_priority: 100,
            jacobian_callback: None,
            stacked_design: None,
            stacked_offset: None,
        };
        (base_family, blockspec, derivative_blocks)
    }

    #[test]
    fn exact_spatial_adaptive_joint_hypergradient_matches_finite_difference() {
        let n = 36usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x0 = i as f64 / (n as f64 - 1.0);
            let x1 = (0.31 * i as f64).sin();
            data[[i, 0]] = x0;
            data[[i, 1]] = x1;
            y[i] = (4.0 * x0).sin() + 0.35 * x1 + 0.2 * ((x0 - 0.55) * 18.0).tanh();
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 8 },
                        length_scale: 0.6,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let baseline = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 30,
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("baseline fit");
        let runtime_caches = extract_spatial_operator_runtime_caches(&spec, &baseline.design)
            .expect("runtime caches");
        assert_eq!(runtime_caches.len(), 1);

        let adaptive_opts = AdaptiveRegularizationOptions::default();
        let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
            &baseline.fit.beta,
            &runtime_caches,
            adaptive_opts.min_epsilon,
        )
        .expect("initial epsilons");
        let (base_family, blockspec, derivative_blocks) =
            build_spatial_adaptive_joint_hyper_scaffold(&baseline, &runtime_caches, &y, n);
        let outer_opts = BlockwiseFitOptions {
            inner_max_cycles: 30,
            inner_tol: 1e-6,
            outer_max_iter: 30,
            outer_tol: 1e-6,
            compute_covariance: false,
            ..BlockwiseFitOptions::default()
        };

        let evaluate_theta = |theta: &Array1<f64>, need_hessian: bool| {
            let family = base_family.with_adaptive_params(
                vec![SpatialAdaptiveTermHyperParams {
                    lambda: [theta[0].exp(), theta[1].exp(), theta[2].exp()],
                    epsilon: [theta[3].exp(), theta[4].exp(), theta[5].exp()],
                }],
                Arc::new(Array2::<f64>::zeros((
                    baseline.design.design.ncols(),
                    baseline.design.design.ncols(),
                ))),
            );
            evaluate_custom_family_joint_hyper(
                &family,
                std::slice::from_ref(&blockspec),
                &outer_opts,
                &Array1::zeros(0),
                &derivative_blocks,
                None,
                if need_hessian {
                    crate::solver::estimate::reml::unified::EvalMode::ValueGradientHessian
                } else {
                    crate::solver::estimate::reml::unified::EvalMode::ValueAndGradient
                },
            )
            .expect("joint hyper eval")
        };

        let theta = array![
            baseline.fit.lambdas[runtime_caches[0].mass_penalty_global_idx]
                .max(1e-6)
                .ln(),
            baseline.fit.lambdas[runtime_caches[0].tension_penalty_global_idx]
                .max(1e-6)
                .ln(),
            baseline.fit.lambdas[runtime_caches[0].stiffness_penalty_global_idx]
                .max(1e-6)
                .ln(),
            eps_0_init.max(1e-6).ln(),
            eps_g_init.max(1e-6).ln(),
            eps_c_init.max(1e-6).ln(),
        ];
        let analytic = evaluate_theta(&theta, true);
        assert_eq!(analytic.gradient.len(), theta.len());
        assert!(
            analytic.outer_hessian.is_analytic(),
            "adaptive joint hyper evaluation must expose exact Hessian curvature"
        );
        assert_eq!(
            analytic.outer_hessian.dim(),
            Some(theta.len()),
            "adaptive joint hyper Hessian must span all lambda/epsilon coordinates"
        );
        let analytic_hessian = analytic
            .outer_hessian
            .clone()
            .materialize_dense()
            .expect("adaptive joint hyper Hessian should materialize")
            .expect("adaptive joint hyper Hessian should be present");
        let h = 1e-5;
        for j in 0..theta.len() {
            let mut plus = theta.clone();
            plus[j] += h;
            let mut minus = theta.clone();
            minus[j] -= h;
            let fd = (evaluate_theta(&plus, false).objective
                - evaluate_theta(&minus, false).objective)
                / (2.0 * h);
            assert!(
                (analytic.gradient[j] - fd).abs() < 5e-3 * (1.0 + fd.abs()),
                "adaptive joint hypergradient mismatch at {j}: analytic={}, fd={fd}",
                analytic.gradient[j]
            );
            let grad_fd = (evaluate_theta(&plus, false).gradient
                - evaluate_theta(&minus, false).gradient)
                / (2.0 * h);
            for i in 0..theta.len() {
                assert!(
                    (analytic_hessian[[i, j]] - grad_fd[i]).abs() < 5e-2 * (1.0 + grad_fd[i].abs()),
                    "adaptive joint hyper-Hessian mismatch at ({i},{j}): analytic={}, fd={}",
                    analytic_hessian[[i, j]],
                    grad_fd[i]
                );
            }
        }
    }

    #[test]
    fn exact_spatial_adaptive_1dobjective_profile_has_finite_gradient_lambda_surface() {
        let n = 96usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = x;
            y[i] = 0.12 * (2.0 * std::f64::consts::PI * x).sin()
                + 0.05 * (5.0 * std::f64::consts::PI * x).cos()
                + 1.4 / (1.0 + (-(x - 0.5) / 0.012).exp());
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 31 },
                        length_scale: Some(1.0),
                        power: 2.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let baseline = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 20,
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("baseline fit");
        let runtime_caches = extract_spatial_operator_runtime_caches(&spec, &baseline.design)
            .expect("runtime caches");
        assert_eq!(runtime_caches.len(), 1);
        let (eps_0, eps_g, eps_c) =
            compute_initial_epsilons(&baseline.fit.beta, &runtime_caches, 1e-8)
                .expect("initial epsilons");
        let (base_family, blockspec, derivative_blocks) =
            build_spatial_adaptive_joint_hyper_scaffold(&baseline, &runtime_caches, &y, n);
        let outer_opts = BlockwiseFitOptions {
            inner_max_cycles: 20,
            inner_tol: 1e-6,
            outer_max_iter: 20,
            outer_tol: 1e-6,
            compute_covariance: false,
            ..BlockwiseFitOptions::default()
        };

        let evaluate_theta = |log_lambda_g: f64| {
            let family = base_family.with_adaptive_params(
                vec![SpatialAdaptiveTermHyperParams {
                    lambda: [1e-12, log_lambda_g.exp(), 1e-12],
                    epsilon: [eps_0, eps_g, eps_c],
                }],
                Arc::new(Array2::<f64>::zeros((
                    baseline.design.design.ncols(),
                    baseline.design.design.ncols(),
                ))),
            );
            evaluate_custom_family_joint_hyper(
                &family,
                std::slice::from_ref(&blockspec),
                &outer_opts,
                &Array1::zeros(0),
                &derivative_blocks,
                None,
                crate::solver::estimate::reml::unified::EvalMode::ValueAndGradient,
            )
            .expect("joint hyper eval")
        };

        let low = evaluate_theta((1e-8_f64).ln());
        let mid = evaluate_theta((1e-4_f64).ln());
        let high = evaluate_theta((1e-2_f64).ln());

        for (label, eval) in [("low", &low), ("mid", &mid), ("high", &high)] {
            assert!(
                eval.objective.is_finite(),
                "{label} gradient-lambda profile objective is not finite: {}",
                eval.objective
            );
            assert!(
                eval.gradient.iter().all(|v| v.is_finite()),
                "{label} gradient-lambda profile gradient contains non-finite entries: {:?}",
                eval.gradient
            );
        }
        assert!(
            (low.objective - high.objective).abs() > 1e-8,
            "gradient-lambda profile should remain identifiable: low={}, high={}",
            low.objective,
            high.objective
        );
    }

    #[test]
    fn high_center_duchon_fit_ignores_unavailable_spatial_adaptive_overlay() {
        let n = 320usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = x;
            y[i] = 0.12 * (2.0 * std::f64::consts::PI * x).sin()
                + 0.05 * (5.0 * std::f64::consts::PI * x).cos()
                + 1.4 / (1.0 + (-(x - 0.5) / 0.012).exp());
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 120 },
                        length_scale: Some(1.0),
                        power: 2.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let fit = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 40,
                adaptive_regularization: Some(AdaptiveRegularizationOptions {
                    enabled: true,
                    ..AdaptiveRegularizationOptions::default()
                }),
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("high-center adaptive Duchon fit should not fail");

        assert!(fit.fit.beta.iter().all(|v| v.is_finite()));
        assert!(fit.fit.deviance.is_finite());
        assert!(fit.fit.edf_total().is_some_and(f64::is_finite));
        assert!(
            fit.adaptive_diagnostics.is_none(),
            "Duchon does not expose the complete operator triplet required by the runtime adaptive overlay"
        );
    }

    #[test]
    fn binomial_logit_tail_curvature_uses_stable_exact_formula() {
        let eta = array![30.0, 30.0, -30.0, -30.0, 40.0, -40.0];
        let y = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
        let weights = Array1::ones(eta.len());
        let obs = evaluate_standard_familyobservations(
            LikelihoodSpec::binomial_logit(),
            None,
            None,
            None,
            &y,
            &weights,
            &eta,
        )
        .expect("stable logit observations");

        for i in 0..eta.len() {
            let jet = logit_inverse_link_jet5(eta[i]);
            assert!(
                (obs.neghessian_eta[i] - jet.d1).abs() <= 1e-12 * (1.0 + jet.d1.abs()),
                "eta={} y={} curvature={} target={}",
                eta[i],
                y[i],
                obs.neghessian_eta[i],
                jet.d1
            );
            assert!(
                (obs.neghessian_eta_derivative[i] - jet.d2).abs() <= 1e-12 * (1.0 + jet.d2.abs()),
                "eta={} y={} dcurvature={} target={}",
                eta[i],
                y[i],
                obs.neghessian_eta_derivative[i],
                jet.d2
            );
            assert!(
                obs.neghessian_eta[i].is_finite()
                    && obs.neghessian_eta_derivative[i].is_finite()
                    && obs.log_likelihood.is_finite(),
                "expected finite logit tail observation state at eta={} y={}",
                eta[i],
                y[i]
            );
        }
    }

    #[test]
    fn non_logit_binomial_tailobservations_stay_finite() {
        let eta = array![12.0, -12.0, 18.0, -18.0];
        let y = array![0.0, 1.0, 1.0, 0.0];
        let weights = Array1::ones(eta.len());
        for family in [
            LikelihoodSpec::binomial_probit(),
            LikelihoodSpec::binomial_cloglog(),
        ] {
            let obs = evaluate_standard_familyobservations(
                family.clone(),
                None,
                None,
                None,
                &y,
                &weights,
                &eta,
            )
            .expect("tail observations");
            assert!(obs.log_likelihood.is_finite(), "family={family:?}");
            assert!(
                obs.score.iter().all(|v| v.is_finite())
                    && obs.neghessian_eta.iter().all(|v| v.is_finite())
                    && obs.neghessian_eta_derivative.iter().all(|v| v.is_finite()),
                "family={family:?}"
            );
        }
    }

    #[test]
    fn two_block_exact_joint_setup_sanitizes_non_finite_rho_seed() {
        let setup = ExactJointHyperSetup::new(
            array![f64::NEG_INFINITY, 0.25, f64::INFINITY],
            array![-12.0, -12.0, -12.0],
            array![12.0, 12.0, 12.0],
            SpatialLogKappaCoords::new_with_dims(array![0.5], vec![1]),
            SpatialLogKappaCoords::new_with_dims(array![-2.0], vec![1]),
            SpatialLogKappaCoords::new_with_dims(array![2.0], vec![1]),
        );

        let theta0 = setup.theta0();
        assert!(theta0.iter().all(|v| v.is_finite()));
        assert_eq!(theta0[0], 0.0);
        assert_eq!(theta0[1], 0.25);
        assert_eq!(theta0[2], 0.0);
        assert_eq!(theta0[3], 0.5);
    }

    #[test]
    fn extracted_spatial_runtime_cache_matches_normalized_design_penalties() {
        let n = 24usize;
        let mut data = Array2::<f64>::zeros((n, 2));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
            data[[i, 1]] = (0.23 * i as f64).cos();
        }
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "matern".to_string(),
                basis: SmoothBasisSpec::Matern {
                    feature_cols: vec![0, 1],
                    spec: MaternBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 7 },
                        length_scale: 0.8,
                        nu: MaternNu::FiveHalves,
                        include_intercept: false,
                        double_penalty: true,
                        identifiability: MaternIdentifiability::CenterSumToZero,
                        aniso_log_scales: None,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        let caches =
            extract_spatial_operator_runtime_caches(&spec, &design).expect("runtime caches");
        assert_eq!(caches.len(), 1);
        let cache = &caches[0];
        let s0 = {
            let raw = cache.d0.t().dot(&cache.d0);
            (&raw + &raw.t()) * 0.5
        };
        let s1 = {
            let raw = cache.d1.t().dot(&cache.d1);
            (&raw + &raw.t()) * 0.5
        };
        let s2 = {
            let raw = cache.d2.t().dot(&cache.d2);
            (&raw + &raw.t()) * 0.5
        };

        let s0_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s0,
        );
        let s1_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s1,
        );
        let s2_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s2,
        );

        let p_total = design.design.ncols();
        let err0 = (&s0_global
            - &design.penalties[cache.mass_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);
        let err1 = (&s1_global
            - &design.penalties[cache.tension_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);
        let err2 = (&s2_global
            - &design.penalties[cache.stiffness_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);

        assert!(err0 < 1e-8, "mass penalty mismatch too large: {err0}");
        assert!(err1 < 1e-8, "tension penalty mismatch too large: {err1}");
        assert!(err2 < 1e-8, "stiffness penalty mismatch too large: {err2}");
    }

    #[test]
    fn extracted_duchon_spatial_runtime_cache_matches_normalized_design_penalties() {
        let n = 32usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        for i in 0..n {
            data[[i, 0]] = i as f64 / (n as f64 - 1.0);
        }
        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 11 },
                        length_scale: Some(0.8),
                        power: 2.0,
                        nullspace_order: DuchonNullspaceOrder::Zero,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::all_active(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };

        let design = build_term_collection_design(data.view(), &spec).expect("design");
        assert_eq!(design.penalties.len(), 4);
        let caches =
            extract_spatial_operator_runtime_caches(&spec, &design).expect("runtime caches");
        assert_eq!(caches.len(), 1);
        let cache = &caches[0];
        let s0 = {
            let raw = cache.d0.t().dot(&cache.d0);
            (&raw + &raw.t()) * 0.5
        };
        let s1 = {
            let raw = cache.d1.t().dot(&cache.d1);
            (&raw + &raw.t()) * 0.5
        };
        let s2 = {
            let raw = cache.d2.t().dot(&cache.d2);
            (&raw + &raw.t()) * 0.5
        };

        let s0_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s0,
        );
        let s1_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s1,
        );
        let s2_global = penalty_matrixwith_local_block(
            design.design.ncols(),
            cache.coeff_global_range.clone(),
            &s2,
        );

        let p_total = design.design.ncols();
        let err0 = (&s0_global
            - &design.penalties[cache.mass_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);
        let err1 = (&s1_global
            - &design.penalties[cache.tension_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);
        let err2 = (&s2_global
            - &design.penalties[cache.stiffness_penalty_global_idx].to_global(p_total))
            .iter()
            .map(|v| v.abs())
            .fold(0.0_f64, f64::max);

        assert!(
            err0 < 1e-8,
            "Duchon mass penalty mismatch too large: {err0}"
        );
        assert!(
            err1 < 1e-8,
            "Duchon tension penalty mismatch too large: {err1}"
        );
        assert!(
            err2 < 1e-8,
            "Duchon stiffness penalty mismatch too large: {err2}"
        );
    }

    #[test]
    fn spatial_adaptive_explicit_second_order_kind_matches_block_sparsity() {
        let alpha_mass_0 = SpatialAdaptiveHyperSpec {
            cache_index: 0,
            kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
        };
        let alpha_mass_1 = SpatialAdaptiveHyperSpec {
            cache_index: 1,
            kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
        };
        let alpha_grad_0 = SpatialAdaptiveHyperSpec {
            cache_index: 0,
            kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
        };
        let eta_mass = SpatialAdaptiveHyperSpec {
            cache_index: 0,
            kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
        };
        let eta_grad = SpatialAdaptiveHyperSpec {
            cache_index: 0,
            kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
        };

        assert_eq!(
            alpha_mass_0.explicit_second_order_kind(alpha_mass_0),
            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
        );
        assert_eq!(
            alpha_mass_0.explicit_second_order_kind(alpha_mass_1),
            SpatialAdaptiveExplicitSecondOrderKind::StructuralZero
        );
        assert_eq!(
            alpha_mass_0.explicit_second_order_kind(alpha_grad_0),
            SpatialAdaptiveExplicitSecondOrderKind::StructuralZero
        );
        assert_eq!(
            alpha_mass_1.explicit_second_order_kind(eta_mass),
            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
        );
        assert_eq!(
            eta_mass.explicit_second_order_kind(alpha_mass_1),
            SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
        );
        assert_eq!(
            eta_mass.explicit_second_order_kind(eta_mass),
            SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta
        );
        assert_eq!(
            eta_mass.explicit_second_order_kind(eta_grad),
            SpatialAdaptiveExplicitSecondOrderKind::StructuralZero
        );
    }

    /// Parity test for the unified adaptive hyper-derivative dispatch (issue
    /// #426). The unified `adaptive_block_eval` / `adaptive_block_drift_eval`
    /// engines must reproduce, bit-for-bit, the per-(component, derivative-kind)
    /// pieces that the previous hand-rolled `adaptive_block_*` /
    /// `adaptive_block_*_drift` functions assembled. The right-hand reference
    /// below is the same closed-form composition those duplicated functions
    /// performed: pick the matching state accessor, run it through the matching
    /// scalar / grouped operator with the component weight, and embed into the
    /// global coordinate range. Equality is exact (both sides call the identical
    /// accessors and operators), so the bound is a true `==`.
    #[test]
    fn adaptive_hyper_derivative_dispatch_matches_reference() {
        let n = 40usize;
        let mut data = Array2::<f64>::zeros((n, 1));
        let mut y = Array1::<f64>::zeros(n);
        for i in 0..n {
            let x = i as f64 / (n as f64 - 1.0);
            data[[i, 0]] = x;
            y[i] = 0.3 * (2.0 * std::f64::consts::PI * x).sin()
                + 1.1 / (1.0 + (-(x - 0.5) / 0.02).exp());
        }

        let spec = TermCollectionSpec {
            linear_terms: vec![],
            random_effect_terms: vec![],
            smooth_terms: vec![SmoothTermSpec {
                name: "duchon".to_string(),
                basis: SmoothBasisSpec::Duchon {
                    feature_cols: vec![0],
                    spec: DuchonBasisSpec {
                        periodic: None,
                        center_strategy: CenterStrategy::FarthestPoint { num_centers: 15 },
                        length_scale: Some(0.9),
                        power: 2.0,
                        nullspace_order: DuchonNullspaceOrder::Linear,
                        identifiability: SpatialIdentifiability::default(),
                        aniso_log_scales: None,
                        operator_penalties: DuchonOperatorPenaltySpec::default(),
                        boundary: OneDimensionalBoundary::Open,
                    },
                    input_scales: None,
                },
                shape: ShapeConstraint::None,
                joint_null_rotation: None,
            }],
        };
        let baseline = fit_term_collection_forspec(
            data.view(),
            y.view(),
            Array1::ones(n).view(),
            Array1::zeros(n).view(),
            &spec,
            LikelihoodSpec::gaussian_identity(),
            &FitOptions {
                max_iter: 20,
                penalty_shrinkage_floor: None,
                ..FitOptions::default()
            },
        )
        .expect("baseline fit");
        let runtime_caches = extract_spatial_operator_runtime_caches(&spec, &baseline.design)
            .expect("runtime caches");
        assert_eq!(runtime_caches.len(), 1);
        let (eps_0, eps_g, eps_c) =
            compute_initial_epsilons(&baseline.fit.beta, &runtime_caches, 1e-8)
                .expect("initial epsilons");
        let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
        let p = baseline.design.design.ncols();

        let family = SpatialAdaptiveExactFamily {
            family: LikelihoodSpec::gaussian_identity(),
            latent_cloglog_state: None,
            mixture_link_state: None,
            sas_link_state: None,
            y: Arc::new(y.clone()),
            weights: Arc::new(Array1::ones(n)),
            design: baseline.design.design.to_dense_arc(),
            offset: Arc::new(Array1::zeros(n)),
            linear_constraints: baseline.design.linear_constraints.clone(),
            runtime_caches: Arc::new(runtime_caches.clone()),
            adaptive_params: vec![SpatialAdaptiveTermHyperParams {
                lambda: [0.7, 1.3, 0.4],
                epsilon: [eps_0, eps_g, eps_c],
            }],
            fixed_quadratichessian: Arc::new(Array2::<f64>::zeros((p, p))),
            hyperspecs: Arc::new(hyperspecs),
            exact_eval_cache: Arc::new(Mutex::new(None)),
        };

        let beta = baseline.fit.beta.clone();
        let eval = family.exact_evaluation(&beta).expect("exact evaluation");
        let cache = &family.runtime_caches[0];
        let params = &family.adaptive_params[0];
        let state = &eval.adaptive_states[0];
        let range = cache.coeff_global_range.clone();

        // Independent reference assembly for one (component, kind) parts triple.
        let reference_parts = |component: AdaptiveComponent,
                               kind: HyperDerivativeKind|
         -> (f64, Array1<f64>, Array2<f64>) {
            let (objective, grad_local, hess_local) = match component {
                AdaptiveComponent::Magnitude => {
                    let lambda = params.lambda[0];
                    let mag = &state.magnitude;
                    let (obj, gc, hd) = match kind {
                        HyperDerivativeKind::Rho => (
                            mag.penalty_value(),
                            mag.betagradient_coeff(),
                            mag.betahessian_diag(),
                        ),
                        HyperDerivativeKind::LogEpsilonFirst => (
                            mag.log_epsilon_gradient_terms().sum(),
                            mag.log_epsilon_betagradient_coeff(),
                            mag.log_epsilon_betahessian_diag(),
                        ),
                        HyperDerivativeKind::LogEpsilonSecond => (
                            mag.log_epsilon_hessian_terms().sum(),
                            mag.log_epsilon_beta_mixed_second_coeff(),
                            mag.log_epsilon_betahessian_second_diag(),
                        ),
                    };
                    (
                        lambda * obj,
                        lambda * scalar_operatorgradient(&cache.d0, &gc),
                        lambda * scalar_operatorhessian(&cache.d0, &hd),
                    )
                }
                AdaptiveComponent::Gradient => {
                    let lambda = params.lambda[1];
                    let grad = &state.gradient;
                    let (obj, gb, hb) = match kind {
                        HyperDerivativeKind::Rho => (
                            grad.penalty_value(),
                            grad.betagradient_blocks(),
                            grad.betahessian_blocks(),
                        ),
                        HyperDerivativeKind::LogEpsilonFirst => (
                            grad.log_epsilon_gradient_terms().sum(),
                            grad.log_epsilon_betagradient_blocks(),
                            grad.log_epsilon_betahessian_blocks(),
                        ),
                        HyperDerivativeKind::LogEpsilonSecond => (
                            grad.log_epsilon_hessian_terms().sum(),
                            grad.log_epsilon_beta_mixed_second_blocks(),
                            grad.log_epsilon_betahessian_second_blocks(),
                        ),
                    };
                    (
                        lambda * obj,
                        lambda
                            * grouped_operatorgradient(&cache.d1, cache.dimension, &gb)
                                .expect("grouped gradient"),
                        lambda
                            * grouped_operatorhessian(&cache.d1, cache.dimension, &hb)
                                .expect("grouped hessian"),
                    )
                }
                AdaptiveComponent::Curvature => {
                    let lambda = params.lambda[2];
                    let group = cache.dimension * cache.dimension;
                    let curv = &state.curvature;
                    let (obj, gb, hb) = match kind {
                        HyperDerivativeKind::Rho => (
                            curv.penalty_value(),
                            curv.betagradient_blocks(),
                            curv.betahessian_blocks(),
                        ),
                        HyperDerivativeKind::LogEpsilonFirst => (
                            curv.log_epsilon_gradient_terms().sum(),
                            curv.log_epsilon_betagradient_blocks(),
                            curv.log_epsilon_betahessian_blocks(),
                        ),
                        HyperDerivativeKind::LogEpsilonSecond => (
                            curv.log_epsilon_hessian_terms().sum(),
                            curv.log_epsilon_beta_mixed_second_blocks(),
                            curv.log_epsilon_betahessian_second_blocks(),
                        ),
                    };
                    (
                        lambda * obj,
                        lambda
                            * grouped_operatorgradient(&cache.d2, group, &gb)
                                .expect("grouped gradient"),
                        lambda
                            * grouped_operatorhessian(&cache.d2, group, &hb)
                                .expect("grouped hessian"),
                    )
                }
            };
            let mut grad = Array1::<f64>::zeros(p);
            grad.slice_mut(s![range.clone()]).assign(&grad_local);
            let mut hess = Array2::<f64>::zeros((p, p));
            hess.slice_mut(s![range.clone(), range.clone()])
                .assign(&hess_local);
            (objective, grad, hess)
        };

        let components = [
            AdaptiveComponent::Magnitude,
            AdaptiveComponent::Gradient,
            AdaptiveComponent::Curvature,
        ];
        let kinds = [
            HyperDerivativeKind::Rho,
            HyperDerivativeKind::LogEpsilonFirst,
            HyperDerivativeKind::LogEpsilonSecond,
        ];

        for &component in &components {
            for &kind in &kinds {
                let (obj_new, grad_new, hess_new) = family
                    .adaptive_block_eval(&eval, 0, component, kind)
                    .expect("unified block eval");
                let (obj_ref, grad_ref, hess_ref) = reference_parts(component, kind);
                assert_eq!(
                    obj_new, obj_ref,
                    "objective mismatch for {component:?}/{kind:?}"
                );
                assert_eq!(
                    grad_new, grad_ref,
                    "gradient mismatch for {component:?}/{kind:?}"
                );
                assert_eq!(
                    hess_new, hess_ref,
                    "hessian mismatch for {component:?}/{kind:?}"
                );
            }
        }

        // Directional-drift parity: independent reference per (component, drift).
        let direction = {
            let mut d = Array1::<f64>::zeros(p);
            for (i, v) in d.iter_mut().enumerate() {
                *v = 0.05 + 0.01 * (i as f64);
            }
            d
        };
        let reference_drift =
            |component: AdaptiveComponent, drift: HyperDriftKind| -> Array2<f64> {
                let direction_local = direction.slice(s![range.clone()]);
                let local = match component {
                    AdaptiveComponent::Magnitude => {
                        let d0_u = cache.d0.dot(&direction_local);
                        let mag = &state.magnitude;
                        let diag = match drift {
                            HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
                            HyperDriftKind::LogEpsilon => {
                                mag.log_epsilon_betahessian_directional_diag(&d0_u)
                            }
                        };
                        params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
                    }
                    AdaptiveComponent::Gradient => {
                        let d1_u = cache.d1.dot(&direction_local);
                        let blocks_in = collocationgradient_blocks(&d1_u, cache.dimension)
                            .expect("collocation gradient blocks");
                        let grad = &state.gradient;
                        let blocks = match drift {
                            HyperDriftKind::Rho => grad.directionalhessian_blocks(&blocks_in),
                            HyperDriftKind::LogEpsilon => {
                                grad.log_epsilon_betahessian_directional_blocks(&blocks_in)
                            }
                        };
                        params.lambda[1]
                            * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
                                .expect("grouped hessian")
                    }
                    AdaptiveComponent::Curvature => {
                        let group = cache.dimension * cache.dimension;
                        let d2_u = cache.d2.dot(&direction_local);
                        let blocks_in = collocationhessian_blocks(&d2_u, cache.dimension)
                            .expect("collocation hessian blocks");
                        let curv = &state.curvature;
                        let blocks = match drift {
                            HyperDriftKind::Rho => curv.directionalhessian_blocks(&blocks_in),
                            HyperDriftKind::LogEpsilon => {
                                curv.log_epsilon_betahessian_directional_blocks(&blocks_in)
                            }
                        };
                        params.lambda[2]
                            * grouped_operatorhessian(&cache.d2, group, &blocks)
                                .expect("grouped hessian")
                    }
                };
                let mut out = Array2::<f64>::zeros((p, p));
                out.slice_mut(s![range.clone(), range.clone()])
                    .assign(&local);
                out
            };

        for &component in &components {
            for &drift in &[HyperDriftKind::Rho, HyperDriftKind::LogEpsilon] {
                let drift_new = family
                    .adaptive_block_drift_eval(&eval, 0, component, drift, &direction)
                    .expect("unified block drift eval");
                let drift_ref = reference_drift(component, drift);
                assert_eq!(
                    drift_new, drift_ref,
                    "drift mismatch for {component:?}/{drift:?}"
                );
            }
        }

        // Dispatch-surface parity: the public entry points must route to the
        // same engine output. `adaptive_hyper_parts` on a per-term `log lambda`
        // spec equals the `Rho` block eval; on a shared `log epsilon` spec it
        // equals the summed `LogEpsilonFirst` blocks (single cache ⇒ the shared
        // sum is the single block's first-order `log epsilon` eval).
        let check_dispatch_parity =
            |specs: [(SpatialAdaptiveHyperKind, AdaptiveComponent); 3],
             deriv_kind: HyperDerivativeKind| {
                for (kind, component) in specs {
                    let hyper = SpatialAdaptiveHyperSpec {
                        cache_index: 0,
                        kind,
                    };
                    let (obj_disp, grad_disp, hess_disp) = family
                        .adaptive_hyper_parts(&eval, hyper)
                        .expect("hyper parts");
                    let (obj_ref, grad_ref, hess_ref) = reference_parts(component, deriv_kind);
                    assert_eq!(obj_disp, obj_ref);
                    assert_eq!(grad_disp, grad_ref);
                    assert_eq!(hess_disp, hess_ref);
                }
            };

        check_dispatch_parity(
            [
                (
                    SpatialAdaptiveHyperKind::LogLambdaMagnitude,
                    AdaptiveComponent::Magnitude,
                ),
                (
                    SpatialAdaptiveHyperKind::LogLambdaGradient,
                    AdaptiveComponent::Gradient,
                ),
                (
                    SpatialAdaptiveHyperKind::LogLambdaCurvature,
                    AdaptiveComponent::Curvature,
                ),
            ],
            HyperDerivativeKind::Rho,
        );

        check_dispatch_parity(
            [
                (
                    SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
                    AdaptiveComponent::Magnitude,
                ),
                (
                    SpatialAdaptiveHyperKind::LogEpsilonGradient,
                    AdaptiveComponent::Gradient,
                ),
                (
                    SpatialAdaptiveHyperKind::LogEpsilonCurvature,
                    AdaptiveComponent::Curvature,
                ),
            ],
            HyperDerivativeKind::LogEpsilonFirst,
        );
    }

    #[test]
    fn scalar_charbonnier_exact_derivatives_match_finite_difference() {
        let signal = array![0.7, -1.1];
        let epsilon = 0.3;
        let state = CharbonnierScalarBlockState::from_signal(signal.clone(), epsilon);
        let h = 1e-5;
        let value = |x: &Array1<f64>| {
            CharbonnierScalarBlockState::from_signal(x.clone(), epsilon).penalty_value()
        };
        for i in 0..signal.len() {
            let mut plus = signal.clone();
            plus[i] += h;
            let mut minus = signal.clone();
            minus[i] -= h;
            let gradfd = (value(&plus) - value(&minus)) / (2.0 * h);
            let hessfd = (value(&plus) - 2.0 * value(&signal) + value(&minus)) / (h * h);
            assert!((state.betagradient_coeff()[i] - gradfd).abs() < 1e-6);
            assert!((state.betahessian_diag()[i] - hessfd).abs() < 1e-4);
        }
    }

    #[test]
    fn grouped_charbonnier_exactgradient_matches_finite_difference() {
        let blocks = array![[0.8, -0.4], [0.3, 0.9]];
        let epsilon = 0.25;
        let state = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), epsilon);
        let h = 1e-6;
        let value = |x: &Array2<f64>| {
            CharbonnierGroupedBlockState::from_signal_blocks(x.clone(), epsilon).penalty_value()
        };
        let analytic = state.betagradient_blocks();
        for k in 0..blocks.nrows() {
            for axis in 0..blocks.ncols() {
                let mut plus = blocks.clone();
                plus[[k, axis]] += h;
                let mut minus = blocks.clone();
                minus[[k, axis]] -= h;
                let gradfd = (value(&plus) - value(&minus)) / (2.0 * h);
                assert!((analytic[[k, axis]] - gradfd).abs() < 1e-6);
            }
        }
    }

    #[test]
    fn scalar_charbonnier_log_epsilon_derivatives_match_finite_difference() {
        let signal = array![0.4, -0.9];
        let epsilon = 0.35_f64;
        let eta = epsilon.ln();
        let state = CharbonnierScalarBlockState::from_signal(signal.clone(), epsilon);
        let h = 1e-5;
        let value = |eta_value: f64| {
            CharbonnierScalarBlockState::from_signal(signal.clone(), eta_value.exp())
                .penalty_value()
        };
        let gradfd = (value(eta + h) - value(eta - h)) / (2.0 * h);
        let hessfd = (value(eta + h) - 2.0 * value(eta) + value(eta - h)) / (h * h);
        assert!((state.log_epsilon_gradient_terms().sum() - gradfd).abs() < 1e-6);
        assert!((state.log_epsilon_hessian_terms().sum() - hessfd).abs() < 1e-4);

        let eval_grad = |eta_value: f64| {
            CharbonnierScalarBlockState::from_signal(signal.clone(), eta_value.exp())
                .betagradient_coeff()
        };
        let mixedfd = (&eval_grad(eta + h) - &eval_grad(eta - h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!((state.log_epsilon_betagradient_coeff()[i] - mixedfd[i]).abs() < 1e-6);
        }

        let eval_hess = |eta_value: f64| {
            CharbonnierScalarBlockState::from_signal(signal.clone(), eta_value.exp())
                .betahessian_diag()
        };
        let betahessfd = (&eval_hess(eta + h) - &eval_hess(eta - h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!((state.log_epsilon_betahessian_diag()[i] - betahessfd[i]).abs() < 1e-5);
        }

        let eval_log_grad = |eta_value: f64| {
            CharbonnierScalarBlockState::from_signal(signal.clone(), eta_value.exp())
                .log_epsilon_betagradient_coeff()
        };
        let second_mixedfd = (&eval_log_grad(eta + h) - &eval_log_grad(eta - h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!(
                (state.log_epsilon_beta_mixed_second_coeff()[i] - second_mixedfd[i]).abs() < 1e-5
            );
        }

        let eval_log_hess = |eta_value: f64| {
            CharbonnierScalarBlockState::from_signal(signal.clone(), eta_value.exp())
                .log_epsilon_betahessian_diag()
        };
        let second_hessfd = (&eval_log_hess(eta + h) - &eval_log_hess(eta - h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!(
                (state.log_epsilon_betahessian_second_diag()[i] - second_hessfd[i]).abs() < 1e-4
            );
        }
    }

    #[test]
    fn scalar_charbonnier_directionalhessian_matches_finite_difference() {
        let signal = array![0.5, -0.6];
        let epsilon = 0.2;
        let direction = array![0.3, -0.1];
        let state = CharbonnierScalarBlockState::from_signal(signal.clone(), epsilon);
        let h = 1e-6;
        let analytic = state.directionalhessian_diag(&direction);
        let evalhess = |step: f64| {
            let shifted = &signal + &(direction.mapv(|v| step * v));
            CharbonnierScalarBlockState::from_signal(shifted, epsilon).betahessian_diag()
        };
        let fd = (&evalhess(h) - &evalhess(-h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!((analytic[i] - fd[i]).abs() < 1e-5);
        }
    }

    #[test]
    fn scalar_charbonnier_log_epsilon_directionalhessian_matches_finite_difference() {
        let signal = array![0.5, -0.6];
        let epsilon = 0.2;
        let direction = array![0.3, -0.1];
        let state = CharbonnierScalarBlockState::from_signal(signal.clone(), epsilon);
        let h = 1e-6;
        let analytic = state.log_epsilon_betahessian_directional_diag(&direction);
        let evalhess = |step: f64| {
            let shifted = &signal + &(direction.mapv(|v| step * v));
            CharbonnierScalarBlockState::from_signal(shifted, epsilon)
                .log_epsilon_betahessian_diag()
        };
        let fd = (&evalhess(h) - &evalhess(-h)) / (2.0 * h);
        for i in 0..signal.len() {
            assert!((analytic[i] - fd[i]).abs() < 1e-4);
        }
    }

    #[test]
    fn grouped_charbonnier_log_epsilon_derivatives_match_finite_difference() {
        let blocks = array![[0.7, -0.2], [0.1, 0.8]];
        let epsilon = 0.3_f64;
        let eta = epsilon.ln();
        let state = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), epsilon);
        let h = 1e-5;
        let value = |eta_value: f64| {
            CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), eta_value.exp())
                .penalty_value()
        };
        let gradfd = (value(eta + h) - value(eta - h)) / (2.0 * h);
        let hessfd = (value(eta + h) - 2.0 * value(eta) + value(eta - h)) / (h * h);
        assert!((state.log_epsilon_gradient_terms().sum() - gradfd).abs() < 1e-6);
        assert!((state.log_epsilon_hessian_terms().sum() - hessfd).abs() < 1e-4);

        let eval_grad = |eta_value: f64| {
            CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), eta_value.exp())
                .betagradient_blocks()
        };
        let mixedfd = (&eval_grad(eta + h) - &eval_grad(eta - h)) / (2.0 * h);
        let analytic_mixed = state.log_epsilon_betagradient_blocks();
        for k in 0..blocks.nrows() {
            for axis in 0..blocks.ncols() {
                assert!((analytic_mixed[[k, axis]] - mixedfd[[k, axis]]).abs() < 1e-6);
            }
        }

        let eval_hess = |eta_value: f64| {
            CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), eta_value.exp())
                .betahessian_blocks()
        };
        let plus_hess = eval_hess(eta + h);
        let minus_hess = eval_hess(eta - h);
        let analytic_hess = state.log_epsilon_betahessian_blocks();
        for k in 0..analytic_hess.len() {
            let fd = (&plus_hess[k] - &minus_hess[k]) / (2.0 * h);
            for i in 0..fd.nrows() {
                for j in 0..fd.ncols() {
                    assert!((analytic_hess[k][[i, j]] - fd[[i, j]]).abs() < 1e-5);
                }
            }
        }

        let eval_log_grad = |eta_value: f64| {
            CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), eta_value.exp())
                .log_epsilon_betagradient_blocks()
        };
        let second_mixedfd = (&eval_log_grad(eta + h) - &eval_log_grad(eta - h)) / (2.0 * h);
        let analytic_second_mixed = state.log_epsilon_beta_mixed_second_blocks();
        for k in 0..blocks.nrows() {
            for axis in 0..blocks.ncols() {
                assert!(
                    (analytic_second_mixed[[k, axis]] - second_mixedfd[[k, axis]]).abs() < 1e-5
                );
            }
        }

        let eval_log_hess = |eta_value: f64| {
            CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), eta_value.exp())
                .log_epsilon_betahessian_blocks()
        };
        let plus_log_hess = eval_log_hess(eta + h);
        let minus_log_hess = eval_log_hess(eta - h);
        let analytic_second_hess = state.log_epsilon_betahessian_second_blocks();
        for k in 0..analytic_second_hess.len() {
            let fd = (&plus_log_hess[k] - &minus_log_hess[k]) / (2.0 * h);
            for i in 0..fd.nrows() {
                for j in 0..fd.ncols() {
                    assert!((analytic_second_hess[k][[i, j]] - fd[[i, j]]).abs() < 1e-4);
                }
            }
        }
    }

    #[test]
    fn grouped_charbonnier_directionalhessian_matches_finite_difference() {
        let blocks = array![[0.6, -0.2], [0.4, 0.5]];
        let direction = array![[0.1, -0.3], [0.2, 0.15]];
        let epsilon = 0.4;
        let state = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), epsilon);
        let analytic = state.directionalhessian_blocks(&direction);
        let h = 1e-6;
        let evalhess = |step: f64| {
            let shifted = &blocks + &(direction.mapv(|v| step * v));
            CharbonnierGroupedBlockState::from_signal_blocks(shifted, epsilon).betahessian_blocks()
        };
        let plus = evalhess(h);
        let minus = evalhess(-h);
        for k in 0..analytic.len() {
            let fd = (&plus[k] - &minus[k]) / (2.0 * h);
            for i in 0..fd.nrows() {
                for j in 0..fd.ncols() {
                    assert!((analytic[k][[i, j]] - fd[[i, j]]).abs() < 1e-5);
                }
            }
        }
    }

    #[test]
    fn grouped_charbonnier_log_epsilon_directionalhessian_matches_finite_difference() {
        let blocks = array![[0.6, -0.2], [0.4, 0.5]];
        let direction = array![[0.1, -0.3], [0.2, 0.15]];
        let epsilon = 0.4;
        let state = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), epsilon);
        let analytic = state.log_epsilon_betahessian_directional_blocks(&direction);
        let h = 1e-6;
        let evalhess = |step: f64| {
            let shifted = &blocks + &(direction.mapv(|v| step * v));
            CharbonnierGroupedBlockState::from_signal_blocks(shifted, epsilon)
                .log_epsilon_betahessian_blocks()
        };
        let plus = evalhess(h);
        let minus = evalhess(-h);
        for k in 0..analytic.len() {
            let fd = (&plus[k] - &minus[k]) / (2.0 * h);
            for i in 0..fd.nrows() {
                for j in 0..fd.ncols() {
                    assert!((analytic[k][[i, j]] - fd[[i, j]]).abs() < 1e-4);
                }
            }
        }
    }

    #[test]
    fn grouped_charbonnier_directionalhessian_blocks_are_symmetric() {
        let blocks = array![[0.6, -0.2], [0.4, 0.5]];
        let direction = array![[0.1, -0.3], [0.2, 0.15]];
        let epsilon = 0.4;
        let analytic = CharbonnierGroupedBlockState::from_signal_blocks(blocks, epsilon)
            .directionalhessian_blocks(&direction);
        for (k, block) in analytic.iter().enumerate() {
            for i in 0..block.nrows() {
                for j in 0..block.ncols() {
                    assert!(
                        (block[[i, j]] - block[[j, i]]).abs() < 1e-12,
                        "directional Hessian block {k} is not symmetric at ({i},{j})"
                    );
                }
            }
        }
    }

    #[test]
    fn scalar_charbonnier_local_quadratic_curvature_matches_transition_scale() {
        let signal = array![0.0, 0.0, 0.0];
        let small = CharbonnierScalarBlockState::from_signal(signal.clone(), 1e-3);
        let large = CharbonnierScalarBlockState::from_signal(signal, 1e3);
        for (&a, &b) in small
            .betahessian_diag()
            .iter()
            .zip(large.betahessian_diag().iter())
        {
            assert!(
                (a - 1e3).abs() < 1e-7,
                "small-epsilon curvature should be 1/eps, got {a}"
            );
            assert!(
                (b - 1e-3).abs() < 1e-13,
                "large-epsilon curvature should be 1/eps, got {b}"
            );
        }
    }

    #[test]
    fn grouped_charbonnier_local_quadratic_curvature_matches_transition_scale() {
        let blocks = array![[0.0, 0.0], [0.0, 0.0]];
        let small = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), 1e-3);
        let large = CharbonnierGroupedBlockState::from_signal_blocks(blocks, 1e3);
        for (small_block, large_block) in small
            .betahessian_blocks()
            .into_iter()
            .zip(large.betahessian_blocks().into_iter())
        {
            let eye = Array2::<f64>::eye(small_block.nrows());
            assert!(
                (&small_block - &eye.mapv(|v| 1e3 * v)).mapv(f64::abs).sum() < 1e-7,
                "small-epsilon grouped curvature should equal I/eps"
            );
            assert!(
                (&large_block - &eye.mapv(|v| 1e-3 * v))
                    .mapv(f64::abs)
                    .sum()
                    < 1e-13,
                "large-epsilon grouped curvature should equal I/eps"
            );
        }
    }

    #[test]
    fn scalar_charbonnier_small_signal_matches_local_half_quadratic() {
        let signal = array![1e-5, -2e-5, 3e-5];
        for &epsilon in &[1e-3, 1e-1, 1.0, 1e2] {
            let state = CharbonnierScalarBlockState::from_signal(signal.clone(), epsilon);
            let value = state.penalty_value();
            let target = 0.5 * signal.iter().map(|v| v * v).sum::<f64>() / epsilon;
            let rel = (value - target).abs() / target.max(1e-20);
            assert!(
                rel < 5e-3,
                "scalar Charbonnier should match 0.5*t^2/eps locally: eps={epsilon}, value={value}, target={target}, rel={rel}"
            );
        }
    }

    #[test]
    fn grouped_charbonnier_small_signal_matches_local_half_quadratic() {
        let blocks = array![[1e-5, -2e-5], [3e-5, 4e-5]];
        for &epsilon in &[1e-3, 1e-1, 1.0, 1e2] {
            let state = CharbonnierGroupedBlockState::from_signal_blocks(blocks.clone(), epsilon);
            let value = state.penalty_value();
            let target = 0.5 * blocks.iter().map(|v| v * v).sum::<f64>() / epsilon;
            let rel = (value - target).abs() / target.max(1e-20);
            assert!(
                rel < 5e-3,
                "grouped Charbonnier should match 0.5*||v||^2/eps locally: eps={epsilon}, value={value}, target={target}, rel={rel}"
            );
        }
    }

    #[test]
    fn adaptive_diagnostics_json_roundtrip_preserves_shapes() {
        let diag = AdaptiveRegularizationDiagnostics {
            epsilon_0: 0.01,
            epsilon_g: 0.02,
            epsilon_c: 0.03,
            epsilon_outer_iterations: 2,
            mm_iterations: 3,
            converged: true,
            maps: vec![AdaptiveSpatialMap {
                termname: "matern".to_string(),
                feature_cols: vec![0, 1],
                collocation_points: array![[1.0, 2.0], [3.0, 4.0]],
                inv_magweight: array![0.05, 0.15],
                invgradweight: array![0.1, 0.2],
                inv_lapweight: array![0.3, 0.4],
            }],
        };
        let payload = serde_json::to_value(&diag).expect("serialize diagnostics");
        assert_eq!(payload["mm_iterations"].as_u64(), Some(3));
        assert_eq!(
            payload["maps"][0]["collocation_points"]["dim"]
                .as_array()
                .map(|v| v.len()),
            Some(2)
        );
        let decoded: AdaptiveRegularizationDiagnostics =
            serde_json::from_value(payload).expect("deserialize diagnostics");
        assert_eq!(decoded.epsilon_outer_iterations, 2);
        assert_eq!(decoded.mm_iterations, 3);
        assert!(decoded.converged);
        assert_eq!(decoded.maps.len(), 1);
        assert_eq!(decoded.maps[0].collocation_points.nrows(), 2);
        assert_eq!(decoded.maps[0].collocation_points.ncols(), 2);
        assert_eq!(decoded.maps[0].invgradweight.len(), 2);
        assert_eq!(decoded.maps[0].inv_lapweight.len(), 2);
    }
}
