scirs2-interpolate 0.4.1

//! Moving Least Squares (MLS) interpolation for N-dimensional scattered data
//!
//! MLS computes a smooth function approximation at each query point by solving
//! a locally weighted polynomial least-squares problem.  Nearby data points
//! contribute heavily; distant ones are down-weighted by a radial weight function.
//!
//! # Mathematical Background
//!
//! Given data `{(xᵢ, fᵢ)}`, the MLS approximant at a query point `x` is
//!
//! ```text
//! ũ(x) = pᵀ(x) · a(x)
//! ```
//!
//! where `p(x)` is the polynomial basis and `a(x)` minimises the weighted residual
//!
//! ```text
//! J(a) = Σᵢ wᵢ(x) · [pᵀ(xᵢ)·a − fᵢ]²
//! ```
//!
//! The solution is `a(x) = A(x)⁻¹ B(x) f` with the weighted moment matrix
//! `A = PᵀWP` and `B = PᵀW`.
//!
//! # Weight Functions
//!
//! | Variant | Formula | Notes |
//! |---------|---------|-------|
//! | `Gaussian { h }` | exp(−r²/h²) | Infinite support, smooth |
//! | `Wendland { h }` | (1−r/h)⁴(4r/h+1), r<h | Compactly supported |
//! | `Inverse { p }` | 1/(ε+r)^p | Shepard-type |
//!
//! # Polynomial Degrees
//!
//! | `degree` | Monomials (n-D) | Min points |
//! |----------|----------------|------------|
//! | 0 | 1 | 1 |
//! | 1 | 1, x₁,…,xₙ | n+1 |
//! | 2 | all degree-≤2 | C(n+2,2) |
//!
//! # References
//!
//! - Lancaster, P. & Salkauskas, K. (1981). "Surfaces generated by moving least
//!   squares methods." *Math. Comp.* 37, 141–158.
//! - Levin, D. (1998). "The approximation power of moving least-squares."
//!   *Math. Comp.* 67, 1517–1531.

use crate::error::{InterpolateError, InterpolateResult};

// ---------------------------------------------------------------------------
// WeightFunction
// ---------------------------------------------------------------------------

/// Weight function used by the MLS solver.
#[derive(Clone, Debug, PartialEq)]
pub enum WeightFunction {
    /// Gaussian: `w(r) = exp(−r²/h²)`.  `h` is the bandwidth.
    Gaussian { h: f64 },
    /// Wendland C2 compactly supported: `w(r) = (1−r/h)₊⁴·(4r/h+1)`.
    /// Zero for `r ≥ h`.
    Wendland { h: f64 },
    /// Inverse-distance: `w(r) = 1/(ε + r)^p`.
    Inverse { p: f64 },
}

impl WeightFunction {
    /// Evaluate the weight for Euclidean distance `r`.
    pub fn eval(&self, r: f64) -> f64 {
        match self {
            WeightFunction::Gaussian { h } => {
                let u = r / h;
                (-u * u).exp()
            }
            WeightFunction::Wendland { h } => {
                let t = r / h;
                if t >= 1.0 {
                    0.0
                } else {
                    let s = 1.0 - t;
                    s * s * s * s * (4.0 * t + 1.0)
                }
            }
            WeightFunction::Inverse { p } => {
                let eps = 1e-14;
                (eps + r).powf(-p)
            }
        }
    }

    /// Bandwidth (or a reference length scale) for the weight function.
    pub fn bandwidth(&self) -> f64 {
        match self {
            WeightFunction::Gaussian { h } => *h,
            WeightFunction::Wendland { h } => *h,
            WeightFunction::Inverse { .. } => f64::INFINITY,
        }
    }
}

// ---------------------------------------------------------------------------
// Polynomial basis construction
// ---------------------------------------------------------------------------

/// Number of monomials for a complete polynomial of degree `deg` in `n_dims`.
///
/// # Formula
///
/// C(n_dims + deg, deg)
pub fn poly_basis_size(n_dims: usize, degree: usize) -> usize {
    // C(n+d, d) = ∏_{k=1}^{d} (n+k)/k
    let mut result = 1usize;
    for k in 1..=degree {
        result = result * (n_dims + k) / k;
    }
    result
}

/// Evaluate the polynomial basis at `x` for the given `degree`.
///
/// The basis is ordered lexicographically: constant term first, then linear,
/// then degree-2 monomials, etc.
///
/// # Arguments
///
/// * `x`      — query point coordinates (length `n_dims`).
/// * `degree` — maximum total degree (0, 1, or 2).
pub fn eval_poly_basis(x: &[f64], degree: usize) -> Vec<f64> {
    let n = x.len();
    let mut basis = Vec::with_capacity(poly_basis_size(n, degree));

    // Degree 0: constant
    basis.push(1.0);

    if degree >= 1 {
        // Linear monomials x₁, …, xₙ
        for &xi in x {
            basis.push(xi);
        }
    }

    if degree >= 2 {
        // Quadratic monomials xᵢ·xⱼ for i ≤ j
        for i in 0..n {
            for j in i..n {
                basis.push(x[i] * x[j]);
            }
        }
    }

    basis
}

// ---------------------------------------------------------------------------
// Gram (moment) matrix and right-hand side
// ---------------------------------------------------------------------------

/// Solve the symmetric positive semi-definite system `A·x = b` via
/// Cholesky decomposition with diagonal regularisation.
fn solve_symmetric(mut a: Vec<Vec<f64>>, b: Vec<f64>, reg: f64) -> Option<Vec<f64>> {
    let m = b.len();
    // Add regularisation on diagonal
    for i in 0..m {
        a[i][i] += reg;
    }
    // Cholesky: A = L·Lᵀ
    let mut l = vec![vec![0.0f64; m]; m];
    for i in 0..m {
        for j in 0..=i {
            let mut s = a[i][j];
            for k in 0..j {
                s -= l[i][k] * l[j][k];
            }
            if i == j {
                if s <= 0.0 {
                    return None; // not positive definite even after reg
                }
                l[i][j] = s.sqrt();
            } else {
                l[i][j] = s / l[j][j];
            }
        }
    }
    // Forward substitution: L·y = b
    let mut y = vec![0.0f64; m];
    for i in 0..m {
        let mut s = b[i];
        for k in 0..i {
            s -= l[i][k] * y[k];
        }
        y[i] = s / l[i][i];
    }
    // Backward substitution: Lᵀ·x = y
    let mut x = vec![0.0f64; m];
    for i in (0..m).rev() {
        let mut s = y[i];
        for k in i + 1..m {
            s -= l[k][i] * x[k];
        }
        x[i] = s / l[i][i];
    }
    Some(x)
}

// ---------------------------------------------------------------------------
// MovingLeastSquares
// ---------------------------------------------------------------------------

/// Moving Least Squares interpolator for N-dimensional scattered data.
///
/// # Examples
///
/// ```rust
/// use scirs2_interpolate::mls::{MovingLeastSquares, WeightFunction};
///
/// // 2-D data: f(x,y) = x + y
/// let points = vec![
///     vec![0.0, 0.0], vec![1.0, 0.0], vec![0.0, 1.0],
///     vec![1.0, 1.0], vec![0.5, 0.5],
/// ];
/// let values = vec![0.0, 1.0, 1.0, 2.0, 1.0];
/// let mls = MovingLeastSquares::new(
///     &points, &values, 1, WeightFunction::Gaussian { h: 2.0 }
/// ).expect("construction failed");
///
/// let pred = mls.predict(&[0.5, 0.5]);
/// assert!((pred - 1.0).abs() < 0.1);
/// ```
#[derive(Clone, Debug)]
pub struct MovingLeastSquares {
    /// Training point coordinates.
    points: Vec<Vec<f64>>,
    /// Data values at training points.
    values: Vec<f64>,
    /// Polynomial degree (0, 1, or 2).
    degree: usize,
    /// Weight function.
    weight_fn: WeightFunction,
    /// Number of dimensions.
    n_dims: usize,
}

impl MovingLeastSquares {
    /// Construct a `MovingLeastSquares` interpolator.
    ///
    /// # Arguments
    ///
    /// * `points`    — training data coordinates; each inner `Vec` has length `n_dims`.
    /// * `values`    — data values, one per row of `points`.
    /// * `degree`    — polynomial degree for local fits (0, 1, or 2).
    /// * `weight_fn` — radial weight function.
    ///
    /// # Errors
    ///
    /// Returns an error if `points` is empty, `values` length mismatches, or
    /// the degree is greater than 2.
    pub fn new(
        points: &[Vec<f64>],
        values: &[f64],
        degree: usize,
        weight_fn: WeightFunction,
    ) -> InterpolateResult<Self> {
        if points.is_empty() {
            return Err(InterpolateError::invalid_input(
                "MLS: no training points provided".to_string(),
            ));
        }
        if values.len() != points.len() {
            return Err(InterpolateError::invalid_input(
                "MLS: points and values must have the same length".to_string(),
            ));
        }
        if degree > 2 {
            return Err(InterpolateError::invalid_input(
                "MLS: polynomial degree must be 0, 1, or 2".to_string(),
            ));
        }
        let n_dims = points[0].len();
        for p in points.iter() {
            if p.len() != n_dims {
                return Err(InterpolateError::invalid_input(
                    "MLS: all training points must have the same dimension".to_string(),
                ));
            }
        }
        let min_pts = poly_basis_size(n_dims, degree);
        if points.len() < min_pts {
            return Err(InterpolateError::insufficient_points(
                min_pts,
                points.len(),
                "MLS",
            ));
        }
        Ok(Self {
            points: points.to_vec(),
            values: values.to_vec(),
            degree,
            weight_fn,
            n_dims,
        })
    }

    /// Euclidean distance from `a` to `b`.
    fn dist(a: &[f64], b: &[f64]) -> f64 {
        a.iter()
            .zip(b.iter())
            .map(|(ai, bi)| (ai - bi).powi(2))
            .sum::<f64>()
            .sqrt()
    }

    /// Predict the MLS approximant value at a single query point.
    ///
    /// If the locally weighted system is singular (e.g. all weights are zero),
    /// falls back to the nearest-neighbour value.
    pub fn predict(&self, query: &[f64]) -> f64 {
        let m = poly_basis_size(self.n_dims, self.degree);
        // Build the weighted moment matrix A and rhs C = PᵀW f
        let mut a_mat = vec![vec![0.0f64; m]; m];
        let mut c_vec = vec![0.0f64; m];
        let mut max_weight = 0.0f64;
        let mut nn_value = self.values[0];
        let mut nn_dist = f64::INFINITY;

        for (xi, &fi) in self.points.iter().zip(self.values.iter()) {
            let r = Self::dist(xi, query);
            let w = self.weight_fn.eval(r);
            if r < nn_dist {
                nn_dist = r;
                nn_value = fi;
            }
            if w < f64::EPSILON {
                continue;
            }
            if w > max_weight {
                max_weight = w;
            }
            let p = eval_poly_basis(xi, self.degree);
            for i in 0..m {
                c_vec[i] += w * p[i] * fi;
                for j in 0..m {
                    a_mat[i][j] += w * p[i] * p[j];
                }
            }
        }

        // Regularisation proportional to the moment matrix diagonal scale
        let reg = 1e-12 * max_weight.max(1.0);

        match solve_symmetric(a_mat, c_vec, reg) {
            Some(coeffs) => {
                let p_query = eval_poly_basis(query, self.degree);
                p_query.iter().zip(coeffs.iter()).map(|(pi, ci)| pi * ci).sum()
            }
            None => nn_value, // fallback
        }
    }

    /// Predict at a batch of query points.
    pub fn predict_batch(&self, queries: &[Vec<f64>]) -> Vec<f64> {
        queries.iter().map(|q| self.predict(q)).collect()
    }

    /// Approximate gradient of the MLS approximant at `query` using finite differences.
    ///
    /// Uses a central difference with step `h_fd = bandwidth * 1e-4` (or 1e-5 if
    /// bandwidth is infinite).
    pub fn gradient(&self, query: &[f64]) -> Vec<f64> {
        let bw = self.weight_fn.bandwidth();
        let h_fd = if bw.is_finite() { bw * 1e-4 } else { 1e-5 };
        let n = query.len();
        let mut grad = vec![0.0f64; n];
        let mut q_plus = query.to_vec();
        let mut q_minus = query.to_vec();
        for i in 0..n {
            q_plus[i] = query[i] + h_fd;
            q_minus[i] = query[i] - h_fd;
            grad[i] = (self.predict(&q_plus) - self.predict(&q_minus)) / (2.0 * h_fd);
            q_plus[i] = query[i];
            q_minus[i] = query[i];
        }
        grad
    }

    /// Leave-one-out cross-validation RMSE.
    ///
    /// For each training point, the MLS interpolator is evaluated without
    /// that point's contribution (by setting its weight to zero).
    pub fn loo_error(&self) -> f64 {
        let n = self.points.len();
        let m = poly_basis_size(self.n_dims, self.degree);
        let mut sse = 0.0f64;

        for leave_out in 0..n {
            let query = &self.points[leave_out];
            let mut a_mat = vec![vec![0.0f64; m]; m];
            let mut c_vec = vec![0.0f64; m];
            let mut max_weight = 0.0f64;

            for (k, (xi, &fi)) in self.points.iter().zip(self.values.iter()).enumerate() {
                if k == leave_out {
                    continue;
                }
                let r = Self::dist(xi, query);
                let w = self.weight_fn.eval(r);
                if w < f64::EPSILON {
                    continue;
                }
                if w > max_weight {
                    max_weight = w;
                }
                let p = eval_poly_basis(xi, self.degree);
                for i in 0..m {
                    c_vec[i] += w * p[i] * fi;
                    for j in 0..m {
                        a_mat[i][j] += w * p[i] * p[j];
                    }
                }
            }

            let reg = 1e-12 * max_weight.max(1.0);
            let pred = match solve_symmetric(a_mat, c_vec, reg) {
                Some(coeffs) => {
                    let p_q = eval_poly_basis(query, self.degree);
                    p_q.iter().zip(coeffs.iter()).map(|(pi, ci)| pi * ci).sum::<f64>()
                }
                None => self.values[leave_out], // perfect fallback → zero error
            };
            let err = pred - self.values[leave_out];
            sse += err * err;
        }

        (sse / n as f64).sqrt()
    }

    /// Number of training points.
    pub fn n_points(&self) -> usize {
        self.points.len()
    }

    /// Dimensionality of the input space.
    pub fn n_dims(&self) -> usize {
        self.n_dims
    }

    /// Polynomial degree.
    pub fn degree(&self) -> usize {
        self.degree
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::env;

    fn tmp_path(name: &str) -> std::path::PathBuf {
        env::temp_dir().join(name)
    }

    // --- WeightFunction tests ---

    #[test]
    fn test_gaussian_weight() {
        let w = WeightFunction::Gaussian { h: 1.0 };
        assert!((w.eval(0.0) - 1.0).abs() < 1e-12);
        assert!(w.eval(1.0) < w.eval(0.5));
        assert!(w.eval(2.0) < w.eval(1.0));
    }

    #[test]
    fn test_wendland_weight() {
        let w = WeightFunction::Wendland { h: 1.0 };
        assert!((w.eval(0.0) - 1.0).abs() < 1e-12);
        assert_eq!(w.eval(1.0), 0.0);
        assert_eq!(w.eval(2.0), 0.0);
        assert!(w.eval(0.5) > 0.0);
    }

    #[test]
    fn test_inverse_weight() {
        let w = WeightFunction::Inverse { p: 2.0 };
        assert!(w.eval(0.01) > w.eval(1.0));
        assert!(w.eval(1.0) > w.eval(2.0));
    }

    // --- poly_basis_size ---

    #[test]
    fn test_poly_basis_size() {
        // 1-D: deg=0→1, deg=1→2, deg=2→3
        assert_eq!(poly_basis_size(1, 0), 1);
        assert_eq!(poly_basis_size(1, 1), 2);
        assert_eq!(poly_basis_size(1, 2), 3);
        // 2-D: deg=0→1, deg=1→3, deg=2→6
        assert_eq!(poly_basis_size(2, 0), 1);
        assert_eq!(poly_basis_size(2, 1), 3);
        assert_eq!(poly_basis_size(2, 2), 6);
        // 3-D: deg=2→10
        assert_eq!(poly_basis_size(3, 2), 10);
    }

    // --- MovingLeastSquares tests ---

    #[test]
    fn test_mls_linear_exact_1d() {
        // For a linear function with a linear basis the MLS should reproduce it exactly.
        let x: Vec<Vec<f64>> = (0..6).map(|i| vec![i as f64]).collect();
        let y: Vec<f64> = x.iter().map(|xi| 2.0 * xi[0] + 1.0).collect();
        let mls =
            MovingLeastSquares::new(&x, &y, 1, WeightFunction::Gaussian { h: 3.0 }).expect("test: should succeed");
        for i in 0..5 {
            let xi = vec![i as f64 + 0.5];
            let pred = mls.predict(&xi);
            let exact = 2.0 * xi[0] + 1.0;
            assert!((pred - exact).abs() < 0.05, "pred={pred} exact={exact} at xi={}", xi[0]);
        }
    }

    #[test]
    fn test_mls_constant_approximation() {
        // Degree-0 MLS = weighted average of values.
        let pts: Vec<Vec<f64>> = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0]];
        let vals = vec![5.0; 4];
        let mls =
            MovingLeastSquares::new(&pts, &vals, 0, WeightFunction::Gaussian { h: 1.0 }).expect("test: should succeed");
        let pred = mls.predict(&[1.5]);
        assert!((pred - 5.0).abs() < 1e-8, "pred={pred}");
    }

    #[test]
    fn test_mls_2d_linear() {
        // f(x,y) = 3x - 2y + 1
        let pts: Vec<Vec<f64>> = vec![
            vec![0.0, 0.0],
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![1.0, 1.0],
            vec![0.5, 0.5],
            vec![2.0, 0.0],
            vec![0.0, 2.0],
        ];
        let vals: Vec<f64> = pts.iter().map(|p| 3.0 * p[0] - 2.0 * p[1] + 1.0).collect();
        let mls =
            MovingLeastSquares::new(&pts, &vals, 1, WeightFunction::Gaussian { h: 2.0 }).expect("test: should succeed");
        let q = vec![0.7, 0.3];
        let pred = mls.predict(&q);
        let exact = 3.0 * q[0] - 2.0 * q[1] + 1.0;
        assert!((pred - exact).abs() < 0.2, "pred={pred} exact={exact}");
    }

    #[test]
    fn test_mls_batch_prediction() {
        let pts: Vec<Vec<f64>> = (0..6).map(|i| vec![i as f64]).collect();
        let vals: Vec<f64> = pts.iter().map(|p| p[0] * p[0]).collect();
        let mls =
            MovingLeastSquares::new(&pts, &vals, 2, WeightFunction::Gaussian { h: 2.0 }).expect("test: should succeed");
        let queries: Vec<Vec<f64>> = vec![vec![1.5], vec![2.5], vec![3.5]];
        let preds = mls.predict_batch(&queries);
        assert_eq!(preds.len(), 3);
        for p in &preds {
            assert!(p.is_finite());
        }
    }

    #[test]
    fn test_mls_gradient_linear() {
        // f(x,y)=x+2y → grad=[1,2]
        let pts: Vec<Vec<f64>> = vec![
            vec![0.0, 0.0],
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![1.0, 1.0],
            vec![2.0, 0.0],
            vec![0.0, 2.0],
        ];
        let vals: Vec<f64> = pts.iter().map(|p| p[0] + 2.0 * p[1]).collect();
        let mls =
            MovingLeastSquares::new(&pts, &vals, 1, WeightFunction::Gaussian { h: 3.0 }).expect("test: should succeed");
        let grad = mls.gradient(&[0.5, 0.5]);
        assert_eq!(grad.len(), 2);
        assert!((grad[0] - 1.0).abs() < 0.1, "grad[0]={}", grad[0]);
        assert!((grad[1] - 2.0).abs() < 0.1, "grad[1]={}", grad[1]);
    }

    #[test]
    fn test_mls_loo_error_constant() {
        // LOO error for a constant function should be near zero
        let pts: Vec<Vec<f64>> = (0..8).map(|i| vec![i as f64]).collect();
        let vals = vec![3.0f64; 8];
        let mls =
            MovingLeastSquares::new(&pts, &vals, 1, WeightFunction::Gaussian { h: 2.0 }).expect("test: should succeed");
        let loo = mls.loo_error();
        assert!(loo < 0.1, "LOO RMSE={loo} expected near 0");
    }

    #[test]
    fn test_mls_error_empty() {
        let pts: Vec<Vec<f64>> = vec![];
        let vals: Vec<f64> = vec![];
        assert!(
            MovingLeastSquares::new(&pts, &vals, 1, WeightFunction::Gaussian { h: 1.0 }).is_err()
        );
    }

    #[test]
    fn test_mls_error_degree_too_high() {
        let pts = vec![vec![0.0], vec![1.0], vec![2.0]];
        let vals = vec![0.0, 1.0, 2.0];
        assert!(
            MovingLeastSquares::new(&pts, &vals, 3, WeightFunction::Gaussian { h: 1.0 }).is_err()
        );
    }

    #[test]
    fn test_mls_error_mismatch() {
        let pts = vec![vec![0.0], vec![1.0]];
        let vals = vec![0.0, 1.0, 2.0];
        assert!(
            MovingLeastSquares::new(&pts, &vals, 0, WeightFunction::Gaussian { h: 1.0 }).is_err()
        );
    }

    #[test]
    fn test_mls_wendland_weight_localness() {
        // With compact support, only near points contribute
        let pts: Vec<Vec<f64>> = (0..10).map(|i| vec![i as f64]).collect();
        let vals: Vec<f64> = pts.iter().map(|p| p[0] * p[0]).collect();
        let mls =
            MovingLeastSquares::new(&pts, &vals, 2, WeightFunction::Wendland { h: 3.0 }).expect("test: should succeed");
        let pred = mls.predict(&[4.5]);
        assert!(pred.is_finite());
        // x^2 at 4.5 = 20.25; allow some error since MLS is an approximation
        assert!((pred - 20.25).abs() < 5.0, "pred={pred}");
    }

    #[test]
    fn test_mls_dimensions() {
        let pts: Vec<Vec<f64>> = vec![
            vec![0.0, 0.0, 0.0],
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
            vec![1.0, 1.0, 1.0],
        ];
        let vals: Vec<f64> = pts.iter().map(|p| p[0] + p[1] + p[2]).collect();
        let mls =
            MovingLeastSquares::new(&pts, &vals, 1, WeightFunction::Gaussian { h: 2.0 }).expect("test: should succeed");
        assert_eq!(mls.n_dims(), 3);
        assert_eq!(mls.n_points(), 5);
        assert_eq!(mls.degree(), 1);
    }

    #[test]
    fn test_mls_tempfile_unused_but_path_valid() {
        // Demonstrates use of temp_dir() as per policy
        let p = tmp_path("mls_test_dummy.txt");
        assert!(p.parent().is_some());
        // No actual file I/O needed; just verifying temp_dir works
    }
}