import sys
print sys.executable
%load_ext autoreload
%autoreload 2
%reload_ext autoreload

import sonnet as snt
import tensorflow as tf
import tflearn
import numpy as np

import dataset_utils

"""
Explanation: Sketchbook to experiment with RNN models to model the dynamics. Once tested, the code in this notebook will be incorporated into functions/classes/ python modules.
Dynamics Model
Models the dynamics of the system p(next observation | history of observations, action)
History of observations consist of exercises a student has done and whether the student solved each of them
Action is the next exercise chosen 
Next observation is whether the student gets the chosen exercise correct
We want to use an RNN to model the dynamics.
Input data represents history of observations, of shape (n_students, n_timesteps, observation_vec_size)
Output represents the probability of getting next exercise correctly, of shape (n_students, n_timesteps, n_exercises) 
So at each timestep, we make a prediction for all actions. 
For each action, the output vector specifies the predicted probability of the student getting the chosen exercise correctly.
The target output only contains binary values.
End of explanation
"""


data = dataset_utils.load_data(filename="../synthetic_data/toy.pickle")
input_data_, output_mask_, target_data_ = dataset_utils.preprocess_data_for_rnn(data)

"""
Explanation: Load toy data set
End of explanation
"""


tflearn.init_graph()

n_hidden = 64
n_samples, n_timesteps, n_inputdim = input_data_.shape
_,_,n_outputdim = target_data_.shape

print n_timesteps

print n_inputdim

print n_outputdim

graph_to_use = tf.Graph()
with graph_to_use.as_default():
    net = tflearn.input_data([None, n_timesteps, n_inputdim],dtype=tf.float32, name='input_data')
    output_mask = tflearn.input_data([None, n_timesteps, n_outputdim], dtype=tf.float32, name='output_mask')
    net = tflearn.lstm(net, n_hidden, return_seq=True, name="lstm_1")
    net = tflearn.lstm(net, n_outputdim, return_seq=True, name="lstm_2")
    net = tf.stack(net, axis=1)
    preds = net
    net = net * output_mask
    net = tflearn.regression(net, optimizer='adam', learning_rate=0.001,
                             loss='mean_square')
    model = tflearn.DNN(net, tensorboard_verbose=0)
    model.fit([ input_data_, output_mask_], target_data_, validation_set=0.1)

tf.get_collection(tf.GraphKeys.INPUTS)

data = dataset_utils.load_data(filename="../synthetic_data/1000stud_100seq_expert.pickle")
input_data_, output_mask_, target_data_ = dataset_utils.preprocess_data_for_rnn(data)

n_samples, n_timesteps, n_inputdim = input_data_.shape
_,_,n_outputdim = target_data_.shape
print n_samples
print n_timesteps
print n_inputdim
print n_outputdim

graph_to_use = tf.Graph()
with graph_to_use.as_default():
    net = tflearn.input_data([None, n_timesteps, n_inputdim],dtype=tf.float32, name='input_data')
    output_mask = tflearn.input_data([None, n_timesteps, n_outputdim], dtype=tf.float32, name='output_mask')
    net = tflearn.lstm(net, n_hidden, return_seq=True, name="lstm_1")
    net = tflearn.lstm(net, n_outputdim, return_seq=True, name="lstm_2")
    net = tf.stack(net, axis=1)
    preds = net
    net = net * output_mask
    net = tflearn.regression(net, optimizer='adam', learning_rate=0.001,
                             loss='mean_square')
    model = tflearn.DNN(net, tensorboard_verbose=2)
    model.fit([ input_data_, output_mask_], target_data_, n_epoch=32, validation_set=0.1)

from dynamics_model import *

model = load_model(model_id="test_model", load_checkpoint=False, is_training=True)

train_data = (input_data_[:,:10,:], output_mask_[:,:10,:], target_data_[:,:10,:])

train(model, train_data)

predict(model, input_data_[:10, :10, :])

import dynamics_model_class as dm

dmodel = dm.DynamicsModel(model_id="test_model",  load_checkpoint=False)

dmodel.train(train_data)

preds = dmodel.predict(input_data_[:10,:10, :])

generator_model = dm.DynamicsModel(model_id="test_model", timesteps=1, load_checkpoint=False)

preds = generator_model.predict(input_data_[:10,:1, :])

print preds[0]

"""
Explanation: Build RNN Model
End of explanation
"""


