import numpy
from numpy import sin, cos
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as pyplot

# Federkonstante
k = 1

# Zeitschritt
dt = 0.01
# Speicher
speicher = 10
# Anzahl Schritte vor replot
steps = 10

# Laufende Variablen

# (Start-)Position
x = 0.3
# (Start-)Geschwindigkeit
v = 0
# Zeit
t = 0

# Tabellen fuer die Ausgabe
tn, xn, en, ekn, epn = [], [], [], [], []

# Erzeugen der (leeren) Graphen
xmax = 1.1*x

ausgabe = pyplot.figure(figsize=(8,4))

plot_x = ausgabe.add_subplot(121)
plot_x.set_xlabel("Zeit T")
plot_x.set_ylabel("Position x")
data_x, = plot_x.plot(tn, xn, '-')

plot_e = ausgabe.add_subplot(122)
plot_e.set_xlabel("Zeit T")
plot_e.set_ylabel("Energie E")
data_e, data_ek, data_ep = plot_e.plot(tn, en, '-', tn, ekn, '-', tn, epn, '-')

# Funktion, 
def animate():
    global x, v, t, tn, xn, en, epn, ekn
        
    F = -k*x
        
    # einige Schritte propagieren
    for i in range(steps):
        # velocity verlet

        # Geschwindigkeits-Halbschritt
        v += 0.5*F*dt

        # neue Positionen
        x += v*dt

        # neue Kraft
        F = -k*x

        # zweiter Geschwindigkeits-Halbschritt
        v += 0.5*F*dt

        t += dt

        # Observablen
        ek = 0.5*v*v
        ep = 0.5*k*x*x
        e = ek+ep

	tn.append(t)
	xn.append(x)
	en.append(ek)
	epn.append(ep)
	ekn.append(e)

	# Daten vom Anfang wegschmeissen
	if tn[-1] - tn[0] > speicher:
		del tn[0], xn[0], en[0], ekn[0], epn[0]

	tmin, tmax = tn[0], tn[-1]
	emin, emax = -0.1*max(en), 1.1*max(en)

	data_x.set_data(tn, xn)
	plot_x.axis((tmin, tmax, -xmax, xmax))

	data_e.set_data(tn, en)
	data_ep.set_data(tn, ekn) 
	data_ek.set_data(tn, epn) 
	plot_e.axis((tmin, tmax, emin, emax))

	ausgabe.canvas.draw()

	# periodisch die Funktion animate aufrufen
	ausgabe.canvas.manager.window.after(100, animate)

ausgabe.canvas.manager.window.after_idle(animate)
pyplot.show()