# showRoots.py
# Andrew Davison, ad@coe.psu.ac.th, Nov. 2025

# inputs: 5-1j 3;  5+1j 3
'''
Visualize the n-th roots of a complex number using matplotlib.
Shows how roots form regular polygons in the complex plane.
'''

import matplotlib.pyplot as plt
import math, cmath
import complexUtils


# -----------------------------
re = float(input("Real part: "))
im = float(input("Imaginary part: "))
n = int(input("Root number (n): "))
z = complex(re, im)

rs = complexUtils.getRoots(z, n)
print(f"\nThe {n} roots:")
for root in rs:
  print(f"  {root: .4f}")


fig, ax = plt.subplots()

# Plot the roots
reals = [root.real for root in rs]
imags = [root.imag for root in rs]
ax.scatter(reals, imags, s=100, c='red', alpha=0.6,
           marker='o', edgecolors='darkred', linewidths=2, 
           label=f'{n}-th roots', zorder=3)

# Connect roots to form a regular polygon
reals.append(reals[0])
imags.append(imags[0])
ax.plot(reals, imags, 'b--', alpha=0.4, lw=2)

# Draw circle for reference
theta = [2*math.pi * t/100 for t in range(101)]
rootMag = abs(z) ** (1/n)
circleX = [rootMag * math.cos(t) for t in theta]
circleY = [rootMag * math.sin(t) for t in theta]
ax.plot(circleX, circleY, 'g--', alpha=0.2, lw=1.5)

# label each root with its value
for root in rs:
  ax.annotate(f'({root.real:.2f}{root.imag:+.2f}j)', 
             (root.real, root.imag), 
             xytext=(7, 7), textcoords='offset points',
             fontsize=9, fontweight='bold')


ax.axhline(y=0, color='k', lw=0.5, alpha=0.5)
ax.axvline(x=0, color='k', lw=0.5, alpha=0.5)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
ax.set_title(f'{n} Roots of {z:.2f}')
ax.set_xlabel('Real')
ax.set_ylabel('Imaginary')

plt.show()