#! /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
import pyximport
pyximport.install(setup_args={"include_dirs":get_include()})
from tridiagonal import *
#diffusion

def test(args):
	n = args.nx; dx = args.xs/n
	alpha = args.dt/dx/dx
	u = zeros((args.nt+2,n))
	#u[0,n/2] = 1
	tx = linspace(-5,5,n-2)
	u[0:2,1:-1] = exp(-.5*tx**2)
	for ti in range(0,args.nt+1):
		if args.method == 'explicit':
			u[ti+1,1:-1] = alpha*(u[ti,:-2]-4*u[ti,1:-1]+u[ti,2:])+u[ti,1:-1]
		elif args.method == 'implicit':
			a = -alpha*ones(n)
			a[-1]=0.
			b = (1+2*alpha)*ones(n)
			b[0]=1.; b[-1]=1.
			c = -alpha*ones(n)
			c[0] = 0.
			u[ti+1] = invert_tridiagonal(a,b,c,u[ti])
		else:
			print "unrecognized method"

	fig = plt.figure()
	ax = plt.axes(xlim=(-.5*args.xs,.5*args.xs), ylim=(-.3, 1))
	line, = ax.plot([], [], lw=2)
	plt.title('Diffusion equation $\partial u /\partial t = \partial^2 u / \partial x^2$ \nwith $\Delta x=$%.3f and $\Delta t=$%.3f, $2\Delta t/\Delta x^2=$%.2f'%(dx,args.dt,2*args.dt/dx/dx))
	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('diffusion.gif', writer='imagemagick', fps=5)
	plt.show()


if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument("-nx","--nx",  type=int, default=500, help="grid size")
	parser.add_argument("-xs","--xs", type=float, default=500., help="space extent")
	parser.add_argument("-dt","--dt", type=float, default=.01, help="time res")
	parser.add_argument("-nt","--nt", type=int, default=50, help="num time steps")
	parser.add_argument("-method","--method", default="explicit", help="method to use: explicit or implicit")
	args = parser.parse_args()
	test(args)