#code to calculate local L functions, copied from the program PSS

def local_normal_form_with_change_vars(Q,p):
    #basically copied from the built-in function local_normal_form(), but it also computes a change-of-variables to the local normal form
    I = range(Q.dim())
    assert p != 2
    M = identity_matrix(QQ,Q.dim())
    Q_Jordan = DiagonalQuadraticForm(ZZ,[])
    while Q.dim() > 0:
        n = Q.dim()
        (min_i, min_j) = Q.find_entry_with_minimal_scale_at_prime(p)
        if min_i == min_j:
            min_val = valuation(2 * Q[min_i, min_j], p)
        else:
            min_val = valuation(Q[min_i, min_j], p)
        if (min_i == min_j):
            Q.swap_variables(0, min_i, in_place = True)
            M.swap_rows(I[0],I[min_i])
        else:
            Q.swap_variables(0, min_i, in_place = True)
            Q.swap_variables(1, min_j, in_place = True)
            M.swap_rows(I[0],I[min_i])
            M.swap_rows(I[1],I[min_j])
            Q.add_symmetric(1, 0, 1, in_place = True)
            M.add_multiple_of_row(I[1],I[0],-1)
        min_scale = p ** min_val
        a = 2 * Q[0,0]
        for j in range(1, n):
            b = Q[0, j]
            g = GCD(a, b)
            Q.multiply_variable(ZZ(a/g), j, in_place = True)
            Q.add_symmetric(ZZ(-b/g), j, 0, in_place = True)
            M.rescale_row(I[j],g/a)
            M.add_multiple_of_row(I[0],I[j],b/g)
        Q_Jordan = Q_Jordan + Q.extract_variables(range(1))
        I.remove(I[0])
        Q = Q.extract_variables(range(1, n))
    return [Q_Jordan,M]


def IsospectralNormalForm(Q,L,p):
    copyQ = QuadraticForm(Q.matrix())
    [D,P] = local_normal_form_with_change_vars(copyQ,p)
    L = L*P.inverse()
    lins = []
    quads = []
    consts = []
    for j in range(Q.dim()):
        a = D[j,j]
        b = L[0,j] / 2
        if b.valuation(p) < a.valuation(p):
            lins.append(b)
        else:
            quads.append(a)
            consts.append(-b^2 / (4*a))
    const = sum(consts)
    if len(lins) > 0:
        linsgcd = gcd(lins)
        return [quads,linsgcd,const]
    else:
        return [quads,0,const]


def IARD(a,r,d,p):
    assert p != 2
    R.<t> = PolynomialRing(QQ)
    if r % 2 == 0:
        kron = kronecker_symbol((-1)^(r/2) * d,p)
        iard = 1 - p^(-r/2) * kron
        if a % p == 0:
            iard = iard * (1 + p^(-r/2) * kron * t) *(1 - 1/p) / (1 - t/p)
        else:
            iard = iard * ((1 - 1/p) / (1 - t/p) + p^(-r/2) * kron)
    else:
        iard = 1
        if a % p == 0:
            iard = iard * (1 - t*p^(-r)) * (1 - p^(-1)) / (1 - p^(-1) * t)
        else:
            kron = kronecker_symbol(a*d*(-1)^((1+r)/2),p)
            iard = iard * (1 + p^(-(1+r)/2)*kron*t)*(1 - p^(-1)) / (1 - p^(-1) * t) - p^(-r) - p^(-(1+r)/2) * kron
    return iard

def IgusaZeta(Q,L,c,p):
    [wrongquads,linsgcd,const] = IsospectralNormalForm(Q,L,p)
    c = const + c/2
    u = denominator(c/2)
    c = c*u
    linsgcd = linsgcd*u
    quads = []
    for a in wrongquads:
        quads.append(a*u)
    R.<t> = PolynomialRing(QQ)
    Z = 0*t
    dold = []
    rold = []
    for j in range(len(quads)):
        a = quads[j]
        pi = a.valuation(p)
        while pi > len(dold) - 1:
            dold.append(1)
            rold.append(0)
        dold[pi] = dold[pi]* ZZ(a / p^pi)
        rold[pi] = rold[pi] + 1
    d = []
    r = []
    qpower = []
    if linsgcd == 0:
        if c != 0:
            maxrange = max(len(dold),c.valuation(p)+1)
        else:
            maxrange = len(dold)
    else:
        if c != 0:
            maxrange = max(len(dold),c.valuation(p)+1,linsgcd.valuation(p))
        else:
            maxrange = max(len(dold),linsgcd.valuation(p))
    for k in range(maxrange+1):
        if k >= len(dold):
            dold.append(1)
            rold.append(0)
        l = k
        while l <= maxrange:
            if l > len(d) - 1:
                d.append(1)
                r.append(0)
                qpower.append(0)
            if (l % 2 == k % 2):
                d[l] = d[l] * dold[k]
                r[l] = r[l] + rold[k]
            if l != k:
                qpower[l] = qpower[l] + r[k]
            l = l + 1
    w = len(dold) - 1
    if linsgcd == 0 and c == 0:
            R = sum(rold)
            Z = (t^(w-1)*IARD(0,r[w-1],d[w-1],p) / p^(qpower[w-1]) + t^w *IARD(0,r[w],d[w],p) / p^(qpower[w]) )/ (1 - t^2 / p^R)
            for i in range(w-1):
                Z = Z + IARD(0,r[i],d[i],p)*t^i / p^(qpower[i])
            return Z
    elif linsgcd.valuation(p) > c.valuation(p):
        kappa = c.valuation(p)
        Z = t^kappa / p^(qpower[kappa + 1])
        for j in range(kappa + 1):
            Z = Z + t^j * IARD(ZZ(c/p^j),r[j],d[j],p) / p^(qpower[j])
        return Z
    elif linsgcd != 0:
        lamda = linsgcd.valuation(p)
        Z = (t^lamda / p^(qpower[lamda])) * ((1 - 1/p) / (1 - t/p))
        for i in range(lamda):
            Z = Z + t^i * IARD(0,r[i],d[i],p) / p^(qpower[i])
        return Z
    
    

def TwoadicClassify(Q):
    J = twonf_with_change_vars_correct(Q)
    J = J[0].matrix()
    #Classifies unimodular quadratic forms over Z2. Every form is equivalent to a direct sum of at most two squares ax^2, and at most one "elliptic plane" 2(x^2 + xy + y^2), and any number of hyperbolic planes 2xy.
    Squares = []
    Ell = 0
    Hyp = 0
    k = 0
    while k < J.nrows():
        if k < J.nrows() - 1:
            if J[k,k+1] != 0:
                if J[k,k]*J[k+1,k+1] % 8 == 0:
                    Hyp = Hyp + 1
                else:
                    Ell = Ell + 1
                k = k + 2
            else:
                Squares.append((J[k,k] / 2) % 8)
                Squares.append((J[k+1,k+1] / 2) % 8)
                k = k + 2
        else:
            Squares.append((J[k,k] / 2) % 8)
            k = k + 1
    while len(Squares) > 2:
        [a,b,c] = sorted(Squares[0:3])
        Squares = Squares[3:len(Squares)]
        if [a,b,c] in [[1,3,5],[1,1,7],[5,5,7]]:
            Squares.append(1)
            Hyp = Hyp + 1
        elif [a,b,c] in [[1,3,7],[3,3,5],[5,7,7]]:
            Squares.append(3)
            Hyp = Hyp + 1
        elif [a,b,c] in [[1,5,7],[3,5,5],[1,1,3]]:
            Squares.append(5)
            Hyp = Hyp + 1
        elif [a,b,c] in [[3,5,7],[1,7,7],[1,3,3]]:
            Squares.append(7)
            Hyp = Hyp + 1
        elif [a,b,c] in [[3,3,3],[3,7,7]]:
            Squares.append(1)
            Ell = Ell + 1
        elif [a,b,c] in [[1,1,1],[1,5,5]]:
            Squares.append(3)
            Ell = Ell + 1
        elif [a,b,c] in [[7,7,7],[3,3,7]]:
            Squares.append(5)
            Ell = Ell + 1
        elif [a,b,c] in [[5,5,5],[1,1,5]]:
            Squares.append(7)
            Ell = Ell + 1
    r = floor(Ell/2)
    Hyp = Hyp + 2*r
    Ell = Ell % 2
    return [Squares,Ell,Hyp]
    
def IaQQQrank(Q0):
    return len(Q0[0]) + 2*Q0[1] + 2*Q0[2]
    
def Ig(a,u,j):
    R.<t> = PolynomialRing(QQ)
    if u == Infinity:
        if a.valuation(2) >= j:
            return t^j / (2-t)
        else:
            return t^(a.valuation(2))
    else:
        return Ig(a,Infinity,min(u,j))
    
def hatHQ(a,u,Q0): #H1
    R.<t> = PolynomialRing(QQ)
    r = IaQQQrank(Q0)
    if len(Q0[0]) == 0:
        return (1 - 2^(-r))*Ig(a,u,1)
    else:
        return Ig(a,u,0) - 2^(-r)*Ig(a,u,1)
    
def tildeHQ(a,u,Q0): #H2
    R.<t> = PolynomialRing(QQ)
    eps = (-1)^Q0[1]
    r = IaQQQrank(Q0)
    if len(Q0[0]) == 0:
        return (1 - eps*2^(-r/2))*(Ig(a,u,1) + eps*2^(-r/2)*Ig(a,u,2))
    elif len(Q0[0]) == 1:
        b = Q0[0][0]
        return (1 - eps*2^(-(r-1)/2))*Ig(a,u,0) - 2^(-r) * Ig(a,u,2) + eps*2^(-(r+1)/2) * (Ig(a,u,2) + Ig(a+b,u,2))
    else:
        b = Q0[0][0]
        c = Q0[0][1]
        if (b+c) % 4 == 0:
            return Ig(a,u,0) - eps*2^(-r/2) * Ig(a,u,1) + (eps*2^(-r/2) - 2^(-r)) * Ig(a,u,2)
        else:
            return (1 - eps*2^(-(r-2)/2))*Ig(a,u,0) + eps*2^(-r/2)*Ig(a,u,1) + eps*2^(-r/2)*Ig(a+b,u,2) - 2^(-r)*Ig(a,u,2)
    
def tildeHQdiff(a,u,Q0): #H3
    R.<t> = PolynomialRing(QQ)
    eps = (-1)^Q0[1]
    r = IaQQQrank(Q0)
    if len(Q0[0]) == 0:
        return t - t
    if len(Q0[0]) == 1:
        b = Q0[0][0]
        return 2^(-r)*(Ig(a+b,u,3) - Ig(a+b,u,2))
    if len(Q0[0]) == 2:
        b = Q0[0][0]
        c = Q0[0][1]
        if (b+c) % 4 == 0:
            sig = (-1)^((b+c)/4)
            return sig * 2^(-r) * (Ig(a,u,3) - Ig(a,u,2))
        else:
            return -2^(1-r) * Ig(a,u,1) + 2^(-r) * (Ig(a,u,2) + Ig(a+b+c,u,3))
    

def IaQQQ(a,u,U0,U1,U2):
    Q0 = TwoadicClassify(U0)
    Q1 = TwoadicClassify(U1)
    Q2 = TwoadicClassify(U2)
    R.<t> = PolynomialRing(QQ)
    if len(Q1[0]) > 0 and len(Q2[0]) > 0:
        return hatHQ(a,u,Q0)
    elif len(Q2[0]) > 0:
        return tildeHQ(a,u,Q0)
    eps1 = (-1)^Q1[1]
    r1 = IaQQQrank(Q1)
    if len(Q1[0]) == 0:
        return tildeHQ(a,u,Q0) + eps1*2^(-r1/2) * tildeHQdiff(a,u,Q0)
    elif len(Q1[0]) == 1:
        c = Q1[0][0]
        return hatHQ(a,u,Q0) + eps1*2^(-(r1+1)/2) * (tildeHQdiff(a,u,Q0) + tildeHQdiff(a+2*c,u,Q0))
    elif len(Q1[0]) == 2:
        c = Q1[0][0]
        d = Q1[0][1]
        if (c + d) % 4 == 0:
            return hatHQ(a,u,Q0) + eps1 * 2^(-r1/2) * tildeHQdiff(a,u,Q0)
        else:
            return hatHQ(a,u,Q0) + eps1 * 2^(-r1/2) * tildeHQdiff(a+2*c,u,Q0)


def twonf_with_change_vars_correct(Q):
    #basically copied from the built-in function local_normal_form(), specialized to p=2, but it also computes a change-of-variables to the local normal form. it skips the "Cassels' proof" step. this is accounted for later.
    #for large lattices this seems to be by far the slowest part of the algorithm. local_normal_form(2) is also slow (for example, try using it on the Leech lattice). this could probably be fixed
    I = range(Q.dim())
    M = identity_matrix(Q.dim())
    Q_Jordan = DiagonalQuadraticForm(ZZ,[])
    while Q.dim() > 0:
        n = Q.dim()
        (min_i,min_j) = Q.find_entry_with_minimal_scale_at_prime(2)
        if min_i == min_j:
            min_val = valuation(2 * Q[min_i,min_j],2)
        else:
            min_val = valuation(Q[min_i, min_j],2)
        if (min_i == min_j):
            block_size = 1
            Q.swap_variables(0,min_i,in_place = True)
            M.swap_rows(I[0],I[min_i])
        else:
            Q.swap_variables(0, min_i, in_place = True)
            Q.swap_variables(1, min_j, in_place = True)
            M.swap_rows(I[0],I[min_i])
            M.swap_rows(I[1],I[min_j])
            block_size = 2
        min_scale = 2^(min_val)
        if (block_size == 1):
            a = 2 * Q[0,0]
            for j in range(block_size, n):
                b = Q[0, j]
                g = GCD(a, b)
                Q.multiply_variable(ZZ(a/g), j, in_place = True)
                Q.add_symmetric(ZZ(-b/g), j, 0, in_place = True)
                P = identity_matrix(QQ,M.nrows())
                P[I[j],I[j]] = g/a
                M = P*M
                P = identity_matrix(ZZ,M.nrows())
                P[I[0],I[j]] = ZZ(b/g)
                M = P*M
        else:
            a1 = 2*Q[0,0]
            a2 = Q[0,1]
            b2 = 2*Q[1,1]
            big_det = (a1*b2 - a2^2)
            small_det = big_det / (min_scale * min_scale)
            for j in range(block_size,n):
                a = Q[0,j]
                b = Q[1,j]
                Q.multiply_variable(big_det,j,in_place = True)
                Q.add_symmetric(ZZ(-(a*b2 - b*a2)),j,0,in_place = True)
                Q.add_symmetric(ZZ(-(-a*a2 + b*a1)),j,1,in_place = True)
                Q.divide_variable(ZZ(min_scale^2),j,in_place = True)
                P = identity_matrix(QQ,M.nrows())
                P[I[j],I[j]] = 1/big_det
                M = P*M
                P = identity_matrix(M.nrows())
                P[I[0],I[j]] = (a*b2 - b*a2)
                M = P*M
                P = identity_matrix(M.nrows())
                P[I[1],I[j]] = (-a*a2 + b*a1)
                M = P*M
                P = identity_matrix(M.nrows())
                P[I[j],I[j]] = min_scale^2
                M = P*M
        Q_Jordan = Q_Jordan + Q.extract_variables(range(block_size))
        for j in range(block_size):
            I.remove(I[0])
        Q = Q.extract_variables(range(block_size, n))
    return [Q_Jordan,M]


def TwoadicIsospectralNormalForm(Q,L):
    copyQ = QuadraticForm(Q.matrix())
    [D,P] = twonf_with_change_vars_correct(copyQ)
    L = L*P.inverse()
    lins = []
    NewQ = DiagonalQuadraticForm(ZZ,[])
    consts = []
    j = 0
    while j < Q.dim():
        if j < Q.dim() - 1 and D[j,j+1] != 0:
            l = L[0,j]
            m = L[0,j+1]
            A = 2*D[j,j]
            B = 2*D[j,j+1]
            C = 2*D[j+1,j+1]
            assert min(A.valuation(2),C.valuation(2)) >= B.valuation(2)
            if min(l.valuation(2),m.valuation(2)) >= B.valuation(2):
                NewQ = NewQ + D.extract_variables([j,j+1])
                U = matrix([[2*A,B],[B,2*C]])
                v = U.inverse()*vector([-l,-m])
                consts.append(A*v[0]^2 + B*v[0]*v[1] + C*v[1]^2 + l*v[0] + m*v[1])
            else:
                lins.append(l)
                lins.append(m)
            j = j + 2
        else:
            a = 2*D[j,j]
            b = L[0,j]
            if b.valuation(2) < a.valuation(2):
                lins.append(b)
            elif b.valuation(2) == a.valuation(2):
                lins.append(2*b)
            else:
                NewQ = NewQ + D.extract_variables([j])
                consts.append(-b^2 / (4*a))
            j = j + 1
    const = sum(consts)
    if len(lins) > 0:
        linsgcd = gcd(lins)
        return [NewQ,linsgcd,const]
    else:
        return [NewQ,0,const]


def twoadic_jordan_blocks_correct(Q):
    Q1 = QuadraticForm(Q.matrix())
    Jordans = []
    Jordanlist = []
    valuations = []
    i = 0
    n = Q.dim()
    while (i < n):
        if (i == n-1) or (Q1[i,i+1] == 0):
            block_size = 1
        else:
            block_size = 2
        if block_size == 1:
            scale = valuation(Q1[i,i], 2) + 1
            if len(valuations) > 0 and scale == valuations[len(valuations) -1]:
                e = valuations.index(scale)
                Jordans[e] = Jordans[e] + QuadraticForm(ZZ,Q1.extract_variables([i]).matrix() / 2^(scale - 1))
            else:
                valuations.append(scale)
                Jordans.append(QuadraticForm(ZZ,Q1.extract_variables([i]).matrix() / 2^(scale - 1)))
            i = i + 1
        else:
            scale = valuation(Q1[i,i+1], 2)
            if len(valuations) > 0 and scale == valuations[len(valuations) -1]:
                e = valuations.index(scale)
                Jordans[e] = Jordans[e] + QuadraticForm(ZZ,Q1.extract_variables([i,i+1]).matrix() / 2^(scale))
            else:
                valuations.append(scale)
                Jordans.append(QuadraticForm(ZZ,Q1.extract_variables([i,i+1]).matrix() / 2^(scale)))
            i = i + 2
    for i in range(len(valuations)):
        Jordanlist.append([valuations[i],Jordans[i]])
    return Jordanlist


def TwoadicIgusaZeta(Q,L,c):
    bigr = Q.dim()
    [NewQ,linsgcd,const] = TwoadicIsospectralNormalForm(Q,L)
    const = const + c
    R.<t> = PolynomialRing(QQ)
    Z = 0*t
    jordanblocks = twoadic_jordan_blocks_correct(NewQ)
    Qs = []
    vb = linsgcd.valuation(2)
    vc = const.valuation(2)
    if len(jordanblocks) > 0:
        maxrange = jordanblocks[len(jordanblocks)-1][0] + 1
    else:
        maxrange = 2
    if linsgcd != 0:
        maxrange = max(maxrange,vb+1)
    if const != 0:
        maxrange = max(maxrange,vc+3)
    maxrange = max(maxrange,2)
    jordans = [DiagonalQuadraticForm(ZZ,[])] * maxrange
    for k in range(len(jordanblocks)):
        jordans[jordanblocks[k][0]] = jordans[jordanblocks[k][0]] + jordanblocks[k][1]
    rs = []
    Zeroform = DiagonalQuadraticForm(ZZ,[])
    Sqform = DiagonalQuadraticForm(ZZ,[1])
    qpowers = []
    for k in range(maxrange):
        Qs.append(DiagonalQuadraticForm(ZZ,[]))
        rs.append(0)
        l = k
        while l >= 0:
            Qs[k] = Qs[k] + jordans[l]
            rs[k] = rs[k] + jordans[l].dim()
            l = l - 2
        if k == 0:
            qpowers.append(1)
        else:
            qpowers.append(qpowers[k-1] * 2^(rs[k-1]))
    if vb == Infinity and vc == Infinity:
        r = sum(rs)
        w = max(len(Qs) - 1,1)
        Z = (t^(w-1) * IaQQQ(0,Infinity,Qs[w-1],Qs[w],Zeroform) * qpowers[w-1]^(-1) + t^w * IaQQQ(0,Infinity,Qs[w],Qs[w-1],Zeroform) * qpowers[w]^(-1)) / (1 - t^2 * 2^(-bigr))
        for i in range(0,w-1):
            Z = Z + t^i * IaQQQ(0,Infinity,Qs[i],Qs[i+1],jordans[i+2]) / qpowers[i]
        return Z
    elif vb < Infinity and vb <= vc:
        Z = t^vb / (qpowers[vb] * (2 - t))
        for i in range(vb - 2):
            Z = Z + t^i * IaQQQ(0,Infinity,Qs[i],Qs[i+1],jordans[i+2]) / qpowers[i]
        if vb >= 2:
            Z = Z + t^(vb-2) * IaQQQ(0,2,Qs[vb-2],Qs[vb-1],Sqform) / qpowers[vb-2]
        if vb >= 1:
            Z = Z + t^(vb-1) * IaQQQ(0,1,Qs[vb-1],Sqform,Sqform) / qpowers[vb-1]
        return Z
    elif vb < Infinity and vc < vb and vb <= vc + 2:
        Z = t^vc / qpowers[vc+1]
        for i in range(vb - 2):
            Z = Z + t^i * IaQQQ(const/2^i,Infinity,Qs[i],Qs[i+1],jordans[i+2]) / qpowers[i]
        if vb >= 2:
            Z = Z + t^(vb-2) * IaQQQ(const/2^(vb-2),2,Qs[vb-2],Qs[vb-1],Sqform) / qpowers[vb-2]
        if vb >= 1 and vb <= vc + 1:
            Z = Z + t^(vb-1) * IaQQQ(const/2^(vb-1),1,Qs[vb-1],Sqform,Sqform) / qpowers[vb-1]
        return Z
    else:
        Z = t^vc / qpowers[vc+1]
        for i in range(vc+1):
            Z = Z + t^i * IaQQQ(const/2^i,Infinity,Qs[i],Qs[i+1],jordans[i+2]) / qpowers[i]
        return Z

def L_Function_t(n,g,S,p):
    if not is_prime(p):
        raise ValueError('p must be prime')
    R.<t> = PolynomialRing(QQ)
    e = S.nrows()
    L = 2*matrix(g*S)
    c = 2*n + g*S*g
    if p == 2:
        f = (1 - TwoadicIgusaZeta(QuadraticForm(S),L,c)) / (1 - t)
        return f(2^e * t)
    else:
        f = (1 - t*IgusaZeta(QuadraticForm(S),L,c,p)) / (1 - t)
        return f(p^e * t)

def L_Function(n,g,S,p):
    R.<t> = PolynomialRing(QQ)
    s = var('s')
    L = L_Function_t(n,g,S,p)
    return L.subs({t:p^(-s)}).simplify_full()

def L_Function_Series(n,g,S,p,prec=20):
    PS.<q> = PowerSeriesRing(QQ,prec)
    return PS(L_Function_t(n,g,S,p))

#EXAMPLE

#L_Function(n,gamma,S,p) returns the local L-function L_p(n,gamma,s). Here S is the Gram matrix of a quadratic form.


#Here are local L_functions from example 12 in "Vector-valued Eisenstein series of small weight". 
S = matrix([[-2,-1],[-1,-2]])
n = 0
g = vector([0,0])
L_Function(n,g,S,2)
L_Function(n,g,S,3)

#L_Function_Series writes the local L-function as a power series in the variable q = p^(-s). Useful for convincing yourself that the program computes things correctly.
L_Function_Series(n,g,S,2)
#check that the coefficient of q^6 is correct "by hand"
count = 0
for a in range(2^6):
    for b in range(2^6):
        if (-a^2 - a*b - b^2) % 2^6 == 0:
            count = count + 1
print count