import numpy as np

#SI units
AU = 149597870700.
GM_sun = 1.32712440018e20
GM_earth = 3.986004418e14

def proj_T(alpha, beta, gamma):
    return np.array([np.cos(alpha)*np.cos(beta)-np.sin(alpha)*np.sin(beta)*np.cos(gamma), np.cos(alpha)*np.sin(beta)+np.sin(alpha)*np.cos(beta)*np.cos(gamma), np.sin(alpha)*np.sin(gamma)])


def Keplerian2Cartesian(p, e, i, Omega, w, theta, GM = GM_sun):
    r = p/(1.+e*np.cos(theta))
    vbar = np.sqrt(GM/p)
    T = proj_T(theta+w, Omega, i)
    Torth = proj_T(theta+w+np.pi/2., Omega, i)
    r_vec = r*T
    v_vec = (vbar*e*np.sin(theta))*T + (vbar*(1.+e*np.cos(theta)))*Torth
    return r_vec, v_vec
        

def Cartesian2Keplerian( r_vec, v_vec, GM = GM_sun):
    r = np.linalg.norm(r_vec)
    v = np.linalg.norm(v_vec)
    assert(r > 0. and v > 0.)
    L_vec = np.cross(r_vec, v_vec)
    L = np.linalg.norm(L_vec)
    Energy = -GM/r + v**2/2.
    p = L**2/GM
    vbar = np.sqrt(GM/p)    
    e = np.sqrt(max(0., 1.+2.*p*Energy/GM)) #take max to avoid round off error
    n_r = r_vec/r
    n_L = L_vec/L
    n_x = np.array([1., 0., 0.])
    n_z = np.array([0., 0., 1.])
    asc_vec = np.cross(n_z, L_vec)
    asc = np.linalg.norm(asc_vec)
    if(asc > 1.e-10):
        n_asc = asc_vec/asc
    else:
        n_asc = n_x
    if(e > 1.e-6):
        n_peri = np.cross(v_vec, n_L)/(vbar*e) - n_r/e
    else:
        n_peri = n_r
    i = np.arccos(np.dot(n_L, n_z))
    Omega = np.arctan2(np.dot(np.cross(n_x, n_asc), n_z), np.dot(n_x, n_asc))
    w = np.arctan2(np.dot(np.cross(n_asc, n_peri), n_L), np.dot(n_asc, n_peri))
    theta = np.arctan2(np.dot(np.cross(n_peri, n_r), n_L), np.dot(n_peri, n_r))
    return p, e, i, Omega, w, theta


r_vec = np.random.normal(size=3) * AU
v_vec = np.random.normal(size=3) * 1.e4

p, e, i, Omega, w, theta = Cartesian2Keplerian(r_vec, v_vec)
print("Semi-latus rectum (p) = ", p/AU)
print("Eccentricity (e) = ", e)
print("Inclination (i) = ", i)
print("Ascending-node longitude (Omega) = ", Omega)
print("Argument of periapsis (omega) = ", w)
print("True anomaly (theta) = ", theta)

#transform back to check consistency
r_vec_back, v_vec_back = Keplerian2Cartesian(p, e, i, Omega, w, theta)
print("transform back:")
print("coordinate difference (AU):", (r_vec_back-r_vec)/AU)
print("velocity difference (km/s):", (v_vec_back-v_vec)/1.e3)