class ModTor(object):
    def msg(self, *args):
        if self.verbose:
            for x in args:
                print x,
            print
    def __init__(self, group, assume_integral=False, verbose=True):
        self.group = group
        self.verbose = verbose
        t = cputime(); self.M = ModularSymbols(group); self.msg("got M", cputime(t))
        t = cputime(); self.S = self.M.cuspidal_subspace(); self.msg("got S", cputime(t))
        if assume_integral:
            self.is_integral = True
        else:
            t = cputime(); self.SZ = self.S.integral_structure(); self.msg("got S_Z", cputime(t))
            self.is_integral = (self.S.free_module().basis_matrix() == self.SZ.basis_matrix())

    def __repr__(self):
        return "Torsion on the modular Jacobian associated to %s"%self.group
            
    def change_basis(self, A):
        if self.is_integral:
            return A
        else:
            return A.restrict(self.SZ)
        
    def dbd(self, d):
        return self.change_basis(self.S.diamond_bracket_operator(d).matrix())       
        
    def eta(self, ell):
        return self.change_basis(self.S.hecke_matrix(ell)) - (1 + self.dbd(ell)*ell)

    def star(self):
        return self.change_basis(self.S.star_involution().matrix())

    def C_lattice(self):
        n = self.S.dimension()
        V = self.M.free_module()
        m = V.dimension() - n
        V0 = V.span_of_basis(self.S.free_module().basis() + self.M.eisenstein_submodule().free_module().basis())
        if self.is_integral:
            pi = lambda x: V0.coordinates(x.element())[:n]
        else:
            pi = lambda x: self.SZ.coordinates(V0.coordinates(x.element())[:n] + [0]*m)
        C = (QQ^n).span([pi(x) for x in self.M.gens()], ZZ)
        return C 
        
    def E_lattice(self, primes, plus=True):
        W = None
        for ell in primes:
            assert 2*self.group.level()%ell != 0
            V = (self.eta(ell)**(-1)).row_module(ZZ)
            if W is None:
                W = V
            else:
                W = W.intersection(V)

        if plus:
            d = W.dimension()
            M = W / (ZZ^d)
            S = self.star() - 1
            d = S.nrows()
            T = M.hom([M(a.lift()*S) for a in M.gens()]).kernel()
            W = T.V()
        return W

@parallel
def J0_torsion_prob_cuspidal(N, B=20):
    if not os.path.exists('data'):
        os.makedirs('data')
    file = 'data/%s-%s.sobj'%(N,B)
    if os.path.exists(file):
        return load(file)
    t = cputime()
    M = ModTor(Gamma0(N), verbose=False)
    C = M.C_lattice()
    E = M.E_lattice([p for p in primes(B) if (2*M.group.level())%p])
    ans = E.is_submodule(C)
    data = {'ans':ans, 'B':B, 'cputime':cputime(t), 'C':C, 'E':E, 'C_order':C.basis_matrix().det()**(-1),
            'E_order':E.basis_matrix().det()**(-1)}
    save(data, file)
    return data

def J0_torsion_prob_cuspidal_go(N1,N2, B=20):
    for X in J0_torsion_prob_cuspidal([(n,B) for n in [N1..N2]]):
        print X


########################################
#

        
def J1_torsion(p, B=13, assume_integral=False):
    if not os.path.exists('data'):
        os.makedirs('data')
    file = 'data/%s-%s-%s.sobj'%(p,B,assume_integral)
    if os.path.exists(file):
        return load(file)
    
    t0=cputime(); T = ModTor(Gamma1(p), assume_integral=assume_integral); t0=cputime(t0)
    tC=cputime(); C = T.C_lattice(); tC=cputime(tC)
    tE=cputime(); E = T.E_lattice([p for p in primes(B+1) if p%2]); tE=cputime(tE)
    
    ans = E.is_submodule(C)
    data = {'ans':ans, 'B':B, 'p':p, 't0':t0, 'tC':tC, 'tE':tE, 'C':C, 'E':E, 'C_order':C.basis_matrix().det()**(-1), 'E_order':E.basis_matrix().det()**(-1)}

    save(data, file)
    return data