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

plot = False
errEst = True

def solve(n):
    h = 1./(n+1)
    #Defining coefficient matrix A = a_ij
    B = 4*np.eye(n)-np.eye(n,k=1)-np.eye(n,k=-1)
    A = np.kron(np.eye(n),B)-np.eye(n*n,k=n)-np.eye(n*n,k=-n)

    #Defining b (in Aw = b)
    x = np.linspace(h,1-h,n)
    y = np.linspace(h,1-h,n)
    X,Y = np.meshgrid(x,y)
    g = 2*h*h*(X*(1-X)+Y*(1-Y))
    b = g.flatten() #Converts from matrix to array.

    #Solving system (Remains to implement algorithm 7.4)
    w = np.linalg.solve(A,b)
    v = np.reshape(w,(n,n)) #Converts from array to matrix
    return X,Y,v

#plot (Warning???)
if(plot):
    X,Y,v = solve(20)
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X, Y, v, rstride=2, cstride=2, color='b')
    plt.show()


#Compare solution with exact
if(errEst):
    for n in [2,5,10,30]:
        X,Y,v = solve(n)
        h = 1./(n+1)
        u = X*(1-X)*Y*(1-Y)
        err = np.amax(np.absolute(u-v))
        errEs = h*h/12 #Error est by theorem 7.2
        print "n = %d: error: %g error estimate %g" % (n,err, errEs)