r"""
Subspaces of modular forms for Hecke triangle groups

AUTHORS:

- Jonas Jermann (2013): initial version

"""
from __future__ import absolute_import

#*****************************************************************************
#       Copyright (C) 2013-2014 Jonas Jermann <jjermann2@gmail.com>
#
#  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.all import ZZ, QQ, infinity

from sage.modules.module import Module
from sage.categories.all import Modules
from sage.modules.free_module_element import vector
from sage.structure.unique_representation import UniqueRepresentation
from sage.misc.cachefunc import cached_method
from sage.matrix.constructor import matrix

from .hecke_triangle_groups import HeckeTriangleGroup
from .abstract_space import FormsSpace_abstract


def canonical_parameters(ambient_space, basis, check=True):
    r"""
    Return a canonical version of the parameters.
    In particular the list/tuple ``basis`` is replaced by a
    tuple of linearly independent elements in the ambient space.

    If ``check=False`` (default: ``True``) then ``basis``
    is assumed to already be a basis.

    EXAMPLES::

        sage: from sage.modular.modform_hecketriangle.subspace import canonical_parameters
        sage: from sage.modular.modform_hecketriangle.space import ModularForms
        sage: MF = ModularForms(n=6, k=12, ep=1)
        sage: canonical_parameters(MF, [MF.Delta().as_ring_element(), MF.gen(0), 2*MF.gen(0)])
        (ModularForms(n=6, k=12, ep=1) over Integer Ring,
         (q + 30*q^2 + 333*q^3 + 1444*q^4 + O(q^5),
          1 + 26208*q^3 + 530712*q^4 + O(q^5)))
    """

    if check:
        coord_matrix = matrix([ambient_space(v).ambient_coordinate_vector() for v in basis])
        pivots = coord_matrix.transpose().pivots()
        new_basis = [ambient_space(basis[l]) for l in pivots]
        basis = tuple(new_basis)
    else:
        basis = [ambient_space(v) for v in basis]
        basis = tuple(basis)

    return (ambient_space, basis)

def ModularFormsSubSpace(*args, **kwargs):
    r"""
    Create a modular forms subspace generated by the supplied arguments if possible.
    Instead of a list of generators also multiple input arguments can be used.
    If ``reduce=True`` then the corresponding ambient space is choosen as small as possible.
    If no subspace is available then the ambient space is returned.

    EXAMPLES::

        sage: from sage.modular.modform_hecketriangle.subspace import ModularFormsSubSpace
        sage: from sage.modular.modform_hecketriangle.space import ModularForms
        sage: MF = ModularForms()
        sage: subspace = ModularFormsSubSpace(MF.E4()^3, MF.E6()^2+MF.Delta(), MF.Delta())
        sage: subspace
        Subspace of dimension 2 of ModularForms(n=3, k=12, ep=1) over Integer Ring
        sage: subspace.ambient_space()
        ModularForms(n=3, k=12, ep=1) over Integer Ring
        sage: subspace.gens()
        [1 + 720*q + 179280*q^2 + 16954560*q^3 + 396974160*q^4 + O(q^5), 1 - 1007*q + 220728*q^2 + 16519356*q^3 + 399516304*q^4 + O(q^5)]
        sage: ModularFormsSubSpace(MF.E4()^3-MF.E6()^2, reduce=True).ambient_space()
        CuspForms(n=3, k=12, ep=1) over Integer Ring
        sage: ModularFormsSubSpace(MF.E4()^3-MF.E6()^2, MF.J_inv()*MF.E4()^3, reduce=True)
        WeakModularForms(n=3, k=12, ep=1) over Integer Ring
    """

    generators = []
    for arg in args:
        if isinstance(arg, list) or isinstance(arg, tuple):
            generators += arg
        else:
            generators.append(arg)
    if ("reduce" in kwargs) and kwargs["reduce"]:
        generators = [gen.full_reduce() for gen in generators]

    if len(generators) == 0:
        raise ValueError("No generators specified")

    el = False
    for gen in generators:
        if el:
            el += gen
        else:
            el = gen

    ambient_space = el.parent()

    try:
        # This works if and only if ambient_space supports subspaces
        ambient_space.coordinate_vector(el)

        generators = [ambient_space(gen) for gen in generators]
        return SubSpaceForms(ambient_space, generators)
    except (NotImplementedError, AttributeError):
        return ambient_space


class SubSpaceForms(FormsSpace_abstract, Module, UniqueRepresentation):
    r"""
    Submodule of (Hecke) forms in the given ambient space for the given basis.
    """

    @staticmethod
    def __classcall__(cls, ambient_space, basis=(), check=True):
        r"""
        Return a (cached) instance with canonical parameters.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.subspace import (canonical_parameters, SubSpaceForms)
            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=12, ep=1)
            sage: (ambient_space, basis) = canonical_parameters(MF, [MF.Delta().as_ring_element(), MF.gen(0)])
            sage: SubSpaceForms(MF, [MF.Delta().as_ring_element(), MF.gen(0)]) == SubSpaceForms(ambient_space, basis)
            True
        """

        (ambient_space, basis) = canonical_parameters(ambient_space, basis, check)

        # we return check=True to ensure only one cached instance
        return super(SubSpaceForms,cls).__classcall__(cls, ambient_space=ambient_space, basis=basis, check=True)

    def __init__(self, ambient_space, basis, check):
        r"""
        Return the Submodule of (Hecke) forms in ``ambient_space`` for the given ``basis``.

        INPUT:

        - ``ambient_space``  -- An ambient forms space.

        - ``basis``          -- A tuple of (not necessarily linearly independent)
                                elements of ``ambient_space``.

        - ``check``          -- If ``True`` (default) then a maximal linearly
                                independent subset of ``basis`` is choosen. Otherwise
                                it is assumed that ``basis`` is linearly independent.

        OUTPUT:

        The corresponding submodule.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms, QuasiCuspForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: MF
            ModularForms(n=6, k=20, ep=1) over Integer Ring
            sage: MF.dimension()
            4
            sage: subspace = MF.subspace([MF.Delta()*MF.E4()^2, MF.gen(0), 2*MF.gen(0)])
            sage: subspace
            Subspace of dimension 2 of ModularForms(n=6, k=20, ep=1) over Integer Ring
            sage: subspace.analytic_type()
            modular
            sage: subspace.category()
            Category of modules over Integer Ring
            sage: subspace in subspace.category()
            True
            sage: subspace.module()
            Vector space of degree 4 and dimension 2 over Fraction Field of Univariate Polynomial Ring in d over Integer Ring
            Basis matrix:
            [            1             0             0             0]
            [            0             1     13/(18*d) 103/(432*d^2)]
            sage: subspace.ambient_module()
            Vector space of dimension 4 over Fraction Field of Univariate Polynomial Ring in d over Integer Ring
            sage: subspace.ambient_module() == MF.module()
            True
            sage: subspace.ambient_space() == MF
            True
            sage: subspace.basis()
            [q + 78*q^2 + 2781*q^3 + 59812*q^4 + O(q^5), 1 + 360360*q^4 + O(q^5)]
            sage: subspace.basis()[0].parent() == MF
            True
            sage: subspace.gens()
            [q + 78*q^2 + 2781*q^3 + 59812*q^4 + O(q^5), 1 + 360360*q^4 + O(q^5)]
            sage: subspace.gens()[0].parent() == subspace
            True
            sage: subspace.is_ambient()
            False

            sage: MF = QuasiCuspForms(n=infinity, k=12, ep=1)
            sage: MF.dimension()
            4
            sage: subspace = MF.subspace([MF.Delta(), MF.E4()*MF.f_inf()*MF.E2()*MF.f_i(), MF.E4()*MF.f_inf()*MF.E2()^2, MF.E4()*MF.f_inf()*(MF.E4()-MF.E2()^2)])
            sage: subspace.default_prec(3)
            sage: subspace
            Subspace of dimension 3 of QuasiCuspForms(n=+Infinity, k=12, ep=1) over Integer Ring
            sage: subspace.gens()
            [q + 24*q^2 + O(q^3), q - 24*q^2 + O(q^3), q - 8*q^2 + O(q^3)]
        """

        FormsSpace_abstract.__init__(self, group=ambient_space.group(), base_ring=ambient_space.base_ring(), k=ambient_space.weight(), ep=ambient_space.ep(), n=ambient_space.hecke_n())
        Module.__init__(self, base=ambient_space.base_ring())

        self._ambient_space = ambient_space
        self._basis = [v for v in basis]
        # self(v) instead would somehow mess up the coercion model
        self._gens = [self._element_constructor_(v) for v in basis]
        self._module = ambient_space._module.submodule([ambient_space.coordinate_vector(v) for v in basis])
        # TODO: get the analytic type from the basis
        #self._analytic_type=self.AT(["quasi", "mero"])
        self._analytic_type = ambient_space._analytic_type

    def _repr_(self):
        r"""
        Return the string representation of ``self``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([MF.Delta()*MF.E4()^2, MF.gen(0)])
            sage: subspace
            Subspace of dimension 2 of ModularForms(n=6, k=20, ep=1) over Integer Ring
        """

        # If we list the basis the representation usually gets too long...
        # return "Subspace with basis {} of {}".format([v.as_ring_element() for v in self.basis()], self._ambient_space)
        return "Subspace of dimension {} of {}".format(len(self._basis), self._ambient_space)

    def change_ring(self, new_base_ring):
        r"""
        Return the same space as ``self`` but over a new base ring ``new_base_ring``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([MF.Delta()*MF.E4()^2, MF.gen(0)])
            sage: subspace.change_ring(CC)
            Subspace of dimension 2 of ModularForms(n=6, k=20, ep=1) over Complex Field with 53 bits of precision
        """

        return self.__class__.__base__(self._ambient_space.change_ring(new_base_ring), self._basis, check=False)

    def change_ambient_space(self, new_ambient_space):
        r"""
        Return a new subspace with the same basis but inside a different ambient space
        (if possible).

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms, QuasiModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([MF.Delta()*MF.E4()^2, MF.gen(0)])
            sage: new_ambient_space = QuasiModularForms(n=6, k=20, ep=1)
            sage: subspace.change_ambient_space(new_ambient_space)    # long time
            Subspace of dimension 2 of QuasiModularForms(n=6, k=20, ep=1) over Integer Ring
        """
        return self.__class__.__base__(new_ambient_space, self._basis, check=False)

    @cached_method
    def contains_coeff_ring(self):
        r"""
        Return whether ``self`` contains its coefficient ring.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(k=0, ep=1, n=8)
            sage: subspace = MF.subspace([1])
            sage: subspace.contains_coeff_ring()
            True
            sage: subspace = MF.subspace([])
            sage: subspace.contains_coeff_ring()
            False
            sage: MF = ModularForms(k=0, ep=-1, n=8)
            sage: subspace = MF.subspace([])
            sage: subspace.contains_coeff_ring()
            False
        """

        return (super(SubSpaceForms, self).contains_coeff_ring() and self.dimension()==ZZ(1))

    @cached_method
    def basis(self):
        r"""
        Return the basis of ``self`` in the ambient space.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.basis()
            [q + 78*q^2 + 2781*q^3 + 59812*q^4 + O(q^5), 1 + 360360*q^4 + O(q^5)]
            sage: subspace.basis()[0].parent() == MF
            True
        """

        return self._basis

    @cached_method
    def gens(self):
        r"""
        Return the basis of ``self``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.gens()
            [q + 78*q^2 + 2781*q^3 + 59812*q^4 + O(q^5), 1 + 360360*q^4 + O(q^5)]
            sage: subspace.gens()[0].parent() == subspace
            True
        """

        return self._gens

    @cached_method
    def dimension(self):
        r"""
        Return the dimension of ``self``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.dimension()
            2
            sage: subspace.dimension() == len(subspace.gens())
            True
        """
        return len(self.basis())

    @cached_method
    def degree(self):
        r"""
        Return the degree of ``self``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.degree()
            4
            sage: subspace.degree() == subspace.ambient_space().degree()
            True
        """
        return self._ambient_space.degree()

    @cached_method
    def rank(self):
        r"""
        Return the rank of ``self``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.rank()
            2
            sage: subspace.rank() == subspace.dimension()
            True
        """
        return len(self.gens())

    @cached_method
    def coordinate_vector(self, v):
        r"""
        Return the coordinate vector of ``v`` with respect to
        the basis ``self.gens()``.

        INPUT:

        - ``v`` -- An element of ``self``.

        OUTPUT:

        The coordinate vector of ``v`` with respect
        to the basis ``self.gens()``.

        Note: The coordinate vector is not an element of ``self.module()``.

        EXAMPLES::

            sage: from sage.modular.modform_hecketriangle.space import ModularForms, QuasiCuspForms
            sage: MF = ModularForms(n=6, k=20, ep=1)
            sage: subspace = MF.subspace([(MF.Delta()*MF.E4()^2).as_ring_element(), MF.gen(0)])
            sage: subspace.coordinate_vector(MF.gen(0) + MF.Delta()*MF.E4()^2).parent()
            Vector space of dimension 2 over Fraction Field of Univariate Polynomial Ring in d over Integer Ring
            sage: subspace.coordinate_vector(MF.gen(0) + MF.Delta()*MF.E4()^2)
            (1, 1)

            sage: MF = ModularForms(n=4, k=24, ep=-1)
            sage: subspace = MF.subspace([MF.gen(0), MF.gen(2)])
            sage: subspace.coordinate_vector(subspace.gen(0)).parent()
            Vector space of dimension 2 over Fraction Field of Univariate Polynomial Ring in d over Integer Ring
            sage: subspace.coordinate_vector(subspace.gen(0))
            (1, 0)

            sage: MF = QuasiCuspForms(n=infinity, k=12, ep=1)
            sage: subspace = MF.subspace([MF.Delta(), MF.E4()*MF.f_inf()*MF.E2()*MF.f_i(), MF.E4()*MF.f_inf()*MF.E2()^2, MF.E4()*MF.f_inf()*(MF.E4()-MF.E2()^2)])
            sage: el = MF.E4()*MF.f_inf()*(7*MF.E4() - 3*MF.E2()^2)
            sage: subspace.coordinate_vector(el)
            (7, 0, -3)
            sage: subspace.ambient_coordinate_vector(el)
            (7, 21/(8*d), 0, -3)
        """

        return self._module.coordinate_vector(self.ambient_coordinate_vector(v))

