import numpy as np
import matplotlib.pyplot as plt

# Densité de probabilité à l'instant t d'un paquet d'onde gaussien de vecteur
# d'onde centrale k0 et de largeur paramétrée par a.
# La fonction d'onde n'est pas normalisée.
# On prend hbar = 1 et m = 1

def densite_proba(x,t, k0, a) :

    A = ((x-k0*t)**2)*(a**2)/2/(a**2+t**2/4)
    proba = np.exp(-A)/(a**2+t**2/4)
    return proba

k0 = 100
a = 0.02
xmin = -3
xmax = 30
t = 0.0

X = np.linspace(xmin,xmax,500)
probaEchantillon = [ densite_proba(x,t, k0, a) for x in X]

# plt.figure()
# plt.xlabel("x")
# plt.ylabel("|psi(x,0)|^2")
# plt.title("Densité de probabilité à t = 0")
# plt.plot(X,probaEchantillon)
# plt.show()
# plt.grid()


import matplotlib.animation as animation

fig = plt.figure()
plt.title("Progression et étalement du paquet d'onde Q au cours du temps")
line, = plt.plot(X,probaEchantillon)
plt.xlabel("x")
plt.ylabel("|psi(x,t)|^2")
plt.grid()
t = 0.0
dt = 2.0e-4

def animate(i):
    global t,xmin,xmax,N, X, a, k0
    t += dt
    probaEchantillon = [ densite_proba(x,t, k0, a) for x in X]
    line.set_xdata(X)
    line.set_ydata(probaEchantillon)
    return line,

ani = animation.FuncAnimation(fig,animate,1000,interval=40, repeat = False)
plt.show()