from __future__ import division, print_function

import os

from tools.plot import quicktitle
from tools.radians import xy_to_rad_vec
from tools.filters import quick_boxcar
from tools.images import tiling_dims

from trajectory import *

from brian2 import *

%matplotlib inline

N_LVC = 16
LVC_width_min = 15 * cm
LVC_width_max = 30 * cm

lvc_eqs = """
    f = exp(-D**2 / (2 * width**2)) : 1
    D = sqrt((X - Xf)**2 + (Y - Yf)**2) : meter
    rpref : meter
    width : meter
    Xf : meter
    Yf : meter
    X : meter (linked)
    Y : meter (linked)
"""

LVC = NeuronGroup(N_LVC, model=lvc_eqs)

LVC.X = linked_var(Trajectory, 'X')
LVC.Y = linked_var(Trajectory, 'Y')

LVC.Xf, LVC.Yf = arena_sample(N_LVC)
LVC.width = permutation(linspace(LVC_width_min, LVC_width_max, N_LVC)) * meter

def plot_lvc_parameters():
    f, ax = subplots(1,2, figsize=(13,4))
    ax1, ax2 = ax
    f.suptitle('LVC neuron parameters')

    ax1.scatter(LVC.Xf / cm, LVC.Yf / cm, marker='.', s=49, color='r')
    draw_arena(ax1)
    ax1.axis('equal')
    ax1.axis([-70,70,-50,50])
    ax1.set_xlabel('Landmark locations')
    print('unit(Xf) =', get_unit(LVC.Xf))

    ax2.hist(LVC.width / cm)
    ax2.set_xlabel('Landmark width')
    print('unit(width) =', get_unit(LVC.width))
    return f

f = plot_lvc_parameters()

imagefn = 'LVC_parameters.pdf'
posterdir = '/Users/joe/projects/poster'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath)

def lvc(D, rpref=20 * cm, width=10*cm):
    width = sqrt(width * rpref)
    return exp(-D**2 / (2 * width**2)) / sqrt(2 * pi * width)

# Plot bvc(d) for different values of g_radius
def plot_lvc_function():
    x = linspace(-R_arena, R_arena, 100)
    for r in np.linspace(LVC_width_min, LVC_width_max, 8):
        plt.plot(x / cm, lvc(x, rpref=r), label='%d cm' % (r / cm))
    plt.xlabel('distance (cm)')
    plt.ylabel('bvc(distance)')
    plt.legend(loc='upper right');
    return plt.gcf()
f = plot_lvc_function()

imagefn = 'LVC_width_profile.pdf'
posterdir = '/Users/joe/projects/poster'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath)

mon = StateMonitor(LVC, ('f', 'D'), record=True)
posmon = StateMonitor(LVC, ('X', 'Y'), record=[0])

defaultclock.dt = 10 * ms
defaultclock.start = 0.0 * second

run(300 * second)

plot(posmon.X[0] / cm, posmon.Y[0] / cm);
draw_arena()
axis('equal')
axis([-70,70,-50,50])
xlabel('X (cm)')
ylabel('Y (cm)')
f = plt.gcf()

imagefn = 'plain_trajectory.pdf'
posterdir = '/Users/joe/projects/poster'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath)

f, ax = subplots(2, 1, sharex=True, figsize=(8,5))
ax1, ax2 = ax

# Plot distance from preferred location
ax1.plot(mon.t, mon.D[0])
ax1.set_ylabel('distance (cm)')
ax1.set_ylim(bottom=0)

# Plot firing rate
ax2.plot(mon.t, mon.f[0], 'r-')
ax2.set_ylabel('firing rate (au)');

r, c = (2, 8) # tiling_dims(N_LVC)

f, ax = subplots(r, c, sharey=True, sharex=True, figsize=(16,4))
ax = ax.flatten()

f.suptitle('Landmark Vector Cell Layer')

# Phase-coded spatial trajectory plot
def plot_phase_trajectory(i, ax):
    ax.plot(LVC.Xf[i] / cm, LVC.Yf[i] / cm, 'm+', ms=20, mew=3)
    ax.scatter(posmon.X / cm, posmon.Y / cm, c=mon.f[i], cmap='jet', linewidths=0)
    ax.axis('equal')
    ax.set_axis_off()
    draw_arena(ax)

for i in xrange(N_LVC):
    plot_phase_trajectory(i, ax[i])
    quicktitle(ax[i], 'LVC %d' % i)

imagefn = 'LVC_examples.png'
posterdir = '/Users/joe/projects/poster'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath, dpi=600)