class superelliptic_witt:
    def __init__(self, t, f):
        ''' Define Witt function on C of the form [t] + V(f). '''
        self.curve = t.curve
        C = t.curve
        p = C.characteristic
        self.t = t #superelliptic_function
        self.f = f #superelliptic_function
    
    def __repr__(self):
        f = self.f
        t = self.t
        if f.function == 0:
            return "[" + str(t) + "]"
        if t.function == 0:
            return "V(" + str(f) + ")"
        return "[" + str(t) + "] + V(" + str(f) + ")"
    
    def __neg__(self):
        f = self.f
        t = self.t
        return superelliptic_witt(-t, -f)
    
    def __add__(self, other):
        C = self.curve
        second_coor = 0*C.x
        X = self.t
        Y = other.t
        for i in range(1, p):
            second_coor -= binomial_prim(p, i)*X^i*Y^(p-i)
        return superelliptic_witt(self.t + other.t, self.f + other.f + second_coor)
    
    def __sub__(self, other):
        return self + (-other)
    
    def __rmul__(self, other):
        p = self.curve.characteristic
        if other in ZZ:
            if other == 0:
                return superelliptic_witt(0*C.x, 0*C.x)
            if other > 0:
                return self + (other-1)*self
            if other < 0:
                return (-other)*(-self)
        if other in QQ:
            other_integer = Integers(p^2)(other)^(-1)
            return other_integer*self
    
    def __mul__(self, other):
        if isinstance(other, superelliptic_witt):
            t1 = self.t
            f1 = self.f
            p = self.curve.characteristic
            t2 = other.t
            f2 = other.f
            return superelliptic_witt(t1*t2, t1^p*f2 + t2^p*f1)
        if isinstance(other, superelliptic_drw_form):
            h1 = other.h1
            h2 = other.h2
            omega = other.omega
            p = other.curve.characteristic
            t = self.t
            f = self.f
            aux_form = t^p*omega - h2*t^(p-1)*t.diffn() + f*h1^p*(C.x)^(p-1)*C.dx
            return superelliptic_drw_form(t*h1, aux_form, t^p*h2)
    
    def __eq__(self, other):
        return self.t == other.t and self.f == other.f
    
    def diffn(self):
        C = self.curve
        t = self.t
        f = self.f
        fC = C.polynomial
        F = C.base_ring
        Rxy.<x, y> = PolynomialRing(F, 2)
        if t.function == 0:
            return superelliptic_drw_form(0*C.x, 0*C.dx, f)
        t_polynomial = t.function
        num = t_polynomial.numerator()
        den = t_polynomial.denominator()
        num_t_fct = superelliptic_function(C, num)
        den_t_fct = superelliptic_function(C, den)
        inv_den_t_fct = superelliptic_function(C, 1/den)
        if den != 1:
            # d([N/D] + V(f)) = [1/D]*d([N]) - [N]*[D^(-2)]*d([D]) + dV(f)
            return ((den_t_fct)^(-1)).teichmuller()*num_t_fct.teichmuller().diffn() - ((den_t_fct)^(-2)).teichmuller()*num_t_fct.teichmuller()*den_t_fct.teichmuller().diffn() + superelliptic_drw_form(0*C.x, 0*C.dx, f)
        t_polynomial = Rxy(t_polynomial)
        fC = superelliptic_function(C, fC)
        fC = fC.teichmuller()
        dy_w = 1/2* ((C.y)^(-1)).teichmuller()*auxilliary_derivative(fC)
        M = t_polynomial.monomials()[0]
        a = t_polynomial.monomial_coefficient(M)
        #[P] = [aM] + Q, where Q = ([P] - [aM])
        aM_fct = superelliptic_function(C, a*M)
        Q = self - aM_fct.teichmuller()
        exp_x = M.exponents()[0][0]
        exp_y = M.exponents()[0][1]
        return Q.diffn() + exp_x*superelliptic_drw_form(aM_fct/C.x, 0*C.dx, 0*C.x) + exp_y*(aM_fct/C.y).teichmuller()*dy_w

def binomial_prim(p, i):
    return binomial(p, i)/p
    
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
    return superelliptic_witt(fct, 0*C.x)

superelliptic_function.teichmuller = teichmuller

#dy = [f(x)]'/2*y dx
#[f1 + M] = [f1] + [M] + V(cos)
#d[f1 + M] = d[f1] + d[M] + dV(f1*M)
#M = b x^a
#d[M] = a*[b x^(a-1)]

def auxilliary_derivative(P):
    '''Return "derivative" of P, where P depends only on x. '''
    P0 = P.t.function
    P1 = P.f.function
    C = P.curve
    F = C.base_ring
    Rx.<x> = PolynomialRing(F)
    P0 = Rx(P0)
    P1 = Rx(P1)
    if P0 == 0:
        return superelliptic_drw_form(0*C.x, 0*C.dx, P.f)
    M = P0.monomials()[0]
    a = P0.monomial_coefficient(M)
    #[P] = [aM] + Q, where Q = ([P] - [aM])
    aM_fct = superelliptic_function(C, a*M)
    Q = P - aM_fct.teichmuller()
    exp = M.exponents()[0]
    return auxilliary_derivative(Q) + exp*superelliptic_drw_form(aM_fct/C.x, 0*C.dx, 0*C.x)
    
class superelliptic_drw_form:
    def __init__(self, h1, omega, h2):
        '''Form [h1] d[x] + V(omega) + dV([h])'''
        self.curve = h1.curve
        self.h1 = h1
        self.omega = omega
        self.h2 = h2
        
    def __eq__(self, other):
        eq1 = (self.h1 == self.h1)
        try:
            H = (self.h2 - other.h2).pth_root()
        except:
            return False
        eq2 = (self.omega - other.omega).cartier() - H.diffn()
        if eq1 and eq2:
            return True
        return False

    def __repr__(self):
        C = self.curve
        h1 = self.h1
        omega = self.omega
        h2 = self.h2
        result = ""
        if h1.function != 0:
            result += "[" + str(h1) + "] d[x]"
        if h1.function !=0 and omega.form != 0:
            result += " + "
        if omega.form != 0:
            result += "V(" + str(omega) + ")"
        if h2.function !=0 and omega.form != 0:
            result += " + "
        if h2.function != 0:
            result += "dV([" + str(h2) +"])"
        if result == "":
            result += "0"
        return result

    def __rmul__(self, other):
        h1 = self.h1
        h2 = self.h2
        omega = self.omega
        p = self.curve.characteristic
        if other in ZZ:
            if other == 0:
                return superelliptic_drw_form(0*C.x, 0*C.dx, 0*C.x)
            if other > 0:
                return self + (other-1)*self
            if other < 0:
                return (-other)*(-self)
        if other in QQ:
            other_integer = Integers(p^2)(other)^(-1)
            return other_integer*self
        t = other.t
        f = other.f
        aux_form = t^p*omega - h2*t^(p-1)*t.diffn() + f*h1^p*(C.x)^(p-1)*C.dx
        return superelliptic_drw_form(t*h1, aux_form, t^p*h2)

    def __neg__(self):
        C = self.curve
        h1 = self.h1
        h2 = self.h2
        omega = self.omega
        return superelliptic_drw_form(-h1, -omega, -h2)
    
    def __add__(self, other):
        C = self.curve
        h1 = self.h1
        h2 = self.h2
        omega = self.omega
        H1 = other.h1
        H2 = other.h2
        OMEGA = other.omega
        aux = (teichmuller(h1) + teichmuller(H1))*superelliptic_drw_form(C.one, 0*C.dx, 0*C.x)
        h1_aux = aux.h1
        h2_aux = aux.h2
        omega_aux = aux.omega
        return superelliptic_drw_form(h1_aux, omega + OMEGA + omega_aux, h2 + H2 + h2_aux)
    
    def frobenius(self):
        C = self.curve
        h1 = self.h1
        h2 = self.h2
        p = C.characteristic
        return h1^p*C.x^(p-1)*C.dx + h2.diffn()

def mult_by_p(omega):
    C = omega.curve
    fct = omega.form
    Fxy, Rxy, x, y = C.fct_field
    omega = superelliptic_form(C, fct^p * x^(p-1))
    result = superelliptic_drw_form(C, 0*C.dx, omega, 0*C.x)
    return result