from __future__ import division, print_function

import os
import time
import subprocess

import numpy as np
import scipy as sp

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 *

os.chdir(datadir)
LCO_f = np.load('LCO-trace.npy')
LCO_x = np.load('LCO-trace-x.npy')
VCO_f = np.load('VCO-trace.npy')
VCO_x = np.load('VCO-trace-x.npy')
print('All data loaded!')
print('LCO matrix:', LCO_f.shape)
print('VCO matrix:', VCO_f.shape)

assert np.all(LCO_x == VCO_x), 'position traces were not the same!'
print('Trajectory verification passed!')

L = LCO_f.T
V = 1 - VCO_f.T
print('L matrix:', L.shape)
print('V matrix:', V.shape)

Lplus = sp.linalg.pinv(L)
print('L-inverse:', Lplus.shape)

W = np.dot(Lplus, V)
print('Weight matrix:', W.shape)

import matplotlib.pyplot as plt
from tools.colormaps import diffmap

%matplotlib inline
plt.pcolor(W,cmap=diffmap(mid=-W.min()/(W.max()-W.min())))
plt.axis("tight")
plt.colorbar()
plt.xlabel('VCO')
plt.ylabel('LCO')
f = plt.gcf()

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

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

which = 'W-matrix'
save_fn = os.path.join(datadir, '%s.npy' % which)

if os.path.exists(save_fn):
    backup_fn = os.path.join(datadir, '%s-%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)

np.save(save_fn, W)
print('Saved weight matrix to: ', save_fn)