#!/usr/bin/python
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode
import matplotlib.animation as anim

#****************************
#*********parameetrid ***********************
#*******************************************************

m_maa=5.9722e+24 # maa mass, kg
r_maa=6378.      # maa raadius, m

m_keha=2800.     # keha makk, kg
r_keha=2.        # keha raadius, m

t_ohu=300.       # ohu temperatuur, K
p_ohu_0=101000.  # ohu rohk, Pa
mu_ohu=29.       # ohu uhe mooli mass, kg/kmool

vx=0.            # mooduli algkiiruse x-projektsioon, km/s
vy=8.5           # mooduli algkiiruse y-projektsioon, km/s
y0=0.            # mooduli y-algkoordinaat, km
h0=100.          # mooduli algkorgus

x0=h0+r_maa      # mooduli x-algkoordinaat, km
aeg=80000        # kogu liikumise aeg sekundites
animkiirus=10    # animatsiooni kiirus

#********************************************************
#********************************************
#*****************************
asa=50000
time=np.linspace(0,aeg,num=asa)
gamma=6.67408e-11
r_gaas=8314.
s_keha=np.pi*r_keha**2
c_keha=0.47
g_rask=9.814
gm=gamma*m_maa*1.e-9
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
y0=[x0,y0,vx,vy]

def func(y,t):
   yp=np.zeros_like(y)
   r=np.sqrt(y[0]**2+y[1]**2)
   h=r-r_maa
   ro_ohu=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_ohu*v*y[2]
   yp[3]=-gm*y[1]/r3-p2*ro_ohu*v*y[3]
   return yp

res=ode.odeint(func,y0,time)

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
c_text=""
for i in np.arange(0,len(time)):
    if r[i] < r_maa:
       itime=i+1
       resr=res[0:itime,:]
       c_text="Kukkus Maale kiirusega"+str("%10.3f" % (vv[itime]*1000.))+" m/s"
       break
    else:
       resr=res
       itime=i

#*********** visualization ************
xmin=min(resr[:,0])*1.2
xmax=max(resr[:,0])*1.2
ymin=min(resr[:,1])*1.2
ymax=max(resr[:,1])*1.2
if xmin >0:
    xmin=-xmin
if ymin >0:
    ymin=-ymin

fig2=plt.figure(2,figsize=(10,10))
fig2.canvas.set_window_title('Trajektoor')
p2 = plt.subplot()
p2=plt.axes(xlim=(xmin,xmax), ylim=(ymin, ymax))
p2.set_aspect("equal")
maa=plt.Circle((0,0),6378,color="green",label="Maa")
atm=plt.Circle((0,0),6478,color="blue",fill=False,label="O~hu")
p2.add_artist(maa)
p2.add_artist(atm)
p2.text(0.02,0.95,c_text,transform=p2.transAxes)
p2.plot(resr[:,0],resr[:,1],color="red",label="moodul")
p2.legend(loc="upper right")
p2.grid()

#**** anim *****
fig1=plt.figure(1,figsize=(10,10))
fig1.canvas.set_window_title('Animatsioon')
p1 = plt.subplot()
p1=plt.axes(xlim=(xmin,xmax), ylim=(ymin, ymax))
p1.set_aspect("equal")
p1.grid()

def animf(i):
    sx = [resr[i,0]]
    sy = [resr[i,1]]
    sat.set_data(sx, sy)
    vel_txt.set_text('Kiirus = %10.3f km/s' % (vv[i]))
    dist_txt.set_text('Korgus = %10.3f km' % (r[i]-r_maa))
    kiirend_txt.set_text('Kiirendus = %10.3f m/s2' % (aa[i]))
    time_txt.set_text('Aeg = %8.3f min = %8.3f tund' % (time[i]/60.,time[i]/3600.))
    return sat,vel_txt,dist_txt,kiirend_txt,time_txt

maa=plt.Circle((0,0),6378,color="green",label="Maa")
atm=plt.Circle((0,0),6478,color="blue",fill=False,label="Maa")
p1.add_artist(maa)
p1.add_artist(atm)

sat, = p1.plot([], [], '.',color="red",label="moodul")
vel_txt = p1.text(0.02,0.92,'',transform=p1.transAxes)
dist_txt = p1.text(0.02,0.95,'',transform=p1.transAxes)
kiirend_txt = p1.text(0.02,0.89,'',transform=p1.transAxes)
time_txt = p1.text(0.02,0.86,'',transform=p1.transAxes)

p1.legend(loc="upper right")
ani = anim.FuncAnimation(fig1, animf, np.arange(0, itime,animkiirus),interval=10,repeat=False,blit=True)

plt.show()