"""
This script makes the list of number fields K such that Gal(K) = (Z/ell)^r where ell is a prime and r an integer, with conductor comprised in an interval [E1,E2].
"""

# DATA ################################################################

E1=1
E2=10000
ell=3
r=1
gd=open("cubic-"+str(E1)+"-"+str(E2)+"-"+str(ell)+"-"+str(r)+".txt","w")

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


"""
Given an integer k, this function makes the list of partitions of r terms.
"""
def integer_partition(k,r):
    total=[]
    # aux is the auxiliary function which computes the particiants
    # of k using r terms, out of which i-1 are fixed, bs, of sum s
    def aux(k,r,i,s,bs):
        for bi in range(k-s+1):
            if i < r-1:
                aux(k,r,i+1,s+bi,[e for e in bs]+[bi])
            else:
                 total.append([e for e in bs]+[bi])
    if r == 0:
        return [[]]
    else:
        aux(k,r,0,0,[])
    return total


"""
Compute all GF(ell)-subspaces of GF(ell)^r of dimension k.
"""
def subgroupsZ(ell,k,r):
    from itertools import product
    if k < r:
        return [],[]
    ms=[]
    co_ms=[]
    for ks in integer_partition(k-r,r):
        for m in product(range(ell),repeat=k*r-r^2-sum([(r-j)*ks[j] for j in range(r)])):
            rows_m=[]
            c=0
            for i in range(r):
                mi=[0]*(i+sum(ks[:i+1]))+[1]
                for j in range(i+1,r):
                    for _ in range(ks[j]):
                        mi=mi+[m[c]]
                        c+=1
                    mi+=[0]
                for _ in range(k-r-sum(ks)):
                    mi=mi+[m[c]]
                    c+=1
                rows_m.append(mi) 
                ms.append(rows_m) 
            if rows_m != []:
                Ms=Matrix(rows_m)
                com=Ms.right_kernel().basis()
                co_ms.append(com)
            else:
                co_ms.append(identity_matrix(ZZ,k).rows())
    return ms,co_ms

"""
Given a prime q, compute a multiplicative generator, in a deterministic manner.
"""
def multiplicative_generator(q):
    Fq=GF(q)
    a=2 
    while gcd(a,q) != 1 or Fq(a).multiplicative_order() != q-1:
        a=next_prime(a)
    return a    

"""
Given an integer n and its factorization, compute gs multiplicative generators of (Z/p^val_p(n))* and ell-th roots.
"""
def find_structure(nfact,ell):
    k=len(nfact)
    mods=[]
    ts=[]
    tells=[]
    for q in nfact:
        if q != ell:
            t=multiplicative_generator(q)
            tell=t^ell
            ts.append(t)
            tells.append(tell)
            mods.append(q)
        else:
            t=(ell+multiplicative_generator(ell)) 
            tell=t^ell
            ts.append(t)
            tells.append(tell)
            mods.append(ell^2)
    bs=[]
    gs=[]
    for i in range(k):
        vi=[1]*i+[tells[i]]+[1]*(k-i-1)
        bi=crt(vi,mods)
        bs.append(bi)    
        wi=[1]*i+[ts[i]]+[1]*(k-i-1)
        gi=crt(wi,mods)
        gs.append(gi)
    return bs,gs

"""
Write a list of integers bs into the pari/Gp format.
"""
def pretty_mod(bs,n):
    str_bs=""
    for b in bs:
        if str_bs != "":
            str_bs=str_bs+","
        str_bs=str_bs+"Mod("+str(b)+","+str(n)+")"
    return str_bs

"""
Given an integer n, list all subfields of Q(zeta_n), where zeta_n is an n-th cyclotomic field, such that Gal(K) = (Z/ell)^r. The fields are listed by a defining polynomial, in file gd.
"""
def list_subcyclo(n,ell,r,gd=sys.stdout):
    nfact=n.prime_factors() 
    k=len(nfact)
    if k < r:
        return None
    gd.write(str(n)+"\n")
    sgps,cobs=subgroupsZ(ell,k,r)
    bs,gs=find_structure(nfact,ell)
    for sgp in cobs:
        bsgp=[e for e in bs]
        for m in sgp:
            bm=prod([gs[i]^m[i] for i in range(k)]) % n
            bsgp.append(bm)    
        str_bsgp=pretty_mod(bsgp,n)        
        s=pari("Z=znstar("+str(n)+");galoissubcyclo(Z,["+str_bsgp+"])")
        gd.write(str(s)+"\n")
        gd.flush()
    return None

"""
This function outputs integers in the interval [E1,E2] which are products of primes =1%ell and or ell^2 (multiplicities not accepted) 
"""
def sieve(E1,E2,ell,r,gd):
    logs=[log(a,2).n() for a in range(E1,E2)]
    p=ell^2
    a0=p*((E1 // p)+1)
    lgp=log(p,2).n()
    for a in range(a0,E2+1,p):
        logs[a-E1]-=lgp
    for p in primes(E2):
        if p % ell == 1:
            a0=p*((E1 // p)+1)
            lgp=log(p,2).n()
            for a in range(a0,E2+1,p):
                logs[a-E1]-=lgp
    results=[]
    for n in range(E1,E2):
        if logs[n-E1] < 0.5:
            list_subcyclo(ZZ(n),ell,r,gd)
    return results

# Main computation.
sieve(E1,E2,ell,r,gd)