#!/usr/bin/env python
#coding: utf-8 

import numpy as np

Mearth, Msun = 5.97e24, 1.98e30        # <-- these are realistic
#Mearth, Msun = 1.97e29, 1.98e30         # <-- unrealistic, but better for visualisation
Xearth = 149e9
Xsun = - Xearth * (Mearth/Msun)
kappa = 6.67e-11

year = 365*24*60*60

# Generate the points in the space  
xaxis = np.linspace(-300e9, 300e9, 2000)
yaxis = np.linspace(-300e9, 300e9, 2000)
x,y = np.meshgrid(xaxis,yaxis)
print ("Allocated two big arrays with size:", x.shape)


# Centrifugal force
rC = (x**2+y**2)**.5
FC = (2*np.pi / year)**2 * rC
FCx = FC * (x/rC)
FCy = FC * (y/rC)

# Gravity of Sun
rS = ((x-Xsun)**2+y**2)**.5
FS = -kappa * Msun / (rS**2)
FSx = FS * (x-Xsun)/rS
FSy = FS * y/rS

# Gravity of Earth
rE = ((x-Xearth)**2+y**2)**.5
FE = -kappa * Mearth / (rE**2)
FEx = FE * (x-Xearth)/rE
FEy = FE * y/rE

# All three forces together
Fax = FCx+FSx+FEx
Fay = FCy+FSy+FEy
Fa = (Fax**2+Fay**2)**.5

import matplotlib.pyplot as plt
# Plot the magnitude of the force:
plt.contourf(xaxis, yaxis, np.log10((Fax**2+Fay**2)**.5), levels=np.linspace(-4, 0, 200), extend='both')
# ... and its direction:
plt.quiver(x[::30, ::30],y[::30, ::30],Fax[::30, ::30]/Fa[::30, ::30], Fay[::30, ::30]/Fa[::30, ::30])

plt.show()