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

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

X = load(x_fn)
Y = load(y_fn)

print('Loaded X trace: ', X.shape)
print('Loaded Y trace: ', Y.shape)

sigma = 0.01  # noise level to use
alpha = 0.1  # feedback gain to use

traces = [
    { 'name': 'nonoise',  'sigma': 0.00,  'alpha': 0.00 },
    { 'name': 'noise',    'sigma': sigma, 'alpha': 0.00 },
    { 'name': 'feedback', 'sigma': sigma, 'alpha': alpha },
]

def data_path(d):
    which = 'VCO-sigma%0.2f-alpha%0.2f' % (d['sigma'], d['alpha'])
    fn = '%s-trace.npy' % which
    return os.path.join(datadir, fn)

for trace in traces:
    trace['data'] = load(data_path(trace))
    print('Loaded \'%s\' trace:' % trace['name'], trace['data'].shape)

grids = []
for trace in traces:
    data = trace['data']
    grids.append(data[0::3].sum(axis=0))
    grids.append(data[1::3].sum(axis=0))
    grids.append(data[2::3].sum(axis=0))

# Normalize grid responses
grid_min_clip = -0.5
for g in grids:
    g += 3 - g.max()
    g[:] = clip(g, grid_min_clip, 3.0)
    
print('min(grids) =', [g.min() for g in grids])
print('max(grids) =', [g.max() for g in grids])

gridfigsize = (6,6)

r, c = tiling_dims(len(grids))
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
ax = ax.flatten()

trace = traces[0]
f.suptitle('%s\nsigma = %0.2f, alpha=%0.2f' % (trace['name'], trace['sigma'], trace['alpha']))

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

for i in xrange(len(grids)):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], 'VCO %d' % i)

f_nonoise = f

r, c = tiling_dims(len(grids))
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
ax = ax.flatten()

trace = traces[1]
f.suptitle('%s\nsigma = %0.2f, alpha=%0.2f' % (trace['name'], trace['sigma'], trace['alpha']))

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

for i in xrange(len(grids)):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], 'VCO %d' % i)

f_noise = f

r, c = tiling_dims(len(grids))
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
ax = ax.flatten()

trace = traces[2]
f.suptitle('%s\nsigma = %0.2f, alpha=%0.2f' % (trace['name'], trace['sigma'], trace['alpha']))

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

for i in xrange(len(grids)):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], 'VCO %d' % i)

f_feedback = f

r, c = tiling_dims(len(grids))
f, ax = subplots(r, c, sharey=True, sharex=True, figsize=gridfigsize)
f.suptitle('Grid Cells')
ax = ax.flatten()

# Firing-rate spatial trajectory plot
def plot_rate_trajectory(i, ax):
    ax.scatter(X, Y, c=grids[i], cmap='jet', vmin=grid_min_clip, vmax=3, linewidths=0)
    ax.axis('equal')
    ax.set_axis_off()
    draw_arena(ax)

for i in xrange(len(grids)):
    plot_rate_trajectory(i, ax[i])
    quicktitle(ax[i], 'Grid Cell %d' % i)

f_grids = f

posterdir = '/Users/joe/projects/poster'

# nonoise
desc = 'sigma%0.2f_alpha%0.2f' % (traces[0]['sigma'], traces[0]['alpha'])
f_nonoise.savefig(os.path.join(posterdir, 'VCO_layer_%s.png' % desc), dpi=600)

# noise
desc = 'sigma%0.2f_alpha%0.2f' % (traces[1]['sigma'], traces[1]['alpha'])
f_noise.savefig(os.path.join(posterdir, 'VCO_layer_%s.png' % desc), dpi=600)

# feedback
desc = 'sigma%0.2f_alpha%0.2f' % (traces[2]['sigma'], traces[2]['alpha'])
f_feedback.savefig(os.path.join(posterdir, 'VCO_layer_%s.png' % desc), dpi=600)

# grids
desc = 'sigma%0.2f_alpha%0.2f' % (sigma, alpha)
f_grids.savefig(os.path.join(posterdir, 'Grids-%s.png' % desc), dpi=600)