def mumford_gp(p):
    '''We want m | p-1 and b to be of order m in F_p.'''
    name = "Mumford group (Z/"+str(p)+")^2⋊ D"+str(p-1)
    short_name = "(Z/"+str(p)+")^2 ⋊ D"+str(p-1)
    elts = [(i, j, s, t) for i in range(p) for j in range(p) for s in range(1) for t in range(p-1)]
    mult = lambda elt1, elt2: (0, 0, 0, 0)
    inv = lambda elt1 : (0, 0, 0, 0)
    gens = [(1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0,0,0,1)]
    one = (0, 0, 0, 0)
    gp = group(name, short_name, elts, one, mult, inv, gens)
    return gp

def mumford_template(p):
    group = mumford_gp(p)
    field = GF(p)
    n = 2
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i) + ','
    variable_names += 'x, y'        
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    height = n
    fcts = [x + y, x - y]
    gp_action = [[z[j] + (i == j) for j in range(n)]+[x, y] for i in range(n)]
    gp_action += [[z[1], z[0], -x, y]]
    d = field(primitive_root(p))
    gp_action += [[d*z[0], d^(-1)*z[1], 1/d*(y+x) - d*(y-x), 1/d*(y+x) + d*(y-x)]]
    return template(height, field, group, fcts, gp_action)
    
def mumford_cover(p, prec=10):
    field = GF(p)
    R.<x> = PolynomialRing(field)
    C = superelliptic(x^2 + 1, 2)
    return as_cover(C, mumford_template(p), [C.x, C.x], branch_points = [], prec = prec)    
    
ASM = mumford_cover(3, prec = 400)


###############
###############
def metacyclic_gp(p, c):
    name = "metacyclic group Z/"+str(p)+"⋊ Z/"+str(c)
    short_name = "Z/"+str(p)+" ⋊ Z/"+str(c)
    elts = [(i, j) for i in range(p) for j in range(c)]
    mult = lambda elt1, elt2: (0, 0)
    inv = lambda elt1 : (0, 0)
    gens = [(1, 0), (0, 1)]
    one = (0, 0)
    gp = group(name, short_name, elts, one, mult, inv, gens)
    return gp

def metacyclic_template(p, m, nn):
    group = metacyclic_gp(p, m*(p-1))
    N = Integers(m*(p-1))(p).multiplicative_order()
    field.<a> = GF(p^N)
    zeta = a^((p^N - 1)/((p-1)*m))
    n = 1
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i) + ','
    variable_names += 'x, y'        
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    height = n
    fcts = [x]
    gp_action = [[z[j] + (i == j) for j in range(n)]+[x, y] for i in range(n)]
    gp_action += [[zeta^m * z[0], zeta^m * x, zeta * y]]
    return template(height, field, group, fcts, gp_action)
    
def metacyclic_cover(p, m, nn, prec=10):
    N = Integers(m*(p-1))(p).multiplicative_order()
    field.<a> = GF(p^N)
    R.<x> = PolynomialRing(field)
    C = superelliptic(sum(x^(p^i) for i in range(0, nn)), m)
    return as_cover(C, metacyclic_template(p, m, nn), [C.x], branch_points = [], prec = prec)

p = 3
m = 2
nn = 2
N = Integers(m*(p-1))(p).multiplicative_order()
field.<a> = GF(p^N)
zeta = a^((p^N - 1)/((p-1)*m))
AS = metacyclic_cover(p, m, nn, prec=100)
Bh = AS.holomorphic_differentials_basis(threshold = 30)
Bs = AS.cohomology_of_structure_sheaf_basis(threshold = 30)
BdR = AS.de_rham_basis(threshold = 30)
print(Bh[2].group_action((1, 0)).coordinates(basis = Bh))


#H = CyclicPermutationGroup(4)

#alpha = PermutationGroupMorphism(A,A,[A.gens()[0]^3*A.gens()[1],A.gens()[1]])

#phi = [[(1,2)],[alpha]]

p = 3
m = 2
nn = 2
N = Integers(m*(p-1))(p).multiplicative_order()
field.<a> = GF(p^N)
zeta = a^((p^N - 1)/((p-1)*m))

def gp_action(fct, elt):
    if elt == (1, 0):
        if isinstance(fct, superelliptic_function):
            f = fct.function
            C = fct.curve
            Fxy, Rxy, x, y = C.fct_field
            f = f.subs({x: x+1, y: y})
            return superelliptic_function(C, f)

        if isinstance(fct, superelliptic_form):
            f = fct.form
            C = fct.curve
            Fxy, Rxy, x, y = C.fct_field
            f = f.subs({x: x+1, y: y})
            return superelliptic_form(C, f)

        if isinstance(fct, superelliptic_cech):
            omega = fct.omega0
            ff = fct.f
            C = fct.curve
            return superelliptic_cech(C, gp_action(omega, elt), gp_action(ff, elt))
    if elt == (0, 1):
        if isinstance(fct, superelliptic_function):
            f = fct.function
            C = fct.curve
            Fxy, Rxy, x, y = C.fct_field
            f = f.subs({x : zeta^m*x, y:zeta*y})
            return superelliptic_function(C, f)

        if isinstance(fct, superelliptic_form):
            f = fct.form
            C = fct.curve
            Fxy, Rxy, x, y = C.fct_field
            f = f.subs({x : zeta^m*x, y:zeta*y})
            return superelliptic_form(C, f*zeta^m)

        if isinstance(fct, superelliptic_cech):
            omega = fct.omega0
            ff = fct.f
            C = fct.curve
            return superelliptic_cech(C, gp_action(omega, elt), gp_action(ff, elt))
        
        
superelliptic_function.group_action = gp_action
superelliptic_form.group_action = gp_action
superelliptic_cech.group_action = gp_action

R.<x> = PolynomialRing(field)
C = superelliptic(x^(p^nn) - x, m)
Bh = C.holomorphic_differentials_basis()
Bs = C.cohomology_of_structure_sheaf_basis()
BdR = C.de_rham_basis()
mat = as_group_action_matrices(field, BdR, [(1, 0), (0, 1)], BdR)