import numpy as np
from stumpff import C, S
from CelestialBody import BODIES
from numerical import newton, laguerre
from lagrange import calc_f, calc_fd, calc_g, calc_gd


def kepler_chi(chi, alpha, r0, vr0, mu, dt):
    ''' Kepler's Equation of the universal anomaly, modified
    for use in numerical solvers. '''
    z = alpha*chi**2
    return (r0*vr0/np.sqrt(mu))*chi**2*C(z) + \
           (1 - alpha*r0)*chi**3*S(z) + \
           r0*chi - np.sqrt(mu)*dt


def dkepler_dchi(chi, alpha, r0, vr0, mu, dt):
    ''' Derivative of Kepler's Equation of the universal anomaly,
    modified for use in numerical solvers. '''
    z = alpha*chi**2
    return (r0*vr0/np.sqrt(mu))*chi*(1 - alpha*chi**2*S(z)) + \
           (1 - alpha*r0)*chi**2*C(z) + r0


def d2kepler_dchi2(chi, alpha, r0, vr0, mu, dt):
    ''' Second derivative of Kepler's Equation of the universal
    anomaly, modified for use in numerical solvers. '''
    z = alpha*chi**2
    S_ = S(z)
    return (r0*vr0/np.sqrt(mu))*(1 - 3*z*S_ + z*(C(z) - 3*S_)) + \
           chi*(1 - z*S_)*(1 - alpha*r0)


def solve_kepler_chi(r_0, v_0, dt, body=BODIES['Earth'], method='laguerre', tol=1e-7, max_iters=100):
    ''' Solve Kepler's Equation of the universal anomaly chi using the specified
    numerical method. Applies Algorithm 3.4 from Orbital Mechanics for Engineering
    Students, 4 ed, Curtis.

    :param r_0: `iterable` (km) initial position 3-vector
    :param v_0: `iterable` (km/s) initial velocity 3-vector
    :param dt: `float` (s) time after initial state to solve for r, v as 3-vectors
    :param body: `CelestialBody` (--) the celestial body to use for orbital parameters
    :param method: `str` (--) which numerical method to use to solve Kepler's Equation
    :param tol: `float` (--) decimal tolerance for numerical method (default 1e-7 is IEEE 745 single precision)
    :param max_iters: `int` (--) maximum number of iterations in numerical method before breaking
    :return: (km) final position 3-vector, (km/s) final velocity 3-vector
    '''
    VALID_METHODS = ('laguerre', 'newton')

    mu = body.mu  # (km**3/s**2) gravitational parameter of the specified primary body

    r0 = np.linalg.norm(r_0)  # (km) initial position magnitude
    v0 = np.linalg.norm(v_0)  # (km/s) initial velocity magnitude
    vr0 = np.dot(v_0, r_0)/r0  # (km/s) initial radial velocity magnitude
    alpha = 2/r0 - v0**2/mu  # (1/km) inverse of semi-major axis

    chi0 = np.sqrt(mu)*np.abs(alpha)*dt

    if method not in VALID_METHODS:
        print(f'Method \'{method}\' is not valid, must be one of {VALID_METHODS}.\nDefaulting to laguerre method.')
        chi, _, _ = laguerre(chi0, kepler_chi, dkepler_dchi, d2kepler_dchi2, alpha, r0, vr0, mu, dt)
    elif method == 'newton':
        chi, _, _ = newton(chi0, kepler_chi, dkepler_dchi, alpha, r0, vr0, mu, dt)
    else:  # method == 'laguerre'
        chi, _, _ = laguerre(chi0, kepler_chi, dkepler_dchi, d2kepler_dchi2, alpha, r0, vr0, mu, dt)

    f = calc_f(chi, r0, alpha)
    g = calc_g(dt, mu, chi, alpha)
    r_1 = f*r_0 + g*v_0
    r1 = np.linalg.norm(r_1)

    fd = calc_fd(mu, r1, r0, alpha, chi)
    gd = calc_gd(chi, r1, alpha)
    v_1 = fd*r_0 + gd*v_0

    return r_1, v_1


def solve_kepler_E(e, Me, tol=1e-7, max_iters=100):
    ''' Solve Kepler's Equation in the form containing Eccentric Anomaly (E),
    eccentricity (e), and Mean Anomaly of Ellipse (Me). Uses Algorithm 3.1 from Orbital
    Mechanics for Engineering Students, 4 ed, Curtis. '''

    # TODO: have this function make use of one of the numerical methods in numerical.py

    def f(E, e, Me):
        return E - e*np.sin(E) - Me

    def fp(E, e):
        return 1 - e*np.cos(E)

    E = Me + e/2 if Me < np.pi else Me - e/2
    ratio = f(E, e, Me)/fp(E, e)
    iters = 0
    while abs(ratio) > tol and iters < max_iters:
        E -= ratio
        ratio = f(E, e, Me)/fp(E, e)
        iters += 1

    E -= ratio
    converged = np.abs(ratio) <= tol

    return E, iters, converged


def test():
    ''' Test the functionality of solve_kepler_chi 
    and solve_kepler_laguerre using Problem 3.20 from
    Orbital Mechanics for Engineering Students, 4 ed, Curtis.
    '''
    # given starting information
    Earth = BODIES['Earth']  # `CelestialBody` (--) Earth and all the Earth things
    r_0 = np.array([20000, -105000, -19000])  # (km) initial position vector
    v_0 = np.array([0.9, -3.4, -1.5])  # (km/s) initial velocity vector
    dt = 2*60*60  # (s) time of interest after initial time

    # given correct answer from textbook
    correct_r_1 = np.array([26338, -128750, -29656])  # (km) final position vector
    correct_v_1 = np.array([0.86280, -3.2116, -1.4613])  # (km/s) final velocity vector

    # solve using above methods
    r_n, v_n = solve_kepler_chi(r_0, v_0, dt, Earth, method='newton')
    r_l, v_l = solve_kepler_chi(r_0, v_0, dt, Earth, method='laguerre')

    # check correctness
    # tolerance based on significant figures of given answers
    newton_valid = np.allclose(r_n, correct_r_1, atol=1) and np.allclose(v_n, correct_v_1, atol=1e-4)
    laguerre_valid = np.allclose(r_l, correct_r_1, atol=1) and np.allclose(v_l, correct_v_1, atol=1e-4)

    return all([newton_valid, laguerre_valid])


if __name__ == '__main__':
    print(test())

[((29, 9, 29, 13), 'stumpff.S', 'S', ({(29, 11, 29, 12): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((52, 9, 52, 28), 'numpy.linalg.norm', 'np.linalg.norm', ({(52, 24, 52, 27): 'r_0'}, {}), '(r_0)', True, 'import numpy as np\n'), ((53, 9, 53, 28), 'numpy.linalg.norm', 'np.linalg.norm', ({(53, 24, 53, 27): 'v_0'}, {}), '(v_0)', True, 'import numpy as np\n'), ((67, 8, 67, 30), 'lagrange.calc_f', 'calc_f', ({(67, 15, 67, 18): 'chi', (67, 20, 67, 22): 'r0', (67, 24, 67, 29): 'alpha'}, {}), '(chi, r0, alpha)', False, 'from lagrange import calc_f, calc_fd, calc_g, calc_gd\n'), ((68, 8, 68, 34), 'lagrange.calc_g', 'calc_g', ({(68, 15, 68, 17): 'dt', (68, 19, 68, 21): 'mu', (68, 23, 68, 26): 'chi', (68, 28, 68, 33): 'alpha'}, {}), '(dt, mu, chi, alpha)', False, 'from lagrange import calc_f, calc_fd, calc_g, calc_gd\n'), ((70, 9, 70, 28), 'numpy.linalg.norm', 'np.linalg.norm', ({(70, 24, 70, 27): 'r_1'}, {}), '(r_1)', True, 'import numpy as np\n'), ((72, 9, 72, 40), 'lagrange.calc_fd', 'calc_fd', ({(72, 17, 72, 19): 'mu', (72, 21, 72, 23): 'r1', (72, 25, 72, 27): 'r0', (72, 29, 72, 34): 'alpha', (72, 36, 72, 39): 'chi'}, {}), '(mu, r1, r0, alpha, chi)', False, 'from lagrange import calc_f, calc_fd, calc_g, calc_gd\n'), ((73, 9, 73, 32), 'lagrange.calc_gd', 'calc_gd', ({(73, 17, 73, 20): 'chi', (73, 22, 73, 24): 'r1', (73, 26, 73, 31): 'alpha'}, {}), '(chi, r1, alpha)', False, 'from lagrange import calc_f, calc_fd, calc_g, calc_gd\n'), ((113, 10, 113, 44), 'numpy.array', 'np.array', ({(113, 19, 113, 43): '[20000, -105000, -19000]'}, {}), '([20000, -105000, -19000])', True, 'import numpy as np\n'), ((114, 10, 114, 37), 'numpy.array', 'np.array', ({(114, 19, 114, 36): '[0.9, -3.4, -1.5]'}, {}), '([0.9, -3.4, -1.5])', True, 'import numpy as np\n'), ((118, 18, 118, 52), 'numpy.array', 'np.array', ({(118, 27, 118, 51): '[26338, -128750, -29656]'}, {}), '([26338, -128750, -29656])', True, 'import numpy as np\n'), ((119, 18, 119, 55), 'numpy.array', 'np.array', ({(119, 27, 119, 54): '[0.8628, -3.2116, -1.4613]'}, {}), '([0.8628, -3.2116, -1.4613])', True, 'import numpy as np\n'), ((54, 10, 54, 26), 'numpy.dot', 'np.dot', ({(54, 17, 54, 20): 'v_0', (54, 22, 54, 25): 'r_0'}, {}), '(v_0, r_0)', True, 'import numpy as np\n'), ((61, 20, 61, 100), 'numerical.laguerre', 'laguerre', ({(61, 29, 61, 33): 'chi0', (61, 35, 61, 45): 'kepler_chi', (61, 47, 61, 59): 'dkepler_dchi', (61, 61, 61, 75): 'd2kepler_dchi2', (61, 77, 61, 82): 'alpha', (61, 84, 61, 86): 'r0', (61, 88, 61, 91): 'vr0', (61, 93, 61, 95): 'mu', (61, 97, 61, 99): 'dt'}, {}), '(chi0, kepler_chi, dkepler_dchi, d2kepler_dchi2, alpha, r0, vr0, mu, dt\n    )', False, 'from numerical import newton, laguerre\n'), ((101, 16, 101, 29), 'numpy.abs', 'np.abs', ({(101, 23, 101, 28): 'ratio'}, {}), '(ratio)', True, 'import numpy as np\n'), ((127, 19, 127, 56), 'numpy.allclose', 'np.allclose', (), '', True, 'import numpy as np\n'), ((127, 61, 127, 101), 'numpy.allclose', 'np.allclose', (), '', True, 'import numpy as np\n'), ((128, 21, 128, 58), 'numpy.allclose', 'np.allclose', (), '', True, 'import numpy as np\n'), ((128, 63, 128, 103), 'numpy.allclose', 'np.allclose', (), '', True, 'import numpy as np\n'), ((14, 20, 14, 31), 'numpy.sqrt', 'np.sqrt', ({(14, 28, 14, 30): 'mu'}, {}), '(mu)', True, 'import numpy as np\n'), ((57, 11, 57, 22), 'numpy.sqrt', 'np.sqrt', ({(57, 19, 57, 21): 'mu'}, {}), '(mu)', True, 'import numpy as np\n'), ((57, 23, 57, 36), 'numpy.abs', 'np.abs', ({(57, 30, 57, 35): 'alpha'}, {}), '(alpha)', True, 'import numpy as np\n'), ((63, 20, 63, 82), 'numerical.newton', 'newton', ({(63, 27, 63, 31): 'chi0', (63, 33, 63, 43): 'kepler_chi', (63, 45, 63, 57): 'dkepler_dchi', (63, 59, 63, 64): 'alpha', (63, 66, 63, 68): 'r0', (63, 70, 63, 73): 'vr0', (63, 75, 63, 77): 'mu', (63, 79, 63, 81): 'dt'}, {}), '(chi0, kepler_chi, dkepler_dchi, alpha, r0, vr0, mu, dt)', False, 'from numerical import newton, laguerre\n'), ((65, 20, 65, 100), 'numerical.laguerre', 'laguerre', ({(65, 29, 65, 33): 'chi0', (65, 35, 65, 45): 'kepler_chi', (65, 47, 65, 59): 'dkepler_dchi', (65, 61, 65, 75): 'd2kepler_dchi2', (65, 77, 65, 82): 'alpha', (65, 84, 65, 86): 'r0', (65, 88, 65, 91): 'vr0', (65, 93, 65, 95): 'mu', (65, 97, 65, 99): 'dt'}, {}), '(chi0, kepler_chi, dkepler_dchi, d2kepler_dchi2, alpha, r0, vr0, mu, dt\n    )', False, 'from numerical import newton, laguerre\n'), ((22, 33, 22, 37), 'stumpff.C', 'C', ({(22, 35, 22, 36): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((30, 19, 30, 30), 'numpy.sqrt', 'np.sqrt', ({(30, 27, 30, 29): 'mu'}, {}), '(mu)', True, 'import numpy as np\n'), ((90, 21, 90, 30), 'numpy.cos', 'np.cos', ({(90, 28, 90, 29): 'E'}, {}), '(E)', True, 'import numpy as np\n'), ((12, 39, 12, 43), 'stumpff.C', 'C', ({(12, 41, 12, 42): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((13, 33, 13, 37), 'stumpff.S', 'S', ({(13, 35, 13, 36): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((87, 21, 87, 30), 'numpy.sin', 'np.sin', ({(87, 28, 87, 29): 'E'}, {}), '(E)', True, 'import numpy as np\n'), ((21, 19, 21, 30), 'numpy.sqrt', 'np.sqrt', ({(21, 27, 21, 29): 'mu'}, {}), '(mu)', True, 'import numpy as np\n'), ((21, 54, 21, 58), 'stumpff.S', 'S', ({(21, 56, 21, 57): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((30, 49, 30, 53), 'stumpff.C', 'C', ({(30, 51, 30, 52): 'z'}, {}), '(z)', False, 'from stumpff import C, S\n'), ((12, 19, 12, 30), 'numpy.sqrt', 'np.sqrt', ({(12, 27, 12, 29): 'mu'}, {}), '(mu)', True, 'import numpy as np\n')]