'''Example script to generate text from Nietzsche's writings.

At least 20 epochs are required before the generated text
starts sounding coherent.

It is recommended to run this script on GPU, as recurrent
networks are quite computationally intensive.

If you try this script on new data, make sure your corpus
has at least ~100k characters. ~1M is better.
'''

from keras.models import Sequential, load_model
from keras.layers import Dense, Activation
from keras.layers import LSTM, SimpleRNN, Dropout, Bidirectional, Embedding, GRU
from keras.callbacks import ModelCheckpoint
from keras.optimizers import RMSprop
from keras.utils.data_utils import get_file
import numpy as np
import random
import sys
import pprint

embeddings_path = "glove.840B.300d-char.txt"
embedding_dim = 300


text_file_kapital = open("Das_Kapital.txt", 'rb')

def preprocess(text_file):

    ##get rid of line breaks and non-ASCII
    lines = []
    for line in text_file:
        line = line.strip().lower()
        line = line.decode("ascii", "ignore")
        if len(line) == 0:
            continue
        lines.append(line)
    text_file.close()
    text = " ".join(lines)
    return text

text = preprocess(text_file_kapital)


chars = sorted(list(set(text)))
print('total chars:', len(chars))
char_indices = dict((c, i) for i, c in enumerate(chars))
indices_char = dict((i, c) for i, c in enumerate(chars))

# cut the text in semi-redundant sequences of maxlen characters
maxlen = 40
step = 10
sentences = []
next_chars = []
for i in range(0, len(text) - maxlen, step):
    sentences.append(text[i: i + maxlen])
    next_chars.append(text[i + maxlen])
print('nb sequences:', len(sentences))

print('Vectorization...')

#working
# X = np.zeros((len(sentences), maxlen, len(chars)), dtype=np.bool)
# y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
# for i, sentence in enumerate(sentences):
#     for t, char in enumerate(sentence):
#         X[i, t, char_indices[char]] = 1
#     y[i, char_indices[next_chars[i]]] = 1

X = np.zeros((len(sentences), maxlen), dtype=np.int)
y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
for i, sentence in enumerate(sentences):
    for t, char in enumerate(sentence):
        X[i, t] = char_indices[char]
y[i, char_indices[next_chars[i]]] = 1

print('Processing pretrained character embeds...')
embedding_vectors = {}
with open(embeddings_path, 'r') as f:
    for line in f:
        line_split = line.strip().split(" ")
        vec = np.array(line_split[1:], dtype=float)
        char = line_split[0]
        embedding_vectors[char] = vec

embedding_matrix = np.zeros((len(chars), 300))
#embedding_matrix = np.random.uniform(-1, 1, (len(chars), 300))
for char, i in char_indices.items():
    #print ("{}, {}".format(char, i))
    embedding_vector = embedding_vectors.get(char)
    if embedding_vector is not None:
        embedding_matrix[i] = embedding_vector



# build the model: a bidirectional GRU

def build_model():

    ## working
    # print('Build model...')
    # model = Sequential()
    # #model.add(Embedding(len(chars), embedding_dim, input_length=maxlen, weights=[embedding_matrix]))
    # model.add(Bidirectional(GRU(128, unroll=True, return_sequences=True), input_shape=(maxlen, len(chars))))
    # model.add(Dropout(0.2))
    # model.add(Bidirectional(GRU(128, unroll=True)))
    # model.add(Dropout(0.2))
    # model.add(Dense(len(chars)))
    # model.add(Activation('softmax'))
    #
    # model.compile(loss='categorical_crossentropy', optimizer="RMSprop")
    # return model

    print('Build model...')
    model = Sequential()
    model.add(Embedding(len(chars), embedding_dim, input_length=maxlen, weights=[embedding_matrix]))
    model.add(Bidirectional(GRU(32, unroll=True), input_shape=(maxlen, len(chars))))
    model.add(Dropout(0.2))
    # model.add(Bidirectional(GRU(128, unroll=True)))
    # model.add(Dropout(0.2))
    model.add(Dense(len(chars)))
    model.add(Activation('softmax'))

    model.compile(loss='categorical_crossentropy', optimizer="RMSprop")
    return model


#model = load_model("models/weights-improvement-00-1.2807-biggggger.hdf5")

model = build_model()
model.summary()

def sample(preds, temperature=1.0):
    # helper function to sample an index from a probability array
    preds = np.asarray(preds).astype('float64')
    preds = np.log(preds + 1e-6) / temperature
    exp_preds = np.exp(preds)
    preds = exp_preds / np.sum(exp_preds)
    probas = np.random.multinomial(1, preds, 1)
    return np.argmax(probas)

# train the model, output generated text after each iteration

def train():
    for iteration in range(1, 60):
        filepath = "models/weights-improvement-{epoch:02d}-{loss:.4f}-embeddings.hdf5"
        checkpoint = ModelCheckpoint(filepath, monitor='loss', verbose=1, save_best_only=True, mode='min')
        callbacks_list = [checkpoint]
        print()
        print('-' * 50)
        print('Iteration', iteration)
        model.fit(X, y,
                  batch_size=128,
                  epochs=1, callbacks = callbacks_list)

        start_index = random.randint(0, len(text) - maxlen - 1)
        #model.save("models/testmodel.h5")

        for diversity in [0.2, 0.5, 1.0, 1.2]:
            print()
            print('----- diversity:', diversity)

            generated = ''
            sentence = text[start_index: start_index + maxlen]
            generated += sentence
            print('----- Generating with seed: "' + sentence + '"')
            sys.stdout.write(generated)

            for i in range(400):
                # x = np.zeros((1, maxlen, len(chars)))
                # for t, char in enumerate(sentence):
                #     x[0, t, char_indices[char]] = 1.
                #
                # preds = model.predict(x, verbose=0)[0]
                # next_index = sample(preds, diversity)
                # next_char = indices_char[next_index]

                x = np.zeros((1, maxlen), dtype=np.int)
                for t, char in enumerate(sentence):
                    x[0, t] = char_indices[char]

                preds = model.predict(x, verbose=0)[0][0]
                next_index = sample(preds, diversity)
                next_char = indices_char[next_index]

                generated += next_char
                sentence = sentence[1:] + next_char

                sys.stdout.write(next_char)
                sys.stdout.flush()
            print()

def predict(num_to_predict, temperatrue, seed):
    start_index = random.randint(0, len(text) - maxlen - 1)

    if len(seed) > 40:
        print("Type fewer characters, you typed this man characters")
        print(len(seed))
        return 0
    newstring = ""
    space = 40 - len(seed)
    for i in range(space):
        newstring += " "
    seed = newstring + seed
    sentence = seed
    generated = ''
    #sentence = text[start_index: start_index + maxlen]
    generated += sentence
    print('----- Hey Marx, Tell me about:  "' + sentence.strip() + '"')
    sys.stdout.write(generated)

    for i in range(num_to_predict):
        x = np.zeros((1, maxlen, len(chars)))
        for t, char in enumerate(sentence):
            x[0, t, char_indices[char]] = 1.

        preds = model.predict(x, verbose=0)[0]
        next_index = sample(preds, temperatrue)
        next_char = indices_char[next_index]

        generated += next_char
        sentence = sentence[1:] + next_char

        #sys.stdout.write(next_char)
        #sys.stdout.flush()
    print()
    pprint.pprint(generated)

train()

# ################# PREDICT
# cont = 0
#
# while cont == 0:
#     newtext = input("What do you want me to tell you about?")
#
#     if newtext == "1":
#         cont = 1
#     predict(500, 0.5, newtext.lower())