"""
Damped spring-mass oscillator (2nd order oscillator)

y'' + 2 eps wo  y' + wo^2 y = 0 with wo^2 = k/m

with the initial conditions y(t=0) = 1, y'(t=0) = 0.
"""
import numpy as np
import matplotlib.pyplot as plt

from scipy.integrate import odeint

plt.rcParams['font.size'] = 14

mass = 0.5
kspring = 4
cviscous = 0.4

nu = cviscous / mass
omega = kspring / mass

def get_derivative(vec_yy, time, nu, omega):
    return (vec_yy[1], -nu * vec_yy[1] - omega * vec_yy[0])

vec_t = np.linspace(0, 10, 100)
vec_yy = odeint(get_derivative, (1, 0), vec_t, args=(nu, omega))

fig, ax = plt.subplots()
ax.grid()

ax.plot(vec_t, vec_yy[:, 0], label=r'$y(t)$ [m]')
ax.plot(vec_t, vec_yy[:, 1], label=r'$\dot{y}(t)$ [m/s]')
ax.set_xlabel(r'$t$ [s]')
ax.legend()

t_mean = vec_t.mean()

spring = ax.plot([t_mean] * 2, [0, vec_yy[0, 0]], 'k', lw=3)[0]
big_dot = ax.plot(t_mean, vec_yy[0, 0], 'ko', ms=20)[0]
dot1 = ax.plot(vec_t[0], vec_yy[0, 0], 'r*')[0]
dot2 = ax.plot(vec_t[0], vec_yy[0, 1], 'g*')[0]

def animate(it):
    it = np.clip(int(round(it)), 0, len(vec_t)-1)
    tt = vec_t[it]
    spring.set_data([t_mean]*2, [0, vec_yy[it, 0]])
    big_dot.set_data([t_mean, vec_yy[it, 0]])
    dot1.set_data([tt, vec_yy[it, 0]])
    dot2.set_data([tt, vec_yy[it, 1]])

    fig.canvas.draw_idle()

if 1:
    from matplotlib.widgets import Slider

    plt.subplots_adjust(bottom=0.2)

    axts = fig.add_axes([0.25, .03, 0.50, 0.02])
    ts = Slider(axts, 'step', 0, len(vec_t)-1, valinit=0, valfmt='%d')
    ts.on_changed(animate)

else:
    import matplotlib.animation as animation

    ani = animation.FuncAnimation(fig, animate, np.arange(len(vec_t)),
                                  interval=(vec_t[1] - vec_t[0]) * 1000)

plt.show()