VERB = 0 # verbosity


def Wmatrix(S, Q):
    # compute action of W_Q on S
    # S should be a space of cusp forms

    N = S.level()
    Q = ZZ(Q)
    if not Q.divides(N): raise ValueError("Q should divide N")
    Q = N // N.prime_to_m_part(Q)
    
    # reduce to prime power case
    if Q == 1: return identity_matrix(QQ, S.rank())
    elif len(Q.prime_divisors()) > 1:
        p,e = Q.factor()[0]
        Q0 = p**e
        Q1 = Q//Q0
        W0 = Wmatrix(S, Q0).change_ring(CyclotomicField(Q))
        W1 = Wmatrix(S, Q1).change_ring(CyclotomicField(Q))
        return S.diamond_bracket_matrix(crt(Q0, 1, Q1, N//Q1)) * W0 * W1

    p = Q.prime_divisors()[0]

    vecs = []
    matrows = []
    for r in [0..N.valuation(p)]:
        rvecs, rmat = Wmatrix_oldspace(S, p, r)
        matrows += [[0]*len(vecs) + list(v) + [0]*(S.dimension() - rmat.ncols() - len(vecs)) for v in rmat.rows()]
        vecs += rvecs
    basmat = matrix(vecs, ncols=S.ambient().rank())
    wmat = matrix(matrows)

    # convert to echelon form
    to_ech = basmat.solve_left(S.basis_matrix())
    verbose("Wmatrix done",level=VERB)
    return to_ech*wmat*~to_ech

def Wmatrix_oldspace(S, p, r):
    # compute a basis for the part of S coming from p-new forms of level N/p^r,
    # and the matrix of W_p on this basis. 
    # basis is given in terms of S.ambient() (not S itself!)

    k = S.weight()
    N = S.level()

    t=verbose("W_%s on level %s oldspace from %s-new forms of level %s" % (p**N.valuation(p), N, p, N//p**r), level=VERB)
    A = S.ambient()

    if S.character() is None:
        Ar = ModularForms(S.group().restrict(N//p**r), S.weight())
    else:
        if not S.character().conductor().divides(N//p**r):
            return (zero_matrix(ZZ, 0, A.rank()), zero_matrix(ZZ, 0, 0))
        else:
            Ar = ModularForms(S.character().restrict(N//p**r), S.weight())
    Sr = Ar.cuspidal_submodule()

    verbose("computing %s-new submodule at level %s" % (p, Sr.group()._latex_().replace('\\', '')), level=VERB)
    t=verbose("dimension should be %s" % dimension_new_cusp_forms(Sr.character() or Sr.group(), S.weight(), p), level=VERB)
    verbose("Sturm bound is %s" % Sr.sturm_bound(), level=VERB)
    if S.level().valuation(p) > r:
        Sr_new = Sr.new_submodule(p)
    else:
        Sr_new = Sr
    verbose("done", t, level=VERB)

    dim = Sr_new.dimension()
    if dim == 0:
        verbose("dimension is 0", level=VERB)
        return (zero_matrix(ZZ, 0, A.rank()), zero_matrix(ZZ, 0, 0))
    
    Sr_new_mat = Wmatrix_newspace(Sr, p)

    if r == 0:
        basmat0 = matrix(f.element() for f in Sr_new.basis())
        wmat0 = Sr_new_mat
        return basmat0, wmat0

    my_Sold_basis = []
    my_Sold_matrix_rows = []
    pinv = ~Zmod(N)(crt(p, 1, N.prime_to_m_part(p), p**(N.valuation(p))))
    
    # this next part is quite slow for some reason

    t=verbose("assembling matrix at level %s" % N, level=VERB)
    # want to compute Sr_new.diamond_bracket_matrix(pinv) but work around a bug
    verbose(" computing diamond operator", level=VERB)
    Dpinv = p**(k-2) * (Sr.diamond_bracket_matrix(pinv).block_sum(zero_matrix(Ar._dim_eisenstein()))).restrict(Sr_new.free_module())
    verbose(" ...done", cputime()-t, level=VERB)
    for s in [0..r]:
        verbose("computing degen map p^%s" % s,level=VERB)
        D = Ar.degeneracy_map(A, p**s)
        for f in Sr_new.basis():
            verbose(" applying to a basis vector", level=VERB) 
            my_Sold_basis.append(p**(s*(k-1))*A(D(f), check=False).element())
        frag = (Dpinv**s)*Sr_new_mat
        my_Sold_matrix_rows += [ [0]*((r-s)*dim) + list(v) + [0]*(s*dim) for v in frag.rows()]
    basmat = matrix(my_Sold_basis)
    wmat = matrix(my_Sold_matrix_rows)
    verbose("...done", t, level=VERB)
    return basmat, wmat

def Wmatrix_newspace(S, p):
    # compute matrix of W on the echelon-form basis of p-new cuspforms.
    assert S == S.ambient().cuspidal_submodule()
    N = S.level()
    r = N.valuation(p)
    Q = p**r
    t0 = verbose("W_%s on newspace at level %s" % (Q, N), level=VERB) 
    
    if r == 0:
        # everything is "p-new", and Wp is trivial
        return identity_matrix(S.rank())
    t=verbose("Computing mod symbols", level=VERB) 
    symbs = S.ambient().modular_symbols(sign=1)
    nsymbs = symbs.cuspidal_submodule().new_submodule(p)
    t=verbose("Done", t, level=VERB)

    cond_parts = decompose_by_conductor(nsymbs, p)
    msg = ", ".join( "%s: %s" % (c, cond_parts[c].rank()) for c in cond_parts.keys())
    verbose("Decomposed %s-new space at level %s into conductor parts, ranks: " % (p, N) + msg, t, level=VERB)
   
    basis = []
    matrix_rows = []
    from sage.modular.modform.cuspidal_submodule import _convert_matrix_from_modsyms
    for c in cond_parts.keys():
        verbose("Computing with conductor p^%s part" % c, level=VERB)
        t=verbose("Computing WR", level=VERB)
        lpad = len(basis)
        rpad = nsymbs.rank() - lpad - cond_parts[c].rank()

        # First compute the matrix of W*R, where R is a twisting operator such
        # that W*R is defined over QQ and commutes with the Hecke operators.
        if c < r:
            W = S.group().atkin_lehner_matrix(Q)
            if c == 0: 
                L = [W]
            else:
                L = [ ]
                for a in xsrange(p**c):
                    if a.gcd(p) > 1: continue
                    aa = crt(a, 1, Q, N//Q)
                    diam = matrix(ZZ, 2, lift_to_sl2z(0, aa, N))
                    L.append( (W * diam * matrix(QQ, 2, [1,-a/p**c,0, 1]) ).change_ring(ZZ) )
            tt=verbose(" Computing matrix of WR on mod symbs", level=VERB)
            msw = symbs._matrix_of_operator_on_modular_symbols(symbs, [x.list() for x in L])
            verbose(" Done", tt, level=VERB)
            symb_A = msw.restrict(cond_parts[c].free_module())
        else:
            tt=verbose(" Computing matrix of Up^{-1} on mod symbs", level=VERB)
            symb_A = p**( (S.weight() - 1) *c) * ~cond_parts[c].hecke_matrix(p**c)
            verbose(" Done", tt, level=VERB)
            
        tt=verbose(" Calling convert_matrix_from_modsyms", level=VERB)
        WR = _convert_matrix_from_modsyms(cond_parts[c], symb_A)[0]
        verbose(" done", tt, level=VERB)
        verbose("done computing WR", t, level=VERB)

        # Now compute the matrix of ~R, which is defined over QQ(zeta_Q) and has eigenvalues of norm p^{-c/2}.
        t = verbose("computing ~R",level=VERB)
        if c > 0:
            d = cond_parts[c].dimension()
            zeta = CyclotomicField(p**c).gen()
            
            tt=verbose(" computing q-expansion module", level=VERB)
            e = d+1
            while 1:
                M = matrix([f.padded_list(e) for f in cond_parts[c].q_expansion_basis(prec=e)])
                for i in range(M.ncols()//p): 
                    M.set_col_to_multiple_of_col(p*i, p*i, 0)
                if M.rank() == d: 
                    e = M.pivots()[-1] + 1
                    M = M[:,:e]
                    break
                e += 5
                if e > p*S.sturm_bound(): raise ArithmeticError("Can't get here -- multiplicity 1 violated")

            Maug = M.augment(identity_matrix(d))
            ech = Maug.echelon_form()
            pivs = ech.pivots()
            T = ech[:,e:]

            tt = verbose(" done; need precision %s" % e, tt, level=VERB)

            tt=verbose(" computing diamond matrices on S_c", level=VERB)
            # We precompute these matrices using the fact that Zmod(p**c)^* is (almost) cyclic.
            dmats = {}
            if p > 2:
                a = Zmod(p**c).unit_gens()[0]
                aa = crt(a.lift(), 1, p**c, N//p**r)
                da = _convert_matrix_from_modsyms(cond_parts[c], cond_parts[c].diamond_bracket_matrix(aa))[0]
                for i in range(p**(c-1) * (p-1)):
                    dmats[ (a**i).lift() ] = da**i
            else:
                # c > 0, hence c >= 2!
                if c == 2: d5 = identity_matrix(cond_parts[c].rank())
                else: d5 = _convert_matrix_from_modsyms(cond_parts[c], cond_parts[c].diamond_bracket_matrix(crt(5, 1, p**c, N//p**r)))[0]
                dm1 = _convert_matrix_from_modsyms(cond_parts[c], cond_parts[c].diamond_bracket_matrix(crt(-1, 1, p**c, N//p**r)))[0]
                for i in range(p**(c-2)):
                    dmats[ (5**i) % (p**c) ] = d5**i
                    dmats[ (- (5**i) ) % p**c ] = d5**i * dm1
            verbose(" done", tt, level=VERB)

            qexp_mat = matrix([f.padded_list(e) for f in cond_parts[c].q_expansion_basis(prec=e)])
            L = []
            for i in range(d):
                t=verbose("  computing twist of %s-th gen" % i, level=VERB)
                v = [zeta.parent()(0)]*d
                for a in xsrange(p**c):
                    if not a%p: continue
                    fa = qexp_mat.linear_combination_of_rows(dmats[a].row(i))
                    for n in xsrange(d):
                        v[n] += fa[pivs[n]]*zeta**(a*pivs[n])

                # why is this step so damn slow?
                L.append(sum(v[i] * T.row(i) for i in xsrange(d)))
                verbose("  done", t, level=VERB)
            invR = matrix(L) / p**c
        else:
            invR = identity_matrix(cond_parts[c].rank())
        
        verbose("done computing ~R; fcp of W: %s" % (WR*invR).fcp(), t, level=VERB)
        # Now we're sorted.
        basis += [S(f).element() for f in cond_parts[c].q_expansion_basis(prec=S.sturm_bound())]
        matrix_rows += [ [0]*lpad + list(v) + [0]*rpad for v in (WR*invR).rows() ]
    
    t = verbose("assembling conductor parts into echelon form", level=VERB)
    basmat = matrix(basis)
    wmat = matrix(matrix_rows)
    to_ech = basmat.solve_left(S.new_submodule(p).basis_matrix())
    x = to_ech*wmat*~to_ech
    verbose("done", t, level=VERB)
    verbose("done computing W_%s on newspace at level %s" % (Q, N), t0, level=VERB)
    return x
            
def decompose_by_conductor(M, p):
    r"""
    Given a Hecke module M, and a prime p, decompose M into a direct sum of
    submodules on which the diamond operators act via characters of a given
    conductor at p.
    """
    N = M.level()
    r = N.valuation(p)
    Q = p**r
    if r == 0:
        return {0: M.free_module() }

    cond_parts = {}
    if p > 2:
        d = M.ambient().diamond_bracket_matrix(crt(1+p, 1, Q, N//Q))
        e = M.ambient().diamond_bracket_matrix(crt(Zmod(Q).unit_gens()[0].lift()**p^(r-1), 1, Q, N//Q))
        for c in [2..r]:
            # cond p^c part is where <1+p> acts as a primive
            # root of unity of order p^(c-1)
            poly = cyclotomic_polynomial( p^(c-1) )
            cond_parts[c] = poly(d).kernel_on(M.free_module())
        C = (d-1).kernel_on(M.free_module())
        x = polygen(QQ)
        poly = (x**(p-1) - 1) // (x-1)
        cond_parts[1] = C.intersection(poly(e).kernel_on(M.free_module()))
        cond_parts[0] = C.intersection((e-1).kernel_on(M.free_module()))
    else: 
        # p = 2 case
        fi = M.ambient().diamond_bracket_matrix(crt(5, 1, Q, N//Q))
        m1 = M.ambient().diamond_bracket_matrix(crt(-1, 1, Q, N//Q))
        for c in [3..r]:
            # cond p^c part is where <5> acts as primitive root of unity
            # of order 2^(c-2)
            cond_parts[c] = cyclotomic_polynomial(2**(c-2))(fi).kernel_on(M.free_module())
        C = (fi - 1).kernel_on(M.free_module())
        cond_parts[2] = C.intersection( (m1+1).kernel_on(M.free_module()) )
        cond_parts[0] = C.intersection( (m1-1).kernel_on(M.free_module()) )
        # no character has conductor 2^1!

    for s in list(cond_parts.keys()):
        if cond_parts[s].rank() == 0:
            cond_parts.pop(s)
        else:
            cond_parts[s] = M.submodule(cond_parts[s], check=False)
    assert sum(cond_parts[k].rank() for k in cond_parts.keys()) == M.rank()    
    return cond_parts

def p_depletion_map(S, p, base_ring=None):
    # S is a space of modular forms, new at p
    if base_ring is None: base_ring=S.base_ring()
    d = S.dimension()
    e = d
    L = S.base_ring()
    while e < 200:
        V = base_ring**e
        vecs = []
        for f in S.q_expansion_basis(prec=e):
            v = [ (n%p and f[n] or 0) for n in range(e) ]
            vecs.append(v)
        D = S.free_module().base_extend(base_ring).Hom(V)(vecs)
        if D.rank() == d: return D
        else: e += 1
    # can't get here!
    raise ValueError("multiplicity one violated!")

@cached_function
def STmatrices(N, k):
    # return matrices of [0,-1,1,0] and [1,1,0,1] on Sk(Gamma(N))
    G = GammaH(N^2, [1 + N*x for x in range(N)])
    C = CuspForms(G, k)
    W = Wmatrix(C, N^2)
    L,z = CyclotomicField(N).objgen()
    verbose("Computing T", level=VERB)
    pivs = C._q_expansion_module().basis_matrix().pivots()
    T = diagonal_matrix([ z**pivs[i] for i in range(C.rank()) ])
    verbose("...Done", level=VERB)
    assert (W*T)^3 == 1
    return (W, T)

@cached_function
def LRmatrices(N, k):
    S, T = STmatrices(N, k)
    return T, S*~T*~S

@cached_function
def action(N, k, g):
    gens = LRmatrices(N, k)
    g = MatrixSpace(Zmod(N), 2)(g)
    assert g.det() == 1
    g = sage.modular.local_comp.liftings.lift_for_SL(g)
    if g == 1: return gens[0]**0
    return prod([gens[i] ** j for (i,j) in sage.modular.arithgroup.arithgroup_perm.sl2z_word_problem(g) ])

def forms_of_level(G):
    # G should be a subgroup of SL(2, Z/N)
    N = G.base_ring().characteristic()
    gens = [x.matrix() for x in G._gap_().MinimalGeneratingSet()]
    assert gens != []
    V = CyclotomicField(N)**Gamma(N).dimension_cusp_forms(2)
    for h in gens:
        verbose("Finding invariants for %s" % (str(h).replace("]\n[","; ")) , level=VERB)
        a = action(N, 2, h)
        V = V.intersection( (a-1).kernel() )
    C = CuspForms(GammaH(N^2, [1 + N*x for x in range(N)]),2, base_ring=V.base_ring())
    return [C.linear_combination_of_basis(x) for x in V.gens()]

def cartannormalizer(p, u=None):
    x = polygen(GF(p))
    if not u: u = GF(p).quadratic_nonresidue()
    k.<s> = GF(p).extension(x^2 - u)
    z = k.multiplicative_generator()
    M = MatrixGroup([ z.matrix(), matrix(GF(p), 2,2,[1,0,0,-1]) ])
    N = MatrixGroup([x for x in M if x.matrix().det() == 1])
    assert N.order() == 2*(p+1)
    return N