# -*- coding: utf-8 -*-
"""
2022
@author: neil
"""

# This is the 1st program in the Monte Carlo series, this is the basic problem.
# We use our very simple FD heat flow code in 1D, with constant everything.

# Solves a shallow heat flow problem, with 2 layers, where we know the temperature at depth
# We are modelling a 100m depth of ground, with a 50m layer of high conductivity wet clay
# underneath a 50m layer of dry sand sand.  We use set values of conductivity
# for the each layer, the wet soil has a conductivity of 2, while the dry has a conductivity
# of 0.5. Note dry and wet soils have highly variable conductivities, which will will investigate.
# THe surface BC is T=0 and the basal temperature is T=3

# This shows a simple non-statistical approach (deterministic, single values)
# the result is the temperature profile and the final output is the resulting surface heat flux

import numpy as np
import matplotlib.pyplot as plt

# Using a Matrix implicit finite difference approach
upBC = 0           # stc at 0C
dnBC = 3           # base at 3C
K1 = .6           # conductivity of clay (these are not reasonable numbers!)
K2 = 2.0            # conductivity for sand
Z = 100            # thickness of problem, 100 meters
n = 5              # number of nodes in Z
delz = Z/(n-1)     # delta z, length of problem / number of nodes
z = np.linspace(0,Z,n) # 'z' is a vector for plotting, and a template for making other arrays

# And we need the transition depths
zK12 = 49       # transition depths in meters

# *****below here there are no hidden numbers *********************************************
Kvec = np.zeros(n-1)   # there is one less layer than nodes
# this is easy to read but not very efficient way of constructing the vector of K values with depth
# there are fancy 'pythonic' ways of doing this, but for us:
j = 1
for j in np.arange(len(Kvec)):        # look at the depth of each node
    Kvec[j] = K2
    if z[j] < zK12:                   # if less than transition depth, make K = to K1
        Kvec[j] = K1                  # and assign a K to the layer below the node
# Kvec is one less than nodes, and each K applies to the layer below the node

a       = np.zeros(len(z))
a[1:-1] = -Kvec[:-1] -Kvec[1:]   # making the A matrix, 'a' is -K[i-1]-K[i]
b = c   = Kvec                   # make b and c the correct length
# because the b and c vectors are placed as below, this works
A = np.diag(a,0) + np.diag(b,1) + np.diag(c,-1)  # b and c end up as 'b' K[i], 'c' K[i-1]

# add boundary conditions to the A matrix
A[0,0] = 1            # adjust the A matrix for the top BC 
A[0,1] = 0
# NOTE the negative index below counts back from end!
A[-1,-2] = 0         # 
A[-1,-1] = 1          # 
# view the A and C arrays to see the difference between the position of the BCs

C=np.zeros(n);       # make the BC vector for the right hand side
C[0]  = upBC            # upper BC is 0 degrees
C[-1] = dnBC            # temperature BC on the bottom node

T = np.linalg.solve(A,C)   # direct solver in numpy

depth = z
# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()
ax1 = fig.add_subplot(1,1,1)
ax1.grid(True)
ax1.plot(T,depth, 'r-o')  
ax1.plot([0,T[-1]],[0, 0], 'g')
ax1.plot([0,T[-1]],[zK12+1, zK12+1], 'b:')
ax1.set_title("Geothermal Temperature Profile, no heat sources")
ax1.set_xlabel("Temperature")  
ax1.set_ylabel("Depth [meters]")  
ax1.invert_yaxis()
heatflux = K1 * (T[1]-T[0])/delz 
s = "Heat Flux at surface {:.3f}".format(heatflux) 
ax1.text(.5,10,s, color='b')
s = "Wet Clay" 
ax1.text(1,80,s,fontsize=18)
s = "Dry Sand" 
ax1.text(.5,30,s,fontsize=18)
plt.show()