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

col = {0.2: 'orange', 0.5: 'purple', 0.8: 'lightblue'}

plt.clf()
plt.figure(figsize=(7.5,3.2))
ax1 = plt.axes([0.08,0.13,0.4,0.8])
ax2 = plt.axes([0.57,0.13,0.4,0.8])
plt.hold(True)

X1 = np.arange(1,11)
X2 = np.arange(0,11)

A1,A2 = [],[]
for p in 0.2,0.5,0.8:
    Y1 = (X1-1)*np.log(1-p) + np.log(p)
    Y1 = np.exp(Y1)
    Y2 = X2*np.log(1-p) + np.log(p)
    Y2 = np.exp(Y2)
    ax1.plot(X1, Y1, '-', color='grey', lw=1)
    ax2.plot(X2, Y2, '-', color='grey', lw=1)
    a1 = ax1.plot(X1, Y1, 'o', color=col[p], lw=3)
    a2 = ax2.plot(X2, Y2, 'o', color=col[p], lw=3)
    A1.append(a1)
    A2.append(a2)

ax1.set_xlabel("x")
ax1.set_ylabel("P(X=x)")
ax2.set_xlabel("x")
ax2.set_ylabel("P(X=x)")

bx1 = ax1.legend(A1, ("p=0.2", "p=0.5", "p=0.8"),\
                 numpoints=1, handletextpad=0, loc="upper right")
bx2 = ax2.legend(A2, ("p=0.2", "p=0.5", "p=0.8"),\
                numpoints=1, handletextpad=0, loc="upper right")
bx1.draw_frame(False)
bx2.draw_frame(False)
ax1.set_xlim(-0.2,10.2)
ax2.set_xlim(-0.2,10.2)
ax1.set_ylim(-0.05,1.05)
ax2.set_ylim(-0.05,1.05)

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