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_prod(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]
    WP = []
    for k in range(0, n):
        aux = witt_prod(p, k)
        Rold = aux.parent()
        Xold = Rold.gens()[:k+1]
        Yold = Rold.gens()[k+1:]
        WP+= [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*WP[k]^(p^(n-k)) for k in range(0, n)))

p = 5
n = 1
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:]
    
WS = []
WP = []
for k in range(0, n+1):
        aux1 = witt_sum(p, k)
        aux2 = witt_prod(p, k)
        Rold = aux1.parent()
        aux2 = Rold(aux2)
        Xold = Rold.gens()[:k+1]
        Yold = Rold.gens()[k+1:]
        WS+= [RQ(aux1.subs({Xold[i] : X[i] for i in range(0, k)} | {Yold[i] : Y[i] for i in range(0, k)}))]
        WP+= [RQ(aux2.subs({Xold[i] : X[i] for i in range(0, k)} | {Yold[i] : Y[i] for i in range(0, k)}))]

Rp2n = PolynomialRing(GF(p), variables, 2*(n+1))
X = Rp2n.gens()[:n+1]
Y = Rp2n.gens()[n+1:]
WS = [Rp2n(a) for a in WS]
WP = [Rp2n(a) for a in WP]