import numpy as np
import matplotlib        # these two lines ensure that the
matplotlib.use("TkAgg")  # matplotlib backend is configured properly
import matplotlib.pyplot as plt

# PARAMETERS
k = 1.0      # spring constant
dt = 0.1     # time step
tmax = 10.0  # maximum simulation time

# VARIABLES
x = 0.3  # (initial) position
v = 0.0  # (initial) velocity
t = 0.0  # time

# OBSERVABLES
observables = []  # empty list for observables

# integrate equations of motion until t >= t_max:
while t < tmax:
    t += dt      # increase time
    F = -k * x   # compute new force
    v += F * dt  # compute new velocity
    x += v * dt  # compute new position
    # calculate observables:
    E_kin = 0.5 * v ** 2      # kinetic energy
    E_pot = 0.5 * k * x ** 2  # potential energy
    E_tot = E_kin + E_pot     # total energy
    # append all observables to the list:
    observables.append([t, x, E_tot, E_kin, E_pot])

# transform the observables to a 2d-NumPy-Array:
observables = np.array(observables)
ts = observables[:, 0]
xs = observables[:, 1]
E_tots = observables[:, 2]
E_kins = observables[:, 3]
E_pots = observables[:, 4]

# create an empty figure
plt.figure()

# create a plot of position vs. time
plt.subplot(1, 2, 1)
plt.xlabel(r"time $t$")
plt.ylabel(r"position $x$")
plt.plot(ts, xs, 'o-')

# create a plot of the different energy components vs. time 
plt.subplot(1, 2, 2)
plt.xlabel(r"time $t$")
plt.ylabel("energies")
plt.plot(ts, E_tots, 'o-', label=r"$E^{tot}$")
plt.plot(ts, E_kins, 'o-', label=r"$E^{kin}$")
plt.plot(ts, E_pots, 'o-', label=r"$E^{pot}$")
plt.legend()

# now show the plots
plt.show()