import math
import numpy
from matplotlib.pyplot import *
import scipy.integrate
            

def equation(z,t):
    x1 = z[0]-1
    x2 = z[0]+1
    yc = z[2]*z[2]
    d1 = math.pow(x1*x1+yc,1.5)
    d2 = math.pow(x2*x2+yc,1.5)
    return [z[1],-x1/d1-x2/d2,z[3],-z[2]/d1-z[2]/d2]
            

tmax=50
te=0.01
t = numpy.arange(start=0,stop=tmax,step=te)
            

atol=1e-6
rtol=1e-6
            

z0 = [0.2,0,5,0]
            

z = scipy.integrate.odeint(equation,z0,t,rtol=rtol,atol=atol)
            

x = z[:,0]
vx = z[:,1]
y = z[:,2]
vy = z[:,3]
             

figure(figsize=(6,6))
plot(x,y)
plot([1,-1],[0,0],'r.')
xlabel('x')
ylabel('y')
axis([-6,6,-6,6])
             

tmax = 200
t = numpy.arange(start=0,stop=tmax,step=te)
z = scipy.integrate.odeint(equation,z0,t,rtol=rtol,atol=atol)
zt = z.transpose()
x = zt[0]
vx = zt[1]
y = zt[2]
vy = zt[3]
figure(figsize=(6,6))
plot(x,y)
plot([1,-1],[0,0],'r.')
xlabel('x')
ylabel('y')
axis([-6,6,-6,6])
             

tmax = 2000
t = numpy.arange(start=0,stop=tmax,step=te)
z = scipy.integrate.odeint(equation,z0,t,rtol=rtol,atol=atol)
zt = z.transpose()
x = zt[0]
vx = zt[1]
y = zt[2]
vy = zt[3]
figure(figsize=(12,4))
plot(t,x)
xlabel('t')
ylabel('x')
axis([0,tmax,-6,6])
             

def energie():
    e = numpy.zeros(t.size)
    for i in range(t.size):
        x1 = x[i]-1
        x2 = x[i]+1
        yc = y[i]*y[i]
        d1 = math.pow(x1*x1+yc,0.5)
        d2 = math.pow(x2*x2+yc,0.5)
        e[i] = 0.5*(vx[i]*vx[i]+vy[i]*vy[i])-1.0/d1-1.0/d2
    return e
figure(figsize=(12,4))
plot(t,energie())
xlabel('t')
ylabel('e')