from __future__ import division, print_function

import os
import time
import subprocess

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

from trajectory import *

%matplotlib inline

from brian2 import *

N_BVC = 64
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)

N_LVC = 64
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

# Create the state monitors
mon_BVC = StateMonitor(BVC, 'f', record=True)
mon_LVC = StateMonitor(LVC, 'f', record=True)
mon_pos = StateMonitor(LVC, ('X', 'Y'), record=[0])

# Reset the clock and run the simulation
defaultclock.dt = 10 * ms
run(300.0 * second)

N_LCO = 128
thresh = 0.4 * mA

C_BVC = 6
C_LVC = 2

def v_LCO(I, thresh=0.5 * mA):
    return tanh(2 * (I - thresh) / thresh)

# Plot v_LCO(d) for different slope values
current = linspace(0,2.5,100) * mA
for t in array([0.2,0.4,0.6,0.8,1.0]) * mA:
    plt.plot(current / mA, v_LCO(current, thresh=t), label='%.1f mA' % (t / mA))
plt.xlabel('current (mA)')
plt.ylabel('V')
plt.legend(loc='lower right');
quicktitle(gca(), 'LCO V-I curve');
f = plt.gcf()

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

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

def make_weights(n, c):
    W = zeros((n, N_LCO), 'd')
    for j in xrange(N_LCO):
        con = permutation(N_BVC)[:c]
        W[con,j] = 1.0
    print('C = ', W.sum(axis=0), '(should be all %d\'s)' % c)
    return W

W_BVC = make_weights(N_BVC, C_BVC) / C_BVC
W_LVC = make_weights(N_BVC, C_LVC) / C_LVC

print('W_BVC = ', W_BVC.sum(axis=0), '(should be all 1\'s)')
print('W_LVC = ', W_LVC.sum(axis=0), '(should be all 1\'s)')

I_BVC = dot(mon_BVC.f.T, W_BVC).T 
I_LVC = dot(mon_LVC.f.T, W_LVC).T 

g_BL = 0.25
I = g_BL * I_BVC + (1 - g_BL) * I_LVC

V = tanh(2 * (I - (thresh / mA)) / (thresh / mA))
f_LCO = clip(V, 0, 1)
Theta = arccos(V)

f = figure(figsize=(12,4))
plot(mon_BVC.t, I[0], c='c', label='I')
plot(mon_BVC.t, V[0], c='b', label='V')
plot(mon_BVC.t, f_LCO[0], 'b--', lw=3, label='f')
plot(mon_BVC.t, Theta[0], 'r-', lw=2, alpha=0.5, label=r'$\theta$-phase')
xlabel('time (s)')
ylabel(r'I (mA), V (mV), f (a.u.), $\theta$-phase (rad)')
legend(loc='upper left')

print('I ~', I.min(), 'to', I.max())
print('V ~', V.min(), 'to', V.max())
print('f ~', f_LCO.min(), 'to', f_LCO.max())

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

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

N_LCO_toshow = min(N_LCO, 16)
N_plots = 2 * N_LCO_toshow

from tools.images import tiling_dims
r, c = (4, 8) # tiling_dims(N_plots)

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

def plot_phase_trajectory(cdata, ax, cmap='jet', **kwds):
    ax.scatter(mon_pos.X[0] / cm, mon_pos.Y[0] / cm, c=cdata, cmap=cmap, linewidths=0, **kwds)
    ax.axis('equal')
    ax.set_axis_off()
    draw_arena(ax)

j = 0
for i in xrange(N_LCO_toshow):
    for ttl, val, cmap in [('rate', f_LCO, 'jet'), ('phase', Theta, 'copper_r')]:
        scatter_kwds = {}
        if ttl == 'phase':
            scatter_kwds.update(vmin=0, vmax=pi)
        plot_phase_trajectory(val[i], ax[j], cmap, **scatter_kwds)
        quicktitle(ax[j], 'LCO %d %s' % (i, ttl))
        j += 1

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

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

which = 'LCO'
save_fn = os.path.join(datadir, '%s-trace.npy' % which)
save_fn_x = os.path.join(datadir, '%s-trace-x.npy' % which)

if os.path.exists(save_fn):
    backup_fn = os.path.join(datadir, '%s-trace-%s.npy' % (which, time.strftime('%Y-%m-%d-%H-%M-%S')))
    if subprocess.call(['mv', save_fn, backup_fn]) == 0:
        print('Saved backup to: ', backup_fn)

save(save_fn, f_LCO)
save(save_fn_x, mon_pos.X[0] / cm)