import numpy as np
import matplotlib.pyplot as plt

a = 17.8    # velka polos
e = 0.967    # excentricita
P = a**(1.5)    # perioda
n = np.arange(0, P)    # casovy vektor v rokoch
M = 2*np.pi/P*n   # stredna anomalia v radianoch
E = M*1
v = M*1
r = M*1
for i in range(len(n)):
    E[i] = M[i]
    # Riesenie Keplerovej rovnice pre excentricku anomaliu E.
    for c in range(10):
        E[i] = M[i] + e*np.sin(E[i])
    # Vypocet pravej anomalie.
    v[i] = 2*np.arctan(np.sqrt((1 + e)/(1 - e))*np.tan(E[i]/2))
    if v[i] < 0:    # korekcia pre uhly vacsie ako pi
        v[i] = v[i] + 2*np.pi
    r[i] = a*(1 - e**2)/(1 + e*np.cos(v[i]))    # r = vzdialenost od ohniska
    
plt.figure()
plt.plot(n, M)
plt.plot(n, E)
plt.plot(n, v)
plt.xlabel('n [rok]')
plt.ylabel('anomalia [rad]')
plt.legend(('stredna anomalia M', 'excentricka anomalia E', 'prava anomalia v'), loc=4)
plt.show()
plt.figure()
plt.polar(v, r, '.')
plt.show()
