"""Wellenausbreitung auf einer gespannten Masse-Feder-Kette.

In diesem Beispiel wird nur eine Longitudinalwelle angeregt,
wobei ein einzelner relativ kurzer Anregungspuls mit kleiner
Amplitude verwendet wird.
"""

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation
import scipy.integrate

# Simulationszeit und Zeitschrittweite [s].
t_max = 10.0
dt = 0.01
# Dimension des Raumes.
n_dim = 2
# Anzahl der Teilchen ohne das anregende Teilchen ganz links.
n_teilchen = 70
# Federkonstante [N/m].
D = 100
# Masse [kg].
m = 0.05
# Länge der ungespannten Federn [m].
federlaenge = 0.05
# Abstand benachbarter Massen in der Ruhelage [m].
abstand = 0.15
# Amplitude der longitudinalen und transversalen Anregung [m].
A_long = 0.01             # sehr kleine Amplitude (longitudinal).
A_tran = 0.0              # keine Transversalwelle.

# Lege die Ruhelage der Teilchen auf der x-Achse fest.
r0 = np.zeros((n_teilchen, n_dim))
r0[:, 0] = np.linspace(abstand, n_teilchen * abstand, n_teilchen)


def anregung(t, t0=0.5, delta_t=0.05):
    """Ortsvektor der anregenden Masse zum Zeitpunkt t."""
    pos = np.array([A_long * np.exp(-((t - t0) / delta_t) ** 2),
                    A_tran * np.exp(-((t - t0) / delta_t) ** 2)])
    return pos


def federkraft(r1, r2):
    """Kraft auf die Masse am Ort r1 durch die Masse am Ort r2."""
    abstand = np.linalg.norm(r2 - r1)
    einheitsvektor = (r2 - r1) / abstand
    return D * (abstand - federlaenge) * einheitsvektor


def dgl(t, u):
    """Berechne die rechte Seite der Differentialgleichung."""
    r, v = np.split(u, 2)
    r = r.reshape(n_teilchen, n_dim)
    a = np.zeros((n_teilchen, n_dim))

    # Addiere die Beschleunigung durch die jeweils linke Feder.
    for i in range(1, n_teilchen):
        a[i] += federkraft(r[i], r[i-1]) / m

    # Addiere die Beschleunigung durch die jeweils rechte Feder.
    for i in range(n_teilchen - 1):
        a[i] += federkraft(r[i], r[i+1]) / m

    # Addiere die Beschleunigung durch die anregende Masse.
    a[0] += federkraft(r[0], anregung(t)) / m

    # Die letzte Masse soll festgehalten werden.
    a[-1] = 0

    return np.concatenate([v, a.reshape(-1)])


# Lege den Zustandsvektor zum Zeitpunkt t=0 fest. Alle Teilchen
# ruhen in der Ruhelage.
v0 = np.zeros(n_teilchen * n_dim)
u0 = np.concatenate((r0.reshape(-1), v0))

# Löse die Bewegungsgleichung bis zum Zeitpunkt t_max.
result = scipy.integrate.solve_ivp(dgl, [0, t_max], u0,
                                   t_eval=np.arange(0, t_max, dt))
t = result.t
r, v = np.split(result.y, 2)

# Wandle r und v in ein 3-dimensionals Array um:
#    1. Index - Teilchennummer
#    2. Index - Koordinatenrichtung
#    3. Index - Zeitpunkt
r = r.reshape(n_teilchen, n_dim, -1)
v = v.reshape(n_teilchen, n_dim, -1)

# Erzeuge eine Figure.
fig = plt.figure(figsize=(12, 4))

# Erzeuge eine Axes für die animierte Darstellung der
# Masse-Feder-Kette.
ax_teilchen = fig.add_subplot(2, 1, 1)
ax_teilchen.set_xlim(-abstand, (n_teilchen + 1) * abstand)
ax_teilchen.set_ylim(-1.2 * A_tran, 1.2 * A_tran)
ax_teilchen.set_ylabel('$y$ [m]')
ax_teilchen.tick_params(labelbottom=False)
ax_teilchen.grid()

# Erzeuge Punktplots für die Teilchenpositionen.
plot_teilchen, = ax_teilchen.plot([], [], 'ob')
plot_teilchen_anregung, = ax_teilchen.plot([], [], 'or')

# Erzeuge ein Textfeld für die Angabe des Zeitpunkts.
text_zeit = ax_teilchen.text(0.97, 0.97, '',
                             transform=ax_teilchen.transAxes,
                             horizontalalignment='right',
                             verticalalignment='top')

# Erzeuge eine zweite Axes für die animierte Darstellung der
# transversalen und longitudinalen Auslenkung.
ax_auslenkung = fig.add_subplot(2, 1, 2)
ax_auslenkung.set_xlim(-abstand, (n_teilchen + 1) * abstand)
ax_auslenkung.set_ylim(-1.2 * max(A_tran, A_long),
                       +1.2 * max(A_tran, A_long))
ax_auslenkung.set_ylabel('$u$ [m]')
ax_auslenkung.set_xlabel('$x$ [m]')
ax_auslenkung.grid()

# Erzeuge je einen Linienplot für die Momentanauslenkung.
plot_welle_trans, = ax_auslenkung.plot([], [], label='trans')
plot_welle_long, = ax_auslenkung.plot([], [], label='long')

# Füge eine Legende hinzu.
ax_auslenkung.legend(loc='upper left')


def update(n):
    """Aktualisiere die Grafik zum n-ten Zeitschritt."""
    # Aktualisiere die Position der simulierten Teilchen.
    plot_teilchen.set_data(r[:, :, n].T)

    # Aktualisiere die Position des anregenden Teilchens.
    plot_teilchen_anregung.set_data(anregung(t[n]).reshape(-1, 1))

    # Erzeuge ein Array der Auslenkungen aller Teilchen und ein
    # Array der x-Positionen der Ruhelagen.
    auslenkungen = np.concatenate([anregung(t[n]).reshape(1, n_dim),
                                   r[:, :, n] - r0])
    ruhelage_x = np.concatenate(([0], r0[:, 0]))

    # Aktualisiere den Plot für die Transversal- und
    # Longitudinalwelle.
    plot_welle_long.set_data(ruhelage_x, auslenkungen[:, 0])
    plot_welle_trans.set_data(ruhelage_x, auslenkungen[:, 1])

    # Aktualisiere die Zeitangabe.
    text_zeit.set_text(f'$t$ = {t[n]:.2f} s')

    return (plot_teilchen, plot_teilchen_anregung,
            plot_welle_trans, plot_welle_long, text_zeit)


# Erstelle die Animation und zeige die Grafik an.
ani = mpl.animation.FuncAnimation(fig, update, frames=t.size,
                                  interval=30, blit=True)
plt.show()