import numpy as np
import scipy.integrate as ode
import matplotlib.pyplot as plt
import matplotlib.animation as anim

alpha=0.1
m=1.
L=1.
u0=1.

A=u0*np.pi/(m*L)
beta=alpha/m
gamma=2*np.pi/L

def func(y,t,A,beta,gamma):
    dy=np.zeros_like(y)
    dy[0]=y[1]
    dy[1]=-A*np.sin(gamma*y[0])-beta*y[1]
    return dy

y0=[0.,np.sqrt(2.5)]
time=np.arange(0.,15.,0.001)
res=ode.odeint(func,y0,time,atol=1.e-10,rtol=1.e-10,args=(A,beta,gamma))

x=np.linspace(0,max(res[:,0]),1000)
epot=0.5*A*m*(1-np.cos(gamma*res[:,0]))
epot_plot=0.5*A*m*(1-np.cos(gamma*x))

fig1,[p1,p2]=plt.subplots(1,2)

p1.plot(time,res[:,0])
p2.plot(time,res[:,1])
p1.grid()
p2.grid()

fig2,p3=plt.subplots()
sat, = p3.plot([], [], 'o',color="red")
def animf(i):
    sx = [res[i,0]]
    sy = [epot[i]]
    sat.set_data(sx, sy)
    return sat,
p3.plot(x,epot_plot)
ani = anim.FuncAnimation(fig2, animf, np.arange(1, len(res[:,0])),interval=1,repeat=False,blit=True)
p3.grid()
plt.show()