#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Evolution du paquet d'ondes pour le plasma
avec dispersion et étalement
"""

#Importation des bibliothèques
from ondes_plasma_relation_dispersion import *
import matplotlib.animation as animation 


#Définition des fonctions
def tau_c(sigma_0,P):
    """
        Temps caractéristique d'étalement d'un paquet d'onde gaussien
    """
    return sigma_0**2/(2*np.abs(P))

def sigma(sigma_0, P, tau):
    """
        Écart type d'un paquet d'onde gaussien au cours du temps
    """
    return sigma_0*np.sqrt(1+(tau/tau_c( sigma_0, P))**2)

def A( A_0,sigma_0, P, xi, tau):
    """
        Enveloppe d'un paquet d'onde gaussien au cours du temps
    """
    s = sigma(sigma_0, P, tau)
    return A_0*np.sqrt(sigma_0)/np.sqrt(s)*np.exp(-xi**2/(2*s**2))

def P(k):
    """
        P = 1/2 d2omega/dk2 adimensionné
    """
    return 1/2*1/(omega(k)**3)

def paquet( A_0, sigma_0, k_0, x, t):
    """
        Paquet d'onde total
    """
    p = P(k_0)
    omega_0 = omega(K_0)
    v_g_var = v_g(omega_0)
    xi = x-v_g_var*t
    tau = t
    a = A( A_0, sigma_0, p, xi, tau)
    return a*np.cos(omega_0*t-k_0*x)



# Programme principal
if __name__ == "__main__":
    A_0 = 1
    SIGMA_0 = 5
    K_0 = 1
    p = P(K_0)
    OMEGA_0 = omega(K_0)
    V_PHI = v_phi(OMEGA_0)
    V_G = v_g(OMEGA_0)

    #facteur de vitesse de l'animation
    f = 5
    t_max = 5000 # en ms
    l_0 = 2*np.pi/K_0*0.1
    x = np.linspace(-5*SIGMA_0, l_0+V_PHI*t_max*1e-3*f, 1000)
    paquet_0 = paquet(A_0, SIGMA_0, K_0, x, 0)
    enveloppe_0 = A(A_0, SIGMA_0, K_0, x, 0)

    fig, ax = plt.subplots()
    ax.set_ylim(bottom = -1.2*A_0, top=1.2*A_0)
    XLABEL = "temps adimensionné"
    YLABEL = "Amplitude"
    TITLE = "Propagation d'un paquet d'onde avec dispersion"
    ax.set_xlabel(XLABEL)
    ax.set_ylabel(YLABEL)
    ax.set_title(TITLE)
    
    
    signal, = ax.plot(x, paquet_0)
    enveloppe_sup, = ax.plot(x,enveloppe_0)
    enveloppe_inf, = ax.plot(x,-enveloppe_0, color="C1")
    max_signal, = ax.plot(0,A(A_0, SIGMA_0, p, 0, 0), marker='o', color='C1')
    
    phase_fixed_point, = ax.plot(l_0,paquet(A_0, SIGMA_0, K_0, l_0, 0), marker='o', color='C0')

    dt = 20 # en ms
    def animate(i): 
        t = f*i*dt*1e-3
        xi = x-V_G*t
        
        signal.set_ydata(paquet(A_0, SIGMA_0, K_0, x, t))
        enveloppe = A(A_0, SIGMA_0, p, xi, t)
        enveloppe_sup.set_ydata(enveloppe)
        enveloppe_inf.set_ydata(-enveloppe)
        x_m = V_G*t
        max_signal.set_xdata(x_m)
        max_signal.set_ydata(A(A_0, SIGMA_0, p, 0, t))
        x_phi = l_0 + V_PHI*t
        phase_fixed_point.set_xdata(x_phi)
        phase_fixed_point.set_ydata(paquet(A_0, SIGMA_0, K_0, x_phi, t))
        return signal, enveloppe_sup, enveloppe_inf,max_signal, phase_fixed_point,

    ani = animation.FuncAnimation(
        fig, animate, interval=dt, blit=True, frames=t_max//dt, repeat=True)

    writer = animation.FFMpegWriter(
      fps=1000//20, metadata=dict(artist='Antoine Chauchat'))
    ani.save("movie.mp4", writer=writer, dpi=200)
    plt.show()