class heisenberg_function:
    def __init__(self, C, g):
        self.curve = C
        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)
        x, y = Rxyz.gens()[:2]
        z = Rxyz.gens()[2:]
        RxyzQ = FractionField(Rxyz)
        self.function = RxyzQ(g)
        #self.function = heisenberg_reduction(AS, RxyzQ(g))
        
    def __repr__(self):
        return str(self.function)
    
    def __eq__(self, other):
        AS = self.curve
        RxyzQ, Rxyz, x, y, z = AS.fct_field
        aux = self - other
        aux = RxyzQ(aux.function)
        aux = aux.numerator()
        aux = heisenberg_function(AS, aux)
        aux = aux.expansion_at_infty()
        if aux.valuation() >= 1:
            return True
        return False

    def __add__(self, other):
        C = self.curve
        g1 = self.function
        g2 = other.function
        return heisenberg_function(C, g1 + g2)

    def __sub__(self, other):
        C = self.curve
        g1 = self.function
        g2 = other.function
        return heisenberg_function(C, g1 - g2)

    def __rmul__(self, constant):
        C = self.curve
        g = self.function
        return heisenberg_function(C, constant*g)
    
    def __neg__(self):
        C = self.curve
        g = self.function
        return heisenberg_function(C, -g)
    
    def __mul__(self, other):
        if isinstance(other, heisenberg_function):
            C = self.curve
            g1 = self.function
            g2 = other.function
            return heisenberg_function(C, g1*g2)
        if isinstance(other, heisenberg_form):
            C = self.curve
            g1 = self.function
            g2 = other.form
            return heisenberg_form(C, g1*g2)
    
    def __truediv__(self, other):
        C = self.curve
        g1 = self.function
        g2 = other.function
        return heisenberg_function(C, g1/g2)
    
    def __pow__(self, exponent):
        C = self.curve
        g1 = self.function
        return heisenberg_function(C, g1^(exponent))
    
    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]
        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.function
        g = RxyzQ(g)
        sub_list = {x : x_series, y : y_series} | {z[j] : z_series[j] for j in range(n)}
        return g.substitute(sub_list)
    
    def expansion(self, pt = 0):
        C = self.curve
        delta = C.nb_of_pts_at_infty
        F = C.base_ring
        x_series = C.x_series[pt]
        y_series = C.y_series[pt]
        z_series = C.z_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.function
        g = RxyzQ(g)
        sub_list = {x : x_series, y : y_series} | {z[j] : z_series[j] for j in range(n)}
        return g.substitute(sub_list)

    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.function
            return heisenberg_function(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.function
            return heisenberg_function(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.function
            return heisenberg_function(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 trace(self):
        print('trace ')
        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_function(C, 0)
        G = C.group
        for a in G:
            result += self.group_action(a)
        result = result.function
        Rxy.<x, y> = PolynomialRing(F, 2)
        Qxy = FractionField(Rxy)
        result = heisenberg_reduction(C, result)
        #print(result)
        return superelliptic_function(C_super, Qxy(result))
    
    def coordinates(self, prec = 100, basis = 0):
        "Return coordinates in H^1(X, OX)."
        AS = self.curve
        if basis == 0:
            basis = [AS.holomorphic_differentials_basis(), AS.cohomology_of_structure_sheaf_basis()]
        holo_diffs = basis[0]
        coh_basis =  basis[1]
        f_products = []
        for f in coh_basis:
            f_products += [[omega.serre_duality_pairing(f) for omega in holo_diffs]]
        product_of_fct_and_omegas = []
        product_of_fct_and_omegas = [omega.serre_duality_pairing(self) for omega in holo_diffs]

        V = (F^(AS.genus())).span_of_basis([vector(a) for a in f_products])
        coh_coordinates = V.coordinates(product_of_fct_and_omegas)
        return coh_coordinates
    
    def diffn(self):
        C = self.curve
        C_super = C.quotient
        n = C.height
        RxyzQ, Rxyz, x, y, z = C.fct_field
        fcts = C.functions
        f = self.function
        y_super = superelliptic_function(C_super, y)
        dy_super = y_super.diffn().form
        dz = []
        for i in range(n):
            dfct = fcts[i].diffn().form
            dz += [-dfct]
        result = f.derivative(x)
        result += f.derivative(y)*dy_super
        for i in range(n):
            result += f.derivative(z[i])*dz[i]
        return heisenberg_form(C, result)
        
    def valuation(self, place = 0):
        '''Return valuation at i-th place at infinity.'''
        C = self.curve
        F = C.base_ring
        Rt.<t> = LaurentSeriesRing(F)
        return Rt(self.expansion_at_infty(place = place)).valuation()