#!/usr/bin/env python
# - * - coding: utf-8 - * -
"""
Created on Sun May 12 20:08:52 2019

@author: MAC
"""

import numpy as np
from numpy import linalg as LA


   
# Constants 
T_SIDE = 86164.10   # sidereal day (s) 
MU = 3.9860e14      # earth standard gravitational parameter (=G * M_earth) (m3.s-2)
R_T = 6378e3        # earth radius (m)
G0 = 9.81 # standard gravity (m.s-2)
R = 287 # specific gas constant (J.Kg-1.K-1)

   
def initial_vel(latitude,retrograde=False):
    """
    Initial velocity (m.s-1) depending on the launch point latitude (°)
    
    :param latitude: the launch point latitude  in °
    :return: the initial velocity in m.s-1
    """
    
    initial_velocity = (2  *  np.pi / T_SIDE)  *  R_T  *  np.cos(np.deg2rad(np.abs(latitude)))
    if retrograde:
        return -initial_velocity
    else:
        return initial_velocity
    
def final_target_vel(h, a):
    """
    
    .. math:: \\sqrt {2 \\cdot MU \\cdot (\\frac {1}{R_T + h} - \\frac {1}{2  *  a})}
    
    :param h: injection altitude (m)
    :type h: float
    :param a: orbit semi major axis (m)
    :type a: float
    :return: orbit absolute injection velocity (m.s-1)
    :rtype: float
    """

    v_square = 2  *  MU  *  ((1 / (R_T + h)) - (1 / (2  *  a)))
    print('{} > {}'.format((R_T + h), (2  *  a)))
    assert v_square >= 0, 'error: square root only on positiv values'
    return np.sqrt(v_square)
   
def staging_optim(isp1,isp2,eps1,eps2,Vc,Mp):

    """Optimise mass loading in the case of two stage rocket
    
    .. math:: structural_coefficient  = \\frac {strcture_mass} {structure_mass + propellant_mass}
    
    :param isp1: specific impulse of the first stage (s)
    :type isp1: float
    :param isp2: specific impulse of the second stage (s)
    :type isp2: float
    :param eps1: structural coefficient of the first stage
    :type eps1: float
    :param eps2: structural coefficient of the second stage
    :type eps2: float
    :param Vc: target critical velocity (m.s-1)
    :type Vc: float
    :param Mp: payload mass (Kg)
    :type Mp: float
    :return: mass of the first and second stage (Kg)
    :rtype: float
    """
    
    omega_tab=np.linspace(0.001,0.9,100)  # array of omega to test
    qp_tab=[]  
    for omega in omega_tab:
        q2 = omega + eps1  *  (1 - omega)
        q2 = np.power(q2,isp1 / isp2)  *  np.exp(Vc / (G0  *  isp2))
        q2 = 1 / q2
        q2 = 1 - q2
        q2 *= omega / (1 - eps2)
        qp = omega - q2
        qp_tab.append(qp)
    
    qp=max(qp_tab)
    omega=omega_tab[qp_tab.index(qp)]
    q2=omega - qp
    q1=1 - q2 - qp
    M1=(Mp / qp)  *  q1
    M2=(Mp / qp)  *  q2
    
    return M1, M2
    
def standard_atmosphere(altitude):
    """    Compute atmospheric properties according to the International \
    Standard Atmosphere (ISA) model
    
    :param altitude: altitude in meter
    :type altitude: float
    :return: T, P, rho : temperature (K), pressure (Pa), density (Kg.m-3)
    :rtype: float, float, float
    """
    
    # ISA model
    assert altitude < 84852 , str(altitude) + ' is higher than max altitude of the ISA model' 
    
    h0_tab = np.array([0,      11e3,    20e3,    32e3,    47e3,    51e3,    71e3,    84852]) # base altitude (m)
    a_tab = np.array([-6.5e-3, 0,       1e-3,    2.8e-3,  0,       -2.8e-3, -2.0e-3]) # lapse rate (K.m-1)
    # the base properties of each range are tabulated to reduce computational cost
    T0_tab = np.array([288.15, 216.65,  216.65,  228.65,  270.65,  270.65,  214.65]) # base temperature (K)
    P0_tab = np.array([101325, 22632,   5474.9,  868.02,  110.91,  66.939,  3.9564]) # base pressure (Pa)
    
    index = 0
    while altitude > h0_tab[index + 1]:
        index  += 1
    
    a = a_tab[index]
    T0 = T0_tab[index]
    P0 = P0_tab[index]
    T = T0 + a * (altitude - h0_tab[index])     # Toussaint equation
    
    if T == T0 : # isohermal region
        P = P0 * np.exp(-G0 * (altitude - h0_tab[index]) / (R * T))
    else : # gradient region
        P = P0 * np.power(T/T0, -G0/(a*R))
    
    rho = P / (R * T)    
    
    return T,P,rho

def laval_nozzle(Pt,Tt,Pamb,Ath,Ae,gamma,R,tolerance=1e-5):
    """Compute the thrust assuming a started nozzle

    .. math::  Thrust = massFlowRate \cdot exitVelocity + (exitPressure - ambiantPressure) \cdot exitArea

    :param Pt: total pressure in the combustion chambre (Pa)
    :type Pt: [type]
    :param Tt: total temperature in the combustion chambre (K)
    :type Tt: [type]
    :param Pamb: ambient pressure (Pa)
    :type Pamb: [type]
    :param Ath: throat section area (m2)
    :type Ath: [type]
    :param Ae: exit section area (m2)
    :type Ae: [type]
    :param gamma: ratio of specific heats (adimentional)
    :type gamma: [type]
    :param R: specific gas constant (J.Kg-1.K-1)
    :type R: [type]
    :param tolerance: [description], defaults to 1e-5
    :type tolerance: [type], optional
    :return: Thrust (N), Specific impulse (s)
    :rtype: [type]
    """

    
    assert gamma > 1, 'gamma must be > 1'
    
    
    A_ratio = lambda M,gamma: (1 / M) * np.power((2 / (gamma + 1)) * (1 + ((gamma - 1) / 2) * M * M), (gamma + 1) / (2 * gamma - 2))
     
    # compute the nozzle exit section area corresponding to the adapted regime
    Madapted = np.sqrt((2 / (gamma - 1)) * (np.power(Pamb / Pt,(1 - gamma) / gamma) - 1))
    Aadapted = Ath * A_ratio(Madapted,gamma)
        
    if Aadapted > Ae: # the nozzle is under-expanded
        # compute the exit Mach number iteratively using dichotomy algorithm
        A_ratio_target = Ae / Ath
        Me_upper = Madapted
        Me_lower = 1
        Me = (Me_upper + Me_lower) / 2
        error = A_ratio_target - A_ratio(Me,gamma)
        while np.abs(error) > tolerance:
            if  error > 0:
                Me_lower = Me
            else:
                Me_upper = Me
            Me = (Me_upper + Me_lower) / 2
            error = A_ratio_target - A_ratio(Me,gamma)
        Te = Tt / (1 + ((gamma - 1) / 2) * Me * Me) # exit static temperature
        Pe = Pt * np.power(Tt / Te,gamma / (1 - gamma)) # exit static pressure
        
    elif Aadapted < Ae: # the nozzle is over-expanded
        Mshock = 1
        step = 1
        A_ratio_target = Ae / Ath
        Me_upper = 1
        Me_lower = 0
        Me = (Me_upper + Me_lower) / 2
        error = A_ratio_target - A_ratio(Me,gamma)
        while np.abs(error) > tolerance:
            if  error < 0:
                Me_lower = Me
            else:
                Me_upper = Me
            Me = (Me_upper + Me_lower) / 2
            error = A_ratio_target - A_ratio(Me,gamma)
        Te = Tt / (1 + ((gamma - 1) / 2) * Me * Me)
        Pe = Pt * np.power(Tt / Te,gamma / (1 - gamma))
        while (Pe - Pamb) * 1e-5 > tolerance:
            Mshock_old = Mshock
            Mshock  += step
            Pe_old = Pe
            Pt_ratio = np.power(((gamma + 1) * Mshock * Mshock) / ((gamma - 1) * Mshock * Mshock + 2),gamma / (gamma - 1)) * np.power((gamma + 1) / (2 * gamma * Mshock * Mshock - gamma + 1),1 / (gamma - 1))
            Pi = Pt * Pt_ratio
            Ai = Ath / Pt_ratio
            A_ratio_target = Ae / Ai
            Me_upper = 1
            Me_lower = 0
            Me = (Me_upper + Me_lower) / 2
            error = A_ratio_target - A_ratio(Me,gamma)
            while np.abs(error) > tolerance:
                if  error < 0:
                    Me_lower = Me
                else:
                    Me_upper = Me
                Me = (Me_upper + Me_lower) / 2
                error = A_ratio_target - A_ratio(Me,gamma)
            Te = Tt / (1 + ((gamma - 1) / 2) * Me * Me)
            Pe = Pi * np.power(Tt / Te,gamma / (1 - gamma))
            if (Pe > Pamb) and (A_ratio(Mshock,gamma) > Ae / Ath):
                # the normal shock is outside the divergent and all the flow inside is supersonic
                A_ratio_target = Ae / Ath
                Me_upper = Mshock
                Me_lower = 1
                Me = (Me_upper + Me_lower) / 2
                error = A_ratio_target - A_ratio(Me,gamma)
                while np.abs(error) > tolerance:
                    if  error > 0:
                        Me_lower = Me
                    else:
                        Me_upper = Me
                    Me = (Me_upper + Me_lower) / 2
                    error = A_ratio_target - A_ratio(Me,gamma)
                Te = Tt / (1 + ((gamma - 1) / 2) * Me * Me) # exit static temperature
                Pe = Pt * np.power(Tt / Te,gamma / (1 - gamma)) # exit static pressure
                break
            if Pe < Pamb:
                Pe = Pe_old
                Mshock = Mshock_old
                step /= 2
        
    else: # the nozzle is adapted
        Me = Madapted
        Te = Tt / (1 + ((gamma - 1) / 2) * Me * Me) # exit static temperature
        Pe = Pamb
    
    Ve = Me * np.sqrt(gamma * R * Te) # exit velocity
    
    # maximum mass flow rate in the case of a started nozzle
    Q = np.power(2 / (gamma + 1),(gamma + 1) / (2 * (gamma - 1))) * np.sqrt(gamma / (Tt * R)) * Pt * Ath
    Thrust = Q * Ve + (Pe - Pamb) * Ae
    isp = Thrust / (Q * G0)

    return Thrust,isp
    
def adiabatic_flame(mech,T,P,X):
    """Compute adiabatic flame temperature
    
    :param mech: chemical mechanism in cti format
    :type mech: [type]
    :param T: initial temperature (K)
    :type T: [type]
    :param P: pressure (Pa)
    :type P: [type]
    :param X: initial chemical composition
    :type X: [type]
    """
    # work in progress
    

#%% Trajectory tools

def solve_traj():
    """
    """
    
# .......
    
#%% Astrodynamics tools
    
def state_to_orbital_elem(r,v):
    """
    Convert state vector to orbital elements
    inputs:
        r : coordinate vector (m)
        v : velocity vector (m.s-1)
    output: {'a','e','i','capital_omega','small_omega','nu'}
    where:
        a : semi_major axis (m)
        e : eccentricity (adimentional)
        i : inclination (deg)                                
        capital_omega : argument of the ascending node (deg) 
        small_omega   : argument of the perigee (deg)
        nu : true anomaly (deg)
    NB: - An internal quantity (small) is considered for the case of undefined 
          classical orbital elements. When an element is undefined it is set to
          zero and other(s) following it are measured according to this convention
        - The function stops when r and v are aligned
        - This function is widely inspired by the Matlab function written by 
          Prof. Josep Masdemart from UPC 
    """
    
    orbel={'a':0,'e':0,'i':0,'capital_omega':0,'small_omega':0,'nu':0}
    
    small = 1e-10
    
    h = np.cross(r,v) # angular momentum
    rm = LA.norm(r) # module of the coordinate vector
    vm = LA.norm(v) # module of the velocity vector
    hm = LA.norm(h) # module of the angular momentum vector
    if hm < small:
        print('rectilinar or collision trajectory')
        return
    n = np.array([-h[1],h[0],0])
    nm = LA.norm(n)
    if nm < small:
        n = np.array([1,0,0])
        nm = 0
    else:
        n = n / nm
        nm = 1
    dotrv = np.dot(r,v)
    e = ((vm * vm - MU / rm) * r - dotrv * v) / MU
    em = LA.norm(e)
    if em < small:
        e = n
        em = 0
    else:
        e = e / em
    orbel['e'] = em
    orbel['a'] = (hm * hm / MU) / (1 - em * em)
    aux = h[2] / hm
    if np.abs(aux) > 1: aux = np.sign(aux)
    # all the angles are first computed in rad
    orbel['i'] = np.arccos(aux)
    if nm > small:
        aux = n[0]
        if np.abs(aux) > 1: aux = np.sign(aux)
        orbel['capital_omega'] = np.arccos(aux)
    if n[1] < 0: orbel['capital_omega'] = 2 * np.pi - orbel['capital_omega']
    if em > small:
        aux = np.dot(e,n)
        if np.abs(aux) > 1: aux = np.sign(aux)
        orbel['small_omega'] = np.arccos(aux)
        if e[2] < 0: orbel['small_omega'] = 2 * np.pi - orbel['small_omega']
    aux = np.dot(e,r) / rm
    if np.abs(aux) > 1: aux = np.sign(aux)
    orbel['nu'] = np.arccos(aux)
    if em > small:
        if dotrv < 0: orbel['nu'] = 2 * np.pi - orbel['nu']
    else:
        ha = np.cross(e,r)
        if np.dot(h,ha) < 0: orbel['nu'] = 2 * np.pi - orbel['nu']
    # convert the angles to deg
    orbel['i'] = np.rad2deg(orbel['i'])
    orbel['capital_omega'] = np.rad2deg(orbel['capital_omega'])
    orbel['small_omega'] = np.rad2deg(orbel['small_omega'])
    orbel['nu'] = np.rad2deg(orbel['nu'])
    return orbel

def propagate_orbit():
    """
    """

#%% Thermodynamic cycles 
# the following functions provide the pressure, the mixuture composition 
# and the initial temperature in the combustion chambre and properties of some 
# components of the cycle.

def pressure_fed():
    """
    """
    
def electric_turbo():
    """
    """
    
def open_expander():
    """
    """
    
def closed_expander():
    """
    """
    
def dual_expander():
    """
    """
    
def tap_off():
    """
    """
    
def gas_generator():
    """
    """
    
def ox_rich_staged_combust():
    """
    """
    
def full_flow_staged_combust():
    """
    """
    
    
if __name__ == "__main__":

    pass