#! /usr/bin/env python
from numpy import *
import argparse
from matplotlib import pyplot as plt
from matplotlib import animation

def test(args):
	n = args.nx; dx = args.xs/n
	c = (args.dt*args.v/dx)**2
	u = zeros((args.nt+2,n))
	#first two steps have single nonzero value in mid
	u[0:2,n/2] = args.amp 
	#u[0:2,1:-1] = args.amp*exp(-.5*linspace(-5,5,n-2)**2)
	for ti in range(1,args.nt+1):
		u[ti+1,1:-1] = c*(u[ti,:-2]-2*u[ti,1:-1]+u[ti,2:]) + 2*u[ti,1:-1] - u[ti-1,1:-1]
		if args.g != 0.: #if damping
			u[ti+1,1:-1] += (args.g*args.dt/dx/dx)*(u[ti,2:]-u[ti-1,2:]-2*(u[ti,1:-1]-u[ti-1,1:-1])+u[ti,:-2]-u[ti-1,:-2])
	
	fig = plt.figure()
	ax = plt.axes(xlim=(-.5*args.xs,.5*args.xs), ylim=(-1, 1))
	plt.title('Damped wave equation $\partial^2 u /\partial t^2d = \partial^2 u / \partial x^2 + \gamma \partial/\partial t \partial^2 u / \partial x^2$\nwith $\Delta x=$%.3f and $\Delta t=$%.3f, $\gamma=$%.2f'%(dx,args.dt,args.g))
	line, = ax.plot([], [], lw=2)
	def init():
		line.set_data([], []);return line,
	def animate(i):
		x = linspace(-.5*args.xs, .5*args.xs, n); y = u[i]
		line.set_data(x, y); return line,
	anim = animation.FuncAnimation(fig, animate, init_func=init,
			frames=args.nt+2, interval=20, blit=False)
	anim.save('wave.gif', writer='imagemagick', fps=5)
	plt.show()

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument("-nx","--nx",  type=int, default=50, help="grid size")
	parser.add_argument("-xs","--xs", type=double, default=1., help="space extent")
	parser.add_argument("-amp","--amp", type=double, default=1., help="initial amplitude")
	parser.add_argument("-v","--v",  type=double, default=1., help="speed")
	parser.add_argument("-g","--g", type=double, default=0., help="gamma, damping")
	parser.add_argument("-dt","--dt", type=double, default=.01, help="time res")
	parser.add_argument("-nt","--nt", type=int, default=10, help="num time steps")
	args = parser.parse_args()
	test(args)