import numpy as np
from scipy import integrate

import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import cnames
from matplotlib import animation

## HOMOGENEOUS:
# m*ddx + y*dx + k*x = 0
# x1 = x
# x2 = dx = dx1
# m*dx2 + y*x2 + k*x1 = 0
# dx1 = x2
# dx2 = (-y*x2 - k*x1)/m

## INHOMOGENEOUS:
# m*ddx + y*dx + k*x = e^iwt = s/(s^2+w^2) + i*w/(s^2+w^2)
# x1 = x
# x2 = dx = dx1
# m*dx2 + y*x2 + k*x1 = cos(w*t) + i*sin(w*t)
# dx1 = x2
# dx2 = (cos(w*t) + i*sin(w*t) - y*x2 - k*x1)/m


def deriv((x1,x2),t0,m=1.,k=1.,y=0.1):
	return [x2, (-y*x2 - k*x1)/m]

x0 = [0, 10] 
t = np.linspace(0, 100, 10000)
x_t = np.asarray(integrate.odeint(deriv, x0, t))

# I'm sure there's a better way to do this...
x1s = []
x2s = []
for xi in x_t:
	x1, x2 = xi[:].T
	x1s.append(x1)
	x2s.append(x2)

# Set up figure
fig = plt.figure()
ax = fig.add_subplot(111)
ax.autoscale(True,'both')

# Plot it
ax.plot(t,x1s)#,t,x2s)
ax.set_title('Damped Harmonic Oscillator - Homogeneous Solution')
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude')

# ANIMAITON (NOT CURRENTLY WORKING)
# fig = plt.figure()
# ax = fig.add_subplot(111)
# point, = ax.plot([], [], lw=5)
# ax.autoscale(True,'both')

# def init():
# 	point.set_data([],[])
# 	return point,

# def animate(i):
# 	i = (2 * i) % x_t.shape[1]
# 	point.set_data(t[i],x1s[i])
# 	return point,

# anim = animation.FuncAnimation(fig, animate, init_func=init,
#                                frames=100, interval=30, blit=True)

plt.show()

# # initialization function: plot the background of each frame
# def init():
#     for line, pt in zip(lines, pts):
#         line.set_data([], [])

#         pt.set_data([], [])
#     return lines + pts

# # animation function.  This will be called sequentially with the frame number
# def animate(i):
#     # we'll step two time-steps per frame.  This leads to nice results.
#     i = (2 * i) % x_t.shape[1]

#     for line, pt, xi in zip(lines, pts, x_t):
#         x1, x2 = x_t[:i].T
#         line.set_data(x1, x2)

#         pt.set_data(x1[-1:], x2[-1:])

#     fig.canvas.draw()
#     return lines + pts

# # instantiate the animator.
# anim = animation.FuncAnimation(fig, animate, init_func=init,
#                                frames=500, interval=30, blit=True)

# # Save as mp4. This requires mplayer or ffmpeg to be installed
# #anim.save('lorentz_attractor.mp4', fps=15, extra_args=['-vcodec', 'libx264'])
# # anim.save('lorentz.gif', writer='imagemagick', fps=30)

# plt.show()