import sys
import pickle
import matplotlib.pyplot as plt
from keras.optimizers import RMSprop
from keras.layers.core import Dense, Activation, Dropout, RepeatVector
from keras.layers import TimeDistributed
from keras.layers import LSTM, Dropout
from keras.layers import Dense, Activation
from keras.models import Sequential, load_model
import tensorflow as tf
import numpy as np
import os
import heapq

np.random.seed(42)
tf.set_random_seed(42)

char_indices = {}
indices_char = {}
sentences = []
next_chars = []
X = None
y = None
SEQUENCE_LENGTH = 3
chars = None


def preprocessing(text):
    global sentences
    global chars
    global char_indices
    global indices_char

    chars = sorted(list(set(text)))
    char_indices = dict((c, i) for i, c in enumerate(chars))
    indices_char = dict((i, c) for i, c in enumerate(chars))
    print(f'unique chars: {len(chars)}')
    print(char_indices)
    step = 1
    for i in range(0, len(text) - SEQUENCE_LENGTH, step):
        sentences.append(text[i: i + SEQUENCE_LENGTH])
        next_chars.append(text[i + SEQUENCE_LENGTH])
    print(f'num training examples: {len(sentences)}')


def feature_build():
    global X
    global y

    X = np.zeros((len(sentences), SEQUENCE_LENGTH, 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
    print(X[0][0])
    # return X, y


def build():
    # build model
    model = Sequential()
    model.add(LSTM(128, input_shape=(SEQUENCE_LENGTH, len(chars))))
    model.add(Dense(len(chars)))
    model.add(Activation('softmax'))
    optimizer = RMSprop(lr=0.01)
    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizer, metrics=['accuracy'])
    history = model.fit(X, y, validation_split=0.05,
                        batch_size=128, epochs=20, shuffle=True).history
    print(history)
    if os.path.exists("model"):
        model.save('model/keras_model.h5')
        pickle.dump(history, open("model/history.p", "wb"))


def load():
    model = load_model('model/keras_model.h5')
    # history = pickle.load(open("model/history.p", "rb"))
    # return model, history
    return model


def prepare_input(text, chars):
    x = np.zeros((1, SEQUENCE_LENGTH, len(chars)))
    for t, char in enumerate(text):
        x[0, t, char_indices[char]] = 1
    return x


def sample(preds, top_n=5):
    preds = np.asarray(preds).astype('float64')
    preds = np.log(preds)
    exp_preds = np.exp(preds)
    preds = exp_preds / np.sum(exp_preds)

    return heapq.nlargest(top_n, range(len(preds)), preds.take)


def predict_completion(model, text, chars):
    original_text = text
    generated = text
    completion = ''
    while True:
        x = prepare_input(text, chars)
        preds = model.predict(x, verbose=0)[0]
        next_index = sample(preds, top_n=1)[0]
        next_char = indices_char[next_index]
        text = text[1:] + next_char
        completion += next_char

        if len(original_text + completion) + 2 > len(original_text) and next_char == ' ':
            return completion


def predict_completions(model, text, chars, n=5):
    x = prepare_input(text, chars)
    preds = model.predict(x, verbose=0)[0]
    next_indices = sample(preds, n)
    return [indices_char[idx] + predict_completion(model, text[1:] + indices_char[idx], chars) for idx in next_indices]
