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_BVC = 16
BVC_width = 10 * cm
BVC_rpref_min = 10 * cm
BVC_rpref_max = 25 * cm

bvc_eqs = """
    f = (D / D_arena) * exp(-(D - D_arena + rpref)**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)
"""

BVC = NeuronGroup(N_BVC, model=bvc_eqs)

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

BVC.Xf, BVC.Yf = rim_sample(N_BVC)
BVC.rpref = permutation(linspace(BVC_rpref_min, BVC_rpref_max, N_BVC)) * meter
BVC.width = sqrt(BVC_width * BVC.rpref)

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

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

    ax2.hist(BVC.rpref / cm)
    ax2.set_xlabel('Preferred distance (cm)')
    print('unit(width) =', get_unit(BVC.width))
    return f

f = plot_bvc_parameters()

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

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

def bvc(D, rpref=20 * cm, width=10*cm):
    width = sqrt(width * rpref)
    return 1.2 * (D / D_arena) * exp(-(D - D_arena + rpref)**2 / (2 * width**2))

# Plot bvc(d) for different values of g_radius
def plot_bvc_function():
    x = linspace(0,100,100) * cm
    for r in np.linspace(BVC_rpref_min, BVC_rpref_max, 8):
        plt.plot(x / cm, bvc(x, rpref=r), label='%d cm' % (r / cm))
    plt.vlines([80], 0, 1.2, linestyles='--', zorder=-10)
    plt.xlabel('Reference Distance (cm)')
    plt.ylabel('Firing Rate')
    plt.legend(loc='upper left');
    return plt.gcf()
f = plot_bvc_function()

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

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

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

defaultclock.dt = 10 * ms
run(300 * second)

plot(posmon.X[0] / cm, posmon.Y[0] / cm);
draw_arena()
axis('equal')
axis([-70,70,-50,50])

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_BVC)

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

f.suptitle('Boundary Vector Cell Layer')

# Phase-coded spatial trajectory plot
def plot_phase_trajectory(i, ax):
    ax.plot(BVC.Xf[i] / cm, BVC.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_BVC):
    plot_phase_trajectory(i, ax[i])
    quicktitle(ax[i], 'BVC %d' % i)

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

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