polyfit 0.11.0

use nalgebra::{Complex, ComplexField, Dim, MatrixViewMut};
use num_traits::{One, Zero};

use crate::{
    basis::{Basis, DifferentialBasis, IntoMonomialBasis, OrthogonalBasis, RootFindingBasis},
    display::{self, format_coefficient, PolynomialDisplay, Sign, Term, DEFAULT_PRECISION},
    error::Result,
    statistics::DomainNormalizer,
    value::Value,
};

/// Normalized Laguerre basis for polynomial curves.
///
/// This basis uses the (generalized) Laguerre polynomials, which form an
/// orthogonal family of polynomials on the interval [0, ∞) with weight `exp(-x)`.
/// Orthogonality improves numerical stability for higher-degree polynomials.
///
/// # When to use
/// - Useful for problems with positive-valued domains or exponentially decaying data.
/// - Preferred when standard monomial fits become unstable.
///
/// # Why Laguerre?
/// - Minimizes numerical issues for high-degree polynomial approximations on [0, ∞).
/// - Naturally models decaying or exponential-type behavior.
/// - Orthogonal under weight `exp(-x)`, making coefficient estimation more stable.
#[derive(Debug, Clone, Copy)]
pub struct LaguerreBasis<T: Value = f64> {
    normalizer: DomainNormalizer<T>,
}
impl<T: Value> LaguerreBasis<T> {
    /// Creates a new Laguerre basis that normalizes inputs from the given range to [0, ∞).
    pub fn new(x_min: T, x_max: T) -> Self {
        let normalizer = DomainNormalizer::new((x_min, x_max), (T::zero(), T::infinity()));
        Self { normalizer }
    }

    /// Creates a new Laguerre polynomial with the given coefficients over the specified x-range.
    ///
    /// # Parameters
    /// - `x_range`: The range of x-values over which the Laguerre basis is defined.
    /// - `coefficients`: The coefficients for the Laguerre basis functions.
    ///
    /// # Returns
    /// A polynomial defined in the Laguerre basis.
    ///
    /// # Errors
    /// Returns an error if the polynomial cannot be created with the given basis and coefficients.
    ///
    /// # Example
    /// ```rust
    /// use polyfit::basis::LaguerreBasis;
    /// let laguerre_poly = LaguerreBasis::new_polynomial((-1.0, 1.0), &[1.0, 0.0, -0.5]).unwrap();
    /// ```
    pub fn new_polynomial(
        x_range: (T, T),
        coefficients: &[T],
    ) -> Result<crate::Polynomial<'_, Self, T>> {
        let basis = Self::new(x_range.0, x_range.1);
        crate::Polynomial::<Self, T>::from_basis(basis, coefficients)
    }
}
impl<T: Value> Basis<T> for LaguerreBasis<T> {
    fn from_range(x_range: std::ops::RangeInclusive<T>) -> Self {
        let normalizer = DomainNormalizer::from_range(x_range, (T::zero(), T::infinity()));
        Self { normalizer }
    }

    #[inline(always)]
    fn normalize_x(&self, x: T) -> T {
        self.normalizer.normalize(x)
    }

    #[inline(always)]
    fn denormalize_x(&self, x: T) -> T {
        self.normalizer.denormalize(x)
    }

    #[inline(always)]
    fn solve_function(&self, j: usize, x: T) -> T {
        match j {
            0 => T::one(),
            1 => T::one() - x,
            _ => {
                let mut l0 = T::one();
                let mut l1 = T::one() - x;
                for n in 1..j {
                    let n_t = T::from_positive_int(n);
                    let l2 = ((T::two() * n_t + T::one() - x) * l1 - n_t * l0) / (n_t + T::one());
                    l0 = l1;
                    l1 = l2;
                }
                l1
            }
        }
    }

    #[inline(always)]
    fn fill_matrix_row<R: Dim, C: Dim, RS: Dim, CS: Dim>(
        &self,
        start_index: usize,
        x: T,
        mut row: MatrixViewMut<T, R, C, RS, CS>,
    ) {
        for j in start_index..row.ncols() {
            row[j] = self.solve_function(j, x);
        }
    }
}

impl<T: Value> PolynomialDisplay<T> for LaguerreBasis<T> {
    fn format_term(&self, degree: i32, coef: T) -> Option<Term> {
        let sign = Sign::from_coef(coef);
        let coef = format_coefficient(coef, degree, DEFAULT_PRECISION)?;

        if degree == 0 {
            return Some(Term { sign, body: coef });
        }

        // Lₙ(x) for Laguerre polynomial
        let rank = display::unicode::subscript(&degree.to_string());
        let func = format!("L{rank}(x)");

        let glue = if coef.is_empty() || func.is_empty() {
            ""
        } else {
            "·"
        };

        let body = format!("{coef}{glue}{func}");
        Some(Term { sign, body })
    }
}

impl<T: Value> DifferentialBasis<T> for LaguerreBasis<T> {
    type B2 = LaguerreBasis<T>;

    fn derivative(&self, a: &[T]) -> Result<(Self::B2, Vec<T>)> {
        let n = a.len();
        let mut b = Vec::with_capacity(n);

        for k in 0..n {
            // sum A_{k+1} to A_{N-1} and negate
            let mut sum = T::zero();
            for i in k + 1..n {
                sum += a[i];
            }
            b.push(-sum);
        }

        Ok((*self, b))
    }
}

impl<T: Value> IntoMonomialBasis<T> for LaguerreBasis<T> {
    fn as_monomial(&self, coefficients: &mut [T]) -> Result<()> {
        let n = coefficients.len();
        let mut result = vec![T::zero(); n];

        for j in 0..n {
            let c_j = coefficients[j];
            for k in 0..=j {
                let sign = if k % 2 == 0 { T::one() } else { -T::one() };
                let binom = T::factorial(j) / (T::factorial(k) * T::factorial(j - k));
                let factor = binom / T::factorial(k); // divide by k!
                result[k] += c_j * sign * factor;
            }
        }

        //
        // Phase 2 - Un-normalize over x
        result = self.normalizer.denormalize_coefs(&result);

        coefficients.copy_from_slice(&result);
        Ok(())
    }
}

impl<T: Value> OrthogonalBasis<T> for LaguerreBasis<T> {
    fn gauss_nodes(&self, n: usize) -> Vec<(T, T)> {
        // Laguerre polynomials are orthogonal on [0, ∞) with weight e^{-x}
        // The Golub–Welsch algorithm constructs a symmetric tridiagonal matrix
        // whose eigenvalues are the roots, and whose eigenvectors give weights.
        if n == 0 {
            return vec![(T::zero(), T::one())];
        }

        // α = 0 for standard Laguerre
        let alpha = T::zero();

        // Build Jacobi matrix
        let mut a = nalgebra::DMatrix::<T>::zeros(n, n);
        for i in 0..n {
            let ti = T::from_positive_int(i);
            a[(i, i)] = T::two() * ti + T::one() + alpha;
            if i > 0 {
                a[(i, i - 1)] = -ti;
                a[(i - 1, i)] = -ti;
            }
        }

        // Diagonalize
        let eig = a.symmetric_eigen();
        let mut nodes = Vec::with_capacity(n);

        // weights: w_i = x_i / [(n+1)^2 * (L_{n+1}(x_i))^2]
        // Golub–Welsch gives normalized weights directly from the first eigenvector component:
        // w_i = v_0^2
        for (xi, vi0) in eig.eigenvalues.iter().zip(eig.eigenvectors.row(0).iter()) {
            let x = *xi;
            let w = *vi0 * *vi0;
            nodes.push((x, w));
        }

        // Sort by node
        nodes.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
        nodes
    }

    fn gauss_weight(&self, x: T) -> T {
        // w(x) = e^{-x}
        (-x).exp()
    }

    fn gauss_normalization(&self, _n: usize) -> T {
        // ∫₀^∞ e^{-x} dx = 1
        T::one()
    }
}

impl<T: Value> RootFindingBasis<T> for LaguerreBasis<T> {
    fn complex_y(&self, z: nalgebra::Complex<T>, coefs: &[T]) -> nalgebra::Complex<T> {
        if coefs.is_empty() {
            return Complex::zero();
        }

        let mut l_nm1 = Complex::one(); // L0
        if coefs.len() == 1 {
            return l_nm1 * Complex::from_real(coefs[0]);
        }

        let mut l_n = l_nm1 - z; // L1
        let mut result = Complex::from_real(coefs[0]) * l_nm1 + Complex::from_real(coefs[1]) * l_n;

        for n in 1..coefs.len() - 1 {
            let n_t = Complex::from_real(T::from_positive_int(n));
            let n1_t = Complex::from_real(T::from_positive_int(n + 1));

            let t_2n1 = Complex::from_real(T::two() * T::from_positive_int(n) + T::one());
            let a = (t_2n1 - z) * l_n;
            let b = n_t * l_nm1;
            let l_np1 = (a - b) / n1_t;

            result += Complex::from_real(coefs[n + 1]) * l_np1;

            l_nm1 = l_n;
            l_n = l_np1;
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        assert_close, assert_fits,
        score::Aic,
        statistics::DegreeBound,
        test::basis_assertions::{self, assert_basis_orthogonal},
        LaguerreFit, Polynomial,
    };

    use super::*;

    fn get_poly() -> Polynomial<'static, LaguerreBasis<f64>> {
        let basis = LaguerreBasis::new(0.0, 100.0);
        Polynomial::from_basis(basis, &[1.0, 2.0, -0.5]).unwrap()
    }

    #[test]
    fn test_laguerre() {
        // Recover polynomial
        let poly = get_poly();
        let data = poly.solve_range(0.0..=100.0, 1.0);
        let fit = LaguerreFit::new_auto(&data, DegreeBound::Relaxed, &Aic).unwrap();
        assert_fits!(&poly, &fit);

        // Monomial conversion
        let mono = fit.as_monomial().unwrap();
        assert_fits!(mono, fit);

        // Solve known functions
        let basis = LaguerreBasis::new(-1.0, 1.0);
        assert_close!(basis.solve_function(0, 1.0), 1.0);
        assert_close!(basis.solve_function(1, 1.0), 0.0);
        assert_close!(basis.solve_function(2, 1.0), -0.5);

        // Calculus tests
        let poly = LaguerreBasis::new_polynomial((0.0, 100.0), &[3.0, 2.0, 1.5]).unwrap();
        basis_assertions::test_derivation(&poly, &poly.basis().normalizer);

        // Orthogonality test points
        assert_basis_orthogonal(&basis, 4, 100, 1e-12);

        basis_assertions::test_complex_y(&poly, 0.0..=100.0);
    }
}