from pylab import *             # library for making plots
from decimal import *           # library for exact arithmetic with large numbers

# make a plot of the probability mass function y(x)
def pmfplot(x, y):
    bar(x, y, width=0.5, align='center')
    xticks(x, ticklabels(x, 10))
    ylim(0, 1.1 * max(y))
    show()
	
# calculate the p.m.f. of the binomial random variable with parameters n and p
def binomialpmf(n, p):
    q = 1.0 - p
    B = [0]*(n + 1)
    B[0] = q**n
    for k in range(1, n + 1):
        B[k] = B[k - 1] * p * (n - k + 1) / (k * q)
    return B

# make a plot of the binomial p.m.f.
def binomial_plot(n, p):
    pmfplot(range(n + 1), binomialpmf(n, p))

# calculate the p.m.f. of the poisson random variable with mean mu for values 0 up to L
def poissonpmf(mu, L):
    P = [0]*(L + 1)
    P[0] = exp(-mu)
    for i in range(1, L + 1):
        P[i] = P[i-1] * mu / i
    return P

# make a plot of the poisson p.m.f.
def poisson_plot(mu, L):
    pmfplot(range(L + 1), poissonpmf(mu, L))


# make tick labels so there are at most l of them 
def ticklabels(seq, l):
    k = len(seq) / l
    if k == 0:
        return seq
    else:
        sp = []
        for i in range(len(seq)):
            if i % k == 0:
                sp.append(seq[i])
            else:
                sp.append('')
        return sp