scirs2-spatial 0.4.0

//! Radial Basis Function interpolation
//!
//! This module provides Radial Basis Function (RBF) interpolation, a flexible
//! technique for interpolating scattered data in any number of dimensions.
//!
//! RBF interpolation works by representing the interpolant as a weighted sum of
//! radial basis functions centered at each data point. The weights are determined
//! by solving a linear system to enforce that the interpolant passes through all
//! data points.
//!
//! Various radial basis functions are provided, each with different smoothness
//! and locality properties:
//!
//! - Gaussian: Infinitely smooth but highly local
//! - Multiquadric: Moderately smooth and less local
//! - Inverse Multiquadric: Infinitely smooth with moderate locality
//! - Thin Plate Spline: Minimizes curvature, very smooth
//! - Linear: Simplest RBF, not smooth at data points
//! - Cubic: Good compromise between smoothness and locality

use crate::error::{SpatialError, SpatialResult};
use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use scirs2_core::random::{thread_rng, Rng};
// Simple linear system solver
#[allow(dead_code)]
fn solve_linear_system(a: Array2<f64>, b: Array1<f64>) -> SpatialResult<Array1<f64>> {
    // We should use a proper linear algebra library, but for now we'll use a simple approach
    // This is not numerically stable for ill-conditioned matrices
    let n = a.nrows();
    if n != a.ncols() {
        return Err(SpatialError::DimensionError(
            "Matrix A must be square".to_string(),
        ));
    }

    if n != b.len() {
        return Err(SpatialError::DimensionError(
            "Matrix A and vector b dimensions must match".to_string(),
        ));
    }

    // Very simple implementation - in production code, use a proper linear algebra library
    let mut x = Array1::zeros(n);

    // Add a small value to the diagonal to improve stability (regularization)
    let mut a_reg = a.clone();
    for i in 0..n {
        a_reg[[i, i]] += 1e-10;
    }

    // Simple Gaussian elimination - not suitable for large or ill-conditioned systems
    let mut aug = Array2::zeros((n, n + 1));
    for i in 0..n {
        for j in 0..n {
            aug[[i, j]] = a_reg[[i, j]];
        }
        aug[[i, n]] = b[i];
    }

    // Forward elimination
    for i in 0..n {
        let mut max_row = i;
        let mut max_val = aug[[i, i]].abs();

        // Partial pivoting
        for j in i + 1..n {
            if aug[[j, i]].abs() > max_val {
                max_row = j;
                max_val = aug[[j, i]].abs();
            }
        }

        if max_val < 1e-10 {
            return Err(SpatialError::ComputationError(
                "Matrix is singular or nearly singular".to_string(),
            ));
        }

        // Swap rows if needed
        if max_row != i {
            for j in 0..=n {
                let temp = aug[[i, j]];
                aug[[i, j]] = aug[[max_row, j]];
                aug[[max_row, j]] = temp;
            }
        }

        // Eliminate below
        for j in i + 1..n {
            let factor = aug[[j, i]] / aug[[i, i]];
            aug[[j, i]] = 0.0;

            for k in i + 1..=n {
                aug[[j, k]] -= factor * aug[[i, k]];
            }
        }
    }

    // Back substitution
    for i in (0..n).rev() {
        x[i] = aug[[i, n]];

        for j in i + 1..n {
            x[i] -= aug[[i, j]] * x[j];
        }

        x[i] /= aug[[i, i]];
    }

    Ok(x)
}

/// Available radial basis function kernels
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RBFKernel {
    /// Gaussian: φ(r) = exp(-ε²r²)
    /// Infinitely smooth but highly local
    Gaussian,

    /// Multiquadric: φ(r) = sqrt(1 + (εr)²)
    /// Moderately smooth and less local
    Multiquadric,

    /// Inverse Multiquadric: φ(r) = 1/sqrt(1 + (εr)²)
    /// Infinitely smooth with moderate locality
    InverseMultiquadric,

    /// Thin Plate Spline: φ(r) = r² ln(r)
    /// Minimizes curvature, very smooth
    ThinPlateSpline,

    /// Linear: φ(r) = r
    /// Simplest RBF, not smooth at data points
    Linear,

    /// Cubic: φ(r) = r³
    /// Good compromise between smoothness and locality
    Cubic,
}

impl RBFKernel {
    /// Apply the kernel function to a distance
    ///
    /// # Arguments
    ///
    /// * `r` - Distance
    /// * `epsilon` - Shape parameter (for kernels that use it)
    ///
    /// # Returns
    ///
    /// Value of the kernel function at distance r
    fn apply(&self, r: f64, epsilon: f64) -> f64 {
        match self {
            RBFKernel::Gaussian => (-epsilon * epsilon * r * r).exp(),
            RBFKernel::Multiquadric => (1.0 + (epsilon * r).powi(2)).sqrt(),
            RBFKernel::InverseMultiquadric => 1.0 / (1.0 + (epsilon * r).powi(2)).sqrt(),
            RBFKernel::ThinPlateSpline => {
                if r < 1e-10 {
                    0.0
                } else {
                    r * r * r.ln()
                }
            }
            RBFKernel::Linear => r,
            RBFKernel::Cubic => r.powi(3),
        }
    }
}

/// Radial Basis Function interpolator for scattered data
///
/// # Examples
///
/// ```
/// use scirs2_spatial::interpolate::{RBFInterpolator, RBFKernel};
/// use scirs2_core::ndarray::array;
///
/// // Create sample points and values
/// let points = array![
///     [0.0, 0.0],
///     [1.0, 0.0],
///     [0.0, 1.0],
///     [1.0, 1.0],
/// ];
/// let values = array![0.0, 1.0, 2.0, 3.0];
///
/// // Create interpolator with Gaussian kernel
/// let interp = RBFInterpolator::new(
///     &points.view(),
///     &values.view(),
///     RBFKernel::Gaussian,
///     Some(1.0),
///     None,
/// ).expect("Operation failed");
///
/// // Interpolate at a point
/// let query_point = array![0.5, 0.5];
/// let result = interp.interpolate(&query_point.view()).expect("Operation failed");
///
/// // For this simple example, should be close to 1.5
/// ```
#[derive(Debug, Clone)]
pub struct RBFInterpolator {
    /// Input points (N x D)
    points: Array2<f64>,
    /// Input values (N)
    _values: Array1<f64>,
    /// Dimensionality of the input points
    dim: usize,
    /// Number of input points
    n_points: usize,
    /// RBF kernel function
    kernel: RBFKernel,
    /// Shape parameter for the kernel
    epsilon: f64,
    /// Whether to include polynomial terms
    polynomial: bool,
    /// Weights for the RBF terms
    weights: Array1<f64>,
    /// Coefficients for the polynomial terms (if included)
    poly_coefs: Option<Array1<f64>>,
}

impl RBFInterpolator {
    /// Create a new RBF interpolator
    ///
    /// # Arguments
    ///
    /// * `points` - Input points with shape (n_samples, n_dims)
    /// * `values` - Input values with shape (n_samples,)
    /// * `kernel` - RBF kernel function to use
    /// * `epsilon` - Shape parameter for the kernel (default depends on kernel)
    /// * `polynomial` - Whether to include polynomial terms (default: false)
    ///
    /// # Returns
    ///
    /// A new RBFInterpolator
    ///
    /// # Errors
    ///
    /// * If points and values have different lengths
    /// * If fewer than d+1 points are provided (where d is the dimensionality)
    /// * If the system of equations is singular
    pub fn new(
        points: &ArrayView2<'_, f64>,
        values: &ArrayView1<f64>,
        kernel: RBFKernel,
        epsilon: Option<f64>,
        polynomial: Option<bool>,
    ) -> SpatialResult<Self> {
        // Check input dimensions
        let n_points = points.nrows();
        let dim = points.ncols();

        if n_points != values.len() {
            return Err(SpatialError::DimensionError(format!(
                "Number of points ({}) must match number of values ({})",
                n_points,
                values.len()
            )));
        }

        if n_points < dim + 1 {
            return Err(SpatialError::ValueError(format!(
                "At least {} points required for {}D interpolation",
                dim + 1,
                dim
            )));
        }

        // Set default epsilon based on kernel
        let epsilon = epsilon.unwrap_or_else(|| Self::default_epsilon(kernel, points));

        // Set default polynomial option
        let polynomial = polynomial.unwrap_or(false);

        // Build the interpolation system
        let (weights, poly_coefs) =
            Self::solve_rbf_system(points, values, kernel, epsilon, polynomial)?;

        Ok(Self {
            points: points.to_owned(),
            _values: values.to_owned(),
            dim,
            n_points,
            kernel,
            epsilon,
            polynomial,
            weights,
            poly_coefs,
        })
    }

    /// Get a default shape parameter based on the kernel and data
    ///
    /// # Arguments
    ///
    /// * `kernel` - The RBF kernel
    /// * `points` - The input points
    ///
    /// # Returns
    ///
    /// A reasonable default value for epsilon
    fn default_epsilon(kernel: RBFKernel, points: &ArrayView2<'_, f64>) -> f64 {
        match kernel {
            RBFKernel::Gaussian => {
                // For Gaussian, a typical choice is 1 / (2 * average distance^2)
                let avg_dist = Self::average_distance(points);
                if avg_dist > 0.0 {
                    1.0 / (2.0 * avg_dist * avg_dist)
                } else {
                    1.0
                }
            }
            RBFKernel::Multiquadric | RBFKernel::InverseMultiquadric => {
                // For multiquadrics, a typical choice is 1 / average distance
                let avg_dist = Self::average_distance(points);
                if avg_dist > 0.0 {
                    1.0 / avg_dist
                } else {
                    1.0
                }
            }
            // Other kernels don't use epsilon (or it's absorbed into the coefficients)
            _ => 1.0,
        }
    }

    /// Calculate average distance between points
    ///
    /// # Arguments
    ///
    /// * `points` - The input points
    ///
    /// # Returns
    ///
    /// The average distance between points
    fn average_distance(points: &ArrayView2<'_, f64>) -> f64 {
        let n_points = points.nrows();

        if n_points <= 1 {
            return 0.0;
        }

        // Sample a subset of pairs for efficiency if there are too many _points
        let max_pairs = 1000;
        let mut total_dist = 0.0;
        let mut n_pairs = 0;

        // Calculate average distance
        if n_points * (n_points - 1) / 2 <= max_pairs {
            // Use all pairs for small datasets
            for i in 0..n_points {
                for j in (i + 1)..n_points {
                    let pi = points.row(i);
                    let pj = points.row(j);
                    total_dist += Self::euclidean_distance(&pi, &pj);
                    n_pairs += 1;
                }
            }
        } else {
            // Sample pairs for large datasets
            let mut rng = thread_rng();
            let mut seen_pairs = std::collections::HashSet::new();

            for _ in 0..max_pairs {
                let i = rng.random_range(0..n_points);
                let j = rng.random_range(0..n_points);

                if i != j {
                    let pair = if i < j { (i, j) } else { (j, i) };
                    if !seen_pairs.contains(&pair) {
                        seen_pairs.insert(pair);
                        let pi = points.row(i);
                        let pj = points.row(j);
                        total_dist += Self::euclidean_distance(&pi, &pj);
                        n_pairs += 1;
                    }
                }
            }
        }

        if n_pairs > 0 {
            total_dist / (n_pairs as f64)
        } else {
            1.0
        }
    }

    /// Solve the RBF interpolation system
    ///
    /// # Arguments
    ///
    /// * `points` - Input points
    /// * `values` - Input values
    /// * `kernel` - RBF kernel function
    /// * `epsilon` - Shape parameter
    /// * `polynomial` - Whether to include polynomial terms
    ///
    /// # Returns
    ///
    /// A tuple (weights, poly_coefs) where poly_coefs is None if polynomial=false
    ///
    /// # Errors
    ///
    /// * If the system of equations is singular
    fn solve_rbf_system(
        points: &ArrayView2<'_, f64>,
        values: &ArrayView1<f64>,
        kernel: RBFKernel,
        epsilon: f64,
        polynomial: bool,
    ) -> SpatialResult<(Array1<f64>, Option<Array1<f64>>)> {
        let n_points = points.nrows();
        let dim = points.ncols();

        if !polynomial {
            // Without polynomial terms, we just need to solve A * w = y
            // where A_ij = kernel(||p_i - p_j||, epsilon)
            let mut a = Array2::zeros((n_points, n_points));

            for i in 0..n_points {
                let pi = points.row(i);
                for j in 0..n_points {
                    let pj = points.row(j);
                    let dist = Self::euclidean_distance(&pi, &pj);
                    a[[i, j]] = kernel.apply(dist, epsilon);
                }
            }

            // Manually solve using pseudo-inverse (not ideal but works for now)
            let trans_a = a.t();
            let ata = trans_a.dot(&a);
            let atb = trans_a.dot(&values.to_owned());
            let weights = solve_linear_system(ata, atb);
            match weights {
                Ok(weights) => Ok((weights, None)),
                Err(e) => Err(SpatialError::ComputationError(format!(
                    "Failed to solve RBF system: {e}"
                ))),
            }
        } else {
            // With polynomial terms, we need to set up an augmented system
            // [ A  P ] [ w ]   [ y ]
            // [ P' 0 ] [ c ] = [ 0 ]
            // where P contains the polynomial basis

            // For a linear polynomial, we need [1, x, y, z, ...]
            let poly_terms = dim + 1;

            // Set up the augmented matrix
            let mut aug_matrix = Array2::zeros((n_points + poly_terms, n_points + poly_terms));
            let mut aug_values = Array1::zeros(n_points + poly_terms);

            // Fill in the RBF part (top-left block)
            for i in 0..n_points {
                let pi = points.row(i);
                for j in 0..n_points {
                    let pj = points.row(j);
                    let dist = Self::euclidean_distance(&pi, &pj);
                    aug_matrix[[i, j]] = kernel.apply(dist, epsilon);
                }
            }

            // Fill in the polynomial part (top-right and bottom-left blocks)
            for i in 0..n_points {
                // Constant term
                aug_matrix[[i, n_points]] = 1.0;
                aug_matrix[[n_points, i]] = 1.0;

                // Linear terms
                for j in 0..dim {
                    aug_matrix[[i, n_points + 1 + j]] = points[[i, j]];
                    aug_matrix[[n_points + 1 + j, i]] = points[[i, j]];
                }
            }

            // Fill in the values
            for i in 0..n_points {
                aug_values[i] = values[i];
            }

            // Manually solve using pseudo-inverse (not ideal but works for now)
            let trans_a = aug_matrix.t();
            let ata = trans_a.dot(&aug_matrix);
            let atb = trans_a.dot(&aug_values);
            let solution = solve_linear_system(ata, atb);
            match solution {
                Ok(solution) => {
                    // Extract weights and polynomial coefficients
                    let weights = solution
                        .slice(scirs2_core::ndarray::s![0..n_points])
                        .to_owned();
                    let poly_coefs = solution
                        .slice(scirs2_core::ndarray::s![n_points..])
                        .to_owned();
                    Ok((weights, Some(poly_coefs)))
                }
                Err(e) => Err(SpatialError::ComputationError(format!(
                    "Failed to solve augmented RBF system: {e}"
                ))),
            }
        }
    }

    /// Interpolate at a single point
    ///
    /// # Arguments
    ///
    /// * `point` - Query point with shape (n_dims,)
    ///
    /// # Returns
    ///
    /// Interpolated value at the query point
    ///
    /// # Errors
    ///
    /// * If the point dimensions don't match the interpolator
    pub fn interpolate(&self, point: &ArrayView1<f64>) -> SpatialResult<f64> {
        // Check dimension
        if point.len() != self.dim {
            return Err(SpatialError::DimensionError(format!(
                "Query _point has dimension {}, expected {}",
                point.len(),
                self.dim
            )));
        }

        // Evaluate the RBF interpolant at the query _point
        let mut result = 0.0;

        // Sum over all RBF terms
        for i in 0..self.n_points {
            let pi = self.points.row(i);
            let dist = Self::euclidean_distance(&pi, point);
            result += self.weights[i] * self.kernel.apply(dist, self.epsilon);
        }

        // Add polynomial terms if present
        if let Some(ref poly_coefs) = self.poly_coefs {
            // Constant term
            result += poly_coefs[0];

            // Linear terms
            for j in 0..self.dim {
                result += poly_coefs[j + 1] * point[j];
            }
        }

        Ok(result)
    }

    /// Interpolate at multiple points
    ///
    /// # Arguments
    ///
    /// * `points` - Query points with shape (n_queries, n_dims)
    ///
    /// # Returns
    ///
    /// Interpolated values with shape (n_queries,)
    ///
    /// # Errors
    ///
    /// * If the points dimensions don't match the interpolator
    pub fn interpolate_many(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array1<f64>> {
        // Check dimensions
        if points.ncols() != self.dim {
            return Err(SpatialError::DimensionError(format!(
                "Query _points have dimension {}, expected {}",
                points.ncols(),
                self.dim
            )));
        }

        let n_queries = points.nrows();
        let mut results = Array1::zeros(n_queries);

        // Interpolate each point
        for i in 0..n_queries {
            let point = points.row(i);
            results[i] = self.interpolate(&point)?;
        }

        Ok(results)
    }

    /// Get the kernel used by this interpolator
    pub fn kernel(&self) -> RBFKernel {
        self.kernel
    }

    /// Get the shape parameter (epsilon) used by this interpolator
    pub fn epsilon(&self) -> f64 {
        self.epsilon
    }

    /// Check if this interpolator includes polynomial terms
    pub fn has_polynomial(&self) -> bool {
        self.polynomial
    }

    /// Compute the Euclidean distance between two points
    ///
    /// # Arguments
    ///
    /// * `p1` - First point
    /// * `p2` - Second point
    ///
    /// # Returns
    ///
    /// Euclidean distance between the points
    fn euclidean_distance(p1: &ArrayView1<f64>, p2: &ArrayView1<f64>) -> f64 {
        let mut sum_sq = 0.0;
        for i in 0..p1.len().min(p2.len()) {
            let diff = p1[i] - p2[i];
            sum_sq += diff * diff;
        }
        sum_sq.sqrt()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_rbf_interpolation_basic() {
        // Create a simple grid of points
        let points = array![
            [0.0, 0.0], // 0: bottom-left
            [1.0, 0.0], // 1: bottom-right
            [0.0, 1.0], // 2: top-left
            [1.0, 1.0], // 3: top-right
        ];

        // Set up a simple function z = x + y
        let values = array![0.0, 1.0, 1.0, 2.0];

        // Test different kernels
        let kernels = [
            RBFKernel::Gaussian,
            RBFKernel::Multiquadric,
            RBFKernel::InverseMultiquadric,
            RBFKernel::ThinPlateSpline,
            RBFKernel::Linear,
            RBFKernel::Cubic,
        ];

        for kernel in &kernels {
            // Create the interpolator
            let interp = RBFInterpolator::new(&points.view(), &values.view(), *kernel, None, None)
                .expect("Operation failed");

            // Test at the data points (should interpolate exactly)
            let val_00 = interp
                .interpolate(&array![0.0, 0.0].view())
                .expect("Operation failed");
            let val_10 = interp
                .interpolate(&array![1.0, 0.0].view())
                .expect("Operation failed");
            let val_01 = interp
                .interpolate(&array![0.0, 1.0].view())
                .expect("Operation failed");
            let val_11 = interp
                .interpolate(&array![1.0, 1.0].view())
                .expect("Operation failed");

            assert_relative_eq!(val_00, 0.0, epsilon = 1e-6);
            assert_relative_eq!(val_10, 1.0, epsilon = 1e-6);
            assert_relative_eq!(val_01, 1.0, epsilon = 1e-6);
            assert_relative_eq!(val_11, 2.0, epsilon = 1e-6);

            // Test at the center - we don't check exact value as it varies by kernel
            let val_center = interp
                .interpolate(&array![0.5, 0.5].view())
                .expect("Operation failed");

            // Instead of checking against 1.0, just make sure the value is finite
            assert!(val_center.is_finite());
        }
    }

    #[test]
    fn test_rbf_with_polynomial() {
        // Create data points on a line
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],];

        // Linear function z = 2*x + 3*y + 1
        let values = array![1.0, 3.0, 4.0, 6.0];

        // Create interpolator with polynomial
        let interp = RBFInterpolator::new(
            &points.view(),
            &values.view(),
            RBFKernel::Gaussian,
            Some(1.0),
            Some(true),
        )
        .expect("Operation failed");

        assert!(interp.has_polynomial());

        // Test at data points
        let val_00 = interp
            .interpolate(&array![0.0, 0.0].view())
            .expect("Operation failed");
        let val_10 = interp
            .interpolate(&array![1.0, 0.0].view())
            .expect("Operation failed");
        let val_01 = interp
            .interpolate(&array![0.0, 1.0].view())
            .expect("Operation failed");
        let val_11 = interp
            .interpolate(&array![1.0, 1.0].view())
            .expect("Operation failed");

        assert_relative_eq!(val_00, 1.0, epsilon = 1e-6);
        assert_relative_eq!(val_10, 3.0, epsilon = 1e-6);
        assert_relative_eq!(val_01, 4.0, epsilon = 1e-6);
        assert_relative_eq!(val_11, 6.0, epsilon = 1e-6);

        // Test at a new point - should follow linear pattern
        let val_new = interp
            .interpolate(&array![2.0, 2.0].view())
            .expect("Operation failed");
        // 2*x + 3*y + 1 = 2*2 + 3*2 + 1 = 11
        assert_relative_eq!(val_new, 11.0, epsilon = 0.1);
    }

    #[test]
    fn test_interpolate_many() {
        // Create a simple grid of points
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],];

        // Set up a simple function z = x + y
        let values = array![0.0, 1.0, 1.0, 2.0];

        // Create the interpolator
        let interp = RBFInterpolator::new(
            &points.view(),
            &values.view(),
            RBFKernel::Gaussian,
            None,
            None,
        )
        .expect("Operation failed");

        // Test multiple points at once
        let query_points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5],];

        let results = interp
            .interpolate_many(&query_points.view())
            .expect("Operation failed");

        assert_eq!(results.len(), 5);
        assert_relative_eq!(results[0], 0.0, epsilon = 1e-6);
        assert_relative_eq!(results[1], 1.0, epsilon = 1e-6);
        assert_relative_eq!(results[2], 1.0, epsilon = 1e-6);
        assert_relative_eq!(results[3], 2.0, epsilon = 1e-6);
        assert_relative_eq!(results[4], 1.0, epsilon = 0.1);
    }

    #[test]
    fn test_error_handling() {
        // Not enough points
        let points = array![[0.0, 0.0]];
        let values = array![0.0];

        let result = RBFInterpolator::new(
            &points.view(),
            &values.view(),
            RBFKernel::Gaussian,
            None,
            None,
        );
        assert!(result.is_err());

        // Mismatched lengths
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]];
        let values = array![0.0, 1.0];

        let result = RBFInterpolator::new(
            &points.view(),
            &values.view(),
            RBFKernel::Gaussian,
            None,
            None,
        );
        assert!(result.is_err());

        // Valid interpolator but wrong dimension for query
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]];
        let values = array![0.0, 1.0, 2.0];

        let interp = RBFInterpolator::new(
            &points.view(),
            &values.view(),
            RBFKernel::Gaussian,
            None,
            None,
        )
        .expect("Operation failed");

        let result = interp.interpolate(&array![0.0, 0.0, 0.0].view());
        assert!(result.is_err());
    }
}