import numpy as np
import matplotlib.pyplot as plt

G = 6.672*10**(-11)
M = 5.974*10**(24)
C = 0.48
r = 0.1
ro_Fe = 7870
ro_0 = 1.24
H = 9150
h = 500000
R = 6378000

t = [0]
dt = 0.01
x = [R + h]   # vzdialenost od stredu Zeme v pripade bez odporu vzduchu
xx = [R + h]   # vzdialenost od stredu Zeme v pripade s odporom vzduchu
v = [-1000]   # rychlost v pripade bez odporu vzduchu
vv = [-1000]   # rychlost v pripade s odporom vzduchu
g = [-G*M/x[0]**2]   # gravitacne zrychlenie v pripade bez odporu vzduchu
a = [-G*M/x[0]**2]   # gravitacne zrychlenie v pripade bez odporu vzduchu
# Pad meteoritu si rozdelime na male casove useky, v ramci ktorych mozme
# povazovat jeho pohyb za rovnomerne zrychleny.
for i in range(1, 10**6):
    if x[i - 1] < R:   # moment dopadu meteoritu na zemsky povrch
        break
    t.append(i*dt)
    ro = ro_0*np.exp(-(x[i - 1] - R)/H)   # hustota vzduchu meniaca sa s vyskou
    a_0 = 3/8*C/r*ro/ro_Fe*v[i - 1]**2   # brzdne zrychlenie odporu vzduchu
    x.append(x[i - 1] + v[i - 1]*dt + 1/2*a[i - 1]*dt**2)
    xx.append(xx[i - 1] + vv[i - 1]*dt + 1/2*g[i - 1]*dt**2)
    v.append(v[i - 1] + a[i - 1]*dt)
    vv.append(vv[i - 1] + g[i - 1]*dt)
    g.append(-G*M/x[i - 1]**2)
    a.append(a_0 + g[i - 1])
plt.plot(t, x, 'r', t, xx, 'g')
plt.plot([0, max(t)], [6378000, 6378000], 'k:')
plt.xlabel('cas [s]')
plt.ylabel('poloha [m]')
plt.legend(['s odporom','bez odporu'])
plt.figure()
plt.plot(t, v, 'r', t, vv, 'g')
plt.xlabel('cas [s]')
plt.ylabel('rychlost [m/s]')
plt.legend(['s odporom','bez odporu'])
plt.figure()
plt.plot(t, a, 'r', t, g, 'g')
plt.xlabel('cas [s]')
plt.ylabel('zrychlenie [ms^-2]')
plt.legend(['s odporom','bez odporu'])