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

# ## License
# ```
# SpaceX/Starlink satellite constellation: orbit visualization.
# Copyright (C) 2019+  Benjamin Winkel (bwinkel@mpifr.de)
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# 
# Note: parts of this software were adapted from Cees Bassa (ASTRON);
#       see https://github.com/cbassa/satellite_analysis
# ```

# In[1]:


import datetime
import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib
from matplotlib import animation, rc

from astropy import time
from astropy import units as u
from astropy import constants as const
from astropy.coordinates import EarthLocation
from pycraf import satellite

rc('animation', html='html5')
matplotlib.rcParams['animation.embed_limit'] = 2 ** 128


# In[2]:


print(const.GM_earth, const.R_earth, sep='\n')


# In[3]:


def satellite_mean_motion(altitude, mu=const.GM_earth, r_earth=const.R_earth):
    '''
    Compute mean motion of satellite at altitude in Earth's gravitational field.
    
    See https://en.wikipedia.org/wiki/Mean_motion#Formulae
    '''
    no = np.sqrt(4.0 * np.pi ** 2 * (altitude + r_earth) ** 3 / mu).to(u.day)
    return 1 / no


# In[4]:


satellite_mean_motion(500 * u.km)


# In[5]:


def tle_from_orbital_parameters(
        sat_name, sat_nr, epoch, inclination, raan, mean_anomaly, mean_motion
        ):
    '''
    Generate TLE strings from orbital parameters.
    
    Note: epoch has a very strange format: first two digits are the year, next three
    digits are the day from beginning of year, then fraction of a day is given, e.g.
    20180.25 would be 2020, day 180, 6 hours (UT?)
    '''
    
    # Note: RAAN = right ascention (or longitude) of ascending node
    
    def checksum(line):
        s = 0
        for c in line[:-1]:
            if c.isdigit():
                s += int(c)
            if c == "-":
                s += 1
        return '{:s}{:1d}'.format(line[:-1], s % 10)
    
    tle0 = sat_name
    tle1 = checksum(
        '1 {:05d}U 20001A   {:14.8f}  .00000000  00000-0  50000-4 '
        '0    0X'.format(sat_nr, epoch)
        )
    tle2 = checksum(
        '2 {:05d} {:8.4f} {:8.4f} 0001000   0.0000 {:8.4f} '
        '{:11.8f}    0X'.format(
            sat_nr, inclination.to_value(u.deg), raan.to_value(u.deg),
            mean_anomaly.to_value(u.deg), mean_motion.to_value(1 / u.day)
        ))
    
    return '\n'.join([tle0, tle1, tle2])


# In[6]:


my_sat_tle = tle_from_orbital_parameters(
    'TEST', 99, 19050.1, 10 * u.deg, 20 * u.deg, 45 * u.deg,
    satellite_mean_motion(500 * u.km)
    )
print(my_sat_tle)


# In[7]:


my_sat = satellite.get_sat(my_sat_tle)
my_sat


# Starlink constellation

# In[8]:


altitudes = np.array([550, 1110, 1130, 1275, 1325, 345.6, 340.8, 335.9]) * u.km
inclinations = np.array([53.0, 53.8, 74.0, 81.0, 70.0, 53.0, 48.0, 42.0]) * u.deg
nplanes = np.array([72, 32, 8, 5, 6, 2547, 2478, 2493])
sats_per_plane = np.array([22, 50, 50, 75, 75, 1, 1, 1])


# In[9]:


def create_constellation(altitudes, inclinations, nplanes, sats_per_plane):
    
    my_sat_tles = []
    sat_nr = 80000
    for alt, inc, n, s in zip(
            altitudes, inclinations, nplanes, sats_per_plane
            ):
        
        if s == 1:
            # random placement for lower orbits
            mas = np.random.uniform(0, 360, n) * u.deg
            raans = np.random.uniform(0, 360, n) * u.deg
        else:
            mas = np.linspace(0.0, 360.0, s, endpoint=False) * u.deg
            mas += np.random.uniform(0, 360, 1) * u.deg
            raans = np.linspace(0.0, 360.0, n, endpoint=False) * u.deg
            mas, raans = np.meshgrid(mas, raans)
            mas, raans = mas.flatten(), raans.flatten()
        
        mm = satellite_mean_motion(alt)
        for ma, raan in zip(mas, raans):
            my_sat_tles.append(
                tle_from_orbital_parameters(
                    'TEST {:d}'.format(sat_nr), sat_nr, 19050.1,
                    inc, raan, ma, mm
                    ))
            sat_nr += 1
    
    return my_sat_tles


# In[10]:


my_sat_tles = create_constellation(altitudes, inclinations, nplanes, sats_per_plane)


# In[11]:


len(my_sat_tles)


# In[12]:


my_sats = [satellite.get_sat(sat_tle) for sat_tle in my_sat_tles]


# In[13]:


def _propagate(sat, dt):
    '''
    True equator mean equinox (TEME) position from `sgp4` at given time.
    Parameters
    ----------
    sat : `sgp4.io.Satellite` instance
        Satellite object filled from TLE
    dt : `~datetime.datetime`
        Time
    Returns
    -------
    xs, ys, zs : float
        TEME (=True equator mean equinox) position of satellite [km]
    '''

    # pos [km], vel [km/s]
    position, velocity = sat.propagate(
        dt.year, dt.month, dt.day,
        dt.hour, dt.minute, dt.second + dt.microsecond / 1e6
        )

    if position is None:
        raise ValueError('Satellite propagation error')

    return position


vec_propagate = np.vectorize(
    _propagate, excluded=['sat'], otypes=[np.float64] * 3
    )


# In[14]:


start_time = time.Time(datetime.datetime.utcnow(), 'utc')
print('Start time:', start_time)
td = time.TimeDelta(np.arange(0, 600, 1), format='sec')
times = start_time + td  # 10 min in steps of 1 s


# In[15]:


pos_eci = np.array([vec_propagate(sat[1], times.datetime) for sat in my_sats])
# pos_eci


# In[16]:


plim = 8000

fig = plt.figure(figsize=(12, 9))
ax = Axes3D(fig)
ax.view_init(azim=60, elev=30)
ax.set_xlim((-plim, plim))
ax.set_ylim((-plim, plim))
ax.set_zlim((-plim, plim))
ax.set_aspect('equal')
ax.set_xlabel('x [km]')
ax.set_ylabel('y [km]')
ax.set_zlabel('z [km]')

rads = np.sqrt(pos_eci[:, 0, 0] ** 2 + pos_eci[:, 1, 0] ** 2 + pos_eci[:, 2, 0] ** 2)
points = ax.scatter(
    pos_eci[:, 0, 0], pos_eci[:, 1, 0], pos_eci[:, 2, 0],
    c=rads, s=1, vmin=rads.min(), vmax=rads.max(), marker='o'
    )
title = ax.set_title('{:%y/%m/%d %H:%M:%S}'.format(times[0].datetime), loc='center', fontsize=20)

def init():
    points._offsets3d = (pos_eci[:, 0, 0], pos_eci[:, 1, 0], pos_eci[:, 2, 0])
    title.set_text('{:%y/%m/%d %H:%M:%S}'.format(times[0].datetime))
    return points, title

def animate(i):
    points._offsets3d = (pos_eci[:, 0, i], pos_eci[:, 1, i], pos_eci[:, 2, i])
    title.set_text('{:%y/%m/%d %H:%M:%S}'.format(times[i].datetime))
    return points, title

# call the animator. blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(
    fig, animate, init_func=init, frames=pos_eci.shape[2], interval=20, blit=True
    )

# this takes a while!
plt.close(anim._fig)
anim


# In[17]:


obs_loc_parkes = EarthLocation(148.25738, -32.9933, 414.8)
obs_loc_effbg = EarthLocation(6.88375, 50.525, 366.)

# create a SatelliteObserver instance
sat_obs_parkes = satellite.SatelliteObserver(obs_loc_parkes)
sat_obs_effbg = satellite.SatelliteObserver(obs_loc_effbg)
pos_horiz_parkes = np.array([
    sat_obs_parkes.azel_from_sat(sat[1], times)
    for sat in my_sats
    ])
pos_horiz_effbg = np.array([
    sat_obs_effbg.azel_from_sat(sat[1], times)
    for sat in my_sats
    ])


# In[18]:


vmin, vmax = np.log10(100), np.log10(10000)

fig = plt.figure(figsize=(12, 7))
ax1 = fig.add_axes((0.1, 0.5, 0.8, 0.35))
ax2 = fig.add_axes((0.1, 0.1, 0.8, 0.35))
cax = fig.add_axes((0.91, 0.2, 0.02, 0.5))
ax2.set_xlabel('Azimuth [deg]')
ax1.set_ylabel('Elevation [deg]')
for ax in [ax1, ax2]:
    ax.set_xlim((-180, 180))
    ax.set_ylim((0, 90))
    ax.set_xticks(range(-150, 180, 30))
    ax.set_yticks(range(0, 91, 30))
    ax.set_aspect('equal')
    ax.grid()

points1 = ax1.scatter(
    pos_horiz_parkes[:, 0, 0], pos_horiz_parkes[:, 1, 0],
    c=np.log10(pos_horiz_parkes[:, 2, 0]),
    cmap='viridis', vmin=vmin, vmax=vmax,
    )
points2 = ax2.scatter(
    pos_horiz_effbg[:, 0, 0], pos_horiz_effbg[:, 1, 0],
    c=np.log10(pos_horiz_effbg[:, 2, 0]),
    cmap='viridis', vmin=vmin, vmax=vmax,
    )
cbar = fig.colorbar(points1, cax=cax)
cbar.set_label('Distance (km)')
cbar.set_ticks([2, 3, 4])
cbar.set_ticklabels([100, 1000, 10000])

ax1.text(-170, 75, 'Parkes 64-m', fontsize=16)
ax2.text(-170, 75, 'Effelsberg 100-m', fontsize=16)
title = ax1.text(
    174, 75, '{:%y/%m/%d %H:%M:%S}'.format(times[0].datetime),
    fontsize=15, ha='right'
    )

def init():
    points1.set_offsets(pos_horiz_parkes[:, 0:2, 0])
    points1.set_array(np.log10(pos_horiz_parkes[:, 2, 0]))
    points2.set_offsets(pos_horiz_effbg[:, 0:2, 0])
    points2.set_array(np.log10(pos_horiz_effbg[:, 2, 0]))
    title.set_text('{:%y/%m/%d %H:%M:%S}'.format(times[0].datetime))
    return points, title

def animate(i):
    points1.set_offsets(pos_horiz_parkes[:, 0:2, i])
    points1.set_array(np.log10(pos_horiz_parkes[:, 2, i]))
    points2.set_offsets(pos_horiz_effbg[:, 0:2, i])
    points2.set_array(np.log10(pos_horiz_effbg[:, 2, i]))
    title.set_text('{:%y/%m/%d %H:%M:%S}'.format(times[i].datetime))
    return points, title

anim = animation.FuncAnimation(
    fig, animate, init_func=init, frames=pos_horiz_parkes.shape[2], interval=20, blit=True
    )

# this takes a while!
plt.close(anim._fig)
anim


# In[ ]: