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)

ms = 4.9

A1,A2 = [],[]
for p in 0.2,0.5,0.8:
    Y1 = 1 - (1-p)**X1
    Y2 = 1 - (1-p)**(X2+1)
    for k in range(len(Y1)-1):
        ax1.plot([X1[k],X1[k]+1], [Y1[k],Y1[k]], '-', color='grey')
        ax2.plot([X2[k],X2[k]+1], [Y2[k],Y2[k]], '-', color='grey')

    ax1.plot(X1[1:], Y1[:-1], 'o', mec='k', mfc='w', ms=ms)
    ax2.plot(X2[1:], Y2[:-1], 'o', mec='k', mfc='w', ms=ms)
    a1 = ax1.plot(X1, Y1, 'o', color=col[p], ms=ms)
    a2 = ax2.plot(X2, Y2, 'o', color=col[p], ms=ms)
    A1.append(a1)
    A2.append(a2)

ax1.set_xlabel("x")
ax1.set_ylabel(u"P(X\N{Less-THAN OR EQUAL TO}x)")
ax2.set_xlabel("x")
ax2.set_ylabel(u"P(X\N{Less-THAN OR EQUAL TO}x)")

bx1 = ax1.legend(A1, ("p=0.2", "p=0.5", "p=0.8"),\
                 numpoints=1, handletextpad=0, loc="lower right")
bx2 = ax2.legend(A2, ("p=0.2", "p=0.5", "p=0.8"),\
                numpoints=1, handletextpad=0, loc="lower 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_cdf.pdf")
plt.savefig("geometric_cdf.svg")