class superelliptic_witt:
    def __init__(self, C, t, f0, f1):
        ''' Define Witt function on C of the form [t] + f0([x], [y]) + V(f1). '''
        self.curve = C
        p = C.characteristic
        self.t = t #superelliptic_function
        self.f0 = reduce_rational_fct(f0, p) #polynomial/rational function over Z/p^2
        self.f1 = f1 #superelliptic_function
    
    def __repr__(self):
        f0 = self.f0
        f1 = self.f1
        t = self.t
        return "[" + str(t) + "] + " + str(f0).replace("x", "[x]").replace("y", "[y]") + " + V(" + str(f1) + ")"
    
    def __add__(self, other):
        C = self.curve
        return superelliptic_witt(C, self.t + other.t, self.f0 + other.f0, self.f1 + other.f1)

    def teichmuller_representation(self):
        '''Represents as [f] + V(g), i.e. f0 = 0.'''
        C = self.curve
        Fxy, Rxy, x, y = self.fct_field
        F = C.base_ring
        function = Rxy(self.f0)
        if self.f0 == 0:
            return self
        M = Rxy(function.monomials()[0])
        a = F(function.monomial_coefficient(M))
        fct1 = fct - superelliptic_function(C, a*M)
        function1 = fct1.function
        return teichmuller(fct1) + superelliptic_witt(C, (a.lift())^p*M.change_ring(QQ), superelliptic_function(C, function1^2*a*M + function1*(a*M)^2))
    
    def antiteichmuller_representation(self):
        '''Represents as f([x], [y]) + V(g), i.e. teichmuller part is zero.'''
        return 0
    
def reduce_rational_fct(fct, p):
    Rxy.<x, y> = PolynomialRing(QQ)
    Fxy = FractionField(Rxy)
    fct = Fxy(fct)
    num = Rxy(fct.numerator())
    den = Rxy(fct.denominator())
    num1 = Rxy(0)
    for m in num.monomials():
        a = num.monomial_coefficient(m)
        num1 += (a%p^2)*m
    den1 = Rxy(0)
    for m in den.monomials():
        a = den.monomial_coefficient(m)
        den1 += (a%p^2)*m
    return num1/den1

def teichmuller(fct):
    C = fct.curve
    Fxy, Rxy, x, y = C.fct_field
    F = C.base_ring
    function = Rxy(fct.function)
    if function == 0:
        return superelliptic_witt(C, 0, 0*C.x)
    M = Rxy(function.monomials()[0])
    a = F(function.monomial_coefficient(M))
    fct1 = fct - superelliptic_function(C, a*M)
    function1 = fct1.function
    return teichmuller(fct1) + superelliptic_witt(C, (a.lift())^p*M.change_ring(QQ), superelliptic_function(C, function1^2*a*M + function1*(a*M)^2))

class superelliptic_drw_form:
    def __init__(self, C, omega_x, omega_y, omega, h):
        '''Form [omega_x] d[x] + [omega_y] d[y] + V(omega) + dV([h])'''
        self.curve = C
        self.omega_x = omega_x
        self.omega_y = omega_y
        self.omega = omega
        self.h = h
        
    def __eq__(self, other):
        eq1 = (self.omega1 == self.omega1)
        try:
            H = (self.h - other.h).pthroot()
        except:
            return False
        eq2 = (self.omega2 - other.omega2).cartier() - H.diffn()
        if eq1 and eq2:
            return True
        return False
    
    def __repr__(self):
        C = self.curve
        omega_x = self.omega_x
        omega_y = self.omega_y
        h = self.h
        return  str(omega_x) + "] dx + [" + str(omega_x.form) + "] d[x] " + "+ V(" + str(omega2) + ") + dV([" + str(h) +"])"

def mult_by_p(omega):
    C = omega.curve
    fct = omega.form
    Fxy, Rxy, x, y = C.fct_field
    omega2 = superelliptic_form(C, fct^p * x^(p-1))
    result = superelliptic_drw_form(C, 0*C.dx, omega2, 0*C.x)
    return result
    
def basis_W2Omega(C):
    basis = C.holomorphic_differentials_basis()
    result = []
    for omega in basis:
        result += [mult_by_p(omega)]
        
    image_of_cartier = []
    
    return result