fheanor 0.11.1

use std::cell::RefCell;
use std::cmp::max;
use serde::Serialize;
use serde::de::DeserializeSeed;

use feanor_math::algorithms::discrete_log::Subgroup;
use feanor_math::group::AbelianGroupStore;
use serialization::{DeserializeSeedPlaintextCircuit, SerializablePlaintextCircuit};
use evaluator::CircuitEvaluator;
use evaluator::HomEvaluator;
use evaluator::HomEvaluatorGal;
use feanor_math::homomorphism::Homomorphism;
use feanor_math::ring::*;
use feanor_math::rings::zn::*;
use feanor_math::serialization::SerializableElementRing;
use tracing::instrument;
use fhe_ir::Program;

use crate::circuit::ir::*;
use crate::cache::*;
use crate::number_ring::galois::*;
use crate::number_ring::hypercube::isomorphism::BaseRing;
use crate::number_ring::*;

mod serialization;

///
/// Contains the trait [`evaluator::CircuitEvaluator`] and different implementations, which describe
/// how to evaluate an arithmetic circuit.
/// 
pub mod evaluator;
///
/// Contains an interface to the FHE intermediate representation format specified by [`fhe_ir`]. 
/// 
pub mod ir;

#[derive(Copy, Clone, Debug)]
pub struct CircuitEvaluatorCosts {
    pub cost_mul: f64,
    pub cost_sqr: f64,
    pub cost_single_gal: f64,
    pub cost_setup_hoisted_gal: f64,
    pub cost_hoisted_gal: f64
}

///
/// A coefficient used in a [`PlaintextCircuit`].
/// 
/// Generally speaking, this always represents an element of `R`
/// (which can be retrieved via [`Coefficient::to_ring_el()`]), but
/// special cases are stored separately for a more efficient evaluation.
/// 
pub enum Coefficient<R: ?Sized + RingBase> {
    Zero, One, NegOne, Integer(i32), Other(R::Element)
}

impl<R> Clone for Coefficient<R>
    where R: ?Sized + RingBase,
        R::Element: Clone
{
    fn clone(&self) -> Self {
        match self {
            Coefficient::Zero => Coefficient::Zero,
            Coefficient::One => Coefficient::One,
            Coefficient::NegOne => Coefficient::NegOne,
            Coefficient::Integer(x) => Coefficient::Integer(*x),
            Coefficient::Other(x) => Coefficient::Other(x.clone())
        }
    }
}

impl<R> Copy for Coefficient<R>
    where R: ?Sized + RingBase,
        R::Element: Copy
{}

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

    pub fn clone<S: RingStore<Type = R>>(&self, ring: S) -> Self {
        match self {
            Coefficient::Zero => Coefficient::Zero,
            Coefficient::One => Coefficient::One,
            Coefficient::NegOne => Coefficient::NegOne,
            Coefficient::Integer(x) => Coefficient::Integer(*x),
            Coefficient::Other(x) => Coefficient::Other(ring.clone_el(x))
        }
    }

    pub fn eq<S: RingStore<Type = R> + Copy>(&self, other: &Self, ring: S) -> bool {
        ring.eq_el(&self.clone(ring).to_ring_el(ring), &other.clone(ring).to_ring_el(ring))
    }

    ///
    /// Computes `self + x`, but avoids a full ring addition if `self` is zero.
    /// 
    pub fn add_to<S: RingStore<Type = R> + Copy>(&self, x: El<S>, ring: S) -> El<S> {
        match self {
            Coefficient::Zero => x,
            Coefficient::One => ring.add(x, ring.one()),
            Coefficient::NegOne => ring.add(x, ring.neg_one()),
            Coefficient::Integer(y) => ring.add(x, ring.int_hom().map(*y)),
            Coefficient::Other(y) => ring.add_ref_snd(x, y)
        }
    }

    ///
    /// Computes `self * x`, but avoids a full ring multiplication if `self`
    /// is stored as an integer.
    /// 
    pub fn mul_to<S: RingStore<Type = R> + Copy>(&self, x: El<S>, ring: S) -> El<S> {
        match self {
            Coefficient::Zero => ring.zero(),
            Coefficient::One => x,
            Coefficient::NegOne => ring.negate(x),
            Coefficient::Integer(y) => ring.int_hom().mul_map(x, *y),
            Coefficient::Other(y) => ring.mul_ref_snd(x, y)
        }
    }

    pub fn is_zero(&self) -> bool {
        match self {
            Coefficient::Zero => true,
            _ => false
        }
    }

    fn from<S: RingStore<Type = R> + Copy>(el: El<S>, ring: S) -> Self {
        if ring.is_zero(&el) {
            Coefficient::Zero
        } else if ring.is_one(&el) {
            Coefficient::One
        } else {
            Coefficient::Other(el)
        }
    }

    pub fn to_ring_el<S: RingStore<Type = R>>(self, ring: S) -> El<S> {
        match self {
            Coefficient::Zero => ring.zero(),
            Coefficient::One => ring.one(),
            Coefficient::NegOne => ring.neg_one(),
            Coefficient::Integer(x) => ring.int_hom().map(x),
            Coefficient::Other(x) => x
        }
    }

    pub fn negate<S: RingStore<Type = R>>(self, ring: S) -> Self {
        match self {
            Coefficient::Zero => Coefficient::Zero,
            Coefficient::One => Coefficient::NegOne,
            Coefficient::NegOne => Coefficient::One,
            Coefficient::Integer(x) => Coefficient::Integer(-x),
            Coefficient::Other(x) => Coefficient::Other(ring.negate(x))
        }
    }

    pub fn add<S: RingStore<Type = R> + Copy>(self, other: Self, ring: S) -> Self {
        match (self, other) {
            (Coefficient::Zero, rhs) => rhs,
            (lhs, Coefficient::Zero) => lhs,
            (Coefficient::One, Coefficient::Integer(x)) => Coefficient::Integer(x + 1),
            (Coefficient::NegOne, Coefficient::Integer(x)) => Coefficient::Integer(x - 1),
            (Coefficient::Integer(x), Coefficient::One) => Coefficient::Integer(x + 1),
            (Coefficient::Integer(x), Coefficient::NegOne) => Coefficient::Integer(x - 1),
            (lhs, rhs) => Coefficient::Other(ring.add(lhs.to_ring_el(ring), rhs.to_ring_el(ring)))
        }
    }

    pub fn mul<S: RingStore<Type = R> + Copy>(self, other: Self, ring: S) -> Self {
        match (self, other) {
            (Coefficient::Zero, _) => Coefficient::Zero,
            (_, Coefficient::Zero) => Coefficient::Zero,
            (Coefficient::One, rhs) => rhs,
            (lhs, Coefficient::One) => lhs,
            (lhs, Coefficient::NegOne) => lhs.negate(ring),
            (Coefficient::NegOne, rhs) => rhs.negate(ring),
            (lhs, rhs) => Coefficient::Other(ring.mul(lhs.to_ring_el(ring), rhs.to_ring_el(ring)))
        }
    }

    ///
    /// Preserves integer coefficients, and maps ring elements as specified
    /// by the given function.
    /// 
    pub fn change_ring<S, F>(self, mut f: F) -> Coefficient<S>
        where F: FnMut(R::Element) -> S::Element,
            S: ?Sized + RingBase
    {
        match self {
            Coefficient::Integer(x) => Coefficient::Integer(x),
            Coefficient::NegOne => Coefficient::NegOne,
            Coefficient::Zero => Coefficient::Zero,
            Coefficient::One => Coefficient::One,
            Coefficient::Other(x) => Coefficient::Other(f(x))
        }
    }

    pub fn as_integer(&self) -> Option<i32> {
        match self {
            Coefficient::Integer(x) => Some(*x),
            Coefficient::NegOne => Some(-1),
            Coefficient::Zero => Some(0),
            Coefficient::One => Some(1),
            Coefficient::Other(_) => None
        }
    }
}

///
/// A "linear combination" gate, which takes an affine linear combination
/// of an arbitrary number of inputs with coefficients in the ring, and produces
/// a single output.
/// 
struct LinearCombination<R: ?Sized + RingBase> {
    factors: Vec<Coefficient<R>>,
    constant: Coefficient<R>
}

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

    fn clone<S: RingStore<Type = R> + Copy>(&self, ring: S) -> Self {
        Self {
            factors: self.factors.iter().map(|c| c.clone(ring)).collect(),
            constant: self.constant.clone(ring)
        }
    }

    fn evaluate_generic<'a, T, E>(&'a self, first_inputs: &[T], second_inputs: &[T], evaluator: &mut E) -> T
        where E: CircuitEvaluator<'a, T, R>
    {
        assert_eq!(self.factors.len(), first_inputs.len() + second_inputs.len());
        let result = evaluator.inner_prod(
            self.factors.iter().zip(first_inputs.iter().chain(second_inputs.iter())).filter(|(l, _)| !l.is_zero())
        );
        if self.constant.is_zero() {
            result
        } else {
            evaluator.add_constant(result, &self.constant)
        }
    }

    fn compose<S>(self, input_transforms: &[LinearCombination<R>], ring: S) -> LinearCombination<R>
        where S: RingStore<Type = R> + Copy
    {
        assert_eq!(self.factors.len(), input_transforms.len());
        if input_transforms.len() == 0 {
            return self.clone(ring);
        }
        let new_input_count = input_transforms[0].factors.len();
        assert!(input_transforms.iter().all(|t| t.factors.len() == new_input_count));
        let mut result_factors = (0..new_input_count).map(|_| Coefficient::Zero).collect::<Vec<_>>();
        let mut result_constant = self.constant.clone(ring);
        for (factor, t) in self.factors.into_iter().zip(input_transforms.iter()) {
            for i in 0..new_input_count {
                take_mut::take(&mut result_factors[i], |x| x.add(factor.clone(ring).mul(t.factors[i].clone(ring), ring), ring));
            }
            result_constant = result_constant.add(factor.mul(t.constant.clone(ring), ring), ring);
        }
        return LinearCombination {
            constant: result_constant,
            factors: result_factors
        };
    }
    
    fn change_ring<S, F1, F2>(self, change_summand: &mut F1, change_factor: &mut F2) -> LinearCombination<S>
        where F1: FnMut(Coefficient<R>) -> Coefficient<S>,
            F2: FnMut(Coefficient<R>) -> Coefficient<S>,
            S: ?Sized + RingBase
    {
        LinearCombination {
            constant: change_summand(self.constant),
            factors: self.factors.into_iter().map(|c| change_factor(c)).collect()
        }
    }
    
    fn eq<S: RingStore<Type = R> + Copy>(&self, other: &Self, ring: S) -> bool {
        assert_eq!(self.factors.len(), other.factors.len());
        return self.constant.eq(&other.constant, &ring) &&
            self.factors.iter().zip(other.factors.iter()).all(|(lhs, rhs)| lhs.eq(rhs, &ring));
    }
}

///
/// A nonlinear gate of a [`PlaintextCircuit`].
/// 
enum PlaintextCircuitGate<R: ?Sized + RingBase> {
    Mul(LinearCombination<R>, LinearCombination<R>),
    Square(LinearCombination<R>),
    Gal(Vec<GaloisGroupEl>, LinearCombination<R>)
}

impl<R: ?Sized + RingBase> PlaintextCircuitGate<R> {
    
    fn clone<S: RingStore<Type = R> + Copy>(&self, ring: S) -> Self {
        match self {
            PlaintextCircuitGate::Mul(lhs, rhs) => PlaintextCircuitGate::Mul(lhs.clone(ring), rhs.clone(ring)),
            PlaintextCircuitGate::Gal(gs, t) => PlaintextCircuitGate::Gal(gs.clone(), t.clone(ring)),
            PlaintextCircuitGate::Square(t) => PlaintextCircuitGate::Square(t.clone(ring))
        }
    }
    
    fn eq<S: RingStore<Type = R> + Copy>(&self, other: &Self, ring: S, galois_group: Option<&Subgroup<CyclotomicGaloisGroup>>) -> bool {
        match (self, other) {
            (PlaintextCircuitGate::Mul(self_lhs, self_rhs), PlaintextCircuitGate::Mul(other_lhs, other_rhs)) => self_lhs.eq(other_lhs, ring) && self_rhs.eq(other_rhs, ring),
            (PlaintextCircuitGate::Square(self_t), PlaintextCircuitGate::Square(other_t)) => self_t.eq(other_t, ring),
            (PlaintextCircuitGate::Gal(self_gs, self_t), PlaintextCircuitGate::Gal(other_gs, other_t)) => self_gs.len() == other_gs.len() && self_gs.iter().zip(other_gs).all(|(l, r)| galois_group.expect("a galois group must be given when comparing circuits that contain galois gates").eq_el(l, r)) && self_t.eq(other_t, ring),
            _ => false
        }
    }
}

///
/// Represents an arithmetic circuit (possibly including Galois gates) over
/// a ring of type `R`. 
/// 
/// In a sense, such a circuit can be considered to be a "HE program" that can
/// be run on encrypted inputs from the ring `R`. Of course, using a circuit for
/// HE computations is completely optional, the computations can also be performed
/// by directly operating on ciphertexts. At the very least, explicitly creating
/// a circuit will allow for much simpler testing, since it can also be executed
/// on unencrypted data.
/// 
/// Simple ways to create circuits are by using [`poly_to_circuit()`]
/// and [`crate::lin_transform::matmul::MatmulTransform::matmul1d()`]. However, you 
/// can also manually build a circuit using the functions of [`PlaintextCircuit`], in
/// particular [`PlaintextCircuit::linear_transform()`], [`PlaintextCircuit::select()`],
/// [`PlaintextCircuit::tensor()`] and [`PlaintextCircuit::compose()`].
/// This allows specifying a circuit exactly, but is usually much more complicated than
/// computing a circuit from a linear transform or a set of polynomials.
/// 
/// Note that the ring is not stored by the circuit, but the same ring must be provided 
/// with every circuit operation that requires ring arithmetic. 
/// 
/// [`poly_to_circuit()`]: crate::poly_eval::to_circuit::poly_to_circuit()
/// 
pub struct PlaintextCircuit<R: ?Sized + RingBase> {
    input_count: usize,
    gates: Vec<PlaintextCircuitGate<R>>,
    output_transforms: Vec<LinearCombination<R>>
}

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

    fn check_invariants(&self) {
        let mut current_count = self.input_count;
        for gate in &self.gates {
            match gate {
                PlaintextCircuitGate::Mul(lhs, rhs) => {
                    assert_eq!(current_count, lhs.factors.len());
                    assert_eq!(current_count, rhs.factors.len());
                    current_count += 1;
                },
                PlaintextCircuitGate::Gal(gs, t) => {
                    assert_eq!(current_count, t.factors.len());
                    current_count += gs.len();
                },
                PlaintextCircuitGate::Square(t) => {
                    assert_eq!(current_count, t.factors.len());
                    current_count += 1;
                }
            }
        }
        for out in &self.output_transforms {
            assert_eq!(current_count, out.factors.len());
        }
    }

    pub fn clone<S: RingStore<Type = R> + Copy>(&self, ring: S) -> Self {
        Self {
            gates: self.gates.iter().map(|gate| gate.clone(ring)).collect(),
            input_count: self.input_count,
            output_transforms: self.output_transforms.iter().map(|t| t.clone(ring)).collect()
        }
    }

    pub fn eq<S: RingStore<Type = R> + Copy>(&self, other: &Self, ring: S, galois_group: Option<&Subgroup<CyclotomicGaloisGroup>>) -> bool {
        self.input_count == other.input_count && 
        self.gates.len() == other.gates.len() && 
        self.gates.iter().zip(other.gates.iter()).all(|(l, r)| l.eq(r, ring, galois_group)) && 
        self.output_transforms.len() == other.output_transforms.len() &&
        self.output_transforms.iter().zip(other.output_transforms.iter()).all(|(l, r)| l.eq(r, ring))
    }

    fn computed_wire_count(&self) -> usize {
        self.gates.iter().map(|gate| match gate {
            PlaintextCircuitGate::Mul(_, _) => 1,
            PlaintextCircuitGate::Square(_) => 1,
            PlaintextCircuitGate::Gal(gs, _) => gs.len()
        }).sum()
    }

    ///
    /// Creates the empty circuit, without in-or outputs.
    /// 
    pub fn empty() -> Self {
        Self {
            input_count: 0,
            gates: Vec::new(),
            output_transforms: Vec::new()
        }
    }

    ///
    /// Creates the constant circuit that always outputs the given constant
    /// ```text
    ///  |‾‾‾|
    ///  |___|
    ///    |
    /// ```
    /// 
    pub fn constant_i32<S: RingStore<Type = R>>(c: i32, _ring: S) -> Self {
        let result = Self {
            input_count: 0,
            gates: Vec::new(),
            output_transforms: vec![LinearCombination {
                constant: if c == 0{
                    Coefficient::Zero
                } else if c == 1 {
                    Coefficient::One
                } else {
                    Coefficient::Integer(c)
                },
                factors: Vec::new()
            }]
        };
        result.check_invariants();
        return result;
    }

    /// 
    /// Creates the constant circuit that always outputs the given constant
    /// ```text
    ///  |‾‾‾|
    ///  |___|
    ///    |
    /// ```
    /// 
    pub fn constant<S: RingStore<Type = R>>(el: El<S>, ring: S) -> Self {
        let result = Self {
            input_count: 0,
            gates: Vec::new(),
            output_transforms: vec![LinearCombination {
                constant: Coefficient::from(el, &ring),
                factors: Vec::new()
            }]
        };
        result.check_invariants();
        return result;
    }

    ///
    /// Computes a new circuit, which has the same graph structure as this circuit,
    /// and its constants are derived from this circuit's constants by the given function.
    /// 
    pub fn change_ring<S, F1, F2>(self, mut change_summand: F1, mut change_factor: F2) -> PlaintextCircuit<S>
        where F1: FnMut(Coefficient<R>) -> Coefficient<S>,
            F2: FnMut(Coefficient<R>) -> Coefficient<S>,
            S: ?Sized + RingBase
    {
        PlaintextCircuit {
            input_count: self.input_count,
            gates: self.gates.into_iter().map(|gate| match gate {
                PlaintextCircuitGate::Gal(gs, t) => PlaintextCircuitGate::Gal(gs, t.change_ring(&mut change_summand, &mut change_factor)),
                PlaintextCircuitGate::Mul(l, r) => PlaintextCircuitGate::Mul(l.change_ring(&mut change_summand, &mut change_factor), r.change_ring(&mut change_summand, &mut change_factor)),
                PlaintextCircuitGate::Square(t) => PlaintextCircuitGate::Square(t.change_ring(&mut change_summand, &mut change_factor))
            }).collect(),
            output_transforms: self.output_transforms.into_iter().map(|t| t.change_ring(&mut change_summand, &mut change_factor)).collect()
        }
    }

    ///
    /// Computes a new circuit, which has the same graph structure as this circuit,
    /// and its constants are derived from this circuit's constants by the given function.
    /// 
    pub fn change_ring_uniform<S, F>(self, f: F) -> PlaintextCircuit<S>
        where F: FnMut(Coefficient<R>) -> Coefficient<S>,
            S: ?Sized + RingBase
    {
        let f_refcell = RefCell::new(f);
        return self.change_ring(|x| (f_refcell.borrow_mut())(x), |x| (f_refcell.borrow_mut())(x));
    }

    /// 
    /// Creates the circuit that computes the linear transform of many input elements,
    /// w.r.t. the given list of coefficients.
    /// ```text
    ///        |   |   |   ...
    ///  |‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾|
    ///  | c[0] x0 + c[1] x1 + ... |
    ///  |_________________________|
    ///              |
    /// ```
    /// If you want to pass the list of coefficients as ring elements, consider using
    /// [`PlaintextCircuit::linear_transform_ring()`].
    /// 
    pub fn linear_transform<S: RingStore<Type = R>>(coeffs: &[Coefficient<R>], ring: S) -> Self {
        let result = Self {
            input_count: coeffs.len(),
            gates: Vec::new(),
            output_transforms: vec![LinearCombination {
                constant: Coefficient::Zero,
                factors: coeffs.iter().map(|c| c.clone(&ring)).collect()
            }]
        };
        result.check_invariants();
        return result;
    }

    /// 
    /// Creates the circuit that computes the linear transform of many input elements,
    /// w.r.t. the given list of coefficients.
    /// ```text
    ///        |   |   |   ...
    ///  |‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾|
    ///  | c[0] x0 + c[1] x1 + ... |
    ///  |_________________________|
    ///              |
    /// ```
    /// 
    pub fn linear_transform_ring<S: RingStore<Type = R>>(coeffs: &[El<S>], ring: S) -> Self {
        let result = Self {
            input_count: coeffs.len(),
            gates: Vec::new(),
            output_transforms: vec![LinearCombination {
                constant: Coefficient::Zero,
                factors: coeffs.iter().map(|c| Coefficient::from(ring.clone_el(c), &ring)).collect()
            }]
        };
        result.check_invariants();
        return result;
    }

    /// 
    /// Creates the circuit consisting of a single addition gate
    /// ```text
    ///   | |
    ///  |‾‾‾|
    ///  | + |
    ///  |___|
    ///    |
    /// ```
    /// This is a special case of [`PlaintextCircuit::linear_transform()`], in many cases
    /// the latter is more convenient to use.
    /// 
    pub fn add<S: RingStore<Type = R>>(ring: S) -> Self {
        Self::vec_add(1, ring)
    }

    /// 
    /// Creates the circuit consisting of `len` addition gates, which perform
    /// a component-wise addition of two vectors
    /// ```text
    ///               | | | | ... | |  
    ///  |‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾| 
    ///  | (a, b) -> (a[0], b[0], a[1], b[1], ...) |
    ///  |_________________________________________|
    ///           | |     | |           | |  
    ///          |‾‾‾|   |‾‾‾|         |‾‾‾| 
    ///          | + |   | + |   ...   | + | 
    ///          |___|   |___|         |___| 
    ///            |       |             |   
    /// ```
    /// This is a special case of [`PlaintextCircuit::linear_transform()`], in many cases
    /// the latter is more convenient to use.
    /// 
    pub fn vec_add<S: RingStore<Type = R>>(len: usize, _ring: S) -> Self {
        let result = Self {
            input_count: 2 * len,
            gates: Vec::new(),
            output_transforms: (0..len).map(|i| LinearCombination {
                constant: Coefficient::Zero,
                factors: (0..(2 * len)).map(|j|
                    if (j % len) == i { Coefficient::One } else { Coefficient::Zero }).collect()
            }).collect()
        };
        return result;
    }

    ///
    /// Creates the circuit consisting of a single squaring gate
    /// ```text
    ///     |
    ///  |‾‾‾‾‾|
    ///  |  ^2 |
    ///  |_____|
    ///     |
    /// ```
    /// which is for most purposes equivalent to
    /// 
    /// ```text
    ///    |
    ///   |‾|
    ///  |‾‾‾|
    ///  | * |
    ///  |___|
    ///    |
    /// ```
    /// However, a square gate is stored separately, and may be evaluated
    /// more efficiently during circuit evaluation.
    /// 
    pub fn square<S: RingStore<Type = R>>(_ring: S) -> Self {
        let result = Self {
            input_count: 1,
            gates: vec![PlaintextCircuitGate::Square(LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::One]
            })],
            output_transforms: vec![LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::Zero, Coefficient::One]
            }]
        };
        return result;
    }

    /// 
    /// Creates the circuit consisting of a single subtraction gate
    /// ```text
    ///   | |
    ///  |‾‾‾|
    ///  | - |
    ///  |___|
    ///    |
    /// ```
    /// This is a special case of [`PlaintextCircuit::linear_transform()`], in many cases
    /// the latter is more convenient to use.
    /// 
    pub fn sub<S: RingStore<Type = R>>(_ring: S) -> Self {
        let result = Self {
            input_count: 2,
            gates: Vec::new(),
            output_transforms: vec![LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::One, Coefficient::NegOne]
            }]
        };
        return result;
    }

    /// 
    /// Creates the circuit consisting of a single multiplication gate
    /// ```text
    ///   | |
    ///  |‾‾‾|
    ///  | * |
    ///  |___|
    ///    |
    /// ```
    /// 
    pub fn mul<S: RingStore<Type = R>>(_ring: S) -> Self {
        let result = Self {
            input_count: 2,
            gates: vec![PlaintextCircuitGate::Mul(
                LinearCombination {
                    constant: Coefficient::Zero,
                    factors: vec![Coefficient::One, Coefficient::Zero]
                },
                LinearCombination {
                    constant: Coefficient::Zero,
                    factors: vec![Coefficient::Zero, Coefficient::One]
                }
            )],
            output_transforms: vec![LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::Zero, Coefficient::Zero, Coefficient::One]
            }]
        };
        result.check_invariants();
        return result;
    }

    /// 
    /// Creates the circuit consisting of a single Galois gate
    /// ```text
    ///    |
    ///  |‾‾‾|
    ///  | σ |
    ///  |___|
    ///    |
    /// ```
    /// 
    /// If the given automorphism is the identity, this will just
    /// give a circuit consisting of a single wire.
    /// 
    pub fn gal<S: RingStore<Type = R>>(g: GaloisGroupEl, galois_group: &Subgroup<CyclotomicGaloisGroup>, ring: S) -> Self {
        Self::gal_many(std::slice::from_ref(&g), galois_group, ring)
    }

    /// 
    /// Creates the circuit consisting of a single multiple-Galois gate
    /// ```text
    ///         |
    ///  |‾‾‾‾‾‾‾‾‾‾‾‾‾|
    ///  | σ1, σ2, ... |
    ///  |_____________|
    ///    |   |  ...
    /// ```
    /// 
    /// All occurrences of the identity in the list of Galois automorphisms
    /// are automatically filtered out, and replaced by a direct wire to the
    /// input.
    /// 
    pub fn gal_many<S: RingStore<Type = R>>(gs: &[GaloisGroupEl], galois_group: &Subgroup<CyclotomicGaloisGroup>, _ring: S) -> Self {
        let non_identity_gs = gs.iter().filter(|g| !galois_group.is_identity(*g)).cloned().collect::<Vec<_>>();
        if non_identity_gs.len() == 0 {
            return PlaintextCircuit {
                gates: Vec::new(),
                input_count: 1,
                output_transforms: gs.iter().map(|_| LinearCombination {
                    constant: Coefficient::Zero,
                    factors: vec![Coefficient::One]
                }).collect()
            };
        }
        let len = non_identity_gs.len() + 1;
        let result = Self {
            input_count: 1,
            gates: vec![PlaintextCircuitGate::Gal(
                non_identity_gs, 
                LinearCombination {
                    constant: Coefficient::Zero,
                    factors: vec![Coefficient::One]
                }
            )],
            output_transforms: gs.iter().scan(0, |nonzero_gs, g| if galois_group.is_identity(g) {
                Some(LinearCombination {
                    constant: Coefficient::Zero,
                    factors: [Coefficient::One].into_iter().chain((1..len).map(|_| Coefficient::Zero)).collect()
                })
            } else {
                *nonzero_gs += 1;
                Some(LinearCombination {
                    constant: Coefficient::Zero,
                    factors: (0..len).map(|j| if j == *nonzero_gs { Coefficient::One } else { Coefficient::Zero }).collect()
                })
            }).collect()
        };
        result.check_invariants();
        return result;
    }

    ///
    /// Copies all the output wires of this circuit, i.e. given a circuit
    /// ```text
    ///    | | | |
    ///   |‾‾‾‾‾‾‾|
    ///   |   C1  |
    ///   |_______|
    ///     | | |
    /// ```
    /// this computes
    /// ```text
    ///    | | | |
    ///   |‾‾‾‾‾‾‾|
    ///   |   C1  |
    ///   |_______|
    ///    |  ┊  ┊
    ///    |‾‾┊‾‾┊‾‾|
    ///    |  |‾‾┊‾‾┊‾‾|
    ///    |  |  |‾‾┊‾‾┊‾‾|
    ///    |  |  |  |  |  |
    /// ```
    /// 
    pub fn output_twice<S: RingStore<Type = R> + Copy>(self, ring: S) -> Self {
        self.output_times(2, ring)
    }

    ///
    /// Creates the circuit that drops the given number of wires, i.e. the circuit
    /// ```text
    ///   |  |  |  ...
    ///   ┴  ┴  ┴
    /// ```
    /// which has no output wires.
    /// 
    pub fn drop(wire_count: usize) -> Self {
        let result = Self {
            input_count: wire_count,
            gates: Vec::new(),
            output_transforms: Vec::new()
        };
        result.check_invariants();
        return result;
    }

    ///
    /// Creates the circuit that leaves all wires unchanged, i.e.
    /// ```text
    ///   |  |  |  |  ...
    ///   |  |  |  |  ...
    /// ```
    /// 
    pub fn identity<S: RingStore<Type = R>>(wire_count: usize, _ring: S) -> Self {
        let result = Self {
            input_count: wire_count,
            gates: Vec::new(),
            output_transforms: (0..wire_count).map(|i| LinearCombination {
                constant: Coefficient::Zero,
                factors: (0..wire_count).map(|j| if j == i { Coefficient::One } else { Coefficient::Zero }).collect()
            }).collect()
        };
        result.check_invariants();
        return result;
    }

    ///
    /// Creates the circuit that outputs the input wires at the indices given in `output_wires`.
    /// An input wire can be mentioned multiple times.
    /// 
    pub fn select<S: RingStore<Type = R>>(input_wire_count: usize, output_wires: &[usize], _ring: S) -> Self {
        let result = Self {
            input_count: input_wire_count,
            gates: Vec::new(),
            output_transforms: output_wires.iter().map(|i| {
                assert!(*i < input_wire_count);
                LinearCombination {
                    constant: Coefficient::Zero,
                    factors: (0..input_wire_count).map(|j| if *i == j { Coefficient::One } else { Coefficient::Zero }).collect()
                }
            }).collect()
        };
        result.check_invariants();
        return result;
    }

    pub fn output_times<S: RingStore<Type = R> + Copy>(self, times: usize, ring: S) -> Self {
        let result = Self {
            input_count: self.input_count,
            gates: self.gates.iter().map(|gate| gate.clone(ring)).collect(),
            output_transforms: (0..times).flat_map(|_| self.output_transforms.iter()).map(|lin| lin.clone(ring)).collect()
        };
        result.check_invariants();
        return result;
    }
    
    ///
    /// "Puts" `n` copies of this circuit "next to each other", i.e. given
    /// ```text
    ///    | | | |
    ///   |‾‾‾‾‾‾‾|
    ///   |   C1  |
    ///   |_______|
    ///     | | | 
    /// ```
    /// this function computes the circuit
    /// ```text
    ///    | | | | | | | |       | | | |
    ///   |‾‾‾‾‾‾‾|‾‾‾‾‾‾‾|‾‾‾‾‾|‾‾‾‾‾‾‾|
    ///   |   C1  |   C1  | ... |   C1  |
    ///   |_______|_______|_____|_______|
    ///     | | |   | | |         | | |
    /// ```
    /// 
    pub fn tensor_pow<S: Copy + RingStore<Type = R>>(self, n: usize, ring: S) -> Self {
        (0..n).fold(Self::empty(), |current, _| current.tensor(self.clone(ring), ring))
    }

    ///
    /// "Puts" two circuits "next to each other", i.e. given
    /// ```text
    ///    | | | |                  |
    ///   |‾‾‾‾‾‾‾|             |‾‾‾‾‾‾|
    ///   |   C1  |     and     |  C2  |
    ///   |_______|             |______|
    ///     | | |                 |  |
    /// ```
    /// this function computes the circuit
    /// ```text
    ///    | | | |  |
    ///   |‾‾‾‾‾‾‾|‾‾‾‾‾‾|
    ///   |   C1  |  C2  |
    ///   |_______|______|
    ///      | | |  | |
    /// ```
    /// 
    pub fn tensor<S: RingStore<Type = R>>(self, rhs: Self, ring: S) -> Self {
        let add_zeros = |vec: &[Coefficient<R>], index: usize, count: usize| 
            vec[0..index].iter().map(|c| c.clone(&ring))
                .chain(std::iter::repeat_with(|| Coefficient::Zero).take(count))
                .chain(vec[index..].iter().map(|c| c.clone(&ring)))
                .collect::<Vec<_>>();

        let map_self_transform = |t: &LinearCombination<R>| LinearCombination {
            constant: t.constant.clone(&ring),
            factors: add_zeros(&t.factors, self.input_count, rhs.input_count)
        };
        let map_rhs_transform = |t: &LinearCombination<R>| LinearCombination {
            constant: t.constant.clone(&ring),
            factors: add_zeros(&add_zeros(&t.factors, rhs.input_count, self.computed_wire_count()), 0, self.input_count)
        };
        let result = Self {
            input_count: self.input_count + rhs.input_count,
            gates: self.gates.iter().map(|gate| match gate {
                PlaintextCircuitGate::Mul(lhs, rhs) => PlaintextCircuitGate::Mul(
                    map_self_transform(&lhs),
                    map_self_transform(&rhs)
                ),
                PlaintextCircuitGate::Gal(gs, t) => PlaintextCircuitGate::Gal(
                    gs.clone(), 
                    map_self_transform(t)
                ),
                PlaintextCircuitGate::Square(t) => PlaintextCircuitGate::Square(
                    map_self_transform(t)
                )
            }).chain(
                rhs.gates.iter().map(|gate| match gate {
                    PlaintextCircuitGate::Mul(lhs, rhs) => PlaintextCircuitGate::Mul(
                        map_rhs_transform(&lhs),
                        map_rhs_transform(&rhs)
                    ),
                    PlaintextCircuitGate::Gal(gs, t) => PlaintextCircuitGate::Gal(
                        gs.clone(), 
                        map_rhs_transform(t)
                    ),
                    PlaintextCircuitGate::Square(t) => PlaintextCircuitGate::Square(
                        map_rhs_transform(t)
                    )
                })
            ).collect(),
            output_transforms: self.output_transforms.iter().map(|t| {
                assert_eq!(self.computed_wire_count() + self.input_count, t.factors.len());
                let added_inputs_t = map_self_transform(t);
                LinearCombination {
                    factors: add_zeros(&added_inputs_t.factors, self.input_count + rhs.input_count + self.computed_wire_count(), rhs.computed_wire_count()),
                    constant: added_inputs_t.constant
                }
            }).chain(rhs.output_transforms.iter().map(|t| {
                assert_eq!(rhs.computed_wire_count() + rhs.input_count, t.factors.len());
                map_rhs_transform(t)
            })).collect()
        };
        result.check_invariants();
        return result;
    }

    ///
    /// "Concatentates" two circuits, i.e. connects the output wires of the given circuit
    /// to the inputs of this circuit.
    /// 
    /// In other words, given
    /// ```text
    ///     |   |                  |
    ///   |‾‾‾‾‾‾‾|             |‾‾‾‾‾‾|
    ///   |   C1  |     and     |  C2  |
    ///   |_______|             |______|
    ///     | | |                 |  |
    /// ```
    /// this function computes the circuit
    /// ```text
    ///      |
    ///   |‾‾‾‾‾‾|
    ///   |  C2  |
    ///   |______|
    ///     |  |
    ///   |‾‾‾‾‾‾|
    ///   |  C1  |
    ///   |______|
    ///     | | |  
    /// ```
    /// 
    #[instrument(skip_all)]
    pub fn compose<S: RingStore<Type = R> + Copy>(self, first: Self, ring: S) -> Self {
        assert_eq!(first.output_count(), self.input_count());

        let map_transform = |t: &LinearCombination<R>| {
            let input_transform = LinearCombination {
                constant: t.constant.clone(&ring),
                factors: t.factors[0..self.input_count].iter().map(|c| c.clone(&ring)).collect()
            };
            let mut result = input_transform.compose(&first.output_transforms, ring);
            result.factors.extend(t.factors[self.input_count..].iter().map(|c| c.clone(&ring)));
            return result;
        };
        let result = Self {
            input_count: first.input_count,
            gates: first.gates.iter().map(|gate| gate.clone(ring)).chain(
                self.gates.iter().map(|gate| match gate {
                    PlaintextCircuitGate::Mul(lhs, rhs) => PlaintextCircuitGate::Mul(
                        map_transform(lhs),
                        map_transform(rhs),
                    ),
                    PlaintextCircuitGate::Gal(gs, t) => PlaintextCircuitGate::Gal(
                        gs.clone(),
                        map_transform(t)
                    ),
                    PlaintextCircuitGate::Square(t) => PlaintextCircuitGate::Square(
                        map_transform(t)
                    )
                })
            ).collect(),
            output_transforms: self.output_transforms.iter().map(map_transform).collect()
        };
        result.check_invariants();
        return result;
    }

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

    pub fn output_count(&self) -> usize {
        self.output_transforms.len()
    }
    
    pub fn to_ir<S: RingStore<Type = R> + Copy>(&self, ring: S, galois_group: Option<&Subgroup<CyclotomicGaloisGroup>>) -> Program<<R as ElToIRRing>::ElRepr>
        where R: ElToIRRing
    {
        circuit_to_ir(ring, galois_group, self)
    }

    pub fn from_ir<S: RingStore<Type = R> + Copy>(ring: S, galois_group: Option<&Subgroup<CyclotomicGaloisGroup>>, program: &Program<<R as ElToIRRing>::ElRepr>) -> Self
        where R: ElToIRRing
    {
        ir_to_circuit(ring, galois_group, program)
    }
    
    ///
    /// Evaluates the circuit on inputs of type `T`, which in some sense encrypt/encode/represent
    /// elements of a ring, into which we can also embed the circuit constants.
    /// 
    /// More concretely, the ring whose elements are represented by `T` should support
    /// the following operations:
    ///  - `constant(c)` should return a `T` representing the ring element `c`
    ///  - `add_prod(y, c, x)` should return a `T` representing `y + c * x`
    ///  - `mul(x, y)` should return a `T` representing `x * y`
    ///  - `gal(gs, x)` should return a list of `T`s, with the `i`-th one representing `σ(x)` for `σ: 𝝵 -> 𝝵^gs[i]` 
    ///    the Galois automorphism defined by `gs[i]`  
    /// 
    /// Naturally, if the circuit does not contain multiplication gates (can be checked e.g. via [`PlaintextCircuit::has_multiplication_gates()`]),
    /// `add_prod` will never be called, and similarly for galois gates.
    /// 
    /// Note also that `constant` and `add_prod` are called even if the underlying element is `0` resp. `1`.
    /// Hence, it is recommended that the given functions for `constant` and `add_prod` perform a match on
    /// the given [`Coefficient`] and treat the `0` resp. `1` case differently.
    /// 
    /// Of course, this example could have been more easily implemented using [`PlaintextCircuit::evaluate()`], since
    /// the operations used here exactly match the ones of `StaticRing::<i32>::RING`.
    /// 
    pub fn evaluate_generic<'a, T, E>(&'a self, inputs: &[T], mut evaluator: E) -> Vec<T>
        where E: CircuitEvaluator<'a, T, R>
    {
        assert_eq!(self.input_count, inputs.len());
        assert!(evaluator.supports_gal() || !self.has_galois_gates());
        assert!(evaluator.supports_mul() || !self.has_multiplication_gates());
        let mut current = Vec::new();
        for gate in &self.gates {
            match gate {
                PlaintextCircuitGate::Mul(lhs, rhs) => {
                    let lhs = lhs.evaluate_generic(inputs, &current, &mut evaluator);
                    let rhs = rhs.evaluate_generic(inputs, &current, &mut evaluator);
                    current.push(evaluator.mul(lhs, rhs));
                },
                PlaintextCircuitGate::Gal(gs, t) => {
                    let val = t.evaluate_generic(inputs, &current, &mut evaluator);
                    current.extend(evaluator.gal(val, gs));
                },
                PlaintextCircuitGate::Square(t) => {
                    let val = t.evaluate_generic(inputs, &current, &mut evaluator);
                    current.push(evaluator.square(val));
                }
            }
        }
        return self.output_transforms.iter().map(|t| t.evaluate_generic(inputs, &current, &mut evaluator)).collect()
    }

    ///
    /// Evaluates the given circuit on inputs from a ring into which we can
    /// embed the circuit constants.
    /// 
    /// Panics if the circuit contains galois gates.
    /// 
    pub fn evaluate_no_galois<S, H>(&self, inputs: &[S::Element], hom: H) -> Vec<S::Element>
        where S: ?Sized + RingBase,
            H: Homomorphism<R, S>
    {
        assert!(!self.has_galois_gates());
        return self.evaluate_generic(inputs, HomEvaluator::new(hom));
    }

    ///
    /// Evaluates the given circuit on inputs from a ring into which we can
    /// embed the circuit constants.
    /// 
    /// Note that the circuit might contain Galois gates, thus the given ring
    /// must support evaluation of Galois automorphisms. In case that it doesn't,
    /// you can use [`PlaintextCircuit::evaluate_no_galois()`].
    /// 
    pub fn evaluate<S, H>(&self, inputs: &[S::Element], hom: H) -> Vec<S::Element>
        where S: ?Sized + RingBase + NumberRingQuotient,
            H: Homomorphism<R, S>
    {
        return self.evaluate_generic(inputs, HomEvaluatorGal::new(hom));
    }

    pub fn has_galois_gates(&self) -> bool {
        self.gates.iter().any(|gate| match gate {
            PlaintextCircuitGate::Gal(_, _) => true,
            PlaintextCircuitGate::Mul(_, _) => false,
            PlaintextCircuitGate::Square(_) => false
        })
    }

    pub fn has_multiplication_gates(&self) -> bool {
        self.gates.iter().any(|gate| match gate {
            PlaintextCircuitGate::Gal(_, _) => false,
            PlaintextCircuitGate::Mul(_, _) => true,
            PlaintextCircuitGate::Square(_) => true
        })
    }

    pub fn multiplication_gate_count(&self) -> usize {
        self.gates.iter().filter(|gate| match gate {
            PlaintextCircuitGate::Gal(_, _) => false,
            PlaintextCircuitGate::Mul(_, _) => true,
            PlaintextCircuitGate::Square(_) => true
        }).count()
    }

    pub fn cost(&self, costs: &CircuitEvaluatorCosts) -> f64 {
        self.gates.iter().map(|gate| match gate {
            PlaintextCircuitGate::Gal(gs, _) => if gs.len() == 1 { costs.cost_single_gal } else { costs.cost_setup_hoisted_gal + gs.len() as f64 * costs.cost_hoisted_gal },
            PlaintextCircuitGate::Mul(_, _) => costs.cost_mul,
            PlaintextCircuitGate::Square(_) => costs.cost_sqr
        }).sum()
    }

    ///
    /// Returns the sum of the number of outputs of each Galois gate in the circuit.
    /// 
    pub fn galois_gate_output_sum(&self) -> usize {
        self.gates.iter().map(|gate| match gate {
            PlaintextCircuitGate::Gal(gs, _) => gs.len(),
            PlaintextCircuitGate::Mul(_, _) => 0,
            PlaintextCircuitGate::Square(_) => 0
        }).sum()
    }

    ///
    /// Returns all galois automorphisms which are evaluated by some
    /// gate in this circuit.
    /// 
    /// This directly corresponds to those Galois automorphisms for which
    /// we require a Galois key when evaluating the circuit on encrypted
    /// inputs.
    /// 
    pub fn required_galois_keys(&self, galois_group: &Subgroup<CyclotomicGaloisGroup>) -> Vec<GaloisGroupEl> {
        let mut result = self.gates.iter().flat_map(|gate| match gate {
            PlaintextCircuitGate::Gal(gs, _) => gs.iter().cloned(),
            PlaintextCircuitGate::Mul(_, _) => [].iter().cloned(),
            PlaintextCircuitGate::Square(_) => [].iter().cloned()
        }).collect::<Vec<_>>();
        result.sort_unstable_by_key(|g| galois_group.representative(g));
        result.dedup_by_key(|g| galois_group.representative(g));
        assert!(result.iter().all(|g| galois_group.contains(g)));
        return result;
    }
    
    ///
    /// Returns whether the circuit computes a linear map.
    /// 
    /// This is false if the circuit contains any multiplication gates.
    /// 
    pub fn is_linear(&self) -> bool {
        !self.has_multiplication_gates()
    }

    ///
    /// Returns the multiplicative depth of the `i`-th output, i.e.
    /// the maximal number of multiplication gates on a path from some input
    /// to the given output.
    /// 
    pub fn mul_depth(&self, i: usize) -> usize {
        let mut multiplicative_depths = Vec::new();
        multiplicative_depths.resize(self.input_count(), 0);
        let mult_depth_of_linear_combination = |lin_combination: &LinearCombination<_>, multiplicative_depths: &[usize]| {
            assert_eq!(lin_combination.factors.len(), multiplicative_depths.len());
            lin_combination.factors.iter().zip(multiplicative_depths.iter()).filter(|(factor, _)| !factor.is_zero()).map(|(_, d)| *d).max().unwrap_or(0)
        };
        for gate in &self.gates {
            let (new_depth, count) = match gate {
                PlaintextCircuitGate::Mul(lhs, rhs) => (max(mult_depth_of_linear_combination(lhs, &multiplicative_depths), mult_depth_of_linear_combination(rhs, &multiplicative_depths)) + 1, 1),
                PlaintextCircuitGate::Gal(gs, t) => (mult_depth_of_linear_combination(t, &multiplicative_depths), gs.len()),
                PlaintextCircuitGate::Square(t) => (mult_depth_of_linear_combination(t, &multiplicative_depths) + 1, 1)
            };
            multiplicative_depths.extend((0..count).map(|_| new_depth));
        }
        return mult_depth_of_linear_combination(&self.output_transforms[i], &multiplicative_depths);
    }

    ///
    /// Returns the maximal multiplicative depth of an output, i.e.
    /// the maximal number of multiplication gates on a path from some input
    /// to some output.
    /// 
    pub fn max_mul_depth(&self) -> usize {
        (0..self.output_count()).map(|i| self.mul_depth(i)).max().unwrap_or(0)
    }
}

impl<'a, R> SerializeDeserializeWith<(R, &'a Subgroup<CyclotomicGaloisGroup>)> for PlaintextCircuit<R::Type>
    where R: RingStore + Copy,
        R::Type: SerializableElementRing
{
    fn deserialize_with_data<'de, D: serde::Deserializer<'de>>(data: (R, &'a Subgroup<CyclotomicGaloisGroup>), deserializer: D) -> Result<Self, D::Error> {
        DeserializeSeedPlaintextCircuit {
            galois_group: Some(data.1),
            ring: data.0
        }.deserialize(deserializer)
    }

    fn serialize_with_data<S: serde::Serializer>(&self, data: &(R, &'a Subgroup<CyclotomicGaloisGroup>), serializer: S) -> Result<S::Ok, S::Error> {
        SerializablePlaintextCircuit {
            circuit: self,
            galois_group: Some(data.1),
            ring: data.0
        }.serialize(serializer)
    }
}

impl<R> SerializeDeserializeWith<(R, )> for PlaintextCircuit<R::Type>
    where R: RingStore + Copy,
        R::Type: SerializableElementRing
{
    fn deserialize_with_data<'de, D: serde::Deserializer<'de>>(data: (R, ), deserializer: D) -> Result<Self, D::Error> {
        DeserializeSeedPlaintextCircuit {
            galois_group: None,
            ring: data.0
        }.deserialize(deserializer)
    }

    fn serialize_with_data<S: serde::Serializer>(&self, data: &(R, ), serializer: S) -> Result<S::Ok, S::Error> {
        SerializablePlaintextCircuit {
            circuit: self,
            galois_group: None,
            ring: data.0
        }.serialize(serializer)
    }
}

pub fn create_circuit_cached<R, F, const LOG: bool>(ring: R, keys: &[CachedDataKey], cache_dir: Option<&str>, create: F) -> PlaintextCircuit<R::Type>
    where R: RingStore + Copy,
        R::Type: NumberRingQuotient + SerializableElementRing,
        BaseRing<R>: ZnRing,
        F: FnOnce() -> PlaintextCircuit<R::Type>
{
    create_cached::<_, _, _, LOG>((ring, ring.acting_galois_group()), create, keys, cache_dir, if cache_dir.is_none() { StoreAs::None } else { StoreAs::AlwaysJson })
}

#[cfg(test)]
use feanor_math::assert_el_eq;
#[cfg(test)]
use feanor_math::primitive_int::*;
#[cfg(test)]
use feanor_math::rings::zn::zn_64::Zn;
#[cfg(test)]
use feanor_math::rings::extension::FreeAlgebraStore;
#[cfg(test)]
use crate::number_ring::pow2_cyclotomic::Pow2CyclotomicNumberRing;
#[cfg(test)]
use crate::number_ring::quotient_by_int::NumberRingQuotientByIntBase;
#[cfg(test)]
use crate::ZZi64;

#[test]
fn test_circuit_tensor_compose() {
    let ring = StaticRing::<i64>::RING;
    let x = PlaintextCircuit::linear_transform_ring(&[1], ring);
    let x_sqr = PlaintextCircuit::mul(ring).compose(x.output_twice(ring), ring);
    assert!(PlaintextCircuit {
        input_count: 1,
        gates: vec![PlaintextCircuitGate::Mul(
            LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::One]
            },
            LinearCombination {
                constant: Coefficient::Zero,
                factors: vec![Coefficient::One]
            }
        )],
        output_transforms: vec![LinearCombination {
            constant: Coefficient::Zero,
            factors: vec![Coefficient::Zero, Coefficient::One]
        }]
    }.eq(&x_sqr, ring, None));

    let x = PlaintextCircuit::identity(1, ring);
    let y = PlaintextCircuit::identity(1, ring);
    let x_y_x_y = x.clone(&ring).tensor(y, ring).output_twice(ring);
    // z = 2 * x + 3 * y
    let x_y_z = x.clone(ring).tensor(x.clone(ring), ring).tensor(PlaintextCircuit::linear_transform_ring(&[2, 3], ring), ring).compose(x_y_x_y, ring);
    let xy_z = PlaintextCircuit::mul(ring).tensor(x, ring).compose(x_y_z, ring);
    // w = x * y * (2 * x + 3 * y)
    let w = PlaintextCircuit::mul(ring).compose(xy_z, ring);
    for x in -5..5 {
        for y in -5..5 {
            assert_eq!(x * y * (2 * x + 3 * y), w.evaluate_no_galois(&[x, y], ring.identity()).into_iter().next().unwrap());
        }
    }

    let w_1_sqr = PlaintextCircuit::mul(ring).compose(PlaintextCircuit::add(ring).compose(w.tensor(PlaintextCircuit::constant(1, ring), ring), ring).output_twice(ring), ring);
    for x in -5..5 {
        for y in -5..5 {
            assert_eq!(StaticRing::<i64>::RING.pow(x * y * (2 * x + 3 * y) + 1, 2), w_1_sqr.evaluate_no_galois(&[x, y], ring.identity()).into_iter().next().unwrap());
        }
    }
}

#[test]
fn test_circuit_tensor_compose_with_galois() {
    let number_ring: Pow2CyclotomicNumberRing = Pow2CyclotomicNumberRing::new(16);
    let ring = NumberRingQuotientByIntBase::new(number_ring, Zn::new(17));

    let x = PlaintextCircuit::identity(1, &ring);
    let y = PlaintextCircuit::identity(1, &ring);
    let xy = PlaintextCircuit::mul(&ring).compose(x.tensor(y, &ring), &ring);
    let conj_xy = PlaintextCircuit::gal(ring.acting_galois_group().from_representative(-1), &ring.acting_galois_group(), &ring).compose(xy.clone(&ring), &ring);
    let partial_trace_xy = PlaintextCircuit::add(&ring).compose(xy.tensor(conj_xy, &ring), &ring).compose(PlaintextCircuit::identity(2, &ring).output_twice(&ring), &ring);

    for x_e in 0..8 {
        for y_e in 0..8 {
            let x = ring.pow(ring.canonical_gen(), x_e);
            let y = ring.pow(ring.canonical_gen(), y_e);
            let xy = ring.mul_ref(&x, &y);
            let conj_xy = ring.mul(ring.pow(ring.canonical_gen(), 16 - x_e), ring.pow(ring.canonical_gen(), 16 - y_e));
            assert_el_eq!(
                &ring,
                ring.add(xy, conj_xy),
                partial_trace_xy.evaluate(&[x, y], ring.identity()).into_iter().next().unwrap()
            );
        }
    }
}

#[test]
fn test_giant_step_circuit() {
    let ring = StaticRing::<i64>::RING;
    let powers = PlaintextCircuit::identity(1, ring).tensor(PlaintextCircuit::mul(ring), ring).tensor(PlaintextCircuit::mul(ring), ring).compose(
        PlaintextCircuit::mul(ring).output_times(4, ring).tensor(PlaintextCircuit::identity(1, ring), ring),
        ring
    ).compose(
        PlaintextCircuit::identity(1, ring).output_times(3, ring),
        ring
    );
    assert_eq!(vec![4, 16, 8], powers.evaluate_no_galois(&[2], ring.identity()));

    let permuted_baby_step_dupl_input = PlaintextCircuit::constant(1, ring).tensor(PlaintextCircuit::identity(1, ring), ring).tensor(powers, ring);
    assert_eq!(vec![1, 2, 4, 16, 8], permuted_baby_step_dupl_input.evaluate_no_galois(&[2, 2], ring.identity()));

    let copy_input = PlaintextCircuit::identity(1, ring).output_twice(ring);
    assert_eq!(vec![2, 2], copy_input.evaluate_no_galois(&[2], ring.identity()));

    let permuted_baby_steps = permuted_baby_step_dupl_input.compose(copy_input, ring);
    assert_eq!(vec![1, 2, 4, 16, 8], permuted_baby_steps.evaluate_no_galois(&[2], ring.identity()));

    let baby_steps = PlaintextCircuit::select(5, &[0, 1, 2, 4, 3], ring).compose(permuted_baby_steps, ring);
    assert_eq!(1, baby_steps.input_count());
    assert_eq!(5, baby_steps.output_count());
    assert_eq!(vec![1, 2, 4, 8, 16], baby_steps.evaluate_no_galois(&[2], ring.identity()));

    let giant_steps_before_baby_steps = PlaintextCircuit::constant(1, ring).tensor(PlaintextCircuit::identity(1, ring), ring);
    let baby_and_giant_steps = PlaintextCircuit::identity(4, ring).tensor(giant_steps_before_baby_steps, ring).compose(baby_steps, ring);
    assert_eq!(vec![1, 2, 4, 8, 1, 16], baby_and_giant_steps.evaluate_no_galois(&[2], ring.identity()));
}

#[test]
fn test_serialization() {
    let ring = StaticRing::<i64>::RING;
    let x = PlaintextCircuit::linear_transform_ring(&[1], ring);
    let neg_x = PlaintextCircuit::linear_transform_ring(&[-1], ring);
    let x_neg_x = PlaintextCircuit::mul(ring).compose(x.clone(ring).tensor(neg_x, ring), ring).compose(x.output_twice(ring), ring);
    let two_minus_x_neg_x = PlaintextCircuit::add(ring).compose(x_neg_x.tensor(PlaintextCircuit::constant(2, ring), ring), ring);
    let circuit = PlaintextCircuit::square(ring).compose(two_minus_x_neg_x, ring);

    for x in -100..100 {
        assert_eq!((2 - x * x) * (2 - x * x), circuit.evaluate_no_galois(&[x], ring.identity()).into_iter().next().unwrap());
    }

    let serializer = serde_assert::Serializer::builder().is_human_readable(true).build();
    let tokens = circuit.serialize_with_data(&(ring, ), &serializer).unwrap();
    let mut deserializer = serde_assert::Deserializer::builder(tokens).is_human_readable(true).build();
    let deserialized_circuit = PlaintextCircuit::deserialize_with_data((ring, ), &mut deserializer).unwrap();
    assert!(deserialized_circuit.eq(&circuit, ring, None));

    let serializer = serde_assert::Serializer::builder().is_human_readable(false).build();
    let tokens = circuit.serialize_with_data(&(ring, ), &serializer).unwrap();
    let mut deserializer = serde_assert::Deserializer::builder(tokens).is_human_readable(false).build();
    let deserialized_circuit = PlaintextCircuit::deserialize_with_data((ring, ), &mut deserializer).unwrap();
    assert!(deserialized_circuit.eq(&circuit, ring, None));
}

#[test]
fn test_identity_galois() {
    let Gal = CyclotomicGaloisGroupBase::new(5).into().full_subgroup();
    let ring = StaticRing::<i64>::RING;
    let circuit = PlaintextCircuit::gal(Gal.identity(), &Gal, ring);
    assert!(circuit.eq(&PlaintextCircuit::identity(1, ring), &ring, None));
}

#[test]
fn test_evaluate() {
    let ring = StaticRing::<i64>::RING;
    let x = PlaintextCircuit::linear_transform_ring(&[1], ring);
    let neg_x = PlaintextCircuit::linear_transform_ring(&[-1], ring);
    let x_neg_x = PlaintextCircuit::mul(ring).compose(x.clone(ring).tensor(neg_x, ring), ring).compose(x.output_twice(ring), ring);
    let two_minus_x_neg_x = PlaintextCircuit::add(ring).compose(x_neg_x.tensor(PlaintextCircuit::constant(2, ring), ring), ring);
    let circuit = PlaintextCircuit::square(ring).compose(two_minus_x_neg_x, ring); // (2 - x * x) * (2 - x * x)

    for x in -10..=10 {
        assert_eq!(vec![(2 - x * x) * (2 - x * x)], circuit.evaluate_generic(&[x], HomEvaluator::new(ZZi64.identity())));
    }
}

#[test]
#[ignore]
fn generate_slots_to_coeffs() {
    use feanor_math::integer::int_cast;
    use crate::ZZbig;
    use crate::filename_keys;

    let (chrome_layer, _guard) = tracing_chrome::ChromeLayerBuilder::new().build();
    let filtered_chrome_layer = tracing_subscriber::Layer::with_filter(chrome_layer, tracing_subscriber::filter::filter_fn(|metadata| !["small_basis_to_mult_basis", "mult_basis_to_small_basis", "small_basis_to_coeff_basis", "coeff_basis_to_small_basis"].contains(&metadata.name())));
    tracing_subscriber::util::SubscriberInitExt::init(tracing_subscriber::prelude::__tracing_subscriber_SubscriberExt::with(tracing_subscriber::registry(), filtered_chrome_layer));
    
    use std::cell::LazyCell;
    use std::fs::File;
    use std::io::{BufWriter, Write};
    use crate::lin_transform::pow2;
    use crate::number_ring::hypercube::isomorphism::HypercubeIsomorphism;
    use crate::number_ring::hypercube::structure::HypercubeStructure;
    use crate::number_ring::pow2_cyclotomic::Pow2CyclotomicNumberRing;
    use crate::number_ring::NumberRingQuotientStore;
    use crate::circuit::create_circuit_cached;

    let m = 1 << 14;
    let e = 1;

    let P = NumberRingQuotientByIntBase::new(Pow2CyclotomicNumberRing::new(m), zn_big::Zn::new(ZZbig, ZZbig.pow(int_cast(65537, ZZbig, ZZi64), e + 1)));
    let H = LazyCell::new(|| {
        let hypercube = HypercubeStructure::default_pow2_hypercube(P.acting_galois_group(), int_cast(65537, ZZbig, ZZi64));
        HypercubeIsomorphism::new::<true>(&P, &hypercube, Some("."))
    });
    let coeffs_to_slots = create_circuit_cached::<_, _, true>(
        &P, 
        &filename_keys![coeffs2slots, m: m, p: 65537, e: e + 1, levels: 4], 
        Some("."), 
        || pow2::coeffs_to_slots_thin(&H, 4)
    );
    let program = coeffs_to_slots.to_ir(&P, Some(P.acting_galois_group()));
    write!(BufWriter::new(File::create(format!("./coeffs_to_slots_m{}_p65537_e{}_levels4.fheir", m, e + 1).as_str()).unwrap()), "{}", program).unwrap();
    assert!(coeffs_to_slots.eq(&PlaintextCircuit::from_ir(&P, Some(P.acting_galois_group()), &program), &P, Some(P.acting_galois_group())));

    let P = NumberRingQuotientByIntBase::new(Pow2CyclotomicNumberRing::new(m), zn_big::Zn::new(ZZbig, ZZbig.pow(int_cast(65537, ZZbig, ZZi64), e)));
    let H = LazyCell::new(|| H.change_modulus(&P));
    let slots_to_coeffs = create_circuit_cached::<_, _, true>(
        &P, 
        &filename_keys![slots2coeffs, m: m, p: 65537, e: e, levels: 4], 
        Some("."), 
        || pow2::slots_to_coeffs_thin(&H, 4)
    );
    let program = slots_to_coeffs.to_ir(&P, Some(P.acting_galois_group()));
    write!(BufWriter::new(File::create(format!("./slots_to_coeffs_m{}_p65537_e{}_levels4.fheir", m, e).as_str()).unwrap()), "{}", program).unwrap();
    assert!(slots_to_coeffs.eq(&PlaintextCircuit::from_ir(&P, Some(P.acting_galois_group()), &program), &P, Some(P.acting_galois_group())));
}