polyfit 0.11.0

use nalgebra::{ComplexField, 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::{FloatClampedCast, IntClampedCast, Value},
};

/// Normalized Legendre basis for polynomial curves.
///
/// This basis uses the Legendre polynomials, which form an
/// orthogonal family of polynomials on the interval [-1, 1].
/// Orthogonality ensures numerical stability and clean separation
/// of polynomial terms, especially useful for least-squares fitting.
///
/// Inputs are normalized so that the evaluation domain
/// [`x_min`, `x_max`] is mapped onto [-1, 1]. This allows Legendre
/// polynomials to be used naturally with arbitrary input ranges
/// while retaining their orthogonality properties.
///
/// # When to use
/// - Use when fitting polynomials and you want orthogonal basis functions.
/// - Prefer for higher-degree polynomials to reduce numerical error
///   in coefficient estimation.
/// - Useful for physics or engineering applications where
///   Legendre expansions are natural (e.g., potential fields, spherical harmonics).
///
/// # Why Legendre?
/// - Orthogonal over [-1, 1] with uniform weight, simplifying least-squares fitting.
/// - Minimizes coefficient correlation, improving numerical stability.
/// - Provides exact integral properties for weighted projections.
#[derive(Debug, Clone, Copy)]
pub struct LegendreBasis<T: Value = f64> {
    /// Normalizer to map input domain to [-1, 1]
    pub normalizer: DomainNormalizer<T>,
}
impl<T: Value> LegendreBasis<T> {
    /// Creates a new Legendre basis that normalizes inputs from the given range to [-1, 1].
    pub fn new(x_min: T, x_max: T) -> Self {
        let normalizer = DomainNormalizer::new((x_min, x_max), (-T::one(), T::one()));
        Self { normalizer }
    }

    /// Creates a new Legendre polynomial with the given coefficients over the specified x-range.
    ///
    /// # Parameters
    /// - `x_range`: The range of x-values over which the Legendre basis is defined.
    /// - `coefficients`: The coefficients for the Legendre basis functions.
    ///
    /// # Returns
    /// A polynomial defined in the Legendre basis.
    ///
    /// # Errors
    /// Returns an error if the polynomial cannot be created with the given basis and coefficients.
    ///
    /// # Example
    /// ```rust
    /// use polyfit::basis::LegendreBasis;
    /// let legendre_poly = LegendreBasis::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)
    }

    /// Evaluate `P_n(x)` and `P_n'(x)` at once
    /// Warning: does not normalize x
    pub fn solve_and_derive(&self, n: usize, x: T) -> (T, T) {
        let mut p0 = T::one();
        let mut p1 = x;
        for k in 2..=n {
            let k = T::from_positive_int(k);
            let p2 = ((T::two() * k - T::one()) * x * p1 - (k - T::one()) * p0) / k;
            p0 = p1;
            p1 = p2;
        }
        let p = if n == 0 { p0 } else { p1 };
        // Derivative using recurrence
        let dp = T::from_positive_int(n) * (x * p - p0) / (x * x - T::one());
        (p, dp)
    }
}

impl<T: Value> Basis<T> for LegendreBasis<T> {
    fn from_range(x_range: std::ops::RangeInclusive<T>) -> Self {
        let normalizer = DomainNormalizer::from_range(x_range, (-T::one(), T::one()));
        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 => x,
            _ => {
                let mut p0 = T::one();
                let mut p1 = x;

                for n in 1..j {
                    let p_next = ((T::from_positive_int(2 * n + 1) * x * p1)
                        - (T::from_positive_int(n) * p0))
                        / T::from_positive_int(n + 1);
                    p0 = p1;
                    p1 = p_next;
                }

                p1
            }
        }
    }

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

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

        let x = display::unicode::subscript("s");
        let x = format!("x{x}");

        let rank = display::unicode::subscript(&degree.to_string());
        let func = if degree > 0 {
            format!("P{rank}({x})")
        } else {
            String::new()
        };
        let coef = format_coefficient(coef, degree, DEFAULT_PRECISION)?;

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

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

    fn format_scaling_formula(&self) -> Option<String> {
        let x = display::unicode::subscript("s");
        let x = format!("x{x}");

        Some(format!("{x} = {}", self.normalizer))
    }
}

impl<T: Value> IntoMonomialBasis<T> for LegendreBasis<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 / 2) {
                let sign = if k % 2 == 0 { T::one() } else { -T::one() };
                let numerator = T::factorial(2 * j - 2 * k);
                let denom = Value::powi(T::two(), j.clamped_cast())
                    * T::factorial(k)
                    * T::factorial(j - k)
                    * T::factorial(j - 2 * k);
                let x_power = j - 2 * k;
                result[x_power] += c_j * sign * numerator / denom;
            }
        }

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

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

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

    fn derivative(&self, a: &[T]) -> Result<(Self::B2, Vec<T>)> {
        if a.is_empty() {
            return Ok((*self, vec![T::zero()]));
        }

        let mut b = vec![T::zero(); a.len() - 1];
        for n in 1..a.len() {
            for m in 0..n {
                let tm1 = T::two() * T::from_positive_int(m) + T::one();
                if (n - m) % 2 == 1 {
                    b[m] += tm1 * a[n];
                }
            }
        }

        // scale coefficients to account for original domain
        let scale = self.normalizer.scale();
        for coeff in &mut b {
            *coeff *= scale;
        }

        Ok((*self, b))
    }
}

impl<T: Value> OrthogonalBasis<T> for LegendreBasis<T> {
    fn gauss_nodes(&self, n: usize) -> Vec<(T, T)> {
        let tol = T::epsilon().sqrt();
        let mut nodes = Vec::with_capacity(n);

        let m = n.div_ceil(2); // roots in [0,1]
        for i in 0..m {
            // Initial guess via cosine
            let theta = T::pi() * (T::from_positive_int(i) + 0.75f64.clamped_cast::<T>())
                / (T::from_positive_int(n) + T::half());
            let mut x = theta.cos();

            // Newton–Raphson
            loop {
                let (p, dp) = self.solve_and_derive(n, x);
                let dx = -p / dp;
                x += dx;
                if Value::abs(dx) < tol {
                    break;
                }
            }

            let (_, dp) = self.solve_and_derive(n, x);
            let w = T::two() / ((T::one() - x * x) * dp * dp);

            // Mirror pair
            nodes.push((-x, w));
            if i != m - 1 || n % 2 == 0 {
                nodes.push((x, w));
            }
        }

        nodes.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        nodes
    }

    fn gauss_weight(&self, _: T) -> T {
        T::one()
    }

    fn gauss_normalization(&self, n: usize) -> T {
        T::two() / (T::two() * T::from_positive_int(n) + T::one())
    }
}

impl<T: Value> RootFindingBasis<T> for LegendreBasis<T> {
    fn complex_y(&self, z: nalgebra::Complex<T>, coefs: &[T]) -> nalgebra::Complex<T> {
        use nalgebra::Complex;

        if coefs.is_empty() {
            return Complex::zero();
        }

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

        // P1
        let mut p_n = z;
        let mut result = Complex::from_real(coefs[0]) * p_nm1 + Complex::from_real(coefs[1]) * p_n;

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

            let a = Complex::from_real(T::from_positive_int(2 * n + 1)) * z * p_n;
            let b = Complex::from_real(n_t) * p_nm1;
            let p_np1 = (a - b) / Complex::from_real(n1_t);

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

            p_nm1 = p_n;
            p_n = p_np1;
        }

        result
    }
}

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

    use super::*;

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

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

        // Orthogonality test points
        assert_basis_orthogonal(fit.basis(), 3, 100, 1e-12);

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

        // Solve known functions
        let basis = LegendreBasis::new(-1.0, 1.0);
        assert_close!(basis.solve_function(0, 0.5), 1.0);
        assert_close!(basis.solve_function(1, 0.5), 0.5);
        assert_close!(basis.solve_function(2, 0.5), -0.125);
        assert_close!(basis.solve_function(3, 0.5), -0.4375);

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

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