class PlanetaryObject():
    """
    A simple class used to store pertinant information about the plantary object
    """
    def __init__(self, date, L, e, SMA, i, peri, asc, r, v, anom, fp, mu):
        self.date = date   # Event Date
        self.L = L         # Longitude
        self.e = e         # Eccentricity
        self.SMA = SMA     # SMA
        self.i = i         # Inclination
        self.peri = peri   # Longitude of Perihelion
        self.asc = asc     # Longitude of Ascending Node
        self.r = r         # Radius
        self.v = v         # Velocity
        self.anom = anom   # True Anomaly
        self.fp = fp       # Flight Path Angle
        self.mu = mu       # Gravitation parameter

earth = PlanetaryObject(
    datetime.date(2021, 3, 22),
    181.44,     # Longitude
    0.0167,     # Eccentricity
    149598020,  # SMA
    0,          # Inclination
    102.958,    # Longitude of Perihelion
    0,          # Longitude of Ascending Node
    14905909.7, # Radius
    29.89,      # Velocity
    78.48,      # True Anomaly
    0.9348,     # Flight Path Angle
    398600.4    # Gravitation parameter
)
mars = PlanetaryObject(
    datetime.date(2021, 10, 8),
    333.22,     # Longitude
    0.0934,     # Eccentricity
    227939133,  # SMA
    1.849,      # Inclination
    336.093,    # Longitude of Perihelion
    49.572,     # Longitude of Ascending Node
    206671197,  # Radius
    26.94,      # Velocity
    357.128,    # True Anomaly
    -0.2452,    # Flight Path Angle
    42828.3     # Gravitation parameter
)

"""
Explanation: SIE 552 HW #3
This is a space mission design assignment for interplantary travel.  The homework was to create a Type-1 trajectory from Earth to Mars.  Departure from earth is set to be March 22, 2021 with an anticipated arrival on Mars at October 8, 2021.  We're also given the following details about Earth and Mars:
When designing the transfer ellipse, the overall mission has to be considered first.  We know that we want to create an interplanetary travel mission between Earth and Mars.  In order to plan this, we have to know more about those planets:


Earth at Launch: (March 22, 2021)

Longitude = 181.44 deg 
Eccentricity = 0.0167
SMA =  149,598,020 Km
Inclination = 0 deg 
Longitude of the Perihelion = 102.958 deg 
Longitude of the Ascending Node = 0 deg 
Radius = 149,905,909.7 km 
Velocity = 29.89 km/sec
True Anomaly = 78.48 deg 
Flight Path Angle = 0.9348 deg 



Mars at Arrival: (October 8, 2021)

Longitude = 333.22 deg 
Eccentricity = 0.0934
SMA =  227,939,133 Km
Inclination = 1.849 deg 
Longitude of the Perihelion = 336.093 deg 
Longitude of the Ascending Node = 49.572 deg 
Radius = 206,671,197 km
Velocity = 26.94 km/sec
True Anomaly = 357.128 deg 
Flight Path Angle = - 0.2452 deg



We'll store these as PlanetaryObejct's in python to help us later.
End of explanation
"""


mu_sun = 132712439935.5
        
def eccentricity(r_1, r_2, theta_1, theta_2):
    """
    Calculates the eccentricity of the transfer ellipse.  This is calculated through 
    the following equation:
    
    .. math::
        \frac {r_2 - r_1} {r_1 * \cos{\theta_1} - r_2 * \cos{\theta_2}}
    
    :param r_1: radius of the departing planetary object
    :param r_2: radius of the arriving planetary object
    :param theta_1: True anomaly of the departing planetary object in degrees
    :param theta_2: True anomaly of the arriving planetary object in degrees
    """
    return (r_2 - r_1) / ((r_1 * math.cos(math.radians(theta_1))) - (r_2 * math.cos(math.radians(theta_2))))

def periapsis_radius(r, e, theta):
    """
    Calculates the periapsis radius of the transfer ellipse.  This is calculated 
    using the following equation:
    
    .. math::
        \frac {r_1 * [1 + e \cos{\theta]}} {1 + e}
    
    :param r: radius of the departing planetary object
    :param e: eccentricity of the transfer ellipse
    """
    return (r * (1 + e * math.cos(math.radians(theta)))) / (1 + e)

def semimajor_axis(r=None, r_a=None, r_p=None, mu=None, V=None, e=None):
    """
    Calculates the semi-major axis of the transfer ellipse.  This is calculated 
    using one of the following equations:
    
    .. math::
        \frac {r_a + r_p} {2}
        
        \frac {\mu r} {2 \mu - V^2 r}
        
        \frac {r_p} {1 - e}
        
        \frac {r_a} {1 + e}
    
    :param r: general radius of the elliptical orbit
    :param r_a: Radius of apoapsis
    :param r_p: Radius of periapsis
    :param mu: gravitation parameter
    :param V: Velocity of the orbiting object
    :param e: Eccentricity of the elliptical orbit
    """
    if r_a != None and r_p != None:
        return (r_a + r_p) / 2
    if mu != None and r !=None and V != None:
        return (mu * r) / (2 * mu - V ** 2 * r)
    if r_p != None and e != None:
        return r_p / (1 - e)
    if r_a != None and e != None:
        return r_a / (1 + e)
    
    # If we reach this point, then the passed in arguments doesn't match
    #    any equations we have defined.  Raise an Error
    raise TypeError("Invalid arguments!")
    

def time_since_periapsis(e, n, theta=None, E=None):
    """
    Calculates the time since the periapsis.  This is calculated using the
    following equation:
    
    .. math::
        \frac {E - e \sin{E}} {n}
        
    If E, isn't defined, it will be calculated using the param theta and 
    the following equation:
    
    ..math:: 
        
        \cos {E} = \frac {e + \cos{\theta}} {1 + e \cos{\theta}}
    
    :param e: eccentricity of the transfer ellipse
    :param n: mean motion
    :param theta: degrees to periapsis
    :param E: eccentric anomaly in radians
    """
    if theta == None and E == None:
        raise TypeError("theta or E MUST be defined")
    if theta != None and E != None:
        raise TypeError("theta OR E must be defined.  Not both")
        
    if E == None:
        cos_E = (e + math.cos(math.radians(theta))) / (1 + e * math.cos(math.radians(theta)))
        E = math.acos(cos_E)
        
    return (E - e * math.sin(E)) / n

def mean_motion(mu, a):
    """
    Calculates the mean motion of an elliptical orbit.  This is calculated 
    using the following equation:
    
    .. math::
        \sqrt{\frac{\mu} {a^3}}
    
    :param mu: gravitation parameter (Mass * Gravitation constant)
    :param a: semimajor axis
    """
    
    return math.sqrt(mu / a ** 3)

def velocity(mu, r, a):
    """
    Calculates the Velocity (V) of an object based on the elliptical orbit.  
    This is calculated using the following equation:
    
    .. math::
        \sqrt{\frac{2 * \mu} {r} - \frac{\mu} {a}}
        
    :param mu: gravitation parameter (Mass * Gravition constant)
    :param a: semimajor axis
    """
    return math.sqrt(2 * mu / r - mu / a)

def flight_path_angle(e, theta):
    """
    Calculates the Flight Path Angle (γ).  This is calculated using
    the following equation:
        
    .. math::
        \tan{γ} = {\frac{e * \sin{\theta}}{1 + 3 * \cos{\theta}}
        
    :param e: eccentricity of the elliptical orbit
    :param theta: 
    """
    tan_y = (e * math.sin(math.radians(theta))) / (1 + e * math.cos(math.radians(theta)))
    return math.atan(tan_y)

def inclination(Omega, L_s, L_t, i):
    a = math.radians(Omega + 180 - L_s)
    b = math.radians(L_t - (180 + Omega))
    alpha = math.radians(180 - i)
    cos_c = math.cos(a) * math.cos(b) + math.sin(a) * math.sin(b) * math.cos(alpha)
    c = math.acos(cos_c)
    sin_i_t = (math.sin(alpha) * math.sin(b)) / math.sin(c)
    return math.asin(sin_i_t)

"""
Explanation: There are also a few fundamental equations we need to know.  These are captured below as python functions.
End of explanation
"""


def transfer_ellipse(start_planet, end_planet, return_trials=False):
    time_of_flight = end_planet.date - start_planet.date
    time_of_flight = time_of_flight.days
    
    longs = []
    tofs = []
    
    line_of_apisides = 180  # trial start
    tof = 9999999999 # large number to get us started
    
    while  tof / 3600 / 24 > time_of_flight:
        true_anom = line_of_apisides + (end_planet.L - start_planet.L)
        longs.append((line_of_apisides, true_anom))
        e = eccentricity(start_planet.r, end_planet.r, line_of_apisides, true_anom)
        r_p = periapsis_radius(start_planet.r, e, line_of_apisides)
        a = semimajor_axis(r_p=r_p, e=e)
        
        n = mean_motion(mu_sun, a)
        
        peri_to_start = time_since_periapsis(e, n, theta=line_of_apisides)
        end_to_peri = time_since_periapsis(e, n, theta=true_anom)
        
        tof = peri_to_start - end_to_peri
        tofs.append(tof / 3600 / 24)
        line_of_apisides += 1
        
    # Calculate the Relative Velocities
    V_start = velocity(mu_sun, start_planet.r, a)
    V_end = velocity(mu_sun, end_planet.r, a)
    
    y_start = flight_path_angle(e, line_of_apisides)
    y_end = flight_path_angle(e, true_anom)
    
    r_dict = {
        'line_of_apisides': line_of_apisides - 1,  # subtract the 1 we added during the loop
        'true_anom': true_anom,
        'eccentricity': e,
        'SMA': a,
        'time_of_flight': tof,
        'V_start': V_start,
        'V_end': V_end,
        'y_start': math.degrees(y_start),
        'y_end': math.degrees(y_end)
    }
        
        
    if return_trials:
        r_dict.update({'runs':{'longs': longs, 'tofs':tofs}})
        
    return r_dict

tf = transfer_ellipse(earth, mars, return_trials=True)
tf

"""
Explanation: We'll split this problem up into 3 different sections:

Transfer Ellipse
Departure Trajectory
Arrival Trajectory

This approach is called the "patched conics approach".  
Designing the Transfer Ellipse
In order to create the transfer ellipse between Earth and Mars, we're going to do the following:

Determine the required # of days based off departure and arrival date
Calculate the angular difference between Earth and Mars
Design the Eccentricity, Periapsis Axis and Semi-Major Axis using an iterative approach
Finally, calculating the respective velocities and flight path angles for each planetary object

To find the required number of days to travel to Mars, we simply subtract the launch and arrival dates:
End of explanation
"""


i_t = inclination(mars.asc, earth.L, mars.L, mars.i)
print("i_t = {:.2f}°".format(math.degrees(i_t)))

"""
Explanation: The Departure Trajectory
The departure trajectory entails the following:

Determining the plane change
Calculating $V_{HE}$ and $C3$

$V_{HE}$ is the vector required for the spacecraft to leave the Earth's gravitational pull, and enter into the calculated transfer ellipse above.
We'll use the function defined above to find the inclination of the transfer plane $i_t$:
End of explanation
"""


cos_alpha_2 = math.cos(i_t) * math.cos(earth.fp + abs(y_ne))
alpha_2 = math.acos(cos_alpha_2)
C3 = earth.v ** 2 + V_ne ** 2 - 2 * earth.v * V_ne * math.cos(alpha_2)
V_he = math.sqrt(C3)
print("C3 = {:.2f} km^2/s^2; V_he = {:.2f} km/s".format(C3, V_he))

"""
Explanation: Using $i_t$, we can now determine the $V_{HE}$ and $C3$ of the departing trajectory:
End of explanation
"""


