class superelliptic:
    """Class of a superelliptic curve. Given a polynomial f(x) with coefficient field F, it constructs
       the curve y^m = f(x)"""
    def __init__(self, f, m):
        Rx = f.parent()
        x = Rx.gen()
        F = Rx.base()        
        Rx.<x> = PolynomialRing(F)
        Rxy.<x, y> = PolynomialRing(F, 2)
        Fxy = FractionField(Rxy)
        self.polynomial = Rx(f)
        self.exponent = m
        self.base_ring = F
        self.characteristic = F.characteristic()
        self.fct_field = Fxy, Rxy, x, y
        r = Rx(f).degree()
        delta = GCD(r, m)
        self.nb_of_pts_at_infty = delta
        self.x = superelliptic_function(self, Rxy(x))
        self.y = superelliptic_function(self, Rxy(y))
        self.dx = superelliptic_form(self, Rxy(1))
        self.one = superelliptic_function(self, Rxy(1))
        
    def __repr__(self):
        f = self.polynomial
        m = self.exponent
        F = self.base_ring
        return 'Superelliptic curve with the equation y^' + str(m) + ' = ' + str(f)+' over ' + str(F)
    
    #Auxilliary algorithm that returns the basis of holomorphic differentials
    #of the curve and (as a second argument) the list of pairs (i, j)
    #such that x^i dx/y^j is holomorphic.
    def basis_holomorphic_differentials_degree(self):
        f = self.polynomial
        m = self.exponent
        r = f.degree()
        delta = GCD(r, m)
        F = self.base_ring
        Rx.<x> = PolynomialRing(F)
        Rxy.<x, y> = PolynomialRing(F, 2)
        Fxy = FractionField(Rxy)
        #########basis of holomorphic differentials and de Rham

        basis_holo = []
        degrees0 = {}
        k = 0

        for j in range(1, m):
            for i in range(1, r):
                if (r*j - m*i >= delta):
                    basis_holo += [superelliptic_form(self, Fxy(x^(i-1)/y^j))]
                    degrees0[k] = (i-1, j)
                    k = k+1

        return(basis_holo, degrees0)
    
    #Returns the basis of holomorphic differentials using the previous algorithm.
    def holomorphic_differentials_basis(self):
        basis_holo, degrees0 = self.basis_holomorphic_differentials_degree()
        return basis_holo
    #Returns the list of pairs (i, j) such that x^i dx/y^j is holomorphic.
    def degrees_holomorphic_differentials(self):
        basis_holo, degrees0 = self.basis_holomorphic_differentials_degree()
        return degrees0

    def basis_de_rham_degrees(self):
        f = self.polynomial
        m = self.exponent
        r = f.degree()
        delta = GCD(r, m)
        F = self.base_ring
        Rx.<x> = PolynomialRing(F)
        Rxy.<x, y> = PolynomialRing(F, 2)
        Fxy = FractionField(Rxy)
        basis_holo = self.holomorphic_differentials_basis()
        basis = []
        #First g_X elements of basis are holomorphic differentials.
        for k in range(0, len(basis_holo)):
            basis += [superelliptic_cech(self, basis_holo[k], superelliptic_function(self, 0))]

        ## Next elements do not come from holomorphic differentials.
        t = len(basis)
        degrees0 = {}
        degrees1 = {}
        for j in range(1, m):
            for i in range(1, r):
                if (r*(m-j) - m*i >= delta): 
                    s = Rx(m-j)*Rx(x)*Rx(f.derivative()) - Rx(m)*Rx(i)*f
                    psi = Rx(cut(s, i))
                    basis += [superelliptic_cech(self, superelliptic_form(self, Fxy(psi/y^j)), superelliptic_function(self, Fxy(m*y^(m-j)/x^i)))]
                    degrees0[t] = (psi.degree(), j)
                    degrees1[t] = (-i, m-j)
                    t += 1
        return basis, degrees0, degrees1

    def de_rham_basis(self):
        basis, degrees0, degrees1 = self.basis_de_rham_degrees()
        return basis

    def degrees_de_rham0(self):
        basis, degrees0, degrees1 = self.basis_de_rham_degrees()
        return degrees0

    def degrees_de_rham1(self):
        basis, degrees0, degrees1 = self.basis_de_rham_degrees()
        return degrees1    
    
    def is_smooth(self):
        f = self.polynomial
        if f.discriminant() == 0:
            return 0
        return 1
    
    def genus(self):
        r = self.polynomial.degree()
        m = self.exponent
        delta = GCD(r, m)
        return 1/2*((r-1)*(m-1) - delta + 1)
    
    def verschiebung_matrix(self):
        basis = self.de_rham_basis()
        g = self.genus()
        p = self.characteristic
        F = self.base_ring
        M = matrix(F, 2*g, 2*g)
        for i in range(0, len(basis)):
            w = basis[i]
            v = w.verschiebung().coordinates()
            M[i, :] = v
        return M
    
    def dr_frobenius_matrix(self):
        basis = self.de_rham_basis()
        g = self.genus()
        p = self.characteristic
        F = self.base_ring
        M = matrix(F, 2*g, 2*g)
        
        for i in range(0, len(basis)):
            w = basis[i]
            v = w.frobenius().coordinates()
            M[i, :] = v
        return M

    def cartier_matrix(self):
        basis = self.holomorphic_differentials_basis()
        g = self.genus()
        p = self.characteristic
        F = self.base_ring
        M = matrix(F, g, g)
        for i in range(0, len(basis)):
            w = basis[i]
            v = w.cartier().coordinates()
            M[:, i] = vector(v)
        return M

    def frobenius_matrix(self, prec=20):
        g = self.genus()
        F = self.base_ring
        p = F.characteristic()
        M = matrix(F, g, g)
        for i, f in enumerate(self.cohomology_of_structure_sheaf_basis()):
            M[i, :] = vector((f^p).coordinates(prec=prec))
        M = M.transpose()
        return M
    
    def p_rank(self):
        if self.exponent != 2:
            raise ValueError('No implemented yet.')
        f = self.polynomial()
        F = self.base_ring
        Rt.<t> = PolynomialRing(F)
        f = Rt(f)
        H = HyperellipticCurve(f, 0)
        return H.p_rank()
    
    def a_number(self):
        g = self.genus()
        return g - self.cartier_matrix().rank()
    
    def final_type(self, test = 0):
        Fr = self.frobenius_matrix()
        V = self.verschiebung_matrix()
        p = self.characteristic
        return flag(Fr, V, p, test)
    
    def cohomology_of_structure_sheaf_basis(self):
        '''Basis of cohomology of structure sheaf H1(X, OX).'''
        m = self.exponent
        f = self.polynomial
        r = f.degree()
        F = self.base_ring
        delta = self.nb_of_pts_at_infty
        Rx.<x> = PolynomialRing(F)
        Rxy.<x, y> = PolynomialRing(F, 2)
        Fxy = FractionField(Rxy)
        basis = []
        for j in range(1, m):
            for i in range(1, r):
                if (r*(m-j) - m*i >= delta): 
                    basis += [superelliptic_function(self, Fxy(m*y^(m-j)/x^i))]
        return basis

    def uniformizer(self):
        m = self.exponent
        r = self.polynomial.degree()
        delta, a, b = xgcd(m, r)
        a = -a
        M = m/delta
        R = r/delta
        while a<0:
            a += R
            b += M
        return (C.x)^a/(C.y)^b

def reduction(C, g):
    '''Auxilliary. Given a superelliptic curve C : y^m = f(x) and a polynomial g(x, y)
    it replaces repeteadly all y^m's in g(x, y) by f(x). As a result
    you obtain \sum_{i = 0}^{m-1} y^i g_i(x).'''
    p = C.characteristic
    F = C.base_ring
    Rxy.<x, y> = PolynomialRing(F, 2)
    Fxy = FractionField(Rxy)
    f = C.polynomial
    r = f.degree()
    m = C.exponent
    g = Fxy(g)
    g1 = g.numerator()
    g2 = g.denominator()
    
    Rx.<x> = PolynomialRing(F)
    Fx = FractionField(Rx)
    FxRy.<y> = PolynomialRing(Fx)    
    (A, B, C) = xgcd(FxRy(g2), FxRy(y^m - f))
    g = FxRy(g1*B/A)
    
    while(g.degree(Rxy(y)) >= m):
        d = g.degree(Rxy(y))
        G = coff(g, d)
        i = floor(d/m)
        g = g - G*y^d + f^i * y^(d%m) *G
    
    return(FxRy(g))

#Auxilliary. Given a superelliptic curve C : y^m = f(x) and a polynomial g(x, y)
#it replaces repeteadly all y^m's in g(x, y) by f(x). As a result
#you obtain \sum_{i = 0}^{m-1} g_i(x)/y^i. This is needed for reduction of
#superelliptic forms.
def reduction_form(C, g):
    F = C.base_ring
    Rxy.<x, y> = PolynomialRing(F, 2)
    Fxy = FractionField(Rxy)
    f = C.polynomial
    r = f.degree()
    m = C.exponent
    g = reduction(C, g)

    g1 = Rxy(0)
    Rx.<x> = PolynomialRing(F)
    Fx = FractionField(Rx)
    FxRy.<y> = PolynomialRing(Fx)
    
    g = FxRy(g)
    for j in range(0, m):
        if j==0:
            G = coff(g, 0)
            g1 += FxRy(G)
        else:
            G = coff(g, j)
            g1 += Fxy(y^(j-m)*f*G)
    return(g1)