import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint as ODEint

def deriv(X, t):
    x1, x2, v1, v2 = X.reshape(4, -1)
    a1 = -x1 * GM * ((x1**2).sum())**-1.5
    a2 = -x2 * GM * ((x2**2).sum())**-1.5
    return np.hstack((v1, v2, a1, a2))

GM  = 3.986E+14  # m^3/s^2
a   = 384399000. # meters
twopi = 2*np.pi

ecc  = 0.8
apo, peri          = a * (1+ecc), a * (1-ecc)
v_apo, v_a, v_peri = [np.sqrt(GM*(2./r - 1./a)) for r in (apo, a, peri)] # vis-viva equation

X0 = np.array([-apo, 0, -a, 0] + [0, -v_apo, 0, -v_a])

T  = twopi * np.sqrt(a**3/GM)
times = np.linspace(0, T, 10001)

answer, info = ODEint(deriv, X0, times, full_output=True)

dirt_pos, moon_pos = answer.T[:4].reshape(2, 2, -1)
r_dirt,   r_moon   = [np.sqrt((pos**2).sum(axis=0)) for pos in (dirt_pos, moon_pos)]
dirt_ang, moon_ang = [np.arctan2(p[1], p[0]) for p in (dirt_pos, moon_pos)]

i_dirt_1 = np.argmax(r_dirt < a)
ang_1 = dirt_ang[i_dirt_1]
i_moon_1 = np.argmax(moon_ang[1:] > ang_1) + 1

i_dirt_2 = np.argmax(r_dirt[i_dirt_1+2:] > a) + i_dirt_1 + 2
i_moon_2 = i_dirt_2 - i_dirt_1 + i_moon_1
print('ang_1: '   , ang_1)
print('i_dirt_1, i_dirt_2, diff: ', i_dirt_1, i_dirt_2, i_dirt_2-i_dirt_1)
print('i_moon_1, i_moon_2, diff: ', i_moon_1, i_moon_2, i_moon_2-i_moon_1)


if False:
    fig = plt.figure()
    ax  = fig.add_subplot(1, 1, 1)
    for a in (dirt_ang, moon_ang):
        plt.plot(a)
        plt.plot([i_dirt_1], a[i_dirt_1:i_dirt_1+1], 'ok')
    plt.show()


if True:
    fig = plt.figure()
    ax  = fig.add_subplot(1, 1, 1)
    # the moon
    x, y = moon_pos
    ax.plot(x, y, '-g')
    ax.plot(x[i_moon_1:i_moon_1+1], y[i_moon_1:i_moon_1+1], 'og', markersize=14)
    ax.plot(x[i_moon_2:i_moon_2+1], y[i_moon_2:i_moon_2+1], 'og', markersize=14)
    ax.plot(x[i_moon_1:i_moon_2], y[i_moon_1:i_moon_2], '-g', linewidth=3)
    # ax.plot(x[:1], y[:1], 'ok', markersize=8)
    # the dirt
    x, y = dirt_pos
    ax.plot(x, y, '-r')
    ax.plot(x[i_dirt_1:i_dirt_1+1], y[i_dirt_1:i_dirt_1+1], 'or', markersize=14)
    ax.plot(x[i_dirt_2:i_dirt_2+1], y[i_dirt_2:i_dirt_2+1], 'or', markersize=14)
    ax.plot(x[i_dirt_1:i_dirt_2], y[i_dirt_1:i_dirt_2], '-r', linewidth=3)
    # ax.plot(x[:1], y[:1], 'ok', markersize=8)
    # now mark the Earth's position
    ax.set_aspect('equal')
    ax.plot([0], [0], 'ob', markersize=16)
    plt.show()