import numpy as np
import cPickle as pickle
import joblib
from moseq.train import ARHMM, train_model
from moseq.train.util import whiten_all
from collections import OrderedDict
from syllables import analysis

# Load the data
with open("/data/efs/drugs/alldoses/dataset.pkl","r") as f:
    dataset = pickle.load(f)
mouse_names = dataset.keys()

# Load the labels
with open('/data/efs/drugs/alldoses/syllablelabels-kappa=18036000-niter=1000-nstates=160.pkl','r') as f:
    syllable_labels = pickle.load(f)
syllable_labels = analysis.relabel_by_usage(syllable_labels)

# Load the labels into our dataset array
split_points = np.cumsum([len(v['data']) for v in dataset.values()])[:-1]
split_syllable_labels = np.array_split(syllable_labels,split_points)
for mouse_name,_syllable_labels in zip(mouse_names,split_syllable_labels):
    dataset[mouse_name]['syllable_labels'] = _syllable_labels


# Make a dictionary data structure that the ARHMM expects
data_dict = OrderedDict((k,v['data']) for k,v in dataset.items())

# Whiten the data
data_dict = whiten_all(data_dict)

# Build AR matrices from labels and data.
# either scott linderman or matt johnson have code for this
# I went diving, and didn't find it.