p = 3
m = 1
F = GF(p)
Rx.<x> = PolynomialRing(F)
f = x
C_super = superelliptic(f, m)

Rxy.<x, y> = PolynomialRing(F, 2)
f1 = superelliptic_function(C_super, x^2)
f2 = superelliptic_function(C_super, x^4)
f3 = superelliptic_function(C_super, x^5)
AS = as_cover(C_super, [f1, f2, f3], prec=1000)
print(AS)
zmag = AS.magical_element(threshold = 20)[0]
zvee = dual_elt(AS, zmag)

### DEFINE THE POLYNOMIALS
n = AS.height
variable_names = 'x, y'
for i in range(n):
    variable_names += ', z' + str(i)
Rxyz = PolynomialRing(F, n+2, variable_names)
x, y = Rxyz.gens()[:2]
z = Rxyz.gens()[2:]

###############
def val_of_components(omega, zvee):
    result = []
    AS = omega.curve
    for i in range(p^2):
        omega_i = ith_magical_component(omega, zvee, i)
        val = omega_i.expansion_at_infty().valuation()
        val = val*p^2 + AS.exponent_of_different()
        result += [val]
    return result
#############
G = AS.group

g = AS.genus()
print(AS.exponent_of_different_prim())
v_x = as_function(AS, x).expansion_at_infty().valuation()
n = AS.height
v_z = [as_function(AS, z[i]).expansion_at_infty().valuation() for i in range(n)]
omega = AS.holomorphic_differentials_basis()[-16]
from itertools import product
pr = [list(range(p)) for _ in range(n)]
for i in range(0, 30):
    for k in product(*pr):
        v_w = i*v_x+ sum(k[i]*v_z[i] for i in range(n))
        if (v_w < - AS.exponent_of_different_prim() + 10 and v_w > - AS.exponent_of_different_prim()):
            w = as_function(AS, x^i * prod(z[i1]^(k[i1]) for i1 in range(n)))
            result = as_form(AS, 0)
            for g in G:
                component = ith_magical_component(omega, zvee, g).form
                result += as_form(AS, w.group_action(g).function*component)
            print(result.expansion_at_infty().valuation(), w.expansion_at_infty().valuation() - zmag.expansion_at_infty().valuation())