import numpy as np
from scipy import integrate
from scipy import signal
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 =  s/(s^2+w^2) + i*w/(s^2+w^2)
# dx1 = x2
# dx2 = (cos(w*t) + i*sin(w*t) - y*x2 - k*x1)/m

# e^i*w*t = 1 / (s - i*w)


# X = e^i*w*t / (m*(i*w)^2 + y*i*w + k)
#   = e^i*w*t / ((-m*w^2 + k) + (y*w)i)
#   = e^i*w*t * ((-m*w^2 + k) - (y*w)i) / ((-m*w^2 + k)^2 + (y*w)^2)
#   = ((-m*w^2 + k)*e^i*w*t - (y*w)i*e^i*w*t) / ((-m*w^2 + k)^2 + (y*w)^2)
#   = 


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

m = 1
k = 1
y = 0.1


# FORCED
def forced():
	global m, k, y
	freq = np.arange(10e-1,10e1,10e-2)
	mags = []
	phases = []
	for w in freq:
		num = -np.power(w,2) + w*1j
		denum = m*np.power(w,4) - m*np.power(w,3) + k*np.power(w,2) - k*w

		mag = np.absolute(num)/np.absolute(denum)
		phase = np.angle(num/denum, True)

		mags.append(20*np.log10(mag))
		phases.append(phase)
		print(phase)

	
	fig = plt.figure()
	ax1 = fig.add_subplot(211)
	ax2 = fig.add_subplot(212, ylim=(130,180))
	ax1.semilogx(freq, mags)    # Bode magnitude plot
	ax2.semilogx(freq, phases)  # Bode phase plot

	ax1.set_title('Damped Harmonic Oscillator - With Forcing')
	ax2.set_xlabel('Frequency (rad/s)')
	ax1.set_ylabel('Amplitude')
	ax2.set_ylabel('Phase')


# UNFORCED

def unforced():
	global m, k, y

	sys = signal.lti(1, [m, y, k])

	w, mag, phase = signal.bode(sys)

	fig2 = plt.figure()
	ax3 = fig2.add_subplot(211)
	ax4 = fig2.add_subplot(212)
	ax3.semilogx(w,mag)
	ax4.semilogx(w,phase)

	ax3.set_title('Damped Harmonic Oscillator - Without Forcing')
	ax4.set_xlabel('Frequency (rad/s)')
	ax3.set_ylabel('Amplitude')
	ax4.set_ylabel('Phase (degrees)')


unforced()
plt.show()

# 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()