class superelliptic_drw_cech:
    def __init__(self, omega0, f):
        self.curve = omega0.curve
        self.omega0 = omega0
        self.omega8 = omega0 - f.diffn()
        self.f = f
        
        
    def reduce(self):
        C = self.curve
        fct = self.f
        f_first_comp = fct.t
        f_second_comp = fct.f
        decomp_first_comp = decomposition_g0_g8(f_first_comp)
        decomp_second_comp = decomposition_g0_g8(f_second_comp)
        new = superelliptic_drw_cech(0*C.dx.verschiebung(), 0*C.x.verschiebung())
        new.omega0 = self.omega0
        new.omega8 = self.omega8
        new.f = self.f
        new.omega0 -= decomposition_g0_g8(f_first_comp)[0].teichmuller().diffn()
        new.omega0 -= decomposition_g0_g8(f_second_comp)[0].verschiebung().diffn()
        new.f = decomposition_g0_g8(f_first_comp)[2].teichmuller() + decomposition_g0_g8(f_second_comp)[2].verschiebung()
        new.omega8 = new.omega0 - new.f.diffn()
        return new
        
    def __repr__(self):
        return("(" + str(self.omega0) + ", "+ str(self.f) + ", " + str(self.omega8) + ")")

    def __add__(self, other):
        C = self.curve
        omega0 = self.omega0
        f = self.f
        omega0_1 = other.omega0
        f_1 = other.f
        return superelliptic_drw_cech(omega0 + omega0_1, f + f_1)
    
    def __sub__(self, other):
        C = self.curve
        omega0 = self.omega0
        f = self.f
        omega0_1 = other.omega0
        f_1 = other.f
        return superelliptic_drw_cech(omega0 - omega0_1, f - f_1)
        
    def __neg__(self):
        C = self.curve
        omega0 = self.omega0
        f = self.f
        return superelliptic_drw_cech(-omega0, -f)
    
    def __rmul__(self, other):
        omega0 = self.omega0
        f = self.f
        return superelliptic_drw_cech(other*omega0, other*f)
    
    def r(self):
        omega0 = self.omega0
        f = self.f
        C = self.curve
        return superelliptic_cech(C, omega0.h1*C.dx, f.t)
    
    def div_by_p(self, info = 0):
        '''Given a regular cocycle of the form (V(omega) + dV(h), [f] + V(t), ...), where [f] = 0 in H^1(X, OX),
        find de Rham cocycle (xi0, f, xi8) such that (V(omega) + dV(h), [f] + V(t), ...) = p*(xi0, f, xi8).'''
        #
        print('div by p', self)
        if info: print("Computing " + str(self) + " divided by p.")
        #
        C = self.curve
        F = C.base_ring
        p = F.characteristic()
        aux = self
        Fxy, Rxy, x, y = C.fct_field
        aux_f_t_0 = decomposition_g0_g8(aux.f.t, prec=50)[0]
        #
        if info: print("Computed decomposition_g0_g8 of self.f.t.")
        #
        aux.f.t = 0*C.x
        aux.omega0 -= aux_f_t_0.teichmuller().diffn()
        aux.omega8 = aux.omega0 - aux.f.diffn()
        #
        # Now omega = (regular on U0) + h^(p-1) dh, so that Cartier(omega) = (regular on U0) + dh.
        # We replace omega by regular on U0
        omega = aux.omega0.omega
        aux.omega0.omega, fct = decomposition_omega0_hpdh(aux.omega0.omega)
        aux.omega0.h2 += fct^p
        #
        if info: print("Computed decomposition_omega0_hpdh of self.omega0.omega.")
        #
        # Now we have to ensure that aux.omega0.h2.function in Rxy...
        # In other words, we decompose h2 = (regular on U0) + A^p.
        # Then we replace h2 by (regular on U0) and omega by omega + inverse Cartier(dA)
        aux.omega0.h2, A = decomposition_g0_pth_power(aux.omega0.h2)
        aux.omega0.omega += (A.diffn()).inv_cartier()
        #
        if info: print("Computed decomposition_g0_p2th_power(aux.omega0.h2).")
        #
        # Now we can reduce: (... + dV(h2), V(f), ...) --> (..., V(f - h2), ...)
        aux.f -= aux.omega0.h2.verschiebung()
        aux.omega0.h2 = 0*C.x
        # Now the same for aux.omega8
        aux.omega8.omega, fct = decomposition_omega8_hpdh(aux.omega8.omega)
        aux.omega8.h2 += fct^p
        aux.omega8.h2 = decomposition_g8_p2th_power(aux.omega8.h2)[0]
        aux.f += aux.omega8.h2.verschiebung()
        aux.omega8.h2 = 0*C.x
        aux_divided_by_p = superelliptic_cech(C, aux.omega0.omega.cartier(), aux.f.f.pth_root())
        return aux_divided_by_p
        
    def coordinates(self, basis = 0, prec = 50, info = 0):
        if info: print("Computing coordinates of " + str(self))
        C = self.curve
        F = C.base_ring
        p = F.characteristic()
        g = C.genus()
        coord_mod_p = self.r().coordinates()
        if info: print("Computed coordinates mod p.")
        coord_lifted = [lift(a) for a in coord_mod_p]
        if basis == 0:
            basis = C.crystalline_cohomology_basis()
        aux = self
        for i, a in enumerate(basis):
            aux -= coord_lifted[i]*a
        print("aux!!!", aux)
        aux_divided_by_p = aux.div_by_p(info = info)
        print("aux by p!!!", aux_divided_by_p)
        print("check", mult_by_p(aux_divided_by_p)-aux)
        coord_aux_divided_by_p = aux_divided_by_p.coordinates()
        if info: print("Computed coordinates mod p of (self - lift)/p.")
        coord_aux_divided_by_p = [ZZ(a) for a in coord_aux_divided_by_p]
        coordinates = [ (coord_lifted[i] + p*coord_aux_divided_by_p[i])%p^2 for i in range(2*g)]
        return coordinates
    
    def is_regular(self):
        eq1 = self.omega0.r().is_regular_on_U0() and self.omega8.r().is_regular_on_Uinfty()
        eq2 = self.omega0.frobenius().is_regular_on_U0() and self.omega8.frobenius().is_regular_on_Uinfty()
        return eq1 and eq2