import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import lil_matrix
from scipy.sparse import diags
from scipy.sparse.linalg import splu

# Step 1
N = 10
faces = np.linspace(0, 1., N+1)
nodes = 0.5*(faces[1:] + faces[:-1])
delta = faces[1]-faces[0]

Tb = 100.0
hP_over_kA = 25
n2 = hP_over_kA
L = 1.0
T_inf = 20.0

# Set up linear system of equations
#A = np.zeros((N, N))      # use a dense matrix
A = lil_matrix((N, N))     # or more memory efficient, use a sparse
b = np.zeros(N)

# ap*T_P - aw*T_W - ae*T_E = 0
for i in range(1, N-1):
    aw = 1. / delta
    ae = 1. / delta
    sp = -n2*delta
    su = n2*delta*T_inf
    ap = aw + ae - sp
    su = n2*delta*T_inf
    A[i, i-1] = -aw
    A[i, i]   = ap
    A[i, i+1] = -ae
    b[i] = su

# Node 1
ae = 1. / delta
aw = 0
sp = -n2*delta -2./delta
su = n2*delta*T_inf + 2.*Tb/delta
ap = aw + ae - sp
A[0, 0] = ap
A[0, 1] = -ae
b[0] = su

# Node N
aw = 1. / delta
ae = 0
sp = -n2*delta
su = n2*delta*T_inf
ap = aw + ae - sp
A[-1, -1] = ap
A[-1, -2] = -aw
b[-1] = su

# Now solve
x = nodes
exact = T_inf + (Tb-T_inf)*np.cosh(np.sqrt(n2)*(L-x))/np.cosh(np.sqrt(n2)*L)

#T = np.linalg.solve(A, b) # dense matrix solve
lu = splu(A.tocsc())       # sparse matrix solve through LU factorization
T = lu.solve(b)            #

plt.plot(nodes, T, 'b', x, exact, 'r:')
plt.show()