# Poisson equation solved iteratively by SOR
from numpy import ones_like,linspace,ones,empty,array, zeros
from pylab import imshow, show, colorbar, clf, ion, ioff, draw, pause, title
from time import sleep
def f(xxx,yyy): return 0.*xxx
def left  (y): return 0.*ones_like(y)#y
def right (y): return 0.*ones_like(y)#y*y
def top   (x): return 100*array(x<a+(b-a)/2,dtype=float)
def bottom(x): return 0.*ones_like(x)#x+1
M = 100; a = 0.; b = float(M+1)  
N = 100; c = 0.; d = float(N+1)
h = (b-a)/float(M); m = M-1; x = linspace(a,b,m+2); xinterior = x[1:-1]
k = (d-c)/float(N); n = N-1; y = linspace(c,d,n+2); yinterior = y[1:-1]
assert(h==k) # code below assumes this
print('h,k',h,k)
mn = m*n
uinterior = zeros((n,m))
u = empty((n+2,m+2))
u[1:-1,1:-1] = uinterior
u[:, 0] = left  (y)
u[:,-1] = right (y)
u[ 0,:] = top   (x)
u[-1,:] = bottom(x)

#w = 2.001
w = float(input('Choose w: '))
ion()

clf()
imshow(u,interpolation='nearest',cmap='rainbow')
title(f'w = {w}, iteration 0')
colorbar()
draw()
pause(1e-6)
print(k+1)
sleep(5)

for k in range(int(1e5)):

	# the following 3 lines should be done in C for 50+ times speedup
	for i in range(1,m+1):
		for j in range(1,n+1):
			u[i,j] = w*(u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1])*0.25 + (1.-w)*u[i,j]


	if k%1==0:
		clf()
		imshow(u,interpolation='nearest',cmap='rainbow')
		title(f'w = {w}, iteration {k+1}')
		colorbar()
		draw()
		pause(1e-6)
		print(k+1)
		#if k==0: sleep(10)
ioff()
show()