import numpy 
import math
from matplotlib import pyplot
%matplotlib inline

# constants
g = 9.81
l = 1.0

def simulate_pendulum(phi_init, omega_init, numSteps, deltaT):
    t = numpy.zeros(numSteps)

    phi = numpy.zeros(numSteps)
    phi[0] = phi_init

    omega = numpy.zeros(numSteps)
    omega[0] = omega_init

    for n in range(1,numSteps):
        t[n] = deltaT * n # store the time for plotting.

        alpha_t = -(g/l) * math.sin(phi[n-1])
        omega_t = omega[n-1] + (delta_t * alpha_t)
        phi_t = phi[n-1] + (delta_t * omega_t)
        omega[n] = omega_t
        phi[n] = phi_t

    return (t, phi, omega)

# initial conditions
phi_0 = math.pi / 6.  # initial angle
omega_0  = 0.         # pendulum initially at rest

N = 1000              # time steps to run the simulation for
delta_t = 0.01        # simulation step size

(t, phi, omega) = simulate_pendulum(phi_0, omega_0, N, delta_t)

pyplot.figure(figsize=(10,4))
pyplot.xlabel('t', fontsize=14)
pyplot.ylabel(r'$\phi$', fontsize=14)
pyplot.hold(True)
pyplot.plot(t, phi)

# this gives the trajectories - the derivative of the state-space matrix
def dX_dt(X):
    #phi = X[0]
    #omega = X[1] = phi_dot
    #omega_dot = -(g/l)*math.sin(phi)
    #return phi_dot, omega_dot
    return numpy.array([ X[1], -(g/l)*numpy.sin(X[0]) ]);

x = numpy.linspace(-math.pi*1.25, math.pi*1.25, 21)
y = numpy.linspace(-8, 8, 20)

(X,Y) = numpy.meshgrid(x, y)
(DX, DY) = dX_dt([X, Y])

pyplot.figure(figsize=(10,5))
pyplot.xlabel(r'$\phi$', fontsize=14)
pyplot.ylabel(r'$\omega$', fontsize=14)
pyplot.quiver(X, Y, DX, DY)

# simulate a second trajectory with a larger initial angle
(t2, phi_large, omega_large) = simulate_pendulum(math.pi * 0.99, 0, N, delta_t)

coords = zip(omega, phi)
pyplot.hold(True)
pyplot.plot(phi, omega, color='b', label='$\phi_0=\pi/6$')
pyplot.plot(phi_large, omega_large, color='g', label='$\phi_0=0.99\pi$')
pyplot.legend()

from JSAnimation import IPython_display
from matplotlib import animation

fig = pyplot.figure(figsize=(8,5))
ax = pyplot.axes(xlim=(0,2), ylim=(-0.2,2))
line, = pyplot.plot([],[])
time_text = pyplot.text(0.02, 1.90, '')

def animate(i):
    phi_step = phi[i]
    line.set_data([1,1-math.sin(phi_step)],[1,1-math.cos(phi_step)])
    time_text.set_text('time = %.1f' % t[i])
    

animation.FuncAnimation(fig, animate,
                        frames=200, interval=20)

def simulate_pendulum_damped(phi_init, omega_init, numSteps, deltaT, damping):
    t = numpy.zeros(numSteps)

    phi = numpy.zeros(numSteps)
    phi[0] = phi_init

    omega = numpy.zeros(numSteps)
    omega[0] = omega_init

    for n in range(1,numSteps):
        t[n] = deltaT * n # store the time for plotting.

        alpha_t = (-(g/l) * math.sin(phi[n-1])) - (damping * omega[n-1])
        omega_t = omega[n-1] + (delta_t * alpha_t)
        phi_t = phi[n-1] + (delta_t * omega_t)
        omega[n] = omega_t
        phi[n] = phi_t

    return (t, phi, omega)

damping = 0.5
(t_damped, phi_damped, omega_damped) = simulate_pendulum_damped(phi_0, omega_0, N, delta_t, damping)

pyplot.figure(figsize=(10,4))
pyplot.xlabel('t', fontsize=14)
pyplot.ylabel(r'$\phi$', fontsize=14)
pyplot.hold(True)
pyplot.plot(t_damped, phi_damped)

# this gives the trajectories - the derivative of the state-space matrix
def dX_dt_damped(X, damping):
    #phi = X[0]
    #omega = X[1] = phi_dot
    #omega_dot = -(g/l)*math.sin(phi)
    #return phi_dot, omega_dot
    return numpy.array([ X[1], (-(g/l)*numpy.sin(X[0]))-(damping * X[1]) ]);

x = numpy.linspace(-1.25, 1.25, 21)
y = numpy.linspace(-2, 2, 20)

(X,Y) = numpy.meshgrid(x, y)
(DX, DY) = dX_dt_damped([X, Y], 0.5)

pyplot.figure(figsize=(10,5))
pyplot.xlabel(r'$\phi$', fontsize=14)
pyplot.ylabel(r'$\omega$', fontsize=14)
pyplot.quiver(X, Y, DX, DY)

coords = zip(omega_damped, phi_damped)
pyplot.hold(True)
pyplot.plot(phi_damped, omega_damped, label='b=0.5')
pyplot.legend()