import numpy as np
import matplotlib.pyplot as plt
import scipy.special as sp

col = {1 : 'orange', 4: 'purple', 10: 'lightblue'}

X = np.arange(0, 21)

plt.rc('text', usetex = True)
plt.rc('font', family = 'serif', size = 12)

##                                                                                
## PMF                                                                            
## 

plt.clf()
plt.figure(figsize=(4,3.2))

A = []
for L in 1,4,10:

    P = -L + X*np.log(L) - sp.gammaln(X+1)
    P = np.exp(P)

    plt.plot(X, P, '-', color='grey', label='_nolegend_')
    a = plt.plot(X, P, 'o', color=col[L], markeredgecolor='k', markeredgewidth=0.5)
    A.append(a)

plt.xlabel("$k$")
plt.ylabel(r"$P(x=k)$")
plt.ylim([0,0.4])
plt.yticks(np.arange(0,0.42,.05))
plt.tick_params(direction='in', top=True, right=True)

plt.legend((r"$\lambda=1$", r"$\lambda=4$", r"$\lambda=10$"),
                numpoints=1, handlelength=0.75, handletextpad=0.5,\
                loc="upper right", frameon=False)
plt.xlim(-1,21)

plt.savefig("poisson_pmf.pdf")
plt.savefig("poisson_pmf.svg")