import numpy as np   
import numpy.linalg as la   
import matplotlib.pyplot as plt 
from scipy import sparse
from scipy.sparse.linalg import spsolve 
from mpl_toolkits.mplot3d import Axes3D  # For 3-d plot
from matplotlib import cm
import time
newparams = {'figure.figsize': (8.0, 4.0), 'axes.grid': True,
             'lines.markersize': 8, 'lines.linewidth': 2,
             'font.size': 14}
plt.rcParams.update(newparams)

def A_laplace(M):
    # Construct the discrete laplacian, that is
    # A = blocktridiag(I,T,I) is a M^2xM^2 matrix where
    # T = tridiag(1,-4,1) is a MxM matrix
    M2 = M**2
    A = -4*np.eye(M2)            # The diagonal matrix
    for i in range(M2-1):
        A[i,i+1] = 1
        A[i+1,i] = 1
    for i in range(M-1,M2-1,M):
        A[i,i+1] = 0
        A[i+1,i] = 0
    for i in range(M2-M):      # The block sub- and sup diagonal
        A[i,i+M] = 1
        A[i+M,i] = 1
    return A

def plot_solution(x, y, U, txt='Solution', nfig = 1):
    # Plot the solution of the heat equation on the unit square
    fig = plt.figure(nfig)
    plt.clf()
    ax = fig.gca(projection='3d')
    X,Y = np.meshgrid(x, y)
    ax.plot_surface(X, Y, U, cmap=cm.coolwarm)
    ax.view_init(azim=30)              # Rotate the figure
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title(txt);
    plt.show()

def poisson_dirichlet(M):
    # Solve the Poisson equation by the 5-point formula on the unit square, with
    # Dirichlet boundary conditions.
    # Input: M, the number of interval

    # Set up the problem
    # Boundary conditions 
    def gs(x):                  # y=0
        return x**3
    def gn(x):                  # y=1
        return 2+x**3
    def gw(y):                  # x=0
        return 2*y**2
    def ge(y):                  # x=1
        return 1+2*y**2   

    # Source function
    def f(x,y):
        return 6*x+4

    # The grid
    x = np.linspace(0, 1, M+1)
    y = np.linspace(0, 1, M+1) 
    h = 1/M

    # Inner grid
    xi = x[1:-1]       
    yi = y[1:-1] 
    Xi, Yi = np.meshgrid(xi, yi)
    Mi = M-1       # Number of inner gridpoints
    Mi2 = Mi**2    # Total number of inner gridpoints (and unknowns)

    A = A_laplace(M-1)

    # The right hand side
    b = np.zeros(Mi2)
    # Boundary conditions
    b[0:Mi] = b[0:Mi]-gs(xi)                         # y=0
    b[Mi2-Mi:Mi2] = b[Mi2-Mi:Mi2] - gn(xi)           # y=1
    b[0:Mi2:Mi] = b[0:Mi2:Mi] - gw(yi)               # x=0
    b[Mi-1:Mi2:Mi] = b[Mi-1:Mi2:Mi] - ge(yi)         # x=1

    # Source function
    b = b+h**2*f(Xi,Yi).flatten()

    # Solve the linear system
    Ui = la.solve(A, b)      

    # Make an (M+1)x(M+1) array to store the solution,
    # including boundary
    U = np.zeros((M+1, M+1))

    # Reshape the solution vector, and insert into the solution matrix
    U[1:-1,1:-1] = np.reshape(Ui, (Mi,Mi)) 

    # include the boundaries
    U[0,:] = gs(x)
    U[M,:] = gn(x)
    U[:,0] = gw(y)
    U[:,M] = ge(y)

    return x, y, U

def sparse_demo(M):
    # Demonstate the benefits of using a sparse solver, by constructing the
    # discrete laplacian A, solve the system Ax=b for some vector b, and
    # compare the time used for solving the system by a full and a sparse
    # solver. 

    # NB! This solves a system of M^2 equations. 

    A = A_laplace(M)
    b = np.ones(M**2)   # Or some other arbitrary b-vector

    start = time.time()
    U = la.solve(A, b)            
    finish= time.time()
    print('Time used for solving a full system:    {:8.4f} sec'.format(finish-start))


    # Sparse solver
    As = sparse.csr_matrix(A)   # Convert the matrix to a sparse format (csr) 
    start = time.time()
    U = spsolve(As, b)          # Use a sparse linear solver
    finish = time.time()
    print('Time used for solving a sparse system:    {:8.4f} sec'.format(finish-start))

if __name__ == "__main__":
    x, y, U = poisson_dirichlet(10)
    X, Y = np.meshgrid(x,y)
    U_exact = X**3+2*Y**2
    plot_solution(x, y, U)
    plot_solution(x, y, U-U_exact, txt='Error', nfig=2)