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)
            if i!=n-1:
                variable_names += ','
        Rzf = PolynomialRing(field, 2*n, variable_names)
        z = Rzf.gens()[:n]
        f = Rzf.gens()[n:]
        self.fct_field = Rzf, z, f
        self.fcts = [Rzf(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)
        if i!=n-1:
            variable_names += ','
    R = PolynomialRing(field, 2*n, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    height = n
    fcts = [f[i] for i in range(n)]
    gp_action = [[z[j] + (i == j) for j in range(n)] for i in range(n)]
    return template(height, field, group, fcts, gp_action)

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)
        if i!=n-1:
            variable_names += ','
    R = PolynomialRing(field, 2*n, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    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]], [z[0] + 1, z[1] + 1, z[2]], [z[0], z[1], z[2] - 1]]
    return template(height, field, group, fcts, gp_action)


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)
        if i!=n-1:
            variable_names += ','
    R = PolynomialRing(field, 2*n, variable_names)
    z = R.gens()[:n]
    f = R.gens()[n:]
    ###########
    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]
    return template(height, field, group, fcts, gp_action)