from nn.AnimationPlotLines import animation_plot

from keras.layers import *
from keras.layers.normalization import BatchNormalization
from keras.layers.advanced_activations import ELU
from keras.models import Model
from keras import backend as K
from keras import objectives
from keras.callbacks import EarlyStopping
from keras.datasets import mnist
from keras.optimizers import Nadam
from tools.utils import load_locomotion

rng = np.random.RandomState(23455)

datasets, std, mean = load_locomotion(rng)

x_train = datasets[0][0][:320]
x_train = x_train.swapaxes(1, 2)

print x_train.shape

I = np.arange(len(x_train))
rng.shuffle(I)
x_train = x_train[I]

batch_size = 10
original_dim = 66*240
latent_dim = 100
intermediate_dim = 500
epsilon_std = 0.1
from nn.ActivationLayer import ActivationLayer
from nn.AnimationPlotLines import animation_plot
from nn.DropoutLayer import DropoutLayer
from nn.Pool1DLayer import Pool1DLayer
from nn.Conv1DLayer import Conv1DLayer
from nn.ReshapeLayer import ReshapeLayer
from nn.HiddenLayer import HiddenLayer
from nn.LSTM1DLayer import LSTM1DLayer
from nn.Network import Network, AutoEncodingNetwork, InverseNetwork

from tools.utils import load_locomotion

rng = np.random.RandomState(23455)

shared = lambda d: theano.shared(d, borrow=True)
dataset, std, mean = load_locomotion(rng)

train_motion_dataset = dataset[0][0][:100]

print "motion dataset shape = ", train_motion_dataset.shape

E = shared(train_motion_dataset)

BATCH_SIZE = 100

network = Network(
    Conv1DLayer(rng, (64, 66, 25), (BATCH_SIZE, 66, 240)),
    Pool1DLayer(rng, (2,), (BATCH_SIZE, 64, 240)),
    ActivationLayer(rng, f='elu'),

    Conv1DLayer(rng, (128, 64, 25), (BATCH_SIZE, 64, 120)),
from nn.DropoutLayer import DropoutLayer
from nn.Pool1DLayer import Pool1DLayer
from nn.Conv1DLayer import Conv1DLayer
from nn.ReshapeLayer import ReshapeLayer
from nn.HiddenLayer import HiddenLayer
from nn.VariationalLayerUnflattened import VariationalLayer
from nn.LSTM1DHiddenInitLayer import LSTM1DLayer
from nn.Network import Network, AutoEncodingNetwork, InverseNetwork

from tools.utils import load_locomotion
from tools.utils import load_terrain

rng = np.random.RandomState(23455)

shared = lambda d: theano.shared(d, borrow=True)
dataset1, std1, mean1 = load_locomotion(rng)
dataset2, std2, mean2 = load_terrain(rng)

dataset = np.concatenate([dataset1[0][0][:300], dataset2[0][0]], axis=0)

train_control_dataset = dataset[:1800]
train_control_dataset = train_control_dataset.swapaxes(0, 1)[-3:]
train_control_dataset = train_control_dataset.swapaxes(0, 1)

print "control dataset shape = ", train_control_dataset.shape

E = shared(train_control_dataset)

dataset1, std1, mean1 = load_locomotion(rng)
dataset2, std2, mean2 = load_terrain(rng)