oxieml 0.1.2

//! Sturm sequence and real root isolation for univariate polynomials.

use super::PolyError;
use super::univariate::Poly;

impl Poly {
    /// Build the Sturm sequence for this polynomial.
    pub fn sturm_sequence(&self) -> Result<Vec<Self>, PolyError> {
        let p = self.normalized();
        if p.is_zero() {
            return Ok(Vec::new());
        }
        let dp = p.diff()?;
        if dp.is_zero() {
            return Ok(vec![p]);
        }

        let mut seq = vec![p, dp];
        loop {
            let n = seq.len();
            let (_, rem) = seq[n - 2].div_rem(&seq[n - 1])?;
            if rem.is_zero() {
                break;
            }
            let neg_rem = rem.neg()?;
            seq.push(neg_rem);
        }
        Ok(seq)
    }

    /// Count sign changes in the Sturm sequence evaluated at `x`.
    pub fn sign_changes(sturm: &[Self], x: f64) -> usize {
        let vals: Vec<f64> = sturm.iter().map(|p| p.eval_f64(x)).collect();
        let nonzero: Vec<f64> = vals.into_iter().filter(|v| *v != 0.0).collect();
        let mut changes = 0usize;
        for w in nonzero.windows(2) {
            if w[0].signum() != w[1].signum() {
                changes += 1;
            }
        }
        changes
    }

    /// Count the number of distinct real roots in `(a, b)` using Sturm's theorem.
    pub fn count_real_roots(&self, a: f64, b: f64) -> Result<usize, PolyError> {
        let sf = self.square_free()?;
        let sturm = sf.sturm_sequence()?;
        if sturm.is_empty() {
            return Ok(0);
        }
        let va = Self::sign_changes(&sturm, a);
        let vb = Self::sign_changes(&sturm, b);
        Ok(va.saturating_sub(vb))
    }

    /// Isolate all real roots of this polynomial in `[lo, hi]`.
    pub fn isolate_real_roots(&self, lo: f64, hi: f64, tol: f64) -> Result<Vec<f64>, PolyError> {
        let sf = self.square_free()?;
        let sturm = sf.sturm_sequence()?;
        if sturm.is_empty() {
            return Ok(Vec::new());
        }

        let mut brackets: Vec<(f64, f64)> = Vec::new();
        bisect_isolate(&sturm, lo, hi, tol, &mut brackets)?;

        let mut roots: Vec<f64> = Vec::new();
        for (a, b) in brackets {
            if (b - a).abs() < tol {
                // Narrow bracket: the root is at or near a (and b ≈ a).
                let fa = sf.eval_f64(a);
                let fb = sf.eval_f64(b);
                if fa.abs() <= fb.abs() {
                    roots.push(a);
                } else {
                    roots.push(b);
                }
                continue;
            }

            let fa = sf.eval_f64(a);
            let fb = sf.eval_f64(b);

            if fa == 0.0 {
                roots.push(a);
                continue;
            }
            if fb == 0.0 {
                roots.push(b);
                continue;
            }

            // Normal sign-change bracket: refine with bisection.
            if fa.is_finite() && fb.is_finite() && fa * fb < 0.0 {
                let root = bisect_refine(&sf, a, b, fa, tol);
                roots.push(root);
            } else if fa.is_finite() && fa.abs() < tol {
                // Endpoint near root; it was generated by exact-root detection.
                roots.push(a);
            } else if fb.is_finite() && fb.abs() < tol {
                roots.push(b);
            }
        }

        roots.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        roots.dedup_by(|a, b| (*a - *b).abs() < tol);
        Ok(roots)
    }
}

fn bisect_isolate(
    sturm: &[Poly],
    a: f64,
    b: f64,
    tol: f64,
    brackets: &mut Vec<(f64, f64)>,
) -> Result<(), PolyError> {
    // Sturm's theorem is stated for points that are not roots of p0. When an
    // endpoint is an exact root, record it directly and skip it in the count.
    let p0_a = sturm.first().map(|p| p.eval_f64(a)).unwrap_or(1.0);
    let p0_b = sturm.first().map(|p| p.eval_f64(b)).unwrap_or(1.0);

    if p0_a == 0.0 {
        brackets.push((a, a));
    }
    if p0_b == 0.0 {
        brackets.push((b, b));
    }

    // Work with a (half-open) interior interval that excludes exact-root endpoints.
    let span = b - a;
    let a_eff = if p0_a == 0.0 { a + span * 1e-12 } else { a };
    let b_eff = if p0_b == 0.0 { b - span * 1e-12 } else { b };

    if a_eff >= b_eff {
        return Ok(());
    }

    bisect_isolate_inner(sturm, a_eff, b_eff, tol, brackets)
}

fn bisect_isolate_inner(
    sturm: &[Poly],
    a: f64,
    b: f64,
    tol: f64,
    brackets: &mut Vec<(f64, f64)>,
) -> Result<(), PolyError> {
    let count = {
        let va = Poly::sign_changes(sturm, a);
        let vb = Poly::sign_changes(sturm, b);
        va.saturating_sub(vb)
    };

    if count == 0 {
        return Ok(());
    }

    if count == 1 || (b - a) < tol {
        brackets.push((a, b));
        return Ok(());
    }

    let mid = (a + b) / 2.0;

    // If mid is an exact root of p0, push it as its own bracket and search
    // (a, mid-eps) and (mid+eps, b) to avoid unreliable Sturm counts at mid.
    if let Some(p0) = sturm.first() {
        let fm = p0.eval_f64(mid);
        if fm == 0.0 && (b - a) > 2.0 * tol {
            brackets.push((mid, mid));
            let eps = (b - a) * 1e-9;
            let lo = mid - eps;
            let hi = mid + eps;
            if lo > a {
                bisect_isolate_inner(sturm, a, lo, tol, brackets)?;
            }
            if hi < b {
                bisect_isolate_inner(sturm, hi, b, tol, brackets)?;
            }
            return Ok(());
        }
    }

    bisect_isolate_inner(sturm, a, mid, tol, brackets)?;
    bisect_isolate_inner(sturm, mid, b, tol, brackets)?;
    Ok(())
}

fn bisect_refine(poly: &Poly, mut a: f64, mut b: f64, mut fa: f64, tol: f64) -> f64 {
    for _ in 0..200 {
        if (b - a).abs() < tol {
            break;
        }
        let mid = (a + b) / 2.0;
        let fm = poly.eval_f64(mid);
        if fm == 0.0 {
            return mid;
        }
        if fa * fm < 0.0 {
            b = mid;
        } else {
            a = mid;
            fa = fm;
        }
    }
    (a + b) / 2.0
}