"""
Types of Jacobi forms of fixed index and weight.

AUTHOR :
    - Martin Raum (2010 - 04 - 07) Initial version.
"""

#===============================================================================
# 
# Copyright (C) 2010 Martin Raum
# 
# 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 3
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
#
#===============================================================================

from operator import xor
from psage.modform.fourier_expansion_framework.gradedexpansions.gradedexpansion_grading import TrivialGrading
from psage.modform.fourier_expansion_framework.modularforms.modularform_ambient import ModularFormsModule_generic
from psage.modform.fourier_expansion_framework.modularforms.modularform_types import ModularFormType_abstract
from psage.modform.jacobiforms.jacobiformd1nn_fegenerators import jacobi_form_by_taylor_expansion,\
    _jacobi_forms_by_taylor_expansion_coords
from psage.modform.jacobiforms.jacobiformd1nn_fourierexpansion import JacobiD1NNFourierExpansionModule, \
                                                                      JacobiFormD1NNFilter
from sage.categories.number_fields import NumberFields
from sage.matrix.constructor import diagonal_matrix, matrix, zero_matrix,\
    identity_matrix
from sage.misc.cachefunc import cached_method
from sage.misc.mrange import mrange
from sage.modular.modform.constructor import ModularForms
from sage.rings.all import Integer
from sage.rings.integer_ring import ZZ
from sage.rings.number_field.number_field import CyclotomicField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.rational_field import QQ
from sage.structure.sequence import Sequence


#===============================================================================
# JacobiFormsD1NN
#===============================================================================

_jacobiforms_cache = dict()

def JacobiFormsD1NN(A, type, precision, *args, **kwds) :
    global _jacobiforms_cache
    
    if isinstance(precision, (int, Integer)) :
        precision = JacobiFormD1NNFilter(precision, type.index())
    
    k = (A, type, precision)
    
    try :
        return _jacobiforms_cache[k]
    except KeyError :
        if isinstance(type, JacobiFormD1NN_Gamma) :
            M = ModularFormsModule_generic(A, type, precision)
        else :
            raise TypeError, "%s must be a Jacobi form type" % (type,)
        
        _jacobiforms_cache[k] = M
        return M
    
#===============================================================================
# JacobiFormD1NN_Gamma
#===============================================================================

class JacobiFormD1NN_Gamma ( ModularFormType_abstract ) :
    r"""
    Type of Jacobi forms of degree `1` and index in `\mathbb{N}` associated with
    the full modular group.
    
    TESTS::
    
        sage: from psage.modform.jacobiforms import *
        sage: from psage.modform.jacobiforms.jacobiformd1nn_fourierexpansion import JacobiFormD1NNFilter 
        sage: JR = JacobiFormsD1NN(QQ, JacobiFormD1NN_Gamma(3, 6), JacobiFormD1NNFilter(10, 3))
        sage: JR.gens()
        (Graded expansion TDE_0, Graded expansion TDE_1)
        sage: JR.0 + 2 * JR.1
        Graded expansion TDE_0 + 2*TDE_1
    """
    def __init__(self, index, weight) :
        if weight % 2 != 0 :
            raise NotImplementedError, "Only even weight forms are implemented."
  
        self.__index = index
        self.__weight = weight        
    
    def index(self) :
        return self.__index
    
    def weight(self) :
        return self.__weight
    
    def _ambient_construction_function(self) :
        return JacobiFormsD1NN
    
    def group(self) :
        return "Sp(2, ZZ)_\infty"
    
    @cached_method
    def _rank(self, K) :
        
        if K is QQ or K in NumberFields() :
            return len(_jacobi_forms_by_taylor_expansion_coords(self.__index, self.__weight, 0))

            ## This is the formula used by Poor and Yuen in Paramodular cusp forms
            if self.__weight == 2 :
                delta = len(self.__index.divisors()) // 2 - 1
            else :
                delta = 0
                
            return sum( ModularForms(1, self.__weight + 2 * j).dimension() + j**2 // (4 * self.__index)
                        for j in xrange(self.__index + 1) ) \
                   + delta
            

            ## This is the formula given by Skoruppa in 
            ## Jacobi forms of critical weight and Weil representations
            ##FIXME: There is some mistake here
            if self.__weight % 2 != 0 :
                ## Otherwise the space X(i**(n - 2 k)) is different
                ## See: Skoruppa, Jacobi forms of critical weight and Weil representations
                raise NotImplementedError
            
            m = self.__index
            K = CyclotomicField(24 * m, 'zeta')
            zeta = K.gen(0)
            
            quadform = lambda x : 6 * x**2
            bilinform = lambda x,y : quadform(x + y) - quadform(x) - quadform(y)
            
            T = diagonal_matrix([zeta**quadform(i) for i in xrange(2*m)])
            S =   sum(zeta**(-quadform(x)) for x in xrange(2 * m)) / (2 * m) \
                * matrix([[zeta**(-bilinform(j,i)) for j in xrange(2*m)] for i in xrange(2*m)])
            subspace_matrix_1 = matrix( [ [1 if j == i or j == 2*m - i else 0 for j in xrange(m + 1) ]
                                        for i in xrange(2*m)] )
            subspace_matrix_2 = zero_matrix(ZZ, m + 1, 2*m)
            subspace_matrix_2.set_block(0,0,identity_matrix(m+1))
            
            T = subspace_matrix_2 * T * subspace_matrix_1
            S = subspace_matrix_2 * S * subspace_matrix_1
            
            sqrt3 = (zeta**(4*m) - zeta**(-4*m)) * zeta**(-6*m) 
            rank =   (self.__weight - 1/2 - 1) / 2 * (m + 1) \
                   + 1/8 * (   zeta**(3*m * (2*self.__weight - 1)) * S.trace()
                             + zeta**(3*m * (1 - 2*self.__weight)) * S.trace().conjugate() ) \
                   + 2/(3*sqrt3) * (   zeta**(4 * m * self.__weight) * (S*T).trace()
                                     + zeta**(-4 * m * self.__weight) * (S*T).trace().conjugate() ) \
                   - sum((j**2 % (m+1))/(m+1) -1/2 for j in range(0,m+1))
            
            if self.__weight > 5 / 2 :
                return rank
            else :
                raise NotImplementedError
            
        raise NotImplementedError
    
    @cached_method
    def generators(self, K, precision) :
        if K is QQ or K in NumberFields() :
            return Sequence( [ jacobi_form_by_taylor_expansion(i, self.__index, self.__weight, precision)
                               for i in xrange(self._rank(K)) ],
                             universe = JacobiD1NNFourierExpansionModule(QQ, self.__index) )
        
        raise NotImplementedError
    
    def grading(self, K) :
        if K is QQ or K in NumberFields() :
            return TrivialGrading( self._rank(K),
                                   (self.__index, self.__weight) )
        
        raise NotImplementedError

    def _generator_names(self, K) :
        if K is QQ or K in NumberFields() :
            return [ "TDE_%s" % (i,) for i in xrange(self._rank(K)) ]
        
        raise NotImplementedError

    def _generator_by_name(self, K, name) :
        if K is QQ or K in NumberFields() :
            R = self.generator_relations(K).ring()
            try :
                return R.gen(self._generator_names(K).index(name))
            except ValueError :
                raise ValueError, "name %s doesn't exist for %s" % (name, K)
        
        raise NotImplementedError
    
    @cached_method
    def generator_relations(self, K) :
        r"""
        An ideal I in a polynomial ring R, such that the associated module
        is (R / I)_1. 
        """
        if K is QQ or K in NumberFields() :
            R = PolynomialRing(K, self._generator_names(K))
            return R.ideal(0)
            
        raise NotImplementedError

    def weights(self, K) :
        r"""
        A list of integers corresponding to the weights.
        """
        if K is QQ or K in NumberFields() :
            return len(self._theta_decomposition_indices()) \
                    * [(self.__index, self.__weight)]
            
        raise NotImplementedError
    
    def graded_submodules_are_free(self) :
        return True

    def __cmp__(self, other) :
        c = cmp(type(self), type(other))
        
        if c == 0 :
            c = cmp(self.__index, other.__index)
        if c == 0 :
            c = cmp(self.__weight, other.__weight)
            
        return c

    def __hash__(self) :
        return reduce(xor, map(hash, [type(self), self.__index, self.__weight]))

7e7e31216cb74ac921e0098d8963c9fb94a065a0
"""
Types of Jacobi forms of fixed index and weight.

AUTHOR :
    - Martin Raum (2010 - 04 - 07) Initial version.
"""

#===============================================================================
# 
# Copyright (C) 2010 Martin Raum
# 
# 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 3
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
#
#===============================================================================

from operator import xor
from psage.modform.fourier_expansion_framework.gradedexpansions.gradedexpansion_grading import TrivialGrading
from psage.modform.fourier_expansion_framework.modularforms.modularform_ambient import ModularFormsModule_generic
from psage.modform.fourier_expansion_framework.modularforms.modularform_types import ModularFormType_abstract
from psage.modform.jacobiforms.jacobiformd1nn_fegenerators import jacobi_form_by_taylor_expansion,\
    _jacobi_forms_by_taylor_expansion_coords
from psage.modform.jacobiforms.jacobiformd1nn_fourierexpansion import JacobiD1NNFourierExpansionModule, \
                                                                      JacobiFormD1NNFilter
from sage.categories.number_fields import NumberFields
from sage.matrix.constructor import diagonal_matrix, matrix, zero_matrix,\
    identity_matrix
from sage.misc.cachefunc import cached_method
from sage.misc.mrange import mrange
from sage.modular.modform.constructor import ModularForms
from sage.rings.all import Integer
from sage.rings.integer_ring import ZZ
from sage.rings.number_field.number_field import CyclotomicField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.rational_field import QQ
from sage.structure.sequence import Sequence


#===============================================================================
# JacobiFormsD1NN
#===============================================================================

_jacobiforms_cache = dict()

def JacobiFormsD1NN(A, type, precision, *args, **kwds) :
    global _jacobiforms_cache
    
    if isinstance(precision, (int, Integer)) :
        precision = JacobiFormD1NNFilter(precision, type.index())
    
    k = (A, type, precision)
    
    try :
        return _jacobiforms_cache[k]
    except KeyError :
        if isinstance(type, JacobiFormD1NN_Gamma) :
            M = ModularFormsModule_generic(A, type, precision)
        else :
            raise TypeError, "%s must be a Jacobi form type" % (type,)
        
        _jacobiforms_cache[k] = M
        return M
    
#===============================================================================
# JacobiFormD1NN_Gamma
#===============================================================================

class JacobiFormD1NN_Gamma ( ModularFormType_abstract ) :
    r"""
    Type of Jacobi forms of degree `1` and index in `\mathbb{N}` associated with
    the full modular group.
    
    TESTS::
    
        sage: from psage.modform.jacobiforms import *
        sage: from psage.modform.jacobiforms.jacobiformd1nn_fourierexpansion import JacobiFormD1NNFilter 
        sage: JR = JacobiFormsD1NN(QQ, JacobiFormD1NN_Gamma(3, 6), JacobiFormD1NNFilter(10, 3))
        sage: JR.gens()
        (Graded expansion TDE_0, Graded expansion TDE_1)
        sage: JR.0 + 2 * JR.1
        Graded expansion TDE_0 + 2*TDE_1
    """
    def __init__(self, index, weight) :
        if weight % 2 != 0 :
            raise NotImplementedError, "Only even weight forms are implemented."
  
        self.__index = index
        self.__weight = weight        
    
    def index(self) :
        return self.__index
    
    def weight(self) :
        return self.__weight
    
    def _ambient_construction_function(self) :
        return JacobiFormsD1NN
    
    def group(self) :
        return "Sp(2, ZZ)_\infty"
    
    @cached_method
    def _rank(self, K) :
        
        if K is QQ or K in NumberFields() :
            return len(_jacobi_forms_by_taylor_expansion_coords(self.__index, self.__weight, 0))

            ## This is the formula used by Poor and Yuen in Paramodular cusp forms
            if self.__weight == 2 :
                delta = len(self.__index.divisors()) // 2 - 1
            else :
                delta = 0
                
            return sum( ModularForms(1, self.__weight + 2 * j).dimension() + j**2 // (4 * self.__index)
                        for j in xrange(self.__index + 1) ) \
                   + delta
            

            ## This is the formula given by Skoruppa in 
            ## Jacobi forms of critical weight and Weil representations
            ##FIXME: There is some mistake here
            if self.__weight % 2 != 0 :
                ## Otherwise the space X(i**(n - 2 k)) is different
                ## See: Skoruppa, Jacobi forms of critical weight and Weil representations
                raise NotImplementedError
            
            m = self.__index
            K = CyclotomicField(24 * m, 'zeta')
            zeta = K.gen(0)
            
            quadform = lambda x : 6 * x**2
            bilinform = lambda x,y : quadform(x + y) - quadform(x) - quadform(y)
            
            T = diagonal_matrix([zeta**quadform(i) for i in xrange(2*m)])
            S =   sum(zeta**(-quadform(x)) for x in xrange(2 * m)) / (2 * m) \
                * matrix([[zeta**(-bilinform(j,i)) for j in xrange(2*m)] for i in xrange(2*m)])
            subspace_matrix_1 = matrix( [ [1 if j == i or j == 2*m - i else 0 for j in xrange(m + 1) ]
                                        for i in xrange(2*m)] )
            subspace_matrix_2 = zero_matrix(ZZ, m + 1, 2*m)
            subspace_matrix_2.set_block(0,0,identity_matrix(m+1))
            
            T = subspace_matrix_2 * T * subspace_matrix_1
            S = subspace_matrix_2 * S * subspace_matrix_1
            
            sqrt3 = (zeta**(4*m) - zeta**(-4*m)) * zeta**(-6*m) 
            rank =   (self.__weight - 1/2 - 1) / 2 * (m + 1) \
                   + 1/8 * (   zeta**(3*m * (2*self.__weight - 1)) * S.trace()
                             + zeta**(3*m * (1 - 2*self.__weight)) * S.trace().conjugate() ) \
                   + 2/(3*sqrt3) * (   zeta**(4 * m * self.__weight) * (S*T).trace()
                                     + zeta**(-4 * m * self.__weight) * (S*T).trace().conjugate() ) \
                   - sum((j**2 % (m+1))/(m+1) -1/2 for j in range(0,m+1))
            
            if self.__weight > 5 / 2 :
                return rank
            else :
                raise NotImplementedError
            
        raise NotImplementedError
    
    @cached_method
    def generators(self, K, precision) :
        if K is QQ or K in NumberFields() :
            return Sequence( [ jacobi_form_by_taylor_expansion(i, self.__index, self.__weight, precision)
                               for i in xrange(self._rank(K)) ],
                             universe = JacobiD1NNFourierExpansionModule(QQ, self.__index) )
        
        raise NotImplementedError
    
    def grading(self, K) :
        if K is QQ or K in NumberFields() :
            return TrivialGrading( self._rank(K),
                                   (self.__index, self.__weight) )
        
        raise NotImplementedError

    def _generator_names(self, K) :
        if K is QQ or K in NumberFields() :
            return [ "TDE_%s" % (i,) for i in xrange(self._rank(K)) ]
        
        raise NotImplementedError

    def _generator_by_name(self, K, name) :
        if K is QQ or K in NumberFields() :
            R = self.generator_relations(K).ring()
            try :
                return R.gen(self._generator_names(K).index(name))
            except ValueError :
                raise ValueError, "name %s doesn't exist for %s" % (name, K)
        
        raise NotImplementedError
    
    @cached_method
    def generator_relations(self, K) :
        r"""
        An ideal I in a polynomial ring R, such that the associated module
        is (R / I)_1. 
        """
        if K is QQ or K in NumberFields() :
            R = PolynomialRing(K, self._generator_names(K))
            return R.ideal(0)
            
        raise NotImplementedError

    def weights(self, K) :
        r"""
        A list of integers corresponding to the weights.
        """
        if K is QQ or K in NumberFields() :
            return len(self._theta_decomposition_indices()) \
                    * [(self.__index, self.__weight)]
            
        raise NotImplementedError
    
    def graded_submodules_are_free(self) :
        return True

    def __cmp__(self, other) :
        c = cmp(type(self), type(other))
        
        if c == 0 :
            c = cmp(self.__index, other.__index)
        if c == 0 :
            c = cmp(self.__weight, other.__weight)
            
        return c

    def __hash__(self) :
        return reduce(xor, map(hash, [type(self), self.__index, self.__weight]))
