feanor-math 3.5.18

use dense_poly::DensePolyRing;
use squarefree_part::poly_power_decomposition_monic_local;

use super::{evaluate_aX, unevaluate_aX};
use crate::MAX_PROBABILISTIC_REPETITIONS;
use crate::algorithms::poly_gcd::hensel::*;
use crate::algorithms::poly_gcd::*;
use crate::seq::*;

const HOPE_FOR_SQUAREFREE_TRIES: usize = 3;

/// Describes the relationship of `f, g, gcd(f, g)` modulo a single maximal ideal
#[derive(PartialEq, Eq)]
struct Signature {
    /// the degree of `gcd(f, g) mod m`
    gcd_deg: usize,
    /// whether `f/d` is coprime to `d`, where `d = gcd(f, g) mod m`
    coprime_to_f_over_d: bool,
}

/// Tries to compute the gcd of monic polynomials `f, g in R[X]` over `Frac(R)` locally, without
/// ever explicitly working in `Frac(R)`. This function will fail in two cases
///  - both `d, f/d` and `d, g/d` are not coprime
///  - the gcd of `f, g` cannot be reconstructed from its reduction modulo `p^e`, where `e` is (as
///    usual) the result of the "heuristic exponent", scaled exponentially by the current attempt
/// If neither is the case, this function is likely to succeed
///
/// More precisely, computes some `d in R[X]` of maximal degree with the property that there exists
/// `a in R \ {0}` such that `d | af, ag`.
fn poly_gcd_monic_coprime_local<P, F, Controller>(
    poly_ring: P,
    f: &El<P>,
    g: &El<P>,
    rng: F,
    current_attempt: usize,
    controller: Controller,
) -> Option<El<P>>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain,
    F: FnMut() -> u64,
    Controller: ComputationController,
{
    assert!(poly_ring.base_ring().is_one(poly_ring.lc(f).unwrap()));
    assert!(poly_ring.base_ring().is_one(poly_ring.lc(g).unwrap()));

    let ring = poly_ring.base_ring().get_ring();

    let ideal = ring.random_suitable_ideal(rng);
    let heuristic_e = ring.heuristic_exponent(&ideal, poly_ring.degree(f).unwrap(), poly_ring.terms(f).map(|(c, _)| c));
    assert!(heuristic_e >= 1);
    let e = (heuristic_e as f64 * INCREASE_EXPONENT_PER_ATTEMPT_CONSTANT.powi(current_attempt.try_into().unwrap()))
        .floor() as usize;
    let reduction = ReductionContext::new(ring, &ideal, e);

    log_progress!(
        controller,
        "(mod={}^{})(parts={})",
        IdealDisplayWrapper::new(ring, &ideal),
        e,
        reduction.len()
    );

    let mut signature: Option<Signature> = None;
    let mut poly_rings_mod_me = Vec::new();
    let mut gcds_mod_me = Vec::new();

    for idx in 0..reduction.len() {
        let S_to_F = reduction.intermediate_ring_to_field_reduction(idx);
        let F_iso = reduction.base_ring_to_field_iso(idx);
        let F = F_iso.codomain();
        let R_to_F = (&F_iso).compose(reduction.main_ring_to_field_reduction(idx));
        let FX = DensePolyRing::new(&F, "X");
        let RX_to_FX = FX.lifted_hom(poly_ring, &R_to_F);

        let d = controller.clone().run_computation(format_args!("local_gcd."), |_| {
            FX.normalize(FX.ideal_gen(&RX_to_FX.map_ref(f), &RX_to_FX.map_ref(g)))
        });

        let f_over_d = FX.checked_div(&RX_to_FX.map_ref(f), &d).unwrap();
        let g_over_d = FX.checked_div(&RX_to_FX.map_ref(g), &d).unwrap();
        let deg_d = FX.degree(&d).unwrap();
        let (poly, factor1, factor2, new_signature) = if FX.is_unit(&FX.ideal_gen(&d, &f_over_d)) {
            (
                f,
                d,
                f_over_d,
                Signature {
                    gcd_deg: deg_d,
                    coprime_to_f_over_d: true,
                },
            )
        } else if FX.is_unit(&FX.ideal_gen(&d, &g_over_d)) {
            (
                g,
                d,
                g_over_d,
                Signature {
                    gcd_deg: deg_d,
                    coprime_to_f_over_d: false,
                },
            )
        } else {
            log_progress!(controller, "(not_coprime)");
            return None;
        };
        if signature.is_some() && signature.as_ref().unwrap() != &new_signature {
            log_progress!(controller, "(signature_mismatch)");
            return None;
        }
        signature = Some(new_signature);
        let SX = DensePolyRing::new(*S_to_F.domain(), "X");
        let RX_to_SX = SX.lifted_hom(poly_ring, reduction.main_ring_to_intermediate_ring_reduction(idx));

        let factors = [factor1, factor2];
        let [d, _] = hensel_lift_factorization(
            &S_to_F,
            &SX,
            &FX,
            &RX_to_SX.map_ref(poly),
            &factors[..],
            controller.clone(),
        )
        .try_into()
        .ok()
        .unwrap();

        poly_rings_mod_me.push(SX);
        gcds_mod_me.push(d);
    }

    let signature = signature.unwrap();
    let mut result = controller.clone().run_computation(format_args!("reconstruct."), |_| {
        poly_ring.from_terms((0..=signature.gcd_deg).map(|i| {
            (
                reduction.reconstruct_ring_el(
                    (0..reduction.len()).map_fn(|j| poly_rings_mod_me[j].coefficient_at(&gcds_mod_me[j], i)),
                ),
                i,
            )
        }))
    });

    let divides_f_and_g = controller
        .clone()
        .run_computation(format_args!("check_division."), |_| {
            poly_ring.divides(f, &result) && poly_ring.divides(g, &result)
        });
    if !divides_f_and_g {
        log_progress!(controller, "(no_divisor)");
        return None;
    } else {
        _ = poly_ring.balance_poly(&mut result);
        return Some(result);
    }
}

/// Tries to compute the gcd of polynomials `f, g in R[X]` over `Frac(R)` locally, without ever
/// explicitly working in `Frac(R)`. This function is likely to succeed if either `d, f/d` or `d,
/// g/d` are coprime (over `Frac(R)`), but will never succeed if neither is the case.
///
/// More precisely, computes some `d in R[X]` of maximal degree with the property that there exists
/// `a in R \ {0}` such that `d | af, ag`.
fn poly_gcd_coprime_local<P, F, Controller>(
    poly_ring: P,
    mut f: El<P>,
    mut g: El<P>,
    rng: F,
    attempt: usize,
    controller: Controller,
) -> Option<El<P>>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain,
    F: FnMut() -> u64,
    Controller: ComputationController,
{
    if poly_ring.is_zero(&f) {
        return Some(poly_ring.clone_el(&g));
    } else if poly_ring.is_zero(&g) {
        return Some(poly_ring.clone_el(&f));
    }
    _ = poly_ring.balance_poly(&mut f);
    _ = poly_ring.balance_poly(&mut g);
    let lcf = poly_ring.base_ring().clone_el(poly_ring.lc(&f).unwrap());
    let lcg = poly_ring.base_ring().clone_el(poly_ring.lc(&g).unwrap());
    let ring = poly_ring.base_ring();
    let a = ring.mul_ref(&lcf, &lcg);
    let f_monic = evaluate_aX(poly_ring, &poly_ring.inclusion().mul_map(f, lcg), &a);
    let g_monic = evaluate_aX(poly_ring, &poly_ring.inclusion().mul_map(g, lcf), &a);

    let d_monic = poly_gcd_monic_coprime_local(poly_ring, &f_monic, &g_monic, rng, attempt, controller)?;

    let mut result = unevaluate_aX(poly_ring, &d_monic, &a);
    _ = poly_ring.balance_poly(&mut result);
    return Some(result);
}

/// Computes the gcd of monic polynomials `f, g in R[X]` over `Frac(R)` locally, without ever
/// explicitly working in `Frac(R)`.
///
/// More precisely, computes some `d in R[X]` of maximal degree with the property that there exists
/// `a in R \ {0}` such that `d | af, ag`.
///
/// The result can be assumed to be "balanced", according to the contract of
/// [`DivisibilityRing::balance_factor()`] of the underlying ring.
#[stability::unstable(feature = "enable")]
pub fn poly_gcd_monic_local<'a, P, Controller>(
    poly_ring: P,
    mut f: &'a El<P>,
    mut g: &'a El<P>,
    controller: Controller,
) -> El<P>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain,
    Controller: ComputationController,
{
    assert!(poly_ring.base_ring().is_one(poly_ring.lc(f).unwrap()));
    assert!(poly_ring.base_ring().is_one(poly_ring.lc(g).unwrap()));

    controller.run_computation(
        format_args!(
            "gcd_local(deg_l={}, deg_r={})",
            poly_ring.degree(f).unwrap(),
            poly_ring.degree(g).unwrap()
        ),
        |controller| {
            let mut rng = oorandom::Rand64::new(1);
            for attempt in 0..HOPE_FOR_SQUAREFREE_TRIES {
                if let Some(result) =
                    poly_gcd_monic_coprime_local(poly_ring, f, g, || rng.rand_u64(), attempt, controller.clone())
                {
                    return result;
                }
            }
            if poly_ring.degree(g).unwrap_or(0) <= poly_ring.degree(f).unwrap_or(0) {
                std::mem::swap(&mut f, &mut g);
            }
            let f_power_decomposition = poly_power_decomposition_monic_local(poly_ring, f, controller.clone());
            let mut g = poly_ring.clone_el(g);
            let mut d = poly_ring.one();

            'extract_part_i: for i in 1.. {
                let squarefree_part_i = poly_ring.prod(
                    f_power_decomposition
                        .iter()
                        .filter(|(_, j)| *j >= i)
                        .map(|(fj, _)| poly_ring.clone_el(fj)),
                );
                if poly_ring.is_one(&squarefree_part_i) {
                    return d;
                }
                for attempt in 0..MAX_PROBABILISTIC_REPETITIONS {
                    if let Some(di) = poly_gcd_coprime_local(
                        poly_ring,
                        poly_ring.clone_el(&squarefree_part_i),
                        poly_ring.clone_el(&g),
                        || rng.rand_u64(),
                        attempt,
                        controller.clone(),
                    ) {
                        g = poly_ring.checked_div(&g, &di).unwrap();
                        poly_ring.mul_assign(&mut d, di);
                        continue 'extract_part_i;
                    }
                }
                unreachable!()
            }
            unreachable!()
        },
    )
}

/// Computes the gcd of polynomials `f, g in R[X]` over `Frac(R)` locally, without ever
/// explicitly working in `Frac(R)`.
///
/// More precisely, computes some `d in R[X]` of maximal degree with the property that there exists
/// `a in R \ {0}` such that `d | af, ag`.
///
/// The result can be assumed to be "balanced", according to the contract of
/// [`DivisibilityRing::balance_factor()`] of the underlying ring.
#[stability::unstable(feature = "enable")]
pub fn poly_gcd_local<P, Controller>(poly_ring: P, mut f: El<P>, mut g: El<P>, controller: Controller) -> El<P>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain,
    Controller: ComputationController,
{
    if poly_ring.is_zero(&f) {
        return poly_ring.clone_el(&g);
    } else if poly_ring.is_zero(&g) {
        return poly_ring.clone_el(&f);
    }
    _ = poly_ring.balance_poly(&mut f);
    _ = poly_ring.balance_poly(&mut g);
    let lcf = poly_ring.base_ring().clone_el(poly_ring.lc(&f).unwrap());
    let lcg = poly_ring.base_ring().clone_el(poly_ring.lc(&g).unwrap());
    let ring = poly_ring.base_ring();
    let a = ring.mul_ref(&lcf, &lcg);
    let f_monic = evaluate_aX(poly_ring, &poly_ring.inclusion().mul_map(f, lcg), &a);
    let g_monic = evaluate_aX(poly_ring, &poly_ring.inclusion().mul_map(g, lcf), &a);

    let d_monic = poly_gcd_monic_local(poly_ring, &f_monic, &g_monic, controller);

    let mut result = unevaluate_aX(poly_ring, &d_monic, &a);
    _ = poly_ring.balance_poly(&mut result);
    if let Some(lc_inv) = poly_ring.base_ring().invert(poly_ring.lc(&result).unwrap()) {
        poly_ring.inclusion().mul_assign_map(&mut result, lc_inv);
    }
    return result;
}

#[cfg(test)]
use crate::RANDOM_TEST_INSTANCE_COUNT;
#[cfg(test)]
use crate::algorithms::poly_gcd::make_primitive;
#[cfg(test)]
use crate::integer::*;

#[test]
fn test_poly_gcd_local() {
    let ring = BigIntRing::RING;
    let poly_ring = DensePolyRing::new(ring, "X");
    let irred_polys = poly_ring.with_wrapped_indeterminate(|X| {
        [
            X - 1,
            X + 1,
            X.pow_ref(2) + X + 1,
            X.pow_ref(3) + X + 100,
            X.pow_ref(4) + X.pow_ref(3) + X.pow_ref(2) + X + 1,
        ]
    });
    let poly = |powers: [usize; 5], scale: i32| {
        poly_ring.int_hom().mul_map(
            poly_ring.prod(
                powers
                    .iter()
                    .zip(irred_polys.iter())
                    .map(|(e, f)| poly_ring.pow(poly_ring.clone_el(f), *e)),
            ),
            scale,
        )
    };

    assert_el_eq!(
        &poly_ring,
        poly([1, 0, 0, 0, 0], 1),
        poly_gcd_local(
            &poly_ring,
            poly([1, 0, 1, 0, 0], 1),
            poly([1, 0, 0, 1, 0], 2),
            TEST_LOG_PROGRESS
        )
    );

    assert_el_eq!(
        &poly_ring,
        poly([1, 0, 1, 0, 1], 1),
        poly_gcd_local(
            &poly_ring,
            poly([1, 1, 1, 0, 1], 20),
            poly([1, 0, 1, 1, 1], 12),
            TEST_LOG_PROGRESS
        )
    );

    assert_el_eq!(
        &poly_ring,
        poly([1, 0, 2, 0, 1], 1),
        poly_gcd_local(
            &poly_ring,
            poly([1, 1, 3, 0, 1], 20),
            poly([3, 0, 2, 0, 3], 12),
            TEST_LOG_PROGRESS
        )
    );

    assert_el_eq!(
        &poly_ring,
        poly([1, 0, 0, 5, 0], 1),
        poly_gcd_local(
            &poly_ring,
            poly([2, 1, 3, 5, 1], 20),
            poly([1, 0, 0, 7, 0], 12),
            TEST_LOG_PROGRESS
        )
    );
}

#[test]
fn random_test_poly_gcd_local() {
    let ring = BigIntRing::RING;
    let poly_ring = dense_poly::DensePolyRing::new(ring, "X");
    let mut rng = oorandom::Rand64::new(1);
    let bound = ring.int_hom().map(10000);
    for _ in 0..RANDOM_TEST_INSTANCE_COUNT {
        let f = poly_ring.from_terms((0..=20).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        let g = poly_ring.from_terms((0..=20).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        let h = poly_ring.from_terms((0..=10).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        // println!("Testing gcd on ({}) * ({}) and ({}) * ({})", poly_ring.format(&f),
        // poly_ring.format(&h), poly_ring.format(&g), poly_ring.format(&h));
        let lhs = poly_ring.mul_ref(&f, &h);
        let rhs = poly_ring.mul_ref(&g, &h);
        let gcd = make_primitive(
            &poly_ring,
            &poly_gcd_local(
                &poly_ring,
                poly_ring.clone_el(&lhs),
                poly_ring.clone_el(&rhs),
                TEST_LOG_PROGRESS,
            ),
        )
        .0;
        // println!("Result {}", poly_ring.format(&gcd));

        assert!(poly_ring.divides(&lhs, &gcd));
        assert!(poly_ring.divides(&rhs, &gcd));
        assert!(poly_ring.divides(&gcd, &make_primitive(&poly_ring, &h).0));
    }
}