#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Feb  6 14:41:09 2020

Simulate then animate the simple pendulum motion.
Method: Euler 1st order to indegrate ODE.
This gives large errors: the simulation is inaccurate.

@author: Prof Roy
"""

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
import scipy

########## parameters ####################
g = 9.81 # [m/s^2]
l = 0.5 # [m]
m = 0.5 # [kg]
gamma = 0 # damping coefficient
omega0 = np.sqrt(g/l)

############# initial conditions  ################

theta= 170 # [degree]  # initial angle
theta = theta*np.pi/180
omega = 0  # initial angular velocity

# given the initial condition, we can compute the period of oscillations
# by a quadrature

def f(x):
    return 1/np.sqrt(np.cos(x)-np.cos(theta))

Tper = scipy.integrate.quad(f ,0,theta )
T0 = 2.0*np.sqrt(2.0)*Tper[0]/omega0
Tmax = 2*T0 # max time of simulation
print('Tmax=',Tmax)

xmax =  1.05*l
ymax =  1.05*l
Nstep= 800   # number of time steps /number of frames
dt= Tmax/Nstep  # time increment
d=0.05
print('period = ',T0 )

# First set up the figure, the axis, and the line object we want to animate

fig = plt.figure()
ax = plt.axes(xlim=(-xmax, xmax), ylim=(-ymax, ymax))

line, = ax.plot([], [], 'o', ls='-', ms=8, markevery=[0,-1])
xR,yR = -d,-d # lower left corner of rectangle
dx,dy= 2*d, 2*d  # dimensions of rectangle
rect = plt.Rectangle((xR,yR), dx,dy, facecolor= 'black', edgecolor= 'black')
ax.add_patch(rect)
plt.axis('off')
ax.set_aspect('equal')

#equations of motion

def F(theta,omega):
    dthetadt = omega
    domegadt = -g/l* np.sin(theta) - gamma/m *omega
    return dthetadt, domegadt

# initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,

# animation function: compute pendulum position in Cartesian coordinate 
# at each time step. Each new position is found by Euler's method.
def animate(i):
    global theta, omega
    Fa,Fb = F(theta,omega)
    theta = theta + Fa*dt # Euler's method
    omega = omega + Fb*dt
    theta = (theta + np.pi) % (2*np.pi) - np.pi # bring theta back to [-pi, pi]
    x,y = l*np.sin(theta), -l*np.cos(theta)
    line.set_data([0, x], [0,y]) # modify the pendulum's position    
    return line,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=Nstep, interval=20, blit=True)

# save the animation as an mp4.  This requires ffmpeg or mencoder to be
# installed.  The extra_args ensure that the x264 codec is used, so that
# the video can be embedded in html5.  
anim.save('Pendulum_animation_Euler.mp4', fps=30, extra_args=['-vcodec', 'libx264'])
plt.show()