import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import scipy.stats as stats

plt.style.use('seaborn-pastel')


n = 100
p_iter = 0.05

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(0, 1))

x_axis = np.arange(0, 1+1/n, 1/n)
line, = ax.plot([], [], lw=2)


def init():
    line.set_data(x_axis, [])
    return line,

def animate(i):
    cur_p = p_iter * i
    y_axis = np.array([stats.binom.pmf(k, n, cur_p) for k in range(0, n+1)])
    line.set_data(x_axis, y_axis)
    plt.title('p = {0}'.format(np.around(cur_p, 2)))
    return line,




cframes = int(1 /p_iter) + 1
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=cframes, interval=250, blit=True)

anim.save('output/Bi(n, p).gif', writer='imagemagick')