class superelliptic_cohomology:
    '''Cohomology of structure sheaf.'''
    def __init__(self, C, f):
        self.curve = C
        self.function = f
        
    def prod(self, form, prec=20):
        result = 0
        C = self.curve
        delta = C.nb_of_pts_at_infty
        fct = superelliptic_function(C, self.function)
        for i in range(delta):
            result += (fct*form).expansion_at_infty(place=i, prec=prec)[-1]
        return -result
    
    def __repr__(self):
        return "[" + str(self.function) + "]"
    
    def __add__(self, other):
        C = self.curve
        g1 = self.function
        g2 = other.function
        g = reduction(C, g1 + g2)
        return superelliptic_cohomology(C, g)
    
    def __sub__(self, other):
        C = self.curve
        g1 = self.function
        g2 = other.function
        g = reduction(C, g1 - g2)
        return superelliptic_cohomology(C, g)
    
    def __rmul__(self, constant):
        C = self.curve
        g = self.function
        return superelliptic_cohomology(C, constant*g)

    def coordinates(self, basis = 0, basis_holo = 0, prec=20):
        C = self.curve
        if basis == 0:
            basis = basis_of_cohomology(C)
        if basis_holo == 0:
            basis_holo = C.holomorphic_differentials_basis()
        g = C.genus()
        coordinates = g*[0]
        for i, omega in enumerate(basis_holo):
            coordinates[i] = self.prod(omega, prec=prec)
        return coordinates
    
    def frobenius(self):
        C = self.curve
        f = self.function
        p = C.base_ring.characteristic()
        return superelliptic_cohomology(C, f^p)
    
def basis_of_cohomology(C):
    m = C.exponent
    f = C.polynomial
    r = f.degree()
    F = C.base_ring
    delta = C.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*j - m*i >= delta):
                    basis += [superelliptic_cohomology(C, Fxy(m*y^(j)/x^i))]
    return basis

def frobenius_matrix(C, prec=20):
    g = C.genus()
    F = C.base_ring
    M = matrix(F, g, g)
    for i, f in enumerate(basis_of_cohomology(C)):
        M[i, :] = vector(f.frobenius().coordinates(prec=prec))
    return M


def frobenius_kernel(C, prec=20):
    M = frobenius_matrix(C, prec=prec)
    K = M.kernel().basis()
    g = C.genus()
    result = []
    basis = basis_of_cohomology(C)
    for v in K:
        coh = superelliptic_cohomology(C, 0)
        for i in range(g):
            coh += v[i] * basis[i]
        result += [coh]
    return result