import pylab as plt
import numpy as np
import numpy.linalg as alg
import matplotlib as mpl
from matplotlib.widgets import Slider, Button

#mpl.rcParams['text.usetex'] = True #True pour mettre du beau LateX mais ça fait laguer ou False c'est moche mais ça lag pas
mpl.rcParams['axes.titlesize'] = 15
mpl.rcParams['axes.labelsize'] = 30
mpl.rcParams['lines.linewidth'] = 3
mpl.rcParams['lines.markersize'] = 8
mpl.rcParams['xtick.labelsize'] = 20
mpl.rcParams['ytick.labelsize'] = 20
mpl.rcParams['legend.fontsize'] = 25

#Résolution de l'effet tunnel pour une particule quantique d'énergie E<V_0 venant de x<0 et arrivant sur une barrière d'épaisseur e et de hauteur V_0 

##Le gros du calcul

# The parametrized function to be plotted
def to_plot(X,V0,E,e) :
    #Paramètres physiques de valeurs non réelles pour ne pas s'embêter avec les puissances de 10 
    m = 1
    hbar = 1
    i = complex(0,1)
    #Vecteurs d'ondes
    k = (2*m*E)**0.5/hbar #particule en dehors de la barrières
    alpha = (2*m*(V0-E))**0.5/hbar #dans la barrière
    
    #Résolution des conditions limites
    A = np.array([[0,-1,1,1],[0,i*k,alpha,-alpha],[-np.exp(i*k*e),0,np.exp(alpha*e),np.exp(-alpha*e)],[-i*k*np.exp(i*k*e),0,alpha*np.exp(alpha*e),-alpha*np.exp(-alpha*e)]])
    B = np.array([1,i*k,0,0])
    [A3,B1,A2,B2] = np.linalg.solve(A,B)
    A1 = 1
    
    #Calcul de la densité de probabilité
    x1 = X[np.where(X<=0)]
    y1 = abs(A1*np.exp(i*k*x1)+B1*np.exp(-i*k*x1))**2
    x2 = X[np.where(np.logical_and( X>0,X<e))]
    y2 = abs(A2*np.exp(alpha*x2)+B2*np.exp(-alpha*x2))**2
    x3 = X[np.where(X>=e)]
    y3 = abs(A3*np.exp(i*k*x3))**2
            
    return np.concatenate((y1,y2,y3))
    
X = np.linspace(-5,5,1000)

##LE plot avec sliders

# Define initial parameters
init_V0 = 20
init_E = 10
init_e = 1

# Create the figure and the line that we will manipulate
fig, ax = plt.subplots()
line, = plt.plot(X, to_plot(X, init_V0, init_E,init_e))
line2 = plt.axvline(init_e,linestyle='--',color='black')
plt.axvline(0,linestyle='--',color='black')
ax.set_xlabel(r'$x$')
ax.set_ylabel('Densité de probabilité '+r'$|\phi(x)|^2$')
ax.set_xlim(-5,5)

# Adjust the main plot to make room for the sliders
#plt.subplots_adjust(left=0.25, bottom=0.25)
plt.subplots_adjust(left=0.1,bottom=0.25,top=0.95)

# Make horizontal sliders to control the parameters.
ax_V0 = plt.axes([0.1, 0.07, 0.8, 0.03])
V0_slider = Slider(
    ax=ax_V0,
    label=r'$V_0$',
    valmin=0.1,
    valmax=30,
    valinit=init_V0,
)
V0_slider.label.set_size(20)
V0_slider.valtext.set_fontsize(20)

ax_E = plt.axes([0.1, 0.1, 0.8, 0.03])
E_slider = Slider(
    ax=ax_E,
    label=r'$E$',
    valmin=0.1,
    valmax=30,
    valinit=init_E,
)
E_slider.label.set_size(20)
E_slider.valtext.set_fontsize(20)

ax_e = plt.axes([0.1, 0.13, 0.8, 0.03])
e_slider = Slider(
    ax=ax_e,
    label=r'$e$',
    valmin=0.001,
    valmax=4.5,
    valstep = 0.001,
    valinit=init_e,
)
e_slider.label.set_size(20)
e_slider.valtext.set_fontsize(20)

# The function to be called anytime a slider's value changes
def update(val):
    line.set_ydata(to_plot(X, V0_slider.val, E_slider.val, e_slider.val ))
    line2.set_xdata(e_slider.val)
    fig.canvas.draw_idle()


# register the update function with each slider
V0_slider.on_changed(update)
E_slider.on_changed(update)
e_slider.on_changed(update)

# Create a `matplotlib.widgets.Button` to reset the sliders to initial values.
resetax = plt.axes([0.8, 0.025, 0.1, 0.04])
button = Button(resetax, 'Reset', hovercolor='0.975')

def reset(event):
    V0_slider.reset()
    E_slider.reset()
    e_slider.reset()
button.on_clicked(reset)

plt.show()