ocas-calc 0.27.3

//! Rational function integration (Bronstein, *Symbolic Integration I*, ch. 2).
//!
//! [`integrate_rational`] computes the integral of a rational function of
//! the integration variable over `ℚ`:
//!
//! 1. the polynomial part by termwise integration,
//! 2. the rational part by Hermite reduction,
//! 3. the logarithmic part by the logarithmic-derivative identity
//!    (`c·f'/f → c·log(f)`), completing the square (degree-2 denominators,
//!    yielding `log` or `atan`), or the Rothstein–Trager resultant for
//!    general squarefree denominators.
//!
//! Terms whose Rothstein–Trager resultant does not split into rational
//! roots cannot be expressed without `RootSum`-style notation; those terms
//! are returned in the unevaluated form `Integral(term, var)`.

use ocas_atom::normalize::normalize;
use ocas_atom::{Atom, AtomArena, Symbol};
use ocas_domain::{Domain, Integer, Rational, RationalDomain};
use ocas_poly::{DenseUnivariatePolynomial, Lex, SparseMultivariatePolynomial};
use ocas_rewrite::rules::default_rules;
use ocas_rewrite::simplify::simplify;

use crate::rules::calculus_rules;
use crate::tower::convert::{
    GeneratorField, atom_to_rational, rational_const_to_atom, rational_to_atom,
};

type DPoly = DenseUnivariatePolynomial<RationalDomain>;
type Sparse = SparseMultivariatePolynomial<RationalDomain, Lex>;

// =========================================================================
// Deterministic work budget
// =========================================================================
//
// Two loops in this module are unbounded in the input degree:
//
// - `reduce_rational`'s `extended_gcd_poly` and the resultant/GCD series in
//   `rothstein_trager`, whose interpolation evaluates `n + 1` degree-`n`
//   resultants with coefficient growth, and
// - `rational_roots`, which runs the integer factorization of the
//   resultant.
//
// Both are charged against a per-entry counter, and `rothstein_trager`
// refuses structurally oversized denominators before the interpolation
// allocates anything. Tripping a gate declines the stage (the chain then
// reports the ordinary `Integral(...)` fallback), so no wrong answer is
// ever emitted.

/// Degree above which the Rothstein–Trager interpolation is declined. The
/// corpus rational cases stay in the single digits; the substitution t-forms
/// that reach here are the ones that used to grind.
const MAX_RT_DEGREE: usize = 32;

/// Degree above which the common-factor cancellation is skipped (the
/// resulting unreduced fraction only makes the later stages decline, which
/// is the documented best-effort behaviour of `reduce_rational`).
const MAX_CANCEL_DEGREE: usize = 128;

/// Work units a single entry into this module may charge.
const MAX_RATIONAL_WORK: u64 = 200_000;

thread_local! {
    static RATIONAL_WORK: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}

/// Reset the work budget; called by the public stage entry point.
fn reset_rational_budget() {
    RATIONAL_WORK.with(|c| c.set(0));
}

/// Charge `units` against the module budget; `true` when exhausted.
fn charge_rational_work(units: u64) -> bool {
    RATIONAL_WORK.with(|c| {
        let v = c.get().saturating_add(units);
        c.set(v);
        v > MAX_RATIONAL_WORK
    })
}

/// Integrate a rational function of `var` over `ℚ`.
///
/// Returns `None` when `expr` is not a rational function of `var` alone
/// (other variables or function applications are present). Terms that
/// would require algebraic numbers beyond `√` are returned in the
/// unevaluated form `Integral(term, var)` inside the sum.
///
/// # Example
///
/// ```
/// use ocas_atom::{AtomArena, Symbol};
/// use ocas_calc::integral::rational::integrate_rational;
/// use ocas_core::arena::Arena;
///
/// let arena = Arena::new();
/// let ctx = AtomArena::new(&arena);
/// let x = ctx.var("x");
/// // ∫ x^-1 dx = log(x)
/// let result = integrate_rational(&ctx, ctx.pow(x, ctx.num(-1)), Symbol::new("x"));
/// assert_eq!(result.unwrap().to_string(), "log(x)");
/// ```
pub fn integrate_rational<'a>(
    ctx: &'a AtomArena<'a>,
    expr: Atom<'a>,
    var: Symbol,
) -> Option<Atom<'a>> {
    // Fresh budget for this invocation: the stage is entered once per chain
    // re-entry, so nothing may accumulate across calls.
    reset_rational_budget();
    let x = ctx.var(var.as_str());
    let gens = [x];
    let rf = atom_to_rational(expr, &gens)?;
    let num = sparse_to_dense(&rf.numerator);
    let den = sparse_to_dense(&rf.denominator);
    let (num, den) = reduce_rational(num, den);
    let raw = integrate_dpoly(ctx, &num, &den, x)?;
    let calc_rules = calculus_rules(ctx, &crate::pattern_alloc::VecAlloc);
    let default_rules = default_rules(ctx, &crate::pattern_alloc::VecAlloc);
    let after_default = simplify(ctx, raw, &default_rules, 20);
    let after_calc = simplify(ctx, after_default, &calc_rules, 10);
    Some(normalize(ctx, after_calc))
}

/// Cancel the polynomial GCD of `num`/`den` before integrating.
///
/// Substitution-based integrators (Euler, Weierstrass) build t-rationals
/// with common factors between numerator and denominator; without
/// reduction the Rothstein–Trager step would see an irreducible factor
/// that the reduced form no longer has and return an unevaluated
/// `Integral` for a rational integral that is in fact elementary.
fn reduce_rational(num: DPoly, den: DPoly) -> (DPoly, DPoly) {
    if num.is_zero() || den.is_zero() {
        return (num, den);
    }
    // Structural gate: the extended Euclidean run is degree-driven, so a
    // pathologically large pair skips the (best-effort) cancellation
    // instead of grinding.
    let deg_sum = num.degree().unwrap_or(0) + den.degree().unwrap_or(0);
    if deg_sum > MAX_CANCEL_DEGREE || charge_rational_work(deg_sum as u64) {
        return (num, den);
    }
    let (g, _, _) = num.extended_gcd_poly(&den);
    if g.is_zero() || g.degree() == Some(0) {
        return (num, den);
    }
    let num = num.div_rem(&g).map(|(q, _)| q).unwrap_or(num);
    let den = den.div_rem(&g).map(|(q, _)| q).unwrap_or(den);
    (num, den)
}

/// Core integrator for `num / den` over `ℚ[x]`, building atoms over `x`.
pub(crate) fn integrate_dpoly<'a>(
    ctx: &'a AtomArena<'a>,
    num: &DPoly,
    den: &DPoly,
    x: Atom<'a>,
) -> Option<Atom<'a>> {
    let s = integrate_dpoly_structured(ctx, num, den, x)?;
    let mut parts: Vec<Atom> = Vec::new();
    if !s.poly.is_zero() {
        parts.push(poly_atom(ctx, &s.poly, x)?);
    }
    if let Some((g_num, g_den)) = &s.rational {
        let g = GeneratorField::from_num_den(dense_to_sparse(g_num), dense_to_sparse(g_den));
        parts.push(rational_to_atom(ctx, &g, &[x])?);
    }
    for (c, v) in &s.log_terms {
        let v_atom = poly_atom(ctx, v, x)?;
        let log = ctx.fun("log", &[v_atom]);
        parts.push(scale_atom(ctx, c, log)?);
    }
    parts.extend(s.extra_atoms.iter().copied());
    Some(match parts.len() {
        0 => ctx.num(0),
        1 => parts[0],
        _ => ctx.add(&parts),
    })
}

/// A structured rational-function integral (used by the Risch tower's
/// base level, which needs the pieces rather than a assembled atom).
pub(crate) struct DPolyIntegral<'a> {
    /// Integral of the polynomial part.
    pub poly: DPoly,
    /// Hermite rational part `(num, den)`, if any.
    pub rational: Option<(DPoly, DPoly)>,
    /// Logarithmic terms `c·log(v)` with `v ∈ ℚ[x]`.
    pub log_terms: Vec<(Rational, DPoly)>,
    /// Other terms (atan, radical logs, unevaluated `Integral`s).
    pub extra_atoms: Vec<Atom<'a>>,
}

/// Structured core of [`integrate_dpoly`].
pub(crate) fn integrate_dpoly_structured<'a>(
    ctx: &'a AtomArena<'a>,
    num: &DPoly,
    den: &DPoly,
    x: Atom<'a>,
) -> Option<DPolyIntegral<'a>> {
    if den.is_zero() {
        return None;
    }
    let (p, r) = num.div_rem(den)?;
    let mut out = DPolyIntegral {
        poly: DPoly::from_coeffs(RationalDomain, vec![]),
        rational: None,
        log_terms: Vec::new(),
        extra_atoms: Vec::new(),
    };
    if !p.is_zero() {
        out.poly = poly_integrate(&p);
    }
    if !r.is_zero() {
        let (g_num, g_den, a1, d1) = hermite_reduce(&r, den);
        if !g_num.is_zero() {
            out.rational = Some((g_num, g_den));
        }
        if !a1.is_zero() {
            for term in log_part(ctx, &a1, &d1, x)? {
                match term {
                    DLogOrAtom::Log(c, v) => out.log_terms.push((c, v)),
                    DLogOrAtom::Atom(a) => out.extra_atoms.push(a),
                }
            }
        }
    }
    Some(out)
}

// ------------------------------------------------------------------
//  Polynomial part
// ------------------------------------------------------------------

/// Termwise integral of a polynomial: `∫ Σ cₖ xᵏ = Σ cₖ/(k+1) xᵏ⁺¹`.
pub(crate) fn poly_integrate(p: &DPoly) -> DPoly {
    let dom = RationalDomain;
    let mut coeffs = vec![dom.zero()];
    for (k, c) in p.coeffs().iter().enumerate() {
        let scale = Rational::new(1, (k + 1) as i64);
        coeffs.push(dom.mul(c, &scale));
    }
    DPoly::from_coeffs(RationalDomain, coeffs)
}

// ------------------------------------------------------------------
//  Hermite reduction
// ------------------------------------------------------------------

/// Hermite reduction of a proper fraction `a / d` with `d` monic.
///
/// Returns `(g_num, g_den, a1, d1)` such that `a/d = (g_num/g_den)' + a1/d1`
/// with `d1` squarefree.
fn hermite_reduce(a: &DPoly, d: &DPoly) -> (DPoly, DPoly, DPoly, DPoly) {
    let factors = d.square_free_factorization();
    let m = factors.iter().map(|&(_, k)| k).max().unwrap_or(1);
    if m <= 1 {
        return (a.zero(), a.one(), a.clone(), d.clone());
    }
    // Pick a factor v of maximal multiplicity m; write d = u * v^m.
    let (v, _) = factors
        .iter()
        .rev()
        .find(|&(_, k)| *k == m)
        .expect("max multiplicity factor exists");
    let vm = v.pow(m as u32);
    let (u, rem) = d.div_rem(&vm).expect("v^m divides d");
    debug_assert!(rem.is_zero());

    // Solve s*B + t*v = a with B = u*v' and deg s < deg v (gcd(B, v) = 1).
    let b = u.mul(&v.derivative());
    let (g, s0, _t0) = b.extended_gcd_poly(v);
    debug_assert!(g.is_one());
    let (_, s) = a.mul(&s0).div_rem(v).expect("field division");
    let (t, rem) = a.sub(&s.mul(&b)).div_rem(v).expect("exact quotient");
    debug_assert!(rem.is_zero());

    // ∫ a/(u v^m) = -(s/(m-1))/v^(m-1) + ∫ (t + u*(s/(m-1))') / (u v^(m-1))
    let inv_m1 = Rational::new(1, (m - 1) as i64);
    let s_scaled = s.mul_scalar(&inv_m1);
    let vm1 = v.pow((m - 1) as u32);
    let new_a = t.add(&u.mul(&s_scaled.derivative()));
    let new_d = u.mul(&vm1);
    let (g2n, g2d, a1, d1) = hermite_reduce(&new_a, &new_d);

    // Combine rational parts: -s_scaled/vm1 + g2n/g2d.
    let g_num = s_scaled.neg().mul(&g2d).add(&g2n.mul(&vm1));
    let g_den = vm1.mul(&g2d);
    (g_num, g_den, a1, d1)
}

// ------------------------------------------------------------------
//  Logarithmic part
// ------------------------------------------------------------------

/// A logarithmic-part term: either a structured `c·log(v)` or an opaque
/// atom (atan, radical forms, unevaluated integrals).
pub(crate) enum DLogOrAtom<'a> {
    /// `c·log(v)` with `v ∈ ℚ[x]`.
    Log(Rational, DPoly),
    /// Any other term as an atom.
    Atom(Atom<'a>),
}

/// Integrate `a / d` with `d` monic and squarefree. Returns the summands;
/// may include an unevaluated `Integral` term.
fn log_part<'a>(
    ctx: &'a AtomArena<'a>,
    a: &DPoly,
    d: &DPoly,
    x: Atom<'a>,
) -> Option<Vec<DLogOrAtom<'a>>> {
    let dom = RationalDomain;
    // Case 1: a == c·d' → c·log(d) (covers deg d == 1).
    let dp = d.derivative();
    if !dp.is_zero() && a.degree() == dp.degree() {
        let c = dom.div(&a.lcoeff(), &dp.lcoeff())?;
        if *a == dp.mul_scalar(&c) {
            return Some(vec![DLogOrAtom::Log(c, d.clone())]);
        }
    }
    // Case 2: degree-2 denominator → completing the square.
    if d.degree() == Some(2) {
        return complete_square(ctx, a, d, x);
    }
    // Case 3: Rothstein–Trager for general squarefree denominators.
    rothstein_trager(ctx, a, d, x)
}

/// Integrate `(p·x + q) / (A·x² + B·x + C)` with `A·x² + B·x + C` squarefree.
fn complete_square<'a>(
    ctx: &'a AtomArena<'a>,
    a: &DPoly,
    d: &DPoly,
    x: Atom<'a>,
) -> Option<Vec<DLogOrAtom<'a>>> {
    let dom = RationalDomain;
    let c0 = d.coeffs().first().cloned().unwrap_or_else(|| dom.zero());
    let c1 = d.coeffs().get(1).cloned().unwrap_or_else(|| dom.zero());
    let c2 = d.coeffs().get(2).cloned().unwrap_or_else(|| dom.zero());
    let big_a = a.coeffs().get(1).cloned().unwrap_or_else(|| dom.zero());
    let big_b = a.coeffs().first().cloned().unwrap_or_else(|| dom.zero());

    // The Hermite recursion can hand back a square-free denominator whose
    // leading coefficient is not 1 (`1 - x²` for `(x² - 1)²`). Assuming a
    // monic denominator flips the sign of the discriminant and selects the
    // `atan` branch for a quadratic that in fact has two rational roots — a
    // silently wrong answer (0.27.1 regression: `∫dx/(x²-1)²` returned
    // `atan(x)/2`). Every formula below therefore uses the *actual* leading
    // coefficient `A = c2` (no rescaling of the operands).
    let two_a = dom.mul(&Rational::new(2, 1), &c2);
    if dom.is_zero(&two_a) {
        return None;
    }

    let mut out = Vec::new();
    // (p/(2A))·log(d)
    if !dom.is_zero(&big_a) {
        let coeff = dom.div(&big_a, &two_a)?;
        out.push(DLogOrAtom::Log(coeff, d.clone()));
    }
    // Remaining constant numerator: (q - p·B/(2A))·∫dx/(A·x²+B·x+C).
    let rem = dom.sub(&big_b, &dom.div(&dom.mul(&big_a, &c1), &two_a)?);
    if dom.is_zero(&rem) {
        return Some(out);
    }
    // Discriminant Δ = B² - 4·A·C (nonzero: d is squarefree).
    let delta = dom.sub(
        &dom.mul(&c1, &c1),
        &dom.mul(&dom.mul(&Rational::new(4, 1), &c2), &c0),
    );
    let lin = linear_atom(ctx, &two_a, &c1, x)?; // 2A·x + B
    if rat_is_negative(&delta) {
        // atan branch: ∫ = rem·(2/s)·atan((2A·x+B)/s), s = √(4AC-B²).
        let s2 = dom.neg(&delta);
        match sqrt_positive_rational(ctx, &s2)? {
            Sqrt::Rat(sr) => {
                let sr_inv = dom.inv(&sr)?;
                let arg = linear_atom(ctx, &dom.mul(&two_a, &sr_inv), &dom.mul(&c1, &sr_inv), x)?;
                let coeff = dom.mul(&rem, &dom.mul(&Rational::new(2, 1), &sr_inv));
                out.push(DLogOrAtom::Atom(scale_atom(
                    ctx,
                    &coeff,
                    ctx.fun("atan", &[arg]),
                )?));
            }
            Sqrt::Rad(s) => {
                let arg = ctx.mul(&[lin, ctx.pow(s, ctx.num(-1))]);
                let atan = ctx.fun("atan", &[arg]);
                let two_rem = dom.mul(&rem, &Rational::new(2, 1));
                out.push(DLogOrAtom::Atom(ctx.mul(&[
                    rational_const_to_atom(ctx, &two_rem)?,
                    ctx.pow(s, ctx.num(-1)),
                    atan,
                ])));
            }
        }
    } else {
        // log branch: ∫ = rem·(1/s)·log((2A·x+B-s)/(2A·x+B+s)), s = √Δ.
        match sqrt_positive_rational(ctx, &delta)? {
            Sqrt::Rat(sr) => {
                let num_atom = linear_atom(ctx, &two_a, &dom.sub(&c1, &sr), x)?;
                let den_atom = linear_atom(ctx, &two_a, &dom.add(&c1, &sr), x)?;
                let ratio = ctx.mul(&[num_atom, ctx.pow(den_atom, ctx.num(-1))]);
                let log = ctx.fun("log", &[ratio]);
                let coeff = dom.mul(&rem, &dom.inv(&sr)?);
                out.push(DLogOrAtom::Atom(scale_atom(ctx, &coeff, log)?));
            }
            Sqrt::Rad(s) => {
                let num_atom = ctx.add(&[lin, ctx.mul(&[ctx.num(-1), s])]);
                let den_atom = ctx.add(&[lin, s]);
                let ratio = ctx.mul(&[num_atom, ctx.pow(den_atom, ctx.num(-1))]);
                let log = ctx.fun("log", &[ratio]);
                out.push(DLogOrAtom::Atom(ctx.mul(&[
                    rational_const_to_atom(ctx, &rem)?,
                    ctx.pow(s, ctx.num(-1)),
                    log,
                ])));
            }
        }
    }
    Some(out)
}

/// Rothstein–Trager: `∫ a/d = Σᵢ cᵢ·log(gcd(d, a - cᵢ·d'))` where the `cᵢ`
/// are the roots of `R(t) = resultantₓ(d, a - t·d')`.
///
/// `R(t)` is recovered by interpolation (resultants are univariate over
/// `ℚ`, so `t` cannot appear symbolically in the coefficients). Falls back
/// to an unevaluated `Integral` when `R` does not split over `ℚ`.
fn rothstein_trager<'a>(
    ctx: &'a AtomArena<'a>,
    a: &DPoly,
    d: &DPoly,
    x: Atom<'a>,
) -> Option<Vec<DLogOrAtom<'a>>> {
    let n = d.degree()?;
    // Structural gate before any resultant work: the interpolation below
    // evaluates `n + 1` resultants of degree-`n` polynomials, so an
    // oversized square-free denominator declines outright.
    if n > MAX_RT_DEGREE {
        return None;
    }
    let dp = d.derivative();

    // Interpolate R(t) = resultant(d, a - t·d') at t = 0, …, n.
    let mut points = Vec::with_capacity(n + 1);
    for j in 0..=(n as i64) {
        if charge_rational_work(1 + (n * n) as u64) {
            return None;
        }
        let tj = Rational::new(j, 1);
        let shifted = a.sub(&dp.mul_scalar(&tj));
        let val = d.resultant(&shifted);
        points.push((tj, val));
    }
    let rt = lagrange_interpolate(&points);

    let (roots, fully_split) = rational_roots(&rt)?;
    if !fully_split {
        return Some(vec![DLogOrAtom::Atom(integral_fallback(ctx, a, d, x)?)]);
    }

    let mut out = Vec::new();
    for c in &roots {
        let v = d.gcd(&a.sub(&dp.mul_scalar(c)));
        if v.degree().unwrap_or(0) == 0 {
            continue;
        }
        out.push(DLogOrAtom::Log(c.clone(), v));
    }
    if out.is_empty() {
        return Some(vec![DLogOrAtom::Atom(integral_fallback(ctx, a, d, x)?)]);
    }
    Some(out)
}

// ------------------------------------------------------------------
//  Helpers
// ------------------------------------------------------------------

/// Lagrange interpolation over `ℚ` from point values.
fn lagrange_interpolate(points: &[(Rational, Rational)]) -> DPoly {
    let dom = RationalDomain;
    let mut result = DPoly::from_coeffs(RationalDomain, vec![]);
    for (j, (xj, yj)) in points.iter().enumerate() {
        let mut basis = DPoly::from_coeffs(RationalDomain, vec![dom.one()]);
        let mut denom = dom.one();
        for (k, (xk, _)) in points.iter().enumerate() {
            if k == j {
                continue;
            }
            if charge_rational_work(points.len() as u64) {
                // Returning the partial interpolation would be wrong. The
                // zero polynomial is not: `rational_roots` reads it as "no
                // rational roots, fully split", so `rothstein_trager` emits
                // its documented unevaluated-`Integral` fallback term — an
                // honest partial result rather than a wrong answer.
                return DPoly::from_coeffs(RationalDomain, vec![]);
            }
            basis = basis.mul(&DPoly::from_coeffs(
                RationalDomain,
                vec![dom.neg(xk), dom.one()],
            ));
            denom = dom.mul(&denom, &dom.sub(xj, xk));
        }
        let scale = dom.div(yj, &denom).expect("distinct interpolation nodes");
        result = result.add(&basis.mul_scalar(&scale));
    }
    result
}

/// Distinct rational roots of a polynomial over `ℚ`, plus whether the
/// polynomial splits completely into linear factors over `ℚ`.
fn rational_roots(f: &DPoly) -> Option<(Vec<Rational>, bool)> {
    let dom = RationalDomain;
    if f.is_zero() || f.degree()? == 0 {
        return Some((Vec::new(), true));
    }
    // Clear denominators to get an integer polynomial.
    let mut lcm: i64 = 1;
    for c in f.coeffs() {
        lcm = num_integer::lcm(lcm, c.denom().to_i64()?);
    }
    let zcoeffs: Option<Vec<Integer>> = f
        .coeffs()
        .iter()
        .map(|c| {
            let scaled = dom.mul(c, &Rational::new(lcm, 1));
            scaled.numer().to_i64().map(Integer::from)
        })
        .collect();
    let zpoly = DenseUnivariatePolynomial::from_coeffs(ocas_domain::IntegerDomain, zcoeffs?);
    // Charge the integer factorization (the expensive step of this
    // function); a decline here surfaces as the unevaluated-`Integral`
    // fallback term, never as a wrong answer.
    if charge_rational_work(1 + zpoly.degree().unwrap_or(0) as u64 * 64) {
        return None;
    }
    let primitive = zpoly.primitive_part();
    let factors = primitive.factor();

    let mut roots = Vec::new();
    let mut split = true;
    for (fac, _mult) in &factors {
        if fac.degree() == Some(1) {
            let c0 = fac.coeffs().first()?.to_i64()?;
            let c1 = fac.coeffs().get(1)?.to_i64()?;
            roots.push(Rational::new(-c0, c1));
        } else {
            split = false;
        }
    }
    Some((roots, split))
}

/// Square root of a positive rational: exact when the squarefree numerator
/// is a perfect square, otherwise a radical atom `√(p·q) / q`.
enum Sqrt<'a> {
    Rat(Rational),
    Rad(Atom<'a>),
}

fn sqrt_positive_rational<'a>(ctx: &'a AtomArena<'a>, r: &Rational) -> Option<Sqrt<'a>> {
    let p = r.numer().to_i64()?;
    let q = r.denom().to_i64()?;
    let n = (p as i128).checked_mul(q as i128)?;
    let m = isqrt_i128(n);
    if m * m == n {
        let m = i64::try_from(m).ok()?;
        Some(Sqrt::Rat(Rational::new(m, q)))
    } else {
        let n = i64::try_from(n).ok()?;
        // √(p/q) = √(p·q)/q — the 1/q factor is part of the value.
        Some(Sqrt::Rad(ctx.mul(&[
            ctx.pow(ctx.num(n), ctx.pow(ctx.num(2), ctx.num(-1))),
            ctx.pow(ctx.num(q), ctx.num(-1)),
        ])))
    }
}

fn isqrt_i128(n: i128) -> i128 {
    // Standard-library integer square root (stabilized in Rust 1.84).
    // The caller guarantees n ≥ 0 via checked_mul on positive numer/denom;
    // checked_isqrt returns None only for n < 0, where we fall back to 0
    // (matching the prior hand-rolled Newton iteration behaviour).
    n.checked_isqrt().unwrap_or(0)
}

fn rat_is_negative(r: &Rational) -> bool {
    use num_traits::Signed;
    r.inner().is_negative()
}

/// `coeff · atom`, eliding a unit coefficient.
fn scale_atom<'a>(ctx: &'a AtomArena<'a>, coeff: &Rational, atom: Atom<'a>) -> Option<Atom<'a>> {
    if RationalDomain.is_one(coeff) {
        return Some(atom);
    }
    Some(ctx.mul(&[rational_const_to_atom(ctx, coeff)?, atom]))
}

/// `a·x + b` with `a ≠ 0`, eliding a zero constant term.
fn linear_atom<'a>(
    ctx: &'a AtomArena<'a>,
    a: &Rational,
    b: &Rational,
    x: Atom<'a>,
) -> Option<Atom<'a>> {
    let dom = RationalDomain;
    let ax = if dom.is_one(a) {
        x
    } else {
        ctx.mul(&[rational_const_to_atom(ctx, a)?, x])
    };
    if dom.is_zero(b) {
        return Some(ax);
    }
    Some(ctx.add(&[ax, rational_const_to_atom(ctx, b)?]))
}

fn poly_atom<'a>(ctx: &'a AtomArena<'a>, p: &DPoly, x: Atom<'a>) -> Option<Atom<'a>> {
    rational_to_atom(
        ctx,
        &GeneratorField::from_polynomial(dense_to_sparse(p)),
        &[x],
    )
}

fn integral_fallback<'a>(
    ctx: &'a AtomArena<'a>,
    a: &DPoly,
    d: &DPoly,
    x: Atom<'a>,
) -> Option<Atom<'a>> {
    let rf = GeneratorField::from_num_den(dense_to_sparse(a), dense_to_sparse(d));
    let expr = rational_to_atom(ctx, &rf, &[x])?;
    Some(ctx.fun("Integral", &[expr, x]))
}

fn sparse_to_dense(p: &Sparse) -> DPoly {
    debug_assert_eq!(p.n_vars(), 1);
    let deg = p.degree_in(0);
    let mut coeffs = vec![RationalDomain.zero(); deg + 1];
    for (exp, coeff) in p.terms_ref() {
        coeffs[exp[0]] = coeff.clone();
    }
    DPoly::from_coeffs(RationalDomain, coeffs)
}

fn dense_to_sparse(p: &DPoly) -> Sparse {
    let terms = p
        .coeffs()
        .iter()
        .enumerate()
        .filter(|&(_, c)| !RationalDomain.is_zero(c))
        .map(|(i, c)| (vec![i], c.clone()))
        .collect();
    Sparse::from_terms(RationalDomain, 1, terms)
}

#[cfg(test)]
mod tests {
    use ocas_atom::AtomArena;
    use ocas_core::arena::Arena;

    use super::*;
    use crate::derivative::diff;
    use ocas_atom::AtomNode;

    /// Verify that `diff(result) == integrand` as canonical rational functions.
    fn assert_antiderivative<'a>(
        ctx: &'a AtomArena<'a>,
        result: Atom<'a>,
        integrand: Atom<'a>,
        x: Atom<'a>,
        var: Symbol,
    ) {
        let d = diff(ctx, result, var);
        let lhs = atom_to_rational(d, &[x]).expect("derivative is rational in x");
        let rhs = atom_to_rational(integrand, &[x]).expect("integrand is rational in x");
        assert_eq!(lhs, rhs, "d/dx(result) != integrand");
    }

    /// Minimal f64 evaluator for numeric checks on radical-containing results.
    fn eval_f64(atom: Atom, x: f64) -> f64 {
        match atom.node() {
            AtomNode::Num(n) => *n as f64,
            AtomNode::Var(_) => x,
            AtomNode::Add(args) => args.iter().map(|a| eval_f64(*a, x)).sum(),
            AtomNode::Mul(args) => args.iter().map(|a| eval_f64(*a, x)).product(),
            AtomNode::Pow(b, e) => eval_f64(*b, x).powf(eval_f64(*e, x)),
            AtomNode::Fun(name, args) => match name.as_str() {
                // |u| keeps the real evaluation finite; log|u| and log(u)
                // differ by a constant on each interval, so derivatives match.
                "log" => eval_f64(args[0], x).abs().ln(),
                "atan" => eval_f64(args[0], x).atan(),
                "sin" => eval_f64(args[0], x).sin(),
                "cos" => eval_f64(args[0], x).cos(),
                "exp" => eval_f64(args[0], x).exp(),
                other => panic!("unsupported function in test eval: {other}"),
            },
        }
    }

    /// Numeric finite-difference check that `result'` matches `integrand`.
    fn assert_numeric_antiderivative(result: Atom, integrand: Atom) {
        let h = 1e-6;
        for px in [0.7f64, 1.3, 2.1] {
            let approx = (eval_f64(result, px + h) - eval_f64(result, px - h)) / (2.0 * h);
            let exact = eval_f64(integrand, px);
            assert!(
                (approx - exact).abs() < 1e-4,
                "numeric derivative mismatch at x={px}: {approx} != {exact}"
            );
        }
    }

    fn integrate_str<'a>(ctx: &'a AtomArena<'a>, expr: Atom<'a>) -> Atom<'a> {
        integrate_rational(ctx, expr, Symbol::new("x")).expect("rational integral")
    }

    #[test]
    fn inverse_gives_log() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        let result = integrate_str(&ctx, ctx.pow(x, ctx.num(-1)));
        assert_eq!(result.to_string(), "log(x)");
    }

    #[test]
    fn atan_branch_sqrt_rational_keeps_denominator() {
        // Regression: √(p/q) must be √(p·q)/q — a dropped 1/q made
        // ∫ dx/(3x²+4x+3) return a wrong coefficient/argument silently.
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        let integrand = ctx.pow(
            ctx.add(&[
                ctx.mul(&[ctx.num(3), ctx.pow(x, ctx.num(2))]),
                ctx.mul(&[ctx.num(4), x]),
                ctx.num(3),
            ]),
            ctx.num(-1),
        );
        let result = integrate_str(&ctx, integrand);
        assert_numeric_antiderivative(result, integrand);
    }

    #[test]
    fn polynomial_part() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ (x^2 + 3x + 2) dx
        let expr = ctx.add(&[
            ctx.pow(x, ctx.num(2)),
            ctx.mul(&[ctx.num(3), x]),
            ctx.num(2),
        ]);
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn atan_over_quadratic() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ 1/(x^2 + 1) dx = atan(x)
        let expr = ctx.pow(ctx.add(&[ctx.pow(x, ctx.num(2)), ctx.num(1)]), ctx.num(-1));
        let result = integrate_str(&ctx, expr);
        assert_eq!(result.to_string(), "atan(x)");
    }

    #[test]
    fn log_over_quadratic_irrational_discriminant() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ 1/(x^2 - 2) dx = (1/(2√2))·log((x-√2)/(x+√2))
        let expr = ctx.pow(ctx.add(&[ctx.pow(x, ctx.num(2)), ctx.num(-2)]), ctx.num(-1));
        let result = integrate_str(&ctx, expr);
        assert!(result.to_string().contains("log"));
        assert_numeric_antiderivative(result, expr);
    }

    #[test]
    fn log_over_quadratic_rational_roots() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ 1/(x^2 - 1) dx = (1/2)·log((x-1)/(x+1))
        let expr = ctx.pow(ctx.add(&[ctx.pow(x, ctx.num(2)), ctx.num(-1)]), ctx.num(-1));
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn logarithmic_derivative_identity() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ (2x + 3)/(x^2 + 3x + 5) dx = log(x^2 + 3x + 5)
        let num = ctx.add(&[ctx.mul(&[ctx.num(2), x]), ctx.num(3)]);
        let den = ctx.add(&[
            ctx.pow(x, ctx.num(2)),
            ctx.mul(&[ctx.num(3), x]),
            ctx.num(5),
        ]);
        let expr = ctx.mul(&[num, ctx.pow(den, ctx.num(-1))]);
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn hermite_repeated_factor() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ 1/(x+1)^2 dx = -(x+1)^-1
        let expr = ctx.pow(ctx.add(&[x, ctx.num(1)]), ctx.num(-2));
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn hermite_mixed_multiplicities() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ x/( (x+1)^2 (x+2) ) dx — repeated factor times simple factor.
        let den = ctx.mul(&[
            ctx.pow(ctx.add(&[x, ctx.num(1)]), ctx.num(2)),
            ctx.add(&[x, ctx.num(2)]),
        ]);
        let expr = ctx.mul(&[x, ctx.pow(den, ctx.num(-1))]);
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn rothstein_trager_cubic() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ 1/(x^3 + x) dx: squarefree cubic, roots 0, ±i — splits over ℚ
        // only partially; RT resultant has one rational root (0)... The
        // integrand = 1/x - x/(x^2+1), so a fully rational-form answer
        // exists via partial fractions with rational terms only after
        // combining conjugates. Verify by differentiation.
        let expr = ctx.pow(ctx.add(&[ctx.pow(x, ctx.num(3)), x]), ctx.num(-1));
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    #[test]
    fn non_rational_returns_none() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // sin(x) is not a rational function of x.
        let expr = ctx.fun("sin", &[x]);
        assert!(integrate_rational(&ctx, expr, Symbol::new("x")).is_none());
    }

    #[test]
    fn higher_multiplicity_hermite() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let x = ctx.var("x");
        // ∫ (x^2 + 1)/(x-1)^3 dx — multiplicity 3 denominator.
        let num = ctx.add(&[ctx.pow(x, ctx.num(2)), ctx.num(1)]);
        let den = ctx.pow(ctx.add(&[x, ctx.num(-1)]), ctx.num(3));
        let expr = ctx.mul(&[num, ctx.pow(den, ctx.num(-1))]);
        let result = integrate_str(&ctx, expr);
        assert_antiderivative(&ctx, result, expr, x, Symbol::new("x"));
    }

    // ------------------- deterministic budget (0.27.1 timeouts) ---------

    /// A high-degree rational function whose Rothstein–Trager step used to
    /// run for the whole per-case budget. The degree gate must decline it
    /// (or return an honest `Integral` term), never a wrong answer.
    #[test]
    fn oversized_rothstein_trager_declines() {
        // Squarefree-ish degree-40 denominator: 1/((x^4+1)(x^5+2)…).
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let expr = ocas_parse::parse(&ctx, "1/((x^4+1)*(x^5+2)*(x^7+3)*(x^9+5)*(x^11+7))").unwrap();
        let var = Symbol::new("x");
        let r = integrate_rational(&ctx, expr, var);
        if let Some(atom) = r {
            let s = atom.to_string();
            assert!(
                s.contains("Integral(") || !s.is_empty(),
                "unexpected empty result"
            );
        }
    }

    /// Every previously-solved rational shape must keep solving.
    #[test]
    fn budget_keeps_solving_normal_inputs() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let var = Symbol::new("x");
        for input in [
            "1/(x^2 + 2*x + 3)",
            "1/(x^3 + x)",
            "x/((x + 1)^2*(x + 2))",
            "(x^2 + 1)/(x - 1)^3",
            "1/(3*x^2 + 4*x + 3)",
        ] {
            let expr = ocas_parse::parse(&ctx, input).unwrap();
            let result = integrate_rational(&ctx, expr, var).expect("solved before");
            assert!(
                !result.to_string().contains("Integral("),
                "{input} regressed: {result}"
            );
            assert_numeric_antiderivative(result, expr);
        }
    }

    /// Regression (0.27.1): the square-free denominator can come back from
    /// the Hermite recursion with a negative leading coefficient
    /// (`1 - x²` for `(x² - 1)²`), which used to flip the discriminant sign
    /// and emit `atan` for a quadratic that has two rational roots —
    /// `∫dx/(x²-1)²` returned `atan(x)/2`, a silently wrong answer.
    #[test]
    fn repeated_quadratic_with_rational_roots_uses_logs() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let var = Symbol::new("x");
        for input in [
            "1/(x^2 - 1)^2",
            "1/(1 - x^2)^2",
            "x^2/(x^2 - 1)^2",
            "1/(x^2 - 4)^2",
            "(x + 1)/(x^2 - 1)^2",
        ] {
            let integrand = ocas_parse::parse(&ctx, input).unwrap();
            let result = integrate_rational(&ctx, integrand, var)
                .unwrap_or_else(|| panic!("declined {input}"));
            let s = result.to_string();
            assert!(
                !s.contains("atan"),
                "{input} took the atan branch although the denominator has \
                 rational roots: {s}"
            );
            assert_numeric_antiderivative(result, integrand);
        }
    }

    /// The genuinely irreducible quadratics must keep their `atan` branch.
    #[test]
    fn irreducible_quadratic_keeps_atan() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let var = Symbol::new("x");
        for input in ["1/(x^2 + 1)^2", "1/(x^2 + 2)^2", "1/((x - 1)^2 + 1)^2"] {
            let integrand = ocas_parse::parse(&ctx, input).unwrap();
            let result = integrate_rational(&ctx, integrand, var)
                .unwrap_or_else(|| panic!("declined {input}"));
            assert!(
                result.to_string().contains("atan"),
                "{input} lost its atan branch: {result}"
            );
            assert_numeric_antiderivative(result, integrand);
        }
    }

    /// The budget resets on every entry: repeating a shape 200 times must
    /// give identical results and stay fast.
    ///
    /// The stress shape is the `1/(x²-1)²` family rather than the degree-36
    /// product used by `oversized_rothstein_trager_declines`: the latter
    /// overflows the stack in the dense-polynomial layer when driven
    /// repeatedly (pre-existing, independent of the budget), and this shape
    /// exercises the same log-part entry point.
    #[test]
    fn budget_does_not_leak_across_calls() {
        let arena = Arena::new();
        let ctx = AtomArena::new(&arena);
        let expr = ocas_parse::parse(&ctx, "1/(x^2 - 1)^2").unwrap();
        let var = Symbol::new("x");
        let start = std::time::Instant::now();
        let first = integrate_rational(&ctx, expr, var).map(|a| a.to_string());
        for i in 0..200 {
            let r = integrate_rational(&ctx, expr, var).map(|a| a.to_string());
            assert_eq!(r, first, "call {i} diverged");
        }
        assert!(
            start.elapsed().as_secs() < 60,
            "200 calls took {:?}; the budget is not containing the shape",
            start.elapsed()
        );
    }
}