fheanor 0.11.9

use std::fmt::Debug;

use feanor_math::algorithms::discrete_log::Subgroup;
use feanor_math::algorithms::miller_rabin::is_prime;
use feanor_math::divisibility::DivisibilityRing;
use feanor_math::integer::{BigIntRing, IntegerRing, IntegerRingStore};
use feanor_math::ordered::OrderedRingStore;
use feanor_math::primitive_int::StaticRing;
use feanor_math::{delegate, ring::*};
use feanor_math::rings::extension::FreeAlgebra;
use feanor_math::rings::finite::FiniteRing;
use feanor_math::rings::poly::PolyRing;
use feanor_math::rings::zn::zn_64;
use feanor_math::integer::int_cast;
use feanor_math::seq::*;

use crate::ZZi64;
use crate::number_ring::galois::*;

pub mod galois;
///
/// Code for fast polynomial division by a cyclotomic polynomial.
/// 
pub mod poly_remainder;
///
/// Contains [`quotient_by_int::NumberRingQuotientByInt`], the standard implementation
/// of a number ring modulo a principal ideal generated by a single integer in `Z`. 
/// 
pub mod quotient_by_int;
///
/// Contains [`quotient_by_ideal::NumberRingQuotientByIdeal`], the standard implementation
/// of a number ring modulo an ideal of the form `I = (p^e, t(ϑ))`, where `ϑ` is the standard
/// generator of the number ring and `p` a prime.
/// 
pub mod quotient_by_ideal;
///
/// Contains [`pow2_cyclotomic::Pow2CyclotomicNumberRing`], an implementation of
/// [`NumberRingDescriptor`] for `Z[X]/(X^m + 1)` with `m` a power of two.
/// 
pub mod pow2_cyclotomic;
///
/// Contains [`tensor_ring::TensorProductNumberRing`], an implementation
/// of [`NumberRingDescriptor`] for `Z[X]/(Phi_(m1 m2))` with `m1, m2` coprime integers.
/// 
pub mod tensor_ring;
///
/// Contains [`general_cyclotomic::OddSquarefreeCyclotomicNumberRing`], an implementation
/// of [`NumberRingDescriptor`] for `Z[X]/(Phi_m)` with `m` an odd and squarefree integer.
/// 
pub mod general_cyclotomic;
///
/// Contains various types to represent the isomorphism `Fp[X]/(Phi_m(X)) ~ F_(p^d)^(phi(m) / d)`
/// via "hypercube structures".
/// 
pub mod hypercube;

///
/// Trait for rings of the form `R/I`, where `R` is a number ring and
/// `I` an ideal of `R`, together with the Galois action
/// ```text
///   { σ in Gal(R/Z) | σ(I) = I } x R/I -> R/I,  (σ, x) -> σx
/// ```
/// 
pub trait NumberRingQuotient: FiniteRing + FreeAlgebra {

    ///
    /// Type of the number ring that this ring is a quotient of.
    /// 
    type NumberRing: NumberRingDescriptor;

    ///
    /// Returns the number ring that this ring is a quotient of.
    /// 
    fn number_ring(&self) -> &Self::NumberRing;

    ///
    /// Returns the part of the Galois group of `R` that acts on `R/I`, i.e.
    /// the group of Galois automorphisms fixing `I`
    /// ```text
    ///   { σ in Gal(R/Z) | σ(I) = I }
    /// ```
    /// 
    fn acting_galois_group(&self) -> &Subgroup<CyclotomicGaloisGroup>;

    ///
    /// Computes `σ(x)` for the given Galois automorphism `σ`.
    /// 
    fn apply_galois_action(&self, x: &Self::Element, g: &GaloisGroupEl) -> Self::Element;

    ///
    /// Computes `σ(x)` for each Galois automorphism `σ` in the given list.
    /// 
    /// This may be faster than using [`NumberRingQuotient::apply_galois_action()`] multiple times.
    /// 
    fn apply_galois_action_many(&self, x: &Self::Element, gs: &[GaloisGroupEl]) -> Vec<Self::Element> {
        gs.iter().map(move |g| self.apply_galois_action(&x, g)).collect()
    }
}

///
/// [`RingStore`] for [`NumberRingQuotient`]
/// 
pub trait NumberRingQuotientStore: RingStore
    where Self::Type: NumberRingQuotient
{
    delegate!{ NumberRingQuotient, fn number_ring(&self) -> &<Self::Type as NumberRingQuotient>::NumberRing }
    delegate!{ NumberRingQuotient, fn acting_galois_group(&self) -> &Subgroup<CyclotomicGaloisGroup> }
    delegate!{ NumberRingQuotient, fn apply_galois_action(&self, x: &El<Self>, g: &GaloisGroupEl) -> El<Self> }
    delegate!{ NumberRingQuotient, fn apply_galois_action_many(&self, x: &El<Self>, gs: &[GaloisGroupEl]) -> Vec<El<Self>> }
}

impl<R: RingStore> NumberRingQuotientStore for R
    where R::Type: NumberRingQuotient
{}

///
/// Trait for objects that describe number rings (more concretely, orders in number
/// fields), endowed with certain information that is necessary to perform HE-related
/// operations efficiently.
///  - The number ring should have a fixed "small basis", which consists of elements that
///    have small canonical norm. This basis will be used for operations that care about
///    "smallness" and noise growth.
///  - The object should also provide functionality to search for primes `p` modulo which
///    `R/p` is isomorphic to `Fp^n`, and thus `R/p` has the "multiplicative basis" corresponding
///    to the unit vectors on the right hand side of `R/p ~ Fp^n`. This basis has the nice
///    property that it allows computing multiplications component-wise, however note that it
///    is not fixed but depends on `p`.
///  - Finally, we also want the order to be generated by a single element `a`, which gives rise
///    to the "coeff basis", given by `1, a, a^2, ..., a^(n - 1)`.
/// 
/// Two [`NumberRingDescriptor`]s that are considered equal as given by [`PartialEq`] should
/// represent the same ring, and also all three basis should coincide (in case of the "multiplicative
/// basis", it should coincide for every prime `p`).
/// 
pub trait NumberRingDescriptor: PartialEq + Clone + Debug {

    type NumberRingQuotientBases: NumberRingQuotientBases;

    fn bases_mod_p(&self, Fp: zn_64::Zn) -> Self::NumberRingQuotientBases;

    fn mod_p_required_root_of_unity(&self) -> u64;

    ///
    /// Returns an upper bound on the value
    /// ```text
    ///   sup_(x, y in R \ {0}) | xy |_sbinf / (| x |_sbinf | y |_sbinf)
    /// ```
    /// where `| x |_sbinf` is the infinity-norm of the coefficients of `x`
    /// when represented w.r.t. the small basis.
    /// 
    fn small_basis_product_expansion_factor(&self) -> f64;

    ///
    /// Returns an upper bound on the value
    /// ```text
    ///   sup_(x, y in R \ {0}) | xy |_coeffinf / (| x |_coeffinf | y |_coeffinf)
    /// ```
    /// where `| x |_coeffinf` is the infinity-norm of the coefficients of `x`
    /// when represented w.r.t. the coefficient basis.
    /// 
    fn coeff_basis_product_expansion_factor(&self) -> f64;

    fn generating_poly<P>(&self, poly_ring: P) -> El<P>
        where P: RingStore,
            P::Type: PolyRing + DivisibilityRing,
            <<P::Type as RingExtension>::BaseRing as RingStore>::Type: IntegerRing;

    fn rank(&self) -> usize;

    fn galois_group(&self) -> &CyclotomicGaloisGroup;
}

///
/// A [`NumberRingDescriptor`] `R` modulo a prime `p` that splits completely in `R`.
/// 
/// This object may define up to three different bases of `R / p`, with the following
/// properties:
///  - the "small basis" should consist of elements whose shortest lift to `R` has small
///    canonical norm
///  - the "multiplicative basis" should allow for component-wise multiplication, i.e. `b_i * b_i = b_i`
///    and `b_i * b_j = 0` for `i != j`
///  - the "coeff basis" should consist of powers of a generator of the ring, which for
///    cyclotomic rings should be the root of unity.
/// 
/// Both "small basis" and "coeff basis" should be the reduction of a corresponding
/// canonical basis of `R`.
/// 
/// Note that it is valid for any of these basis to coincide, and then implement the 
/// corresponding conversions as no-ops.
/// 
/// This design is motivated by the example of `Z[𝝵_m]` for a composite `m = m_1 * m_2`, 
/// since in this case, we need three different basis.
///  - The "small basis" is the tensor-product basis `𝝵^(m_2 * i_1 + m_1 * i_2)` with
///    `0 <= i_1 < phi(m_1)` and `0 <= i_2 < phi(m_2)`.
///  - The "multiplicative basis" is the preimage of the unit vector basis under `Fp[𝝵] -> Fp^phi(m)`
///  - The "coeff basis" is the basis `1, 𝝵, 𝝵^2, ..., 𝝵^phi(m)`
/// While one could choose "small basis" and "coeff basis" to be equal (after all, the
/// elements `𝝵^i` are all "small"), staying in "small basis" whenever possible has
/// performance benefits, because of the tensor-decomposition.
/// 
pub trait NumberRingQuotientBases: PartialEq {

    fn base_ring(&self) -> &zn_64::Zn;

    fn rank(&self) -> usize;

    fn small_basis_to_mult_basis<V>(&self, data: V)
        where V: SwappableVectorViewMut<zn_64::ZnEl>;

    fn mult_basis_to_small_basis<V>(&self, data: V)
        where V: SwappableVectorViewMut<zn_64::ZnEl>;

    fn coeff_basis_to_small_basis<V>(&self, data: V)
        where V: SwappableVectorViewMut<zn_64::ZnEl>;

    fn small_basis_to_coeff_basis<V>(&self, data: V)
        where V: SwappableVectorViewMut<zn_64::ZnEl>;
        
    ///
    /// Permutes the components of an element w.r.t. the mult basis to
    /// obtain its image under the given Galois action.
    /// 
    fn permute_galois_action<V1, V2>(&self, src: V1, dst: V2, g: &GaloisGroupEl)
        where V1: VectorView<zn_64::ZnEl>,
            V2: SwappableVectorViewMut<zn_64::ZnEl>;
            
    fn galois_group(&self) -> &CyclotomicGaloisGroup;
}

///
/// Returns the largest prime number that is `<= leq_than` and `= 1 mod congruent_to_one_mod`,
/// or `None` if there is no such prime number.
/// 
pub fn largest_prime_leq_congruent_to_one(leq_than: i64, congruent_to_one_mod: i64) -> Option<i64> {
    assert!(leq_than > congruent_to_one_mod);
    let mut current = leq_than - (leq_than - 1) % congruent_to_one_mod;
    while !is_prime(&StaticRing::<i64>::RING, &current, 10) {
        current -= congruent_to_one_mod;
        if current <= 0 {
            return None;
        }
    }
    return Some(current);
}

///
/// Attempts to find a list of distinct primes, each smaller than `2^max_bits_each_modulus`, whose 
/// product is between `2^min_bits` and `2^max_bits`.
/// 
/// Only primes that are returned by the given function are used, which allows the caller to sample
/// a list of primes that satisfy additional constraints, like being `= 1 mod m` for some integer `m`.
/// More concretely, the given function `largest_prime_leq` should, on input `B`, return the largest
/// prime that satisfies all desired constraint and is `<= B`, or `None` if no such prime exists.
/// 
/// This function operates on a best-effort basis, it might not find an accepted result in extreme cases
/// where `min_bits` and `max_bits` are very small or very close together, even if a result theoretically
/// exists. It is, however, quite reliable in most situations.
/// 
pub fn sample_primes<F>(min_bits: usize, max_bits: usize, max_bits_each_modulus: usize, largest_prime_leq: F) -> Option<Vec<i64>>
    where F: FnMut(i64) -> Option<i64>
{
    extend_sampled_primes(&[], min_bits, max_bits, max_bits_each_modulus, largest_prime_leq)
}

///
/// Like [`sample_primes()`], but starts with a non-empty list of integers. Additional primes are added
/// to the end of this list, and only primes that do not appear in the list are added. Note that in the
/// case that the list contains non-primes, newly added primes may divide numbers already in the list.
/// 
/// Only primes that are returned by the given function are used, which allows the caller to sample
/// a list of primes that satisfy additional constraints, like being `= 1 mod m` for some integer `m`.
/// More concretely, the given function `largest_prime_leq` should, on input `B`, return the largest
/// prime that satisfies all desired constraint and is `<= B`, or `None` if no such prime exists.
/// 
/// This function operates on a best-effort basis, it might not find an accepted result in extreme cases
/// where `min_bits` and `max_bits` are very small or very close together, even if a result theoretically
/// exists. It is, however, quite reliable in most situations.
/// 
pub fn extend_sampled_primes<F>(begin_with: &[i64], min_bits: usize, max_bits: usize, max_bits_each_modulus: usize, mut largest_prime_leq: F) -> Option<Vec<i64>>
    where F: FnMut(i64) -> Option<i64>
{
    let ZZbig = BigIntRing::RING;
    assert!(max_bits > min_bits);

    let mut result = begin_with.to_vec();
    let mut current_bits = result.iter().map(|p| (*p as f64).log2()).sum::<f64>();
    assert!((current_bits.floor() as usize) < max_bits);
    let mut current_upper_bound = 1 << max_bits_each_modulus;

    while current_bits < min_bits as f64 {

        if min_bits as f64 - current_bits < max_bits_each_modulus as f64 {  
            current_upper_bound = i64::min(current_upper_bound, 1 << (f64::min(max_bits as f64 - current_bits, max_bits_each_modulus as f64).floor() as usize));
        } else {
            let required_number_of_primes = ((min_bits as f64 - current_bits) / max_bits_each_modulus as f64).ceil() as usize;
            current_upper_bound = i64::min(current_upper_bound, 1 << (f64::min((max_bits as f64 - current_bits) / required_number_of_primes as f64, max_bits_each_modulus as f64).floor() as usize));
        }

        let mut prime = largest_prime_leq(current_upper_bound)?;
        current_upper_bound = prime - 1;
        while begin_with.iter().any(|p| *p == prime) {
            prime = largest_prime_leq(current_upper_bound)?;
            current_upper_bound = prime - 1
        }
        let bits = (prime as f64).log2();
        current_bits += bits;
        result.push(prime);
    }
    debug_assert!(ZZbig.is_geq(&ZZbig.prod(result.iter().map(|p| int_cast(*p, ZZbig, ZZi64))), &ZZbig.power_of_two(min_bits)));
    debug_assert!(ZZbig.is_lt(&ZZbig.prod(result.iter().map(|p| int_cast(*p, ZZbig, ZZi64))), &ZZbig.power_of_two(max_bits)));
    return Some(result);
}

#[cfg(test)]
use feanor_math::pid::EuclideanRingStore;

#[test]
fn test_sample_primes() {
    feanor_tracing::DelayedLogger::init_test();
    let ZZbig = BigIntRing::RING;
    let result = sample_primes(60, 62, 58, |b| largest_prime_leq_congruent_to_one(b, 422144)).unwrap();
    assert_eq!(result.len(), 2);
    let prod = ZZbig.prod(result.iter().map(|p| int_cast(*p, ZZbig, ZZi64)));
    assert!(ZZbig.abs_log2_floor(&prod).unwrap() >= 60);
    assert!(ZZbig.abs_log2_ceil(&prod).unwrap() <= 62);
    assert!(result.iter().all(|p| ZZbig.is_one(&ZZbig.euclidean_rem(int_cast(*p, ZZbig, ZZi64), &int_cast(422144, ZZbig, StaticRing::<i64>::RING)))));

    let ZZbig = BigIntRing::RING;
    let result = sample_primes(135, 138, 58, |b| largest_prime_leq_congruent_to_one(b, 422144)).unwrap();
    assert_eq!(result.len(), 3);
    let prod = ZZbig.prod(result.iter().map(|p| int_cast(*p, ZZbig, ZZi64)));
    assert!(ZZbig.abs_log2_floor(&prod).unwrap() >= 135);
    assert!(ZZbig.abs_log2_ceil(&prod).unwrap() <= 138);
    assert!(result.iter().all(|p| ZZbig.is_one(&ZZbig.euclidean_rem(int_cast(*p, ZZbig, ZZi64), &int_cast(422144, ZZbig, StaticRing::<i64>::RING)))));

    let ZZbig = BigIntRing::RING;
    let result = sample_primes(115, 118, 58, |b| largest_prime_leq_congruent_to_one(b, 422144)).unwrap();
    assert_eq!(result.len(), 2);
    let prod = ZZbig.prod(result.iter().map(|p| int_cast(*p, ZZbig, ZZi64)));
    assert!(ZZbig.abs_log2_floor(&prod).unwrap() >= 115);
    assert!(ZZbig.abs_log2_ceil(&prod).unwrap() <= 118);
    assert!(result.iter().all(|p| ZZbig.is_one(&ZZbig.euclidean_rem(int_cast(*p, ZZbig, ZZi64), &int_cast(422144, ZZbig, StaticRing::<i64>::RING)))));
}