import numpy as np
import matplotlib.pyplot as plt
plt.style.use('ff.mplstyle')

x = np.array([0.2, 0.4, 0.6, 0.8, 1.0])       # м
v2_1 = np.array([9, 18, 26, 29, 42]) * 1e-3   # м^2/с^2
v2_2 = np.array([15, 34, 47, 65, 85]) * 1e-3  # м^2/с^2
v2_3 = np.array([20, 49, 75, 95, 127]) * 1e-3 # м^2/с^2
plt.scatter(x, v2_1, label=r'$\alpha = 0.11^\circ$')
plt.scatter(x, v2_2, label=r'$\alpha = 0.23^\circ$')
plt.scatter(x, v2_3, label=r'$\alpha = 0.37^\circ$')

xx = np.linspace(0, 1.1, 100)
for y in [v2_1, v2_2, v2_3]:
    s0 = 1.5e-7
    a = np.sum(x * y) / np.sum(x**2)
    Sa = np.sqrt(s0**2 / (np.sum(x**2))**2)
    print(a, Sa)
    
    ym = np.mean(y)
    SSres = np.sum((y - a*x)**2)
    SStot = np.sum((y - ym)**2)
    R2 = 1 - SSres / SStot
    print(R2)
    
    plt.plot(xx, a*xx, label=fr'{a =: .2e} $\pm$ {Sa:.2e};  $R^2$ = {R2:.3f}')

plt.xlim(xmin=0)
plt.ylim(ymin=1e-6)
plt.xlabel('$x$, м')
plt.ylabel(r'$v^2, \frac{м^2}{c^2}$')
plt.xticks([0, 0.25, 0.5, 0.75, 1])
plt.title('Зависимость квадрата скорости тележки $v^2$ \n от пройденного ей расстояния $x$. Случай тележки без груза')
plt.legend(title='Наклон скамьи')
plt.savefig('example_5_proportionality.pdf')