import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as spo


# Données expérimentales
# angle d'équlibre
theta = np.array([20, 30, 39, 48, 55, 67, 69, 78])*np.pi/180
u_theta = np.array([1 for e in theta])*np.pi/180

# vitesse de rotation en rad/s
omega = np.array([152.5, 159.5, 169.4, 183, 196.4, 226.4, 239, 322.6])*2*np.pi/60
u_omega = np.array([3 for e in omega])*2*np.pi/60

x = 1/omega**2
ux = 2*x*u_omega/omega
y = np.cos(theta)
uy = np.abs(np.sin(theta)*u_theta)

# Point en direct
# angle d'équlibre
thetal = np.array([])*np.pi/180
u_thetal = np.array([1 for e in thetal])*np.pi/180

# vitesse de rotation
omegal = np.array([])*2*np.pi/60
u_omegal = np.array([3 for e in omegal])*2*np.pi/60

xl= 1/omegal**2
uxl = 2*xl*u_omegal/omegal
yl = np.cos(thetal)
uyl = np.abs(np.sin(thetal)*u_thetal)

# données complètes : 1 inclure le point en live, 2 sinon
live = 1
if live == 1 :
    xdata = np.concatenate((x,xl))
    ydata = np.concatenate((y,yl))
    u_xdata = np.concatenate((ux,uxl))
    u_ydata = np.concatenate((uy,uyl))
else :
    xdata = x
    ydata = y
    u_xdata = ux
    u_ydata = uy

# Ajustement
# fonction f décrivant la courbe à ajuster aux données
def f(x,a,b):
    return a*x+b

# estimation initiale des paramètres
params0 = np.array([0,0])

# routine Python réalisant l'ajustement
result = spo.curve_fit(f, xdata, ydata)

# on obtient :
# les paramètres d'ajustement optimaux
a,b = result[0];
# la matrice de variance-covariance estimée des paramètres
pcov = result[1];
# les incertitudes-types sur ces paramètres
ua,ub = np.sqrt(np.abs(np.diagonal(pcov)))


# Représentation graphique

fig = plt.figure()
ax = fig.add_subplot(111)
ax.tick_params(labelsize=18)

plt.grid()

lx=np.linspace(min(xdata),max(xdata),1000)
plt.plot(lx,f(lx,a,b),linewidth=2,color='green',
label='ajustement $f(x) = ax+b$ \n $a$ = {:.3e} $\pm$ {:.3e} \n $b$ = {:.3e} $\pm$ {:.3e}'.format(a, ua, b, ub))

plt.errorbar(x, y, xerr=ux, yerr=uy, fmt='o', ms=5, capsize=2, capthick=2, linewidth=2)

#point en live
if live == 1 :
    plt.errorbar(xl, yl, xerr=uxl, yerr=uyl,fmt='o', ms=5, capsize=2, capthick=2, linewidth=2, color='red')

# légendes des axes
plt.xlabel(r"1/$\Omega ^2$ (rad$^2$.s$^{-2}$)", fontsize=18)
plt.ylabel(r"angle d'équilibre (rad)", fontsize=18)

plt.legend(fontsize=18)
plt.show()

## bifurcation
fig = plt.figure()
ax = fig.add_subplot(111)
ax.tick_params(labelsize=18)

plt.grid()

plt.errorbar(omega, theta, xerr=u_omega, yerr=u_theta, fmt='o', ms=5, capsize=2, capthick=2, linewidth=2)
plt.errorbar(omegal, thetal, xerr=u_omegal, yerr=u_thetal, fmt='o', ms=5, capsize=2, capthick=2, linewidth=2)

plt.xlabel(r"$\Omega$ (rad.s$^{12}$)", fontsize=18)
plt.ylabel(r"angle d'équilibre (rad)", fontsize=18)

plt.ylim(0, 1.5)

plt.title("diagramme de bifurcation", fontsize=18)

plt.show()