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 *

from brian2 import *

%matplotlib inline

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.0 # noise width, radians

vco_eqs = """
    dtheta/dt = 2 * pi * f_theta + beta * speed * v / second + sigma / sqrt(dt) * xi : 1
    v = cos(phi - phi_pref) : 1
    theta_baseline = 2 * pi * f_theta * t : 1 (shared)
    beta : 1
    phi_pref : 1
    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.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('Relative Theta Phase (rad)')
    ylabel('VCO Firing Rate')
    gca().set_xlim(theta.min(), theta.max())
    return plt.gcf()

f = plot_vco_function()

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

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

mon = StateMonitor(VCO, 'theta', 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)

r, c = tiling_dims(N_VCO)
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=(12,12))
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=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], 'VCO %d' % i)

which = 'VCO'
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_VCO)
save(save_fn_x, shmon.X[0] / cm)