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

# This is the 2nd program in the Monte Carlo series
# this is basically the same as the 1st program, in terms of the simple 1D FD code,
# However, this treats the problem from a stocastic viewpoint, recognizing that the conductivities
# of the 2 layers are only known as probability distributions, which are NOT guassian.

# Solves a shallow heat flow problem, with 2 layers, where we know the temperature at depth
# See 1st program for basic description
# In this version instead of using constant values of K for each layer, we assume we know
# that the values are actually given better by a distribution of possible K values.  

# Here we assume that the sand layer is Lognormally distrubuted about the value of .5, while
# the clay has a triangular distribution about 2, but biased towards low values (dry patches)
# the clay and sand distributions are modelled by sampling from lognormal and triangular PDFs
# using multiple samples, the final output is the resulting surface heat flux PMF!

# note the finite difference routine is identical to the deterministic code (monte1.py)

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
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
nsamples = 50000   # this is the number of Monte Carlo iterations

# pre-make the arrays
Kvec    = np.ones(n-1)   # there is one less layer than nodes
a       = np.zeros(len(z))
Q       = np.zeros(nsamples)
C=np.zeros(n)           # make the BC vector for the right hand side
C[0]  = upBC            # upper BC is 0 degrees
C[-1] = dnBC            # lower BC is 3 degrees

K1 = np.random.lognormal(-0.3,.5,size=nsamples)  # sample the lognormal distribution, this is the sand
# 1st parameter is the mean of the underlying normal distribution, while the 2nd is the sd
K2 = np.random.triangular(.5,2,2.5,size=nsamples)  # sample a triangular distribution, this is the clay

# below here there are no hidden numbers *********************************************
# We run our Monte Carlo simulation, sampling from the input PDFS and running the 
# samples thru our deterministic model of heat flow.  If we do this many times we
# should get an estimate of the output PDF
for i in np.arange(0,nsamples):
    Kvec[:]     = K1[i]        # 
    Kvec[n//2:] = K2[i]
    
    # Now we just make our A and C arrays and solve for Temperature
    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
    A[-1,-2] = 0         # 
    A[-1,-1] = 1          # 
        
    T = np.linalg.solve(A,C)   # direct solver in numpy
    # finally we calc the output heat flux for this one run
    Q[i] = K1[i] * (T[1]-T[0])/delz

# when we exit the loop, we have an array of heat flux (Q) values

depth = z

# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()              # set up a plot window
ax1 = fig.add_subplot(2,1,2)
ax1.grid(True)

ax1.hist(Q,bins=200,density=True)        # this calls np.histogram to calculate the bins
# we could 'normalize' the histogram to get a PMF or PDF, or use numpy to do it for us:
# set 'density' option to 'True', to draw and return an (estimate) of probability density (PDF): 
# each bin will display the bin's raw count divided by the total number of counts and the bin width 
# (density = counts / (sum(counts) * np.diff(bins))), so that the area under the histogram integrates to 1 
#ax1.set_xlim([0,6])                     # the lognormal has a LONG positive tail
ax1.set_ylabel("PDF value")
ax1.set_xlabel("Values of Heat flux")
ax1.text(.05,25,"Output heat flux PDF",fontsize='16')

ax2=fig.add_subplot(2,1,1)
ax2.hist(K2,bins=200,density=True)
ax2.hist(K1,bins=200,density=True,facecolor='r')
ax2.set_ylabel("PDF value")
ax2.text(2.5,1,"Values of K1 (red) and K2 (blue)")
ax2.set_title("Input conductivity PDFs")
ax2.grid(True)
ax2.hist(K2,bins=200,density=True,histtype='step',edgecolor='b')
ax2.set_xlim([0,5])


plt.show()