import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode

def rask(g_const,m_maa,r_maa,h):
 return g_const*m_maa/(r_maa+h)**2

def func(t,y):
   yp=np.zeros_like(y)
   yp[0]=y[1]
   g_rask=rask(g_const,m_maa,r_maa,y[0])
   ro=ro_atm_0*np.exp(-p1*y[0]/(1+y[0]/r_maa))
   yp[1]=-g_rask-p2*ro*np.abs(y[1])*y[1]+p3*ro*g_rask
   return yp

g_const=6.67403e-11
m_maa=5.9722e+24
r_maa=6371000.

c_keha=0.47
r_keha=1.8
m_keha=15.
s_keha=np.pi*r_keha**2
v_keha=4.*np.pi*r_keha**3/3

p_atm=100000.
t_atm=293.
mu=29.
r_atm=8314.
ro_atm_0=p_atm*mu/(r_atm*t_atm)

p1=rask(g_const,m_maa,r_maa,0.)*mu/(r_atm*t_atm)
p2=0.5*c_keha*s_keha/m_keha
p3=v_keha/m_keha

y0=[0.,0.]
t0=0.
dt=0.01
nt=400000

r=ode.ode(func)
r.set_initial_value(y0,t0)
res=[]
t=[]
it=0

while r.successful() and r.y[0] >= 0.0 and it<=nt:
 it+=1
 r.integrate(r.t+dt)
 res.append(r.y)
 t.append(r.t)
res=np.array(res)
t=np.array(t)

fig,[p1,p2]=plt.subplots(1,2)
p1.plot(t,res[:,0])
p2.plot(t,res[:,1])
p1.grid()
p2.grid()
p1.set_xlabel("t,sek")
p1.set_ylabel("x,m")
p2.set_xlabel("t,sek")
p2.set_ylabel("v,m/s")
plt.show()