import numpy as np

def kepler(nt,mstar,mp,a,e,ph0=0.0,ang0=0.0,body='planet'):
    """
    Returns one full period of the Kepler orbit of a planet with
    mass mp around a star with mass mstar, separated by semi-major
    axis a, having eccentricity e. The orbit lies in the (x,y) plane.
    For ph0==0 (default) the start of the planet is at periastron,
    while for ph0>0 it starts after periastron with time delay 
    Period*ph0/(2*pi). For ang0==0 (default) the ellipse lies in 
    the x-direction. For ang0>0 the ellipse is rotated
    counterclockwise.

    ARGUMENTS:
    nt      Nr of time steps between start and end of orbit (last 
            time step is at identical position as first).
    mstar   Star mass in gram
    mp      planet mass in gram
    a       semi-major axis in cm
    e       eccentricity

    OPTIONAL KEYWORDS:
    ph0     The phase delay after periastron of the start of the orbit
            Default = 0
    ang0    The rotation of the ellipse in the x,y plane (counterclockwise)
            Default = 0
    body    =='planet' means it returns the orbit of the planet wrt barycenter
            =='star' means it returns the orbit of the star wrt barycenter
            Default = 'planet'

    RETURNS:
    t[0:nt], x[0:nt], y[0:nt]   Time in seconds, and the location
            of the planet from barycenter x and y in cm 

    EXAMPLE:

    from kepler import *
    import matplotlib.pyplot as plt
    MS  = 1.98892e33     # Solar mass              [g]
    Mea = 5.9736e27      # Mass of Earth           [g]
    au  = 1.49598e13     # Astronomical Unit       [cm]
    t,x,y,vx,vy = kepler(100,MS,Mea,au,0.4,ph0=0.6,ang0=0.4,body='planet')
    rmau= 2.0
    fig = plt.figure()
    ax  = plt.axes(xlim=(-rmau, rmau), ylim=(-rmau, rmau))
    ax.set(aspect='equal',xlabel='x [au]',ylabel='y [au]')
    ax.plot(x/au,y/au)
    ax.plot([x[0]/au],[y[0]/au],'o') # Show start
    ax.plot([0],[0],'o')             # Show Sun
    plt.show()
    """
    GG     = 6.67408e-08    # Gravitational constant  [cm^3/g/s^2]
    mtot   = mstar+mp
    omk    = np.sqrt(GG*mtot/a**3)
    meanan = np.linspace(0.,2*np.pi,nt)+ph0
    time   = meanan/omk
    E      = meanan
    errtol = 1e-10
    err    = 1e10
    iter   = 0
    while err > errtol:
        Eold = E.copy()
        E    = meanan + e*np.sin(E)
        err  = np.abs(E-Eold).max()
        iter += 1
        assert iter<100, "Could not converge Kepler equation"
    cosE   = np.cos(E)
    cosnu  = (cosE-e)/(1-e*cosE)
    sinnu  = np.sqrt(1.-cosnu**2)
    sinnu[np.sin(meanan)<0] *= -1
    r      = a*(1.-e**2)/(1.+e*cosnu)
    x0     = r*cosnu
    y0     = r*sinnu
    cosa   = np.cos(ang0)
    sina   = np.sin(ang0)
    x      = cosa*x0 - sina*y0
    y      = sina*x0 + cosa*y0
    v      = np.sqrt(GG*mtot*(2.0/r-1.0/a))
    phi    = np.arctan(e*sinnu/(1.0+e*cosnu))
    vxn    = -v*y/r
    vyn    =  v*x/r
    cosphi = np.cos(phi)
    sinphi = np.sin(phi)
    vx     =  cosphi*vxn + sinphi*vyn
    vy     = -sinphi*vxn + cosphi*vyn
    if body=='planet':
        factor = (mstar/mtot)
    elif body=='star':
        factor = -(mp/mtot)
    else:
        factor = 1.0
    x  = factor * x
    y  = factor * y
    vx = factor * vx
    vy = factor * vy
    return time,x,y,vx,vy