%matplotlib inline
import math
import numpy as np
import matplotlib.pyplot as plt

class Planet:

    def __init__(self,x,y,vx,vy):
        self.x = []
        self.y = []
        self.vx = []
        self.vy = []
        self.vx.append(vx)
        self.vy.append(vy)
        self.x.append(x)
        self.y.append(y)

    def doTimeStep(self,dt):
        k = 4*math.pi**2
        n = 0
        initialAngle = np.sign(self.y[n]) * np.arccos(self.x[n] / np.sqrt(self.x[n] ** 2 + self.y[n] ** 2))
        okToExitLoop = False

        for n in range(1,1000000):
            self.vx.append(self.vx[n-1]+(-k)*self.x[n-1]/((self.x[n-1]**2+self.y[n-1]**2)**(3/2))*dt)
            self.vy.append(self.vy[n-1]+(-k)*self.y[n-1]/((self.x[n-1]**2+self.y[n-1]**2)**(3/2))*dt)
            self.x.append(self.x[n-1]+self.vx[n]*dt)
            self.y.append(self.y[n-1]+self.vy[n]*dt)

            if np.sign(self.y[n]) * np.arccos(self.x[n] / np.sqrt(self.x[n] ** 2 + self.y[n] ** 2)) < initialAngle:
                okToExitLoop = True
            if np.sign(self.y[n]) * np.arccos(self.x[n] / np.sqrt(self.x[n] ** 2 + self.y[n] ** 2)) > initialAngle and okToExitLoop:
                break

            #if self.x[n]>self.x[0] and self.x[n-1]<self.x[0]:
                #break

    def getPeriodAndAxis(self,dt):
        self.a = 0
        b = 0
        self.T = dt*len(self.x)
        for n in range(len(self.x)):
            b = np.sqrt(self.x[n]**2+self.y[n]**2)
            if b > self.a:
                self.a = b
        self.ratio = (self.T**2)/self.a**3

def PlotPlanet(x,y,vx,vy,dt):
    Planet1 = Planet(x,y,vx,vy)
    Planet1.doTimeStep(dt)
    plt.plot(Planet1.x,Planet1.y)
    Planet1.getPeriodAndAxis(dt)
    print('T^2/a^3 =')
    print(Planet1.ratio)

PlotPlanet(1,0,0,1,0.00001)
PlotPlanet(1,0,1,1,0.00001)
PlotPlanet(1,0,2,1,0.00001)
PlotPlanet(1,0,3,1,0.00001)
PlotPlanet(1,0,1,2,0.00001)
PlotPlanet(1,0,1,3,0.00001)