# conicsCarts.py
# Andrew Davison, ad@coe.psu.ac.th, Oct. 2025
'''
Allow the user to draw a parabola, ellipse, or hyperbola. The curves are defined using (x,y) Cartesian coordinates unlike the polar versions used in ellipseString.py and hyperbolaString.py
'''

import matplotlib.pyplot as plt


def ellipsePoints(a, b):
  # Returns points for ellipse x^2/a^2 + y^2/b^2 = 1
  xs = []
  step = 0.01
  x = -a
  while x <= a:
    xs.append(x)
    x += step
  ys = [b * (1 - (x**2/a**2))**0.5 for x in xs]
  return xs, ys


def parabolaPoints(a, yRange):
  # Returns points for parabola y^2 = 4ax
  ys = []
  step = 0.1
  y = -yRange
  while y <= yRange:
    ys.append(y)
    y += step
  xs = [(y**2)/(4*a) for y in ys]
  return xs, ys


def hyperbolaPoints(a, b, xMax):
  # Returns points for hyperbola x^2/a^2 - y^2/b^2 = 1
  xs = []
  step = 0.05
  x = a
  while x <= xMax:
    xs.append(x)
    x += step
  ys = [b * ((x**2/a**2) - 1)**0.5 for x in xs]
  return xs, ys


# --- ---

choice = input("Enter e (ellipse), p (parabola), or h (hyperbola): ").strip().lower()

plt.figure()

if choice == 'e':
  a, b = 5, 3
  xs, ys = ellipsePoints(a, b)
  plt.plot(xs, ys, 'b')
  plt.plot(xs, [-y for y in ys], 'b')
  plt.title("Ellipse: x^2/a^2 + y^2/b^2 = 1")
  
  # Foci and directrices
  c = (a**2 - b**2)**0.5
  plt.scatter([c, -c], [0, 0], color='red', 
                              label='Foci')
  plt.axvline(x=a**2/c, color='gray', ls='--', 
                              label='Directrix')
  plt.axvline(x=-a**2/c, color='gray', ls='--')

elif choice == 'p':
  a = 2
  yRange = 10
  xs, ys = parabolaPoints(a, yRange)
  plt.plot(xs, ys, 'g')
  plt.title("Parabola: y^2 = 4ax")
  
  # Focus and directrix
  plt.scatter([a], [0], color='red', label='Focus')
  plt.axvline(x=-a, color='gray', ls='--', 
                                 label='Directrix')

elif choice == 'h':
  a, b, xMax = 4, 2, 10
  xs, ys = hyperbolaPoints(a, b, xMax)
  plt.plot(xs, ys, 'r')
  plt.plot(xs, [-y for y in ys], 'r')

  # draw both branches
  plt.plot([-x for x in xs], ys, 'r')
  plt.plot([-x for x in xs], [-y for y in ys], 'r')
  plt.title("Hyperbola: x^2/a^2 - y^2/b^2 = 1")

  # Foci and asymptotes
  c = (a**2 + b**2)**0.5
  plt.scatter([c, -c], [0, 0], color='red', 
                                  label='Foci')
  xRange = [-xMax, xMax]
  yAsym1 = [b/a * x for x in xRange]
  yAsym2 = [-b/a * x for x in xRange]
  plt.plot(xRange, yAsym1, 'gray', ls='--', 
                            label='Asymptotes')
  plt.plot(xRange, yAsym2, 'gray', ls='--')

else:
  print("Invalid choice")
  quit()

plt.xlabel('x')
plt.ylabel('y')
plt.axis('equal')
plt.grid(True)
plt.legend()
plt.show()