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


num_t = 200
num_cells = 200
t = arange(num_t)
k1 = zeros(num_cells)
k1_last = zeros(num_cells)
k = arange(num_cells)
u = zeros([num_cells,num_t])

# r = (v**2 * del_t**2 / del_x**2)
r = 0.1
d = 0.05

# def update(j,i,u):
# 	u[j][i] = r*(u[j+1][i-1] - 2*u[j][i-1] + u[j-1][i-1]) + 2*u[j][i-1] - u[j][i-2]
# 	return u

def update(j,i,u):
	global k1, k1_last, d, r
	k1_last[j] = k1[j]
	k1[j] = u[j+1][i-1] - 2.*u[j][i-1] + u[j-1][i-1]
	u[j][i] = r*k1[j] + d*(k1[j]-k1_last[j]) + 2.*u[j][i-1] - u[j][i-2]
	return u

def init(u):
	u[num_cells/2] = 10
	return u

u = init(u)

for i in range(1,num_t):
	for j in range(1,num_cells-1):
		update(j,i,u)
		# u[num_cells/2] = 2*np.exp(-0.5*(np.power(-i/20.,2.)))
		# u[num_cells/2][i] = 3.
		# u[0] = 0.
		# u[num_cells-1] = 0.
	# print u[:,i]


# 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, num_cells), ylim=(-2, 2))
ax.axis("off")
point, = ax.plot([], [],'b')

# initialization function: plot the background of each frame
def init_ani():
    point.set_data([], [])
    return point,
 
# animation function.  This is called sequentially
def animate(i):
	global k
	x = u[:,i]
	point.set_data(k,x)
	return point,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init_ani,
                               frames=num_t, interval=20, blit=True)
 
# this is how you save your animation to file:
anim.save('1d_wave_damped.gif', writer='imagemagick', fps=30)
 
plt.show()