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

def func(t,y):
 dy=np.zeros_like(y)
 rr=np.sqrt(y[0]**2+y[1]**2)
 vv=np.sqrt(y[2]**2+y[3]**2)
 h=rr-r_maa
 ro=ro_atm_0*np.exp(-p1*h/(1+h/r_maa))
 g_rask=rask(g_const,r_maa,m_maa,h)
 dy[0]=y[2]
 dy[1]=y[3]
 dy[2]=-y[0]*gm/rr**3-p2*ro*vv*y[2]+p3*g_rask*ro*y[0]/rr
 dy[3]=-y[1]*gm/rr**3-p2*ro*vv*y[3]+p3*g_rask*ro*y[1]/rr
 return dy

def rask(g_const,r_maa,m_maa,h):
 return g_const*m_maa*1.e-6/(r_maa+h)**2

c_keha=0.47
r_keha=2.
m_keha=2800.
p_atm=101000.
mu_atm=29.
r_gaas=8314.
t_atm=300.
g_const=6.67408e-11
m_maa=5.9722e+24
gm=g_const*m_maa*1.e-9
r_maa=6371.
s_keha=np.pi*r_keha**2
v_keha=4.*np.pi*r_keha**3/3.
ro_atm_0=p_atm*mu_atm/(r_gaas*t_atm)
g_rask=rask(g_const,r_maa,m_maa,0.)
p1=mu_atm*g_rask*1.e3/(r_gaas*t_atm)
p2=s_keha*c_keha*500./m_keha
p3=v_keha*1.e-3/m_keha

print (g_rask)

v0=7.7884079
alpha=0.3
h0=200.
y0=[r_maa+h0,0.,-v0*np.sin(np.pi*alpha/180.),v0*np.cos(np.pi*alpha/180.)]
t0=0.
dt=1.
nt=200000

r = ode.ode(func)
r.set_initial_value(y0, t0)
r.set_integrator('vode',atol=1.e-9,rtol=1.e-9)
res = []
t = []
a=[]
it=0
iflag=0

while r.successful() and np.sqrt(r.y[0]**2+r.y[1]**2) > r_maa and it<nt:
 it=it+1
 r.integrate(r.t+dt)
 res.append(r.y)
 t.append(r.t)
 hh=np.sqrt(r.y[0]**2+r.y[1]**2)-r_maa
#----------
 a.append(func(t,r.y))
#---------
 if (iflag == 0):
  if (hh <= 5):
   iflag=1
   r_keha=4.
   s_keha=np.pi*r_keha**2
   v_keha=4.*np.pi*r_keha**3/3.
   p2=s_keha*c_keha*500./m_keha
   p3=v_keha*1.e-3/m_keha  
res = np.array(res)
a=np.array(a)
t = np.array(t)/60.

h=np.sqrt(res[:,0]**2+res[:,1]**2)-r_maa
vv=np.sqrt(res[:,2]**2+res[:,3]**2)
aa=np.sqrt(a[:,0]**2+a[:,1]**2)

print(t[-1],"min")
fig,p=plt.subplots(2,2,figsize=[10,8])
p[0,0].plot(t,h,label="h(t)")
p[0,0].set_xlabel("t,min")
p[0,0].set_ylabel("h(t),sek")

p[0,1].plot(h,vv,label="v(h)")
p[0,1].set_xlabel("h,km")
p[0,1].set_ylabel("v,km/s")

p[1,0].plot(h,aa*1000,label="a(h)")
p[1,0].set_xlabel("h,km")
p[1,0].set_ylabel("a,m/s2")

p[1,1].plot(t,res[:,0],label="x(t)")
p[1,1].set_xlabel("t,min")
p[1,1].set_ylabel("x,m/s2")

p[0,0].legend()
p[0,1].legend()
p[1,0].legend()
p[1,1].legend()
p[0,0].grid()
p[0,1].grid()
p[1,0].grid()
p[1,1].grid()
plt.show()

fig.savefig("out.png",format="png")