from matplotlib import pyplot

# time increment
dt = 0.001

def f_p1(p1, p2, q1, q2):
    return 2*q1*(2-p2**2/8.0-q2**2*(1+q2**2/(p1**2+p2**2)**2))

def f_p2(p1, p2, q1, q2):
    return 2*q2*(2-p1**2/8.0-q1**2*(1+q1**2/(q1**2+q2**2)**2))

def f_q1(p1, p2, q1, q2):
    return p1*q2**2/4.0 

def f_q2(p1, p2, q1, q2):
    return p2*q1**2/4.0

#	foutrh order Runge-Kutta routine
def rk4(p1, p2, q1, q2):
    global dt
    k1p1 = dt * f_p1(p1, p2, q1, q2)
    k1p2 = dt * f_p2(p1, p2, q1, q2)
    k1q1 = dt * f_q1(p1, p2, q1, q2)
    k1q2 = dt * f_q2(p1, p2, q1, q2)

    k2p1 = dt * f_p1(p1 + k1p1/2.0, p2 + k1p2/2.0, q1 + k1q1/2.0, q2 + k1q2/2.0)
    k2p2 = dt * f_p2(p1 + k1p1/2.0, p2 + k1p2/2.0, q1 + k1q1/2.0, q2 + k1q2/2.0)
    k2q1 = dt * f_q1(p1 + k1p1/2.0, p2 + k1p2/2.0, q1 + k1q1/2.0, q2 + k1q2/2.0)
    k2q2 = dt * f_q2(p1 + k1p1/2.0, p2 + k1p2/2.0, q1 + k1q1/2.0, q2 + k1q2/2.0)

    k3p1 = dt * f_p1(p1 + k2p1/2.0, p2 + k2p2/2.0, q1 + k2q1/2.0, q2 + k2q2/2.0)
    k3p2 = dt * f_p2(p1 + k2p1/2.0, p2 + k2p2/2.0, q1 + k2q1/2.0, q2 + k2q2/2.0)
    k3q1 = dt * f_q1(p1 + k2p1/2.0, p2 + k2p2/2.0, q1 + k2q1/2.0, q2 + k2q2/2.0)
    k3q2 = dt * f_q2(p1 + k2p1/2.0, p2 + k2p2/2.0, q1 + k2q1/2.0, q2 + k2q2/2.0)

    k4p1 = dt * f_p1(p1 + k3p1, p2 + k3p2, q1 + k3q1, q2 + k3q2)
    k4p2 = dt * f_p2(p1 + k3p1, p2 + k3p2, q1 + k3q1, q2 + k3q2)
    k4q1 = dt * f_q1(p1 + k3p1, p2 + k3p2, q1 + k3q1, q2 + k3q2)
    k4q2 = dt * f_q2(p1 + k3p1, p2 + k3p2, q1 + k3q1, q2 + k3q2)

    p1 = p1 + k1p1/6.0 + k2p1/3.0 + k3p1/3.0 +k4p1/6.0
    p2 = p2 + k1p2/6.0 + k2p2/3.0 + k3p2/3.0 +k4p2/6.0
    q1 = q1 + k1q1/6.0 + k2q1/3.0 + k3q1/3.0 +k4q1/6.0
    q2 = q2 + k1q2/6.0 + k2q2/3.0 + k3q2/3.0 +k4q2/6.0

    return [p1, p2, q1, q2]

# list of all points to be plotted
all_p1 = []
all_p2 = []
all_q1 = []
all_q2 = []

i = 0
while i < 10:
	# set initial values
    p1 = 1 + i / 20.0
    p2 = 1 + i / 20.0
    q1 = 2 + i / 20.0
    q2 = 2 + i / 20.0

	# list of points for one orbit
    list_p1 = [p1]
    list_p2 = [p2]
    list_q1 = [q1]
    list_q2 = [q2]

    t = 1
    while t < 20000:
        [p1, p2, q1, q2] = rk4(list_p1[t-1], list_p2[t-1], list_q1[t-1], list_q2[t-1])
        list_p1.append(p1)
        list_p2.append(p2)
        list_q1.append(q1)
        list_q2.append(q2)
        t = t + 1

    all_q1 = all_q1 + list_q1
    all_q2 = all_q2 + list_q2
    i = i + 1

# draw graph using [all_x, all_y, all_z]
all_r1 = [x**2 for x in all_q1]
all_r2 = [x**2 for x in all_q2]
fig = pyplot.figure()
ax = fig.add_subplot(111)
ax.scatter(all_r1, all_r2, s=0.1)
ax.set_xlabel('r1')
ax.set_ylabel('r2')
pyplot.show()