# -*- coding: utf-8 -*-
"""
Sensitivity analysis example
Warning, can take a long time if lots of samples
@author: neil 2022
"""

# This is the 3nd program in the Monte Carlo series
# this is basically the same as the 1st program, in terms of the simple 1D FD code,
# This illustrates a 'sensitivity' analysis.  Lets say you have knowledge of the surface
# and depth BCs, and you know there are 9 layers, of which you know their approximate
# thermal conductivities, +/- 25%.  With no other knowledge, can you put some bounds on the
# expected error of the surface heat flow?

# A big decision needs to be made, does +/-25% mean guassian errors (which implies we know
# this to be true) or should we just assume that it means the conductivity can be random 
# within those bounds, or maybe lognormal?
# (Turns out that when doing lots of calcs with random numbers, there is a tendency for the results
# to head towards guassian, no matter what the input.)

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

import numpy as np
import matplotlib.pyplot as plt

thermKerror  = .25         # fraction error in termal conductivities
BCerror      = .05         # fraction error in basal temperature

# 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    = 101         # 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 = 10000                          # this is the number of Monte Carlo iterations
layers = [8,3,6,12,10,22,7,13,19]         # a list of layer thicknesses in meters
Kl     = [1.,3.,6.,2.,4.,3.,5.,1.,3.]     # list of approx conductivities, each will have errors

fig = plt.figure()           # set up a plot window to place our iterations of the results
ax = fig.add_subplot(1,2,1)  # we will have 2 plots, one of the temperature profiles, and another the histogram of the output heat flow
ax.grid(True)
ax.invert_yaxis() 

# 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 distribution of heat flows
# set up the basic 'A' and 'C' arrays
Kvec    = np.ones(n-1)        # there is one less layer than nodes
a       = np.zeros(len(z))    # a place to form our tridagonals for the 'A' matrix
Q       = np.zeros(nsamples)  # a place to save the heat flows
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

for i in np.arange(0,nsamples):               # number of FD runs
    # this is the hard part, assigning 'random' values to each K layer and making Kvec
    nn=0                                      # nn runs over the nodes, n runs over the layers, m steps over the nodes in each layer
    for n in range(len(layers)):              # step over the layers
        for m in range(layers[n]):            # step over the nodes in each layer
            err = (np.random.rand()-.5)*2     # for each layer get a random number between -1 +1
            Kvec[nn] = Kl[n] + Kl[n]*err*thermCerror   # add or subtract the error from the base value 'Kl'
            nn +=1
               
    # Now we make our A and C arrays and solve for Temperature profile produced by the current K values
    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 for correct K[?]
    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          # 
    C[-1]    = dnBC + ((np.random.rand()-.5)*2)*dnBC * BCerror # finally, add error to BC temperature
        
    T = np.linalg.solve(A,C)   # direct solver in numpy
    ax.plot(T,z)               # plot the temperature profile for this one realization
    # we calc the output heat flux for this one run, and save for the output histogram of heat fluxes
    Q[i] = Kvec[0] * (T[1]-T[0])/delz

# when we exit the loop, we have an array of heat flux (Q) values for our 'nsamples' iterations
ax.set_xlabel("Temperature")
ax.set_ylabel("depth")
ax.set_title("Temperature profile realizations")

ax1 = fig.add_subplot(1,2,2)             # the 2nd plot will be the histogram (normalized to be a PDF)
ax1.grid(True)
ax1.hist(Q,bins=100,density=True)        # this calls np.histogram to calculate the bins
ax1.set_ylabel("PDF value")
ax1.set_xlabel("Heat flux")
ax1.set_title("PDF, heat flux")