#! /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, using ADI operator splitting for 2d
import png,itertools

def read_png(filename):
	r = png.Reader(filename)
	r.chunk()
	phys = r.chunk()
	r = png.Reader(filename)
	data = png.Reader(filename).asDirect()
	bitdepth = data[3]['bitdepth']; assert(bitdepth==8)
	png_array = vstack(itertools.imap(uint8, data[2])).reshape(data[0],data[1],1)
	return png_array[...,0],bitdepth,data[0],data[1]

def test(args):
	n = args.nx; dx = args.xs/n
	alpha = args.dt/dx/dx
	u = zeros((2*args.nt+1,n,n)) #use twice time steps for half-stepping
	if args.filename=="":
		#u[0,n/2,n/2] = 1
		tx = linspace(-5,5,n-2); ty = linspace(-5,5,n-2).reshape(-1,1)
		u[0:1,1:-1,1:-1] = exp(-.5*(tx**2+ty**2))
	else:
		image,bitdepth,nx,ny = read_png(args.filename)
		from scipy.ndimage.interpolation import zoom
		u[0,1:-1,1:-1] = zoom(image[::-1],((n-2)/nx,(n-2)/ny))/float(2**bitdepth-1)
	for ti in range(0,args.nt):
		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.
		d = u[2*ti] + alpha*( roll(u[2*ti],-1,axis=1)-2*u[2*ti]+roll(u[2*ti],1,axis=1) )
		u[2*ti+1] = tridiagonal_system(a,b,c,d)
		d = u[2*ti+1] + alpha*( roll(u[2*ti+1],-1,axis=0)-2*u[2*ti+1]+roll(u[2*ti+1],1,axis=0) )
		u[2*(ti+1)] = tridiagonal_system(a,b,c,d.T).T

	fig = plt.figure()
	ax = plt.axes(xlim=(-.5*args.xs,.5*args.xs), ylim=(-.5*args.xs,.5*args.xs))
	plt.title('2d 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))
	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):
		z = u[1+2*i]
		cont = plt.contourf(x, y, z, linspace(-0.001,1.001,25))
		return cont
	anim = animation.FuncAnimation(fig, animate,
			frames=args.nt, interval=20, blit=False)
	#anim.save('sam-heat.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("-filename","--filename",default="",help="PNG for input")
	args = parser.parse_args()
	test(args)