r"""
Routines for testing Siegel modular forms

"""
import sage.libs.mpmath.all as mpmath
RF=RealField(2100)
mpmath.mp.dps=mpmath.prec_to_dps(RF.prec())

#attach "siegel_modular_form.sage"

## Unpickle the Igusa generators
PREC=Integer(1500)
if(True):
    try:
        A=SiegelModularForm("IgA.sobj")
        B=SiegelModularForm("IgB.sobj")
        C=SiegelModularForm("IgC.sobj")
        D=SiegelModularForm("IgD.sobj")
    except IOError:
        A,B,C,D = siegel_modular_forms_generators(prec=PREC)
        f=open("IgA.sobj","w"); A.pickle(f); f.close()
        f=open("IgB.sobj","w"); B.pickle(f); f.close()
        f=open("IgC.sobj","w"); C.pickle(f); f.close()
        f=open("IgD.sobj","w"); D.pickle(f); f.close()
        
        # making sure that A,B,C,D are correct defined
        grp=A.group()
        if(grp<>'Sp(2,Z)'):
            raise ValueError, "Igusa generators A,B,C,D incorrect defined, A.group()=%s" %(grp)
        
        else:
            print "A.group()=",grp
            
#from minverse import *

# We define several auxiliary functions.
# Dirichlet series $\sum a(n)n^{-s}$ are represented by dictionaries $n \mapsto a(n)$.
def X_( S, gd,NMAX=100):
       r"""
       Return the Spinor L-function of the Hecke eigenform S
       divided by $\zeta(2s-2k+4)$
       within the precision defined by gd (GRAM_DATA).
       It is assumed that S[(1,1,1)] is nonzero.
       """
       Z_S = dict()
       for N in gd:
           if(N>NMAX):
               break
           try:
               T = HeckeOperator(N)
               TS=T(S)
           except NameError:
               TS=S.hecke_image(N)
           if(TS.keys().count((1,1,1))>=1):
               Z_S[N] = TS[(1,1,1)]/S[(1,1,1)]
           else:
               print "Not high enough precision for this coefficient! N=",N
               Z_S[N]=0
       return Z_S

def jac_coeffs( S, gd,NMAX=100):
       r"""
       Return a dictionary (whose keys are the indices defined in gd)
       of vectors containing the coefficients of the Fourier-Jacobi coeficients
       of the Siegel modular form S as required by gd (GRAM_DATA).
       """
       vec_dict = dict()
       for N in gd:
           if(N>NMAX):
               break
           #print "N=",N
           J = S.fourier_jacobi_coeff(N)
           vec_dict[N] = vector( QQ, [ J[pair] for pair in gd[N][1]])
       return vec_dict

def R_( F, G, gd,NMAX=100,print_error=False):
       r"""
       Using gd (Fredrik's GRAM_DATA) for the computation,
       return the Dirichlet series R_{F,G} as in [Kohnen-Skoruppa]. (i.e. without zeta factor)
       """

       jF = jac_coeffs( F, gd,NMAX)
       jG = jac_coeffs( G, gd,NMAX)
       ls = dict()
       #       fak= RF((4.0*pi.n(300))**(18.5))/gamma(RF(18.5))
       fak=RF(1.0)
       for N in gd:
           if(N>NMAX):
               break
           if(print_error):
               print "||F||_2 [",N,"]=",RR(jF[N].norm(2))
               print "||G||_2 [",N,"]=",RR(jG[N].norm(2))
           ls[N] = RF(jF[N] * gd[N][0] * jG[N])
           ls[N]=ls[N]/RF(sqrt(N))
       return ls
   
def printFGcoeff(F,G,Weight,numd=30,NMAX=100,norm=True):
    r"""
    Prints the (normalized) coefficients of D(F,F)
    with numd bits precision
    """
    DFF=D_(F,G,GRAM_DATA,Weight)
    for N in GRAM_DATA:
        if(N>NMAX):
            break
        tmp=DFF[N]/DFF[1]
        #if(norm):
        #    tmp=tmp/(N**((Weight-RF(0))/RF(2)))
        print tmp.n(digits=numd)


def print_cmp_rel_er(X,Y,NMAX=-1):
    r""" Print X[N]/X[1], Y[N] and the relative error
         |X[N]/X[1]-Y[N]|/|Y[N]|
    """
    if(X.keys().count(1)==0 or X[1]==0):
        raise ValueError, "X[1] not defined or equal to 0!"
    for N in X.keys():
        if( (n > NMAX and NMAX >0) or Y.keys().count(N)==0):
            break
        tmp=X[N]/X[1]
        if(Y[N]<>0):
            tmp2=abs(tmp-Y[N])/abs(Y[N])
        else:
            tmp2=abs(tmp-Y[N])
        print " %s \t %s \t %s \t %s" %(N,tmp.n(100),RF(Y[N]).n(100),tmp2.n(2))


def get_cmp_rel_er(X,Y,NMAX=-1):
    r""" Compute the maximal relative error:
         |X[N]/X[1]-Y[N]|/|Y[N]|
    """
    mer=0
    if(X.keys().count(1)==0 or X[1]==0):
        raise ValueError, "X[1] not defined or equal to 0!"
    for N in X.keys():
        if( (n > NMAX and NMAX >0) or Y.keys().count(N)==0):
            break
        tmp=X[N]/X[1]
        if(Y[N]<>0):
            tmp2=abs(tmp-Y[N])/abs(Y[N])
        else:
            tmp2=abs(tmp-Y[N])
        if(tmp2>mer):
            mer=tmp2
    return mer
    

def print_Lcoeff(L,numd=30,NMAX=-1):
    r"""
    Prints the coefficients a[n]/a[1] of L corresopnding to in gd
    with n<nmax if nmax>0
    with numb bits precision
    """
    if(L.keys().count(1)==0 or L[1]==0):
        raise ValueError, "L[1] not defined or equal to 0!"
    for N in L.keys():
        if(N>NMAX and NMAX >=0 ):
            break
        tmp=L[N]/L[1]
        print tmp.n(digits=numd)


    

def D_(F,G,GRAM_DATA,Weight,NMAX=100):
    r""" Computes the D_{F,G} = R_{F,G} * zeta(2s-2k+4)
    """
    R_FG=R_(F,G,GRAM_DATA)
    A=dict()
    #print "F=",F
    #print "G=",G
#    for N in R_FG.keys():
#        R_FG[N]=R_FG[N]/R_FG[1]
    for N in GRAM_DATA:
        #print "R_FG[",N,"]=",R_FG[N]
        A[N]=R_FG[N]/R_FG[1]
        #print "A[",N,"]=",A[N]

    D=mult_dseries_with_zeta2sk(A,RF(4-2*Weight))
    # also normalize the coefficients
    for N in GRAM_DATA:
        #        D[N]=D[N]/RF(N)**RF(Weight-1.5)/D[1]
        tmp=D[N]/D[1]
        D[N]=tmp/RF(N)**RF(Weight-1.5)
        #print D[N].n(200)
    return D

def D_from_R(R,GRAM_DATA):
    r""" Takes a Raankin product series R_{F,G} and returns
         a normalized D_{F,G} = R_{F,G} * Z(2s-2k+4)
    """
    D=mult_dseries_with_zeta2sk(R,RF(4-2*Weight))
    # also normalize the coefficients
    for N in GRAM_DATA:
        #        D[N]=D[N]/RF(N)**RF(Weight-1.5)/D[1]
        tmp=D[N]/D[1]
        D[N]=tmp/RF(N)**RF(Weight-1.5)
        #print N,D[N].n(200)
    return D
    

def mult_dseries_with_zeta2sk(A,k):
    r""" Multiply the Dirichlet series A(s) with zeta(2s+k)
         given by coefficient vector A
     """
    C=dict()
    Z=dict()
    # Note that zeta(2s+k)=\Sum b(n) n^-s
    # where b(n)=sqrt(n)^-k  if n is a square and 0 otherwise
    for n in A.keys():
 #       print "A[",n,"]=",A[n]
        if(n.is_square()):
            q=n.sqrt()
            Z[n]=q**(-k)
        else:
            Z[n]=0
    for r in A.keys():
        C[r]=0
      #  print "r=",r
        for  n in A.keys():
            if( (r % n) == 0):
      #          print "C[r][",n,"]=+",A[n]*Z[Integer(r/n)]
                C[r]=C[r]+A[n]*Z[Integer(r/n)]
    return C
    
def check_linear_independence(L,LB):
    r""" Check if the Lfunction given by coefficients L is in the linear span of the L-functions in LB
    EXAMPLE:
      L=D(Upsilon20,Upsilon20)
      B=[X_1,X_2]
      where X_1=Z(M1), X_2=Z(M2) are the Maass spec
    """
    return True
    n=len(LB)
    FLD=RF  # LB[0][2].base_ring()
    B=matrix(FLD,n)
    BT=matrix(FLD,1,n)
    A=matrix(FLD,n,1)
    X=matrix(FLD,n,1)
    for j in range(1,n+1):
        A[j-1]=L[j]
        for k in range(1,n+1):
            B[k-1,j-1]=LB[j-1][k]
    [d,BI]=mat_inverse(B)
    print "det(Basis)=",d.n(20)
    if(abs(d)<1e-20):
        print "The Basis sbubmitted is not lin. independent!"
        return False
    if(n+1>len(LB[1])):
       print " ERROR: Not enough element in LB to determine Linear independence!"
       return False
    X=BI*A
    # satisfies f=B*X
    test=RF(0)
    for j in range(n):
        test=test+LB[j][n+1]*X[j,0]
    # is L[n]=test?
    #print "test=",test.n(100)
    #print "L[n]=",  L[n+1].n(100)
    diff=abs(test-L[n+1])
    if(abs(diff)<1e-20):
        # If we are here then we are linearly dependent and we should print the linear combination
        #print "diff=",diff.n(20)
        return [True,diff.n(20)]
    else:
        return [True,diff.n(20)]


def L2_norm(F,gd,nmax):
    r"""
    Compute the maximum L_2 norm of the components F_N of F
    for the indices in gd with N<=nmax
    """
    jF = jac_coeffs( F, gd,nmax)
    dpold=mpmath.mp.dps
    norm=0
    for N in gd:
        if(N>nmax):
            break
        tmp=0
        for j in range(len(jF[N])):
            tmp1=mpmath.mpf(jF[N][j])
            #print "|F(",N,")|=", abs(tmp1)
            tmp+=tmp1**2
        tmp=mpmath.sqrt(tmp)
        mpmath.mp.dps=5
        #print "||F",N,"||_2=", tmp
        mpmath.mp.dps=dpold
        if(tmp>norm):
            norm=tmp
    return norm


def L2_norm_N(F,gd,N):
    r"""
    Compute the maximum L_2 norm of the components F_N of F
    for the indices in gd with N<=nmax
    """
    if(N not in gd):
        return 0
    jF = jac_coeffs( F, gd,N)
    dpold=mpmath.mp.dps
    norm=0
    tmp=0
    for j in range(len(jF[N])):
        tmp1=mpmath.mpf(jF[N][j])
        tmp+=tmp1**2
    norm=mpmath.sqrt(tmp)
    return norm




def check_total_indep(LB,eps2):
    r""" Check independece of all vectors given in LB
         where the error in any coefficient is bounded by eps2 
    EXAMPLE:
      L=R_(Upsilon20,Upsilon20)
      B=[X_1,X_2]
      BL=[X_1,X_2,L]
      norm=L2_norm(F2,GRAM_DATA,len(BL))
      eps2=norm**2*eps1
      check_total_indep(LB,eps2):
      where X_1=Z(M1), X_2=Z(M2) are the Maass spec
    """
    #    FLD=LB[0][1].base_ring()

    dpold=mpmath.mp.dps
    print "dpold=",dpold
    n=len(LB)
    if(len(LB[0])<n):
        print " ERROR: Not enough elements to determine Linear independence!"
        return [False,-1]
    B=mpmath.matrix(int(n),int(n))
    for r in range(0,n):
        for k in range(1,n+1):
            tmp1=mpmath.mpf(LB[r][k])/mpmath.mpf(LB[r][1])    
            tmp=mpmath.mpi(tmp1-eps2,tmp1+eps2)
            B[r,k-1]=tmp 
    d=mpmath.mpi(mpmath.det(B))
    mpmath.mp.dps=10
    print "eps2=",eps2
    print "det(B)=",d
    print "mid(det((B))=",d.mid
    print "diam(det((B))=",d.delta
    print "||B||=",mpmath.norm(B)
    mpmath.mp.dps=dpold
    zero=mpmath.mpi(0)
    if( zero in d):
        print "Functions are linearly dependent!"
        return [False,d]        
    else:
        print "Functions are linearly independent!"
        return [True,d]        





def check_total_indep2(LB,eps2):
    r""" Check independece of all vectors given in LB
         where the error in any coefficient LB[i][j] is bounded by eps2[i][j]
         here 
    EXAMPLE:
      L=R_(Upsilon20,Upsilon20)
      B=[X_1,X_2]
      BL=[X_1,X_2,L]
      norm=L2_norm(F2,GRAM_DATA,len(BL))
      eps2=norm**2*eps1
      check_total_indep(LB,eps2):
      where X_1=Z(M1), X_2=Z(M2) are the Maass spec
    """
    #    FLD=LB[0][1].base_ring()

    dpold=mpmath.mp.dps
    n=len(LB)
    if(len(LB[0])<n):
        print " ERROR: Not enough elements to determine Linear independence!"
        return [False,-1]
    B=mpmath.matrix(int(n),int(n))
    for r in range(0,n):
        for k in range(1,n+1):
            tmp1=mpmath.mpf(LB[r][k])/mpmath.mpf(LB[r][1])    
            tmp=mpmath.mpi(tmp1-eps2[r][k],tmp1+eps2[r][k])
            B[r,k-1]=tmp 
    d=mpmath.mpi(mpmath.det(B))
    mpmath.mp.dps=10
    print "eps2=",eps2
    print "det(B)=",d
    print "mid(det((B))=",d.mid
    print "diam(det((B))=",d.delta
    print "||B||=",mpmath.norm(B)
    mpmath.mp.dps=dpold
    zero=mpmath.mpi(0)
    if( zero in d):
        print "Functions are linearly dependent!"
        return [False,d]        
    else:
        print "Functions are linearly independent!"
        return [True,d]        







def test_hecke_form(FF,n):
    r""" test that the form FF is a Hecke eigenform for T(n)
    """
    try: 
        T=HeckeOperator(n)
        TFF=T(FF)
    except NameError:
        TFF=FF.hecke_image(n)
    ev=TFF.coeffs()[(1,1,1)]/FF.coeffs()[(1,1,1)]
    #print "(tentative) Hecke eigenvalue = ",ev
    for t in TFF.coeffs():
        ac=FF.coeffs()[t]
        bc=TFF.coeffs()[t]
        dc=bc-ac*ev
        #print t,d
    if(dc>1E-20):
            return False
    return True


def jac_2norm(F,gd):
    jF = jac_coeffs( F, gd)
    for N in jF.keys():
        print N,(RF(jF[N].norm(2))).n(20)

def mppr(x):
    r"""
    Print fewer digits!

    """
    
    # We don't want to print 500 digits
    # So I truncate at two digits with a stupid simple rounding
    dpold=mpmath.mp.dps
    mpmath.mp.dps=5
    s=str(x)
    mpmath.mp.dps=dpold
    (s1,dot,s2)=s.partition(".")
    (s3,es,s4)=s2.partition("e")
    if(len(s4)==0):  # No "e" format
        return s[0:15]
    try:
        ma=s3[0]
        if(len(s3)>1):
            if(s3[2]<5):
                ma=ma+s3[1]
            elif(s3[2]>=5):
                ma=ma+str(ZZ(s3[1])+1)
        ss=s1+"."+ma+"e"+s4
    except:
        print "index error: s3=",s3
        return str(x)
    return ss


def proper_eps2(Fv,R_FFv,GRAM_DATA,NMAX,eps1):
    r"""
    Compute the true epsilon_2 
    """
    ## Compute the estimated factor for each F_j:
    eps2=dict()
    for i in range(len(Fv)):
        d1=abs(R_FFv[i][1])
        n1=L2_norm_N(Fv[i],GRAM_DATA,1)**2
        print "|<F_1|F_1>|(",i,")=",d1
        #d2=d1-eps1*L2_norm_N(Fv[i],GRAM_DATA,1)**2
        for N in range(1,NMAX+1):        
            nN=L2_norm_N(Fv[i],GRAM_DATA,N)**2
            dN=abs(R_FFv[i][N])
            tmp=(n1*dN+nN*(mpmath.mpf(1)+eps1))/abs(d1-eps1*n1)
            eps2[i,N]=eps1*tmp/n1
            print "eps2[",i,",",N,"]=",mppr(eps2[i,N])
    e2=max(eps2.values())
    print "max eps2=",mppr(e2)
    return e2



def true_eps2_matrix(Fv,R_FFv,GRAM_DATA,NMAX,n_exact,n_est,eps1):
    r"""
    Compute the true epsilon_2
    we assume the first n_exact functions are computed exactly
    """
    ## Compute the estimated factor for each F_j:
    eps2=dict()
    if(n_exact+n_est <> NMAX):
        raise ArithmeticError,"Dimension doesn't match!"
    for i in range(n_exact):
        for j in range(NMAX+1):
            eps2[i,j]=mpmath.mpf(0)
    for N in range(1,NMAX+1):        
        maxn=0
        for i in range(len(Fv)):
            d1=abs(R_FFv[i][1])
            n1=L2_norm_N(Fv[i],GRAM_DATA,1)**2
            #print "|<F_1|F_1>|(",i,")=",d1
            nN=L2_norm_N(Fv[i],GRAM_DATA,N)**2
            dN=abs(R_FFv[i][N])
            tmp=(n1*dN+nN*(mpmath.mpf(1)+eps1))/abs(d1-eps1*n1)
            eps2tmp=eps1*tmp/n1
            #print "eps2[",i,",",N,"]=",mppr(eps2[i,N])
            if(eps2tmp>maxn):
                maxn=eps2tmp
        print "max[",N,"]=",maxn
        for i in range(n_exact,NMAX):
            eps2[i,N]=maxn
    #e2=max(eps2.values())
    #print "max eps2=",mppr(e2)
    return eps2