mathr 0.1.0

//! Equation solvers: Newton–Raphson, bisection, secant, and polynomial
//! root finding via the Durand–Kerner method.

use crate::error::{MathError, Result};

/// Configuration shared by all iterative solvers.
#[derive(Debug, Clone)]
pub struct SolveOptions {
    pub max_iter: usize,
    pub tol: f64,
    pub h: f64, // finite-difference step
}

impl Default for SolveOptions {
    fn default() -> Self {
        Self {
            max_iter: 100,
            tol: 1e-10,
            h: 1e-6,
        }
    }
}

/// Bisection method: finds a root of `f` on `[a, b]` assuming `f(a)*f(b) < 0`.
pub fn bisect<F: Fn(f64) -> f64>(
    f: F,
    a: f64,
    b: f64,
    opts: SolveOptions,
) -> Result<(f64, f64)> {
    let mut lo = a;
    let mut hi = b;
    let fa = f(lo);
    let fb = f(hi);
    if fa == 0.0 {
        return Ok((lo, 0.0));
    }
    if fb == 0.0 {
        return Ok((hi, 0.0));
    }
    if fa.signum() == fb.signum() {
        return Err(MathError::InvalidArgument(format!(
            "bisect: f(a) and f(b) must have opposite signs (got f({})={}, f({})={})",
            lo, fa, hi, fb
        )));
    }
    let mut mid = 0.5 * (lo + hi);
    for _ in 0..opts.max_iter {
        mid = 0.5 * (lo + hi);
        let fm = f(mid);
        if fm.abs() < opts.tol || (hi - lo) * 0.5 < opts.tol {
            return Ok((mid, fm));
        }
        if fm.signum() == fa.signum() {
            lo = mid;
        } else {
            hi = mid;
        }
    }
    Err(MathError::NotConvergent(format!(
        "bisect did not converge in {} iterations",
        opts.max_iter
    )))
}

/// Newton–Raphson method. Requires a derivative `df`; if you don't have one,
/// use [`newton_central`] which approximates it with a centred difference.
pub fn newton<F: Fn(f64) -> f64, D: Fn(f64) -> f64>(
    f: F,
    df: D,
    x0: f64,
    opts: SolveOptions,
) -> Result<(f64, f64)> {
    let mut x = x0;
    let mut fx = f(x);
    for i in 0..opts.max_iter {
        if fx.abs() < opts.tol {
            return Ok((x, fx));
        }
        let dfx = df(x);
        if dfx == 0.0 {
            return Err(MathError::NotConvergent(format!(
                "newton: derivative vanished at iteration {}",
                i
            )));
        }
        let step = fx / dfx;
        x -= step;
        fx = f(x);
        if step.abs() < opts.tol {
            return Ok((x, fx));
        }
    }
    Err(MathError::NotConvergent(format!(
        "newton did not converge in {} iterations",
        opts.max_iter
    )))
}

/// Newton method with a numerically-estimated derivative (central difference).
pub fn newton_central<F>(f: F, x0: f64, opts: SolveOptions) -> Result<(f64, f64)>
where
    F: Fn(f64) -> f64 + Clone,
{
    let h = opts.h;
    let f2 = f.clone();
    let df = move |x: f64| (f2(x + h) - f2(x - h)) / (2.0 * h);
    newton(f, df, x0, opts)
}

/// Secant method: derivative-free, but starts with two initial guesses.
pub fn secant<F: Fn(f64) -> f64>(
    f: F,
    x0: f64,
    x1: f64,
    opts: SolveOptions,
) -> Result<(f64, f64)> {
    let mut x_prev = x0;
    let mut x_curr = x1;
    let mut f_prev = f(x_prev);
    let mut f_curr = f(x_curr);
    for _ in 0..opts.max_iter {
        if f_curr.abs() < opts.tol {
            return Ok((x_curr, f_curr));
        }
        let denom = f_curr - f_prev;
        if denom.abs() < 1e-30 {
            return Err(MathError::NotConvergent(
                "secant: zero denominator".into(),
            ));
        }
        let x_next = x_curr - f_curr * (x_curr - x_prev) / denom;
        if (x_next - x_curr).abs() < opts.tol {
            let fn_ = f(x_next);
            return Ok((x_next, fn_));
        }
        x_prev = x_curr;
        f_prev = f_curr;
        x_curr = x_next;
        f_curr = f(x_next);
    }
    Err(MathError::NotConvergent(format!(
        "secant did not converge in {} iterations",
        opts.max_iter
    )))
}

/// Find all (real) roots of a polynomial with real coefficients given by
/// `coeffs` in descending order: `coeffs[0]*x^n + ... + coeffs[n]`.
///
/// Implemented with the Durand–Kerner method, which converges quickly for
/// well-separated roots and is robust enough for moderate degrees.
pub fn polynomial_roots(coeffs: &[f64]) -> Result<Vec<(f64, f64)>> {
    let n = coeffs.len().saturating_sub(1);
    if n == 0 {
        return Ok(Vec::new());
    }
    // normalize so leading coefficient is 1
    let lead = coeffs[0];
    if lead == 0.0 {
        return Err(MathError::InvalidArgument("polynomial leading coefficient is 0".into()));
    }
    let a: Vec<f64> = coeffs.iter().map(|c| c / lead).collect();

    // initial guesses: complex roots spread around the unit circle
    use crate::complex::Complex;
    let mut z: Vec<Complex<f64>> = (0..n)
        .map(|k| {
            let theta = std::f64::consts::PI * (2.0 * k as f64 + 1.0) / n as f64;
            Complex::from_polar(0.4 + 0.9 * (n as f64 / 10.0), theta)
        })
        .collect();

    for _ in 0..200 {
        let mut converged = true;
        for i in 0..n {
            let mut p = Complex::new(a[0], 0.0);
            for j in 1..a.len() {
                p = p * z[i] + Complex::new(a[j], 0.0);
            }
            let mut denom = Complex::new(1.0, 0.0);
            for j in 0..n {
                if i != j {
                    denom = denom * (z[i] - z[j]);
                }
            }
            let step = if denom.abs() < 1e-30 {
                Complex::new(0.0, 0.0)
            } else {
                p / denom
            };
            z[i] = z[i] - step;
            if step.abs() > 1e-10 {
                converged = false;
            }
        }
        if converged {
            break;
        }
    }

    // collect real roots (within tolerance) and pair them with f(root)
    let mut out = Vec::new();
    for zi in z {
        if zi.im.abs() < 1e-6 {
            let x = zi.re;
            let fx = polynomial_eval(coeffs, x);
            out.push((x, fx));
        }
    }
    out.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
    Ok(out)
}

fn polynomial_eval(coeffs: &[f64], x: f64) -> f64 {
    let mut acc = 0.0;
    for c in coeffs {
        acc = acc * x + c;
    }
    acc
}

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

    #[test]
    fn bisect_root() {
        // x^2 - 4 = 0 on [0, 5]
        let f = |x: f64| x * x - 4.0;
        let (r, _) = bisect(f, 0.0, 5.0, SolveOptions::default()).unwrap();
        assert!((r - 2.0).abs() < 1e-6);
    }

    #[test]
    fn newton_root() {
        // cos(x) = x  -> Newton's with f=cos(x)-x, df=-sin(x)-1
        let f = |x: f64| x.cos() - x;
        let df = |x: f64| -x.sin() - 1.0;
        let (r, _) = newton(f, df, 0.5, SolveOptions::default()).unwrap();
        // known fixed point ≈ 0.7390851332151607
        assert!((r - 0.7390851332151607).abs() < 1e-8);
    }

    #[test]
    fn newton_central_root() {
        // Newton from 3 on sin(x) converges to the closest root, which is π.
        let f = |x: f64| x.sin();
        let (r, _) = newton_central(f, 3.0, SolveOptions::default()).unwrap();
        assert!((r - PI).abs() < 1e-7, "expected root at π, got {}", r);
    }

    #[test]
    fn secant_root() {
        let f = |x: f64| x.powi(3) - 2.0;
        let (r, _) = secant(f, 1.0, 2.0, SolveOptions::default()).unwrap();
        assert!((r - 2f64.cbrt()).abs() < 1e-8);
    }

    #[test]
    fn polynomial_two_roots() {
        // x^2 - 5x + 6 = (x-2)(x-3)
        let r = polynomial_roots(&[1.0, -5.0, 6.0]).unwrap();
        let xs: Vec<f64> = r.iter().map(|(x, _)| *x).collect();
        assert_eq!(xs.len(), 2);
        let mut sorted = xs.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        assert!((sorted[0] - 2.0).abs() < 1e-4);
        assert!((sorted[1] - 3.0).abs() < 1e-4);
    }
}