'''
    Solve Kepler problem at N=100 points along the orbit 
    of a Saturn-class exoplanet K2-261b discovered in 2018 around a
    sun-like star GJ3512b
        see exoplanet.eu/catalog/k2-261_b/
    at uniformly distributed times 
        t_i = (i/100)*P, 
    where P = orbital period, and index i = 0...99.
    Find 100 eccentric anomalies E(t_i), the same number of
    theta's (true anomalies = azimuthal angles), and their
    corresponding star-planet distances
    r = a(1-e^2)/(1+e cos(theta))
    where e = eccentricity = 0.39, a = semi-major axis = 0.102 AU,
    
    Kepler equation
        E - e sin(E) = 2 pi t/P
    cannot be solved on paper, but can be solved iteratively in Python. 
    After that's done, we use a connection formula between E and theta
        tan(theta/2) = sqrt((1+e)/(1-e)) tan(E/2)
    and the equation of ellipse in polar coordinates (r,theta):
        r(t) = a(1-e*e)/(1+e*cos(theta(t))),
    where
    M  = 2 pi t/P = the right-hand side of Kepler equation, 
        a.k.a. mean anomaly (angle),
    E = eccentric anomaly (angle),
    theta = true anomaly (angle), 
    r = radius, or momentary star-planet distance
    (The reason for these names is simply astronomical tradition.)
    
    This allows to start with the phase of the motion (fraction t/P, 
    i.e. time t in units of full orbital periods), obtain E by
    iteration, calculate theta, calculate r, and finally 
    from (r,theta) the Cartesian corrdinates (x,y) 
    which we use in plotting the position in a figure
        x = r cos(theta)
        y = r sin(theta)
'''

from numpy import pi,array,zeros,sqrt,tan,arctan,sin,cos,sum,around
import matplotlib.pyplot as plt
#   input number of points N, eccentricity and semi-maj. axis
N = 180
e = 0.4356; a = 0.338

# Calculate mean temperature on the planet
L = 0.00083 # [L_sun]
T = 280 * L**0.25 * a**(-0.5) *(1-e*e)**(-1/8)


#   the results will be accumulated in 
result = zeros((5,N), dtype=float)
#   prepare plot frame and description
plt.figure(figsize=(10,6.8))
plt.xlabel('X [AU]',fontsize=14)
plt.ylabel('Y [AU]',fontsize=14)
plt.text(.25,.3,'e = 0.4356',fontsize=14)
plt.text(.25,.25,'M$_*$ = 0.123 M$_{sun}$',fontsize=14)
plt.text(.25,.2,'m$_p$ >= 0.463 M$_J$',fontsize=14)
plt.text(.25,.15,'T$_p$ = 84 K$',fontsize=14)

plt.title('Orbit of exoplanet GJ 3512b at '+str(N)+' equal time intervals',   fontsize=14)
plt.axis([-1.75*a, 1.75*a, -1.2*a, 1.2*a])
plt.grid(True,color=(.85,.9,.97))

# compute and plot in x,y coordinates N points of the orbit
for i in range(N) :
    M = 2*pi * i/N
    E = M                # first guess to start with 
    for k in range(99):  # counter of iterations
        E_previous = E   # save previous E for comparison
        E = M + e*sin(E) # Kepler equation iteration
        if (i%30==0):  print('i',i,' k',k," E =",E) # every 30th 
        if (abs(E-E_previous)< 1e-9): break  # accuracy ok
    tanTheta_2 = sqrt((1+e)/(1-e))*tan(E/2)    
    Theta = 2*arctan(tanTheta_2)
    r = a*(1-e*e)/(1+e*cos(Theta))
    x = r*cos(Theta); y = r*sin(Theta)
# run, see format of output, learn C-style formatting below
    print("i=%d k=%d M=%7.4f E=%7.4f Theta=%8.4f r=%7.4f x=%6.3f y=%7.3f" \
          %(i,k,M,E,Theta,r,x,y))
    
    # plot 1 pt of trajectory here
    plt.scatter(x,y,s=7,color=(0.,.7,.8),linewidth=2)
    # save data  
    result[0:4, i] = array([M,E,Theta,r])       
    
plt.scatter(0.,0.,s=20,color=(1.,.2,.2),linewidth=4) # mark the star
plt.show()

        
# additional plot of irradiation
plt.figure(figsize=(9,6))
M_anom = result[0,:]/(2*pi)     # mean anomaly/(2*pi)
radius = result[3,:]
irrad  = (a/radius)**2          # irradiation factor (w.r.t. e=0 case)
# averaging over time; irrad*0+1 is an array holding constant function==1.0,
# and  sum(irrad*0+1)==180, since that's the length of array irrad;
# but I wrote it as a sum() to mimic the integral def of an average.
irr_int = sum(irrad)/N          #  or /sum(irrad*0+1)
irr_e = 1/sqrt(1-e*e)
diff = irr_int - irr_e
print(' average irradiation, integrated = <I>/I0 = ',irr_int)
print(' 1/sqrt(1-e^2) =',irr_e,'  integ-theor =',around(diff,14))
print(" Mean T = ",T,"K")
1
# output
# M_anom =  = time in orbital periods
txt = str(around(irr_int,9))
plt.plot(M_anom,radius/a,color=(.3,.4,1.),linewidth=2)
plt.plot(M_anom,0*radius+1,color=(1,0,1), linewidth=1)
plt.plot(M_anom,irrad,color=(1,.5,0), linewidth=2)
plt.title(" GJ 3512b, e = 0.4356",fontsize=14)
plt.text(.15,2.2,"irradiation = "+txt,color=(1,.5,0),fontsize=14)
plt.text(.01,.75,"r(t)",color=(.3,.4,1.),fontsize=14)
plt.xlabel("time [P]",fontsize=14 )
plt.grid()
plt.show()