# coding=utf-8
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime, timedelta

from nom_client.nom_client import NoMClient

"""
title: Creating a trajectory, basic tutorial
description: Create a trajectory exploring the tool's options.
keywords: trajectory, GEO, sun-synchronous, near-earth interplanetary
"""

"""
The 'Trajectory' tool has a large amount of input parameters. These parameters can be found on the model documentation page on the server web pages.
There are some preset sets of parameters below for the most common trajectories.

All trajectories require a start date, along with both a 'duration' and a 'mission duration'. 
The 'duration' is the length of the representative trajectory, typically few days. This is the duration the actual trajectory is computed for. 
It should be long enough to capture the complete spatial extent of the orbit. 
The 'mission duration' is the total length of the trajectory, typicaly years. 

The orbit data, calculated for the 'duration' will be used in the effects tools to calculate thier results and this data 
is scaled by the 'mission duration' to get the total mission values.
"""


def convert_cmjd_to_datetime(modified_julian_date):
    return datetime(1950, 1, 1, 0, 0, 0) + timedelta(days=modified_julian_date)


def plot(x_data, y_data, plot_types, line_types, xlabel, ylabel, labels, title):
    fig, ax = plt.subplots()

    for data_index in range(len(x_data)):
        if plot_types[data_index] == 'loglog':
            ax.loglog(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                       markersize=12.0,
                       label=labels[data_index])
        elif plot_types[data_index] == 'semilogy':
            ax.semilogy(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        elif plot_types[data_index] == 'semilogx':
            ax.semilogx(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                         markersize=12.0,
                         label=labels[data_index])
        else:
            plt.plot(x_data[data_index], y_data[data_index], linestyle=line_types[data_index], linewidth=2.0,
                     markersize=12.0,
                     label=labels[data_index])

    plt.title(title, fontsize=14)
    plt.grid(True)
    plt.legend(bbox_to_anchor=(0.95, 0.95), loc=1, borderaxespad=0.,
               fancybox=False, fontsize=14, shadow=False)

    ax.xaxis.set_major_formatter(mdates.DateFormatter('%d-%m-%y'))
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)

    fig.autofmt_xdate()
    plt.show()


"""
General orbit
"""
gen_params = {
    "startDate": "2023-01-03",
    "duration": 1,
    "missionDuration": 365,
    "orbitSpecificationCode": 205,
    "perigeeAltitude": 600,
    "apogeeAltitude": 6000,
    "timeSteps": [60, 240, 3600], # it is possible to control the timesteps above and below certain altitudes to keep the file size manageable
    "altitudeSteps": [2000, 20000]
}

"""
Geostationary orbit
"""
geo_params = {
    "startDate": "2023-01-03",
    "duration": 5,
    "missionDuration": 365,
    "orbitSpecificationCode": "GEO",
    "longitude": 123
}


"""
Near-Earth interplanetary orbit
"""
nei_params = {
    "startDate": "2023-01-03",
    "duration": 5,
    "missionDuration": 365,
    "orbitSpecificationCode": "NEI"
}

"""
Sun synchronous circular orbit
"""
hel_params = {
    "startDate": "2023-01-03",
    "duration": 5,
    "missionDuration": 365,
    "orbitSpecificationCode": "HEL",
    "circularAltitude": 2000,
    "localTimeAscendingNode": 12
}


"""
Initialise the NoM Client with a project name (optional) and the configuation data.
"""
nom_client = NoMClient(project_name="Create a trajectory")

"""
We will use the 'Trajectory' model to create our trajectory
"""
trajectory_model = nom_client.get_model(model_name='trajectory')

trajectory_model.set_params(params=gen_params)

trajectory_result = nom_client.run_model(model=trajectory_model)

print(trajectory_result)

trajectory_data = trajectory_result.get_model_result_by_name(result_name="trajectory")
trajectory_mjd = trajectory_data.get_variable_data(variable_name="MJD").squeeze()
# convert CNES MJD to datetime
trajectory_datetime = [convert_cmjd_to_datetime(x) for x in trajectory_mjd]

trajectory_altitude = trajectory_data.get_variable_data(variable_name="altitude").squeeze()

plot(x_data=[trajectory_datetime],
     y_data=[trajectory_altitude],
     plot_types=["plot"],
     line_types=["solid"],
     xlabel="MJD",
     ylabel="Altitude (km)",
     labels=["General orbit"],
     title=" ")