from numpy import *
from scipy.sparse import *
from scipy.sparse import linalg

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm


# n is the size of the grid of varialbles, i.e. not counting the boundary

# most straightforward way to define the matrix, entry by entry with loops
def laplace2d_mat_simp(n):
    n2 = n*n
    L = lil_matrix((n2, n2))
    # -4 for the central point of the finite difference, 1 for neighbors on different sides 
    # index of variable at  grid location (i,j) is ind = i + n*j
    #  for each point (i,j) not on the boundary we need (i,j) = ind  itself, (i-1,j) = ind-1,  (i+1,j) = ind+1, (i,j-1) = ind-n, (i,j+1) = ind +n  
    for j in range(0,n):
        for i in range(0,n):
            L[i+n*j,i+n*j] = -4.0;
    for j in range(0,n):
        for i in range(1,n):
            L[i+n*j,(i-1)+n*j] = 1.0;
    for j in range(0,n):
        for i in range(0,n-1):
            L[i+n*j,(i+1)+n*j] = 1.0; 
    for j in range(1,n):
        for i in range(0,n):
            L[i+n*j,i+n*(j-1)] = 1.0;
    for j in range(0,n-1):
        for i in range(0,n):
            L[i+n*j,i+n*(j+1)] = 1.0;
    return L.tocsr()

# more 
def laplace2d_mat_vec(n):
    vals_center = [-4.0]*(n*n) # list of -4 repeated n*n times
    vals_offcenter = [1.0]*(n*(n-1))
    
    row_ind_center =  [ i   + n*j     for j in range(0,n)   for i in range(0,n)]
    col_ind_center =  [ i   + n*j     for j in range(0,n)   for i in range(0,n)]
    
    row_ind_left   =  [ i   + n*j     for j in range(0,n)   for i in range(1,n)]
    col_ind_left   =  [(i-1)+ n*j     for j in range(0,n)   for i in range(1,n)]
    
    row_ind_right  =  [ i   + n*j     for j in range(0,n)   for i in range(0,n-1)]
    col_ind_right  =  [(i+1)+ n*j     for j in range(0,n)   for i in range(0,n-1)]
    
    row_ind_up     =  [ i   + n*j     for j in range(1,n)   for i in range(0,n)]
    col_ind_up     =  [ i   + n*(j-1) for j in range(1,n)   for i in range(0,n)]
    
    row_ind_down   =  [ i   + n*j     for j in range(0,n-1) for i in range(0,n)]
    col_ind_down   =  [ i   + n*(j+1) for j in range(0,n-1) for i in range(0,n)]
    
    vals    = vals_center +vals_offcenter + vals_offcenter + vals_offcenter + vals_offcenter
    row_ind = row_ind_center + row_ind_left + row_ind_right + row_ind_up + row_ind_down
    col_ind = col_ind_center + col_ind_left + col_ind_right + col_ind_up + col_ind_down

    return coo_matrix( (vals,(row_ind,col_ind))).tocsr()
 
 
def laplace2d_rhs(bnd_fnc,n):
    rhs = array([0.0]*(n*n))
    # indices of variables next to the fixed boundary
    h = 1.0/float(n+1) # grid step, assuming [0,1] x [0,1] domain
    bnd_ind_left   =  [  0  + n*j      for j in range(0,n)]
    bnd_pts_left   =  [ (0, (j+1)*h)      for j in range(0,n)]
    bnd_ind_right  =  [ n-1 + n*j      for j in range(0,n)]        
    bnd_pts_right  =  [ ((n+1)*h, (j+1)*h) for j in range(0,n)]        
    bnd_ind_up     =  [ i   + n*0      for i in range(0,n)]
    bnd_pts_up     =  [ ((i+1)*h,0)    for i in range(0,n)]
    bnd_ind_down   =  [ i   + n*(n-1)  for i in range(0,n)] 
    bnd_pts_down   =  [ ((i+1)*h,(n+1)*h)  for i in range(0,n)] 

    # for each grid point next to the boundary, subtract the boundary value
    # multiplied by the stencil coefficient (1.0 in our case) from the rhs
    # for corners two values are subtracted because they are in two lists 
    # for all other points, only one 
    rhs[bnd_ind_left ] = rhs[bnd_ind_left ]  -1.0*array(map(bnd_fnc,bnd_pts_left))
    rhs[bnd_ind_right] = rhs[bnd_ind_right]  -1.0*array(map(bnd_fnc,bnd_pts_right))
    rhs[bnd_ind_up   ] = rhs[bnd_ind_up   ]  -1.0*array(map(bnd_fnc,bnd_pts_up ))
    rhs[bnd_ind_down ] = rhs[bnd_ind_down ]  -1.0*array(map(bnd_fnc,bnd_pts_down ))
    return reshape(rhs, (n*n,1))
    
# for given boundary function bnd_fnc(x,y) compute solution of  Delta u(x,y) = 0, u(x,y) = bnd_fnc(x,y) on 
# the boundary of the square [0,1] x [0,1] and plot it 
# n is the number of points on the grid in each dimension including boundaries (number of squares is n-1, number of variables (n-2)**2
def laplace2d(bnd_fnc,n):
    h = 1.0/(n-1)
    L=  laplace2d_mat_vec(n-2)
    r = laplace2d_rhs(bnd_fnc, n-2)
    u = linalg.spsolve(L,r)
#    fig = plt.figure()
#    ax = fig.gca(projection='3d')    
#    X =  linspace(0,1,n-2)
#    Y =  linspace(0,1,n-2)
#    X, Y = meshgrid(X, Y)
#    surf = ax.plot_surface(X, Y, reshape(u,(n-2,n-2)), rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0)
#    ax.set_zlim(-2, 2)
#    plt.show()


# example:  use sin(10*(x+y)) as boundary condition, 30x30 grid
laplace2d(lambda (x): sin(10*(x[0]+x[1])),700)