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

rng = np.random.RandomState(4637281)
lmarg=0.1
rmarg=0.1

fig = plt.figure(figsize=cm_size(figure_width, 2.8*figure_height))
spec = gridspec.GridSpec(nrows=4, ncols=1, height_ratios=[4, 4, 1, 3], hspace=0.2,
                         **adjust_fs(fig, left=4.0))
ax = fig.add_subplot(spec[0, 0])
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.yaxis.set_major_locator(plt.NullLocator())
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')
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, **fsBa)
ax.plot(xx, yy, **lsB)

ax = fig.add_subplot(spec[1, 0])
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.0)
ax.yaxis.set_major_locator(plt.NullLocator())
ax.text(-0.2, 0.5*3.5, 'Activity', rotation='vertical', va='center')
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, ls, ps in zip(prefphases, [lsE, lsC, lsD, lsB, lsD, lsC, lsE],
                      [psE, psC, psD, psB, psD, psC, psE]) :
    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))
    ax.plot(xx, yy, **ls)
    y = np.exp(np.cos(2.0*(xresponse+pp)))
    responses.append(y + rng.randn()*0.25*y)
    ax.plot(xresponse, y, **ps)
responses = np.array(responses)

ax = fig.add_subplot(spec[2, 0])
ax.show_spines('')
r = 0.3
ax.set_ylim(-1.1*r, 1.1*r)
for pp in prefphases:
    ax.plot([pp-0.5*r*np.cos(pp), pp+0.5*r*np.cos(pp)], [-r*np.sin(pp), r*np.sin(pp)],
            colors['black'], lw=6, clip_on=False)

ax = fig.add_subplot(spec[3, 0])
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.yaxis.set_major_locator(plt.NullLocator())
ax.set_xlabel('Orientation')
ax.text(-0.2, -800, 'Log-Likelihood', rotation='vertical', va='center')
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, **lsA)

plt.savefig('mlecoding.pdf')