import numpy as np

# Konstanty ze zadání
R_Z = 6378000              #Poloměr Země
M_Z = 5.972 * pow(10,24)   #Hmotnost Země
h = 160000                 #Výsledná výška letu rakety
G = 6.67 * pow(10,-11)     #Gravitační konstanta
C_r = 0.3                  #Součinitel odporu
S = 0.21                   #Plocha průřezu rakety
rho_0 = 1.225              #Hustota standardní atmosféry
H_n = 10.4                 #Konstanta poklesu hustoty
m_I = 540                  #Hmotnost prvního stupně
m_II = 40                  #Hmotnost druhého stupně
m_nakl = 1                 #Hmotnost předmětu
m_p2 = 652                 #Hmotnost paliva druhého stupně
u=2600                     #Výtoková rychlost zplodin
m_dot = 71.2               #Hmotnostní průtok paliva
m_max = 17660              #Horní odhad hmotnosti prvního stupně rakety
m_0 = m_I + m_II + m_nakl + m_p2   #Hmotnost na konci první fáze letu
# Počáteční podmínky
m_pa=0
v=0
x=0
dt=0.1 #Časová změna, čím větší dt bude, tím rychleji program poběží, ale bude méně přesný výsledek
vysl=20000 #libovolné dostatečně vysoké číslo, aby se v průběhu simulace zapsal nejnižší možný výsledek 

#Řešení ODR
while m_pa <= m_max:
    v=0
    x=0
    m_p=m_pa
    m=m_0+m_pa
    while x<=h and v>=0:
        x=x+v*dt
        if m_p>0:    #Průběh změny rychlosti dokud raketa spaluje palivo
            v=v+(u*m_dot/(m_0+m_p) - G*M_Z/(R_Z+x)**2 - 0.5*C_r*rho_0*S/(m_0+m_p) * np.exp(-x/H_n)*v**2)*dt
            m_p=m_p-m_dot*dt
        else:       #Změna rychlosti po spálení veškerého paliva
            v=v+(- G*M_Z/(R_Z+x)**2 - 0.5*C_r*rho_0*S/(m_0) * np.exp(-x/H_n)*v**2)*dt

    if x>=160000 and m_pa <= vysl:     #Zápis výsledku
        vysl = m_pa
        print("Výsledná hmotnost paliva prvního stupně rakety je:", vysl, "kg")
        break            #Pokud raketa vyletí do požadované výšky, získáme tím nejmenší možnou hmotnost a program skončí
    m_pa = m_pa + 1      #Program zkouší rovnici řešit pro všechny přípustné hmotnosti paliva