import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from plotstyle import *

fig = plt.figure(figsize=cm_size(figure_width, 1.8*figure_height))
spec = gridspec.GridSpec(nrows=2, ncols=2, hspace=0.6,
                         **adjust_fs(fig, left=5.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(spec[0, :])
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('Prob. 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=10)
    if mu > rmu :
        ax.text(mu-0.1, 0.04, '?', zorder=10, ha='right')
    else :
        ax.text(mu+0.1, 0.04, '?', zorder=10)
for k, ls in enumerate([lsCm, lsBm, lsDm]) :
    ax.plot(x, g[:,k], zorder=5, **ls)
ax.plot(xd, 0.05*rng.rand(len(xd))+0.2, zorder=10, **psAm)

# 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(spec[1, 0])
ax.set_xlabel(r'Parameter $\mu$')
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, **lsAm)

ax = fig.add_subplot(spec[1, 1])
ax.set_xlabel(r'Parameter $\mu$')
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), **lsAm)

plt.savefig('mlemean.pdf')