# -*- coding: utf-8 -*-
"""
@author: neil 2022
"""
# 2nd program for drilling example of Baysian updating, trying to find the bottom of an ore body while drilling
# this version allows for multiple sequential core samples, and sequentially updating the output PMF
# using procedural and functional programing to do simple Bayes update on our drilling example,
# (this template only works for a set of hypothyses (discrete outcomes) that are exclusive and exhaustive, [allows default normalizing])

# Now we can set up our problem.  Assume we are drilling and we get either a length of core
# that is hi-grade or matrix, we use bayes to update our probablity of being in host or ore rock

# So the question is: if we get a sample, how does that change our view of whether we are in ore or host?
# First we need to find our current belief in the probablities.  Lets say we are ambivalent about
# if we have reached the bottom of the ore body.  Then our PRIOR might best be defined as 50-50 (or you can try other priors)
# note this is a probability distribution of our belief that we are in either host rock or ore rock, and
# these are mutuably exclusive and also complete (nothing else is possible), so the 2 possibilities cover
# everything and don't overlap.  Thus the sum of these probablities equals 1 or the sure event (useful for normalizing).

# this uses 'dictionaries' for storing our PMFs, dictionaries are basically a list of pairs of keys and values.

import numpy as np
import matplotlib.pyplot as plt

# define several functions that make the manipulation of PMFs easier
figX  = plt.figure()    # initialize a figure to be a place to plot all the posterior PMFs, as we go thru the core samples    
figX.suptitle('Prior and posteriors for several samples', fontsize=16)
# since all our plots are going to be 'bar' plots, write a short function to plot any PMF  
def barplot(axes,pmf,title):   
    """plot a bar plot of an pmf on axes that are supplied, input is an axes, a PMF and a title for the plot"""    
    x,y = zip(*pmf.items()) # this is a cute way of extracting the keys and the values of a dictionary (not obvious!) 
    axes.bar(x,y)
    axes.set_title(title)
    axes.set_ylim([0,1.])

# Two more function definitions to make the code simpler and smaller
# define a normalizing routine for a PMF
def normalize(pmf):
    """ normailize a PMF probability distribution stored as a dictionary
     the input pmf is just a dictionary of names and relative probabilities """
    x,y = zip(*pmf.items())    # y is a list of values (x a list of keys)
    # .items returns the LIST of dictionary key-value pairs, the 'unzip' operator '*' makes an ITERATOR
    factor = 1/np.sum(y)       # sum the values, to make the values add to 1, multi each by 1/sum
    for k in x:                # cycle over the keys
        pmf[k] *= factor       # the keys act as indexes into the dictionary
    return pmf

#define a multiplication routine for PMFs stored as dictionaries (element by element multiplication)
def pmfMulti(p1,p2):
    """ given 2 dictionaries with the same keys, multi the dictionary values and return a new dictionary """
    pmf = {}
    for k in p1.keys():        # the keys must be the same in each dictionary, 
        pmf.update( {k: p1[k]*p2[k]} )
    return pmf

#### data section, this is the information we need to determine the 'likelyhood' of finding a particular sample #####
# dictionaries (lookup tables) of probabilities of seeing  a low-grade (LG) or hi-grade (HG) in core sample of ore or host rock
# these are very simple PMFs (dictionaries are convienient since you can store names, but we could have used lists)
host  = {"LG":.75,"HG":.25}  # host rock with some ore, note these PMFs are normalized
ore   = {"LG":.5,"HG":.5}    # profitable ore, hi-grade mixed with some matrix rock
# the above allows us to make the "likelyhood histogram" or PMF for a given core sample

# now define our prior belief, are we in host or Ore, based on our previous experience, OR on previous samples
prior = {'H':.5,'O':.5}      # a somewhat uninformative prior PMF of the probability of being in Ore or Not, 

# now start drilling, coring and taking samples, here we assume we are going to take 8 core samples
# and see how each sample changes our belief in whether we are in host or ore.
ax1    = figX.add_subplot(3,3,1)         # Make 9 plotting windows (axes),
barplot(ax1,prior,"Initial Prior PMF")   # plot our starting belief (PMF)

# here is our list of core samples that we get while drilling 
sampleList  = ['LG','HG','HG','LG','LG','HG','LG','HG']            # data is what is in the core samples, either matrix 'M' or ore 'O'    

# now we need probability of the data (sample) given the 2 different 'hypotheses', that we are sampling host or ore
# We call the first hypotheses 'Host', the 2nd 'Ore', and the data is our sample of either matrix or hi-grade.

# The 'likelyhood' of the observing the data, is constructed as the probability of getting that data with our 2 hypotheses
# If, for example the sample is 'LG', then we look to see what the probability of 'LG' is in host or ore.  These 2 probabilities
# are the likelyhood distribution.. this is NOT a normalized PMF, but the probability ratios are correct.  We normalize later.

# We now take the 8 Baysian steps, updating our posterior, and using that as the prior at each step
plotn = 2  # used to cycle thru the samples and keep track of the sequential plots, plot each new posterior in sequence in a new axis
for data in sampleList:   # this is the loop thru the list of samples, updating our prior and posterior with the data
    if data == 'LG':
        likelyhood = {'H':host['LG'],'O':ore['LG']}  # this is the likelyhood that we observe low-grade if in host or in ore
    else:
        likelyhood = {'H':host['HG'],'O':ore['HG']}  # this is the likelyhood that we observe hi-grade if in host or in ore
    
    # the product of the prior and likelyhood (piecewise) gives the un-normalized posterior. Since the hypotheses are 
    # exclusive and complete, we can just normalize to ensure a total probability of 1
    posterior = pmfMulti(prior,likelyhood)  # multi the current prior by the likelyhood of the sampled data
    posterior = normalize(posterior)        # normalize the product

    prior = posterior.copy()                # update our prior, and do the whole thing again with a new sample
    ax1    = figX.add_subplot(3,3,plotn)    # add a plot of our new posterior, which becomes our next prior
    barplot(ax1,posterior,"Posterior for sample #" + str(plotn-1) + ' ' + sampleList[plotn-2]) 
    plotn += 1