import numpy
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as pyplot

# Federkonstanten
k1 = 1
k2 = 0.7

# Federkonstante der Kopplungsfeder
kc = 0.2

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

# Laufende Variablen

# (Start-)Position
x1 = 0.3
x2 = 0.2

# (Start-)Geschwindigkeit
v1 = 0
v2 = 0
# Zeit
t = 0

# Tabellen fuer die Ausgabe
tn, x1n, x2n, e1n, e2n, ecn, en = [], [], [], [], [], [], []

# Erzeugen der (leeren) Graphen
xmax = 1.1*abs(x1)+abs(x2)

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

plot_x = ausgabe.add_subplot(121)
plot_x.set_xlabel("Zeit T")
plot_x.set_ylabel("Position x")
line_x1, line_x2 = plot_x.plot(tn, x1n, '-', tn, x2n, '-')

plot_e = ausgabe.add_subplot(122)
plot_e.set_xlabel("Zeit T")
plot_e.set_ylabel("Energie E")
line_e1, line_e2, line_ec, line_e = \
    plot_e.plot(tn, e1n, '-', tn, e2n, '-', 
                tn, ecn, '-', tn, en, '-')

def animate():
    global x1, x2, v1, v2, t, tn, x1n, x2n, e1n, e2n, ecn, en

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

    # Kraft, die auf die Gewichte wirkt
    F1 = -k1*x1 - kc*(x1-x2)
    F2 = -k2*x2 - kc*(x2-x1) 

    # einige Schritte propagieren
    for i in range(steps):
        # velocity verlet

        # Geschwindigkeits-Halbschritt
        v1 += 0.5*F1*dt
        v2 += 0.5*F2*dt

        # neue Positionen
        x1 += v1*dt
        x2 += v2*dt
 
        # neue Kraefte
        F1 = -k1*x1 - kc*(x1-x2)
        F2 = -k2*x2 - kc*(x2-x1) 

        # zweiter Geschwindigkeits-Halbschritt
        v1 += 0.5*F1*dt
        v2 += 0.5*F2*dt
        
        t += dt

        # Observablen
        e1 = 0.5*v1*v1 + 0.5*k1*x1*x1
        e2 = 0.5*v2*v2 + 0.5*k2*x2*x2
        ec = 0.5*kc*(x1-x2)**2

        e = e1+e2+ec
                
        tn.append(t)
        x1n.append(x1)
        x2n.append(x2)
        e1n.append(e1)
        e2n.append(e2)
        ecn.append(ec)
        en.append(e)

        # Daten vom Anfang wegschmeissen
        if tn[-1] - tn[0] > speicher:
            del tn[0], x1n[0], x2n[0], e1n[0], e2n[0], ecn[0], en[0]

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

    line_x1.set_data(tn, x1n)
    line_x2.set_data(tn, x2n)
    plot_x.axis((tmin, tmax, -xmax, xmax))

    line_e1.set_data(tn, e1n)
    line_e2.set_data(tn, e2n) 
    line_ec.set_data(tn, ecn) 
    line_e.set_data(tn, en)
    plot_e.axis((tmin, tmax, emin, emax))

    ausgabe.canvas.draw()

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