r"""
<Short one-line summary that ends with no period>

<Paragraph description>

EXAMPLES::

<Lots and lots of examples>

AUTHORS:

- Carlos Alquézar Baeta (2023-02-07): initial version
- Miguel Ángel Marco (2023-02-07): initial version
- Jorge Martín Morales (2023-02-07): initial version

"""

# ****************************************************************************
#       Copyright (C) 2023 Carlos Alquézar Baeta <alquezar@unizar.es>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************


from sage.algebras.commutative_dga import *
from sage.structure.parent import Parent
from sage.structure.element import Element


class AplN(DifferentialGCAlgebra_multigraded, UniqueRepresentation):
    """
      A differential multi-graded commutative algebra

      This class extends the DifferentialGCAlgebra_multigraded class
      to create the algebra of polynomial differential forms with coefficients in QQ

    INPUT:

    - ``A`` -- a multi-graded commutative algebra
    - ``differential`` -- the differential rules

    TESTS::

    sage: 2
    2

    """

    def __init__(self, A, differential):
        """
        Initialize ``self``.

        INPUT:

        - ``A`` -- a multi-graded commutative algebra
        - ``differential`` -- a differential

                TESTS::

                sage: 2
                2
        """

        self.n = len(A._names) // 2 - 1
        DifferentialGCAlgebra_multigraded.__init__(self, A, differential)
        self.t_relation = sum(self.gens()[:self.n+1])-1
        self.y_relation = sum(self.gens()[self.n+1:])

    def _base_repr(self):
        """
        Return the base string representation of ``self``.

                TESTS::

                sage: 2
                2
        """

        s = DifferentialGCAlgebra._base_repr(self)
        old = '{}'.format(self._degrees)
        new = '{}'.format([i[0] for i in self._degrees_multi])
        return s.replace(old, new)

    def integrand(self, value):
        """
        Returns a polynomial whose definite integral is equal to 'value'

        OUTPUT:
            - A polynomial

                TESTS::

                sage: 2
                2
        """
        if (value in QQ) or ((value in A) and (value.degree().is_zero())):
            return value*factorial(self.n)*prod(self.gens()[self.n+2:])
        else:
            raise ValueError("only field elements can be integrated.")

    class Element(GCAlgebra_multigraded.Element, DifferentialGCAlgebra.Element):
        """
        Element class of the commutative differential multi-graded algebra (Apl)n.
        """

        def is_homogeneous(self):
            """
            Return "True" if "self" is homogeneous and "False" otherwise.

                        TESTS::

                        sage: 2
                        2
            """

            degrees = set([k[0] for k in self.homogeneous_parts()])
            return len(degrees) == 1

        def total_degree(self):
            """
            Return the total degree of this element.

                        TESTS::

                        sage: 2
                        2
            """

            degrees = set([k[0] for k in self.homogeneous_parts()])
            return max(degrees)

        def homogeneous_part(self, degree):
            """
            Returns the homogeneous part of the element for a given degree.

                        TESTS::

                        sage: 2
                        2
            """

            if (degree < 0):
                raise ValueError('`degree` must be a non-negative integer')
            parts = [self.homogeneous_parts()[k]
                     for k in self.homogeneous_parts() if list(k)[0] == degree]
            return self.parent()(sum(parts))

        def reduce(self, relations=[]):
            """"
            The method reduces the polynomial by the Ideal <sum(t_i)-1, sum(y_j)>

            INPUT::
                - p: polynomial in variables t_0,..t_n,y_0,...,y_n
                - relations (optional): if relations are given, the polynomial is reduced following
                    this rules. If no relations are given, the polynomial is reduced taking into
                    account the relations of the parent algebra.

            OUTPUT::
                - The polynomial after applying the corresponding relations

                        TESTS::

                        sage: 2
                        2
            """
            apl = self.parent()
            if relations:
                return DifferentialGCAlgebra.Element.reduce(self, relations)
            else:
                return DifferentialGCAlgebra.Element.reduce(self, [apl.t_relation, apl.y_relation])

        def _integrable_part(self):
            """
            The method returns the coefficient in Q[t_1,...,t_n] of the term
            of degree `n`

            OUTPUT:
                - A polynomial in the ring Q<t_1,..t_n>.
                  (After reducing the variables t_{0},y_{0} by the corresponding relations,
                  the algorithm removes the variables y_{1},..,y_{n}).

                        TESTS::

                        sage: 2
                        2
            """

            apl = self.parent()
            # needed in case reduce() returns an integer
            pol = apl(self.reduce())
            relations = []
            for i in range(apl.n+1):
                relations.append(apl.gens()[apl.n+1+i]-1)
            pol = apl(pol.homogeneous_part(degree=apl.n))
            return apl(pol.reduce(relations=relations))

        def integral(self):
            """
            Returns the definite integral of the term of degree `n`

            OUTPUT:
                - The rational number obtained by the definite integral
                  of the term of degree `n`

                        TESTS::

                        sage: 2
                        2
            """

            apl = self.parent()
            pol = self._integrable_part()
            result = 0
            for k in pol.dict().keys():  # k is a monomial
                numerator = 1
                denominator = apl.n
                for e in k:  # e is an exponents
                    numerator *= factorial(e)
                    denominator += e
                coef = pol.dict()[k]
                result += coef*numerator/factorial(denominator)
            # return ((-1) ** ((apl.n*(apl.n-1)) // 2)) * result
            return result

        def _get_term(self, coef, exp, gens):
            """"

            Returns the term specified by its coefficient, exponents and generators

            INPUT::
                - coef: The coefficient of the term
                - exp: A tuple with the exponents of each variable
                - gens: The generators

            OUTPUT::
                - A polynomial

                        TESTS::

                        sage: 2
                        2
            """
            return coef*prod(list(map(lambda exp, gen: gen**exp, exp, gens)))

        def _get_terms(self):
            """"

            Returns self seen as a collection of terms in the variables y_j with coefficients in QQ<t_1,..,t_n>

            OUTPUT::
                - A dictionary with all terms

                        TESTS::

                        sage: 2
                        2

            """
            a = self.parent()
            n = len(a.gens())//2
            ts_only_gens = a.gens()[:n]+tuple(n*[1])
            ys_only_gens = tuple(n*[1])+a.gens()[n:]
            terms = {}
            pol = self.dict()
            for exp in pol.keys():
                key = self._get_term(1, exp, ys_only_gens)
                coef = self._get_term(pol[exp], exp, ts_only_gens)
                if key not in terms.keys():
                    terms[key] = coef
                else:
                    terms[key] += coef
            return terms

        def _get_leading_pol(self):
            """"

            Returns the leading polynomial

            OUTPUT::
                - A polynomial

                        TESTS::

                        sage: 2
                        2
            """
            terms = self._get_terms()
            keys = list(terms.keys())
            keys.sort(reverse=True)
            return terms[keys[0]]*keys[0]

        def _get_idx(self, exp):
            """"

            Returns the idx 'i' of the first generator 'y_i'  that appears in exp

            INPUT::
                - exp: A tuple with the exponents

            OUTPUT::
                - A non-negative number lower or equal to 'n'

                        TESTS::

                        sage: 2
                        2
            """
            n = len(exp)//2
            filter_ts = n*[0]+n*[1]
            diff_forms = list(map(lambda v, f: v*f, exp, filter_ts))
            return next((i for i, x in enumerate(list(diff_forms)) if x), None)-n

        def _monomial_primitive(self, monomial, idx):
            """"

            Returns the primitive of the given monomial

            INPUT::
                - monomial: A monomial
                - idx: the index of the variable y_j that will be used to calculate the pseudo-primitive

            OUTPUT::
                - A polynomial

                        TESTS::

                        sage: 2
                        2
            """
            exp = list(monomial.dict().keys())[0]
            coef = monomial.dict()[exp]
            n = len(exp)//2
            if (exp[n+idx] == 0):
                raise IndexError(
                    "It is not possible to get a primitive using the given idx.")
            primitive = list(deepcopy(exp))
            primitive[n+idx] = 0
            primitive[idx] += 1
            new_coef = coef/primitive[idx]
            a = monomial.parent()
            return self._get_term(new_coef, primitive, a.gens())

        def _leading_primitive(self):
            """"
            Returns the primitive of the leading term of self

            OUTPUT::
                - A polynomial

                        TESTS::

                        sage: 2
                        2
            """
            polynomial = self._get_leading_pol()
            # necessary condition: all monomials have apl.gens[n+idx] factor
            exp = list(polynomial.dict().keys())[0]
            idx = self._get_idx(exp)
            aGens = polynomial.parent().gens()
            primitive = 0
            for exp in polynomial.dict().keys():
                coef = polynomial.dict()[exp]
                monomial = self._get_term(coef, exp, aGens)
                primitive += self._monomial_primitive(monomial, idx)
            terms = primitive._get_terms()
            result = sum([terms[key]*key for key in terms.keys()])
            return result

        def primitive(self):
            """"
            Retuns a polynomial whose differential is self

            OUTPUT::
                - A polynomial

                        TESTS::

                        sage: 2
                        2
            """
            if self.differential() != 0:
                raise ValueError("This polynomial is not closed.")
            w = self._leading_primitive()
            theta = self - w.differential()
            if (theta == 0):
                return w
            else:
                return w+theta.primitive()


def AlgebraPL(n):
    """
    The graded commutative differential cochain algebra Apl of dimension ``n``

    INPUT:

    - ``n`` -- A non-negative integer corresponding to the number of generators of degree 0.

        TESTS::

        sage: 2
        2
    """

    if (n < 0):
        raise ValueError('`n` must be a non-negative integer')
    ts = ['t'+str(t) for t in range(n+1)]
    ys = ['y'+str(t) for t in range(n+1)]
    degrees = tuple((n+1)*[(0, 2)])+tuple((n+1)*[(1, 2)])
    gca = GradedCommutativeAlgebra(QQ, names=ts+ys, degrees=degrees)
    differentials = {gca.gens_dict()[ts[i]]: gca.gens_dict()[
        ys[i]] for i in range(n+1)}
    return AplN(gca, differentials)

#############################################
# TEST
# print("Test AlgebraPL")
#A = AlgebraPL(5)
# A.inject_variables()
#p = t0**2 + t1 + y3 + y1*y2
#q = p.differential()
#############################################