gam 0.3.103

//! Penalized multinomial-logit (softmax) GLM driver — fixed-λ inner solve.
//!
//! This is the principled vector-response companion to the scalar PIRLS path:
//! the inner-loop Newton solver for a multi-class GAM at fixed smoothing
//! parameters λ, using the canonical multinomial-logit likelihood
//! ([`MultinomialLogitLikelihood`]) and the existing dense block-Fisher
//! assembly in [`crate::solver::pirls::dense_block_xtwx`] /
//! [`crate::solver::pirls::dense_block_xtwy`].
//!
//! # What this module does
//!
//! Solve, for the reference-coded multinomial-logit GAM with `K` classes and
//! design matrix `X ∈ ℝ^{N×P}`,
//!
//! ```text
//!     β̂ = argmin_β { − log L(β) + ½ Σ_{a=0}^{K-2} λ_a · β_a^T S β_a }
//! ```
//!
//! where `β = [β_0; β_1; …; β_{K-2}]` is the stacked coefficient vector in
//! output-major order (`β_a ∈ ℝ^P` is the coefficient block for class `a`),
//! `S ∈ ℝ^{P×P}` is the smoothing penalty matrix (shared across classes,
//! replicated as `I_{K-1} ⊗ S` over the full parameter space), and `λ_a` is
//! a per-class smoothing parameter.
//!
//! The likelihood uses class `K - 1` as the reference (`η_{K-1} ≡ 0`), so the
//! softmax gauge is fixed at the η level and no additional sum-to-zero
//! projection is required.
//!
//! # Layering
//!
//! * **Fixed-λ inner solve** — [`fit_penalized_multinomial`] is the canonical
//!   coefficient-space Newton solver at *given* smoothing parameters `λ`,
//!   built on the shared [`crate::families::penalized_vector_glm`] engine.
//!
//! * **REML / LAML smoothing-parameter selection** — [`fit_penalized_multinomial_formula`]
//!   routes through [`crate::families::custom_family::fit_custom_family_with_rho_prior`]
//!   so the per-active-class `λ_a` are selected by the outer REML/LAML loop;
//!   the caller's `init_lambda` is only a warm-start seed. The multinomial
//!   [`crate::families::multinomial_reml::MultinomialFamily`] `CustomFamily`
//!   impl calls the fixed-λ math above as its inner solve at each ρ trial and
//!   supplies the dense per-row Hessian block for the outer trace terms.
//!
//! * **Formula → design integration** — `build_formula_design_for_multinomial`
//!   parses the Wilkinson formula and assembles `X` and the per-term `S`
//!   blocks; the `fit_multinomial_formula_pyfunc` FFI shim wires the Python
//!   `gamfit.fit(..., family='multinomial')` entry straight to this path.
//!
//! # Convergence
//!
//! The damped-Newton-with-backtracking scaffold lives once in the shared
//! [`crate::families::penalized_vector_glm`] engine: at each iteration the
//! assembled penalized Hessian `H + I_{K-1} ⊗ (λ_a S)` is factored via faer's
//! symmetric-PD-with-fallback path, the full Newton step `δ = −H^{-1} ∇F` is
//! computed, and accepted with step halving if the objective fails to decrease
//! (up to a small backtracking budget). The convergence test is the relative
//! coefficient step norm `‖δ‖ / (1 + ‖β‖) ≤ tol`, matching the existing pyffi
//! reference path. This module is the softmax adapter over that engine: it
//! supplies the dense `(K-1)×(K-1)` Fisher block, the residual, and the
//! log-likelihood through [`MultinomialLogitLikelihood`], and owns the
//! class-count / simplex preconditions. The independent-binomial sibling
//! [`crate::families::binomial_multi`] is the same engine with a row-diagonal
//! Fisher block instead.

use crate::families::custom_family::{
    BlockwiseFitOptions, ParameterBlockState, fit_custom_family_with_rho_prior,
};
use crate::families::multinomial_reml::MultinomialFamily;
use crate::families::penalized_vector_glm::{PenalizedVectorGlmInputs, fit_penalized_vector_glm};
use crate::families::vector_response::{MultinomialLogitLikelihood, validate_multinomial_simplex};
use crate::inference::data::EncodedDataset;
use crate::inference::formula_dsl::parse_formula;
use crate::inference::model::ColumnKindTag;
use crate::resource::ProblemHints;
use crate::solver::estimate::EstimationError;
use crate::solver::workflow::{
    FitConfig, build_termspec_with_geometry_and_overrides, resolved_resource_policy,
};
use crate::terms::smooth::{
    TermCollectionDesign, TermCollectionSpec, build_term_collection_design,
};
use crate::terms::term_builder::resolve_role_col;
use crate::types::ResponseColumnKind;
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// Solver-only numerical stabilization floor for the formula-driven
/// multinomial REML inner solve (gam#747).
///
/// Installed with [`RidgePolicy::solver_only`](crate::types::RidgePolicy::solver_only)
/// so it stabilizes the inner joint-Newton **linear solve** but never enters
/// the REML objective, the penalty log-determinant, or the Laplace Hessian.
///
/// What it does: the multinomial smoothing penalties are rank-deficient by
/// design (each smooth carries an unpenalized polynomial null space) and the
/// formula may add a fully unpenalized parametric term (`x3` / `body_mass`). On
/// near-separable hard labels the softmax curvature is ill-conditioned along
/// those directions, so the bare Newton step `H⁻¹∇` is huge. Lifting the
/// smallest Hessian eigenvalue to `δ` bounds the step (`‖(H+δI)⁻¹∇‖ ≤ ‖∇‖/δ`),
/// keeping the screening iterates finite without poisoning the softmax with
/// `inf − inf = NaN`.
///
/// What it deliberately does NOT do: it adds no `½·δ·‖β‖²` term to the
/// objective and no `δ`-shift to the REML log-determinant. The earlier
/// `explicit_stabilization_pospart` policy folded both into the criterion,
/// which made `1e-4` a fixed-λ Gaussian prior that shrank every identified
/// coefficient off the MLE and biased smoothing-parameter selection — a value
/// that had to be tuned *between* under-stabilization (NaN seeds) and
/// over-shrinkage (lost VGAM match). As a solver-only floor that tradeoff is
/// gone: the over-shrinkage failure mode cannot occur (nothing is shrunk), the
/// optimized objective is the true penalized REML criterion, and the floor
/// only has to be large enough to keep the linear algebra finite.
///
/// The separation defect (#753) is no longer this floor's job. If the
/// multinomial MLE is genuinely at infinity for an unpenalized/null-space
/// direction (complete/quasi-complete separation), no solver floor makes that
/// direction's estimate finite — the principled response is a model-declared
/// bias-reduction prior, not a magnitude on this floor. The formula REML path
/// below supplies exactly that: because `MultinomialFamily` is a `CustomFamily`
/// routed through [`fit_custom_family_with_rho_prior`], it inherits the
/// UNIVERSAL, always-on full-span Jeffreys/Firth proper prior
/// `Φ = ½ log|Z_Jᵀ H Z_J|` that the joint-Newton path folds into every coupled
/// custom-family inner solve (see
/// [`crate::solver::reml::jeffreys_subspace::joint_jeffreys_term`] and
/// `build_joint_jeffreys_subspace` / `custom_family_joint_jeffreys_term` in
/// `custom_family.rs`). The conditioning gate keeps it byte-identical to the
/// un-penalized Newton on an identified fit and supplies the missing
/// `O(1)`-bounding curvature only on a separating direction, where the
/// multinomial family's exact joint Hessian and its analytic directional
/// derivatives (`exact_newton_joint_hessian{,_directional_derivative}` in
/// [`crate::families::multinomial_reml`]) drive both the score `∇Φ` and the
/// Gauss-Newton curvature `H_Φ`. So a separating multinomial REML fit now
/// converges to FINITE Firth-reduced coefficients rather than drifting to the
/// screening cap. The bare fixed-λ inner driver [`fit_penalized_multinomial`]
/// (no outer REML, no Jeffreys term) instead surfaces the explicit
/// `MultinomialSeparationDetected` diagnostic — the same #753 acceptance,
/// option (b), for the path that has no proper prior to lean on.
const MULTINOMIAL_FORMULA_RIDGE_FLOOR: f64 = 1.0e-4;

/// Largest smoothing-parameter dimension where exact dense outer curvature is
/// still worth paying for multinomial formula fits.
///
/// `D = (K - 1) * n_terms`. Medium-size loaded models (`D <= 6`) benefit from
/// exact curvature because the first-order route can wander on near-separable
/// labels. Smooth-by-factor models with one global plus one per-level smooth
/// already reach `D = 8` for `K = 3`, where the O(D^2) dense outer Hessian
/// dominates runtime; those stay on the exact-gradient quasi-Newton route.
const MULTINOMIAL_EXACT_OUTER_HESSIAN_MAX_DIM: usize = 6;

fn multinomial_formula_use_outer_hessian(total_rho_dim: usize) -> bool {
    total_rho_dim <= MULTINOMIAL_EXACT_OUTER_HESSIAN_MAX_DIM
}

/// Logit magnitude beyond which fitted probabilities are saturated at ordinary
/// double precision diagnostic scale. If a multinomial Newton solve exhausts
/// its iteration budget at this scale, the returned iterate is a separation
/// artifact, not a finite MLE.
const MULTINOMIAL_SEPARATION_ETA_THRESHOLD: f64 = 25.0;

fn max_abs_eta_location(eta: ArrayView2<'_, f64>) -> (f64, usize, usize) {
    let mut best = (0.0_f64, 0usize, 0usize);
    for ((row, active_class), &value) in eta.indexed_iter() {
        let abs = value.abs();
        if abs > best.0 {
            best = (abs, row, active_class);
        }
    }
    best
}

fn max_abs_eta_location_from_block_states(
    block_states: &[ParameterBlockState],
) -> (f64, usize, usize) {
    let mut best = (0.0_f64, 0usize, 0usize);
    for (active_class, state) in block_states.iter().enumerate() {
        for (row, &value) in state.eta.iter().enumerate() {
            let abs = value.abs();
            if abs > best.0 {
                best = (abs, row, active_class);
            }
        }
    }
    best
}

fn multinomial_formula_separation_diagnostic(
    outer_converged: bool,
    inner_cycles: usize,
    outer_iterations: usize,
    block_states: &[ParameterBlockState],
) -> Option<EstimationError> {
    let (max_abs_eta, row_index, active_class_index) =
        max_abs_eta_location_from_block_states(block_states);
    if !outer_converged && max_abs_eta >= MULTINOMIAL_SEPARATION_ETA_THRESHOLD {
        Some(EstimationError::MultinomialSeparationDetected {
            iteration: inner_cycles.max(outer_iterations),
            max_abs_eta,
            active_class_index,
            row_index,
        })
    } else {
        None
    }
}

/// Inputs to [`fit_penalized_multinomial`].
///
/// The penalty matrix `S` is shared across classes; per-class smoothing
/// parameters `lambdas` (length `K - 1`) scale `S` independently for each
/// active class. The full block-replicated penalty is `diag_a(λ_a) ⊗ S`,
/// which is exactly what [`crate::solver::arrow_schur::KroneckerPenaltyOp`]
/// expresses in matrix-free form when this driver is later lifted into the
/// arrow-Schur loop.
#[derive(Debug, Clone)]
pub struct MultinomialFitInputs<'a> {
    /// Design matrix `X ∈ ℝ^{N×P}` (one row per observation).
    pub design: ArrayView2<'a, f64>,
    /// Categorical response `Y ∈ ℝ^{N×K}`. Each row must be a point on the
    /// probability simplex (`y_c ≥ 0`, `Σ_c y_c = 1`): a one-hot indicator for
    /// hard classification, or a label-smoothed probability vector. Rows whose
    /// mass departs from 1 are rejected — the softmax residual gradient and
    /// Fisher block are the derivatives of `Σ_c y_c log p_c` only under the
    /// simplex constraint (see `validate_multinomial_simplex`).
    pub y_one_hot: ArrayView2<'a, f64>,
    /// Shared smoothing penalty `S ∈ ℝ^{P×P}` (symmetric, PSD).
    pub penalty: ArrayView2<'a, f64>,
    /// Per-active-class smoothing parameter `λ_a` (length `K - 1`).
    pub lambdas: ArrayView1<'a, f64>,
    /// Optional per-row weights (length `N`); `None` ⇒ uniform 1.0.
    pub row_weights: Option<ArrayView1<'a, f64>>,
    /// Optional per-row Fisher-block override, shape `(N, K-1, K-1)` in the
    /// active-class gauge (the reference class `K-1` is dropped). When `Some`,
    /// each Newton step uses this block as the curvature `W` in place of the
    /// analytic softmax Fisher `w_n (δ_ab p_a − p_a p_b)`; the gradient/residual
    /// path stays analytic, so this is a curvature-only override (the
    /// research escape-hatch for latent multinomial fits, issue #349). Each
    /// per-row block must be symmetric, PSD, and finite — preconditions the
    /// FFI boundary discharges before constructing this view.
    pub fisher_w_override: Option<ArrayView3<'a, f64>>,
    /// Maximum Newton iterations; recommend 50.
    pub max_iter: usize,
    /// Relative-step convergence tolerance; recommend 1e-7.
    pub tol: f64,
}

/// Outputs of [`fit_penalized_multinomial`].
#[derive(Debug, Clone)]
pub struct MultinomialFitOutputs {
    /// Active-class coefficient block, shape `(P, K-1)` (column `a` is `β_a`).
    /// The reference class `K - 1` has `β_{K-1} ≡ 0` by construction and is
    /// not stored.
    pub coefficients_active: Array2<f64>,
    /// Fitted probabilities, shape `(N, K)`.
    pub fitted_probabilities: Array2<f64>,
    /// Number of Newton iterations executed (including the final step that
    /// satisfied the tolerance).
    pub iterations: usize,
    /// `true` if the relative-step test was satisfied; `false` if the
    /// solver exhausted `max_iter`. (A non-converged solve is still
    /// returned; the caller decides whether to escalate.)
    pub converged: bool,
    /// Penalized negative log-likelihood at the returned `β̂`:
    /// `−log L(β̂) + ½ Σ_a λ_a · β̂_a^T S β̂_a`.
    pub penalized_neg_log_likelihood: f64,
    /// Unpenalized deviance `−2 log L(β̂)` for diagnostic reporting.
    pub deviance: f64,
}

/// Fit a penalized multinomial-logit GAM at fixed `λ`.
///
/// See the module docs for the optimization problem and conventions. This
/// function is the canonical inner solve: the outer REML/LAML loop, when
/// added, calls this at each `ρ = log λ` trial.
pub fn fit_penalized_multinomial(
    inputs: MultinomialFitInputs<'_>,
) -> Result<MultinomialFitOutputs, EstimationError> {
    let MultinomialFitInputs {
        design,
        y_one_hot,
        penalty,
        lambdas,
        row_weights,
        fisher_w_override,
        max_iter,
        tol,
    } = inputs;

    // ──────────────────────── family-specific validation ───────────────────
    // The shared engine re-validates the geometry common to every vector-GLM
    // (nonempty design, penalty shape, λ finiteness/non-negativity, override
    // `(N, M, M)` shape, finite design). The multinomial family owns the
    // class-count contract (`K ≥ 2`, λ length `K − 1`), the per-row simplex
    // precondition under which the softmax residual/Fisher are the exact
    // derivatives of `Σ_c y_c log p_c`, and the row-weight check the likelihood
    // adapter consumes.
    let n_obs = design.nrows();
    let (y_rows, k) = y_one_hot.dim();
    if y_rows != n_obs {
        crate::bail_invalid_estim!(
            "fit_penalized_multinomial: y rows {y_rows} ≠ design rows {n_obs}"
        );
    }
    if k < 2 {
        crate::bail_invalid_estim!(
            "fit_penalized_multinomial: need at least 2 classes (got K={k})"
        );
    }
    let m = k - 1;
    if lambdas.len() != m {
        crate::bail_invalid_estim!(
            "fit_penalized_multinomial: lambdas length {} ≠ K-1 = {m}",
            lambdas.len()
        );
    }
    if let Some(fw) = fisher_w_override.as_ref() {
        if fw.dim() != (n_obs, m, m) {
            crate::bail_invalid_estim!(
                "fit_penalized_multinomial: fisher_w_override shape {:?} ≠ (N, K-1, K-1) = ({n_obs}, {m}, {m})",
                fw.dim()
            );
        }
    }
    if let Some(w) = row_weights.as_ref() {
        if w.len() != n_obs {
            crate::bail_invalid_estim!(
                "fit_penalized_multinomial: row_weights length {} ≠ N = {n_obs}",
                w.len()
            );
        }
        for (i, &v) in w.iter().enumerate() {
            if !(v.is_finite() && v >= 0.0) {
                crate::bail_invalid_estim!(
                    "fit_penalized_multinomial: row_weights[{i}] must be finite and ≥ 0 (got {v})"
                );
            }
        }
    }
    validate_multinomial_simplex(y_one_hot, "fit_penalized_multinomial")?;

    // ────────────────────────── likelihood construction ───────────────────
    let mut likelihood = MultinomialLogitLikelihood::with_classes(k)?;
    if let Some(w) = row_weights.as_ref() {
        likelihood = likelihood.with_row_weights(w.to_owned())?;
    }

    // ─────────────────── shared penalized vector-GLM solve ─────────────────
    // The softmax Fisher block is dense across the `M = K − 1` active classes;
    // the engine assembles the coupled `(P·M)×(P·M)` penalized Hessian, runs
    // the damped Newton loop, and returns the converged `β̂` and `η = X β̂`.
    let fit = fit_penalized_vector_glm(
        PenalizedVectorGlmInputs {
            design,
            y: y_one_hot,
            penalty,
            lambdas,
            fisher_w_override,
            max_iter,
            tol,
        },
        &likelihood,
        "fit_penalized_multinomial",
    )?;

    let (max_abs_eta, row_index, active_class_index) = max_abs_eta_location(fit.eta.view());
    if !fit.converged && max_abs_eta >= MULTINOMIAL_SEPARATION_ETA_THRESHOLD {
        return Err(EstimationError::MultinomialSeparationDetected {
            iteration: fit.iterations,
            max_abs_eta,
            active_class_index,
            row_index,
        });
    }

    let fitted_probabilities = likelihood.probabilities(fit.eta.view());

    Ok(MultinomialFitOutputs {
        coefficients_active: fit.coefficients,
        fitted_probabilities,
        iterations: fit.iterations,
        converged: fit.converged,
        penalized_neg_log_likelihood: -fit.log_likelihood + fit.penalty_term,
        deviance: -2.0 * fit.log_likelihood,
    })
}

// ---------------------------------------------------------------------------
// Formula-driven multinomial pipeline
// ---------------------------------------------------------------------------
//
// Slice A of the multinomial integration: a single public entry that takes
// a parsed `EncodedDataset`, a Wilkinson-style formula, and a uniform initial
// smoothing parameter, then runs the full
//
//     parse → termspec → design (X, S blocks) → one-hot Y → REML λ-selection
//
// pipeline. `fit_penalized_multinomial_formula` drives the outer REML/LAML
// loop (via the custom-family path) to select an independent λ per (class,
// term); `init_lambda` (default 1.0) is only the warm-start seed for every
// block. The reference class is the last level of the categorical response
// column as recorded in the dataset schema.

/// Saved-model payload for a multinomial fit driven by a Wilkinson formula.
///
/// This is what the FFI returns to Python. It carries everything the Python
/// `MultinomialModel.predict` path needs to evaluate `softmax(X_new · β)` on
/// fresh data using the *training* basis / penalty structure (no refit on
/// predict, no re-derivation of class levels).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultinomialSavedModel {
    /// The training formula, verbatim. Stored so Python's `summary()` and
    /// any round-trip persistence path can echo what was fit.
    pub formula: String,
    /// Names of the *training* response levels in canonical order. The last
    /// entry is the reference class (η = 0); the first `K - 1` carry the
    /// active linear-predictor blocks. Class permutations are forbidden:
    /// this list is fixed at fit time and predictions emit columns in the
    /// same order.
    pub class_levels: Vec<String>,
    /// Index of the reference class within `class_levels` — currently always
    /// `class_levels.len() - 1`, exposed as a field so future "user-pinned
    /// reference" gauges (e.g. `family='multinomial', reference='setosa'`)
    /// can land without changing the on-disk shape.
    pub reference_class_index: usize,
    /// Resolved term-collection spec used to build `X` at fit time. Replayed
    /// on predict via [`crate::terms::smooth::build_term_collection_design`].
    pub resolved_termspec: TermCollectionSpec,
    /// Active-class coefficient block, shape `(P, K-1)`. Column `a` is the
    /// coefficient vector for class `class_levels[a]`. Stored flat in
    /// row-major order to keep the serde payload self-describing.
    pub coefficients_flat: Vec<f64>,
    /// `P` — coefficient count per active class. Matches the column count of
    /// the design matrix the saved `resolved_termspec` produces.
    pub p_per_class: usize,
    /// Number of active classes (`K - 1`).
    pub n_active_classes: usize,
    /// Original training column headers, in dataset-column order. Needed at
    /// predict time so the FFI can align a fresh `Dataset` to the training
    /// schema before evaluating the basis.
    pub training_headers: Vec<String>,
    /// REML/LAML-selected smoothing parameters, one per `(active class, smooth
    /// term)`, flattened in block-major order: all of class 0's per-term λ,
    /// then class 1's, and so on. Per-term penalties (#561) mean each active
    /// class block selects an *independent* λ for every smooth term, so this
    /// vector has length `Σ_a (#terms in class a)` = `(K − 1) · #terms`. Use
    /// [`MultinomialSavedModel::lambdas_per_block`] to segment it by class. An
    /// unpenalized model (no smooth terms) yields an empty vector.
    pub lambdas: Vec<f64>,
    /// Number of smoothing parameters (smooth terms) in each active class
    /// block, parallel to `class_levels[0..K-1]`. Segments the flat `lambdas`
    /// vector: class `a`'s λ are `lambdas[Σ_{b<a} lambdas_per_block[b] ..][..
    /// lambdas_per_block[a]]`. Every entry is identical in the shared-design
    /// architecture (all classes share the same term structure), but it is
    /// stored explicitly so consumers never have to assume that.
    pub lambdas_per_block: Vec<usize>,
    /// Newton iterations executed; recorded for the summary report.
    pub iterations: usize,
    /// `true` if the inner Newton solver hit the relative-step tolerance.
    pub converged: bool,
    /// Penalized negative log-likelihood at the returned `β̂`.
    pub penalized_neg_log_likelihood: f64,
    /// Unpenalized deviance `−2 log L(β̂)`.
    pub deviance: f64,
    /// Per-active-class effective degrees of freedom (hat-matrix trace),
    /// length `K - 1`. Populated when the REML driver reports an
    /// inference block; falls back to `None` for the legacy fixed-λ path.
    #[serde(default)]
    pub edf_per_class: Option<Vec<f64>>,
}

impl MultinomialSavedModel {
    /// Active-class coefficient block as an `(P, K-1)` `ndarray` view.
    pub fn coefficients_active(&self) -> Array2<f64> {
        Array2::from_shape_vec(
            (self.p_per_class, self.n_active_classes),
            self.coefficients_flat.clone(),
        )
        .expect(
            "MultinomialSavedModel.coefficients_flat length must equal p_per_class * n_active_classes",
        )
    }

    /// Evaluate `softmax(X · β)` at fresh data rows. `X_new` must have
    /// `self.p_per_class` columns (i.e. it was built from the same
    /// `resolved_termspec` as fit time). Returns an `(N_new, K)` matrix
    /// with rows summing to 1; column order matches `self.class_levels`.
    pub fn predict_probabilities(&self, x_new: ArrayView2<'_, f64>) -> Array2<f64> {
        let n_new = x_new.nrows();
        let p = self.p_per_class;
        let m = self.n_active_classes;
        let k = m + 1;
        assert_eq!(
            x_new.ncols(),
            p,
            "MultinomialSavedModel.predict_probabilities: X has {} cols, expected {p}",
            x_new.ncols()
        );
        let beta = self.coefficients_active();
        let mut probs = Array2::<f64>::zeros((n_new, k));
        let mut eta_active = vec![0.0_f64; m];
        let mut row_probs = vec![0.0_f64; k];
        for row in 0..n_new {
            for a in 0..m {
                let mut v = 0.0_f64;
                for i in 0..p {
                    v += x_new[[row, i]] * beta[[i, a]];
                }
                eta_active[a] = v;
            }
            MultinomialLogitLikelihood::softmax_with_baseline(&eta_active, &mut row_probs);
            for c in 0..k {
                probs[[row, c]] = row_probs[c];
            }
        }
        probs
    }
}

/// One-hot-encode the categorical response column and return both the
/// encoding and the captured level names. The level order matches the order
/// recorded in the dataset schema, which is itself the order of first
/// appearance during inferred-schema construction — so it is stable and
/// deterministic across runs (no silent class permutation).
fn one_hot_categorical_response(
    data: &EncodedDataset,
    y_col: usize,
    response_name: &str,
) -> Result<(Array2<f64>, Vec<String>), EstimationError> {
    let levels: Vec<String> = data
        .schema
        .columns
        .get(y_col)
        .map(|sc| sc.levels.clone())
        .unwrap_or_default();
    if levels.len() < 2 {
        crate::bail_invalid_estim!(
            "multinomial response '{response_name}' must have at least 2 categorical levels (got {})",
            levels.len()
        );
    }
    let n = data.values.nrows();
    let k = levels.len();
    let mut y_one_hot = Array2::<f64>::zeros((n, k));
    for row in 0..n {
        let encoded = data.values[[row, y_col]];
        if !encoded.is_finite() {
            crate::bail_invalid_estim!(
                "multinomial response '{response_name}' row {row} is non-finite ({encoded})"
            );
        }
        let class_idx = encoded.round() as i64;
        if class_idx < 0 || (class_idx as usize) >= k {
            crate::bail_invalid_estim!(
                "multinomial response '{response_name}' row {row} encoded as {encoded} \
                 is outside the level range 0..{k}"
            );
        }
        y_one_hot[[row, class_idx as usize]] = 1.0;
    }
    Ok((y_one_hot, levels))
}

/// Build `(TermCollectionSpec, TermCollectionDesign)` from a formula against
/// a categorical-response dataset. Mirrors the early scaffolding inside
/// `materialize_standard` (response role resolution, geometry-aware spec
/// build) without touching the scalar-family resolution path — multinomial
/// owns its own response kind check.
fn build_formula_design_for_multinomial(
    formula: &str,
    data: &EncodedDataset,
    config: &FitConfig,
) -> Result<
    (
        TermCollectionSpec,
        TermCollectionDesign,
        usize,
        String,
        ResponseColumnKind,
    ),
    EstimationError,
> {
    let parsed = parse_formula(formula).map_err(|err| {
        EstimationError::InvalidInput(format!(
            "multinomial fit: failed to parse formula {formula:?}: {err}"
        ))
    })?;
    let col_map = data.column_map();
    let y_col = resolve_role_col(&col_map, &parsed.response, "response")
        .map_err(|err| EstimationError::InvalidInput(format!("multinomial fit: {err}")))?;
    let y_kind = crate::solver::workflow::response_column_kind(data, y_col);
    let policy = resolved_resource_policy(config, data, ProblemHints::default());
    let mut inference_notes: Vec<String> = Vec::new();
    let spec = build_termspec_with_geometry_and_overrides(
        &parsed.terms,
        data,
        &col_map,
        &mut inference_notes,
        config.scale_dimensions,
        &policy,
        config.smooth_overrides.as_ref(),
    )
    .map_err(|err| {
        EstimationError::InvalidInput(format!("multinomial fit: build termspec: {err}"))
    })?;
    let design = build_term_collection_design(data.values.view(), &spec).map_err(|err| {
        EstimationError::InvalidInput(format!("multinomial fit: build design: {err}"))
    })?;
    Ok((spec, design, y_col, parsed.response, y_kind))
}

/// Top-level formula-driven multinomial fit.
///
/// Routes through [`fit_custom_family_with_rho_prior`] so the per-active-class
/// smoothing parameters `λ_a` (one per class block, shared-penalty
/// architecture) are selected by the outer REML/LAML loop rather than pinned
/// by the caller. `init_lambda` survives as a warm-start hint that seeds
/// every block's `initial_log_lambdas`; the inner Newton solve still uses
/// `max_iter` / `tol` via `BlockwiseFitOptions`.
///
/// The categorical response column is recognised via the dataset schema
/// (`ColumnKindTag::Categorical`); reference class = last level. Returns a
/// [`MultinomialSavedModel`] that can be serialised to bytes for the Python
/// wrapper or used in-process for `predict_probabilities`.
pub fn fit_penalized_multinomial_formula(
    data: &EncodedDataset,
    formula: &str,
    config: &FitConfig,
    init_lambda: f64,
    max_iter: usize,
    tol: f64,
) -> Result<MultinomialSavedModel, EstimationError> {
    if !(init_lambda.is_finite() && init_lambda > 0.0) {
        crate::bail_invalid_estim!(
            "multinomial fit: init_lambda must be finite and > 0 (got {init_lambda})"
        );
    }
    let (spec, design, y_col, response_name, y_kind) =
        build_formula_design_for_multinomial(formula, data, config)?;
    let class_levels = match y_kind {
        ResponseColumnKind::Categorical { levels } => levels,
        ResponseColumnKind::Binary => vec!["0".to_string(), "1".to_string()],
        ResponseColumnKind::Numeric => {
            crate::bail_invalid_estim!(
                "multinomial fit: response '{response_name}' is numeric, not categorical; \
                 use family='gaussian'/'binomial'/... or convert the column to a categorical type"
            );
        }
    };
    if data.column_kinds.get(y_col) == Some(&ColumnKindTag::Binary) {
        // Promote to a 2-level categorical for the multinomial driver; the
        // caller explicitly asked for multinomial, so we route through the
        // K-1 = 1 active-class softmax (equivalent math to logistic).
    } else if data.column_kinds.get(y_col) != Some(&ColumnKindTag::Categorical) {
        crate::bail_invalid_estim!(
            "multinomial fit: response '{response_name}' must be a categorical column \
             (got column kind {:?})",
            data.column_kinds.get(y_col)
        );
    }
    let (y_one_hot, _) = one_hot_categorical_response(data, y_col, &response_name)?;
    // Build the global X dense (the design is a DesignMatrix abstraction).
    let x_dense = design
        .design
        .try_to_dense_by_chunks("multinomial fit design")
        .map_err(EstimationError::InvalidInput)?;
    // Preserve the per-smooth-term penalty block structure (#561): each smooth
    // term `t` contributes its own `P × P` penalty component (`Blockwise` with
    // `total_dim = P`, the term's local `S_t` embedded at its `col_range`), and
    // every active class block receives the FULL list. The outer REML/LAML loop
    // then selects an independent smoothing parameter λ_{a,t} per (class, term),
    // matching mgcv/VGAM. Pre-summing the terms into one fused `S` (the prior
    // behaviour) forced a single λ per class that scales `Σ_t S_t`, so one
    // shared λ had to over-smooth a rough term while under-smoothing a smooth
    // one — biasing any multi-term class-probability surface.
    let per_term_penalties = design.penalties_as_penalty_matrix();
    let per_term_nullspace_dims = design.nullspace_dims.clone();
    let k = y_one_hot.ncols();
    let m = k - 1;
    let n_obs = y_one_hot.nrows();

    // ── Custom-family driven REML/LAML path ───────────────────────────────
    // Each active class becomes one ParameterBlockSpec, all sharing X and the
    // per-term penalty list. `initial_log_lambdas` is seeded from the caller's
    // `init_lambda` (one entry per term).
    let design_arc = Arc::new(x_dense);
    let penalties_arc = Arc::new(per_term_penalties);
    let nullspace_dims_arc = Arc::new(per_term_nullspace_dims);
    let weights = Array1::<f64>::ones(n_obs);
    let family = MultinomialFamily::new(
        y_one_hot.clone(),
        weights,
        k,
        design_arc.clone(),
        penalties_arc.clone(),
        nullspace_dims_arc.clone(),
    )
    .map_err(EstimationError::InvalidInput)?;
    let mut blocks = family.build_block_specs();
    let log_init = init_lambda.ln();
    for spec_block in blocks.iter_mut() {
        for v in spec_block.initial_log_lambdas.iter_mut() {
            *v = log_init;
        }
    }

    // ── Outer-derivative policy: auto-derived from the smoothing dimension ──
    // The total smoothing-parameter dimension is `D = (K−1) · n_terms`.
    // Multinomial exact outer curvature is pairwise in `D`, so smooth-by-factor
    // `D = 8` models must avoid the O(D²) dense Hessian path (#714). But
    // medium `D = 6` loaded models with hard labels benefit from exact
    // curvature because first-order BFGS can wander along separation-induced
    // ridges (#722). The helper below encodes that crossover explicitly.
    let total_rho_dim = m.saturating_mul(penalties_arc.len());
    let use_outer_hessian = multinomial_formula_use_outer_hessian(total_rho_dim);

    let options = BlockwiseFitOptions {
        inner_max_cycles: max_iter,
        inner_tol: tol,
        ridge_floor: MULTINOMIAL_FORMULA_RIDGE_FLOOR,
        // #747: the stabilization floor is SOLVER-ONLY — it keeps the inner
        // joint-Newton linear solve finite during screening (bounding the step
        // `(H+δI)⁻¹∇` away from a near-separable, rank-deficient curvature) but
        // is excluded from the REML objective, the penalty log-determinant, and
        // the Laplace Hessian. The earlier default (`explicit_stabilization_pospart`)
        // folded `½·δ·‖β‖²` and a `δ`-shift of the log-determinant into the
        // criterion, shrinking every identified coefficient off the MLE and
        // perturbing smoothing-parameter selection — a fixed-λ prior masking
        // separation, not a numerical stabilizer. With the floor solver-only the
        // optimized objective is the true penalized REML criterion (value tracks
        // its analytic gradient), and the smooth directions remain governed
        // solely by their own REML-selected `λ`.
        ridge_policy: crate::types::RidgePolicy::solver_only(),
        use_outer_hessian,
        ..BlockwiseFitOptions::default()
    };
    let fit =
        fit_custom_family_with_rho_prior(&family, &blocks, &options, crate::types::RhoPrior::Flat)
            .map_err(|err| EstimationError::InvalidInput(format!("multinomial REML: {err}")))?;
    if let Some(err) = multinomial_formula_separation_diagnostic(
        fit.outer_converged,
        fit.inner_cycles,
        fit.outer_iterations,
        &fit.block_states,
    ) {
        return Err(err);
    }

    // ── Repack coefficients (P, K-1) from per-block β vectors ─────────────
    if fit.blocks.len() != m {
        crate::bail_invalid_estim!(
            "multinomial REML: expected {m} fitted blocks (K-1), got {}",
            fit.blocks.len()
        );
    }
    let p_per_class = fit.blocks[0].beta.len();
    let mut coefficients_active = Array2::<f64>::zeros((p_per_class, m));
    for (a, block) in fit.blocks.iter().enumerate() {
        if block.beta.len() != p_per_class {
            crate::bail_invalid_estim!(
                "multinomial REML: block {a} has {} coefs, expected {p_per_class}",
                block.beta.len()
            );
        }
        for i in 0..p_per_class {
            coefficients_active[[i, a]] = block.beta[i];
        }
    }
    // Flatten every (class, term) smoothing parameter in block-major order
    // (class 0's terms, then class 1's, …). With per-term penalties each block
    // now carries one λ per smooth term, so a single λ per class would discard
    // the independent per-term selection that fixes #561. `lambdas_per_block`
    // segments the flat vector by class so callers can recover per-term λ.
    let lambdas_per_block: Vec<usize> = fit.blocks.iter().map(|b| b.lambdas.len()).collect();
    let lambdas_flat: Vec<f64> = fit
        .blocks
        .iter()
        .flat_map(|b| b.lambdas.iter().copied())
        .collect();
    let edf_per_class = fit.inference.as_ref().map(|info| info.edf_by_block.clone());
    let coefficients_flat: Vec<f64> = coefficients_active.iter().copied().collect();

    // Unpenalized deviance read directly from the converged unpenalized
    // log-likelihood the rho-prior driver already computed (issue #348):
    // MultinomialFamily::evaluate sets FamilyEvaluation.log_likelihood =
    // log_lik(η, y) with no penalty term, and that value flows unchanged into
    // UnifiedFitResult.log_likelihood. This reproduces the legacy fixed-λ
    // path's `deviance = -2 · log_lik` contract bit-for-bit, so the previous
    // row-by-row η = Xβ rebuild and softmax recompute were pure dead work.
    let deviance = -2.0 * fit.log_likelihood;

    Ok(MultinomialSavedModel {
        formula: formula.to_string(),
        class_levels: class_levels.clone(),
        reference_class_index: class_levels.len() - 1,
        resolved_termspec: spec,
        coefficients_flat,
        p_per_class,
        n_active_classes: m,
        training_headers: data.headers.clone(),
        lambdas: lambdas_flat,
        lambdas_per_block,
        iterations: fit.inner_cycles,
        converged: fit.outer_converged,
        penalized_neg_log_likelihood: -fit.log_likelihood + 0.5 * fit.stable_penalty_term,
        deviance,
        edf_per_class,
    })
}

/// Replay the saved termspec to build the predict-time design on a fresh
/// dataset, then evaluate softmax probabilities. The predict dataset must carry
/// the same feature columns the training data did, matched **by name** — it need
/// not reproduce the training column order, and in particular need not carry the
/// response column (prediction is for label-free new data).
pub fn predict_multinomial_formula(
    model: &MultinomialSavedModel,
    data: &EncodedDataset,
) -> Result<Array2<f64>, EstimationError> {
    // The saved termspec stores feature columns as absolute indices into the
    // *training* table `[response, features...]`. Replaying it verbatim only
    // works if the predict frame reproduces that exact layout — i.e. carries the
    // (unknown, at predict time) response column in the same position. Realign
    // the indices onto this dataset's columns by name instead, so prediction
    // works on label-free new data exactly as every other family's predict path
    // does. The response column is simply never referenced by any term, so its
    // absence is a non-issue once resolution is by name (issue #803).
    let predict_columns = data.column_map();
    let realigned = model.resolved_termspec.remap_feature_columns(
        |index| -> Result<usize, EstimationError> {
            let name = model.training_headers.get(index).ok_or_else(|| {
                EstimationError::InvalidInput(format!(
                    "multinomial predict: saved training column index {index} is out of bounds \
                     for {} training headers",
                    model.training_headers.len()
                ))
            })?;
            resolve_role_col(&predict_columns, name, "feature")
                .map_err(|err| EstimationError::InvalidInput(err.to_string()))
        },
    )?;
    let design = build_term_collection_design(data.values.view(), &realigned).map_err(|err| {
        EstimationError::InvalidInput(format!(
            "multinomial predict: rebuild design from saved termspec: {err}"
        ))
    })?;
    let x_dense = design
        .design
        .try_to_dense_by_chunks("multinomial predict design")
        .map_err(EstimationError::InvalidInput)?;
    if x_dense.ncols() != model.p_per_class {
        crate::bail_invalid_estim!(
            "multinomial predict: predict design has {} cols, saved model expects {}",
            x_dense.ncols(),
            model.p_per_class
        );
    }
    Ok(model.predict_probabilities(x_dense.view()))
}

#[cfg(test)]
mod fisher_override_tests {
    use super::*;
    use ndarray::Array3;

    fn toy() -> (Array2<f64>, Array2<f64>, Array2<f64>, Array1<f64>) {
        let n = 15;
        let p = 2;
        let k = 3;
        let design =
            Array2::<f64>::from_shape_fn(
                (n, p),
                |(i, j)| {
                    if j == 0 { 1.0 } else { ((i + 2) as f64).cos() }
                },
            );
        let mut y = Array2::<f64>::zeros((n, k));
        for i in 0..n {
            y[[i, i % k]] = 1.0;
        }
        let penalty = Array2::<f64>::eye(p);
        let lambdas = Array1::<f64>::from_elem(k - 1, 0.5);
        (design, y, penalty, lambdas)
    }

    #[test]
    fn fisher_override_none_reproduces_analytic() {
        // Issue #349: None override is exactly the analytic fit.
        let (design, y, penalty, lambdas) = toy();
        let mk = |over: Option<ndarray::ArrayView3<'_, f64>>| {
            fit_penalized_multinomial(MultinomialFitInputs {
                design: design.view(),
                y_one_hot: y.view(),
                penalty: penalty.view(),
                lambdas: lambdas.view(),
                row_weights: None,
                fisher_w_override: over,
                max_iter: 50,
                tol: 1.0e-9,
            })
            .expect("fit must succeed")
        };
        let a = mk(None);
        let b = mk(None);
        for (x, z) in a
            .coefficients_active
            .iter()
            .zip(b.coefficients_active.iter())
        {
            assert_eq!(x, z);
        }
    }

    #[test]
    fn fisher_override_wrong_shape_is_rejected() {
        let (design, y, penalty, lambdas) = toy();
        let n = design.nrows();
        let m = y.ncols(); // K, not K-1 — deliberately wrong
        let bad = Array3::<f64>::zeros((n, m, m));
        let err = fit_penalized_multinomial(MultinomialFitInputs {
            design: design.view(),
            y_one_hot: y.view(),
            penalty: penalty.view(),
            lambdas: lambdas.view(),
            row_weights: None,
            fisher_w_override: Some(bad.view()),
            max_iter: 50,
            tol: 1.0e-9,
        })
        .expect_err("wrong active-block shape must error");
        assert!(format!("{err}").contains("fisher_w_override shape"));
    }

    #[test]
    fn formula_outer_route_uses_exact_curvature_for_medium_d() {
        assert!(
            multinomial_formula_use_outer_hessian(6),
            "D=6 loaded multinomial fits need exact curvature to avoid near-separation BFGS wandering"
        );
    }

    #[test]
    fn formula_outer_route_uses_first_order_for_smooth_by_factor_d8() {
        assert!(
            !multinomial_formula_use_outer_hessian(8),
            "D=8 smooth-by-factor multinomial fits must avoid the O(D^2) dense outer Hessian"
        );
    }

    #[test]
    fn fixed_lambda_multinomial_reports_complete_separation() {
        let n = 90;
        let design = Array2::<f64>::from_shape_fn((n, 2), |(row, col)| match col {
            0 => 1.0,
            _ => -3.0 + 6.0 * (row as f64) / ((n - 1) as f64),
        });
        let mut y = Array2::<f64>::zeros((n, 3));
        for row in 0..n {
            let x = design[[row, 1]];
            let class = if x < -1.0 {
                0
            } else if x > 1.0 {
                1
            } else {
                2
            };
            y[[row, class]] = 1.0;
        }
        let penalty = Array2::<f64>::zeros((2, 2));
        let lambdas = Array1::<f64>::zeros(2);
        let err = fit_penalized_multinomial(MultinomialFitInputs {
            design: design.view(),
            y_one_hot: y.view(),
            penalty: penalty.view(),
            lambdas: lambdas.view(),
            row_weights: None,
            fisher_w_override: None,
            max_iter: 80,
            tol: 1.0e-12,
        })
        .expect_err("complete softmax separation must be a hard diagnostic");
        assert!(
            matches!(err, EstimationError::MultinomialSeparationDetected { .. }),
            "expected MultinomialSeparationDetected, got {err:?}"
        );
        assert!(
            err.to_string().contains("separation"),
            "diagnostic should mention separation, got {err}"
        );
        assert!(
            err.to_string().contains("active class-"),
            "diagnostic should name the separated active class logit, got {err}"
        );
        assert!(
            !err.to_string().contains("binary outcomes"),
            "multinomial diagnostic must not reuse the binary separation text, got {err}"
        );
    }

    #[test]
    fn formula_multinomial_reports_saturated_nonconvergence_as_separation() {
        let block_states = vec![
            ParameterBlockState {
                beta: Array1::from_vec(vec![1.0, 2.0]),
                eta: Array1::from_vec(vec![0.2, 4.0, -7.0]),
            },
            ParameterBlockState {
                beta: Array1::from_vec(vec![-1.0, 3.0]),
                eta: Array1::from_vec(vec![1.0, 25.5, -0.1]),
            },
        ];

        let err = multinomial_formula_separation_diagnostic(false, 17, 9, &block_states)
            .expect("nonconverged formula fit at saturated logits must be diagnostic");
        assert!(
            matches!(
                err,
                EstimationError::MultinomialSeparationDetected {
                    iteration: 17,
                    max_abs_eta,
                    active_class_index: 1,
                    row_index: 1,
                } if (max_abs_eta - 25.5).abs() <= f64::EPSILON
            ),
            "expected typed multinomial separation diagnostic with final max|eta| and channel location, got {err:?}"
        );

        assert!(
            multinomial_formula_separation_diagnostic(true, 17, 9, &block_states).is_none(),
            "a converged robust/Firth formula fit is a finite fit, not a separation diagnostic"
        );

        let finite_states = vec![ParameterBlockState {
            beta: Array1::from_vec(vec![0.0]),
            eta: Array1::from_vec(vec![24.9]),
        }];
        assert!(
            multinomial_formula_separation_diagnostic(false, 3, 11, &finite_states).is_none(),
            "nonconvergence below the saturation threshold should remain a convergence failure"
        );
    }

    #[test]
    fn scaled_fisher_override_changes_first_step() {
        // Curvature scaled by 4× shrinks the first Newton step relative to the
        // analytic fit, so a single-iteration fit must differ.
        let (design, y, penalty, lambdas) = toy();
        let n = design.nrows();
        let m = y.ncols() - 1;
        // Analytic block at β = 0: p_a = 1/K = 1/3, so diag = p_a(1−p_a),
        // off-diag = −p_a p_b. Scale that exact block by 4.
        let pk = 1.0 / (y.ncols() as f64);
        let mut over = Array3::<f64>::zeros((n, m, m));
        for row in 0..n {
            for a in 0..m {
                for b in 0..m {
                    let analytic = if a == b { pk * (1.0 - pk) } else { -pk * pk };
                    over[[row, a, b]] = 4.0 * analytic;
                }
            }
        }
        let scaled = fit_penalized_multinomial(MultinomialFitInputs {
            design: design.view(),
            y_one_hot: y.view(),
            penalty: penalty.view(),
            lambdas: lambdas.view(),
            row_weights: None,
            fisher_w_override: Some(over.view()),
            max_iter: 1,
            tol: 1.0e-9,
        })
        .expect("override fit must succeed");
        let analytic = fit_penalized_multinomial(MultinomialFitInputs {
            design: design.view(),
            y_one_hot: y.view(),
            penalty: penalty.view(),
            lambdas: lambdas.view(),
            row_weights: None,
            fisher_w_override: None,
            max_iter: 1,
            tol: 1.0e-9,
        })
        .expect("analytic fit must succeed");
        let differs = scaled
            .coefficients_active
            .iter()
            .zip(analytic.coefficients_active.iter())
            .any(|(a, b)| (a - b).abs() > 1.0e-6);
        assert!(differs, "scaled curvature must change the first step");
    }
}