#!/usr/bin/python3

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

chemin=""
nom="Circulation_accordeon_animation.gif"

nL = 35 # nombre de véhicules dans la longueur
nl = 3 # nombre de voies
L = nL*2 # longueur de voie

N = 200 # nombre d'images générées

v0 = 0.25 # vitesse globale de flux
vsin = 10*v0 # amplitude des perturbations sinusoïdales de vitesse
pesp = 0.3 # amplitude du placement aléatoire

xbase = np.linspace(0, L, nL)
x0 = np.array([xbase, xbase, xbase]) + pesp*np.random.randn(nl, nL)
x0[x0<0] = 0
x0[x0>L] = L

y0 = np.ones(nL)

fig, ax = plt.subplots(figsize = [10, 0.7])
ligne1, = ax.plot(x0[0,:], y0, "sk")
ligne2, = ax.plot(x0[1,:], 2*y0, "sk")
ligne3, = ax.plot(x0[2,:], 3*y0, "sk")
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

plt.show()

def anime(i):
    dx = i*v0
    x = x0 + dx + vsin*np.sin(6*np.pi*(x0+dx)/L)
    x[x>L] = x[x>L]-L
    ligne1.set_xdata(x[0,:])
    ligne2.set_xdata(x[1,:])
    ligne3.set_xdata(x[2,:])
    return ligne1, ligne2, ligne3

ani = animation.FuncAnimation(
    fig, anime, interval=100, blit=True, save_count=N)

ani.save(chemin+nom, writer="imagemagick")