import numpy as np
import matplotlib.pyplot as plt

plt.xkcd()
fig = plt.figure( figsize=(6,5) )

# the data:
n = 40
rng = np.random.RandomState(54637281)
sigma = 0.5
rmu = 2.0
xd = rng.randn(n)*sigma+rmu
# and possible pdfs:
x = np.arange( 0.0, 4.0, 0.01 )
mus = [1.5, 2.0, 2.5]
g=np.zeros((len(x), len(mus)))
for k, mu in enumerate(mus) :
    g[:,k] = np.exp(-0.5*((x-mu)/sigma)**2.0)/np.sqrt(2.0*np.pi)/sigma
# plot it:
ax = fig.add_subplot( 2, 1, 1 )
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.set_xlim(0.5, 3.5)
ax.set_ylim(-0.02, 0.85)
ax.set_xticks( np.arange(0, 5))
ax.set_yticks( np.arange(0, 0.9, 0.2))
ax.set_xlabel('x')
ax.set_ylabel('Probability density')
s = 1
for mu in mus :
    r = 5.0*rng.rand()+2.0
    cs = 'angle3,angleA={:.0f},angleB={:.0f}'.format(90+s*r, 90-s*r)
    s *= -1
    ax.annotate('', xy=(mu, 0.02), xycoords='data',
                xytext=(mu, 0.75), textcoords='data',
                arrowprops=dict(arrowstyle="->", relpos=(0.5,0.5),
                connectionstyle=cs), zorder=1 )
    if mu > rmu :
        ax.text(mu-0.1, 0.04, '?', zorder=1, ha='right')
    else :
        ax.text(mu+0.1, 0.04, '?', zorder=1)
for k in xrange(len(mus)) :
    ax.plot(x, g[:,k], zorder=5)
ax.scatter(xd, 0.05*rng.rand(len(xd))+0.2, s=30, zorder=10)

# likelihood:
thetas=np.arange(1.5, 2.6, 0.01)
ps=np.zeros((len(xd),len(thetas)));
for i, theta in enumerate(thetas) :
    ps[:,i]=np.exp(-0.5*((xd-theta)/sigma)**2.0)/np.sqrt(2.0*np.pi)/sigma
p=np.prod(ps,axis=0)
# plot it:
ax = fig.add_subplot( 2, 2, 3 )
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.set_xlabel(r'Parameter $\theta$')
ax.set_ylabel('Likelihood')
ax.set_xticks( np.arange(1.6, 2.5, 0.4))
ax.annotate('Maximum',
            xy=(2.0, 5.5e-11), xycoords='data',
            xytext=(1.0, 1.1), textcoords='axes fraction', ha='right',
            arrowprops=dict(arrowstyle="->", relpos=(0.0,0.5),
            connectionstyle="angle3,angleA=10,angleB=70") )
ax.annotate('',
            xy=(2.0, 0), xycoords='data',
            xytext=(2.0, 5e-11), textcoords='data',
            arrowprops=dict(arrowstyle="->", relpos=(0.5,0.5),
            connectionstyle="angle3,angleA=90,angleB=80") )
ax.plot(thetas,p)

ax = fig.add_subplot( 2, 2, 4 )
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.yaxis.set_ticks_position('left')
ax.xaxis.set_ticks_position('bottom')
ax.set_xlabel(r'Parameter $\theta$')
ax.set_ylabel('Log-Likelihood')
ax.set_ylim(-50,-20)
ax.set_xticks( np.arange(1.6, 2.5, 0.4))
ax.set_yticks( np.arange(-50, -19, 10.0))
ax.annotate('Maximum',
            xy=(2.0, -23), xycoords='data',
            xytext=(1.0, 1.1), textcoords='axes fraction', ha='right',
            arrowprops=dict(arrowstyle="->", relpos=(0.0,0.5),
            connectionstyle="angle3,angleA=10,angleB=70") )
ax.annotate('',
            xy=(2.0, -50), xycoords='data',
            xytext=(2.0, -26), textcoords='data',
            arrowprops=dict(arrowstyle="->", relpos=(0.5,0.5),
            connectionstyle="angle3,angleA=80,angleB=100") )
ax.plot(thetas,np.log(p))

plt.tight_layout();
plt.savefig('mlemean.pdf')
#plt.show();