#! /usr/bin/env python
from __future__ import division
from numpy import *
import argparse
from matplotlib import pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import axes3d

def test(args):
	n = args.nx; dx = args.xs/n
	c = (args.dt*args.v/dx)**2
	print 'dx=%.4f, dt=%.4f, c=%.4f'%(dx,args.dt,c)
	u = zeros((args.nt+2,n,n))
	tx = linspace(-5,5,n-2); ty = linspace(-5,5,n-2).reshape(-1,1)
	u[0:2,1:-1,1:-1] = args.amp*exp(-.5*(tx**2+ty**2))
	for ti in range(1,args.nt+1):
		u[ti+1,1:-1,1:-1] = c*(u[ti,:-2,1:-1]+u[ti,1:-1,:-2]-4*u[ti,1:-1,1:-1]+u[ti,1:-1,2:]+u[ti,2:,1:-1])+ 2*u[ti,1:-1,1:-1]- u[ti-1,1:-1,1:-1]

	x = linspace(-.5*args.xs, .5*args.xs, n); y = linspace(-.5*args.xs, .5*args.xs, n) 
	X,Y = meshgrid(x, y)
	plt.ion()
	fig = plt.figure()
	ax = fig.add_subplot(111, projection='3d')
	ax.set_zlim(-1,1)
	wframe = None
	for i in range(args.nt+2):
		oldcol = wframe
		z = u[i]
		wframe = ax.plot_wireframe(X, Y, z, rstride=2, cstride=2)
		# Remove old line collection before drawing
		if oldcol is not None:
			ax.collections.remove(oldcol)
		plt.savefig('frames/frame%03d.png'%i)
		plt.draw()


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