feanor-math 3.5.18

use crate::MAX_PROBABILISTIC_REPETITIONS;
use crate::algorithms::poly_factor::FactorPolyField;
use crate::algorithms::poly_gcd::hensel::*;
use crate::algorithms::poly_gcd::squarefree_part::poly_power_decomposition_local;
use crate::algorithms::poly_gcd::*;
use crate::computation::ComputationController;
use crate::divisibility::*;
use crate::homomorphism::*;
use crate::integer::*;
use crate::iters::{clone_slice, powerset};
use crate::reduce_lift::poly_factor_gcd::*;
use crate::ring::*;
use crate::rings::poly::dense_poly::*;
use crate::rings::poly::*;
use crate::seq::VectorView;

fn combine_local_factors_local<'ring, 'data, 'local, R, P1, P2>(
    reduction: &'local ReductionContext<'ring, 'data, R>,
    poly_ring: P1,
    poly: &El<P1>,
    local_poly_ring: P2,
    local_e: usize,
    local_factors: Vec<El<P2>>,
) -> Vec<El<P1>>
where
    R: ?Sized + PolyGCDLocallyDomain,
    P1: RingStore + Copy,
    P1::Type: PolyRing + DivisibilityRing,
    <P1::Type as RingExtension>::BaseRing: RingStore<Type = R>,
    P2: RingStore + Copy,
    P2::Type: PolyRing<BaseRing = &'local R::LocalRing<'ring>>,
    R::LocalRing<'ring>: 'local,
{
    debug_assert!(poly_ring.base_ring().is_one(poly_ring.lc(poly).unwrap()));
    debug_assert!(local_factors.iter().all(|local_factor| {
        local_poly_ring
            .base_ring()
            .is_one(local_poly_ring.lc(local_factor).unwrap())
    }));

    let ring = poly_ring.base_ring().get_ring();
    let reconstruct_poly = |factor| {
        let mut result = poly_ring.from_terms(local_poly_ring.terms(&factor).map(|(c, i)| {
            (
                ring.reconstruct_ring_el(
                    reduction.ideal(),
                    std::slice::from_ref(*local_poly_ring.base_ring()).as_fn(),
                    local_e,
                    std::slice::from_ref(c).as_fn(),
                ),
                i,
            )
        }));
        _ = poly_ring.balance_poly(&mut result);
        return result;
    };

    let mut ungrouped_factors = (0..local_factors.len()).collect::<Vec<_>>();
    let mut current = poly_ring.clone_el(poly);
    let mut result = Vec::new();
    while !ungrouped_factors.is_empty() {
        // Here we use the naive approach to group the factors such that the product of each group
        // is integral - just try all combinations. It might be worth using LLL for this instead;
        // note that powerset yields smaller subsets first
        let (factor, new, factor_group) = powerset(ungrouped_factors.iter().copied(), |indices| {
            if indices.is_empty() {
                return None;
            }
            let factor = local_poly_ring.prod(
                indices
                    .iter()
                    .copied()
                    .map(|i| local_poly_ring.clone_el(&local_factors[i])),
            );
            let lifted_factor = reconstruct_poly(factor);
            if let Some(quo) = poly_ring.checked_div(&current, &lifted_factor) {
                return Some((lifted_factor, quo, clone_slice(indices)));
            } else {
                return None;
            }
        })
        .flatten()
        .next()
        .unwrap();
        current = new;
        result.push(factor);
        ungrouped_factors.retain(|j| !factor_group.contains(j));
    }
    return result;
}

#[stability::unstable(feature = "enable")]
pub enum FactorAndLiftModpeResult<P>
where
    P: RingStore,
    P::Type: PolyRing,
{
    PartialFactorization(Vec<El<P>>),
    Irreducible,
    Unknown,
    NotSquarefreeModpe,
}

/// Factors the given monic polynomial modulo `p^e` and searches for the smallest groups of factors
/// whose product lifts and gives a factor of `f` globally. If all factors of `f` are shortest lifts
/// of polynomials modulo `p^e`, this means that the result is the factorization of `f`.
///
/// For smaller `e`, this might still compute a nontrivial, partial factorization. However, clearly,
/// it might return non-irreducible factors of `f`.
#[stability::unstable(feature = "enable")]
pub fn factor_and_lift_mod_pe<'ring, R, P, Controller>(
    poly_ring: P,
    prime: &R::SuitableIdeal<'ring>,
    e: usize,
    poly: &El<P>,
    controller: Controller,
) -> FactorAndLiftModpeResult<P>
where
    R: ?Sized + PolyGCDLocallyDomain,
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <P::Type as RingExtension>::BaseRing: RingStore<Type = R>,
    Controller: ComputationController,
{
    let ring = poly_ring.base_ring().get_ring();
    assert_eq!(
        1,
        ring.maximal_ideal_factor_count(prime),
        "currently only maximal ideals are supported, got {}",
        IdealDisplayWrapper::new(ring, prime)
    );

    let reduction = ReductionContext::new(ring, prime, e);
    let red_map = reduction.intermediate_ring_to_field_reduction(0);

    let iso = reduction.base_ring_to_field_iso(0);
    let F = iso.codomain();
    let FX = DensePolyRing::new(&F, "X");
    let R_to_F = (&iso).compose(reduction.main_ring_to_field_reduction(0));

    let poly_mod_m = FX.lifted_hom(poly_ring, R_to_F).map_ref(poly);
    let mut factors = Vec::new();
    for (f, k) in <_ as FactorPolyField>::factor_poly_with_controller(&FX, &poly_mod_m, controller.clone()).0 {
        if k > 1 {
            return FactorAndLiftModpeResult::NotSquarefreeModpe;
        }
        factors.push(f);
    }
    if factors.len() == 1 {
        return FactorAndLiftModpeResult::Irreducible;
    }

    let SX = DensePolyRing::new(*red_map.domain(), "X");
    let poly_mod_me = SX
        .lifted_hom(poly_ring, reduction.main_ring_to_intermediate_ring_reduction(0))
        .map_ref(poly);
    let factorization_mod_me = hensel_lift_factorization(&red_map, &SX, &FX, &poly_mod_me, &factors[..], controller);
    let combined_factorization = combine_local_factors_local(&reduction, poly_ring, poly, &SX, e, factorization_mod_me);
    if combined_factorization.len() == 1 {
        return FactorAndLiftModpeResult::Unknown;
    } else {
        return FactorAndLiftModpeResult::PartialFactorization(combined_factorization);
    }
}

fn ln_factor_max_coeff<P>(ZZX: P, f: &El<P>) -> f64
where
    P: RingStore,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: IntegerRing,
{
    assert!(!ZZX.is_zero(f));
    let ZZ = ZZX.base_ring();
    let d = ZZX.degree(f).unwrap();

    // we use Theorem 3.5.1 from "A course in computational algebraic number theory", Cohen,
    // or equivalently Ex. 20 from Chapter 4.6.2 in Knuth's Art
    let log2_poly_norm =
        ZZX.terms(f).map(|(c, _)| ZZ.abs_log2_ceil(c).unwrap()).max().unwrap() as f64 + (d as f64).log2();
    return (log2_poly_norm + d as f64) * 2.0f64.ln();
}

fn factor_squarefree_monic_integer_poly_local<'a, P, Controller>(
    ZZX: P,
    f: &El<P>,
    controller: Controller,
) -> Vec<El<P>>
where
    P: 'a + RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: IntegerRing,
    Controller: ComputationController,
{
    let ZZ = ZZX.base_ring();
    assert!(ZZ.is_one(ZZX.lc(f).unwrap()));
    let mut rng = oorandom::Rand64::new(1);
    let bound = ln_factor_max_coeff(ZZX, f);

    for _ in 0..MAX_PROBABILISTIC_REPETITIONS {
        let prime = ZZ.get_ring().random_suitable_ideal(|| rng.rand_u64());
        assert_eq!(1, ZZ.get_ring().maximal_ideal_factor_count(&prime));
        let prime_f64 = BigIntRing::RING.to_float_approx(&ZZ.get_ring().principal_ideal_generator(&prime));
        let e = (bound / prime_f64.ln()).ceil() as usize + 1;
        log_progress!(
            controller,
            "(mod={}^{})",
            IdealDisplayWrapper::new(ZZ.get_ring(), &prime),
            e
        );
        match factor_and_lift_mod_pe(ZZX, &prime, e, f, controller.clone()) {
            FactorAndLiftModpeResult::Irreducible => return vec![ZZX.clone_el(f)],
            FactorAndLiftModpeResult::PartialFactorization(result) => return result,
            // unknown means irreducible, since we chose `e` large enough
            FactorAndLiftModpeResult::Unknown => return vec![ZZX.clone_el(f)],
            FactorAndLiftModpeResult::NotSquarefreeModpe => {}
        }
    }
    unreachable!()
}

/// Factors the given polynomial over the integers.
///
/// Its factors are returned as primitive polynomials, thus their
/// product is `f` only up to multiplication by a nonzero integer.
#[stability::unstable(feature = "enable")]
pub fn poly_factor_integer<P, Controller>(ZZX: P, f: El<P>, controller: Controller) -> Vec<(El<P>, usize)>
where
    P: PolyRingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: IntegerRing,
    Controller: ComputationController,
{
    assert!(!ZZX.is_zero(&f));
    let power_decomposition = poly_power_decomposition_local(ZZX, ZZX.clone_el(&f), controller.clone());

    controller.run_computation(
        format_args!("factor_int_poly(deg={})", ZZX.degree(&f).unwrap()),
        |controller| {
            let mut result = Vec::new();
            let mut current = ZZX.clone_el(&f);
            for (factor, _k) in power_decomposition {
                log_progress!(controller, "(deg={})", ZZX.degree(&factor).unwrap());
                let lc_factor = ZZX.lc(&factor).unwrap();
                let factor_monic = evaluate_aX(ZZX, &factor, lc_factor);
                let factorization = factor_squarefree_monic_integer_poly_local(&ZZX, &factor_monic, controller.clone());
                for irred_factor in factorization.into_iter().map(|fi| unevaluate_aX(ZZX, &fi, lc_factor)) {
                    let irred_factor_lc = ZZX.lc(&irred_factor).unwrap();
                    let mut power = 0;
                    let irred_factor = make_primitive(ZZX, &irred_factor).0;
                    while let Some(quo) = ZZX.checked_div(
                        &ZZX.inclusion().mul_ref_map(
                            &current,
                            &ZZX.base_ring()
                                .pow(ZZX.base_ring().clone_el(irred_factor_lc), ZZX.degree(&f).unwrap()),
                        ),
                        &irred_factor,
                    ) {
                        current = quo;
                        _ = ZZX.balance_poly(&mut current);
                        power += 1;
                    }
                    assert!(power >= 1);
                    result.push((irred_factor, power));
                }
            }
            debug_assert_eq!(
                ZZX.degree(&f).unwrap(),
                result.iter().map(|(fi, i)| *i * ZZX.degree(fi).unwrap()).sum::<usize>()
            );
            return result;
        },
    )
}

#[cfg(test)]
use crate::algorithms::poly_gcd::make_primitive;
#[cfg(test)]
use crate::computation::TEST_LOG_PROGRESS;
#[cfg(test)]
use crate::primitive_int::*;

#[test]
fn test_factor_int_poly() {
    let ZZX = DensePolyRing::new(StaticRing::<i64>::RING, "X");
    let [f, g] = ZZX.with_wrapped_indeterminate(|X| [X.pow_ref(2) + 1, X + 1]);
    let input = ZZX.mul_ref(&f, &g);
    let actual = poly_factor_integer(&ZZX, input, TEST_LOG_PROGRESS);
    assert_eq!(2, actual.len());
    for (factor, e) in &actual {
        assert_eq!(1, *e);
        assert!(
            ZZX.eq_el(&f, factor) || ZZX.eq_el(&g, factor),
            "Got unexpected factor {}",
            ZZX.format(&factor)
        );
    }

    let [f, g] = ZZX.with_wrapped_indeterminate(|X| [5 * X.pow_ref(2) + 1, 3 * X.pow_ref(2) + 2]);
    let input = ZZX.mul_ref(&f, &g);
    let actual = poly_factor_integer(&ZZX, input, TEST_LOG_PROGRESS);
    assert_eq!(2, actual.len());
    for (factor, e) in &actual {
        assert_eq!(1, *e);
        assert!(
            ZZX.eq_el(&f, factor) || ZZX.eq_el(&g, factor),
            "Got unexpected factor {}",
            ZZX.format(&factor)
        );
    }

    let [f] = ZZX.with_wrapped_indeterminate(|X| [5 * X.pow_ref(2) + 1]);
    let input = ZZX.mul_ref(&f, &f);
    let actual = poly_factor_integer(&ZZX, input, TEST_LOG_PROGRESS);
    assert_eq!(1, actual.len());
    assert_eq!(2, actual[0].1);
    assert_el_eq!(&ZZX, &f, &actual[0].0);
}