#!/usr/bin/env python
# coding: utf-8

# # High-level RNN Keras (CNTK) Example

# In[1]:


import os
import sys
import numpy as np
os.environ['KERAS_BACKEND'] = "cntk"
import keras as K
import cntk
from keras.models import Sequential
from keras.layers import Dense, Embedding, GRU, CuDNNGRU
from common.params_lstm import *
from common.utils import *


# In[2]:


print("OS: ", sys.platform)
print("Python: ", sys.version)
print("Keras: ", K.__version__)
print("Numpy: ", np.__version__)
print("CNTK: ", cntk.__version__)
print(K.backend.backend())


# In[3]:


def create_symbol(CUDNN=True):
    model = Sequential()
    model.add(Embedding(MAXFEATURES, EMBEDSIZE, input_length=MAXLEN))
    # Only return last output
    if not CUDNN:
        model.add(GRU(NUMHIDDEN, return_sequences=False, return_state=False))
    else:
        model.add(CuDNNGRU(NUMHIDDEN, return_sequences=False, return_state=False))
    model.add(Dense(2, activation='softmax'))
    return model


# In[4]:


def init_model(m):
    m.compile(
        loss = "categorical_crossentropy",
        optimizer = K.optimizers.Adam(LR, BETA_1, BETA_2, EPS),
        metrics = ['accuracy'])
    return m


# In[5]:


get_ipython().run_cell_magic('time', '', '# Data into format for library\nx_train, x_test, y_train, y_test = imdb_for_library(seq_len=MAXLEN, max_features=MAXFEATURES, one_hot=True)\nprint(x_train.shape, x_test.shape, y_train.shape, y_test.shape)\nprint(x_train.dtype, x_test.dtype, y_train.dtype, y_test.dtype)\n')


# In[7]:


get_ipython().run_cell_magic('time', '', '# Load symbol\n# CuDNN RNNs are only available with the TensorFlow backend.\nsym = create_symbol(CUDNN=False)\n')


# In[8]:


get_ipython().run_cell_magic('time', '', '# Initialise model\nmodel = init_model(sym)\n')


# In[9]:


model.summary()


# In[10]:


get_ipython().run_cell_magic('time', '', '# 86s\n# Train model\nmodel.fit(x_train,\n          y_train,\n          batch_size=BATCHSIZE,\n          epochs=EPOCHS,\n          verbose=1)\n')


# In[11]:


get_ipython().run_cell_magic('time', '', 'y_guess = model.predict(x_test, batch_size=BATCHSIZE)\ny_guess = np.argmax(y_guess, axis=-1)\ny_truth = np.argmax(y_test, axis=-1)\n')


# In[12]:


print("Accuracy: ", sum(y_guess == y_truth)/len(y_guess))