# -*- coding: utf-8 -*-
"""
Created on Sun Apr 19 20:59:59 2020

@author: neil
"""
# A first program in a Bayesian approach to geo-stats
# illustrates a simple but powerful data structure, a DICTIONARY (a special type of 2D list of key:value pairs)
# also plotting a histogram, and normalizing to create a PMF
# finally a standard PYTHONIC way of splitting and creating 1D lists from 2+D lists (worth knowing, but advanced)

import numpy as np
import matplotlib.pyplot as plt

# Create, and plot a PDF (actually a PMF) of the results of 3 flips of a coin
# HHH, HHT, HTH, HTT, THH, THT, TTH, TTT  (these are all equally likely)
# plot them as a histogram of 0,1,2,3 heads

# make a dictionary of the distrbution.  A python dictionary is a list of key value pairs using '{}'
# brackets, and the keys and values are separated by a ':'.  You can 'index' using the keys.
# create a dictionary with names like 'H2' meaning 'two heads'
PMF = {'H0':1,'H1':3,'H2':3,'H3':1}  # the keys are the names of the histogram bars, values are the freq of heads

# plot the histogram
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(*PMF.items())  # dictionary.items() returns a list of key,value pairs
# .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    
ax1    = figX.add_subplot(1,2,1)        # add 1 of 2 plotting windows 
ax1.bar(x,y)     # x is a list of the key names, y is a list of number of heads for each key
ax1.set_title("histogram of heads in 3 coin flips")
ax1.set_ylabel("Number of heads in 3 flips")
ax1.set_xlabel("Outcome, number of heads in 3 flips")
# to turn into a proper (normalized) PMF, the sum of the events must =1
# then we can consider them to be probabilities

# normailize the histogram to a PMF probability distribution
x,y = zip(*PMF.items())        # x is a list of keys, y is a list of values, '*' is a special unzip command
factor = 1/np.sum(y)           # sum the values, to make the values add to 1, multi each by 1/sum

for k in x:
    PMF[k] = PMF[k]*factor

x,y = zip(*PMF.items())          # the distribution is normalized                  # 
axy1    = figX.add_subplot(1,2,2)      # add two plotting windows 
axy1.bar(x,y)
axy1.set_title("normalized PMF of outcomes of 3 coin flips")
axy1.set_ylabel("probability of outcome")
axy1.set_xlabel("Outcome, number of heads in 3 flips")