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

X = np.arange(18, 81, 2)

P = 1/(1+np.exp(1-0.2*np.sqrt(X-18)))

Y,S,n = [],[],500
for i,x in enumerate(X):
    y = 1*(np.random.uniform(size=n)<P[i])
    p = y.mean()
    Y.append(p)
    S.append(np.sqrt(p*(1-p)/n))
Y = np.array(Y)
S = np.array(S)

## Multiplier for 95% simultaneous confidence band using Bonferroni method.
f = -sp.ndtri(0.025/len(P))

plt.clf()
a = plt.plot(X, Y, '-', color='black')
plt.hold(True)
b = plt.plot(X, Y+2*S, '-', color='royalblue')
plt.plot(X, Y-2*S, '-', color='royalblue')
c = plt.plot(X, Y+f*S, '-', color='orangered')
plt.plot(X, Y-f*S, '-', color='orangered')
d = plt.plot(X, P, '-', color='green')
plt.ylim(0,1)
plt.xlim(18,80)
B = plt.legend((a,b,c,d), ('Estimate', '95% point-wise CB',\
                           '95% simultaneous CB', "True proportions"),\
                           'upper left')
B.draw_frame(False)
plt.xlabel("Age")
plt.ylabel("Proportion supporting candidate A")
plt.savefig("binomial_confidence_band.png")
plt.savefig("binomial_confidence_band.svg")