#!/home/misha/venv/bin/python
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode

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

p_ohu=100000.
t_ohu=293.
mu_ohu=29.
r_gaas=8314.

ro_ohu_0=p_ohu*mu_ohu/(r_gaas*t_ohu)
p1=g_rask*mu_ohu/(r_gaas*t_ohu)
p2=0.5*c_keha*s_keha/m_keha

def func(t,y):
   yp=np.zeros_like(y)
   yp[0]=y[1]
   ro=ro_ohu_0*np.exp(-p1*y[0])
   v=np.abs(y[1])
   yp[1]=g_rask*(ro/ro_keha-1.)-p2*ro*v*y[1]
   return yp

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

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

#while r.successful() and r.y[0] <= 0.0:
while r.successful() and r.t<= 5000:
 r.integrate(r.t+dt)
 res.append(r.y)
 t.append(r.t)
# print (r.t,r.y[0])
res=np.array(res)
t=np.array(t)

fig,[p1,p2]=plt.subplots(1,2)
p1.plot(t,res[:,1])
p1.legend(["h(t)"])
p2.plot(t,res[:,0])
p2.legend(["v(t)"])
p1.grid()
p2.grid()
plt.show()