# conicsPolar.py
# Andrew Davison, ad@coe.psu.ac.th, Oct. 2025
'''
Plots conic sections (ellipse, parabola, hyperbola) 
using Matplotlib polar coordinates
with the polar equation r = l / (1 + e cos(theta))
'''

import math
import matplotlib.pyplot as plt


def conicPts(e, a, step=0.01, thetaRange=None):
  # Returns theta, r points for a conic
  if e < 1:     # Ellipse
    l = a * (1 - e**2)
  elif abs(e - 1) < 1e-6:  # parabola; v.close to 1
    l = a
  else:         # Hyperbola
    l = a * (e**2 - 1)
  
  if thetaRange is None:
    thetaStart = -math.pi + 0.0001
    thetaEnd = math.pi - 0.0001
  else:
    thetaStart, thetaEnd = thetaRange
  
  thetas = []
  rs = []
  theta = thetaStart
  while theta <= thetaEnd:
    denom = 1 + e * math.cos(theta)
    if abs(denom) > 1e-6:  # not close to 0
      r = l / denom
      thetas.append(theta)
      rs.append(r)
    theta += step
  return thetas, rs


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

fig = plt.figure()
ax = fig.add_subplot(111, projection='polar')

if choice == 'e':
  e = 0.6
  a = 5
  thetas, rs = conicPts(e, a)
  label = "Ellipse"

elif choice == 'p':
  e = 1.0
  a = 2
  thetaRange = (-math.pi/2 + 0.1, math.pi/2 - 0.1)
  thetas, rs = conicPts(e, a, thetaRange=thetaRange)
  label = "Parabola"

elif choice == 'h':
  e = 1.5
  a = 5
  thetaLimit = math.acos(-1/e) - 0.1
  thetaRange = (-thetaLimit, thetaLimit)
  thetas, rs = conicPts(e, a, thetaRange=thetaRange)
  label = "Hyperbola"

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

ax.plot(thetas, rs, linewidth=2)
ax.set_title(f"{label} (polar): r = l / (1 + e cosθ), e = {e}", 
                          va='bottom')
ax.plot([0], [0], 'ro', markersize=8, label='Focus')
ax.grid(True)
ax.legend(loc='upper right')
plt.show()