from __future__ import division, print_function

import scipy.io as sio

import os
root = '/Users/joe/data/blab/'
dataset = os.path.join(root, 'OTheta10/OTheta10_Day8')
assert os.path.isdir(dataset)

tracking_file = os.path.join(dataset, 'trackdata.mat')
trackdata = sio.loadmat(tracking_file)

print(*sorted(trackdata.keys()), sep=', ')

start = 1000.0 # s
duration = 1500.0 # s
fs = 30.0 # s
res = 5.25 # pixels/cm

t0 = int(start * fs)
N = int(duration * fs)
t_slice = slice(t0, t0 + N)

print('Starting at sample %d for a total of %d samples.' % (t0, N))

x, t = trackdata['Position_X'][t_slice].T
y = trackdata['Position_Y'][t_slice,0]
print('Data starts at t = %f.' % t[0])

# Convert from pixels to cm
x /= res
y /= res

import numpy as np
pos = np.c_[t, x, y].T
print('New position shape = ', pos.shape)

%matplotlib inline
import matplotlib.pyplot as plt

plt.plot(x, y, 'k-', alpha=0.7)
plt.axis('equal')
plt.xlabel('x')
plt.ylabel('y')

model_path = '/Users/joe/source/phase_model'
data_path = os.path.join(model_path, 'data')
assert os.path.isdir(data_path)

save_file = os.path.join(data_path, 'pos.npy')
np.save(save_file, pos)
print('Saved:', save_file)

pos_saved = np.load(save_file)
plt.plot(pos_saved[1], pos_saved[2], 'b-', alpha=0.7)
plt.axis('equal')