import matplotlib.pyplot as plt
import numpy as np
import math as m
from matplotlib.animation import FuncAnimation

xmin = 0  # abscisse minimale de la fenêtre d'affichage
xmax = 2  # abscisse maximale de la fenêtre d'affichage
Lambda = (xmax - xmin)/2  # longueur d'onde (fenêtre d'affichage = 2 longueurs onde)
N = 15   # nombre de flèches du champ électrique affichées sur une longueur d'onde
h = Lambda/N

fig = plt.figure(1)
ax = fig.gca()
ax.set_xlim(xmin,xmax)
ax.set_ylim(-1.5,1.5)

def init() :
    ax.set_xlim(xmin,xmax)
    ax.set_ylim(-1.5,1.5)
    return a,
    
X = []
x = xmin
while x <= xmax :
    X.append(x)
    x += h

nombrePoints = len(X)
X = np.array(X)
minMax = [10,11]

## ONDE PLANE PROGRESSIVE SINUSOIDALE (ou HARMONIQUE) (décommenter pour voir)

# c = 10  # célérité de l'onde
# T = Lambda/c  # période temporelle de l'onde

# E = np.cos(2*np.pi/Lambda*X)
# 
# for i in range(nombrePoints) :
#     a = ax.arrow(X[i],0,0,E[i], width = 0.005)
# 
# def anime(t) :
#     global ax
#     ax.clear()
#     ax.set_xlim(xmin,xmax)
#     ax.set_ylim(-1.5,1.5)
#     U = 2*np.pi*(X/Lambda - t/T)
#     E = np.cos(U)
#     for i in range(nombrePoints) :   
#         ax.arrow(X[i],0,0,E[i], width = 0.005)
#     return a,

##  ONDE STATIONNAIRE (décommenter pour voir)

# c = 30  # célérité de l'onde
# T = Lambda/c  # période temporelle de l'onde
# 
# E = np.sin(2*np.pi/Lambda*X)
# for i in range(nombrePoints) :
#     a = ax.arrow(X[i],0,0,E[i], width = 0.005)
# 
# def anime(t) :
#     global ax
#     ax.clear()
#     ax.set_xlim(xmin,xmax)
#     ax.set_ylim(-1.5,1.5)
#     E = np.sin(2*np.pi*X/Lambda)*np.sin(2*np.pi*t/T)
#     for i in range(nombrePoints) :   
#         ax.arrow(X[i],0,0,E[i], width = 0.005)
#     return a,


##  PSEUDO-ONDE PLANE PROGRESSIVE HARMONIQUE - ABSORPTION (décommenter pour voir)

# c = 20  # célérité de l'onde
# T = Lambda/c  # période temporelle de l'onde
# k1 = 2*np.pi/Lambda  # partie réelle du vecteur d'onde complexe : propagation
# k2 = 1  # partie imaginaire du vecteur d'onde complexe : atténuation
# 
# E = np.exp(-k2*X)*np.cos(k1*X)
# 
# for i in range(nombrePoints) :
#     a = ax.arrow(X[i],0,0,E[i], width = 0.005)
# 
# def anime(t) :
#     global ax
#     ax.clear()
#     ax.set_xlim(xmin,xmax)
#     ax.set_ylim(-1.5,1.5)
#     U = 2*np.pi*(X/Lambda - t/T)
#     E = np.exp(-k2*X)*np.cos(U)
#     for i in range(nombrePoints) :
#         ax.arrow(X[i],0,0,E[i], width = 0.005)
#     return a,


## ONDE EVANESCENTE (décommenter pour voir)

c = 20  # célérité de l'onde
T = Lambda/c  # période temporelle de l'onde
k2 = 1 # partie imaginaire du vecteur d'onde complexe : atténuation

E = np.exp(-k2*X)

for i in range(nombrePoints) :
    a = ax.arrow(X[i],0,0,E[i], width = 0.005)

def anime(t) :
    global ax
    ax.clear()
    ax.set_xlim(xmin,xmax)
    ax.set_ylim(-1.5,1.5)
    E = np.exp(-k2*X)*m.cos(2*np.pi*t/T)
    for i in range(nombrePoints) :
        ax.arrow(X[i],0,0,E[i], width = 0.005)
    return a,


## PROGRAMME PRINCIPAL

dates = np.linspace(0,2*T,40)
dateInt = T/40*1000

ani=FuncAnimation(fig,anime,frames=dates, blit = False, init_func = init, interval=dateInt, repeat=True)    
fig.show()