# -*- coding: utf-8 -*-
"""
Created on Thu Aug  6 18:27:54 2020

Chapter 6 PDE: Excercise 2: Metallic plate and heat sources

2D Elliptic PDE
"""

"""
Heat distribution in a metallic plate
size: 1m x 1m
distribution: 
    d2T/d2x + d2T/d2y = f
    f = -Q/k (constant)
"""
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter


# define the space
x0, xf , nX = 0.0, 1.0, 100
y0, yf , nY = 0.0, 1.0, 100

# Vectors
X = np.linspace(x0, xf, nX)
Y = np.linspace(y0, yf, nY)
h = X[1] - X[0]
h2 = 1/(h**2)
dh = 1/(2*h) # 1/2dx and 1/2dy

# Size of the system
N = nX*nY
A = sp.lil_matrix((N,N))
b = np.zeros(N)

# Right hand side of the equation
Q = 100 # kW/m3 heat generation
k = 0.01 # W/mC thermal conductivity
f = -Q/k

# nested loop for the node mapping
for i in range(nX):
    for j in range(nY):
        # mapping relation center
        K = j*nY + i
        # neighbours p:plus m:minus
        KpN = (j+1)*nY + i              
        KpN2 = (j+2)*nY + i             
        KmN = (j-1)*nY + i              
        KmN2 = (j-2)*nY + i             
        Kp1 = j*nY + (i+1)              
        Kp2 = j*nY + (i+2)              
        Km1 = j*nY + (i-1)              
        Km2 = j*nY + (i-2)              
        
        # Check boundary conditions
        # Fill matrix
        # K,K is the diagonal value
        # Robin Boundary Condition
        if (i == 0):
            A[K,K] = -3*dh - 5
            A[K,Kp1] = 2/h
            A[K,Kp2] = -dh
            b[K] = -125.0  # degrees Celsius
            
        elif (i == nX-1):
            A[K,K] = -3*dh + 5
            A[K,Km1] = 2/h
            A[K,Km2] = -dh
            b[K] = 125.0  # degrees Celsius
            
        elif (j == 0):
            A[K,K] = -3*dh - 5
            A[K,KpN] = 2/h
            A[K,KpN2] = -dh
            b[K] = -125.0  # degrees Celsius
            
        elif (j == nY-1):
            A[K,K] = -3*dh + 5
            A[K,KmN] = 2/h
            A[K,KmN2] = -dh
            b[K] = 125.0  # degrees Celsius
            
        # Not Boundary Conditions
        else:
            A[K,K] = -4*h2
            A[K,Kp1] = h2
            A[K,Km1] = h2
            A[K,KpN] = h2
            A[K,KmN] = h2
            b[K] = f



# nested loop for the Corner mapping
# The loop is needed again to use the definition of K
for i in range(nX):
    for j in range(nY):
        # mapping relation 
        K = j*nY + i
        # neighbours p:plus m:minus
        KpN = (j+1)*nY + i              
        KmN = (j-1)*nY + i              
        Kp1 = j*nY + (i+1)              
        Km1 = j*nY + (i-1)              
        
        # Corner 1: 0,0; 2: 1,1; 3: -1,0; 4: -1,-1 
        if (i == 0 and j == 0):
            b[K] = (b[Kp1]+b[KpN]+b[11])/3
            
        elif (i == (nX-1) and j == 0):
            b[K] = (b[Km1]+b[KpN]+b[18])/3
            
        elif (i == 0 and j == (nY-1)):
            b[K] = (b[Kp1]+b[KmN]+b[81])/3
            
        elif (i == (nX-1) and j == (nY-1)):
            b[K] = (b[Km1]+b[KmN]+b[88])/3


# Convert the LIL matrix to a CSR matrix
A = A.tocsr()

# Solve using spsolve
z = spsolve(A, b)
Z = z.reshape(nX, nY)
plt.contourf(X, Y, Z, cmap='inferno')
plt.colorbar()
plt.xlabel('X')
plt.ylabel('Y')

# 3d plot
X, Y = np.meshgrid(X, Y)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
                       cmap=cm.plasma, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)


plt.show()