# -*- coding: utf-8 -*-
from numpy import sin, cos, pi, array
import numpy as np
import scipy.integrate as scint
import matplotlib.pyplot as plt
import matplotlib.animation as animation

a = 0
b = 100
w02 = 10
y0 = np.array([np.pi / 2, 0])
ybis0 = np.array([np.pi / 2, 0])
N = 10000

def f_amorti(y, t):
    # y[0] = angle
    # y[1] = vitesse angulaire
    frottement = 0.5 * y[1]
    # On renvoie : (vitesse, accélération)
    return np.array([y[1], -w02 * np.sin(y[0]) - frottement])

def f_libre(y, t):
    # y[0] = angle
    # y[1] = vitesse angulaire
    # On renvoie : (vitesse, accélération)
    return np.array([y[1], -w02 * np.sin(y[0])])

t = np.linspace(a, b, N + 1)
y = scint.odeint(f_libre, y0, t)
ybis = scint.odeint(f_amorti, ybis0, t)

theta = y[:, 0]
thetadot = y[:, 1]

thetabis = ybis[:, 0]
thetabisdot = ybis[:, 1]

fig = plt.figure()
# 111 : découpage 1x1, sous-figure numéro 1
ax = fig.add_subplot(111, autoscale_on=False,
                     xlim=(-1.5, 1.5), ylim=(-1.5, .5))
ax.grid()

plt.title('Pendule')
plt.xlabel('x')
plt.ylabel('y = f(x)')
pendules = plt.plot((0, sin(y0[0])),   (0, -cos(y0[0])),    'o-',
                    (0, sin(ybis0[0])), (0, -cos(ybis0[0])), 'o-')

def animate(i):
    angle = theta[i]
    x = (0, sin(angle))
    y = (0, -cos(angle))
    pendules[0].set_data(x, y)

    anglebis = thetabis[i]
    xbis = (0, sin(anglebis))
    ybis = (0, -cos(anglebis))
    pendules[1].set_data(xbis, ybis)


ani = animation.FuncAnimation(fig, animate, interval=10)

plt.show()