import numpy as np
import matplotlib.pyplot as plt
#Parametry początkowe
h=1    #wysokość
N=1000 #liczba kroków
dt=0.001#krok czasowy
g=9.81 #przyspieszenie
t=np.arange(0,N*dt,dt)
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)*0.91 #prędkośc uderzenia
       x[i]=vz*(dt-dtz)-g*(dt-dtz)**2/2
       v[i]=vz-g*(dt-dtz)
ymax=np.max(x*1.1) #zakres wykresu w osi pionowej
for i in range(N):
    plt.scatter(0,x[i])
    plt.ylim(0,ymax)
    plt.pause(dt)
    plt.clf()

#plt.plot(t,x)
plt.scatter(t[-1],x[-1])
plt.ylim(0,ymax)
plt.show()