class superelliptic_regular_form:
    def __init__(self, A, B):
        self.dx = A
        self.dy = B
        self.curve = A.curve
        
    def __repr__(self):
        if self.dx.function == 0:
            return "(" + str(self.dy) + ") dy"
        if self.dy.function == 0:
            return "("+str(self.dx) + ") dx"
        return "("+str(self.dx) + ") dx + (" + str(self.dy) + ") dy"
    
    def form(self):
        C = self.curve
        return self.dx*C.dx + self.dy*C.y.diffn()
    
    def integral(self):
        '''Regular integral. Works only for hyperelliptics.'''
        C = self.curve
        f = C.polynomial
        if C.exponent != 2:
            raise ValueError("Works only for hyperelliptics.")
        F = C.base_ring
        Rx.<x> = PolynomialRing(F)
        Fx = FractionField(Rx)
        Fxy, Rxy, x, y = C.fct_field
        if self.dx == 0*C.x and self.dy == 0*C.x:
            return 0*C.x
        #which = random.choice([0, 1])
        RxRy.<y> = PolynomialRing(Rx)
        P = RxRy(self.dx.function)
        Q = RxRy(self.dy.function)
        Py, Px = P.quo_rem(y) #P = y*Py + Px
        Qy, Qx = Q.quo_rem(y)
        Py, Px, Qy, Qx, f = Rx(Py), Rx(Px), Rx(Qy), Rx(Qx), Rx(f)
        result = superelliptic_function(C, Rx(Px + F(1/2)*Qy*f.derivative()).integral())
        numerator = Rx(2*f*Py + f.derivative()*Qx)
        # Now numerator = 2W' f + W f'. We are looking for W. Then result is W*y.
        W = Rx(0)
        while(numerator != 0):
            d = numerator.degree()
            r = f.degree()
            n_lead = numerator.leading_coefficient()
            f_lead = Rx(f).leading_coefficient()
            a = d - (r-1)
            if a >= 0:
                W_coeff = F(n_lead/f_lead)*F((2*a + r)^(-1))
                W += W_coeff*Rx(x^a)
                numerator -= 2*f*(W_coeff*Rx(x^a)).derivative() + (W_coeff*Rx(x^a))*f.derivative()
                numerator = Rx(numerator)
            if a < 0:
                W += Rx(numerator/f.derivative())
                numerator = Rx(0)
        result = result + superelliptic_function(C, y*W)
        return result

class superelliptic_regular_drw_form:
    def __init__(self, A, B, omega, h2):
        self.dx = A
        self.dy = B
        self.omega = omega
        self.h2 = h2
        self.curve = A.curve
        
    def form(self):
        C = self.curve
        A = self.dx
        B = self.dy
        h2 = self.h2
        omega = self.omega
        form1 = superelliptic_drw_form(A, omega.form(), h2)
        form2 = B.teichmuller()*C.y.teichmuller().diffn()
        
    def __repr__(self):
        return "[" + str(self.dx) + "] d[x] + [" + str(self.dy) + "] d[y] + V(" + str(self.omega) + ") + dV(" + str(self.h2) + ")"
        
def regular_drw_form(omega):
    C = omega.curve
    p = C.characteristic
    omega_aux = omega.r()
    omega_aux = omega_aux.regular_form()
    aux = omega - omega_aux.dx.teichmuller()*C.x.teichmuller().diffn() - omega_aux.dy.teichmuller()*C.y.teichmuller().diffn()
    aux1 = aux.omega
    aux.omega, fct = decomposition_omega0_hpdh(aux.omega)
    aux.h2 += fct^p
    aux.h2, A = decomposition_g0_pth_power(aux.h2)
    aux.omega += (A.diffn()).inv_cartier()
    result = superelliptic_regular_drw_form(omega_aux.dx, omega_aux.dy, aux.omega.regular_form(), aux.h2)
    return result

superelliptic_drw_form.regular_form = regular_drw_form

def regular_drw_cech(cocycle):
    return("( " + str(cocycle.omega0.regular_form()) + ", " + str(cocycle.f) + " )")

superelliptic_drw_cech.regular_form = regular_drw_cech
    
def regular_form(omega):
    '''Given a form omega regular on U0, present it as P(x, y) dx + Q(x, y) dy for some polynomial P, Q.
    The output is A(x)*y, B(x), where omega = A(x) y dx + B(x) dy'''
    if not omega.is_regular_on_U0():
        raise ValueError("The form " + str(omega) + " is not regular on U0.")
    C = omega.curve
    f = C.polynomial
    Fxy, Rxy, x, y = C.fct_field
    F = C.base_ring
    Rx.<x> = PolynomialRing(F)
    fct = omega.form
    if fct.denominator() == y:
        fct = fct.numerator()
        integral_part, fct = fct.quo_rem(y)
        d, A, B = xgcd(f, f.derivative())
        return superelliptic_regular_form(superelliptic_function(C, integral_part + A*fct*y), superelliptic_function(C,2*B*fct))
    if fct.denominator() == 1:
        return superelliptic_regular_form(superelliptic_function(C, fct), 0*C.x)
        
superelliptic_form.regular_form = regular_form