import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

plt.xkcd()
fig = plt.figure( figsize=(6,6.8) )
rng = np.random.RandomState(4637281)
lmarg=0.1
rmarg=0.1

ax = fig.add_axes([lmarg, 0.75, 1.0-rmarg, 0.25])
ax.spines['bottom'].set_position('zero')
ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.xaxis.set_ticks_position('bottom')
ax.get_yaxis().set_visible(False)
ax.set_xlim(0.0, np.pi)
ax.set_xticks(np.arange(0.125*np.pi, 1.*np.pi, 0.125*np.pi))
ax.set_xticklabels([])
ax.set_ylim(0.0, 3.5)
ax.text(-0.2, 0.5*3.5, 'Activity', rotation='vertical', va='center')
ax.annotate('Tuning curve',
            xy=(0.42*np.pi, 2.5), xycoords='data',
            xytext=(0.3*np.pi, 3.2), textcoords='data', ha='right',
            arrowprops=dict(arrowstyle="->", relpos=(1.0,0.5),
            connectionstyle="angle3,angleA=-10,angleB=110") )
ax.annotate('',
            xy=(0.5*np.pi, 0.1), xycoords='data',
            xytext=(0.5*np.pi, 2.6), textcoords='data',
            arrowprops=dict(arrowstyle="->", relpos=(0.5,0.5),
            connectionstyle="angle3,angleA=80,angleB=90") )
ax.text(0.52*np.pi, 0.7, 'preferred\norientation')
ax.plot([0, 0], [0.0, 3.5], 'k', zorder=10, clip_on=False)
xx = np.arange(0.0, 2.0*np.pi, 0.01)
pp = 0.5*np.pi
yy = np.exp(np.cos(2.0*(xx+pp)))
ax.fill_between(xx, yy+0.25*yy, yy-0.25*yy, color=cm.autumn(0.3, 1), alpha=0.5)
ax.plot(xx, yy, color=cm.autumn(0.0, 1))

ax = fig.add_axes([lmarg, 0.34, 1.0-rmarg, 0.38])
ax.spines['bottom'].set_position('zero')
ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.xaxis.set_ticks_position('bottom')
ax.get_yaxis().set_visible(False)
ax.set_xlim(0.0, np.pi)
ax.set_xticks(np.arange(0.125*np.pi, 1.*np.pi, 0.125*np.pi))
ax.set_xticklabels([])
ax.set_ylim(-1.5, 3.0)
ax.text(0.5*np.pi, -1.8, 'Orientation', ha='center')
ax.text(-0.2, 0.5*3.5, 'Activity', rotation='vertical', va='center')
ax.plot([0, 0], [0.0, 3.0], 'k', zorder=10, clip_on=False)
xx = np.arange(0.0, 1.0*np.pi, 0.01)
prefphases = np.arange(0.125*np.pi, 1.*np.pi, 0.125*np.pi)
responses = []
xresponse = 0.475*np.pi
for pp in prefphases :
    yy = np.exp(np.cos(2.0*(xx+pp)))
    ax.plot(xx, yy, color=cm.autumn(2.0*np.abs(pp/np.pi-0.5), 1))
    y = np.exp(np.cos(2.0*(xresponse+pp)))
    responses.append(y + rng.randn()*0.25*y)
    ax.plot(xresponse, y, '.', markersize=20, color=cm.autumn(2.0*np.abs(pp/np.pi-0.5), 1))
    r=0.3
    y=-0.8
    ax.plot([pp-0.5*r*np.cos(pp), pp+0.5*r*np.cos(pp)], [y-r*np.sin(pp), y+r*np.sin(pp)], 'k', lw=6)
responses = np.array(responses)

ax = fig.add_axes([lmarg, 0.05, 1.0-rmarg, 0.22])
ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.xaxis.set_ticks_position('bottom')
ax.get_yaxis().set_visible(False)
ax.set_xlim(0.0, np.pi)
ax.set_xticks(np.arange(0.125*np.pi, 1.*np.pi, 0.125*np.pi))
ax.set_xticklabels([])
ax.set_ylim(-1600, 0)
ax.set_xlabel('Orientiation')
ax.text(-0.2, -800, 'Log-Likelihood', rotation='vertical', va='center')
ax.plot([0, 0], [-1600, 0], 'k', zorder=10, clip_on=False)
phases = np.linspace(0.0, 1.1*np.pi, 100)
probs = np.zeros((len(responses), len(phases)))
for k, (pp, r) in enumerate(zip(prefphases, responses)) :
    y = np.exp(np.cos(2.0*(phases+pp)))
    sigma = 0.1*y
    probs[k,:] = np.exp(-0.5*((r-y)/sigma)**2.0)/np.sqrt(2.0*np.pi)/sigma
loglikelihood = np.sum(np.log(probs), 0)
maxl = np.max(loglikelihood)
maxp = phases[np.argmax(loglikelihood)]
ax.annotate('',
            xy=(maxp, -1600), xycoords='data',
            xytext=(maxp, -30), textcoords='data',
            arrowprops=dict(arrowstyle="->", relpos=(0.5,0.5),
            connectionstyle="angle3,angleA=80,angleB=90") )
ax.text(maxp+0.05, -1100, 'most likely\norientation\ngiven the responses')
ax.plot(phases, loglikelihood, '-b')

plt.savefig('mlecoding.pdf')
#plt.show();
