# the code for the degree 5 polynomial takes FOREVER to run,
# so I'm not sure if it works. This is definitely correct for
# the degree 3 polynomial here, and terminates (with the incorrect answer :/)
# for the degree 4 polynomial. I'm not sure if that's because of a bug
# in my code or if it's because 4 is not prime...

f = x^5 + x^4 - 4*x^3 - 3*x^2 + 3*x + 1
# f = x^4 -4*x^2 + 2
# f = x^3 + x^2 - 2*x - 1

###################################################

# the coefficients of f.
# we reverse them to make upcoming code cleaner
poly = f.coefficients(sparse=False)[::-1]

n = f.degree(x)

# make a number field where w is a nth root of 1
K.<w> = CyclotomicField(n)

# make vars for the roots.
# a[i] = alpha_i
R = PolynomialRing(K, n, 'a')
a = R.gens()

#v_1^n
v = sum([w^k * a[k] for k in range(n)])^n

# build the conjugates of v.
# it turns out we only need one from each conjugacy class,
# which keeps the degree down.

Sn = SymmetricGroup(n)

print("start conjugates")
conjugates = []
for tau in Sn:
  if v * tau not in conjugates:
    conjugates += [v * tau]
print("finish conjugates")

# It turns out to be WAY less efficient to have iterated polynomial rings
S.<Y> = PolynomialRing(R)
Sf = S.flattening_morphism().codomain()

print("start computing psi")
psi = prod([Sf(Y - conj) for conj in conjugates])
print("finish computing psi")

# the coefficients of psi are symmetric functions in the ai.
# so let's write them as polynomials in the elementary
# symmetric functions.

Sym = SymmetricFunctions(K)
e = Sym.e()

# now, since each elementary symmetric function is one of the coefficients
# of f (up to sign), we can convert the coefficients of psi into numbers!

def coeffToNumber(coeff):
  """
  Convert a sum of elementary terms into a number
  """

  def elemToNumber(elem):
    """
    Convert an elementary term into a number

    the pair (ns = [n1,n2,...,nk], c) corresponds to c * e_{n1} * e_{n2} ... * e_{nk}
    """
    ns, c = elem

    out = c
    for n in ns:
      if n < len(poly):
        # the elementary polys alternate in sign
        # when you view them as the coefficients of
        # a polynomial from the roots.
        out *= (-1)^(n) * poly[n]
      else:
        out *= 0

    return out


  return sum(elemToNumber(elem) for elem in coeff)


print("start computing the coefficients of psi in Q")
psiReduced = 0
for j in range(n+1):
  print(j)

  print("start converting to R")
  coeff = R(psi.coefficient({Y:j})) # convert back to a ring without Y
  print("finish converting to R")

  print("start converting to sym")
  coeff = e(Sym.from_polynomial(coeff)) # convert into a symmetric polynomial
  print("finish converting to sym")

  print("add to psi")
  psiReduced += Y^j * coeffToNumber(list(coeff)) # convert into a member of K
  print("finish adding to psi")

# v1,v2,...v(n-1) are exactly the roots of psi!
# notice these are already in our underlying field,
# so finding the roots isn't actually cheating.
# We could do this by hand by factoring psi, which
# must (by theory) factor into linear terms.

v0 = - poly[1] # v0 = a0 + a1 + ... + a(n-1) = e1
vs = psiReduced.roots(multiplicities=False)

# Finally, we're done!
r = (v0 + sum(v^(1/n) for v in vs))/n

print(f)
print(r)

# and we can verify
print(f(r.n()))