from __future__ import print_function, division
import numpy as np
import matplotlib.pyplot as plt
plt.ion()

def one_step(x_vect, v_vect, dt):
    #Current radius of the satellite
    radius = np.sqrt(np.sum(x_vect**2))
    
    #Calculate the acceleration
    #-x_vect/radius is a unit vector pointing to the origin.
    accel = -x_vect/radius**3
    
    #Use the semi-implicit Euler's method
    v_vect += accel*dt
    x_vect += v_vect*dt
    return x_vect, v_vect
    
def plot_orbit(x_vect = [1,0], v_vect=[0,0.5], dt=0.02, maxt=10):
    """plot an orbit"""
    #Check our inputs.
    x_vect = np.array(x_vect).astype(float)
    v_vect = np.array(v_vect).astype(float)
    nsteps = int(maxt/dt)
    
    #Initialize plotting
    plt.clf()
    plt.xlabel('X (norm units)')
    plt.ylabel('Y (norm units)')
    plt.axis([-1.2,1.2,-1.2,1.2])
    plt.plot(0,0,'x')
    all_x = x_vect.reshape(1,2)
    
    #Loop through and plot the orbit
    for i in range(nsteps):
        x_vect, v_vect = one_step(x_vect, v_vect, dt)
        all_x = np.append(all_x,x_vect.reshape(1,2),axis=0)

    plt.plot(all_x[:,0], all_x[:,1])
        

if __name__ == "__main__":
    plot_orbit()