# Compute and display the translunar gravitational potential
# Graphics based on the example at
# https://matplotlib.org/3.1.0/gallery/mplot3d/surface3d.html
#
# Eric Myers <EricMyers47@gmail.com> - 20 July 2019 - Version 1.0
########################################################################

import numpy as np
import math
import matplotlib.pyplot as plt
from matplotlib import cm           # color maps:
# Required for projection='3d'
from mpl_toolkits.mplot3d import Axes3D
# To customize the z axis
from matplotlib.ticker import LinearLocator, FormatStrFormatter

##
# Choose limits to plot:

Xlow = -10
Xhigh = +70
Yhigh = (Xhigh-Xlow)/4.0
Ylow = -Yhigh

##
# Physical Parameters:

M_earth = 5.9722E+24  # kg          - mass of the Earth
R_earth = 6.371E+6    # m           - radius of the Earth
M_moon = 7.342E+22    # kg          - mass of the Moon
R_moon = 1.7371E+6    # m           - radius of the Moon
d_EM = 3.84399E+8     # m           - Earth/Moon distance
G_Newton = 6.67E-11   # N*m^2/kg^2  - Newton's gravitation constant

##
# Vgrav(X, Y) takes position coordinates x and y, as meshgrids,
# in units of Earth radii, and returns an np.array (meshgrid) of values
# representing the gravitational potential (J/kg) at that position.

def Vgrav(X, Y):
    Xm = X*R_earth      # x position in meters
    Ym = Y*R_earth      # y position in meters

    # Distance from the Earth
    R1 = np.sqrt(Xm**2 + Ym**2)
    # If R1 < R_earth then R1=R_earth
    R0 = R_earth
    R1 = np.maximum(R1, R_earth)

    # Distance from the Moon
    R2 = np.sqrt((Xm-d_EM)**2 + Ym**2)
    # If R2 < R_moon then R2=R_moon
    R2 = R2*(R2 > R_moon) + R_moon*(R2 <= R_moon)

    V1 = -G_Newton*M_earth / R1
    V2 = -G_Newton*M_moon / R2
    V =  V1 + V2
    return V

##
# Layout the grid:

X = np.arange(Xlow, Xhigh, 0.5)
Y = np.arange(Ylow,Yhigh, 0.5)
X, Y = np.meshgrid(X, Y)

# Compute the gravitational potential at all points

Vg = Vgrav(X,Y)
print("Vg ranges from ",Vg.min(),' to ', Vg.max() )

# Vertical of the plot uses a logarithmic scale to make it clearer
Z = -np.log10(-Vg)

# Uncomment these to cut off the potential below that of the Moon's surface
"""
Z_moon = -math.log10(G_Newton*M_moon / R_moon)
print("Moon's surface potential (logarithm) is at ",Z_moon)
Z = np.maximum(Z,Z_moon)
"""

# Shift the bottom of the graph to zero
Z = Z - Z.min()
print("Z now ranges from ",Z.min(),' to ', Z.max() )


##
# Make the plot

# Expect a wide figure
fig = plt.figure(figsize=(5.,2.0),dpi=300)

# This is a 3D plot
ax = fig.gca(projection='3d')

# No margins, use all the space
plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0)

#"""
# Plot the surface.
surf = ax.plot_surface(X, Y, Z,
                       cmap=cm.terrain,   # terrain, prism, flag, gist_earth,
                       #vmin=Z_moon,
                       rcount=200, ccount=200,
                       linewidth=0.0, antialiased=True)

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.50, aspect=20)
#"""

# Add contour plot
ax.contour(X,Y,Z,20)

# Customize the z axis.
ax.set_zlim( Z.min(), Z.max() )
ax.zaxis.set_major_locator(LinearLocator(2))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))


#-- Scale sizes of all 3 axes by different amounts. --#
#See https://stackoverflow.com/questions/30223161/matplotlib-mplot3d-how-to-increase-the-size-of-an-axis-stretch-in-a-3d-plo

x_scale=8
y_scale=4
z_scale=6

scale=np.diag([x_scale, y_scale, z_scale, 1.0])
scale=scale*(1.0/scale.max())
scale[3,3]=1.0

def short_proj():
   return np.dot(Axes3D.get_proj(ax), scale)

ax.get_proj=short_proj

#-- End of axis scaling --#

# Label axes.
ax.set_xlabel(r'x ($R_\oplus$)')
ax.set_ylabel(R'y ($R_\oplus$)')
ax.set_zlabel('Log Potential ')

#plt.savefig("EarthMoonXX.png")
# It's better to view the image, rotate it, and save it manually
plt.show()

##EOF##