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

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

r_keha=0.5
s_keha=np.pi*r_keha**2
c_keha=0.47
v_keha=4.*np.pi*r_keha**3/3
m_keha=5.
p=100400.
temp=280.
mu=29.
R=8314.
m_maa=5.97219e+24
r_maa=6371000.
g_const=6.6743015e-11
ro_0=p*mu/(R*temp)
vtuul=-12.5

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

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

def func(t,y):
 yp=np.zeros_like(y)
 yp[0]=y[2]
 yp[1]=y[3]
 ro=ro_0*np.exp(-p1*y[1])
 g_rask=rask(g_const,r_maa,m_maa,y[1])   
 vv=np.sqrt((y[2]-vtuul)**2+y[3]**2)
 yp[2]=-p2*ro*vv*(y[2]-vtuul)
 yp[3]=-g_rask-p2*ro*vv*y[3]+p3*ro*g_rask
 return yp

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

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

g_rask=rask(g_const,r_maa,m_maa,res[:,1])
vv=np.sqrt((res[:,2]-vtuul)**2+res[:,3]**2)
ro=ro_0*np.exp(-p1*res[:,1])
#vv=np.array(vv)
#ro=np.array(ro)
ax=-p2*ro[:]*vv[:]*(res[:,2]-vtuul)
ay=-g_rask-p2*ro[:]*vv[:]*res[:,3]+p3*ro[:]*g_rask[:]
#ax=np.array(ax)
#ay=np.array(ay)



#print(atx,tt)
#print (len(atx))
scalep=10
widthp=0.002
n=len(t[:])
tt=np.arange(0,n,10)

plt.quiver(res[tt,0],res[tt,1],ax[tt],ay[tt],angles="xy",color="black",scale_units="xy",scale=scalep,width=widthp,label="atot")
plt.quiver(res[tt,0],res[tt,1],res[tt,2],res[tt,3],angles="xy",color="magenta",scale_units="xy",scale=scalep,width=widthp,label="v")
plt.plot(res[:,0],res[:,1],color="gray",label="traj")
plt.legend()
plt.grid()
plt.show()