# ellipseRays.py
# Andrew Davison, ad@coe.psu.ac.th, Oct. 2025

'''
Generate several vectors leaving the left focus
of an ellipse. Calculate the reflected vectors, 
and see them all pass through the other focus of
the ellipse.
'''

import random
import math
import matplotlib.pyplot as plt
from Vec import Vec   


def ellipsePts(n):
  # use parametric form: 
  # (x,y) = (a cos theta, b sin theta)
  pts = []
  for i in range(n+1):
    theta = 2*math.pi * i/n
    x = a * math.cos(theta)
    y = b * math.sin(theta)
    pts.append(Vec(x, y, 0))
  return pts



def intersectPt(origin, dir):
  x0, y0 = origin.x, origin.y
  dx, dy = dir.x, dir.y

  A = (dx*dx)/(a*a) + (dy*dy)/(b*b)
  B = 2*x0*dx/(a*a) + 2*y0*dy/(b*b)
  C = (x0*x0)/(a*a) + (y0*y0)/(b*b) - 1

  # get positive intersection distance
  t = (-B + math.sqrt(B*B - 4*A*C)) / (2*A)
  return origin + dir*t


def reflectRay(hitPt, dir):
  """
  Reflect vector off the ellipse.
  If the ellipse is
     F(x, y) = (x/a)^2 + (y/b)^2 - 1
  then the normal is
     n = (x/a^2, y/b^2)
  """
  nx = hitPt.x /(a*a)
  ny = hitPt.y /(b*b)
  n = Vec(nx, ny, 0).normalize()
  dot = dir.dot(n)
  return dir - n * (2 * dot)


# parameters
a = 5  # semi-major axis length
b = 3  # semi-minor

# foci
c = math.sqrt(a*a - b*b)  
   # distance from center to each focus
focus1 = Vec(-c, 0, 0)
focus2 = Vec(c, 0, 0)


fig, ax = plt.subplots(figsize=(7, 6))

# Draw ellipse and foci
ellipse = ellipsePts(200)
ax.plot([p.x for p in ellipse], 
        [p.y for p in ellipse], 'k')
ax.plot(focus1.x, focus1.y, 'ro')
ax.plot(focus2.x, focus2.y, 'ro')

# Emit 10 random rays from focus 1
for _ in range(10):
  angle = random.uniform(0, 2 * math.pi)
  dir = Vec(math.cos(angle), math.sin(angle), 0)

  # find where ray hits the ellipse
  hitPt = intersectPt(focus1, dir)

  # plot green ray from focus 1
  ax.plot([focus1.x, hitPt.x], 
          [focus1.y, hitPt.y], 'g--')

  # reflected blue ray
  vReflected = reflectRay(hitPt, dir)
  endPt = hitPt + vReflected*10  # arbitrary length
  ax.plot([hitPt.x, endPt.x], 
          [hitPt.y, endPt.y], 'b-')

ax.set_aspect("equal", adjustable="box")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.grid(True, lw=0.6, alpha=0.6)
ax.set_title("Elliptic Reflection")
plt.show()