semiflow 0.15.0-beta

//! ADR-0123 — `SmolyakGridND<F, const D: usize>` sparse-grid quadrature backend.
//!
//! Smolyak combination over the in-tree GH ladder {1,3,5,7,9}-pt.
//! Default level `ℓ = D+3`: D=5 → ℓ=8 → 341 nodes (tensor 5⁵=3125, 9.2× reduction).
//! NO new dependency. Weights carry SIGNS (combination coefficients).
//! Gate: `G_SMOLYAK_D5` — slope ≤ −0.95, nodes < 3125, F(0)=I verified at construction.

use alloc::vec::Vec;
use core::marker::PhantomData;

// `f64` has inherent math methods only when `std` is linked; otherwise they come
// from `num_traits::Float` (libm). Test builds link `std` even without the
// feature (harness, dev-dependencies), hence the `allow`.
#[cfg(not(feature = "std"))]
#[allow(unused_imports)]
use num_traits::Float;

use crate::{
    chernoff::{ChernoffFunction, Growth},
    error::SemiflowError,
    float::{from_f64, SemiflowFloat},
    grid_nd::{GridFnND, GridND},
    grid_nd_sweep::{apply_shift_step, ShiftStep},
    scratch::ScratchPool,
    shift_nd::{
        build_cholesky_cache, SquareMatrix, GH1_NODES_F64, GH1_WEIGHTS_F64, GH3_NODES_F64,
        GH3_WEIGHTS_F64, GH5_NODES_F64, GH5_WEIGHTS_F64, GH7_NODES_F64, GH7_WEIGHTS_F64,
        GH9_NODES_F64, GH9_WEIGHTS_F64,
    },
};

// 1-D GH ladder: level ∈ {1..=5} → (q, nodes, weights).
struct GhLevel {
    q: usize,
    nodes: &'static [f64],
    weights: &'static [f64],
}

const GH_LADDER: [GhLevel; 5] = [
    GhLevel {
        q: 1,
        nodes: &GH1_NODES_F64,
        weights: &GH1_WEIGHTS_F64,
    },
    GhLevel {
        q: 3,
        nodes: &GH3_NODES_F64,
        weights: &GH3_WEIGHTS_F64,
    },
    GhLevel {
        q: 5,
        nodes: &GH5_NODES_F64,
        weights: &GH5_WEIGHTS_F64,
    },
    GhLevel {
        q: 7,
        nodes: &GH7_NODES_F64,
        weights: &GH7_WEIGHTS_F64,
    },
    GhLevel {
        q: 9,
        nodes: &GH9_NODES_F64,
        weights: &GH9_WEIGHTS_F64,
    },
];

/// `GH_LADDER` entry for 1-based level `lev` (1..=5); `None` otherwise.
fn gh_level(lev: usize) -> Option<&'static GhLevel> {
    if lev == 0 || lev > 5 {
        return None;
    }
    Some(&GH_LADDER[lev - 1])
}

/// C(n, k) for n,k ≤ 12; 0 for k > n.  Fits in i64 (max C(12,6)=924).
fn binom(n: usize, k: usize) -> i64 {
    if k > n {
        return 0;
    }
    // Use the smaller k for efficiency.
    let k = k.min(n - k);
    let mut result: i64 = 1;
    for i in 0..k {
        #[allow(clippy::cast_possible_wrap)]
        let num = (n - i) as i64;
        #[allow(clippy::cast_possible_wrap)]
        let den = (i + 1) as i64;
        result = result * num / den;
    }
    result
}

// ---------------------------------------------------------------------------
// Smolyak node/weight builder
// ---------------------------------------------------------------------------

/// Build the Smolyak sparse grid A(ell, D) as a list of (node, weight) pairs,
/// with duplicate nodes merged (nested rules share nodes across sub-rules).
///
/// Nodes rounded to 14 decimal places for deduplication; weights carry SIGNS.
pub(crate) fn build_smolyak<const D: usize>(
    ell: usize,
) -> Result<Vec<([f64; D], f64)>, SemiflowError> {
    if ell < D {
        #[allow(clippy::cast_precision_loss)]
        return Err(SemiflowError::DomainViolation {
            what: "SmolyakGridND: ell must be >= D",
            value: ell as f64,
        });
    }
    let mut raw: Vec<([f64; D], f64)> = Vec::new();
    let mut stack = Vec::new();
    enumerate_multi_index::<D>(ell, &mut stack, &mut raw)?;

    // Sort by node lexicographically, then merge coincident nodes.
    raw.sort_by(|a, b| {
        for d in 0..D {
            let c = a.0[d].total_cmp(&b.0[d]);
            if c != core::cmp::Ordering::Equal {
                return c;
            }
        }
        core::cmp::Ordering::Equal
    });
    let nodes_eq = |a: &[f64; D], b: &[f64; D]| -> bool {
        a.iter().zip(b.iter()).all(|(x, y)| x.total_cmp(y).is_eq())
    };
    let mut merged: Vec<([f64; D], f64)> = Vec::new();
    for (node, w) in raw {
        if let Some(last) = merged.last_mut() {
            if nodes_eq(&last.0, &node) {
                last.1 += w;
                continue;
            }
        }
        merged.push((node, w));
    }
    // Drop near-zero combined weights (cancellation artifact).
    merged.retain(|(_, w)| w.abs() > 1e-15);

    Ok(merged)
}

/// Fill `raw` with all (tensor-product node, weight) pairs for admissible multi-indices.
fn enumerate_multi_index<const D: usize>(
    ell: usize,
    partial: &mut Vec<usize>,
    raw: &mut Vec<([f64; D], f64)>,
) -> Result<(), SemiflowError> {
    let depth = partial.len();
    if depth == D {
        let s: usize = partial.iter().sum();
        if s < ell.saturating_sub(D - 1) || s > ell {
            return Ok(());
        }
        let diff = ell - s;
        let coeff_abs = binom(D - 1, diff);
        if coeff_abs == 0 {
            return Ok(());
        }
        let sign: f64 = if diff % 2 == 0 { 1.0 } else { -1.0 };
        #[allow(clippy::cast_precision_loss)]
        let coeff = sign * coeff_abs as f64;
        let rules: Vec<&GhLevel> = partial
            .iter()
            .map(|&lev| {
                #[allow(clippy::cast_precision_loss)]
                gh_level(lev).ok_or(SemiflowError::DomainViolation {
                    what: "SmolyakGridND: level exceeds GH ladder depth (max 5)",
                    value: lev as f64,
                })
            })
            .collect::<Result<Vec<_>, _>>()?;
        add_tensor_product::<D>(&rules, coeff, raw);
        return Ok(());
    }
    let remaining = D - depth;
    for lev in 1..=ell.min(5) {
        let partial_sum: usize = partial.iter().sum();
        if partial_sum + lev + (remaining - 1) > ell {
            break;
        }
        partial.push(lev);
        enumerate_multi_index::<D>(ell, partial, raw)?;
        partial.pop();
    }
    Ok(())
}

/// Add all tensor-product (node, weight) pairs from `rules` scaled by `coeff` into `raw`.
fn add_tensor_product<const D: usize>(
    rules: &[&GhLevel],
    coeff: f64,
    raw: &mut Vec<([f64; D], f64)>,
) {
    // Total number of tensor combinations.
    let total: usize = rules.iter().map(|r| r.q).product();
    for qi in 0..total {
        let mut node = [0.0f64; D];
        let mut w_prod = coeff;
        let mut rem = qi;
        for d in 0..D {
            let q = rules[d].q;
            let digit = rem % q;
            rem /= q;
            // Round to 14 decimal places for stable deduplication.
            node[d] = (rules[d].nodes[digit] * 1e14).round() / 1e14;
            w_prod *= rules[d].weights[digit];
        }
        raw.push((node, w_prod));
    }
}

// ---------------------------------------------------------------------------
// SmolyakGridND<F, D>
// ---------------------------------------------------------------------------

/// Boxed drift closure type alias (suppress `type_complexity` lint).
type DriftFn<F, const D: usize> = alloc::boxed::Box<dyn Fn(&[F; D], &mut [F; D]) + Send + Sync>;
/// Boxed reaction closure type alias (suppress `type_complexity` lint).
type ReactionFn<F, const D: usize> = alloc::boxed::Box<dyn Fn(&[F; D]) -> F + Send + Sync>;

/// Sparse-grid Smolyak quadrature backend for d-D anisotropic shift (ADR-0123).
///
/// Implements the same `ChernoffFunction<F>` interface as
/// `AnisotropicShiftChernoffND` but uses a Smolyak sparse grid instead of the
/// full tensor product. For D=5 at the default level ℓ=D+3=8 this gives
/// **341 quadrature nodes** vs 3125 for the tensor rule — a 9.2× reduction.
///
/// Weights carry SIGNS. The `F(0)=I` unit witness is verified at construction.
///
/// # Construction
///
/// ```rust,ignore
/// # use semiflow::{Grid1D, smolyak::SmolyakGridND, grid_nd::GridND};
/// # let ax = Grid1D::new(-5.0_f64, 5.0, 8).unwrap();
/// # let grid = GridND::<f64, 5>::new([ax; 5]).unwrap();
/// let kernel = SmolyakGridND::<f64, 5>::new(
///     |_x, a| { for i in 0..5 { a.set(i, i, 1.0); } },
///     |_x, b| { for v in b.iter_mut() { *v = 0.0; } },
///     |_x| 0.0_f64,
///     grid,
/// ).unwrap();
/// ```
pub struct SmolyakGridND<F: SemiflowFloat = f64, const D: usize = 5> {
    /// Drift vector `b_i(x)`.
    b_i: DriftFn<F, D>,
    /// Reaction coefficient `c(x)`.
    c: ReactionFn<F, D>,
    /// Cached Cholesky factors at each grid point.
    cholesky_cache: Vec<SquareMatrix<F, D>>,
    /// Smolyak sparse-grid nodes and signed weights `(η_q, w_q)` (raw η values,
    /// combination coefficients × GH weights), cast to `F` once at construction so
    /// `apply_into` stays allocation-free.
    quad: Vec<([F; D], F)>,
    /// Smolyak level ℓ (default: D+3).
    level: usize,
    /// Reference grid geometry.
    grid: GridND<F, D>,
    _f: PhantomData<F>,
}

impl<F: SemiflowFloat, const D: usize> SmolyakGridND<F, D> {
    /// Construct with default Smolyak level `ℓ = D + 3`.
    ///
    /// # Errors
    /// - `DomainViolation` if `D == 0`.
    /// - `DomainViolation` if grid has fewer than 2 nodes per axis.
    /// - `DomainViolation` if `a_ij` is not SPD at any grid point.
    pub fn new(
        a_ij: impl Fn(&[F; D], &mut SquareMatrix<F, D>) + Send + Sync + 'static,
        b_i: impl Fn(&[F; D], &mut [F; D]) + Send + Sync + 'static,
        c: impl Fn(&[F; D]) -> F + Send + Sync + 'static,
        grid: GridND<F, D>,
    ) -> Result<Self, SemiflowError> {
        Self::with_level(a_ij, b_i, c, grid, D + 3)
    }

    /// Construct with explicit Smolyak level `ell`.
    ///
    /// # Errors
    /// Same as `new` plus `DomainViolation` if `ell < D` or `ell - D + 1 > 5`
    /// (level exceeds the in-tree GH ladder depth of 5).
    pub fn with_level(
        a_ij: impl Fn(&[F; D], &mut SquareMatrix<F, D>) + Send + Sync + 'static,
        b_i: impl Fn(&[F; D], &mut [F; D]) + Send + Sync + 'static,
        c: impl Fn(&[F; D]) -> F + Send + Sync + 'static,
        grid: GridND<F, D>,
        ell: usize,
    ) -> Result<Self, SemiflowError> {
        if D == 0 {
            return Err(SemiflowError::DomainViolation {
                what: "SmolyakGridND: D must be >= 1",
                value: 0.0,
            });
        }
        // Build Smolyak nodes/weights (validates ell ≥ D, level ≤ 5).
        let nw = build_smolyak::<D>(ell)?;
        let n_nodes = nw.len();
        if n_nodes == 0 {
            #[allow(clippy::cast_precision_loss)]
            return Err(SemiflowError::DomainViolation {
                what: "SmolyakGridND: empty node set (ell too small?)",
                value: ell as f64,
            });
        }
        // Verify F(0)=I unit witness: Σ weights = π^{D/2}.
        let wsum: f64 = nw.iter().map(|(_, w)| *w).sum();
        let pi_dhalf = 1.0_f64 / crate::float::inv_pi_pow_half::<f64>(D);
        let rel_err = (wsum - pi_dhalf).abs() / pi_dhalf;
        if rel_err > 1e-10 {
            return Err(SemiflowError::DomainViolation {
                what: "SmolyakGridND: weight sum ≠ π^{D/2} (F(0)=I violated)",
                value: rel_err,
            });
        }

        let cholesky_cache = build_cholesky_cache(&a_ij, &grid)?;
        let quad = cast_quad(&nw);
        Ok(Self {
            b_i: alloc::boxed::Box::new(b_i),
            c: alloc::boxed::Box::new(c),
            cholesky_cache,
            quad,
            level: ell,
            grid,
            _f: PhantomData,
        })
    }

    /// Number of Smolyak quadrature nodes (sparse grid size).
    pub fn n_nodes(&self) -> usize {
        self.quad.len()
    }

    /// Smolyak level `ℓ` used by this kernel.
    pub fn level(&self) -> usize {
        self.level
    }

    /// Return a shared reference to the kernel's grid geometry.
    pub fn grid(&self) -> &GridND<F, D> {
        &self.grid
    }
}

// ---------------------------------------------------------------------------
// ChernoffFunction<F> impl — math §32.4 algorithm with Smolyak quadrature
// ---------------------------------------------------------------------------

/// The Smolyak table `(η_q, w_q)` cast to `F` (once, at construction).
fn cast_quad<F: SemiflowFloat, const D: usize>(nw: &[([f64; D], f64)]) -> Vec<([F; D], F)> {
    nw.iter()
        .map(|(node, w)| {
            (
                core::array::from_fn(|d| from_f64::<F>(node[d])),
                from_f64::<F>(*w),
            )
        })
        .collect()
}

impl<F: SemiflowFloat, const D: usize> ChernoffFunction<F> for SmolyakGridND<F, D> {
    type S = GridFnND<F, D>;

    /// Apply one Chernoff step using the Smolyak sparse-grid quadrature.
    ///
    /// Same formula as `AnisotropicShiftChernoffND::apply_into` (math §32.4
    /// eq 32.3) but the sum over `q_idx` uses the Smolyak sparse grid instead
    /// of the full tensor product. Weights may be negative (combination signs).
    fn apply_into(
        &self,
        tau: F,
        src: &GridFnND<F, D>,
        dst: &mut GridFnND<F, D>,
        scratch: &mut ScratchPool<F>,
    ) -> Result<(), SemiflowError> {
        if !tau.is_finite() || tau < F::zero() {
            return Err(SemiflowError::DomainViolation {
                what: "SmolyakGridND::apply_into: tau must be finite >= 0",
                value: tau.to_f64().unwrap_or(f64::NAN),
            });
        }
        // y_q = x_k + τ·b(x_k) + 2√τ·L_k·η_q, node scale 2√τ and normalisation π^{−D/2}
        // as in AnisotropicShiftChernoffND (ADR-0112); constant coefficients take the
        // bit-identical separable path (ADR-0214).
        let step = ShiftStep {
            tau,
            grid: &self.grid,
            b: &*self.b_i,
            c: &*self.c,
            chol: &self.cholesky_cache,
            quad: &self.quad,
        };
        apply_shift_step(&step, src, dst, scratch);
        Ok(())
    }

    /// Nominal Chernoff order 1: same constructor-level declaration as
    /// `AnisotropicShiftChernoffND`.
    ///
    /// Heavy variable-coefficient gates may still measure only global order 1/2
    /// on specific data, for the same `2√τ` frozen-coefficient reason documented
    /// by the dense anisotropic D=5 gate; callers must not infer every
    /// variable-`A` datum will slope-test at -1.
    fn order(&self) -> u32 {
        1
    }

    /// Growth bound: multiplier 1.5, omega 0 (consistent with ADR-0081).
    fn growth(&self) -> Growth<F> {
        Growth::new(from_f64::<F>(1.5_f64), F::zero())
    }
}

// ---------------------------------------------------------------------------
// Inline unit tests — extracted to sibling file per ≤500-line cap.
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "smolyak_tests.rs"]
mod tests;