#! /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 *
#using sor to solve poisson's equation via diffusion

def test(args):
	n = args.nx; dx = args.xs/n
	u = zeros((args.nt+1,n,n)) #use twice time steps for half-stepping
	u[:,:,0] = zeros(n)
	u[:,0,:] = zeros(n)
	u[:,:,-1] = ones(n)
	u[:,-1,:] = ones(n)
	rho = zeros_like(u[0]) #since this is laplace, rho = 0
	for ti in range(0,args.nt):
		u[ti+1] = sor(u[ti],rho,args.alpha,dx)

	fig = plt.figure()
	ax = plt.axes(xlim=(-.5*args.xs,.5*args.xs), ylim=(-.5*args.xs,.5*args.xs))
	ax.set_aspect(1.)
	x = linspace(-.5*args.xs, .5*args.xs, n); y = linspace(-.5*args.xs, .5*args.xs, n);
	X,Y = meshgrid(x,y)
	def animate(i):
		cont = plt.contourf(x, y, u[i], linspace(-0.001,1.001,25))
		return cont
	#anim = animation.FuncAnimation(fig, animate,
	#		frames=args.nt, interval=20, blit=False)
	#anim.save('laplace.gif', writer='imagemagick', fps=5)
	plt.contourf(x, y, u[-1], linspace(-0.001,1.001,25))
	plt.title('Solving Laplace equation with SOR on $\partial u /\partial td = \partial^2 u / \partial x^2$ with $\\alpha=$%.2f'%(args.alpha))
	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=100., 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("-alpha","--alpha", type=float, default=.1, help="alpha, relaxation parameter")
	args = parser.parse_args()
	test(args)