import numpy as np
import matplotlib.pyplot as plt


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)


nbSteps = 80
x_parabola = np.linspace(-2, 2, nbSteps)
ax.plot(x_parabola, x_parabola**2)


alpha = -np.pi/2 + .1  # angle measured wuth respect to Oy
a = -alpha + np.pi/2
print(f"angle of the incident vector = {a*180/np.pi}")
u = np.array([np.cos(a), np.sin(a)])  # incident vector -- not used


N = 11
X = np.linspace(-2, 2, N)
Y = X**2
title = f"The angle in degrees is {round(alpha, 2)} " + \
    f"and the number of light rays is {N}"
fig.suptitle(title)


for j in range(N):
    print('\nj =', j)
    x = X[j]
    y = Y[j]
    t = np.arctan2(2*x, 1)
    n = np.pi/2 + t  # angle of the normal vector
    print(f"angle of the normal vector = {n*180/np.pi}")
    r = 2*n - a
    print(f"angle of the reflected vector = {r*180/np.pi}")
    refl = np.array([np.cos(r), np.sin(r)])  # reflection vector -- not used
    
    i_xx = [x - np.cos(a), x + np.cos(a)]
    i_yy = [y - np.sin(a), y + np.sin(a)]
    r_xx = [x - 5*np.cos(r), x + 5*np.cos(r)]
    r_yy = [y - 5*np.sin(r), y + 5*np.sin(r)]
    
    ax.plot(i_xx, i_yy, c = 'g', alpha=.2, linewidth=.8)
    ax.plot(r_xx, r_yy, c = 'r', alpha=.4, linewidth=.8)
 
plt.show()