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

# This notebook analyzes the perturbation effects on an orbit due to the Earth's gravity field.

# Parameters:

# In[1]:


import numpy as np
# Keplerian parameters
a = 6578137.0
ex = 0.0
ey = 0.0
i = float(np.deg2rad(90.0))
raan = 0.0
alpha = 0.0  # true latitude argument

import datetime
epoch = datetime.datetime(2020, 1, 1)

n_harmonics = 8  # Number of spheric harmonics to use

dt = 30.0  # s
duration = 5300.0

position_tolerance = 1.0 
prop_min_step = 0.001 # s
prop_max_step = 300.0 # s
prop_position_error = 10.0 # m


# Firing up a JVM for Orekit, and downloading and importing the Orekit data ZIP

# In[2]:


import orekit
orekit.initVM()

from orekit.pyhelpers import download_orekit_data_curdir, setup_orekit_curdir
try:
    download_orekit_data_curdir()
except:
    print('Download failed')
setup_orekit_curdir()


# Setting up models (frames, timescales)

# In[3]:


from org.orekit.frames import FramesFactory
eme2000 = FramesFactory.getEME2000()

from org.orekit.utils import IERSConventions
itrf = FramesFactory.getITRF(IERSConventions.IERS_2010, True)

from org.orekit.models.earth import ReferenceEllipsoid
wgs84_ellipsoid = ReferenceEllipsoid.getWgs84(itrf)

from org.orekit.time import TimeScalesFactory
utc = TimeScalesFactory.getUTC()


# Setting up the initial orbit in Earth-inertial frame.

# In[4]:


from org.orekit.orbits import CircularOrbit, PositionAngle, OrbitType
from orekit.pyhelpers import datetime_to_absolutedate
from org.orekit.utils import Constants as orekit_constants
date_init = datetime_to_absolutedate(epoch)
date_end = date_init.shiftedBy(duration)
orbit = CircularOrbit(a, ex, ey, i, raan, alpha,
                       PositionAngle.TRUE, 
                       eme2000, 
                       date_init,
                       orekit_constants.EIGEN5C_EARTH_MU)

satellite_mass = 100.0 # kg
from org.orekit.propagation import SpacecraftState
initialState = SpacecraftState(orbit, satellite_mass) 


# Setting up numerical propagator and gravity force model

# In[5]:


from org.orekit.propagation.numerical import NumericalPropagator
tol = NumericalPropagator.tolerances(position_tolerance, orbit, OrbitType.CIRCULAR)

from orekit import JArray_double
from org.hipparchus.ode.nonstiff import DormandPrince853Integrator
integrator = DormandPrince853Integrator(prop_min_step, prop_max_step, 
    JArray_double.cast_(tol[0]),  # Double array of doubles needs to be casted in Python
    JArray_double.cast_(tol[1]))
integrator.setInitialStepSize(prop_max_step)

propagator = NumericalPropagator(integrator)
propagator.setOrbitType(OrbitType.CIRCULAR)
propagator.setInitialState(initialState)

# Earth gravity field with degree 64 and order 64
from org.orekit.forces.gravity.potential import GravityFieldFactory
gravity_provider = GravityFieldFactory.getConstantNormalizedProvider(n_harmonics, n_harmonics)
from org.orekit.forces.gravity import HolmesFeatherstoneAttractionModel
gravity_attraction_model = HolmesFeatherstoneAttractionModel(itrf, gravity_provider)
propagator.addForceModel(gravity_attraction_model)


# In[6]:


from org.orekit.propagation.sampling import PythonOrekitFixedStepHandler
from orekit.pyhelpers import absolutedate_to_datetime
from org.orekit.utils import IERSConventions
from org.orekit.frames import FramesFactory
from org.orekit.models.earth import ReferenceEllipsoid
from java.util import ArrayList
import numpy as np
import pandas as pd

class mystephandler(PythonOrekitFixedStepHandler):    
    itrf = FramesFactory.getITRF(IERSConventions.IERS_2010, True)
    wgs84_ellipsoid = ReferenceEllipsoid.getWgs84(itrf)
    df = pd.DataFrame(columns=['a', 'e', 'h', 'h_sphere'])
    intermediate_states = ArrayList()
    
    def init(self, s0, t, step):
        pass
        
    '''
    When a new state is received:
      - Adding the state to the ArrayList
      - Computing and saving the latitude/longitude coordinates
    '''
    def handleStep(self, currentState, isLast):
        self.intermediate_states.add(currentState)
        position = currentState.getPVCoordinates(self.itrf).getPosition()
        geodetic_point = wgs84_ellipsoid.transform(position,
                                                   self.itrf,
                                                   currentState.getDate())
        self.df.loc[absolutedate_to_datetime(currentState.getDate())] = [currentState.getA(),
                                                                        currentState.getE(),
                                                                         geodetic_point.getAltitude(),
                                                                        position.getNorm() - 6378137.0]
        
    def getDf(self):
        return self.df
    
    def getIntermediateStates(self):
        return self.intermediate_states
    
handler = mystephandler()
propagator.resetInitialState(initialState)
propagator.setMasterMode(dt, handler)
finalState = propagator.propagate(date_end)
df = handler.getDf()
display(df)


# In[7]:


from plotly.subplots import make_subplots
import plotly.graph_objects as go

fig = make_subplots(rows=4, cols=1, 
                    shared_xaxes=True, 
                    vertical_spacing=0.05,
                    subplot_titles=("Semi-major axis [km]", "Eccentricity", 
                                    "Altitude above WGS84 ellipsoid [km]", "Altitude above 6378.1km sphere [km] (without perturbations, this should be constant)"))

# Add traces
fig.add_trace(go.Scatter(x=df.index, y=1e-3*df['a'],
                    mode='lines'),
              row=1, col=1)
fig.add_trace(go.Scatter(x=df.index, y=df['e'],
                    mode='lines'),
              row=2, col=1)
fig.add_trace(go.Scatter(x=df.index, y=1e-3*df['h'],
                    mode='lines'),
              row=3, col=1)
fig.add_trace(go.Scatter(x=df.index, y=1e-3*df['h_sphere'],
                    mode='lines'),
              row=4, col=1)

fig.update_layout(height=1000, width=1000,
                  title_text="Gravity field perturbations on 200km orbit at 90° inclination (8*8 spherical harmonics)",
                  showlegend=False)
fig.show()
fig.write_image("perturbated_orbit.png", scale=2)


# In[ ]: