import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation, rc
from matplotlib.animation import FuncAnimation
from IPython.display import HTML, Image # For GIF

rc('animation', html='html5')
np.random.seed(5)

# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)

prob_vals = np.arange(start=0.0, stop=1.01, step=0.05)

plt.style.use('ggplot')
fig = plt.figure(figsize=(10, 5))

def update(i):
    # initialize the graph of the previous frame
    plt.cla()

    p = prob_vals[i]

    # draw bernoulli distribution
    plt.bar([0.0, 1.0], [1.0 - p, p], alpha=0.5) # bar graph
    plt.xlabel('x')
    plt.ylabel('probability')
    plt.suptitle('Bernoulli Distribution', fontsize=20)
    plt.title('$p=' + str(np.round(p, 2)) + '$', loc='left')
    plt.xticks(ticks=[0, 1]) # x axis ticks
    plt.grid()
    plt.ylim(-0.1, 1.1)

anime_prob = FuncAnimation(fig, update, frames=len(prob_vals), interval=1000)

anime_prob.save('bernoulli_dist.gif', writer='pillow', fps=1)
Image(url='bernoulli_dist.gif')