"""
p-Adic Extension Element

A common superclass for all elements of extension rings and field of `\ZZ_p` and
`\QQ_p`.

AUTHORS:

- David Roe (2007): initial version

- Julian Rueth (2012-10-18): added residue
"""

#*****************************************************************************
#       Copyright (C) 2007-2010 David Roe <roed.math@gmail.com>
#                     2012 Julian Rueth <julian.rueth@fsfe.org>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#  as published by the Free Software Foundation; either version 2 of
#  the License, or (at your option) any later version.
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************

from sage.rings.padics.pow_computer cimport PowComputer_class
from sage.rings.integer import Integer
from sage.libs.ntl.ntl_ZZ_p cimport ntl_ZZ_p

cdef class pAdicExtElement(pAdicGenericElement):
    cdef int _set_from_list(self, L) except -1:
        """
        Sets self from a list.

        The list should either be uniform in type, or all of the entries should be coercible to integers.
        If any of the entries in L is a list, L will be cast to a ZZ_pEX

        INPUT:
        L -- a list.
        """
        raise NotImplementedError

    cdef int _set_from_list_rel(self, L, long relprec) except -1:
        raise NotImplementedError

    cdef int _set_from_list_abs(self, L, long absprec) except -1:
        raise NotImplementedError

    cdef int _set_from_list_both(self, L, long absprec, long relprec) except -1:
        raise NotImplementedError

    cdef int _set_from_ZZX(self, ZZX_c poly) except -1:
        """
        Sets from a ZZX_c, choosing how to handle based on the
        precision type of self.parent().

        Fixed modulus elements should override this function.

        This function is not used internally.
        """
        if self.parent().is_capped_relative():
            self._set_from_ZZX_rel(poly, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        elif self.parent().is_capped_absolute():
            self._set_from_ZZX_abs(poly, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        else:
            raise RuntimeError("_set_from_ZZX should have been overridden")

    cdef int _set_from_ZZX_rel(self, ZZX_c poly, long relprec) except -1:
        """
        Set from a ZZX_c with bounded relative precision.

        Capped relative elements should override this function, so the
        default implementation is for capped absolute.

        This function is not used internally.
        """
        self._set_from_ZZX_both(poly, (<PowComputer_class>self.parent().prime_pow).prec_cap, relprec)

    cdef int _set_from_ZZX_abs(self, ZZX_c poly, long absprec) except -1:
        """
        Set from a ZZX_c with bounded absolute precision.

        Capped absolute elements should override this function, so the
        default implementation is for capped relative.

        This function is not used internally.
        """
        self._set_from_ZZX_both(poly, absprec, (<PowComputer_class>self.parent().prime_pow).prec_cap)

    cdef int _set_from_ZZX_both(self, ZZX_c poly, long absprec, long relprec) except -1:
        """
        Set from a ZZX_c with both absolute and relative precisions bounded.

        This function should be overridden for both capped absolute
        and capped relative elements.

        This function is not used internally.
        """
        if self.parent().is_fixed_mod():
            self._set_from_ZZX(poly)
        else:
            raise RuntimeError("_set_from_ZZX_both should have been overridden")

    cdef int _set_from_ZZ_pX(self, ZZ_pX_c* poly, ntl_ZZ_pContext_class ctx) except -1:
        """
        Sets self from a ZZ_pX defined with context ctx.

        This function should be overridden for fixed modulus elements.

        This function is not used internally.
        """
        if self.parent().is_capped_relative():
            self._set_from_ZZ_pX_rel(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        elif self.parent().is_capped_absolute():
            self._set_from_ZZ_pX_abs(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        else:
            raise RuntimeError("_set_from_ZZ_pX should have been overridden")

    cdef int _set_from_ZZ_pX_rel(self, ZZ_pX_c* poly, ntl_ZZ_pContext_class ctx, long relprec) except -1:
        """
        Set from a ZZ_pX_c with bounded relative precision.

        Capped relative rings should override this function, so the
        default implementation is for capped absolute.

        This function is not used internally.
        """
        self._set_from_ZZ_pX_both(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap, relprec)

    cdef int _set_from_ZZ_pX_abs(self, ZZ_pX_c* poly, ntl_ZZ_pContext_class ctx, long absprec) except -1:
        """
        Set from a ZZ_pX_c with bounded absolute precision.

        Capped absolute rings should override this function, so the
        default implementation is for capped relative.

        This function is not used internally.
        """
        self._set_from_ZZ_pX_both(poly, ctx, absprec, (<PowComputer_class>self.parent().prime_pow).prec_cap)

    cdef int _set_from_ZZ_pX_both(self, ZZ_pX_c* poly, ntl_ZZ_pContext_class ctx, long absprec, long relprec) except -1:
        """
        Set from a ZZ_pX_c with both absolute and relative precision bounded.

        This function should be overridden by both capped absolute and capped relative elements.

        This function is not used internally.
        """
        if self.parent().is_fixed_mod():
            self._set_from_ZZ_pX(poly, ctx)
        else:
            raise RuntimeError("_set_from_ZZ_pX_both should have been overridden")

    cdef int _set_from_ZZ_pE(self, ZZ_pE_c* poly, ntl_ZZ_pEContext_class ctx) except -1:
        """
        Set from a ZZ_pE_c.

        This function is not used internally.
        """
        if self.parent().is_capped_relative():
            self._set_from_ZZ_pE_rel(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        elif self.parent().is_capped_absolute():
            self._set_from_ZZ_pE_abs(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        else:
            raise RuntimeError("_set_from_ZZ_pE should have been overridden")

    cdef int _set_from_ZZ_pE_rel(self, ZZ_pE_c* poly, ntl_ZZ_pEContext_class ctx, long relprec) except -1:
        """
        Set from a ZZ_pE_c with bounded relative precision.

        Capped relative rings should override this function, so the
        default implementation is for capped absolute.

        This function is not used internally.
        """
        self._set_from_ZZ_pE_both(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap, relprec)

    cdef int _set_from_ZZ_pE_abs(self, ZZ_pE_c* poly, ntl_ZZ_pEContext_class ctx, long absprec) except -1:
        """
        Set from a ZZ_pE_c with bounded absolute precision.

        Capped absolute elements should override this function, so the
        default implementation is for capped relative.

        This function is not used internally.
        """
        self._set_from_ZZ_pE_both(poly, ctx, absprec, (<PowComputer_class>self.parent().prime_pow).prec_cap)

    cdef int _set_from_ZZ_pE_both(self, ZZ_pE_c* poly, ntl_ZZ_pEContext_class ctx, long absprec, long relprec) except -1:
        """
        Sets from a ZZ_pE_c with both absolute and relative precision bounded.

        Capped absolute and capped relative elements should override
        this function.

        This function is not used internally.
        """
        if self.parent().is_fixed_mod():
            self._set_from_ZZ_pE(poly, ctx)
        else:
            raise RuntimeError("_set_from_ZZ_pE_both should have been overridden")

    cdef int _set_from_ZZ_pEX(self, ZZ_pEX_c* poly, ntl_ZZ_pEContext_class ctx) except -1:
        """
        Sets self from a ZZ_pEX_c.

        Fixed modulus elements should override this function.

        This function is not used internally.
        """
        if self.parent().is_capped_relative():
            self._set_from_ZZ_pEX_rel(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        elif self.parent().is_capped_absolute():
            self._set_from_ZZ_pEX_abs(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap)
        else:
            raise RuntimeError("_set_from_ZZ_pEX should have been overridden")

    cdef int _set_from_ZZ_pEX_rel(self, ZZ_pEX_c* poly, ntl_ZZ_pEContext_class ctx, long relprec) except -1:
        """
        Set from a ZZ_pEX_c with bounded relative precision.

        Capped relative elements should override this function, so the
        default implementation is for capped absolute.

        This function is not used internally.
        """
        self._set_from_ZZ_pEX_both(poly, ctx, (<PowComputer_class>self.parent().prime_pow).prec_cap, relprec)

    cdef int _set_from_ZZ_pEX_abs(self, ZZ_pEX_c* poly, ntl_ZZ_pEContext_class ctx, long absprec) except -1:
        """
        Set from a ZZ_pEX_c with bounded absolute precision.

        Capped absolute elements should override this function, so the
        default implementation is for capped relative.

        This function is not used internally.
        """
        self._set_from_ZZ_pEX_both(poly, ctx, absprec, (<PowComputer_class>self.parent().prime_pow).prec_cap)

    cdef int _set_from_ZZ_pEX_both(self, ZZ_pEX_c* poly, ntl_ZZ_pEContext_class ctx, long absprec, long relprec) except -1:
        """
        Sets from a ZZ_pEX_c with both absolute and relative precision bounded.

        Capped absolute and capped relative elements should override
        this function.

        This function is not used internally.
        """
        if self.parent().is_fixed_mod():
            self._set_from_ZZ_pEX(poly, ctx)
        else:
            raise RuntimeError("_set_from_ZZ_pEX_both should have been overridden")

    cdef long _check_ZZ_pContext(self, ntl_ZZ_pContext_class ctx) except -1:
        raise NotImplementedError

    cdef long _check_ZZ_pEContext(self, ntl_ZZ_pEContext_class ctx) except -1:
        raise NotImplementedError

    cdef ext_p_list(self, bint pos):
        raise NotImplementedError

    cdef ext_p_list_precs(self, bint pos, long prec):
        raise NotImplementedError

    def _const_term_test(self):
        """
        Returns the constant term of a polynomial representing self.

        This function is mainly for troubleshooting, and the meaning
        of the return value will depend on whether self is capped
        relative or otherwise.

        EXAMPLES::

            sage: R = Zp(5,5)
            sage: S.<x> = R[]
            sage: f = x^5 + 75*x^3 - 15*x^2 +125*x - 5
            sage: W.<w> = R.ext(f)
            sage: a = W(566)
            sage: a._const_term_test()
            566
        """
        cdef ntl_ZZ_p ans = ntl_ZZ_p(modulus=self.parent().prime())
        ans.x = self._const_term()
        return ans

    cdef ZZ_p_c _const_term(self):
        raise NotImplementedError

    def _ext_p_list(self, pos):
        """
        Returns a list of integers (in the Eisenstein case) or a list
        of lists of integers (in the unramified case).  self can be
        reconstructed as a sum of elements of the list times powers of
        the uniformiser (in the Eisenstein case), or as a sum of
        powers of the p times polynomials in the generator (in the
        unramified case).

        Note that zeros are truncated from the returned list, so you
        must use the valuation() function to completely recover self.

        INPUT:

            - pos -- bint.  If True, all integers will be in the range [0,p-1],
              otherwise they will be in the range [(1-p)/2, p/2].

        OUTPUT:

            - L -- A list of integers or list of lists giving the
              series expansion of self.

        EXAMPLES::

            sage: R = Zp(5,5)
            sage: S.<x> = R[]
            sage: f = x^5 + 75*x^3 - 15*x^2 +125*x - 5
            sage: W.<w> = R.ext(f)
            sage: y = W(775, 19); y
            w^10 + 4*w^12 + 2*w^14 + w^15 + 2*w^16 + 4*w^17 + w^18 + O(w^19)
            sage: y._ext_p_list(True)
            [1, 0, 4, 0, 2, 1, 2, 4, 1]
            sage: y._ext_p_list(False)
            [1, 0, -1, 0, 2, 1, 2, 0, 1]
        """
        return self.ext_p_list(pos)

    def frobenius(self, arithmetic=True):
        """
        Returns the image of this element under the Frobenius automorphism
        applied to its parent.

        INPUT:

        - ``self`` -- an element of an unramified extension.
        - ``arithmetic`` -- whether to apply the arithmetic Frobenius (acting
          by raising to the `p`-th power on the residue field). If ``False`` is
          provided, the image of geometric Frobenius (raising to the `(1/p)`-th
          power on the residue field) will be returned instead.

        EXAMPLES::

            sage: R.<a> = Zq(5^4,3)
            sage: a.frobenius()
            (a^3 + a^2 + 3*a) + (3*a + 1)*5 + (2*a^3 + 2*a^2 + 2*a)*5^2 + O(5^3)
            sage: f = R.defining_polynomial()
            sage: f(a)
            O(5^3)
            sage: f(a.frobenius())
            O(5^3)
            sage: for i in range(4): a = a.frobenius()
            sage: a
            a + O(5^3)

            sage: K.<a> = Qq(7^3,4)
            sage: b = (a+1)/7
            sage: c = b.frobenius(); c
            (3*a^2 + 5*a + 1)*7^-1 + (6*a^2 + 6*a + 6) + (4*a^2 + 3*a + 4)*7 + (6*a^2 + a + 6)*7^2 + O(7^3)
            sage: c.frobenius().frobenius()
            (a + 1)*7^-1 + O(7^3)

        An error will be raised if the parent of self is a ramified extension::

            sage: K.<a> = Qp(5).extension(x^2 - 5)
            sage: a.frobenius()
            Traceback (most recent call last):
            ...
            NotImplementedError: Frobenius automorphism only implemented for unramified extensions
        """
        R = self.parent()
        if R.e() != 1:
            raise NotImplementedError("Frobenius automorphism only implemented for unramified extensions")
        if self.is_zero(): return self
        L = self.teichmuller_list()
        ppow = R.uniformizer_pow(self.valuation())
        if arithmetic:
            exp = R.prime()
        else:
            exp = R.prime()**(R.degree()-1)
        ans = ppow * L[0]**exp
        for m in range(1,len(L)):
            ppow = ppow << 1
            ans += ppow * L[m]**exp
        return ans

    cpdef bint _is_base_elt(self, p) except -1:
        """
        Return ``True`` if this element is an element of Zp or Qp (rather than
        an extension).

        INPUT:

        - ``p`` -- a prime, which is compared with the parent of this element.

        EXAMPLES::

            sage: K.<a> = Qq(7^3,4)
            sage: a._is_base_elt(5)
            False

        """
        return False

    def residue(self, absprec=1):
        r"""
        Reduces this element modulo `\pi^\mathrm{absprec}`.

        INPUT:

        - ``absprec`` - a non-negative integer (default: ``1``)

        OUTPUT:

        This element reduced modulo `\pi^\mathrm{absprec}`.

        If ``absprec`` is zero, then as an element of `\ZZ/(1)`.

        If ``absprec`` is one, then as an element of the residue field.

        .. NOTE::

            Only implemented for ``absprec`` less than or equal to one.

        AUTHORS:

        - Julian Rueth (2012-10-18): initial version

        EXAMPLES:

        Unramified case::

            sage: R = ZpCA(3,5)
            sage: S.<a> = R[]
            sage: W.<a> = R.extension(a^2 + 9*a + 1)
            sage: (a + 1).residue(1)
            a0 + 1
            sage: a.residue(2)
            Traceback (most recent call last):
            ...
            NotImplementedError: reduction modulo p^n with n>1.

        Eisenstein case::

            sage: R = ZpCA(3,5)
            sage: S.<a> = R[]
            sage: W.<a> = R.extension(a^2 + 9*a + 3)
            sage: (a + 1).residue(1)
            1
            sage: a.residue(2)
            Traceback (most recent call last):
            ...
            NotImplementedError: residue() not implemented in extensions for absprec larger than one.

        TESTS:

            sage: K = Qp(3,5)
            sage: S.<a> = R[]
            sage: W.<a> = R.extension(a^2 + 9*a + 1)
            sage: (a/3).residue(0)
            Traceback (most recent call last):
            ...
            ValueError: element must have non-negative valuation in order to compute residue.

            sage: R = ZpFM(3,5)
            sage: S.<a> = R[]
            sage: W.<a> = R.extension(a^2 + 9*a + 1)
            sage: W.one().residue(0)
            0
            sage: a.residue(-1)
            Traceback (most recent call last):
            ...
            ValueError: cannot reduce modulo a negative power of the uniformizer.
            sage: a.residue(16)
            Traceback (most recent call last):
            ...
            PrecisionError: insufficient precision to reduce modulo p^16.

        """
        if absprec < 0:
            raise ValueError("cannot reduce modulo a negative power of the uniformizer.")
        if absprec > self.precision_absolute():
            from precision_error import PrecisionError
            raise PrecisionError("not enough precision known in order to compute residue.")
        if self.valuation() < 0:
            raise ValueError("element must have non-negative valuation in order to compute residue.")

        if absprec == 0:
            from sage.rings.finite_rings.integer_mod import Mod
            return Mod(0,1)
        elif absprec == 1:
            return self.parent().residue_field()(self[0])
        else:
            raise NotImplementedError("residue() not implemented in extensions for absprec larger than one.")
