# Import required libraries
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing import sequence
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Embedding, SimpleRNN, Dense,GRU
import pickle
import numpy as np
import matplotlib.pyplot as plt
from prettytable import PrettyTable
from pprint import pprint

# We fix a vocabulary size of 5000
# Use the code below to call a small subset of the imdb dataset
# We keep the vocabulary size fixed because it was used to curate the sub-dataset
vocabulary_size = 5000
with open('imdb_mini.pkl','rb') as f:
    X_train, y_train, X_test, y_test = pickle.load(f)

# Run the code below to see the first tokenized review and the label
print('---review---')
print(X_train[0])
print('---label---')
print(y_train[0])

# You can get the word2id mapping by
# using the imdb.get_word_index() function

word2id = imdb.get_word_index()
# We need to adjust the mapping by 3 because of tensorflow.keras preprocessing
# more here: https://stackoverflow.com/questions/42821330/restore-original-text-from-keras-s-imdb-dataset
word2id = {k:(v+3) for k,v in word2id.items()}
word2id[""] = 0
word2id[""] = 1
word2id[""] = 2
word2id[""] = 3

# Reversing the key,value pair will give the id2word
id2word = {i: word for word, i in word2id.items()}
pprint('---review with words---')
pprint(" ".join([id2word[i] for i in X_train[0]]))
pprint('---label---')
pprint(y_train[0]);

### edTest(test_chow1) ###
# Submit an answer choice as a string below (eg. if you choose option C, put 'C')
answer1 = '___'

# For training we need our sequences to be of fixed length, but the reviews
# are of different sizes
print(f'Maximum review length: {len(max([i for i in X_train]+[i for i in X_test], key=len))}')
print(f'Minimum review length: {len(min([i for i in X_train]+[i for i in X_test], key=len))}')

# We also create two indices for short and long reviews
# we will use this later
idx_short = [i for i,val in enumerate(X_train) if len(val)<100]
idx_long = [i for i,val in enumerate(X_train) if len(val)>500]

### edTest(test_chow2) ###
# Submit an answer choice as a string below (eg. if you choose option C, put 'C')
answer2 = '___'

# We will clip large reviews and pad smaller reviews to 500 words
max_words = 500
# We can pad the smaller sequences with 0s before, or after.
# This choice can severely affect network performance
# In the first case we, will pad after the sequence 
# Please utilize sequence.pad_sequences()
postpad_X_train = ___
postpad_X_test = ___

def model_maker(summary=True,gru=False):
# One layer RNN model with 32 rnn cells
    embedding_size=32
    model=Sequential()
    model.add(Embedding(vocabulary_size, embedding_size, input_length=max_words))
    
    # We can specify if we want the GRU cell or the vanilla RNN cell
    if gru:
        model.add(GRU(32))
    else:
        model.add(SimpleRNN(32))
    model.add(Dense(1, activation='sigmoid'))
    if summary:
        print(model.summary())
        
    # model compile step
    model.compile(loss='binary_crossentropy',
                 optimizer='adam',
                 metrics=['accuracy'])
    return model

# We build two new models with vanilla RNNs
model_short = model_maker(summary=False)
model_long = model_maker(summary=False)

# First we train `model_short` with short reviews
# set epochs to 10

X_short = ___
y_short = ___
epochs = ___
history_short = model_short.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);

# Then we train `model_long` with short reviews

X_long = ___
y_long = ___
history_long = model_long.fit(___,___, epochs=epochs,batch_size=640,verbose=0);

### edTest(test_chow2_1) ###
X_short_shape, y_short_shape = X_short.shape, y_short.shape
X_long_shape, y_long_shape = X_long.shape, y_long.shape
print(X_short_shape, y_short_shape,X_long_shape, y_long_shape)

# Helper function to plot the data
# Plot the MSE of the model
plt.rcParams["figure.figsize"] = (8,6)
plt.title("Padding='post'",fontsize=20)
plt.semilogy(history_short.history['loss'], label='Shorter reviews', color='#FF9A98', linewidth=3)
plt.semilogy(history_long.history['loss'],  label='Longer reviews', color='#75B594', linewidth=3)
plt.legend()

# Set the axes labels
plt.xlabel('Epochs',fontsize=14)
plt.xticks(range(1,epochs,4))
plt.ylabel('MSE Loss',fontsize=14)
plt.legend(fontsize=14)
plt.show()

# We can pre-pad by using `sequence.pad_sequences` with `padding='pre'`
max_words = 500
prepad_X_train = ___
prepad_X_test = ___

# Reinitializing the models for the two categories
model_short = model_maker(summary=False)
model_long = model_maker(summary=False)

# Again we train `model_short` with short reviews

X_short = ___
y_short = ___
epochs = 10
history_short = model_short.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);

# Then we train `model_long` with short reviews

X_long = ___
y_long = ___
history_long = model_long.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);

# Helper function to plot the data
# Plot the MSE of the model
plt.rcParams["figure.figsize"] = (8,6)
plt.title("Padding='pre'",fontsize=20)
plt.semilogy(history_short.history['loss'], label='Shorter reviews', color='#FF9A98', linewidth=3)
plt.semilogy(history_long.history['loss'],  label='Longer reviews', color='#75B594', linewidth=3)
plt.legend()

# Set the axes labels
plt.xlabel('Epochs',fontsize=14)
plt.xticks(range(1,epochs,4))
plt.ylabel('MSE Loss',fontsize=14)
plt.legend(fontsize=14)
plt.show()

### edTest(test_chow3) ###
# Submit your answer as a string below
answer3 = '___'