#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 26 20:08:38 2019
@author: MAC
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ---------------- Paramètres à modifier ---------------------------
c_k1 = 0.1
c_k2 = 0.1
c_Vc = 10000
c_lam1 = c_k1 / (c_k1 + 1)
c_lam2 = c_k2 / (c_k2 + 1)
c_g0 = 9.8


def totalMass(p_payload_mass,p_isp2):
    """
        Compute total mass according to payload mass and second stage isp
        
        :param p_payload_mass: The mass of the payload
        :param p_isp2: The isp of the second stage
    """
    v_isp1 = p_isp2 * 0.7 
    
    v_QP = []
    v_q2 = 0
    v_qp = 0
    v_qpMax = 0
    
    v_omega_tab = np.linspace(0.001,0.99,100)
    for v_omega in v_omega_tab:
        v_qp = getQP(v_omega,v_isp1,p_isp2)
        v_QP.append(v_qp)
    
    v_qpMax = max(v_QP)
    v_omega = v_omega_tab[v_QP.index(v_qpMax)]
    v_q2 = v_omega - v_qpMax
    v_q1 = 1 - v_q2 - v_qpMax
    v_m1 = (p_payload_mass / v_qpMax) * v_q1
    v_m2 = (p_payload_mass / v_qpMax) * v_q2
    v_mt =  p_payload_mass + v_m1 + v_m2
    return v_mt
    
    
def getQP(p_omega,p_isp1,p_isp2):
    v_q2 = p_omega + c_lam1 * (1 - p_omega)
    v_q2 = np.power(v_q2,p_isp1 / p_isp2) * np.exp(c_Vc / (c_g0 * p_isp2))
    v_q2 = 1 / v_q2
    v_q2 = 1 - v_q2
    v_q2 *= p_omega / (1 - c_lam2)
    v_qp = p_omega - v_q2
    return v_qp


if __name__ == '__main__':

    v_mpList=[30,40,50] # Mass Payload 

    for v_mp in v_mpList:
        # for each different mass of the payload
        v_Vc_tab = np.linspace(9000,11000,100)
        v_ISP_tab = np.linspace(300,400,100)
        MT=[]
        for v_ISP2 in v_ISP_tab:
            MT.append(totalMass(v_mp,v_ISP2))
        
        Vc_tab_plot = [i*0.001 for i in v_Vc_tab]
        MT_plot = [i*0.001 for i in MT]
        plt.plot(v_ISP_tab,MT_plot,label='Payload='+str(v_mp)+' Kg')
    plt.legend(loc=0)
    plt.xlabel('$Isp_2$ (s)')
    plt.ylabel('Masse totale ($10^3$ kg)')
    plt.show()