def derivxx(x, mu):
    r1 = np.sqrt((x      + mu)**2)
    r2 = np.sqrt((x - 1. + mu)**2)
    term_1 = x
    term_2 = -(1.-mu) * (x + mu) / r1**3
    term_3 =     -mu  * (x - 1. + mu) / r2**3

    xddot  = term_1 + term_2 + term_3

    return xddot

def get_x_L1():

    mu_EM  = GMm / (GMm + GMe)
    mu_SE  = GMe / (GMe + GMs)
    mu_SEb = (GMe + GMm) / (GMe + GMm + GMs)

    mu = mu_SEb

    xtol = 1E-10
    answer, blob = spo.brentq(derivxx, 0.9, 0.9999, args=(mu),
                              full_output = True, xtol=xtol)
    print "converged = ", blob.converged
    if blob.converged:
        x_L1 = answer
    else:
        x_L1 = None

    return x_L1

def ROTIT(X, th):

    sth, cth, zth, oth = ( np.sin(th), np.cos(th),
                           np.zeros_like(th), np.ones_like(th) )

    xR = (X * np.vstack((cth, -sth,  zth))).sum(axis=0)
    yR = (X * np.vstack((sth,  cth,  zth))).sum(axis=0)
    zR = (X * np.vstack((zth,  zth,  oth))).sum(axis=0)

    return np.vstack([xR, yR, zR])

def cosine_ramp(a, b, n):
    x    = np.linspace(0, np.pi, n)
    ramp = 0.5*(1. - np.cos(x))
    y    = a + (b-a)*ramp

    return y

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

halfpi, pi, twopi = [f*np.pi for f in (0.5, 1, 2)]
GMs, GMe, GMm     = (1.32712440018E+20, 3.98600418E+14, 4.9048695E+12)  # m^3 s^-2

import json
import scipy.optimize as spo
from scipy.optimize import minimize as mini

EMB_pos = np.load('EMB_pos.npy')
SOHOH_pos = np.load('SOHOH_pos.npy')
SUN_pos = np.load('SUN_pos.npy')

x_L1 = get_x_L1()  # about 0.99  # I'm going to say the sun doesn't move here

EMB_from_SUN  = EMB_pos - SUN_pos

L1_from_SUN   = x_L1 * EMB_from_SUN

SOHO_from_SUN = SOHOH_pos - SUN_pos

SOHO_from_EMB = SOHO_from_SUN - EMB_from_SUN
SOHO_from_L1  = SOHO_from_SUN - L1_from_SUN

EMB_from_L1   = EMB_from_SUN - L1_from_SUN

theta = np.arctan2(EMB_from_SUN[1], EMB_from_SUN[0])  # OK here is the angle!

SOHO_from_L1_rot  = ROTIT(SOHO_from_L1, -theta)
EMB_from_L1_rot   = ROTIT(EMB_from_L1,  -theta)

fig = plt.figure(figsize=[10, 8])

ax = fig.add_subplot(1, 1, 1, projection='3d')
x, y, z = SOHO_from_L1_rot
ax.plot(x, y, z, '-c', linewidth=0.7)

nn1, nn2 = 20, 20  # first and last twenty point
ax.plot(x[:nn1], y[:nn1], z[:nn1], '.k')
ax.plot(x[-nn2:], y[-nn2:], z[-nn2:], '.k')
ax.plot([0], [0], [0], 'or')
x, y, z = EMB_from_L1_rot
ax.plot(x, y, z, marker='o', fillstyle='full', markeredgecolor='blue',
        color='blue', markeredgewidth=0.0)

if True:
    # Build a set of views
    ns, nm = 5, 31
    azis   = ns*[0] + cosine_ramp(0, -90, nm).tolist() + ns*[-90] + nm*[-90] + ns*[-90]
    altis  = ns*[0] + nm*[0] + ns*[0] + cosine_ramp(0, 90, nm).tolist() + ns*[90]
    views  = zip(altis, azis)

    for i, (alt, az) in enumerate(views[::10]):
        ax.view_init(alt, az)
        plt.savefig('SOHO_3D_' + str(10000+i)[1:])

plt.show()