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

# using only procedural programing to do simple Bayes one step update on drilling example, with a single sample
# trying to find the bottom of an ore body while drilling
# (this template only works for a set of hypothyses 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
# of Hi-grade (lets say copper) or Low-grade (not much good ore), 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 (hypothesis) might best be 
# defined as 50-50, a probability of 0.5 ore, 0.5 host (or you can try other priors),
# in fact, since there are only 2 possibilities (ore or host) our prior is actually a 2 valued
# probability (Mass) 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).

# we use Bayesian updating in the form of:
# [posterior belief (ore or host)] = [likelyhoods of getting the observed data] times prior [belief]
# where each of these are PMFs (or un-normalized PMFs)
# the likelyhoods are the probability of seeing the observed data, based on
# knowledge of the problem, external to the priors, in this case
# this knowledge is the relative composition of the rock in the host and ore bodies)
# technically: posterior belief PMF = likelyhood of observed sample/data PMF * prior belief PMF

import numpy as np
import matplotlib.pyplot as plt

#### data section ##### (you can change this)
data  = 'LowGrade'   # data is what is in the core sample, either 'LowGrade' or 'HiGrade' material 
####

# Define the character of our 'sample space', using PMFs to describe the host and ore rocks
# We need to probability of the data (sample,) given the 2 different possibilities that we are sampling:
# host or ore.  These posibilities are usual refered to as the 
# hypotheses.  We call the first hypotheses 'host', the 2nd 'ore', and the data is our sample.
# We use dictionaries (lookup tables) of probabilities of seeing low-grade or hi-grade in core sample
host  = {"LowGrade":.75,"HiGrade":.25}  # host rock with some ore, note these PMFs are normalized
ore   = {"LowGrade":.5,"HiGrade":.5}    # profitable ore, mixed with some host rock
# a dictionary is an unordered list of pairs, {'key name':value},and you can lookup a value using the 'keys'
# these are very simple PMFs (dictionaries are convienient since you can store names, but we could have used a table)
# We can access a value in a dictionary using this syntax: host['HiGrade'] yields the value 0.25
# and we need to start the process by defining a prior estimate of the result (may be based on other knowledge)
prior = {'inHost':.5,'inOre':.5}    # a somewhat uninformative prior PMF, you can try varying this(?)   

# our data is that our sample is Higrade or lowgrade, therefor the probability of getting that data with our 2 hypotheses
# is the 1st or 2nd column of our host and ore PMFs
# this probability of the data given our 2 hypostheses is usually called the 'likelyhood of the observing the data'
# Note the likelyhood is NOT a normalized PMF, but the probability ratios are correct.  We normalize later.

# this creates a dictionay called likelyhood, based on the sample. The dictionary is really a PMF with the 
# x axis the 'keys' or hypotheses and the vertical axis the 'values' or probabilities
if data == 'LowGrade':
    likelyhood = {'inHost':host['LowGrade'],'inOre':ore['LowGrade']}  # this is the likelyhood that we observe lowgrade if in host or in ore
    title = 'LowGrade'                           # used for plotting
else:
    likelyhood = {'inHost':host['HiGrade'],'inOre':ore['HiGrade']}  # this is the likelyhood that we observe higrade if in host or in ore
    title = 'HiGrade'

# the product of the prior and the likelyhood (piecewise) gives the un-normalized posterior.  We can normalize with the 
# probability of the data, but since the hypotheses are exclusive and complete, we can just normalize to ensure
# a total probability of 1
# The posterior PMF is the piecewise multiplication of the likelyhood and prior PMFs
posterior =  {'inHost':prior['inHost']*likelyhood['inHost'],'inOre':prior['inOre']*likelyhood['inOre']}
# note the above syntax:  posterior dictionary has a key entry 'inHost', with value prior['inHost']*likelyhood['inHost']

# normailize to a PMF probability distribution of 2 items (could be written as a function)
x,y = zip(*posterior.items())    # y is a list of values
factor = 1/np.sum(y)             # sum the values, to make the values add to 1, multi each by 1/sum
posterior['inHost'] = posterior['inHost']*factor
posterior['inOre'] = posterior['inOre']*factor
# we now have the normalized posterior or our result and can decide what to do with it?
    
################# plotting routines #####################################################################    
figX  = plt.figure()                    # make a figure of all the input PMFs
# this is a cute way of extracting the keys and the values of a dictionary (not obvious!) 
x,y = zip(*host.items()) 
# .items returns the list of dictionary pairs, the 'unzip' operator '*' makes an iterator
# over the tuples (pairs), The 'zip' then packs those into 2 tuples placed in x and y respectively   
ax1    = figX.add_subplot(2,2,1)        # add 1 of 4 plotting windows 
ax1.bar(x,y)
ax1.set_title("host PMF for ratio of hi & low grade")
ax1.set_ylim([0,1.])

x,y = zip(*ore.items())
ax2    = figX.add_subplot(2,2,2)        # add 2 of 4 plotting windows 
ax2.bar(x,y)
ax2.set_title("ore PMF for ratio of hi & low grade")
ax2.set_ylim([0,1.])

x,y = zip(*prior.items())
ax3    = figX.add_subplot(2,2,3)        # add 3 of 4 plotting windows 
ax3.bar(x,y)
ax3.set_title("prior PMF")
ax3.set_ylim([0,1.])

x,y = zip(*likelyhood.items())
ax4    = figX.add_subplot(2,2,4)        # add 4 of 4 plotting windows 
ax4.bar(x,y)
ax4.set_title("likelyhood PMF for sample of "+title)
ax4.set_ylim([0,1.])
figX.suptitle('Prior information for a sample of '+title, fontsize=16)


x,y = zip(*posterior.items())          # the distribution is normalized 
figY  = plt.figure()                   # 
axy1    = figY.add_subplot(1,2,2)      # add two plotting windows 
axy1.bar(x,y)
axy1.set_title("posterior PMF")
axy1.set_ylim([0,1.])

# so we can see the difference, also plot the prior again
x,y = zip(*prior.items())
axy3    = figY.add_subplot(1,2,1)        # add two plotting windows 
axy3.bar(x,y)
axy3.set_title("prior PMF")
axy3.set_ylim([0,1.])
figY.suptitle('Bayes Posterior PMF due to a sample of '+title, fontsize=16)