use super::field::{AsField, AsFieldBase};
use super::poly::dense_poly::DensePolyRing;
use super::poly::{PolyRing, PolyRingStore};
use crate::algorithms::extension_ops;
use crate::algorithms::linsolve::LinSolveRing;
use crate::algorithms::poly_factor::FactorPolyField;
use crate::divisibility::DivisibilityRing;
use crate::field::*;
use crate::homomorphism::*;
use crate::pid::PrincipalIdealRing;
use crate::ring::*;
use crate::seq::*;
use crate::wrapper::RingElementWrapper;

/// Contains [`extension_impl::FreeAlgebraImpl`], an implementation of [`FreeAlgebra`] based
/// on polynomial division.
pub mod extension_impl;

/// Contains [`galois_field::GaloisField`], an implementation of Galois fields.
pub mod galois_field;

/// Contains [`number_field::NumberField`], an implementation of number fields.
pub mod number_field;

/// A table of Conway polynomials, for standardized creation of finite fields.
pub mod conway;

/// A ring `R` that is an extension of a base ring `S`, generated by a single element
/// that is algebraic resp. integral over `S`.
///
/// This is equivalent to rings generated by a single element that is a zero of a monic polynomial
/// over the base ring. While sounding quite technical, this includes a wide class of important
/// rings, like number fields or galois fields.
/// One consequence of this is that `R` is a free `S`-module, with a basis given by the powers
/// of [`FreeAlgebra::canonical_gen()`], which is where the name "free" comes from.
///
/// The main implementation is [`extension_impl::FreeAlgebraImpl`].
///
/// # Nontrivial Automorphisms
///
/// Rings of this form very often have nontrivial automorphisms. In order to simplify situations
/// where morphisms or other objects are only unique up to isomorphism, canonical morphisms between
/// rings of this type must also preserve the canonical generator.
///
/// # Examples
///
/// One of the most common use cases seems to be the implementation of finite fields (sometimes
/// called galois fields).
/// ```rust
/// #![feature(allocator_api)]
/// # use std::alloc::Global;
/// # use feanor_math::assert_el_eq;
/// # use feanor_math::ring::*;
/// # use feanor_math::rings::zn::*;
/// # use feanor_math::primitive_int::*;
/// # use feanor_math::rings::extension::*;
/// # use feanor_math::rings::extension::extension_impl::*;
/// # use feanor_math::algorithms::convolution::*;
/// # use feanor_math::field::FieldStore;
/// # use feanor_math::divisibility::*;
/// # use feanor_math::rings::finite::*;
/// // we have to decide for an implementation of the prime field
/// let prime_field = zn_static::Fp::<3>::RING;
/// let galois_field = FreeAlgebraImpl::new(prime_field, 3, [2, 1]);
/// // this is now the finite field with 27 elements, or F_27 or GF(27) since X^3 + 2X + 1 is irreducible modulo 3
/// let galois_field = galois_field.as_field().ok().unwrap();
/// assert_eq!(Some(27), galois_field.size(&StaticRing::<i64>::RING));
/// for x in galois_field.elements() {
///     if !galois_field.is_zero(&x) {
///         let inv_x = galois_field.div(&galois_field.one(), &x);
///         assert_el_eq!(galois_field, galois_field.one(), galois_field.mul(x, inv_x));
///     }
/// }
/// // since galois fields are so important, an efficient construction is provided by feanor-math
/// let galois_field_2 = galois_field::GaloisField::new_with_convolution(prime_field, 3, Global, STANDARD_CONVOLUTION);
/// // note that the generating polynomial might be different, so it is not necessarily the "same" ring
/// assert!(galois_field_2.can_iso(&galois_field).is_none());
/// ```
pub trait FreeAlgebra: RingExtension {
    /// Type of the canonical-basis representation of a ring element, as returned by
    /// [`FreeAlgebra::wrt_canonical_basis()`].
    type VectorRepresentation<'a>: VectorFn<El<Self::BaseRing>>
    where
        Self: 'a;

    /// Returns the fixed element that generates this ring as a free module over the base ring.
    fn canonical_gen(&self) -> Self::Element;

    /// Returns the rank of this ring as a free module over the base ring.
    fn rank(&self) -> usize;

    /// Returns the representation of the element w.r.t. the canonical basis, that is the basis
    /// given by the powers `x^i` where `x` is the canonical generator given by
    /// [`FreeAlgebra::canonical_gen()`] and `i` goes from `0` to `rank - 1`.
    ///
    /// In this sense, this is the opposite function to [`FreeAlgebra::from_canonical_basis()`].
    fn wrt_canonical_basis<'a>(&'a self, el: &'a Self::Element) -> Self::VectorRepresentation<'a>;

    /// Returns the element that has the given representation w.r.t. the canonical basis, that is
    /// the basis given by the powers `x^i` where `x` is the canonical generator given by
    /// [`FreeAlgebra::canonical_gen()`] and `i` goes from `0` to `rank - 1`.
    ///
    /// In this sense, this is the opposite function to [`FreeAlgebra::wrt_canonical_basis()`].
    fn from_canonical_basis<V>(&self, vec: V) -> Self::Element
    where
        V: IntoIterator<Item = El<Self::BaseRing>>,
        V::IntoIter: DoubleEndedIterator,
    {
        let mut given_len = 0;
        let x = self.canonical_gen();
        let mut result = self.zero();
        for c in vec.into_iter().rev() {
            self.mul_assign_ref(&mut result, &x);
            self.add_assign(&mut result, self.from(c));
            given_len += 1;
        }
        assert_eq!(given_len, self.rank());
        return result;
    }

    /// Multiplies the given element by the `power`-th power of the canonical generator
    /// of this ring, as given by [`FreeAlgebra::canonical_gen()`].
    fn mul_assign_gen_power(&self, el: &mut Self::Element, power: usize) {
        self.mul_assign(el, RingRef::new(self).pow(self.canonical_gen(), power));
    }

    /// Like [`FreeAlgebra::from_canonical_basis()`], this computes the sum `sum_i vec[i] * x^i`
    /// where `x` is the canonical generator given by [`FreeAlgebra::canonical_gen()`]. Unlike
    /// [`FreeAlgebra::from_canonical_basis()`], `vec` can return any number elements.
    fn from_canonical_basis_extended<V>(&self, vec: V) -> Self::Element
    where
        V: IntoIterator<Item = El<Self::BaseRing>>,
    {
        extension_ops::from_canonical_basis_extended(self, vec)
    }

    /// Computes the characteristic polynomial of the given element.
    ///
    /// The characteristic polynomial is `det(XI - B)`, where `B` is the
    /// matrix representation of the given element `b`, or equivalently the
    /// matrix representation of the multiplication-by-`b` map `R^n -> R^n`.
    fn charpoly<P, H>(&self, el: &Self::Element, poly_ring: P, hom: H) -> El<P>
    where
        P: RingStore,
        P::Type: PolyRing,
        <<P::Type as RingExtension>::BaseRing as RingStore>::Type: LinSolveRing,
        H: Homomorphism<<Self::BaseRing as RingStore>::Type, <<P::Type as RingExtension>::BaseRing as RingStore>::Type>,
    {
        extension_ops::charpoly(self, el, poly_ring, hom)
    }

    /// Computes the (or a) minimal polynomial of the given element.
    ///
    /// The minimal polynomial is the monic polynomial of minimal degree that
    /// has the given value as a root. Its degree is always at least 1 and at
    /// most [`FreeAlgebra::rank()`]. If the base ring is a principal ideal domain,
    /// then the minimal polynomial is unique.
    ///
    /// Note that the existence of the minimal polynomial is a consequence of the
    /// Cayley-Hamilton theorem, i.e. `det(XI - A)(A) = 0` for a square matrix over
    /// any ring `R`. In particular, representing element of the ring `R[a]` a matrices
    /// over `R`, we find that `det(XI - B)(b) = 0` for any element `b`, thus `b` must
    /// be the root of a monic polynomial of degree `<= n`.
    fn minpoly<P, H>(&self, el: &Self::Element, poly_ring: P, hom: H) -> El<P>
    where
        P: RingStore,
        P::Type: PolyRing,
        <<P::Type as RingExtension>::BaseRing as RingStore>::Type: LinSolveRing,
        H: Homomorphism<<Self::BaseRing as RingStore>::Type, <<P::Type as RingExtension>::BaseRing as RingStore>::Type>,
    {
        extension_ops::minpoly(self, el, poly_ring, hom)
    }

    /// Computes the trace of an element `a` in this ring extension, which is defined as the
    /// matrix trace of the multiplication-by-`a` map.
    ///
    /// In nice extensions, the trace has many characterizations. For example, in a Galois
    /// field extension, it is the sum of `sigma(a)` as `sigma` runs through the Galois group
    /// of the extension. It is also equal to +/- the second largest coefficient of the
    /// characteristic polynomial of the element.
    fn trace(&self, el: Self::Element) -> El<Self::BaseRing> {
        let mut current = el;
        let generator = self.canonical_gen();
        return self.base_ring().sum((0..self.rank()).map(|i| {
            let result = self.wrt_canonical_basis(&current).at(i);
            self.mul_assign_ref(&mut current, &generator);
            return result;
        }));
    }

    /// Computes the discriminant of the canonical basis of this ring extension,
    /// which is defined as the determinant of the trace matrix `(Tr(a^(i + j)))`,
    /// where `a` is the canonical generator of this ring extension.
    ///
    /// Note that the discriminant in this sense depends on the choice of the canonical
    /// generator. In particular, this is usually different from the discriminant of an
    /// algebraic number field, since that is defined as the discriminant w.r.t. the basis
    /// that generates the maximal order. On the other hand, this means that if this ring
    /// is an order in number field, this discriminant is exactly the discriminant of the
    /// order as considered in algebraic number theory.
    fn discriminant(&self) -> El<Self::BaseRing>
    where
        <Self::BaseRing as RingStore>::Type: PrincipalIdealRing,
    {
        extension_ops::discriminant(self)
    }
}

/// [`RingStore`] for [`FreeAlgebra`].
pub trait FreeAlgebraStore: RingStore
where
    Self::Type: FreeAlgebra,
{
    delegate! { FreeAlgebra, fn canonical_gen(&self) -> El<Self> }
    delegate! { FreeAlgebra, fn rank(&self) -> usize }
    delegate! { FreeAlgebra, fn trace(&self, el: El<Self>) -> El<<Self::Type as RingExtension>::BaseRing> }
    delegate! { FreeAlgebra, fn mul_assign_gen_power(&self, el: &mut El<Self>, power: usize) -> () }

    /// See [`FreeAlgebra::wrt_canonical_basis()`].
    fn wrt_canonical_basis<'a>(&'a self, el: &'a El<Self>) -> <Self::Type as FreeAlgebra>::VectorRepresentation<'a> {
        self.get_ring().wrt_canonical_basis(el)
    }

    /// See [`FreeAlgebra::from_canonical_basis()`].
    fn from_canonical_basis<V>(&self, vec: V) -> El<Self>
    where
        V: IntoIterator<Item = El<<Self::Type as RingExtension>::BaseRing>>,
        V::IntoIter: DoubleEndedIterator,
    {
        self.get_ring().from_canonical_basis(vec)
    }

    /// See [`FreeAlgebra::from_canonical_basis_extended()`].
    fn from_canonical_basis_extended<V>(&self, vec: V) -> El<Self>
    where
        V: IntoIterator<Item = El<<Self::Type as RingExtension>::BaseRing>>,
    {
        self.get_ring().from_canonical_basis_extended(vec)
    }

    /// Returns the generating polynomial of this ring, i.e. the monic polynomial `f(X)` such that
    /// this ring is isomorphic to `R[X]/(f(X))`, where `R` is the base ring.
    fn generating_poly<P, H>(&self, poly_ring: P, hom: H) -> El<P>
    where
        P: PolyRingStore,
        P::Type: PolyRing,
        H: Homomorphism<
                <<Self::Type as RingExtension>::BaseRing as RingStore>::Type,
                <<P::Type as RingExtension>::BaseRing as RingStore>::Type,
            >,
    {
        assert!(hom.domain().get_ring() == self.base_ring().get_ring());
        poly_ring.sub(
            poly_ring.from_terms([(poly_ring.base_ring().one(), self.rank())]),
            self.poly_repr(&poly_ring, &self.pow(self.canonical_gen(), self.rank()), hom),
        )
    }

    /// If this ring is a field, returns a wrapper around this ring that implements
    /// [`crate::field::FieldStore`].
    ///
    /// For details, see [`crate::rings::field::AsField`].
    fn as_field(self) -> Result<AsField<Self>, Self>
    where
        Self::Type: DivisibilityRing,
        <<Self::Type as RingExtension>::BaseRing as RingStore>::Type: Field + FactorPolyField,
    {
        let poly_ring = DensePolyRing::new(self.base_ring(), "X");
        if <_ as FactorPolyField>::factor_poly(
            &poly_ring,
            &self.generating_poly(&poly_ring, self.base_ring().identity()),
        )
        .0
        .len()
            > 1
        {
            return Err(self);
        } else {
            return Ok(RingValue::from(AsFieldBase::promise_is_perfect_field(self)));
        }
    }

    /// Returns the polynomial representation of the given element `y`, i.e. the polynomial `f(X)`
    /// of degree at most [`FreeAlgebraStore::rank()`] such that `f(x) = y`, where `y` is the
    /// canonical generator of this ring, as given by [`FreeAlgebraStore::canonical_gen()`].
    fn poly_repr<P, H>(&self, to: P, el: &El<Self>, hom: H) -> El<P>
    where
        P: PolyRingStore,
        P::Type: PolyRing,
        H: Homomorphism<
                <<Self::Type as RingExtension>::BaseRing as RingStore>::Type,
                <<P::Type as RingExtension>::BaseRing as RingStore>::Type,
            >,
    {
        let coeff_vec = self.wrt_canonical_basis(el);
        to.from_terms(
            (0..self.rank())
                .map(|i| coeff_vec.at(i))
                .enumerate()
                .filter(|(_, x)| !self.base_ring().is_zero(x))
                .map(|(j, x)| (hom.map(x), j)),
        )
    }

    /// Computes the discriminant of the canonical basis of this ring extension,
    /// which is defined as the determinant of the trace matrix `(Tr(a^(i + j)))`,
    /// where `a` is the canonical generator of this ring extension.
    ///
    /// See also [`FreeAlgebra::discriminant()`].
    fn discriminant(&self) -> El<<Self::Type as RingExtension>::BaseRing>
    where
        <<Self::Type as RingExtension>::BaseRing as RingStore>::Type: PrincipalIdealRing,
    {
        self.get_ring().discriminant()
    }

    /// See also [`FreeAlgebra::charpoly()`].
    fn charpoly<P, H>(&self, el: &El<Self>, poly_ring: P, hom: H) -> El<P>
    where
        P: RingStore,
        P::Type: PolyRing,
        <<P::Type as RingExtension>::BaseRing as RingStore>::Type: LinSolveRing,
        H: Homomorphism<
                <<Self::Type as RingExtension>::BaseRing as RingStore>::Type,
                <<P::Type as RingExtension>::BaseRing as RingStore>::Type,
            >,
    {
        self.get_ring().charpoly(el, poly_ring, hom)
    }

    /// Temporarily wraps the canonical generator in a [`RingElementWrapper`], for more
    /// natural creation of ring elements.
    #[stability::unstable(feature = "enable")]
    fn with_wrapped_generator<'a, F, const M: usize>(&'a self, f: F) -> [El<Self>; M]
    where
        F: FnOnce(&RingElementWrapper<&'a Self>) -> [RingElementWrapper<&'a Self>; M],
    {
        let wrapped_indet = RingElementWrapper::new(self, self.canonical_gen());
        let mut result_it = f(&wrapped_indet).into_iter();
        return std::array::from_fn(|_| result_it.next().unwrap().unwrap());
    }
}

impl<R: RingStore> FreeAlgebraStore for R where R::Type: FreeAlgebra {}

#[stability::unstable(feature = "enable")]
pub struct FreeAlgebraHom<R, S>
where
    R: RingStore,
    R::Type: FreeAlgebra,
    S: RingStore,
    S::Type: FreeAlgebra,
    <S::Type as RingExtension>::BaseRing: RingStore<Type = <<R::Type as RingExtension>::BaseRing as RingStore>::Type>,
{
    from: R,
    to: S,
    image_of_generator: El<S>,
}

impl<R> FreeAlgebraHom<R, R>
where
    R: RingStore + Clone,
    R::Type: FreeAlgebra,
{
    #[stability::unstable(feature = "enable")]
    pub fn identity(ring: R) -> Self {
        Self {
            image_of_generator: ring.canonical_gen(),
            from: ring.clone(),
            to: ring,
        }
    }
}

impl<R, S> FreeAlgebraHom<R, S>
where
    R: RingStore,
    R::Type: FreeAlgebra,
    S: RingStore,
    S::Type: FreeAlgebra,
    <S::Type as RingExtension>::BaseRing: RingStore<Type = <<R::Type as RingExtension>::BaseRing as RingStore>::Type>,
{
    /// Creates a new [`FreeAlgebraHom`] from `R` to `S`, mapping the canonical
    /// generator of `R` to the given element of `S`. This assumes that the resulting
    /// homomorphism is well-defined, i.e. the generating polynomial of `R` evaluated
    /// at `image_of_generator` gives zero in `S`.
    ///
    /// The checked variant of this function is [`FreeAlgebraHom::new()`].
    #[stability::unstable(feature = "enable")]
    pub fn promise_is_well_defined(from: R, to: S, image_of_generator: El<S>) -> Self {
        assert!(from.base_ring().get_ring() == to.base_ring().get_ring());
        Self {
            from,
            to,
            image_of_generator,
        }
    }

    /// Creates a new [`FreeAlgebraHom`] from `R` to `S`, mapping the canonical
    /// generator of `R` to the given element of `S`.
    ///
    /// As opposed to [`FreeAlgebraHom::promise_is_well_defined()`], this function
    /// checks that the resulting homomorphism is well-defined.
    #[stability::unstable(feature = "enable")]
    pub fn new(from: R, to: S, image_of_generator: El<S>) -> Self {
        assert!(from.base_ring().get_ring() == to.base_ring().get_ring());
        let poly_ring = DensePolyRing::new(to.base_ring(), "X");
        assert!(to.is_zero(&poly_ring.evaluate(
            &from.generating_poly(&poly_ring, poly_ring.base_ring().identity()),
            &image_of_generator,
            to.inclusion()
        )));
        Self {
            from,
            to,
            image_of_generator,
        }
    }

    /// Consumes this object, producing the domain ring store, the codomain ring store
    /// and the image of the canonical generator of the domain number field.
    #[stability::unstable(feature = "enable")]
    pub fn destruct(self) -> (R, S, El<S>) { (self.from, self.to, self.image_of_generator) }
}

impl<R, S> Homomorphism<R::Type, S::Type> for FreeAlgebraHom<R, S>
where
    R: RingStore,
    R::Type: FreeAlgebra,
    S: RingStore,
    S::Type: FreeAlgebra,
    <S::Type as RingExtension>::BaseRing: RingStore<Type = <<R::Type as RingExtension>::BaseRing as RingStore>::Type>,
{
    type DomainStore = R;
    type CodomainStore = S;

    fn domain(&self) -> &Self::DomainStore { &self.from }

    fn codomain(&self) -> &Self::CodomainStore { &self.to }

    fn map(&self, x: El<R>) -> El<S> { self.map_ref(&x) }

    fn map_ref(&self, x: &El<R>) -> El<S> {
        let poly_ring = DensePolyRing::new(self.to.base_ring(), "X");
        return poly_ring.evaluate(
            &self.from.poly_repr(&poly_ring, x, self.to.base_ring().identity()),
            &self.image_of_generator,
            self.to.inclusion(),
        );
    }
}

#[cfg(any(test, feature = "generic_tests"))]
pub mod generic_tests {
    use super::*;

    pub fn test_free_algebra_axioms<R: FreeAlgebraStore>(ring: R)
    where
        R::Type: FreeAlgebra,
    {
        let x = ring.canonical_gen();
        let n = ring.rank();

        let xn_original = ring.pow(ring.clone_el(&x), n);
        let xn_vec = ring.wrt_canonical_basis(&xn_original);
        let xn = ring.sum(Iterator::map(0..n, |i| {
            ring.mul(ring.inclusion().map(xn_vec.at(i)), ring.pow(ring.clone_el(&x), i))
        }));
        assert_el_eq!(ring, xn_original, xn);

        let x_n_1_vec_expected = (0..n).map_fn(|i| {
            if i > 0 {
                ring.base_ring()
                    .add(ring.base_ring().mul(xn_vec.at(n - 1), xn_vec.at(i)), xn_vec.at(i - 1))
            } else {
                ring.base_ring().mul(xn_vec.at(n - 1), xn_vec.at(0))
            }
        });
        let x_n_1 = ring.pow(ring.clone_el(&x), n + 1);
        let x_n_1_vec_actual = ring.wrt_canonical_basis(&x_n_1);
        for i in 0..n {
            assert_el_eq!(ring.base_ring(), x_n_1_vec_expected.at(i), x_n_1_vec_actual.at(i));
        }

        // test basis wrt_canonical_basis linearity and compatibility
        // from_canonical_basis/wrt_canonical_basis
        for i in (0..ring.rank()).step_by(3) {
            for j in (1..ring.rank()).step_by(5) {
                if i == j {
                    continue;
                }
                let element = ring.from_canonical_basis((0..n).map(|k| {
                    if k == i {
                        ring.base_ring().one()
                    } else if k == j {
                        ring.base_ring().int_hom().map(2)
                    } else {
                        ring.base_ring().zero()
                    }
                }));
                let expected = ring.add(
                    ring.pow(ring.clone_el(&x), i),
                    ring.int_hom().mul_map(ring.pow(ring.clone_el(&x), j), 2),
                );
                assert_el_eq!(ring, expected, element);
                let element_vec = ring.wrt_canonical_basis(&expected);
                for k in 0..ring.rank() {
                    if k == i {
                        assert_el_eq!(ring.base_ring(), ring.base_ring().one(), element_vec.at(k));
                    } else if k == j {
                        assert_el_eq!(ring.base_ring(), ring.base_ring().int_hom().map(2), element_vec.at(k));
                    } else {
                        assert_el_eq!(ring.base_ring(), ring.base_ring().zero(), element_vec.at(k));
                    }
                }
            }
        }

        // test basis mul_assign_gen_power
        for i in (0..ring.rank()).step_by(3) {
            for j in (0..ring.rank()).step_by(5) {
                let element = ring.add(ring.pow(ring.canonical_gen(), i), ring.one());
                let mut actual = ring.clone_el(&element);
                ring.mul_assign_gen_power(&mut actual, j);
                assert_el_eq!(ring, ring.mul(element, ring.pow(ring.canonical_gen(), j)), actual);
            }
        }
    }
}
