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

gamma=6.67408e-11
m_maa=5.9722e+24
gm=gamma*m_maa*1.e-9
r_maa=6378.
g_rask=9.814

m_keha=2800.
r_keha=2.8
c_keha=0.47
s_keha=np.pi*r_keha**2
t_ohu=300.
p_ohu_0=101000.
mu_ohu=29.
r_gaas=8314.
ro_ohu_0=p_ohu_0*mu_ohu/(r_gaas*t_ohu)
p1=mu_ohu*g_rask*1000./(r_gaas*t_ohu)
p2=500.*c_keha*s_keha/m_keha

def func(t,y):
   yp=np.zeros_like(y)
   r=np.sqrt(y[0]**2+y[1]**2)
   h=r-r_maa
   ro=ro_ohu_0*np.exp(-p1*h)
   v=np.sqrt(y[2]**2+y[3]**2)
   r3=r**3
   yp[0]=y[2]
   yp[1]=y[3]
   yp[2]=-gm*y[0]/r3-p2*ro*v*y[2]
   yp[3]=-gm*y[1]/r3-p2*ro*v*y[3]
   return yp

v0=7.7884079
alpha=0.5*np.pi/180.
vx=-v0*np.sin(alpha)
vy=v0*np.cos(alpha)

y0=[6578.,0.,vx,vy]
t0=0.
dt=1.
nt=200000
it=0
r=ode.ode(func)
r.set_initial_value(y0,t0)
res=[]
t=[]
h=200.

while r.successful() and h >= 0.0 and it <nt:
  it=it+1
  r.integrate(r.t+dt)
  res.append(r.y)
  t.append(r.t)
  h=np.sqrt(r.y[0]**2+r.y[1]**2)-r_maa
  print (h)
res=np.array(res)
t=np.array(t)

r=np.sqrt(res[:,0]**2+res[:,1]**2)
hh=r-r_maa
vv=np.sqrt(res[:,2]**2+res[:,3]**2)
r3=r**3
ro=ro_ohu_0*np.exp(-p1*hh)
ax=-gm*res[:,0]/r3-p2*ro*vv*res[:,2]
ay=-gm*res[:,1]/r3-p2*ro*vv*res[:,3]
aa=np.sqrt(ax**2+ay**2)*1000

fig,(p1,p2,p3)=plt.subplots(1,3)

p1.set_xlabel("korgus")
p1.set_ylabel("kiirus")
p1.plot(hh,vv,color='blue')
p1.grid()

p1.set_xlabel("korgus")
p2.set_ylabel("kiirendus")
p2.plot(hh,aa,color='red')
p2.grid()

p1.set_xlabel("aeg")
p3.set_ylabel("korgus")
p3.plot(t/60,hh,color='black')
p3.grid()
plt.show()