import numpy as np
import matplotlib.pyplot as plt


### computations
##############################################################

# arguments
nbSteps = 80  # for the parabola
alpha = -np.pi/2 + .1  # angle 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
i_angle = -alpha + np.pi/2 # incidence angle measured with respect to Ox
X = np.linspace(-2, 2, N) # intersections with the parabola
Y = X**2

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

i_xx = np.zeros([2, N])  # matrix of incident vectors' x (in columns) 
i_xx[0, :] = X - np.cos(i_angle)
i_xx[1, :] = X + np.cos(i_angle)

i_yy = np.zeros([2, N])  # matrix of incident vectors' x (in columns) 
i_yy[0, :] = Y - np.sin(i_angle)
i_yy[1, :] = Y + np.sin(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)

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)


### 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 angle in degrees is {round(alpha, 2)} " + \
    f"and the number of light rays is {N}"
fig.suptitle(title)

ax.plot(x_parabola, y_parabola)
# ax.plot(i_xx, i_yy, c = 'g', alpha=.5, linewidth=.8)
ax.plot(r_xx, r_yy, c = 'r', alpha=.3, linewidth=.8)

plt.show()