p = 5
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^3)
f2 = superelliptic_function(C_super, x^11)
AS = as_cover(C_super, [f1, f2], prec=500)
zmag = AS.magical_element(threshold = 20)[0]
zvee = dual_elt(AS, zmag)

### DEFINE THE POLYNOMIALS
n = 2
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
#############
print(zvee.expansion_at_infty().valuation())
g = AS.genus()
print(AS.exponent_of_different_prim())
#for i in range(g):
#    om = AS.holomorphic_differentials_basis(threshold = 30)[i]
#    print(AS.exponent_of_different_prim(), val_of_components(om, zvee))


v_x = as_function(AS, x).expansion_at_infty().valuation()
v_z0 = as_function(AS, z[0]).expansion_at_infty().valuation()
v_z1 = as_function(AS, z[1]).expansion_at_infty().valuation()
n = 2
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+k[0]*v_z0+k[1]*v_z1
        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)))
            tr_wz = (zvee*w).trace()
            val = tr_wz.expansion_at_infty().valuation()
            #val *= p^2
            #print(val)