# -*- coding: utf-8 -*-

"""
Le code initial est issu du site http://scipyscriptrepo.com/wp/?page_id=2
"""

"""
This short code snippet utilizes the new animation package in matplotlib 1.1.0
It's the shortest snippet that I know of that can produce an animate plot in
python. I'm hoping that the animation package can be improved so that one could
more simply animate things. What do you think?
"""
import numpy as np
import math 
from pylab import *
from matplotlib import gridspec
# import matplotlib.pyplot as plt
import matplotlib.animation as animation

##### Fonction de calcul #####
def simData():
# this function is called as the argument for
# the simPoints function. This function contains
# (or defines) and iterator---a device that computes
# a value, passes it back to the main program, and then
# returns to exactly where it left off in the function.
# I believe that one has to use this method to animate a plot
# using the matplotlib animation package.

#initialisation 
    puls = 40  #puls au carre
    lbd = 0.25 #coeff de frottement
    t_max = 10.0 #temps de l anim
    dt = 0.02 # pas de calcul
    theta0 = -1     #condition initial en theta
    omega0 = 15 #condition initial en omega
    t = 0.0 #initialisation du temps 
    while t < t_max: 
        #Calculs des coefficients de RK4:
        j1 = omega0*dt
        k1 = (-2*lbd*omega0-puls*math.sin(theta0))*dt
        j2 = (omega0+k1/2)*dt
        k2 = (-2*lbd*omega0-puls*math.sin(theta0+(j1/2)))*dt
        j3 = (omega0+k2/2)*dt
        k3 = (-2*lbd*omega0-puls*math.sin(theta0+j2/2))*dt
        j4 = (omega0+k3)*dt
        k4 = (-2*lbd*omega0-puls*math.sin(theta0+j3))*dt
        theta1 = theta0 + (1/6.0)*(j1 + 2*j2 + 2*j3 + j4)
        omega1 = omega0 + (1/6.0)*(k1 + 2*k2 + 2*k3 + k4)
        # Passage en cartesien:
        x = -2*math.cos(theta1)
        y = 2*math.sin(theta1)
        #pour boucler:
        theta0 = theta1 
        omega0 = omega1
        t = t + dt
        yield x, y, t, theta1 #renvoie des valeurs 
 #(yield = iterateur, les valeurs sont transmises mais pas stockees)
        
##### Fonction graphique ####
theta = [] #liste qui contiendra tous les theta calcules
vectt = [] #liste qui contiendra les temps correspondants
def simPoints(simData):
    #Recuperation des valeurs de la fonction simdata retourne par yield
    x, y, t, theta1 = simData[0], simData[1], simData[2], simData[3]
    #Passage de theta en degre
    theta.append(round(theta1*180/np.pi,0)) #on ajoute le theta a la liste des theta
    vectt.append(t) #on ajoute le temps a la liste des temps
    theta1 = round(theta1*180/np.pi,0) #passage de l angle en degre
    time_text.set_text(time_template%(t)) #affichage du temps
    angle_text.set_text(angle_template%(theta1)) #affichage de l angle
    x = [0,x] #permet le dessin du point (0,0)
    y = [0,y]
    line.set_data(y, x)  #on tracera y en fonction de x
    linebis.set_data(vectt,theta) #et theta en fonction de t
    return line, linebis, time_text, angle_text
 
#####   set up figure for plotting: #####
fig = plt.figure()
# fig = plt.figure()
gs = gridspec.GridSpec(2, 1,height_ratios=[6,2]) #creation d une grille avec plusieurs plot
ax = plt.subplot(gs[0]) # premier plot
# ax.axis('equal') #axes premiere figure normes
ax.set_aspect('equal',adjustable='box')  #axes orthonorme
ax.set_xlim(-4, 4) #limite en x
ax.set_ylim(-4.0, 4.0) #limite en y
line, = ax.plot([], [], 'bo-',lw=3, ms=10) #trace du premier plot
axbis = plt.subplot(gs[1]) # deuxieme plot
axbis.set_xlim(0, 10) #limite en x
axbis.set_ylim(0, 600) #limite en y
linebis, = axbis.plot([], [], 'bo-',lw=3, ms=1)#trace du deuxieme plot

# option graphique:
# bo : blue ball ; - : trait ; lw : line width ; ms : taille : ball

time_template = 'Time = %.1f s'    # prints running simulation time arrondi a 1 decimal
angle_template = 'Angle = %.1f deg' 
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)# coordonnees du texte
angle_text = ax.text(0.05, 0.8, '', transform=ax.transAxes)
 
#### Now call the animation package: ####
ani = animation.FuncAnimation(fig, simPoints, simData, blit=False, interval=50, repeat=False)
plt.show()