#!/usr/bin/env python3

#
# Bayesian coinflip, 
# numerically discretized distribution functions
# by Daniel Nevermann
#

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# ***
# *** global variables
# ***
trueBias = 0.3                             # true coin bias
nX       = 101                             # numerical discretization
xx = [i*1.0/(nX-1.0) for i in range(nX)]   
pp = [1.0/nX  for _ in range(nX)]          # starting prior
damp_const = 0.

# ***
# *** Bayesian updating
# ***
def updatePrior():
    """Bayesian inference for coin flips;
       marginal corresponds to normalization of posterior"""
    evidence = 0.0 if (np.random.uniform() > trueBias) else 1.0  # coin flip
    pp_old = pp.copy() # full prior
    for iX in range(nX):
        prior = pp[iX]
        likelihood = xx[iX] if (evidence == 1) else 1.0 - xx[iX]
        pp[iX] = likelihood * prior  # Bayes theorem
    norm = sum(pp)
    for iX in range(nX):  # damping
        pp[iX] = pp[iX] / norm
        pp[iX] = pp_old[iX] * damp_const + pp[iX] * (1-damp_const)
    return evidence

# ***
# *** Animation
# ***
fig, ax = plt.subplots()
bar_container = ax.bar(xx, pp, width=0.05)
def update(frame):
    evidence = updatePrior()
    for iX, bar in enumerate(bar_container):
        bar.set_height(pp[iX])
    ax.set_title(f'Observation: {frame + 1}, Evidence: {evidence:.1f}')
    return bar_container

# ***
# *** main
# ***

# Set up animation parameters
ani = FuncAnimation(fig, update, frames=2000, repeat=False, blit=False, interval=10)  

# Display the animated plot
plt.xlabel("Coin Bias Hypothesis")
plt.ylabel("Probability")
plt.title("Bayesian Inference: Updating Prior")
plt.ylim(0, 0.3)
plt.show()