"""
Program symulujący spadek swobodny punktowego ciała i odbicie 
od powierzchni Ziemi (z utratą energii).
"""
import numpy as np
import matplotlib.pyplot as plt

#Parametry początkowe
h   = 1       #wysokość
N   = 840     #liczba kroków (jeśli za duża, to numeryka się rozjedzie)
dt  = 0.01    #długość kroku w sekundach
g   = 9.81    #przyspieszenie ziemskie
coef= 0.9     #jak skaluje się wartość prędkości przy odbiciu

x = np.empty(N)
v = np.empty(N)
x[0] = h
v[0] = 0
for i in range(1,N):
    x[i] = x[i-1] + v[i-1]*dt - g*dt**2/2
    v[i] = v[i-1] - g*dt
    if x[i] < 0: #odbicie
       dtz  = (v[i-1] + np.sqrt(v[i-1]**2 + 2*g*x[i-1]))/g #czas do uderzenia w ziemię
       vz   = -(v[i-1] - g*dtz)*coef #prędkośc po uderzeniu - zmienia znak i zmniejsza się
       dtr  = dt - dtz #pozostała część kroku czasowego
       x[i] = vz*dtr - g*dtr**2/2
       v[i] = vz - g*dtr

print(x)
ymax = 1.1*np.max(x) #Zakres osi y
for i in range(N): #Animacja
    plt.scatter(0,x[i])
    plt.ylim(0,ymax)
    plt.pause(dt) #czas wyświetlania klatki
    plt.clf()     #wyczyszczenie klatki animacji (CLear current Figure)

#Użycie plt.pause(dt) powyżej skutkuje tym, że obecny wykres znika po dt.
#Dlatego na koniec rysujemy jeszcze raz ostatnią klatkę za pomocą plt.show().
plt.scatter(t[-1],x[-1])
plt.ylim(0,ymax)
plt.show()