import numpy as np
from numpy import *
from scipy import interpolate
import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from matplotlib import animation

# PENDULUM 2
# g = 9.81
# A = 0.01
# w = 2.*pi*5
# l = 0.15

# PENDULUM 3
# g = 9.81
# A = 0.02
# w = 2.*pi*10
# l = 0.15

g = 9.81
A = 0.1
w = 2.*pi*10
l = 0.15

# ODE:
# l*ddtheta + (g + ddz)*sin(theta) = 0
# 
# x1 = theta
# x2 = dtheta
# 
# dx1 = x2 		( = dtheta)
# dx2 = -(g+A*sin(w*t))*sin(x1)/l 	( = ddtheta)
# 
# f(x1,x2) = [dx1, dx2]
def f(x1,x2):
	global g, A, w, l
	return [x2, -((g-A*pow(w,2)*sin(w*t))*sin(x1))/l]
	# return [x2,-(g*sin(x1))/l]

# RK4:
# k1 = h*f(x,y(x))
# k2 = h*f(x+h/2,y(x)+k1/2)
# k3 = h*f(x+h/2,y(x)+k2/2)
# k4 = h*f(x+h,y(x)+k3)
# y(x + h) = y(x) + k1/6 + k2/3 + k3/3 + k4/6

h = 0.005
hlast = h
desired_error = 1e-10
[x1, x2] = [pi-0.0001, 0]
# t = arange(0,10*pi,h)
t = 0
tmax = 10.*pi
tt = []
xx1 = []
xx2 = []
zz = []

while (t < tmax):
	dError = 1
	while (dError > 0):
		# full step:
		k1 = h*array(f(x1,x2))
		k2 = h*array(f(x1+k1[0]/2.,x2+k1[1]/2.))
		k3 = h*array(f(x1+k2[0]/2.,x2+k2[1]/2.))
		k4 = h*array(f(x1+k3[0],x2+k3[1]))
		[x1f, x2f] = [x1, x2] + k1/6. + k2/3. + k3/3. + k4/6.

		# two half steps:
		k1 = h/2.*array(f(x1,x2))
		k2 = h/2.*array(f(x1+k1[0]/2.,x2+k1[1]/2.))
		k3 = h/2.*array(f(x1+k2[0]/2.,x2+k2[1]/2.))
		k4 = h/2.*array(f(x1+k3[0],x2+k3[1]))
		[x1h, x2h] = [x1, x2] + k1/6. + k2/3. + k3/3. + k4/6.

		k1 = h/2.*array(f(x1h,x2h))
		k2 = h/2.*array(f(x1h+k1[0]/2.,x2h+k1[1]/2.))
		k3 = h/2.*array(f(x1h+k2[0]/2.,x2h+k2[1]/2.))
		k4 = h/2.*array(f(x1h+k3[0],x2h+k3[1]))
		[x1h, x2h] = [x1h, x2h] + k1/6. + k2/3. + k3/3. + k4/6.

		# update step size
		error_estimate = x1f-x1h
		dError = (error_estimate-desired_error)
		
		if (dError > 0):
			h = h*0.76
			hlast = h
			print("Error too large... Redo'ing step")
		elif (dError < 0):
			hlast = h
			h = h*1.02
			print("Expanding step size")

	print("t %.3f \t eE %.3e \t dE %.3e \t h %.3e" %(t,error_estimate,dError,h))
	t += hlast
	tt.append(t)
	[x1, x2] = [x1f, x2f]

	xx1.append(x1)
	xx2.append(x2)
	z = A*sin(w*t)
	zz.append(z)
	
# plt.figure()
# plt.plot(tt,xx1)
# plt.plot(tt,zz)

# Resample the datapoints in linear time
k = linspace(0,tmax,3000)
theta = interp(k,tt,xx1)
height = interp(k,tt,zz)
xs = []
ys = []

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = fig.add_subplot(111,xlim=(-0.5, 0.5), ylim=(-0.5, 0.5))
frame = plt.gca()
frame.axes.get_xaxis().set_visible(False)
frame.axes.get_yaxis().set_visible(False)
ax.set_frame_on(False)
point, = ax.plot([], [],'b.',markersize=50.0)
line, = ax.plot([],[],lw=2)
platform, = ax.plot([],[],'k',lw=5)
trace, = ax.plot([],[],color='0.3',lw=1)
 
# initialization function: plot the background of each frame
def init():
    point.set_data([], [])
    platform.set_data([],[])
    line.set_data([],[])
    return line, point, platform,
 
# animation function.  This is called sequentially
def animate(i):
	global theta, height, l, xs, ys
	x = l*cos(theta[i]-pi/2.)
	y = l*sin(theta[i]-pi/2.)
	xs.append(x)
	ys.append(y+height[i])
	platform.set_data([-0.1, 0.1], [height[i], height[i]])
	line.set_data([0,x], [height[i],height[i]+y])
	point.set_data([x], [height[i]+y])
	trace.set_data([xs],[ys])
	return line, point, platform, trace
 
# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=500, interval=20, blit=True)
 
# this is how you save your animation to file:
# anim.save('pendulum4.gif', writer='imagemagick', fps=30)
 
plt.show()