# Calculate and plot the 1-d spring pendulum energy and trajectory
# Isaac Lenton
# August 2017
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

mass = 1.0;             # Mass of pendulum (kg)
acceleration = 0.0;     # Gravitational acceleration (m/s/s)
springConst = 1.0;      # Spring constant (kg/s/s)
x0 = -1.0;              # Initial position (m)

dt = 1.0e-3;            # Simulation step size
outputRate = 100;       # Number of steps between output cycles
outputLength = 100;     # Length of simulation

potential = lambda x: -mass*acceleration*x + 0.5*springConst*x**2;
eqmotion = lambda x: mass*acceleration - springConst*x;
kinetic = lambda v: 0.5*mass*v**2;

tt = [0.0];
xt = [x0];
vt = [0.0];
pt = [potential(x0)];
kt = [kinetic(vt[0])];

# plot the potential, initial energy and show the current value
fig, ax = plt.subplots()
plt.subplot(1, 2, 1);
xp = np.linspace(-2, 2, 100);
_, _, line1 = plt.plot(xp, potential(xp), 'b',
    xp, potential(x0)*np.ones(np.size(xp)), 'g', xt, potential(x0), 'ro');
plt.xlabel('Position (m)');
plt.ylabel('Potential Energy (J)');

# Plot the kinetic and potential energy
plt.subplot(1, 2, 2);
tmax = 10;
line2, line3 = plt.plot(tt, pt, tt, kt);
plt.axis([0, tmax, 0, 0.5]);
plt.legend([line2, line3], ['Potential', 'Kinetic']);
plt.xlabel('Time (s)');
plt.ylabel('Energy (J)');

def animate(step_number):

    global tt, xt, vt, pt, kt

    # Evolve the particle state (first order, inaccurate at long times)
    xt.append(xt[-1])
    vt.append(vt[-1])
    for n in xrange(outputRate):
        oldvt = vt[-1];
        vt[-1] = vt[-1] + eqmotion(xt[-1])/mass*dt;
        xt[-1] = xt[-1] + oldvt*dt;
    pt.append(potential(xt[-1]));
    kt.append(kinetic(vt[-1]));
    tt.append(tt[-1]+dt*outputRate);
    
    # Update the plots (could also plot vt and xt too)
    line1.set_data(xt[-1], potential(xt[-1]))
    line2.set_data(tt, pt)
    line3.set_data(tt, kt)

    return line1, line2, line3
    
ani = animation.FuncAnimation(fig, animate, xrange(1, outputLength),
    interval=25, blit=True, repeat=False)
#ani.save('output.mp4')
plt.show()