class template:
    '''Template of a p-group cover'''
    def __init__(self, height, field, group, fcts, gp_action):
        self.height = height
        self.group = group
        self.gp_action = gp_action #action of the generators of the group on z[i]'s
        self.field = field
        n = height
        variable_names = ''
        for i in range(n):
            variable_names += 'z'+str(i)+','
        for i in range(n):
            variable_names += 'f'+str(i)+','
        variable_names += 'x, y'
        Rzf = PolynomialRing(field, 2*n+2, variable_names)
        z = Rzf.gens()[:n]
        f = Rzf.gens()[n:]
        x = Rzf.gens()[-2]
        y = Rzf.gens()[-1]
        RzfQ = FractionField(Rzf) #nowa linijka
        self.fct_field = RzfQ, z, f, x, y #Rzf zmienione na RzfQ
        self.fcts = [RzfQ(ff) for ff in fcts] #RHSs of the Artin-Schreier equations

def elementary_template(p, n):
    group = elementary_gp(p, n)
    field = GF(p)
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i) + ','
    variable_names += 'x, y'        
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    height = n
    fcts = [f[i] for i in range(n)]
    gp_action = [[z[j] + (i == j) for j in range(n)]+[x, y] for i in range(n)]
    return template(height, field, group, fcts, gp_action)

def elementary_cover(list_of_fcts, prec=10):
    n = len(list_of_fcts)
    C = list_of_fcts[0].curve
    return as_cover(C, elementary_template(p, n), list_of_fcts, branch_points = [], prec = prec)

def heisenberg_template(p):
    group = heisenberg(p)
    field = GF(p)
    variable_names = ''
    n = 3
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i)+','
    variable_names += 'x, y'
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    height = n
    fcts = [f[i] for i in range(n)]
    fcts[2] += (z[0] - z[1])*f[1]
    gp_action = [[z[0] + 1, z[1], z[2] + z[1], x, y], [z[0] + 1, z[1] + 1, z[2], x, y], [z[0], z[1], z[2] - 1, x, y]]
    return template(height, field, group, fcts, gp_action)

def heisenberg_cover(list_of_fcts, prec=10):
    n = len(list_of_fcts)
    C = list_of_fcts[0].curve
    return as_cover(C, heisenberg_template(p), list_of_fcts, branch_points = [], prec = prec)

def witt_pol(X, p, n):
    n = len(X)
    return sum(p^i*X[i]^(p^(n-i-1)) for i in range(0, n))

def witt_sum(p, n):
    variables = ''
    for i in range(0, n+1):
        variables += 'X' + str(i) + ','
    for i in range(0, n+1):
        variables += 'Y' + str(i)
        if i!=n:
            variables += ','
    RQ = PolynomialRing(QQ, variables, 2*(n+1))
    X = RQ.gens()[:n+1]
    Y = RQ.gens()[n+1:]
    Rpx.<x> = PolynomialRing(GF(p), 1)
    RQx.<x> = PolynomialRing(QQ, 1)
    if n == 0:
        return X[0] + Y[0]
    WS = []
    for k in range(0, n):
        aux = witt_sum(p, k)
        Rold = aux.parent()
        Xold = Rold.gens()[:k+1]
        Yold = Rold.gens()[k+1:]
        WS+= [aux.subs({Xold[i] : X[i] for i in range(0, k)} | {Yold[i] : Y[i] for i in range(0, k)})]
    return 1/p^n*(witt_pol(X[:n+1], p, n) + witt_pol(Y[:n+1], p, n) - sum(p^k*WS[k]^(p^(n-k)) for k in range(0, n)))

def witt_sum_mod_p(p, n):
    variables = ''
    for i in range(0, n+1):
        variables += 'X' + str(i) + ','
    for i in range(0, n+1):
        variables += 'Y' + str(i)
        if i!=n:
            variables += ','
    RQ = PolynomialRing(QQ, variables, 2*(n+1))
    X = RQ.gens()[:n+1]
    Y = RQ.gens()[n+1:]
    P = RQ(witt_sum(p, n))
    RQp = PolynomialRing(GF(p), variables, 2*(n+1))
    Xp = RQp.gens()[:n+1]
    Yp = RQp.gens()[n+1:]
    return RQp(P)

def witt_template(p, n):
    height = n
    field = GF(p)
    group = cyclic_gp(p, n)
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i)+','
    variable_names += 'x, y'
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    ###########
    rhs = []
    gp_action = []
    for i in range(0, n):
        aux = witt_sum_mod_p(p, i)
        Raux = aux.parent()
        Xpn = Raux.gens()[:i+1]
        Ypn = Raux.gens()[i+1:]
        rhs_aux = aux.subs({Xpn[ii] : z[ii]^p for ii in range(i+1)}|{Ypn[ii] : -z[ii] for ii in range(i+1)})
        rhs += [rhs_aux]
        gp_action_aux = aux.subs({Xpn[ii] : z[ii] for ii in range(i+1)}|{Ypn[ii] : ii == 0 for ii in range(i+1)})
        gp_action += [gp_action_aux]
    fcts = [-rhs[i] + z[i]^p - z[i] + f[i] for i in range(n)]
    ########
    aux
    gp_action = [gp_action+[x, y]]
    return template(height, field, group, fcts, gp_action)

def witt_cover(list_of_fcts, prec=10):
    n = len(list_of_fcts)
    C = list_of_fcts[0].curve
    return as_cover(C, witt_template(p, n), list_of_fcts, branch_points = [], prec = prec)

def quaternion_template():
    field = GF(2)
    height = 3
    n = 3
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i)+','
    variable_names += 'x, y'
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    group = quaternion_gp()
    fcts = [f[0], f[1], f[2] + z[0]*f[0]+z[1]*(f[0] + f[1])]
    gp_action = [[z[0]+1, z[1], z[2] + z[0], x, y], [z[0], z[1] + 1, z[2] + z[1] + z[0], x, y]]
    return template(height, field, group, fcts, gp_action)

def quaternion_cover(list_of_fcts, prec=10):
    n = len(list_of_fcts)
    C = list_of_fcts[0].curve
    return as_cover(C, quaternion_template(), list_of_fcts, branch_points = [], prec = prec)

def hypoelementary_template(p, m, b, zeta):
    '''unfinished'''
    field = GF(p)
    height = 1
    n = 1
    variable_names = ''
    for i in range(n):
        variable_names += 'z'+str(i)+','
    for i in range(n):
        variable_names += 'f'+str(i)+','
    variable_names += 'x, y'
    R = PolynomialRing(field, 2*n+2, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    x = R.gens()[-2]
    y = R.gens()[-1]
    group = hypoelementary(p, m, b)
    fcts = [1/(zeta - b)*f[0]^p*z[0]^p - 1/(zeta - b)*f[0]*z[0]]
    gp_action = []
    gp_action += [b*z[0]+f[0]*y, x, zeta*y]
    gp_action += [z[0]+1, x, y]
    return template(height, field, group, fcts, gp_action)