from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def f():
return 4

a = 0
b = 1

c = 0
d = 2

n = 50
m = 50

h = (b-a)/n
k = (d-c)/m

#x = np.linspace(0.02, 0.98, n-1)
x = np.zeros(n-1)
for i in range(0, n-1):
x[i] = a + (i+1)*h
#y = np.linspace(0.04, 1.96, m-1)
y = np.zeros(m-1)
for i in range(0, m-1):
y[i] = c + (i+1)*k


w = np.zeros((n-1, m-1))
lam = h**2 / k**2
mu = 2*(1 + lam)
l = 1
N = 100

def g(x, y):
return (x - y)**2

while l < N:
z = (-4*h**2 + g(a, y[-1]) + lam*g(x[0], d) + lam*w[0, -2] + w[1, -1])/mu
Norm = abs(z - w[0, -1])
w[0, -1] = z

for i in range(1, n-2):
z = (-4*h**2 + lam*g(x[i], d) + w[i-1, -1] + w[i+1, m-2] + lam*w[i, m-3])/mu
if abs(w[i, m-2] - z) > Norm:
Norm = abs(w[i, m-2] - z)
w[i, m-2] = z

z = (-4*h**2 + g(b, y[-1]) + lam*g(x[-1], d) + w[n-3, m-2] + lam*w[n-2, m-3])/mu
if abs(w[n-2, m-2]-z) > Norm:
Norm = abs(w[n-2, m-2] - z)
w[n-2, m-2] = z

for j in range(m-3, 0, -1):
z = (-4*h**2 + g(a, y[j]) + lam * w[0, j+1] + lam * w[0, j-1] + w[1, j])/mu
if abs(w[0,j] - z) > Norm:
Norm = abs(w[0,j] - z)
w[0,j] = z

for i in range(1, n-3):
z = (-4*h**2 + w[i-1, j] + lam*w[i,j] + w[i+1,j] + lam*w[i,j-1])/mu
if abs(w[i,j] - z) > Norm:
Norm = abs(w[i,j] - z)
w[i,j] = z

z = (-4*h**2 + g(b, y[j]) + w[n-3,j] + lam*w[n-2,j+1] + lam*w[n-2,j-1])/mu
if abs(w[n-2,j] - z) > Norm:
Norm = abs(w[n-2,j] - z)
w[n-2,j] = z

z = (-4*h**2 + g(a, y[0]) + lam*g(x[0],c) + lam*w[0,1] + w[1,0])/mu
if abs(w[0,0]-z) > Norm:
Norm = abs(w[0,0] - z)
w[0,0] = z

for i in range(1, n-3):
z = (-4*h**2 + lam*g(x[i], c) + w[i-1,0] + lam*w[i,1] + w[i+1,0])/mu
if abs(w[i,0] - z) > Norm:
Norm = abs(w[i,0]-z)
w[i,0]=z

z = (-4*h**2 + g(b,y[0]) + lam*g(x[-1],c) + w[-2,0] + lam*w[-1,1])/mu
if abs(w[-1,0] - z) > Norm:
Norm = abs(w[-1,0] - z)
w[-1,0] = z

l += 1

X, Y = np.meshgrid(x, y)
Sol = g(X,Y)

fig = plt.figure()
ax1 = fig.add_subplot(121, projection='3d')

ax2 = fig.add_subplot(122, projection='3d')

ax1.plot_wireframe(X, Y, Sol)
ax1.set_title('Exacta')
plt.xlabel('$x$')
plt.ylabel('$y$')
ax2.plot_wireframe(X, Y, w)
ax2.set_title('Aproximada')

print 'Desviacion promedio: $|\hat{u}-u|$'
print np.std(w - Sol)