import os
import pickle
from collections import Counter

import numpy as np

"""
Explanation: TV script generator
Generate TV script from the Simpsons.
This is longer than the original as I've added intermediate steps. Helper functions are incorporated here rather than being in a separate file.
End of explanation
"""


data_file = './data/simpsons/moes_tavern_lines.txt'

with open(data_file) as inf:
     text= inf.read()

# Ignore the "notice" section
text = text[81:]
print(text[:100])
len(text)


"""
Explanation: Get data
End of explanation
"""


view_sentence_range = (0, 10)

# original:
#num_unique_words = len({word: None for word in text.split()})
#print('Roughly the number of unique words: {}'.format()

## Here we use a counter to find the size of the vocab

wc = Counter(text.lower().split())

print('Vocab:', len(wc))
print('Total words:', sum(wc.values()))
print('Common words:', wc.most_common(10), '...')

scenes = text.split('\n\n')
print('Number of scenes: {}'.format(len(scenes)))

sentence_count_scene = [scene.count('\n') for scene in scenes]
print('Average number of sentences in each scene: {}'.format(np.average(sentence_count_scene)))

sentences = [sentence for scene in scenes for sentence in scene.split('\n')]
print('Number of lines: {}'.format(len(sentences)))

word_count_sentence = [len(sentence.split()) for sentence in sentences]
print('Average number of words in each line: {}'.format(np.average(word_count_sentence)))

print()
print('Sentence {} to {}:'.format(*view_sentence_range))
print()
print('\n'.join(text.split('\n')[view_sentence_range[0]:view_sentence_range[1]]))

"""
Explanation: Explore different parts of the data
End of explanation
"""


# Create vocab_to_int and int_to_vocab.
# This also sorts from most frequent to least frequent.
# For example, integer 0 represents the most frequently used vocab.

vocab_to_int = {}
int_to_vocab = {}

for i, (word, cnt) in enumerate(wc.most_common()):
    vocab_to_int[word] = i
    int_to_vocab[i] = word

[(k, int_to_vocab[k]) for k in list(int_to_vocab.keys())[:5]]

[(k, vocab_to_int[k]) for k in list(vocab_to_int.keys())[:5]]

## make it into a function suitable for the original project

def create_lookup_tables(text):
    """
    Create lookup tables for vocabulary
    :param text: The text of tv scripts split into words
    :return: A tuple of dicts (vocab_to_int, int_to_vocab)
    """
    wc = Counter(text)
    
    vocab_to_int = {}
    int_to_vocab = {}
    
    for i, (word, cnt) in enumerate(wc.most_common()):
        vocab_to_int[word] = i
        int_to_vocab[i] = word
        
    return vocab_to_int, int_to_vocab

## quick test

text2 = "today is an interesting day, is it not?"
text_tokens = text2.lower().split()
v2i, i2v = create_lookup_tables(text_tokens)

v2i

i2v

## test with given unit test

import problem_unittests as tests
tests.test_create_lookup_tables(create_lookup_tables)

"""
Explanation: Convert vocab to integer representation
End of explanation
"""


## the string package has a pre-defined list of punctuations

import string
type(string.punctuation), string.punctuation

## create a dictionary according to project

def token_lookup():
    """
    Generate a dict to turn punctuation into a token.
    :return: Tokenize dictionary where the key is the punctuation and the value is the token
    """
    #import string
    punct_to_token = {'.':'||PERIOD||', 
                      ',':'||COMMA||', 
                      '"':'||QUOTE||', 
                      ';':'||SEMICOLON||', 
                      '!':'||EXCMARK||', 
                      '?':'||Q||', 
                      '(':'||OPENP||', 
                      ')':'||CLOSEP||', 
                      '--':'||DASH||',
                      '\n':'||NEWLINE||'}
    return punct_to_token

# quick test
token_dict = token_lookup()
token_dict

## test with built-in test
tests.test_tokenize(token_lookup)

"""
Explanation: tokenize punctuations
We need to distinguish between "bye" and "bye!".
So we will translate punctuations to tokens. For example, change "!" into "||Exclamation_Mark||". 

Period ( . )
Comma ( , )
Quotation Mark ( " )
Semicolon ( ; )
Exclamation mark ( ! )
Question mark ( ? )
Left Parentheses ( ( )
Right Parentheses ( ) )
Dash ( -- )
Return ( \n )
End of explanation
"""


for key, token in token_dict.items():
    text = text.replace(key, ' {} '.format(token))

text = text.lower()
text = text.split()

vocab_to_int, int_to_vocab = create_lookup_tables(text)
int_text = [vocab_to_int[word] for word in text]

with open('preprocess.p', 'wb') as outf:
    pickle.dump((int_text, vocab_to_int, int_to_vocab, token_dict), outf)

len(int_text)

_ = [print(int_to_vocab[item], end=" ") for item in int_text[:15]]

_ = [print(int_to_vocab[item], end=" ") for item in int_text[16:32]]

## check current directory
os.listdir('.')

"""
Explanation: Preprocess data
End of explanation
"""


with open('preprocess.p', mode='rb') as inf:
    int_text, vocab_to_int, int_to_vocab, token_dict = \
        pickle.load(inf)
        
len(int_text), len(vocab_to_int), len(int_to_vocab), len(token_dict)

"""
Explanation: Check Point
Simply reload saved files from this point forward.
End of explanation
"""


from distutils.version import LooseVersion
import warnings
import tensorflow as tf

# Check TensorFlow Version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer'
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
    warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))

"""
Explanation: Build NN
End of explanation
"""


def get_inputs():
    """
    Create TF Placeholders for input, targets, and learning rate.
    :return: Tuple (input, targets, learning rate)
    """
    input = tf.placeholder(tf.int32, [None, None], name='input')
    targets = tf.placeholder(tf.int32, [None, None], name='targets')
    learning_rate = tf.placeholder(tf.float32, name='learning_rate')
    return (input, targets, learning_rate)

## supplied tests
tests.test_get_inputs(get_inputs)

"""
Explanation: Input
Create placeholder for inputs:
- Input text placeholder named "input" using the TF Placeholder name parameter.
- Targets placeholder.
- Learning Rate placeholder.
Return the placeholders in the following the tuple (Input, Targets, LearingRate)
End of explanation
"""


def get_init_cell(batch_size, rnn_size):
    """
    Create an RNN Cell and initialize it.
    :param batch_size: Size of batches
    :param rnn_size: Size of RNNs
    :return: Tuple (cell, initialize state)
    """
    num_layers = 2

    # this doesn't seem to work
    #lstm_layer = tf.contrib.rnn.BasicLSTMCell(rnn_size)
    #cell = tf.contrib.rnn.MultiRNNCell([lstm_layer] * num_layers)
    
    # use staked version instead:
    stacked_rnn = []
    for i in range(num_layers):
        stacked_rnn.append(tf.nn.rnn_cell.LSTMCell(num_units=512, state_is_tuple=True))

    cell = tf.nn.rnn_cell.MultiRNNCell(cells=stacked_rnn, state_is_tuple=True)
    
    # yet another version
    #cell = tf.contrib.rnn.MultiRNNCell([lstm_layer, lstm_layer])
    
    initial_state = cell.zero_state(batch_size, tf.float32)
    initial_state = tf.identity(initial_state, name='initial_state')
    return (cell, initial_state)


tests.test_get_init_cell(get_init_cell)

"""
Explanation: Build RNN Cell and Initialize
Stack one or more BasicLSTMCells in a MultiRNNCell
The Rnn size should be set using rnn_size
- Initalize Cell State using the MultiRNNCell's zero_state() function
- Apply the name "initial_state" to the initial state using tf.identity()
- Return the cell and initial state in the following tuple (Cell, InitialState)
End of explanation
"""



def get_embed(input_data, vocab_size, embed_dim):
    """
    Create embedding for <input_data>.
    :param input_data: TF placeholder for text input.
    :param vocab_size: Number of words in vocabulary.
    :param embed_dim: Number of embedding dimensions
    :return: Embedded input.
    """
    #embedding = tf.Variable(tf.random_uniform((vocab_size, embed_dim), -1, 1))
    #embed = tf.nn.embedding_lookup(embedding, input_data)
    
    embed = tf.contrib.layers.embed_sequence(input_data, 
                                            vocab_size=vocab_size, 
                                            embed_dim=embed_dim)
    return embed


## supplied test
tests.test_get_embed(get_embed)

"""
Explanation: Apply word embedding to input
End of explanation
"""


def build_rnn(cell, inputs):
    """
    Create a RNN using a RNN Cell
    :param cell: RNN Cell
    :param inputs: Input text data
    :return: Tuple (Outputs, Final State)
    """
    outputs, final_state = tf.nn.dynamic_rnn(cell, inputs, dtype=tf.float32)
    final_state = tf.identity(final_state, name="final_state")
    return outputs, final_state


tests.test_build_rnn(build_rnn)

"""
Explanation: Build RNN
You created a RNN Cell in the get_init_cell() function. Time to use the cell to create a RNN.
- Build the RNN using the tf.nn.dynamic_rnn()
- Apply the name "final_state" to the final state using tf.identity()
- Return the outputs and final_state state in the following tuple (Outputs, FinalState)
End of explanation
"""


def build_nn(cell, rnn_size, input_data, vocab_size, embed_dim):
    """
    Build part of the neural network
    :param cell: RNN cell
    :param rnn_size: Size of rnns
    :param input_data: Input data
    :param vocab_size: Vocabulary size
    :return: Tuple (Logits, FinalState)
    """
    embed = get_embed(input_data, vocab_size, embed_dim=embed_dim)
    outputs, final_state = build_rnn(cell, embed)
    #logits = tf.contrib.layers.fully_connected(outputs, vocab_size)
    logits = tf.contrib.layers.fully_connected(outputs, vocab_size, activation_fn=None)
    return logits, final_state


tests.test_build_nn(build_nn)

def test_get_batches(get_batches):
    with tf.Graph().as_default():
        test_batch_size = 128
        test_seq_length = 5
        test_int_text = list(range(1000*test_seq_length))
        batches = get_batches(test_int_text, test_batch_size, test_seq_length)

        # Check type
        assert isinstance(batches, np.ndarray),\
            'Batches is not a Numpy array'

        # Check shape
        assert batches.shape == (7, 2, 128, 5),\
            'Batches returned wrong shape.  Found {}'.format(batches.shape)

    tests._print_success_message()

"""
Explanation: Build the Neural Network
Apply the functions you implemented above to:
- Apply embedding to input_data using your get_embed(input_data, vocab_size, embed_dim) function.
- Build RNN using cell and your build_rnn(cell, inputs) function.
- Apply a fully connected layer with a linear activation and vocab_size as the number of outputs.
Return the logits and final state in the following tuple (Logits, FinalState)
End of explanation
"""


def get_batches(int_text, batch_size, seq_length):
    """
    Return batches of input and target
    :param int_text: Text with the words replaced by their ids
    :param batch_size: The size of batch
    :param seq_length: The length of sequence
    :return: Batches as a Numpy array
    """
    # TODO: Implement Function
    int_text = int_text[16:]
    n_batches = (len(int_text)-1)//(batch_size * seq_length)    
    int_text = int_text[:n_batches * batch_size * seq_length + 1]
    int_text_sequences = [int_text[i*seq_length:i*seq_length+seq_length] for i in range(0, n_batches * batch_size)]
    int_text = int_text[1:]
    int_text_targets = [int_text[i*seq_length:i*seq_length+seq_length] for i in range(0, n_batches * batch_size)]
    output = []
    for batch in range(n_batches):
        inputs = []
        targets = []
        for size in range(batch_size):
            inputs.append(int_text_sequences[size * n_batches + batch])
            targets.append(int_text_targets[size * n_batches + batch])
        output.append([inputs, targets])
    return np.array(output)

## check out why this fails
#tests.test_get_batches(get_batches)
test_get_batches(get_batches)

"""
Explanation: Batches
Implement get_batches to create batches of input and targets using int_text. The batches should be a Numpy array with the shape (number of batches, 2, batch size, sequence length). Each batch contains two elements:
- The first element is a single batch of input with the shape [batch size, sequence length]
- The second element is a single batch of targets with the shape [batch size, sequence length]
If you can't fill the last batch with enough data, drop the last batch.
For exmple, get_batches([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], 2, 3) would return a Numpy array of the following:
[
  # First Batch <br>
  [
    # Batch of Input
    [[ 1  2  3], [ 7  8  9]],
    # Batch of targets
    [[ 2  3  4], [ 8  9 10]]
  ],
# Second Batch
  [
    # Batch of Input
    [[ 4  5  6], [10 11 12]],
    # Batch of targets
    [[ 5  6  7], [11 12 13]]
  ]
]
End of explanation
"""


num_epochs = 120
batch_size = 101
rnn_size = 256
seq_length = 20
learning_rate = 0.01
embed_dim = 300

# Show stats for every n number of batches
show_every_n_batches = 100
#show_every_n_batches= get_batches(int_text, batch_size, seq_length).shape[0]
print(show_every_n_batches)

save_dir = './save'

## determine batch size (from someone else)

total_words = len(int_text[16:])
print('total words:', total_words)

print("Batch Size -> words missed\n")

for i in range(-5,5):
    try_batch_size = batch_size + i
    batches = get_batches(int_text, try_batch_size, seq_length)
    flag = ""
    if i == 0:
        flag = "<<< Current choice. Should be minimum words missed."
    print("{:>10}   {:>5} {}".format(try_batch_size, total_words - try_batch_size * seq_length * batches.shape[0], flag))

"""
Explanation: NN Training
Tune hyperparameters:
- Set num_epochs to the number of epochs.
- Set batch_size to the batch size.
- Set rnn_size to the size of the RNNs.
- Set seq_length to the length of sequence.
- Set learning_rate to the learning rate.
- Set show_every_n_batches to the number of batches the neural network should print progress.
End of explanation
"""


from tensorflow.contrib import seq2seq

train_graph = tf.Graph()
with train_graph.as_default():
    vocab_size = len(int_to_vocab)
    input_text, targets, lr = get_inputs()
    input_data_shape = tf.shape(input_text)
    cell, initial_state = get_init_cell(input_data_shape[0], rnn_size)
    logits, final_state = build_nn(cell, rnn_size, input_text, vocab_size, embed_dim)
    
    # Probabilities for generating words
    probs = tf.nn.softmax(logits, name='probs')

    # Loss function
    cost = seq2seq.sequence_loss(
        logits,
        targets,
        tf.ones([input_data_shape[0], input_data_shape[1]]))

    # Optimizer
    optimizer = tf.train.AdamOptimizer(lr)

    # Gradient Clipping
    gradients = optimizer.compute_gradients(cost)
    capped_gradients = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in gradients]
    train_op = optimizer.apply_gradients(capped_gradients)
    

"""
Explanation: Build graph
End of explanation
"""


batches = get_batches(int_text, batch_size, seq_length)

with tf.Session(graph=train_graph) as sess:
    sess.run(tf.global_variables_initializer())

    for epoch_i in range(num_epochs):
        state = sess.run(initial_state, {input_text: batches[0][0]})

        for batch_i, (x, y) in enumerate(batches):
            feed = {
                input_text: x,
                targets: y,
                initial_state: state,
                lr: learning_rate}
            train_loss, state, _ = sess.run([cost, final_state, train_op], feed)

            # Show every <show_every_n_batches> batches
            if (epoch_i * len(batches) + batch_i) % show_every_n_batches == 0:
                print('Epoch {:>3} Batch {:>4}/{}   train_loss = {:.3f}'.format(
                    epoch_i,
                    batch_i,
                    len(batches),
                    train_loss))

    # Save Model
    saver = tf.train.Saver()
    saver.save(sess, save_dir)
    print('Model Trained and Saved')

"""
Explanation: Train
End of explanation
"""


# Save parameters for checkpoint
params = (seq_length, save_dir)
pickle.dump(params, open('params.p', 'wb'))

"""
Explanation: Save Parameters
Save seq_length and save_dir for generating a new TV script.
End of explanation
"""


import tensorflow as tf
import numpy as np
import problem_unittests as tests

_, vocab_to_int, int_to_vocab, token_dict = pickle.load(open('preprocess.p', mode='rb'))
seq_length, load_dir = pickle.load(open('params.p', mode='rb'))

"""
Explanation: Check point
End of explanation
"""


def get_tensors(loaded_graph):
    """
    Get input, initial state, final state, and probabilities tensor from <loaded_graph>
    :param loaded_graph: TensorFlow graph loaded from file
    :return: Tuple (InputTensor, InitialStateTensor, FinalStateTensor, ProbsTensor)
    """
    InputTensor = loaded_graph.get_tensor_by_name("input:0")
    InitialStateTensor = loaded_graph.get_tensor_by_name("initial_state:0") 
    FinalStateTensor = loaded_graph.get_tensor_by_name("final_state:0") 
    ProbsTensor = loaded_graph.get_tensor_by_name("probs:0")
    return  (InputTensor, InitialStateTensor, FinalStateTensor, ProbsTensor)


tests.test_get_tensors(get_tensors)

"""
Explanation: Implement Generate Functions
Get tensors from loaded_graph using the function get_tensor_by_name(). Get the tensors using the following names:
- "input:0"
- "initial_state:0"
- "final_state:0"
- "probs:0"
Return the tensors in the following tuple (InputTensor, InitialStateTensor, FinalStateTensor, ProbsTensor)
End of explanation
"""


import random

def weighted_choice(choices):
    # From http://stackoverflow.com/questions/3679694/a-weighted-version-of-random-choice
    total = sum(w for c, w in choices)
    r = random.uniform(0, total)
    upto = 0
    for c, w in choices:
        if upto + w >= r:
            return c
        upto += w
    assert False, "Shouldn't get here"
    
def pick_from_top_5(choices):
    top5 = []
    for i in range(min(len(choices), 5)):
        index = np.argmax(choices)
        top5.append((index, choices[index]))
        choices.itemset(index, 0) # Avoid picking this index as argmax again
    return weighted_choice(top5)

def pick_word(probabilities, int_to_vocab):
    """
    Pick the next word in the generated text
    :param probabilities: Probabilites of the next word
    :param int_to_vocab: Dictionary of word ids as the keys and words as the values
    :return: String of the predicted word
    """
    return int_to_vocab[pick_from_top_5(probabilities)]


tests.test_pick_word(pick_word)

"""
Explanation: Choose Word
Select the next word using probabilities.
End of explanation
"""


gen_length = 100

# homer_simpson, moe_szyslak, or Barney_Gumble
prime_word = 'homer_simpson'

loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    # Load saved model
    loader = tf.train.import_meta_graph(load_dir + '.meta')
    loader.restore(sess, load_dir)

    # Get Tensors from loaded model
    input_text, initial_state, final_state, probs = get_tensors(loaded_graph)

    # Sentences generation setup
    gen_sentences = [prime_word + ':']
    prev_state = sess.run(initial_state, {input_text: np.array([[1]])})

    # Generate sentences
    for n in range(gen_length):
        # Dynamic Input
        dyn_input = [[vocab_to_int[word] for word in gen_sentences[-seq_length:]]]
        dyn_seq_length = len(dyn_input[0])

        # Get Prediction
        probabilities, prev_state = sess.run(
            [probs, final_state],
            {input_text: dyn_input, initial_state: prev_state})
        
        probabilities = probabilities[0]
        
        pred_word = pick_word(probabilities[dyn_seq_length-1], int_to_vocab)

        gen_sentences.append(pred_word)
    
    # Remove tokens
    tv_script = ' '.join(gen_sentences)
    for key, token in token_dict.items():
        ending = ' ' if key in ['\n', '(', '"'] else ''
        tv_script = tv_script.replace(' ' + token.lower(), key)
    tv_script = tv_script.replace('\n ', '\n')
    tv_script = tv_script.replace('( ', '(')
        
    print(tv_script)

"""
Explanation: Generate TV Script
Set gen_length to the length of TV script you want to generate.
End of explanation
"""


