# paraRays.py
# Andrew Davison, ad@coe.psu.ac.th, Oct. 2025
'''
Generate several vectors heading downwards to
hit an upward facing parabola. Calculate the
reflected vectors, and see them all pass through the
parabola's focus.
'''

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


def parabolaY(x, f):
  # Return y on the parabola y = x^2 / (4 f).
  return (x * x) / (4.0 * f)


def reflectVector(v, n):
  # Reflect vector v across unit normal n.
  # v' = v - 2 (v · n) n
  dot = v.dot(n)
  return v - n * (2 * dot)


def extendRay(p0, v, length=5.0):
  # Extend a ray from p0 in direction v by given length.
  dirNorm = v.normalize()
  return p0 + dirNorm * length


# Parameters
f = 2.0
nRays = 6
xMin, xMax = -6.0, 6.0   # ranges for parabola
yMin, yMax = -1.0, 10.0


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

# plot parabola
xs = [xMin + (xMax-xMin) * i/400.0 
                            for i in range(401)]
ys = [parabolaY(x, f) for x in xs]

ax.plot(xs, ys, lw=2, color="black", label="Parabola")
ax.scatter([0], [f], color="black", s=40, zorder=5, 
                                     label="Focus")

downVec = Vec(0, -1, 0)

# random x values for vectors
xVals = [random.uniform(xMin * 0.9, xMax * 0.9) 
                              for _ in range(nRays)]
xVals.sort()

for i, x in enumerate(xVals, start=1):
  yHit = parabolaY(x, f) 
  pHit = Vec(x, yHit, 0)  # contact with parabola

  # Incoming blue dashed ray (vertical down)
  ax.plot([x, x], [yMax, yHit], ls="--", 
                             lw=1.5, color="blue")
  # Place number at top edge
  ax.text(x, yMax + 0.3, f"{i}", color="blue", 
               ha="center", va="bottom", fontsize=10)

  # convert parabola slope at x into a normal
  slope = x / (2.0 * f)
  normal = Vec(-slope, 1, 0).normalize()

  # Reflected red ray
  vReflected = reflectVector(downVec, normal)
  ray = extendRay(pHit, vReflected, length=8.0)
  ax.plot([pHit.x, ray.x], [pHit.y, ray.y], 
                                  color="red", lw=1.5)
  # Place red number at end
  ax.text(ray.x, ray.y, f"{i}", color="red", 
                ha="center", va="bottom", fontsize=10)

ax.set_xlim(xMin, xMax)
ax.set_ylim(yMin, yMax + 1.0)
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("Parabolic Reflection")
plt.show()