import numpy as np
from numpy import *
import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

N = 50
T = 100

def boundary(u):
	global N
	u[:,0] = 0#sin(arange(N)/5.)
	u[0,:] = 0#sin(arange(N)/5.)
	u[N-1,:] = 1
	u[:,N-1] = -1
	return u


u = zeros([N,N])
boundary(u)
uu = zeros([T,N,N])


# alpha = 1.0 # GAUSS-SEIDEL
# alpha = 1.93 # JUST STABLE
alpha = 2.0 # UNSTABLE

# def update(i,j,k,u):
# 	u[i][j][k] = (1-alpha)*u[i-1][j][k] + alpha/4*(u[i-1][j+1][k] + u[i-1][j-1][k] 
# 		+ u[i-1][j][k+1] + u[i-1][j][k-1]) #- alpha/4
# 	return u

def update(j,k,u):
	u[j][k] = (1-alpha)*u[j][k] + alpha/4*(u[j+1][k] + u[j-1][k] 
		+ u[j][k+1] + u[j][k-1]) #- alpha/4
	return u

# def init(u):
# 	u[num_cells/2] = 1.
# 	return u

print(u)
# u = init(u)
for i in range(1,T):
	for j in range(1,N-1):
		for k in range(1,N-1):
			u = update(j,k,u)

	u = boundary(u)
	uu[i] = u
			
	# print(u)


x, y = meshgrid(arange(0,N),arange(0,N))
# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
# ax = fig.add_subplot(111, projection='3d')#xlim=(0, N), ylim=(-0.1, 1))
# ax = axes3d.Axes3D(fig)
ax = Axes3D(fig)
wireframe = ax.plot_surface(x, y, uu[0,x,y], rstride=2, cstride=2,cmap=cm.jet)
# wireframe = ax.plot_wireframe(x, y, uu[0,x,y], rstride=2, cstride=2)
ax.set_zlim(-1,1)
# ax = plt.gca(projection='3d'); ax._axis3don = False
ax.axis("off")
# frame = plt.gca()
# frame.axes.get_xaxis().set_visible(False)
# frame.axes.get_yaxis().set_visible(False)
ax.set_frame_on(False)
# initialization function: plot the background of each frame
# def init_ani():
#     point.set_data([], [])
#     point.set_3d_properties([])
#     return point,
 
# # animation function.  This is called sequentially
def animate(i, ax, fig):
	global x,y,u
	ax.cla()
	wireframe = ax.plot_surface(x, y, uu[i,x,y], rstride=2, cstride=2,cmap=cm.jet)
	# wireframe = ax.plot_wireframe(x, y, uu[i,x,y], rstride=2, cstride=2)
	ax.set_zlim(-1,1)
	# ax = plt.gca(projection='3d'); ax._axis3don = False
	ax.axis("off")
	ax.set_frame_on(False)
	# frame.axes.get_xaxis().set_visible(False)
	# frame.axes.get_yaxis().set_visible(False)
	# ax.set_frame_on(False)
	ax.view_init(elev=25., azim=140)
	return wireframe,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, frames=T, fargs=(ax, fig), interval=20)
 
# this is how you save your animation to file:
anim.save('2d_sor_unstable.gif', writer='imagemagick', fps=30)
 
plt.show()