import numpy as np
import matplotlib.pyplot as plt


### computations
##############################################################
"""
The angle alpha will vary and using this variation we shall define an
animation of the reflected rays, hence of the caustique.
"""
# arguments
nbSteps = 80  # for the parabola
nbImag = 200
alpha = np.linspace(-np.pi/2, np.pi/2, nbImag)  # angles measured with respect to Oy
N = 90  # incident light rays
d = 5  # half length of the reflected ray


# parabola
x_parabola = np.linspace(-2, 2, nbSteps)
y_parabola = x_parabola**2


# light rays
rX_list = []
rY_list = []
X = np.linspace(-2, 2, N) # intersections with the parabola
Y = X**2

for k in range(nbImag):
    i_angle = -alpha[k] + np.pi/2 # incidence angle measured with respect to Ox

    t_angles = np.arctan2(2*X, 1)
    n_angles = t_angles + np.pi/2
    r_angles = 2*n_angles - i_angle

    r_xx = np.zeros([2, N])  # matrix of reflent vectors' x (in columns) 
    r_xx[0, :] = X - d*np.cos(r_angles)
    r_xx[1, :] = X + d*np.cos(r_angles)
    rX_list.append(r_xx)

    r_yy = np.zeros([2, N])  # matrix of reflent vectors' x (in columns) 
    r_yy[0, :] = Y - d*np.sin(r_angles)
    r_yy[1, :] = Y + d*np.sin(r_angles)
    rY_list.append(r_yy)

### drawings
##############################################################

# figure and axes
fig = plt.figure(figsize=(7, 8))
ax = fig.add_subplot(111, aspect='equal')

ax.set_xlim(-2.2, 2.2)
ax.set_ylim(-.5, 6)
title = f"The number of light rays is {N}"
fig.suptitle(title)

ax.plot(x_parabola, y_parabola)

# initialization of the animation
rLines_list = []
for j in range(N):
    x = rX_list[0][:, j]
    y = rY_list[0][:, j]
    rLine = ax.plot(x, y, c='g', alpha=.3, linewidth=.8)[0]
    rLines_list.append(rLine)

# animation
for k in range(1, nbImag):
    for j in range(N):
        rLine = rLines_list[j]
        rLine.set_xdata(rX_list[k][:, j])
        rLine.set_ydata(rY_list[k][:, j])
    plt.pause(0.01)
    
plt.show()