#   OrbitMapper.py   
#   M.Lampton  UCB SSL Nov 2018
#   Be sure to use the appropriate library environment

import matplotlib.pyplot as plt  
import numpy as np
import scipy as sp
from mpl_toolkits.basemap import Basemap
from itertools import chain

def draw_map(m, scale=0.2):
    # ref: VanderPlas, Python Data Science Handbook Chapter 4 2006
    m.shadedrelief(scale=scale)
    
    # lats and longs are returned as a dictionary
    lats = m.drawparallels(np.linspace(-90, 90, 13))
    lons = m.drawmeridians(np.linspace(-180, 180, 13))

    # keys contain the plt.Line2D instances
    lat_lines = chain(*(tup[1][0] for tup in lats.items()))
    lon_lines = chain(*(tup[1][0] for tup in lons.items()))
    all_lines = chain(lat_lines, lon_lines)
    
    # cycle through these lines and set the desired style
    for line in all_lines:
        line.set(linestyle='-', alpha=0.3, color='w')

def put360(deg):
    return np.fmod(deg+720., 360.)
    
#==========MAIN PROGRAM STARTS HERE==================      

INCL      = +38.0           # orbit inclination, constant, degrees
LONZERO   = +150.0          # orbit track longitude at t=0, degrees
ALTKM     = 500.0           # orbit altitude
GM        = 3.984E14        # gravity * EarthMass
RE        = 6.378E6         # Earth radius, meters
SIDEREAL  = 86164.0         # sidereal day, seconds
J2        = 1.082E-3        # Earth quadrupole moment
RORBIT    = RE+1000*ALTKM   # orbit radius, meters
PERIOD    = (6.283/np.sqrt(GM))*np.power(RORBIT, 1.5)
PHIDOT    = 360.0/PERIOD    # angular orbit speed, deg/sec
GSTDOT    = 360.0/SIDEREAL  # Earth rotation rate, deg/sec
RAANDOT   = -1.5*PHIDOT*(RE/RORBIT)**2*J2*np.cos(np.radians(INCL))
TSTART    = 1400.0          # start plotting at this time, sec
TSTOP     = 86000.0         # stop plotting at this time, sec. 1 day = 86400.
NPTS      = 5000            # how many samples over [TSTART,TSTOP]

# Draw the Earth map viewed from over lat=+15deg, elon=-120deg

fig = plt.figure(figsize=(8,8))
map = Basemap(projection='ortho', resolution=None, lat_0=15, lon_0=-120)
draw_map(map)

# Compute the orbit position at NPTS sample times

times = np.linspace(TSTART, TSTOP, NPTS)  # array of sample times, sec
raan = RAANDOT*times                      # array of RAANs; zero when t=0
incl = INCL                               # const INCL, degrees
phi = PHIDOT*times                        # array of PHIs; zero when t=0
cr = np.cos(np.radians(raan))
sr = np.sin(np.radians(raan))
ci = np.cos(np.radians(incl))
si = np.sin(np.radians(incl))
cp = np.cos(np.radians(phi))
sp = np.sin(np.radians(phi))
xeq = cr*cp - sr*ci*sp                    # equatorial coordinate x
yeq = sr*cp + cr*ci*sp                    # equatorial coordinate y
zeq = si*sp                               # equatorial coordinate z
ra = np.degrees(np.arctan2(yeq,xeq))      # right ascension of orbit position
dec = np.degrees(np.arcsin(zeq))          # declination of orbit position

# Transform samples to geographic coordinates and plot them.

lat = dec
lon = put360(ra + LONZERO - GSTDOT*times)
map.scatter(lon, lat, marker='.', color='red', latlon=True)

fig.text(0.1, 0.18, 'INCL = {:4.1f}'.format(INCL))
fig.text(0.1, 0.16, 'ALTKM = {:5.1f}'.format(ALTKM)) 
fig.text(0.1, 0.14, 'TSTART = '+str(int(TSTART)))
fig.text(0.1, 0.12, 'TSTOP = '+ str(int(TSTOP)))
fig.text(0.1, 0.10, 'OrbitMapper.py')
fig.savefig('OrbitMapper.pdf', format='pdf', dpi=1200)
plt.show()