feanor-math 3.5.18

use std::fmt::Debug;

use super::{INCREASE_EXPONENT_PER_ATTEMPT_CONSTANT, evaluate_aX, unevaluate_aX};
use crate::MAX_PROBABILISTIC_REPETITIONS;
use crate::algorithms::poly_gcd::hensel::*;
use crate::algorithms::poly_gcd::*;
use crate::rings::poly::dense_poly::DensePolyRing;
use crate::seq::*;

/// For the power-decomposition `f = f1^e1 ... fr^er`, stores a tuple (ei, deg(fi))
#[derive(PartialEq, Eq, Debug)]
struct Signature {
    perfect_power: usize,
    degree: usize,
}

/// Lifts the power decomposition modulo `I^e` to a potential power decomposition in the ring `R`.
/// If this indeed gives the correct power decomposition, it is returned, otherwise `None` is
/// returned.
///
/// We use the notation
///  - `R` is the main ring
///  - `F` is `R/m` where `m` is a maximal ideal containing the currently considered ideal `I`
///  - `S` is `R/m^e`
fn power_decomposition_from_local_power_decomposition<'ring, 'data, 'local, R, P>(
    reduction: &'local ReductionContext<'ring, 'data, R>,
    RX: P,
    poly: &El<P>,
    signature: &[Signature],
    SXs: &[DensePolyRing<&'local R::LocalRing<'ring>>],
    local_power_decompositions: &[Vec<El<DensePolyRing<&'local R::LocalRing<'ring>>>>],
) -> Option<Vec<(El<P>, usize)>>
where
    R: ?Sized + PolyGCDLocallyDomain,
    P: RingStore + Copy,
    P::Type: PolyRing,
    <P::Type as RingExtension>::BaseRing: RingStore<Type = R>,
{
    assert_eq!(reduction.len(), local_power_decompositions.len());
    assert_eq!(reduction.len(), SXs.len());

    let mut result = Vec::new();
    for (k, sig) in signature.iter().enumerate() {
        let power_factor = RX.from_terms((0..=(sig.degree * sig.perfect_power)).map(|i| {
            (
                reduction.reconstruct_ring_el(
                    (0..reduction.len()).map_fn(|j| SXs[j].coefficient_at(&local_power_decompositions[j][k], i)),
                ),
                i,
            )
        }));
        if let Some(mut root_of_factor) = poly_root(RX, &power_factor, sig.perfect_power) {
            _ = RX.balance_poly(&mut root_of_factor);
            let lc_inv = RX.base_ring().invert(RX.lc(&root_of_factor).unwrap()).unwrap();
            RX.inclusion().mul_assign_map(&mut root_of_factor, lc_inv);
            result.push((root_of_factor, sig.perfect_power));
        } else {
            return None;
        }
    }
    // at first, I thought this could not happen, but actually it can. If we do a faulty lift, the
    // polynomials might after all still turn out to be perfect powers; the alternative to this
    // check here would be to check previously if all "factors" really divide f; I believe this
    // is faster
    if !RX.eq_el(poly, &RX.prod(result.iter().map(|(f, k)| RX.pow(RX.clone_el(f), *k)))) {
        return None;
    }
    return Some(result);
}

/// Computes the power decomposition modulo `m` and lifts it to `m^e`, for the given maximal ideal
/// `m` as specified by `S_to_F`.
///
/// We use the notation
///  - `R` is the main ring
///  - `F` is `R/m` where `m` is a maximal ideal containing the currently considered ideal `I`
///  - `S` is `R/m^e`
fn compute_local_power_decomposition<'ring, 'data, 'local, R, P1, P2, Controller>(
    RX: P1,
    f: &El<P1>,
    S_to_F: &PolyGCDLocallyIntermediateReductionMap<'ring, 'data, 'local, R>,
    SX: P2,
    controller: Controller,
) -> Option<(Vec<Signature>, Vec<El<P2>>)>
where
    R: ?Sized + PolyGCDLocallyDomain,
    P1: RingStore + Copy,
    P1::Type: PolyRing,
    <P1::Type as RingExtension>::BaseRing: RingStore<Type = R>,
    P2: RingStore + Copy,
    P2::Type: PolyRing<BaseRing = &'local R::LocalRing<'ring>>,
    R::LocalRing<'ring>: 'local,
    Controller: ComputationController,
{
    assert!(SX.base_ring().get_ring() == S_to_F.domain().get_ring());
    let R = RX.base_ring().get_ring();
    let F = R.local_field_at(S_to_F.ideal(), S_to_F.max_ideal_idx());
    let FX = DensePolyRing::new(&F, "X");
    let iso = PolyGCDLocallyBaseRingToFieldIso::new(
        R,
        S_to_F.ideal(),
        S_to_F.codomain().get_ring(),
        F.get_ring(),
        S_to_F.max_ideal_idx(),
    );

    let f_mod_m = FX.from_terms(RX.terms(f).map(|(c, i)| {
        (
            iso.map(R.reduce_ring_el(
                S_to_F.ideal(),
                (S_to_F.codomain().get_ring(), 1),
                S_to_F.max_ideal_idx(),
                R.clone_el(c),
            )),
            i,
        )
    }));
    let f_mod_me = SX.from_terms(RX.terms(f).map(|(c, i)| {
        (
            R.reduce_ring_el(
                S_to_F.ideal(),
                (S_to_F.domain().get_ring(), S_to_F.from_e()),
                S_to_F.max_ideal_idx(),
                R.clone_el(c),
            ),
            i,
        )
    }));

    let mut power_decomposition_mod_m = Vec::new();
    let mut signature = Vec::new();
    for (f, k) in <_ as PolyTFracGCDRing>::power_decomposition(&FX, &f_mod_m).into_iter() {
        signature.push(Signature {
            perfect_power: k,
            degree: FX.degree(&f).unwrap(),
        });
        power_decomposition_mod_m.push(FX.pow(f, k));
    }

    let power_decomposition_mod_me =
        hensel_lift_factorization(S_to_F, SX, &FX, &f_mod_me, &power_decomposition_mod_m[..], controller);

    return Some((signature, power_decomposition_mod_me));
}

/// For a monic polynomial `f in R[X]`, computes squarefree polynomials `fi` such that `a f = f1
/// f2^2 f3^3 ...` for some nonzero element `a in R \ {0}`. These polynomials are returned as tuples
/// `(fi, i)` with `fi != 0`.
///
/// The results can all be assumed to be "balanced", according to the contract of
/// [`DivisibilityRing::balance_factor()`] of the underlying ring.
#[stability::unstable(feature = "enable")]
pub fn poly_power_decomposition_monic_local<P, Controller>(
    poly_ring: P,
    poly: &El<P>,
    controller: Controller,
) -> Vec<(El<P>, usize)>
where
    P: RingStore + Copy,
    P::Type: PolyRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain,
    Controller: ComputationController,
{
    assert!(poly_ring.base_ring().is_one(poly_ring.lc(poly).unwrap()));

    controller.run_computation(
        format_args!("power_decomp_local(deg={})", poly_ring.degree(poly).unwrap()),
        |controller| {
            let ring = poly_ring.base_ring().get_ring();
            let mut rng = oorandom::Rand64::new(1);

            'try_random_ideal: for current_attempt in 0..MAX_PROBABILISTIC_REPETITIONS {
                let ideal = ring.random_suitable_ideal(|| rng.rand_u64());
                let heuristic_e = ring.heuristic_exponent(
                    &ideal,
                    poly_ring.degree(poly).unwrap(),
                    poly_ring.terms(poly).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<Vec<_>> = None;
                let mut poly_rings_mod_me = Vec::new();
                let mut power_decompositions_mod_me = Vec::new();

                for idx in 0..reduction.len() {
                    let SX = DensePolyRing::new(*reduction.intermediate_ring_to_field_reduction(idx).domain(), "X");
                    match compute_local_power_decomposition(
                        poly_ring,
                        poly,
                        &reduction.intermediate_ring_to_field_reduction(idx),
                        &SX,
                        controller.clone(),
                    ) {
                        None => {
                            unreachable!("`compute_local_power_decomposition()` currently cannot fail");
                        }
                        Some((new_signature, local_power_decomposition)) => {
                            if new_signature
                                == [Signature {
                                    degree: poly_ring.degree(poly).unwrap(),
                                    perfect_power: 1,
                                }]
                            {
                                return vec![(poly_ring.clone_el(poly), 1)];
                            } else if signature.is_some() && signature.as_ref().unwrap()[..] != new_signature[..] {
                                log_progress!(controller, "(signature_mismatch)");
                                continue 'try_random_ideal;
                            } else {
                                signature = Some(new_signature);
                                power_decompositions_mod_me.push(local_power_decomposition);
                                poly_rings_mod_me.push(SX);
                            }
                        }
                    }
                }

                if let Some(result) = power_decomposition_from_local_power_decomposition(
                    &reduction,
                    poly_ring,
                    poly,
                    &signature.as_ref().unwrap()[..],
                    &poly_rings_mod_me[..],
                    &power_decompositions_mod_me[..],
                ) {
                    return result;
                } else {
                    log_progress!(controller, "(invalid_lift)");
                }
            }
            unreachable!()
        },
    )
}

/// For a polynomial `f in R[X]`, computes squarefree polynomials `fi` such that `a f = f1 f2^2 f3^3
/// ...` for some nonzero ring element `a in R \ {0}`. These polynomials are returned as tuples
/// `(fi, i)` with `fi != 0`.
///
/// The results can all be assumed to be "balanced", according to the contract of
/// [`DivisibilityRing::balance_factor()`] of the underlying ring.
#[stability::unstable(feature = "enable")]
pub fn poly_power_decomposition_local<P, Controller>(
    poly_ring: P,
    mut f: El<P>,
    controller: Controller,
) -> Vec<(El<P>, usize)>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain + DivisibilityRing,
    Controller: ComputationController,
{
    assert!(!poly_ring.is_zero(&f));
    _ = poly_ring.balance_poly(&mut f);
    let lcf = poly_ring.lc(&f).unwrap();
    let f_monic = evaluate_aX(poly_ring, &f, lcf);
    let power_decomposition = poly_power_decomposition_monic_local(poly_ring, &f_monic, controller);
    let result = power_decomposition
        .into_iter()
        .map(|(fi, i)| {
            let mut result = unevaluate_aX(poly_ring, &fi, lcf);
            _ = poly_ring.balance_poly(&mut result);
            return (result, i);
        })
        .collect::<Vec<_>>();
    debug_assert!(
        poly_ring
            .checked_div(
                &poly_ring.prod(result.iter().map(|(fi, i)| poly_ring.pow(poly_ring.clone_el(fi), *i))),
                &f
            )
            .is_some()
    );
    debug_assert_eq!(
        poly_ring.degree(&f).unwrap(),
        result
            .iter()
            .map(|(fi, i)| *i * poly_ring.degree(fi).unwrap())
            .sum::<usize>()
    );
    return result;
}

/// Computes the square-free part of a polynomial `f in R[X]`, up to multiplication by `R \ {0}`.
///
/// More concretely, returns the largest square-free polynomial `g` for which there is `a in R \
/// {0}` with `g | af`.
///
/// 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_squarefree_part_local<P, Controller>(poly_ring: P, f: El<P>, controller: Controller) -> El<P>
where
    P: RingStore + Copy,
    P::Type: PolyRing + DivisibilityRing,
    <<P::Type as RingExtension>::BaseRing as RingStore>::Type: PolyGCDLocallyDomain + DivisibilityRing,
    Controller: ComputationController,
{
    assert!(!poly_ring.is_zero(&f));
    let mut result = poly_ring.prod(
        poly_power_decomposition_local(poly_ring, f, controller)
            .into_iter()
            .map(|(fi, _i)| fi),
    );
    _ = poly_ring.balance_poly(&mut result);
    return result;
}

#[cfg(test)]
use super::make_primitive;
#[cfg(test)]
use crate::RANDOM_TEST_INSTANCE_COUNT;

#[test]
fn test_squarefree_part_local() {
    let ring = BigIntRing::RING;
    let poly_ring = dense_poly::DensePolyRing::new(ring, "X");
    let [f1, f2, f3, f4] = poly_ring.with_wrapped_indeterminate(|X| {
        [
            X - 1,
            X + 1,
            X.pow_ref(3) + X + 100,
            X.pow_ref(4) + X.pow_ref(3) + X.pow_ref(2) + X + 1,
        ]
    });
    let multiply_out = |list: &[(El<DensePolyRing<_>>, usize)]| {
        poly_ring.prod(list.iter().map(|(g, k)| poly_ring.pow(poly_ring.clone_el(g), *k)))
    };
    let assert_eq = |expected: &[(El<DensePolyRing<_>>, usize)], actual: &[(El<DensePolyRing<_>>, usize)]| {
        assert!(expected.is_sorted_by_key(|(_, k)| *k));
        assert!(actual.is_sorted_by_key(|(_, k)| *k));
        assert_eq!(expected.len(), actual.len());
        for ((f_expected, k_expected), (f_actual, k_actual)) in expected.iter().zip(actual.iter()) {
            assert_el_eq!(&poly_ring, f_expected, f_actual);
            assert_eq!(k_expected, k_actual);
        }
    };

    let expected = [(poly_ring.clone_el(&f1), 1)];
    let actual = poly_power_decomposition_monic_local(&poly_ring, &multiply_out(&expected), TEST_LOG_PROGRESS);
    assert_eq(&expected, &actual);

    let expected = [(poly_ring.mul_ref(&f3, &f4), 3)];
    let actual = poly_power_decomposition_monic_local(&poly_ring, &multiply_out(&expected), TEST_LOG_PROGRESS);
    assert_eq(&expected, &actual);

    let expected = [(poly_ring.clone_el(&f2), 2), (poly_ring.mul_ref(&f3, &f4), 3)];
    let actual = poly_power_decomposition_monic_local(&poly_ring, &multiply_out(&expected), TEST_LOG_PROGRESS);
    assert_eq(&expected, &actual);

    let expected = [
        (poly_ring.mul_ref(&f1, &f2), 1),
        (poly_ring.clone_el(&f4), 2),
        (poly_ring.clone_el(&f3), 3),
    ];
    let actual = poly_power_decomposition_monic_local(&poly_ring, &multiply_out(&expected), TEST_LOG_PROGRESS);
    assert_eq(&expected, &actual);

    let expected = [
        (poly_ring.mul_ref(&f1, &f2), 2),
        (poly_ring.clone_el(&f4), 4),
        (poly_ring.clone_el(&f3), 6),
    ];
    let actual = poly_power_decomposition_monic_local(&poly_ring, &multiply_out(&expected), TEST_LOG_PROGRESS);
    assert_eq(&expected, &actual);
}

#[test]
#[ignore]
fn random_test_poly_power_decomposition_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(1000);
    for _ in 0..RANDOM_TEST_INSTANCE_COUNT {
        let f = poly_ring.from_terms((0..=7).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        let g = poly_ring.from_terms((0..=4).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        let h = poly_ring.from_terms((0..=2).map(|i| (ring.get_uniformly_random(&bound, || rng.rand_u64()), i)));
        let poly = make_primitive(
            &poly_ring,
            &poly_ring.prod(
                [&f, &g, &g, &h, &h, &h, &h, &h]
                    .into_iter()
                    .map(|poly| poly_ring.clone_el(poly)),
            ),
        )
        .0;

        let mut power_decomp = poly_power_decomposition_local(&poly_ring, poly_ring.clone_el(&poly), TEST_LOG_PROGRESS);
        for (f, _k) in &mut power_decomp {
            *f = make_primitive(&poly_ring, &f).0;
        }

        assert_el_eq!(
            &poly_ring,
            &poly,
            poly_ring.prod(
                power_decomp
                    .iter()
                    .map(|(poly, k)| poly_ring.pow(poly_ring.clone_el(poly), *k))
            )
        );
        assert!(
            poly_ring.divides(
                &poly_ring.prod(
                    power_decomp
                        .iter()
                        .filter(|(_, k)| k % 5 == 0)
                        .map(|(poly, k)| poly_ring.pow(poly_ring.clone_el(poly), k / 5))
                ),
                &make_primitive(&poly_ring, &h).0
            )
        );
        assert!(
            poly_ring.divides(
                &poly_ring.prod(
                    power_decomp
                        .iter()
                        .filter(|(_, k)| k % 2 == 0)
                        .map(|(poly, k)| poly_ring.pow(poly_ring.clone_el(poly), k / 2))
                ),
                &make_primitive(&poly_ring, &g).0
            )
        );
    }
}