class heisenberg_cover:
    def __init__(self, C, list_of_fcts, branch_points = [], prec = 10):
        self.quotient = C
        self.functions = list_of_fcts
        self.height = len(list_of_fcts)
        F = C.base_ring
        self.base_ring = F
        p = C.characteristic
        self.characteristic = p
        self.prec = prec
        #group acting
        n = self.height
        from itertools import product
        pr = [list(GF(p)) for _ in range(n)]
        self.group = [(a, b, c) for a in range(p) for b in range(p) for c in range(p)]
        #########
        f = C.polynomial
        m = C.exponent
        r = f.degree()
        delta = GCD(m, r)
        self.nb_of_pts_at_infty = delta
        self.branch_points = list(range(delta)) + branch_points
        Rxy.<x, y> = PolynomialRing(F, 2)
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)

        all_x_series = {}
        all_y_series = {}
        all_z_series = {}
        all_dx_series = {}
        all_jumps = {}

        for pt in self.branch_points:
            x_series = superelliptic_function(C, x).expansion(pt=pt, prec=prec)
            y_series = superelliptic_function(C, y).expansion(pt=pt, prec=prec)
            z_series = []
            jumps = []
            n = len(list_of_fcts)
            list_of_power_series = [g.expansion(pt=pt, prec=prec) for g in list_of_fcts]
            for j in range(n):
                if j == 2:
                    list_of_power_series[j] += (z_series[0] - z_series[1])*list_of_power_series[1]
                power_series = list_of_power_series[j]
                jump, correction, t_old, z = artin_schreier_transform(power_series, prec = prec)
                x_series = x_series(t = t_old)
                y_series = y_series(t = t_old)
                z_series = [zi(t = t_old) for zi in z_series]
                z_series += [z]
                jumps += [jump]
                list_of_power_series = [g(t = t_old) for g in list_of_power_series]
            
            all_jumps[pt] = jumps
            all_x_series[pt] = x_series
            all_y_series[pt] = y_series
            all_z_series[pt] = z_series
            all_dx_series[pt] = x_series.derivative()
        self.jumps = all_jumps
        self.x_series = all_x_series
        self.y_series = all_y_series
        self.z_series = all_z_series
        self.dx_series = all_dx_series
        ##############
        #Function field
        variable_names = 'x, y'
        for i in range(n):
            variable_names += ', z' + str(i)
        Rxyz = PolynomialRing(F, n+2, variable_names)
        x, y = Rxyz.gens()[:2]
        z = Rxyz.gens()[2:]
        RxyzQ = FractionField(Rxyz)
        self.fct_field = (RxyzQ, Rxyz, x, y, z)
        self.x = heisenberg_function(self, x)
        self.y = heisenberg_function(self, y)
        self.z = [heisenberg_function(self, z[j]) for j in range(n)]
        self.dx = heisenberg_form(self, 1)
        self.one = heisenberg_function(self, 1)
        
        
    def __repr__(self):
        n = self.height
        p = self.characteristic
        if n==1:
            return "(Z/p)-cover of " + str(self.quotient)+" with the equation:\n z^" + str(p) + " - z = " + str(self.functions[0])
        
        result = "E(p^3)-cover of " + str(self.quotient)+" with the equations:\n"
        for i in range(n):
            if i !=2:
                result += 'z' + str(i) + "^" + str(p) + " - z" + str(i) + " = " + str(self.functions[i]) + "\n"
            else:
                result += 'z2'+ "^" + str(p) + ' - z2 = ' + str(self.functions[2]) + ' + ' + '(z0 - z1) * (' + str(self.functions[1]) + ")"
        return result
        
    def genus(self):
        jumps = self.jumps
        gY = self.quotient.genus()
        n = self.height
        branch_pts = self.branch_points
        p = self.characteristic
        return p^n*gY + (p^n - 1)*(len(branch_pts) - 1) + sum(p^(n-j-1)*(jumps[pt][j]-1)*(p-1)/2 for j in range(n) for pt in branch_pts)
    
    def exponent_of_different(self, place = 0):
        jumps = self.jumps
        n = self.height
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        return sum(p^(n-j-1)*(jumps[place][j]+1)*(p-1) for j in range(n))

    def exponent_of_different_prim(self, place = 0):
        jumps = self.jumps
        n = self.height
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        return sum(p^(n-j-1)*(jumps[place][j])*(p-1) for j in range(n))
    
    def holomorphic_differentials_basis(self, threshold = 8):
        from itertools import product
        x_series = self.x_series
        y_series = self.y_series
        z_series = self.z_series
        dx_series = self.dx_series
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        n = self.height
        prec = self.prec
        C = self.quotient
        F = self.base_ring
        m = C.exponent
        r = C.polynomial.degree()
        RxyzQ, Rxyz, x, y, z = self.fct_field
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        #Tworzymy zbiór S form z^i x^j y^k dx/y o waluacji >= waluacja z^(p-1)*dx/y
        S = []
        pr = [list(GF(p)) for _ in range(n)]
        for i in range(0, threshold*r):
            for j in range(0, m):
                for k in product(*pr):
                    eta = heisenberg_form(self, x^i * prod(z[i1]^(k[i1]) for i1 in range(n))/y^j)
                    eta_exp = eta.expansion(pt=self.branch_points[0])
                    S += [(eta, eta_exp)]

        forms = heisenberg_holomorphic_combinations(S)
        
        for pt in self.branch_points[1:]:
            forms = [(omega, omega.expansion(pt=pt)) for omega in forms]
            forms = heisenberg_holomorphic_combinations(forms)
        
        if len(forms) < self.genus():
            print("I haven't found all forms, only ", len(forms), " of ", self.genus())
            return holomorphic_differentials_basis(self, threshold = threshold + 1)
        if len(forms) > self.genus():
            raise ValueError("Increase precision.")
        return forms
    
    def cartier_matrix(self, prec=50):
        g = self.genus()
        F = self.base_ring
        M = matrix(F, g, g)
        for i, omega in enumerate(self.holomorphic_differentials_basis()):
            M[:, i] = vector(omega.cartier().coordinates())
        return M

    def at_most_poles(self, pole_order, threshold = 8):
        """ Find fcts with pole order in infty's at most pole_order. Threshold gives a bound on powers of x in the function.
           If you suspect that you haven't found all the functions, you may increase it."""
        from itertools import product
        x_series = self.x_series
        y_series = self.y_series
        z_series = self.z_series
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        n = self.height
        prec = self.prec
        C = self.quotient
        F = self.base_ring
        m = C.exponent
        r = C.polynomial.degree()
        RxyzQ, Rxyz, x, y, z = self.fct_field
        F = C.base_ring
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        #Tworzymy zbiór S form z^i x^j y^k dx/y o waluacji >= waluacja z^(p-1)*dx/y
        S = []
        RQxyz = FractionField(Rxyz)
        pr = [list(GF(p)) for _ in range(n)]
        for i in range(0, threshold*r):
            for j in range(0, m):
                for k in product(*pr):
                    eta = heisenberg_function(self, x^i * prod(z[i1]^(k[i1]) for i1 in range(n))*y^j)
                    eta_exp = eta.expansion_at_infty()
                    S += [(eta, eta_exp)]

        forms = holomorphic_combinations_fcts(S, pole_order)

        for i in range(1, delta):
            forms = [(omega, omega.expansion_at_infty(place = i)) for omega in forms]
            forms = holomorphic_combinations_fcts(forms, pole_order)
                  
        return forms
    
    def magical_element(self, threshold = 8):
        list_of_elts = self.at_most_poles(self.exponent_of_different_prim(), threshold)
        result = []
        for a in list_of_elts:
            if a.trace().function != 0:
                result += [a]
        return result

    def pseudo_magical_element(self, threshold = 8):
        list_of_elts = self.at_most_poles(self.exponent_of_different(), threshold)
        result = []
        for a in list_of_elts:
            if a.trace().function != 0:
                result += [a]
        return result

    def at_most_poles_forms(self, pole_order, threshold = 8):
        """Find forms with pole order in all the points at infty equat at most to pole_order. Threshold gives a bound on powers of x in the form.
           If you suspect that you haven't found all the functions, you may increase it."""
        from itertools import product
        x_series = self.x_series
        y_series = self.y_series
        z_series = self.z_series
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        n = self.height
        prec = self.prec
        C = self.quotient
        F = self.base_ring
        m = C.exponent
        r = C.polynomial.degree()
        RxyzQ, Rxyz, x, y, z = self.fct_field
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        #Tworzymy zbiór S form z^i x^j y^k dx/y o waluacji >= waluacja z^(p-1)*dx/y
        S = []
        RQxyz = FractionField(Rxyz)
        pr = [list(GF(p)) for _ in range(n)]
        for i in range(0, threshold*r):
            for j in range(0, m):
                for k in product(*pr):
                    eta = heisenberg_form(self, x^i * prod(z[i1]^(k[i1]) for i1 in range(n))/y^j)
                    eta_exp = eta.expansion_at_infty()
                    S += [(eta, eta_exp)]

        forms = holomorphic_combinations_forms(S, pole_order)

        for pt in self.branch_points[1:]:
            forms = [(omega, omega.expansion(pt=pt)) for omega in forms]
            forms = holomorphic_combinations_forms(forms, pole_order)
                  
        return forms
    
    def uniformizer(self, place = 0):
        '''Return uniformizer of curve self at place-th place at infinity.'''
        p = self.characteristic
        n = self.height
        F = self.base_ring
        RxyzQ, Rxyz, x, y, z = self.fct_field
        fx = heisenberg_function(self, x)
        z = [heisenberg_function(self, zi) for zi in z]
        # We create a list of functions. We add there all variables...
        list_of_fcts = [fx]+z
        vfx = fx.valuation(place)
        vz = [zi.valuation(place) for zi in z]
        
        # Then we subtract powers of variables with the same valuation (so that 1/t^(kp) cancels) and add to this list.
        for  j1 in range(n):
            for j2 in range(n):
                if j1>j2:
                    a = gcd(vz[j1] , vz[j2])
                    vz1 = vz[j1]/a
                    vz2 = vz[j2]/a
                    for b in F:
                        if (z[j1]^(vz2) - b*z[j2]^(vz1)).valuation(place) > (z[j2]^(vz1)).valuation(place):
                            list_of_fcts += [z[j1]^(vz2) - b*z[j2]^(vz1)]
        for j1 in range(n):
            a = gcd(vz[j1], vfx)
            vzj = vz[j1] /a
            vfx = vfx/a
            for b in F:
                if (fx^(vzj) - b*z[j1]^(vfx)).valuation(place) > (z[j1]^(vfx)).valuation(place):
                    list_of_fcts += [fx^(vzj) - b*z[j1]^(vfx)]
        #Finally, we check if on the list there are two elements with the same valuation.
        for f1 in list_of_fcts:
            for f2 in list_of_fcts:
                d, a, b = xgcd(f1.valuation(place), f2.valuation(place))
                if d == 1:
                    return f1^a*f2^b
        raise ValueError("My method of generating fcts with relatively prime valuation failed.")
        

    def ith_ramification_gp(self, i, place = 0):
        '''Find ith ramification group at place at infty of nb place.'''
        G = self.group
        t = self.uniformizer(place)
        Gi = [G[0]]
        for g in G:
            if g != G[0]:
                tg = t.group_action(g)
                v = (tg - t).valuation(place)
                if v >= i+1:
                    Gi += [g]
        return Gi
    
    def ramification_jumps(self, place = 0):
        '''Return list of lower ramification jumps at at place at infty of nb place.'''
        G = self.group
        ramification_jps = []
        i = 0
        while len(G) > 1:
            Gi = self.ith_ramification_gp(i+1, place)
            if len(Gi) < len(G):
                ramification_jps += [i]
            G = Gi
            i+=1
        return ramification_jps
    
    def a_number(self):
        g = self.genus()
        return g - self.cartier_matrix().rank()

    def cohomology_of_structure_sheaf_basis(self, holo_basis = 0, threshold = 8):
        if holo_basis == 0:
            holo_basis = self.holomorphic_differentials_basis(threshold = threshold)
        from itertools import product
        x_series = self.x_series
        y_series = self.y_series
        z_series = self.z_series
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        n = self.height
        prec = self.prec
        C = self.quotient
        F = self.base_ring
        m = C.exponent
        r = C.polynomial.degree()
        RxyzQ, Rxyz, x, y, z = self.fct_field
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        #Tworzymy zbiór S form z^i x^j y^k dx/y o waluacji >= waluacja z^(p-1)*dx/y
        result_fcts = []
        V = VectorSpace(F,self.genus())
        S = V.subspace([])
        RQxyz = FractionField(Rxyz)
        pr = [list(GF(p)) for _ in range(n)]
        i = 0
        while len(result_fcts) < self.genus():
            for j in range(0, m):
                for k in product(*pr):
                    f = heisenberg_function(self, prod(z[i1]^(k[i1]) for i1 in range(n))/x^i*y^j)
                    f_products = [omega.serre_duality_pairing(f) for omega in holo_basis]
                    if vector(f_products) not in S:
                        S = S+V.subspace([V(f_products)])
                        result_fcts += [f]
            i += 1
        return result_fcts

    def lift_to_de_rham(self, fct, threshold = 30):
        '''Given function fct, find form eta regular on affine part such that eta - d(fct) is regular in infty. (Works for one place at infty now)'''
        from itertools import product
        x_series = self.x_series
        y_series = self.y_series
        z_series = self.z_series
        dx_series = self.dx_series
        delta = self.nb_of_pts_at_infty
        p = self.characteristic
        n = self.height
        prec = self.prec
        C = self.quotient
        F = self.base_ring
        m = C.exponent
        r = C.polynomial.degree()
        RxyzQ, Rxyz, x, y, z = self.fct_field
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        #Tworzymy zbiór S form z^i x^j y^k dx/y o waluacji >= waluacja z^(p-1)*dx/y
        S = [(fct.diffn(), fct.diffn().expansion_at_infty())]
        pr = [list(GF(p)) for _ in range(n)]
        holo = self.holomorphic_differentials_basis(threshold = threshold)
        for i in range(0, threshold*r):
            for j in range(0, m):
                for k in product(*pr):
                    eta = heisenberg_form(self, x^i*prod(z[i1]^(k[i1]) for i1 in range(n))/y^j)
                    eta_exp = eta.expansion_at_infty()
                    S += [(eta, eta_exp)]
        forms = heisenberg_holomorphic_combinations(S)
        if len(forms) <= self.genus():
            raise ValueError("Increase threshold!")
        for omega in forms:
            for a in F:
                if (a*omega + fct.diffn()).is_regular_on_U0():
                    return a*omega + fct.diffn()
        raise ValueError("Unknown.")

    def de_rham_basis(self, holo_basis = 0, cohomology_basis = 0, threshold = 30):
        if holo_basis == 0:
            holo_basis = self.holomorphic_differentials_basis(threshold = threshold)
        if cohomology_basis == 0:
            cohomology_basis = self.cohomology_of_structure_sheaf_basis(holo_basis = holo_basis, threshold = threshold)
        result = []
        for omega in holo_basis:
            result += [heisenberg_cech(self, omega, heisenberg_function(self, 0))]
        for f in cohomology_basis:
            omega = self.lift_to_de_rham(f, threshold = threshold)
            result += [heisenberg_cech(self, omega, f)]
        return result
    
def heisenberg_holomorphic_combinations(S):
    """Given a list S of pairs (form, corresponding Laurent series at some pt), find their combinations holomorphic at that pt."""
    C_AS = S[0][0].curve
    p = C_AS.characteristic
    F = C_AS.base_ring
    prec = C_AS.prec
    Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
    RtQ = FractionField(Rt)
    minimal_valuation = min([g[1].valuation() for g in S])
    if minimal_valuation >= 0:
        return [s[0] for s in S]
    list_of_lists = [] #to będzie lista złożona z list współczynników część nieholomorficznych rozwinięcia form z S
    for eta, eta_exp in S:
        a = -minimal_valuation + eta_exp.valuation()
        list_coeffs = a*[0] + eta_exp.list() + (-minimal_valuation)*[0]
        list_coeffs = list_coeffs[:-minimal_valuation]
        list_of_lists += [list_coeffs]
    M = matrix(F, list_of_lists)
    V = M.kernel() #chcemy wyzerować części nieholomorficzne, biorąc kombinacje form z S


    # Sprawdzamy, jakim formom odpowiadają elementy V.
    forms = []
    for vec in V.basis():
        forma_holo = heisenberg_form(C_AS, 0)
        forma_holo_power_series = Rt(0)
        for vec_wspolrzedna, elt_S in zip(vec, S):
            eta = elt_S[0]
            #eta_exp = elt_S[1]
            forma_holo += vec_wspolrzedna*eta
            #forma_holo_power_series += vec_wspolrzedna*eta_exp
        forms += [forma_holo]
    return forms