#!/usr/bin/env python
# -*- coding: utf-8 -*-
#Die zweite Zeile braucht es, damit wir Umlaute verwenden können

"""
   Pythonskript zur Simulation von n gekoppelten Pendeln mit Massen $m_i$ 
   und Reibungskoeffizienten b_i.
   Die Federkonstanten sind hier noch alle gleich. 
"""


from scipy.integrate import odeint
from numpy import *
#from numpy import loadtxt, append
import matplotlib
import matplotlib.pyplot as plt

#Definition von Konstanten des Problems (kann man ändern und damit spielen!)
#Anzahl Pendel
n = 5
#Massen:
m = 1.*ones(n)
#Reibungskoeffizienten
r = 0.00025*ones(n)
#Federkonstante
D = 1.e0

#Verpacken der Parameter in einen "Vektor"
p = append(append(append(m,r),D),n)


def eqs(xin,t,p):
    """
    Hier wird die zu integrierende Funktion definiert.
    x:  Vektor mit den Variablen x_i, die den aktuellen Zustand beschreiben
    t:  Zeit
    p:  Parameter (m_i, b_i, D) (Massen, Dämpfungen, Kopplung)
    """
    #print 'p =',p,', n = ',p[-1]
    n = int(p[-1])
    D = p[-2]
    m = p[0:n]
    b = p[n:2*n]
    x = [xin[i] for i in arange(0,2*n,2)]
    y = [xin[i] for i in arange(1,2*n+1,2)]
   
    #Nun beschreiben wir die Ableitungen
    f = []
    f.append(y[0])
    f.append((-b[0]*y[0] - 2*D*x[0] + D*x[1])/m[0])
    for i in arange(1,n-1,1):
        f.append(y[i])
        f.append((-b[i]*y[i] + D*x[i-1] - 2*D*x[i] + D*x[i+1])/m[i])
    f.append(y[n-1])
    f.append((-b[n-1]*y[n-1] + D*x[n-2] - 2*D*x[n-1])/m[n-1])
    return f


#Definition der Anfangsbedingunen
#Alle Pendel sind in ihrer Ruheposition
#Nur Pendel 1 hat eine Anfangsgeschwindigkeit
x = zeros(n)
y = zeros(n)
x[0] = 0.2

#Verpacke nun die Anfangsbedingungen in einen "Vektor" z0 (Zustand 0)
z0 = []
for i in arange(0,n):
    z0.append(x[i])
    z0.append(y[i])

# Definiere Parameter für die Integrationsroutine
abserr = 1.0e-8
relerr = 1.0e-6
stoptime = 500.0
numpoints = 750

# Erzeuge die Zeitpunkte, für die die Lösungen berechnet werden sollen.
t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]

#und rufe die Integrationsroutine auf

loesung = odeint(eqs, z0, t, args = (p,), atol = abserr, rtol = relerr)

#Speichere nun die Resultate ab

#filename = '{}-oszillatoren.dat'.format(n)
##print 'writing to file ', filename
#with open(filename,'w') as f:
#    for zeit, l in zip(t, loesung):
#        print >> f, zeit, [l[i] for i in arange(0,n)]

fig, ax = plt.subplots()
sum = zeros(numpoints)
#for i in arange(0,2*n,2):
#    ax.plot(t,loesung[:,i],lw=2,label=r'$x_{}$'.format(i/2))
#    #if i == 5: ax.plot(t,loesung[:,i],lw=4,label=r'$x_{}$'.format(i/2))
#    sum += loesung[:,i]**2
ax.plot(t,loesung[:,0],lw=2)
ax.plot(t,sum,lw=3,c='k')
ax.set_xlabel('Zeit',fontsize=16)
ax.set_ylabel('Auslenkung',fontsize=16)
ax.tick_params(labelsize=14)
if n < 8: ax.legend(loc='upper right')
ax.set_title(u'{} gekoppelte Pendel'.format(n),fontsize=16)
plt.savefig(u'{}-oszillatoren.png'.format(n),dpi=100)
plt.savefig(u'{}-oszillatoren.eps'.format(n),dpi=100)
plt.show()