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_vec
from tools.filters import quick_boxcar
from tools.images import tiling_dims

from trajectory import *

from brian2 import *

%matplotlib inline

f_LCO = load(os.path.join(datadir, 'LCO-trace.npy'))
W_LV = load(os.path.join(datadir, 'W-matrix.npy'))
print('f_LCO :', f_LCO.shape)
print('W_LV :', W_LV.shape)

I_LCO = dot(f_LCO.T, W_LV).T
print('I_LCO :', I_LCO.shape)
print('(min, max) =', I_LCO.min(axis=1), '→', I_LCO.max(axis=1))

print((I_LCO / I_LCO.max(axis=1)[:,newaxis]).shape)

I_LCO_norm = I_LCO - I_LCO.min(axis=1)[:,newaxis]
I_LCO_norm /= I_LCO_norm.max(axis=1)[:,newaxis]
# I_LCO_norm = I_LCO_norm**2 # sharpen feedback tuning around theta trough
print('I_LCO_norm :', I_LCO_norm.shape)
print('(min, max) =', I_LCO_norm.min(axis=1), '→', I_LCO_norm.max(axis=1))

N_VCO = 9 # 3 preferred angles x 3 spatial scales

I_LCO_t = [TimedArray(I_LCO_norm[i], 10*ms) for i in xrange(N_VCO)]

@check_units(i=1, t=second, result=1)
def f(i, t):
    return np.array([float(I_LCO_t[j](t)) for j in i])

lco_input_eqs = """
    I = f(i,t) : 1
"""

I_LCO_VCO = NeuronGroup(N_VCO, model=lco_input_eqs, namespace={'f': f})

N_VCO = 9 # 3 preferred angles x 3 spatial scales
L_min = 4 # cm
L_max = 10 # cm
f_theta = 8 * Hz
phi_0 = 2.26216937 # 2 * pi * rand()

sigma = 0.02 # noise width, radians
alpha = 0.0 # feedback gain

@check_units(a=1, result=1)
def pidiff(a=0.0):
    return circle_diff_vec(pi, a)

vco_eqs = """
    dtheta/dt = 2 * pi * f_theta + (feedback + pathint) / second + sigma / sqrt(dt) * xi : 1
    
    # Main subexpressions
    feedback = alpha * I_LCO * target : 1
    pathint = beta * speed * heading : 1
    
    # Quantities used in subexpressions
    target = pidiff(theta - theta_baseline) : 1
    heading = cos(phi - phi_pref) : 1
    
    # Neuron parameters
    beta : 1
    phi_pref : 1
    
    # Shared baseline theta wave
    theta_baseline = 2 * pi * f_theta * t : 1 (shared)
    
    # Total synaptic input from LCO population
    I_LCO : 1 (linked)
    
    # Trajectory inputs
    speed : 1 (linked)
    phi : 1 (linked)
    X : meter (linked)  # include so it is sampled at same dt
    Y : meter (linked)  # include so it is sampled at same dt
"""

VCO = NeuronGroup(N_VCO, model=vco_eqs)

VCO.I_LCO = linked_var(I_LCO_VCO, 'I')

VCO.speed = linked_var(Trajectory, 'speed')
VCO.phi = linked_var(Trajectory, 'angle')
VCO.X = linked_var(Trajectory, 'X')
VCO.Y = linked_var(Trajectory, 'Y')

VCO.theta = linspace(0, 2 * pi, N_VCO, endpoint=False)
VCO.phi_pref = repeat(array([0, 2 * pi / 3, 4 * pi / 3]) + phi_0, 3)
VCO.beta = tile(linspace(L_min, L_max, 3), 3)**-1

print(VCO.phi_pref)
print(VCO.beta)

def vco(theta, baseline=0.0):
#     return clip(cos(theta - baseline), 0, 1)
    return cos(theta - baseline)

# Plot vco(d) for across values of theta
def plot_vco_function():
    theta = linspace(-pi,pi,256)
    plot(theta, vco(theta))
    xlabel('theta (rad)')
    ylabel('firing rate')
    gca().set_xlim(theta.min(), theta.max())

plot_vco_function()

mon = StateMonitor(VCO, ['theta', 'I_LCO'], record=True)
shmon = StateMonitor(VCO, ['X', 'Y', 'theta_baseline'], record=[0])

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

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

delta_theta = mon.theta - shmon.theta_baseline
f_VCO = cos(delta_theta) # vco(mon.theta, baseline=shmon.theta_baseline)
print('min(f) =', f_VCO.min(), ', max(f) =', f_VCO.max())
print(mon.theta.shape)
print(f_VCO.shape)

f, ax = subplots(1, 1, sharex=True, figsize=(6,4))

# Plot firing rate
ax.plot(mon.t, f_VCO[0], 'r-')
ax.set_ylabel('VCO firing rate')
ax.set_ylim(bottom=0)

gridfigsize = (6,6)

r, c = tiling_dims(N_VCO)
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
if hasattr(ax, 'flatten'):
    ax = ax.flatten()
else:
    ax = [ax]

f.suptitle('sigma = %0.2f, alpha = %0.2f' % (sigma, alpha))

# Firing-rate spatial trajectory plot
def plot_rate_trajectory(i, ax):
    ax.scatter(shmon.X / cm, shmon.Y / cm, c=f_VCO[i], cmap='jet', linewidths=0)
    ax.axis('equal')
    ax.set_axis_off()
    draw_arena(ax)

for i in xrange(N_VCO):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], u'VCO %d' % i)
    
f.savefig(os.path.join('/Users/joe/projects/poster/layer-pngs', 'vco-layer-sigma%0.2f-alpha%0.2f.png' % (sigma, alpha)))

r, c = tiling_dims(N_VCO)
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
if hasattr(ax, 'flatten'):
    ax = ax.flatten()
else:
    ax = [ax]

# Firing-rate spatial trajectory plot
def plot_rate_trajectory(i, ax):
    ax.scatter(shmon.X / cm, shmon.Y / cm, c=mon.I_LCO[i], cmap='jet', linewidths=0)
    ax.axis('equal')
    ax.set_axis_off()
    draw_arena(ax)

for i in xrange(N_VCO):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], u'LCO input to VCO %d' % i)

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

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

which = 'VCO-sigma%0.2f-alpha%0.2f' % (sigma, alpha)
fn = '%s-trace.npy' % which
save_fn = os.path.join(datadir, fn)

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_VCO)
print('Saved VCO traces to: ', save_fn)

x_fn = os.path.join(datadir, 'X-trace.npy')
y_fn = os.path.join(datadir, 'Y-trace.npy')

save(x_fn, shmon.X / cm)
save(y_fn, shmon.Y / cm)

print('Saved X trace to: ', x_fn)
print('Saved Y trace to: ', y_fn)