class heisenberg_form:
    def __init__(self, C, g):
        self.curve = C
        n = C.height
        F = C.base_ring
        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.form = RxyzQ(g)
        
    def __repr__(self):
        return "(" + str(self.form)+") * dx"
    
    def expansion_at_infty(self, place = 0):
        C = self.curve
        delta = C.nb_of_pts_at_infty
        F = C.base_ring
        x_series = C.x_series[place]
        y_series = C.y_series[place]
        z_series = C.z_series[place]
        dx_series = C.dx_series[place]
        n = C.height
        variable_names = 'x, y'
        for j in range(n):
            variable_names += ', z' + str(j)
        Rxyz = PolynomialRing(F, n+2, variable_names)
        x, y = Rxyz.gens()[:2]
        z = Rxyz.gens()[2:]
        RxyzQ = FractionField(Rxyz)
        prec = C.prec
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        g = self.form
        sub_list = {x : x_series, y : y_series} | {z[j] : z_series[j] for j in range(n)}
        return g.substitute(sub_list)*dx_series
    
    def expansion(self, pt = 0):
        '''Same code as expansion_at_infty.'''
        C = self.curve
        F = C.base_ring
        x_series = C.x_series[pt]
        y_series = C.y_series[pt]
        z_series = C.z_series[pt]
        dx_series = C.dx_series[pt]
        n = C.height
        variable_names = 'x, y'
        for j in range(n):
            variable_names += ', z' + str(j)
        Rxyz = PolynomialRing(F, n+2, variable_names)
        x, y = Rxyz.gens()[:2]
        z = Rxyz.gens()[2:]
        RxyzQ = FractionField(Rxyz)
        prec = C.prec
        Rt.<t> = LaurentSeriesRing(F, default_prec=prec)
        g = self.form
        sub_list = {x : x_series, y : y_series} | {z[j] : z_series[j] for j in range(n)}
        return g.substitute(sub_list)*dx_series
    
    def __add__(self, other):
        C = self.curve
        g1 = self.form
        g2 = other.form
        return heisenberg_form(C, g1 + g2)
    
    def __sub__(self, other):
        C = self.curve
        g1 = self.form
        g2 = other.form
        return heisenberg_form(C, g1 - g2)
    
    def __neg__(self):
        C = self.curve
        g = self.form
        return heisenberg_form(C, -g)
    
    def __rmul__(self, constant):
        C = self.curve
        omega = self.form
        return heisenberg_form(C, constant*omega)

    def group_action(self, elt):
        AS = self.curve
        RxyzQ, Rxyz, x, y, z = AS.fct_field
        if elt == (0, 0, 0):
            return self
        if elt == (1, 0, 0):
            sub_list = {x : x, y : y} | {z[0] : z[0] + 1, z[1] : z[1], z[2]: z[2] + z[1]}
            g = self.form
            return heisenberg_form(AS, g.substitute(sub_list))
        if elt == (0, 1, 0):
            sub_list = {x : x, y : y} | {z[0] : z[0] + 1, z[1] : z[1] + 1, z[2]: z[2]}
            g = self.form
            return heisenberg_form(AS, g.substitute(sub_list))
        if elt == (0, 0, 1):
            sub_list = {x : x, y : y} | {z[0] : z[0], z[1] : z[1], z[2]: z[2] - 1}
            g = self.form
            return heisenberg_form(AS, g.substitute(sub_list))
        
        if elt[0] > 0:
            elt1 = (elt[0] - 1, elt[1], elt[2])
            return self.group_action(elt1).group_action((1, 0, 0))
        
        if elt[1] > 0:
            elt1 = (elt[0], elt[1] - 1, elt[2])
            return self.group_action(elt1).group_action((0, 1, 0))
        
        if elt[2] > 0:
            elt1 = (elt[0], elt[1], elt[2] - 1)
            return self.group_action(elt1).group_action((0, 0, 1))

    def coordinates(self, basis = 0):
        """Find coordinates of the given holomorphic form self in terms of the basis forms in a list holo."""
        C = self.curve
        if basis == 0:
            basis = C.holomorphic_differentials_basis()
        RxyzQ, Rxyz, x, y, z = C.fct_field
        # We need to have only polynomials to use monomial_coefficients in linear_representation_polynomials,
        # and sometimes basis elements have denominators. Thus we multiply by them.
        fct = heisenberg_function(C, self.form)
        fct = heisenberg_reduction(C, fct)
        self = heisenberg_form(C, fct)
        denom = LCM([denominator(omega.form) for omega in basis])
        basis = [denom*omega for omega in basis]
        self_with_no_denominator = denom*self
        return linear_representation_polynomials(Rxyz(self_with_no_denominator.form), [Rxyz(omega.form) for omega in basis])

    def trace(self):
        C = self.curve
        C_super = C.quotient
        n = C.height
        F = C.base_ring
        variable_names = 'x, y'
        for j in range(n):
            variable_names += ', z' + str(j)
        Rxyz = PolynomialRing(F, n+2, variable_names)
        x, y = Rxyz.gens()[:2]
        z = Rxyz.gens()[2:]
        RxyzQ = FractionField(Rxyz)
        result = heisenberg_form(C, 0)
        G = C.group
        for a in G:
            result += self.group_action(a)
        result = result.form
        Rxy.<x, y> = PolynomialRing(F, 2)
        Qxy = FractionField(Rxy)
        result = heisenberg_reduction(AS, result)
        return superelliptic_form(C_super, Qxy(result))

    def residue(self, place=0):
        return self.expansion_at_infty(place = place).residue()

    def valuation(self, place=0):
        return self.expansion_at_infty(place = place).valuation()

    def serre_duality_pairing(self, fct):
        AS = self.curve
        return sum((fct*self).residue(place = _) for _ in range(AS.nb_of_pts_at_infty))

    def cartier(self):
        C = self.curve
        F = C.base_ring
        n = C.height
        ff = C.functions
        p = F.characteristic()
        C_super = C.quotient
        (RxyzQ, Rxyz, x, y, z) = C.fct_field
        fct = self.form
        Rxy.<x, y> = PolynomialRing(F, 2)
        RxyQ = FractionField(Rxy)
        x, y = Rxyz.gens()[0], Rxyz.gens()[1]
        z = Rxyz.gens()[2:]
        num = Rxyz(fct.numerator())
        den = Rxyz(fct.denominator())
        result = RxyzQ(0)
        #return (num, den, z, fct)
        if den in Rxy:
            sub_list = {x : x, y : y} | {z[j] : (z[j]^p - RxyzQ(ff[j].function)) for j in range(n)}
            num = RxyzQ(num.substitute(sub_list))
            den1 = Rxyz(num.denominator())
            num = Rxyz(num*den1^p)
            for monomial in Rxyz(num).monomials():
                degrees = [monomial.degree(z[i]) for i in range(n)]
                product_of_z = prod(z[i]^(degrees[i]) for i in range(n))
                monomial_divided_by_z = monomial/product_of_z
                product_of_z_no_p = prod(z[i]^(degrees[i]/p) for i in range(n))
                aux_form = superelliptic_form(C_super, RxyQ(monomial_divided_by_z/den))
                aux_form = aux_form.cartier()
                result += product_of_z_no_p * Rxyz(num).monomial_coefficient(monomial) * aux_form.form/den1
            return heisenberg_form(C, result)
        raise ValueError("Please present first your form as sum z^i omega_i, where omega_i are forms on quotient curve.")
        
    def is_regular_on_U0(self):
        AS = self.curve
        C = AS.quotient
        m = C.exponent
        RxyzQ, Rxyz, x, y, z = AS.fct_field
        if y^(m-1)*self.form in Rxyz:
            return True
        return False

def are_forms_linearly_dependent(set_of_forms):
    from sage.rings.polynomial.toy_variety import is_linearly_dependent
    C = set_of_forms[0].curve
    F = C.base_ring
    n = C.height
    variable_names = 'x, y'
    for i in range(n):
        variable_names += ', z' + str(i)
    Rxyz = PolynomialRing(F, n+2, variable_names)
    denominators = prod(denominator(omega.form) for omega in set_of_forms)
    return is_linearly_dependent([Rxyz(denominators*omega.form) for omega in set_of_forms])

def holomorphic_combinations_fcts(S, pole_order):
    '''given a set S of (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([Rt(g[1]).valuation() for g in S])
    if minimal_valuation >= -pole_order:
        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 + Rt(eta_exp).valuation()
        list_coeffs = a*[0] + Rt(eta_exp).list() + (-minimal_valuation)*[0]
        list_coeffs = list_coeffs[:-minimal_valuation - pole_order]
        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_function(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

#given a set S of (form, corresponding Laurent series at some pt), find their combinations holomorphic at that pt
def holomorphic_combinations_forms(S, pole_order):
    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([Rt(g[1]).valuation() for g in S])
    if minimal_valuation >= -pole_order:
        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 + Rt(eta_exp).valuation()
        list_coeffs = a*[0] + Rt(eta_exp).list() + (-minimal_valuation)*[0]
        list_coeffs = list_coeffs[:-minimal_valuation - pole_order]
        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

#print only forms that are log at the branch pts, but not holomorphic
def only_log_forms(C_AS):
    list1 = AS.at_most_poles_forms(0)
    list2 = AS.at_most_poles_forms(1)
    result = []
    for a in list2:
        if not(are_forms_linearly_dependent(list1 + result + [a])):
            result += [a]
    return result