from Tkinter import *
from numpy import *
import time, Image, ImageTk

#window & problem sizes
SIZE = 500
ngrid = 100

#step sizes and diffusion constant
dt = 0.7
dx = 1
D = 1
alpha = D*dt/dx**2

#Vectors to handle the matrix inversion
a = ones(ngrid)*(-alpha)
b = ones(ngrid)*(1+2.0*alpha)
c = ones(ngrid)*(-alpha)
a[0] = 0; a[ngrid-1] = 0
b[0] = 1; b[ngrid-1] = 1
c[0] = 0; c[ngrid-1] = 0

cprime = zeros(ngrid)
dprime = zeros(ngrid)

u = zeros((ngrid, ngrid))
#random.seed(1)
#for i in range(0, ngrid, 1):
#	for j in range(0, ngrid, 1):
#		u[i,j] = random.uniform(0,2)
for i in range((int)(ngrid/2)-2, (int)(ngrid/2)+2, 1):
	for j in range((int)(ngrid/3)-2, (int)(ngrid/3)+2, 1):
		u[i,j] = 4
for i in range((int)(ngrid/4)-2, (int)(ngrid/4)+2, 1):
	for j in range((int)(ngrid/4)-2, (int)(ngrid/4)+2, 1):
		u[i,j] = 4

def paint(parent,canvas):
	global u, alpha   

	for j in range (1, ngrid-1, 1):
		cprime[0] = c[0]/b[0]
		dprime[0] = (u[0,j]*(1-2.0*alpha)+u[0,j+1]*alpha+u[0,j-1]*alpha)/b[0]
		for i in range(1, ngrid, 1):
			cprime[i] = c[i]/(b[i]-a[i]*cprime[i-1])
			dprime[i] = (u[i,j]*(1-2.0*alpha)+u[i,j+1]*alpha+u[i,j-1]*alpha-a[i]*dprime[i-1])/(b[i]-a[i]*cprime[i])
		u[ngrid-1, j] = dprime[ngrid-1]
		for i in range (ngrid-2, 0, -1):
			u[i,j] = dprime[i] - cprime[i] * u[i+1,j]		
	for i in range (1, ngrid-1, 1):
		cprime[0] = c[0]/b[0]
		dprime[0] = (u[i,0]*(1-2.0*alpha)+u[i+1,0]*alpha+u[i-1,0]*alpha)/b[0]
		for j in range(1, ngrid, 1):
			cprime[j] = c[j]/(b[j]-a[i]*cprime[j-1])
			dprime[j] = (u[i,j]*(1-2.0*alpha)+u[i+1,j]*alpha+u[i-1,j]*alpha-a[j]*dprime[j-1])/(b[j]-a[j]*cprime[j])
		u[i, ngrid-1] = dprime[ngrid-1]
		for j in range (ngrid-2, 0, -1):
			u[i,j] = dprime[j] - cprime[j] * u[i,j+1]		
	
	#Tkinter-----
	intensity = uint32(255*(1+u)/2.0)
	grayscale = (1 << 16)*intensity + (1 << 8)*intensity + (1 << 0)*intensity
	img = Image.fromarray(grayscale,mode="RGBX")
	img = img.resize((SIZE,SIZE))
	tkimg = ImageTk.PhotoImage(img)
	canvas.create_image(SIZE/2,SIZE/2,image=tkimg)
	root.update()
	time.sleep(0.0001)
	parent.after_idle(paint,parent,canvas)
	#Tkinter-----
root = Tk()
root.title('Solution of diffusion equation using ADI')
root.bind('q','exit')
canvas = Canvas(root, width=SIZE, height=SIZE, background='white')
canvas.pack()
root.after(100,paint,root,canvas)
root.mainloop()