
use feanor_math::algorithms::miller_rabin::is_prime;
use feanor_math::divisibility::*;
use feanor_math::integer::*;
use feanor_math::ordered::*;
use feanor_math::primitive_int::{StaticRing, StaticRingBase};
use feanor_math::ring::*;
use feanor_math::rings::poly::dense_poly::DensePolyRing;
use feanor_math::rings::poly::PolyRingStore;
use feanor_math::rings::zn::zn_64::Zn;
use feanor_math::homomorphism::*;
use polys::{digit_retain_poly, poly_to_circuit, precomputed_p_2};
use tracing::instrument;

use crate::circuit::PlaintextCircuit;
use crate::digit_extract::polys::bounded_digit_retain_poly;
use crate::{ZZbig, ZZi64};

///
/// Contains various tools to digit extraction polynomials and convert them into circuits,
/// in particular [`polys::poly_to_circuit()`].
/// 
pub mod polys;

///
/// The digit extraction operation, as required during BFV and
/// BGV bootstrapping.
/// 
/// Concretely, this encapsulates an efficient implementation of the
/// per-slot digit extraction function
/// ```text
///   Z/p^eZ -> Z/p^rZ x Z/p^eZ,  x -> (x - (x mod p^v) / p^v, x mod p^v)
/// ```
/// for `v = e - r`. Here `x mod p^v` refers to the smallest positive element
/// of `Z/p^eZ` that is congruent to `x` modulo `p^v`.
/// 
/// This function can also be applied to values in a ring `Z/p^e'Z` for
/// `e' > e`, i.e. it will then have the signature
/// ```text
///   Z/p^e'Z -> Z/p^(e' - e + r)Z x Z/p^e'Z
/// ```
/// In this case, the results are only specified modulo `p^r` resp. `p^e`, i.e.
/// may be perturbed by an arbitrary value `p^r a` resp. `p^e a'`.
/// 
pub struct DigitExtract<R: ?Sized + RingBase = StaticRingBase<i64>> {
    extraction_circuits: Vec<(Vec<usize>, PlaintextCircuit<R>)>,
    /// the one-input, one-output identity circuit
    identity_circuit: PlaintextCircuit<R>,
    /// the two-input, one-output addition circuit
    add_circuit: PlaintextCircuit<R>,
    /// the two-input, one-output subtraction circuit
    sub_circuit: PlaintextCircuit<R>,
    v: usize,
    e: usize,
    p: El<BigIntRing>
}

impl DigitExtract {

    ///
    /// Creates a [`DigitExtract`] for a scalar ring `Z/2^eZ`.
    /// 
    /// Uses the precomputed table of best digit extraction circuits for `e <= 23`.
    /// 
    #[instrument(skip_all)]
    pub fn new_precomputed_p_is_2(p: i64, e: usize, r: usize) -> Self {
        assert_eq!(2, p);
        assert!(is_prime(&StaticRing::<i64>::RING, &p, 10));
        return Self::new_with_circuits(
            int_cast(p, ZZbig, ZZi64), 
            e, 
            r, 
            StaticRing::<i64>::RING, 
            [1, 2, 4, 8, 16, 23].into_iter().map(|e| (
                [1, 2, 4, 8, 16, 23].into_iter().take_while(|i| *i <= e).collect(),
                precomputed_p_2(e)
            )).collect::<Vec<_>>()
        );
    }
    
    ///
    /// Creates a [`DigitExtract`] for a scalar ring `Z/p^eZ`.
    /// 
    /// Uses the Chen-Han digit retain polynomials <https://ia.cr/2018/067> together with
    /// a heuristic method to compile them into an arithmetic circuit, based on the
    /// Paterson-Stockmeyer method.
    /// 
    #[instrument(skip_all)]
    pub fn new_default(p: i64, e: usize, r: usize) -> Self {
        assert!(is_prime(&StaticRing::<i64>::RING, &p, 10));
        assert!(e > r);
        let v = e - r;
        
        let digit_extraction_circuits = (1..=v).rev().map(|i| {
            let required_digits = (2..=(v - i + 1)).chain([r + v - i + 1].into_iter()).collect::<Vec<_>>();
            let poly_ring = DensePolyRing::new(Zn::new(StaticRing::<i64>::RING.pow(p, *required_digits.last().unwrap()) as u64), "X");
            let circuit = poly_to_circuit(&poly_ring, &required_digits.iter().map(|j| digit_retain_poly(&poly_ring, *j)).collect::<Vec<_>>());
            return (required_digits, circuit);
        }).collect::<Vec<_>>();
        assert!(digit_extraction_circuits.is_sorted_by_key(|(digits, _)| *digits.last().unwrap()));
        
        return Self::new_with_circuits(int_cast(p, ZZbig, ZZi64), e, r, StaticRing::<i64>::RING, digit_extraction_circuits);
    }

    ///
    /// Creates a [`DigitExtract`] for a scalar ring `Z/p^eZ`.
    /// 
    /// Uses the Ma, Huang, Wang and Want digit extraction polynomials <https://ia.cr/2024/115> 
    /// for errors bounded by `B` and large `p`, together with a heuristic method to compile them
    /// into an algebraic circuit, based on the Paterson-Stockmeyer method.
    /// 
    #[instrument(skip_all)]
    pub fn new_bounded_error(p: i64, e: usize, B: i64) -> Self {
        assert!(is_prime(&ZZi64, &p, 10));
        assert!(B >= 0);
        assert!(2 * B + 1 <= p);
        
        let poly_ring = DensePolyRing::new(Zn::new(StaticRing::<i64>::RING.pow(p, e) as u64), "X");
        let p_half = poly_ring.inclusion().map(poly_ring.base_ring().coerce(&ZZi64, p / 2));
        let digit_extraction_circuits = vec![
            (vec![e], poly_to_circuit(&poly_ring, &[
                poly_ring.add(poly_ring.evaluate(&bounded_digit_retain_poly(&poly_ring, B), &poly_ring.sub_ref_snd(poly_ring.indeterminate(), &p_half), poly_ring.inclusion()), p_half)
            ]))
        ];
        
        return Self::new_with_circuits(int_cast(p, ZZbig, ZZi64), e, e - 1, StaticRing::<i64>::RING, digit_extraction_circuits);
    }
}

impl<R: ?Sized + RingBase> DigitExtract<R> {

    ///
    /// Creates a new [`DigitExtract`] from the given circuits.
    /// 
    /// This functions expects the list of circuits to contain tuples `(digits, C)`,
    /// where the circuit `C` takes a single input and computes `digits.len()` outputs, 
    /// such that the `i`-th output is congruent to `lift(input mod p)` modulo 
    /// `p^digits[i]`.
    /// 
    /// If you want to use the default choice of circuits, consider using [`DigitExtract::new_default()`].
    /// 
    pub fn new_with_circuits<S: Copy + RingStore<Type = R>>(p: El<BigIntRing>, e: usize, r: usize, ring: S, extraction_circuits: Vec<(Vec<usize>, PlaintextCircuit<R>)>) -> Self {
        assert!(is_prime(ZZbig, &p, 10));
        assert!(e > r);
        for (digits, circuit) in &extraction_circuits {
            assert!(digits.is_sorted());
            assert_eq!(digits.len(), circuit.output_count());
            assert_eq!(1, circuit.input_count());
        }
        assert!(extraction_circuits.iter().any(|(digits, _)| *digits.last().unwrap() >= e));
        Self {
            extraction_circuits: extraction_circuits,
            add_circuit: PlaintextCircuit::add(ring),
            sub_circuit: PlaintextCircuit::sub(ring),
            identity_circuit: PlaintextCircuit::identity(1, ring),
            v: e - r,
            p: p,
            e: e
        }
    }

    pub fn r(&self) -> usize {
        self.e - self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn p(&self) -> &El<BigIntRing> {
        &self.p
    }
    
    ///
    /// Evaluates the digit extraction function over any representation of elements of `Z/p^iZ`, which
    /// supports the evaluation of [`PlaintextCircuit`]s. Since digit extraction requires computations
    /// in all the rings `Z/p^(r - 1)Z, ...., Z/p^eZ`, we also require a `change_space` function, with
    /// the following properties:
    /// ```text
    ///   change_space(e, e', .): Z/p^eZ -> Z/p^e' Z
    ///   change_space(e, e', x mod p^e) = x p^(e' - e) mod p^e'      if e' > e
    ///   change_space(e, e', x mod p^e) = x / p^(e - e') mod p^e'    if e' < e and p^(e - e') | x
    /// ```
    /// If the passed functions behave as specified, `change_space(e, e', x)` will never be called for
    /// `e' < e` and an `x` which is not divisible by `p^(e - e')`.
    /// 
    /// Furthermore, the `eval_circuit` is given the exponent of the current ring we work in as the first
    /// parameter. The result of [`DigitExtract::evaluate_generic()`] is then the tuple `(quo, rem)` with
    /// `quo` in `Z/p^rZ` and `rem` in `Z/p^eZ` such that `x = p^(e - r) * quo + rem` and `rem < p^(e - r)`.
    /// 
    /// If [`DigitExtract`] is used on elements of `Z/p^e'Z` with `e' > e` (as mentioned at the end of
    /// the doc of [`DigitExtract`]), the moduli passed to `eval_circuit()` and `change_space()` remain
    /// nevertheless unchanged - after all, `evaluate_generic()` does not know that we are in a larger
    /// ring. If necessary, you have to manually offset all exponents passed to `eval_circuit` and 
    /// `change_space` by `e' - e`.
    /// 
    pub fn evaluate_generic<T, EvalCircuit, ChangeSpace>(&self, 
        input: T,
        mut eval_circuit: EvalCircuit,
        mut change_space: ChangeSpace
    ) -> (T, T) 
        where EvalCircuit: FnMut(/* exponent of p */ usize, &[T], &PlaintextCircuit<R>) -> Vec<T>,
            ChangeSpace: FnMut(/* input exponent of p */ usize, /* output exponent of p */ usize, T) -> T
    {
        let e = self.e;
        let r = self.e - self.v;

        enum OneOrTwoValues<T> {
            One(T), Two([T; 2])
        }

        impl<T> OneOrTwoValues<T> {

            fn with_first_el<'a>(&'a mut self, first: T) -> &'a mut [T; 2] {
                take_mut::take(self, |value| match value {
                    OneOrTwoValues::One(second) => OneOrTwoValues::Two([first, second]),
                    OneOrTwoValues::Two([_, second]) => OneOrTwoValues::Two([first, second])
                });
                return match self {
                    OneOrTwoValues::One(_) => unreachable!(),
                    OneOrTwoValues::Two(data) => data
                };
            }

            fn get_second<'a>(&'a self) -> &'a T {
                match self {
                    OneOrTwoValues::One(second) => second,
                    OneOrTwoValues::Two([_, second]) => second
                }
            }
        }

        let clone_value = |modulus_exp: usize, value: &T, eval_circuit: &mut EvalCircuit| eval_circuit(modulus_exp, std::slice::from_ref(value), &self.identity_circuit).into_iter().next().unwrap();
        let sub_values = |modulus_exp: usize, params: &[T; 2], eval_circuit: &mut EvalCircuit| eval_circuit(modulus_exp, params, &self.sub_circuit).into_iter().next().unwrap();
        let add_values = |modulus_exp: usize, params: &[T; 2], eval_circuit: &mut EvalCircuit| eval_circuit(modulus_exp, params, &self.add_circuit).into_iter().next().unwrap();

        let mut mod_result: Option<T> = None;
        let mut partial_floor_divs = (0..self.v).map(|_| Some(clone_value(e, &input, &mut eval_circuit))).collect::<Vec<_>>();
        let mut floor_div_result = input;
        for i in 0..self.v {
            let remaining_digits = e - i;
            debug_assert!(self.extraction_circuits.is_sorted_by_key(|(digits, _)| *digits.last().unwrap()));
            let (use_circuit_digits, use_circuit) = self.extraction_circuits.iter().filter(|(digits, _)| *digits.last().unwrap() >= remaining_digits).next().unwrap();
            debug_assert!(use_circuit_digits.is_sorted());

            let current = change_space(e, remaining_digits, partial_floor_divs[i].take().unwrap());
            let digit_extracted = eval_circuit(remaining_digits, std::slice::from_ref(&current), use_circuit);
            let mut digit_extracted = digit_extracted.into_iter().map(|value| OneOrTwoValues::One(change_space(remaining_digits, e, value))).collect::<Vec<_>>();
            
            let last_digit_extracted = digit_extracted.last_mut().unwrap();
            take_mut::take(&mut floor_div_result, |current| sub_values(e, last_digit_extracted.with_first_el(current), &mut eval_circuit));
            if let Some(mod_result) = &mut mod_result {
                take_mut::take(mod_result, |current| add_values(e, last_digit_extracted.with_first_el(current), &mut eval_circuit));
            } else {
                mod_result = Some(clone_value(e, last_digit_extracted.get_second(), &mut eval_circuit));
            }

            for j in (i + 1)..self.v {
                let digit_extracted_index = use_circuit_digits.iter().enumerate().filter(|(_, cleared_digits)| **cleared_digits > j - i).next().unwrap().0;
                take_mut::take(partial_floor_divs[j].as_mut().unwrap(), |current| sub_values(e, digit_extracted[digit_extracted_index].with_first_el(current), &mut eval_circuit));
            }
        }

        return (change_space(e, r, floor_div_result), mod_result.unwrap());
    }

    ///
    /// Computes `(quo, rem)` with `input = quo * p^(e - r) + rem` and `rem < p^(e - r)`.
    /// Note that both `quo` and `rem` are returned as elements of `Z/p^eZ`, which means that
    /// `quo` is defined only up to a multiple of `p^r`.
    /// 
    /// This function is designed to test digit extraction, since `quo` and `rem` will be computed
    /// exactly in the same way as in a homomorphic setting. Note also that performing euclidean
    /// division can be done much easier with [`EuclideanRing::euclidean_div_rem()`]
    /// when you have access to the ring elements.
    /// 
    /// This function does not perform any checks on the underlying ring, in particular, you can
    /// call it on an input in `Z/p^e'Z` with `e' > e` or an input in `Z`. Of course, in any case,
    /// the output will only be correct modulo `p^r` resp. `p^e`.
    /// 
    /// [`EuclideanRing::euclidean_div_rem()`]: feanor_math::pid::EuclideanRing::euclidean_div_rem()
    /// 
    pub fn evaluate<H, S>(&self, input: S::Element, hom: H) -> (S::Element, S::Element)
        where H: Homomorphism<R, S>,
            S: ?Sized + RingBase + DivisibilityRing
    {
        // temporarily copied from feanor-math
        fn map_from_integer_ring<I, R>(from: I, to: R, mut x: El<I>) -> El<R>
            where I: RingStore,
                I::Type: IntegerRing,
                R: RingStore
        {
            let basis = to.int_hom().map(1 << 16);
            let is_neg = if from.is_neg(&x) {
                from.negate_inplace(&mut x);
                true
            } else {
                false
            };
            let mut current = to.zero();
            let mut current_pow = to.one();
            while !from.is_zero(&x) {
                let mut quo = from.clone_el(&x);
                from.euclidean_div_pow_2(&mut quo, 16);
                let mut rem = from.clone_el(&quo);
                from.mul_pow_2(&mut rem, 16);
                from.sub_self_assign(&mut rem, x);
                let rem = int_cast(rem, StaticRing::<i32>::RING, &from);
                to.add_assign(&mut current, to.mul_ref_snd(to.int_hom().map(rem), &current_pow));
                x = quo;
                to.mul_assign_ref(&mut current_pow, &basis);
            }
            if is_neg {
                return to.negate(current);
            } else {
                return current;
            }
        }

        let p = map_from_integer_ring(ZZbig, hom.codomain(), ZZbig.clone_el(&self.p));
        self.evaluate_generic(
            input,
            |_, params, circuit| circuit.evaluate_no_galois(params, &hom),
            |from, to, x| if from < to {
                hom.codomain().mul(x, hom.codomain().pow(hom.codomain().clone_el(&p), to - from))
            } else {
                hom.codomain().checked_div(&x, &hom.codomain().pow(hom.codomain().clone_el(&p), from - to)).unwrap()
            }
        )
    }
}

#[cfg(test)]
use feanor_math::rings::zn::ZnRingStore;
#[cfg(test)]
use feanor_math::assert_el_eq;
#[cfg(test)]
use feanor_math::divisibility::DivisibilityRingStore;

#[test]
fn test_digit_extract_p2_polys() {
    let digitextract = DigitExtract::new_default(2, 17, 9);
    let ring = Zn::new(ZZi64.pow(2, 17) as u64);
    let hom = ring.can_hom(&ZZi64).unwrap();

    for x in 0..*ring.modulus() {
        let (quo, rem) = digitextract.evaluate_generic(
            (17, hom.map(x)),
            |exp, params, circuit| {
                assert!(params.iter().all(|(p_exp, _)| *p_exp == exp));
                circuit.evaluate_no_galois(&params.iter().map(|(_, x)| *x).collect::<Vec<_>>(), &hom).into_iter().map(|x| (exp, x)).collect()
            },
            |from, to, (exp, x)| {
                assert_eq!(from, exp);
                if from < to {
                    (to, ring.mul(x, ring.pow(hom.map(2), to - from)))
                } else {
                    (to, ring.checked_div(&x, &ring.pow(hom.map(2), from - to)).unwrap())
                }
            }
        );
        assert_eq!(17, rem.0);
        assert_el_eq!(&ring, hom.map(x % (1 << 8)), rem.1);
        assert_eq!(9, quo.0);
        assert_eq!(x / (1 << 8), ring.smallest_positive_lift(quo.1) % (1 << 9));
    }
}

#[test]
fn test_digit_extract() {
    let digitextract = DigitExtract::new_default(3, 5, 2);
    let ring = Zn::new(StaticRing::<i64>::RING.pow(3, 5) as u64);
    let hom = ring.can_hom(&StaticRing::<i64>::RING).unwrap();

    for x in 0..*ring.modulus() {
        let (quo, rem) = digitextract.evaluate_generic(
            (5, hom.map(x)),
            |exp, params, circuit| {
                assert!(params.iter().all(|(p_exp, _)| *p_exp == exp));
                circuit.evaluate_no_galois(&params.iter().map(|(_, x)| *x).collect::<Vec<_>>(), &hom).into_iter().map(|x| (exp, x)).collect()
            },
            |from, to, (exp, x)| {
                assert_eq!(from, exp);
                if from < to {
                    (to, ring.mul(x, ring.pow(hom.map(3), to - from)))
                } else {
                    (to, ring.checked_div(&x, &ring.pow(hom.map(3), from - to)).unwrap())
                }
            }
        );
        assert_eq!(5, rem.0);
        assert_el_eq!(&ring, hom.map(x % 27), rem.1);
        assert_eq!(2, quo.0);
        assert_eq!(x / 27, ring.smallest_positive_lift(quo.1) % 9);
    }
}

#[test]
fn test_digit_extract_evaluate() {
    let ring = Zn::new(16);
    let digit_extract = DigitExtract::new_default(2, 4, 2);
    for x in 0..16 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!(x / 4, ring.smallest_positive_lift(actual_high) as i32 % 4);
        assert_eq!(x % 4, ring.smallest_positive_lift(actual_low) as i32);
    }

    let ring = Zn::new(81);
    let digit_extract = DigitExtract::new_default(3, 4, 2);
    for x in 0..81 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!(x / 9, ring.smallest_positive_lift(actual_high) as i32 % 9);
        assert_eq!(x % 9, ring.smallest_positive_lift(actual_low) as i32);
    }

    let ring = Zn::new(125);
    let digit_extract = DigitExtract::new_default(5, 3, 2);
    for x in 0..125 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!(x / 5, ring.smallest_positive_lift(actual_high) as i32 % 25);
        assert_eq!(x % 5, ring.smallest_positive_lift(actual_low) as i32);
    }
}

#[test]
fn test_digit_extract_evaluate_ignore_higher() {
    let ring = Zn::new(64);
    let digit_extract = DigitExtract::new_default(2, 4, 2);
    for x in 0..64 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!((x / 4) % 4, ring.smallest_positive_lift(actual_high) as i32 % 4);
        assert_eq!(x % 4, ring.smallest_positive_lift(actual_low) as i32 % 16);
    }

    let ring = Zn::new(243);
    let digit_extract = DigitExtract::new_default(3, 4, 2);
    for x in 0..243 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!((x / 9) % 9, ring.smallest_positive_lift(actual_high) as i32 % 9);
        assert_eq!(x % 9, ring.smallest_positive_lift(actual_low) as i32 % 81);
    }

    let ring = Zn::new(625);
    let digit_extract = DigitExtract::new_default(5, 3, 2);
    for x in 0..625 {
        let (actual_high, actual_low) = digit_extract.evaluate(ring.int_hom().map(x), ring.can_hom(&StaticRing::<i64>::RING).unwrap());
        assert_eq!((x / 5) % 25, ring.smallest_positive_lift(actual_high) as i32 % 25);
        assert_eq!(x % 5, ring.smallest_positive_lift(actual_low) as i32 % 125);
    }
}
