import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
import cmath
import sys

# Choix de la version à lancer : ver = 1 ou ver = 2 paquet(s) d'ondes
ver = 1


## Discrétisation
# Discrétisation de la coordonnée spatiale
zmin,zmax,nz = -0,900,1200
##zmin,zmax,nz = -0,6000,9000
Z = np.linspace(zmin,zmax,num=nz)
# Discrétisation de la coordonnée temporelle
tmin,tmax,nt = 00,900,400
dt = (tmax-tmin)/nt



## Propriétés du paquet d'onde
# On va prendre un spectre du paquet d'ondes en forme de gaussienne
omega0 = 1 # pulsation centrale
deltaomega = 0.15 # "largeur" spectrale
# On va calculer l'onde en faisant en réalité une somme discrète (au lieu d'une intégrale) d'un très grand nombre d'OPPH
# On définit ci-dessous les pulsations de chacune de ces OPPH
nomega = 1000 # nombre d'OPPH additionnées dans la superposition
omega = np.linspace(omega0-3*deltaomega,omega0+3*deltaomega,num=nomega)



## Spectre du paquet d'ondes
# En fait cette fonction reçoit en entrée la pulsation oméga et renvoie l'amplitude (complexe) de la composante de pulsation oméga
# Ici il s'agit d'un spectre en forme de gaussienne
def s(w):
    return cmath.exp(-(w-omega0)**2/deltaomega**2) # amplitude *complexe* d'où le cmath

def st(w):  # et là c'est le même paquet d'ondes, mais translaté en z de 70
    if ver == 1:
        return 0+0j
    return cmath.exp(-(w-omega0)**2/deltaomega**2)*cmath.exp(1j*w*70)


## On définit ici la relation de dispersion
def k(w):
    wp = 0.5
    if w<wp:   # tel quel le programme n'a pas été écrit pour prendre en compte les k complexes
        print('Attention vous demandez une pulsation pour laquelle k est imaginaire pur')
        sys.exit()
    return np.sqrt(w**2-wp**2) # on a pris c=1
    

## On fabrique le paquet d'onde avec k(w) qui vient de la relation de dispersion
def paquetdondes(z,t):
    # En entrée, un tableau numpy contenant les coordonnées z, et un scalaire contenant la date t
    # Cette fonction renvoie la valeur (complexe) du paquet d'ondes à la date t, en tous les points de z, sous forme d'un tableau numpy de longueur len(z)
    h = np.zeros_like(z,dtype=complex)   # il faut préciser que l'on veut un tableau de complexes
    for p in range(nomega):
        h += (s(omega[p])+st(omega[p]))*np.exp(1j*(omega[p]*t-k(omega[p])*z))
    return h




# Création de la figure "vide"
fig = plt.figure(figsize=(18,5))
line, = plt.plot([],[])
# On précise les limites du tracé en abscisse et en ordonnée
plt.xlim(zmin,zmax)
plt.ylim(-300,300)

def animate(i):
    # C'est ici que l'on rafraîchit, pour chaque date t, la figure à tracer
    t = tmin+i*dt
    y = paquetdondes(Z,t)
    line.set_data(Z,y.real)  # attention on doit tracer la partie réelle
    return line,

# Tracé de l'animation
ani = animation.FuncAnimation(fig,animate,frames=nt,blit=True,interval=10,repeat=False,save_count=1500)
writergif = animation.PillowWriter(fps=30)
#ani.save("etalement1paquetdondes.gif",writer=writergif)  # au cas où on souhaite enregistrer l'animation produite

plt.show()