import warnings
from typing import List

import numpy as np
from keras import Sequential, Model
from keras.callbacks import Callback, History
from keras.layers import Embedding, LSTM, Dropout, Dense
from keras.utils import to_categorical
from keras_preprocessing.sequence import pad_sequences
from keras_preprocessing.text import Tokenizer

from glossolalia.tokens import PoemTokenizer

warnings.filterwarnings("ignore")
warnings.simplefilter(action='ignore', category=FutureWarning)


class LisSansTaMaman(object):
    """ A LSTM model adapted for french lyrical texts."""

    def __init__(self, nb_layers: int = 100,
                 dropout: float = 0.1, validation_split: float = 0.0,
                 tokenizer=PoemTokenizer(),
                 debug: bool = False):
        self.validation_split = validation_split
        self.dropout = dropout
        self.nb_layers = nb_layers
        self.tokenizer = tokenizer

        # Model state
        self.model: Model = None
        self.predictors = None
        self.label = None
        self.max_sequence_len = None

        if debug:
            debug_unrandomize()

    def create_model(self, corpus: List[str]):
        inp_sequences, total_words = self.tokenizer.get_sequence_of_tokens(corpus)

        self.predictors, self.label, self.max_sequence_len = generate_padded_sequences(inp_sequences, total_words)
        model = create_model(self.max_sequence_len, total_words, layers=self.nb_layers, dropout=self.dropout)
        model.summary()

        self.model = model

    # TODO: Batch fit? splitting nb_epoch into N step
    def fit(self, epochs: int, initial_epoch: int = 0,
            callbacks: List[Callback] = None,
            validation_split: float = 0
            ) -> History:
        return self.model.fit(self.predictors, self.label,
                              verbose=2,
                              callbacks=callbacks,
                              validation_split=validation_split,
                              epochs=epochs, initial_epoch=initial_epoch)

    def predict_seeds(self, seeds: List[str] = None, nb_words=None):
        if seeds is None:
            seeds = ["", "Je", "Tu", "Le", "La", "Les", "Un", "On", "Nous"]
        return [self.predict(seed, nb_words) for seed in seeds]

    def predict(self, seed="", nb_words=None):
        if nb_words is None:
            nb_words = 20  # TODO: Guess based on model a good number of words
        return generate_text(self.model, self.tokenizer, seed, nb_words, self.max_sequence_len)


def generate_padded_sequences(input_sequences, total_words: int):
    max_sequence_len = max([len(x) for x in input_sequences])
    input_sequences = np.array(pad_sequences(input_sequences, maxlen=max_sequence_len, padding='pre'))
    predictors, label = input_sequences[:, :-1], input_sequences[:, -1]
    label = to_categorical(label, num_classes=total_words)
    return predictors, label, max_sequence_len


def create_model(max_sequence_len: int, total_words: int, layers: int, dropout: float):
    print("Creating model across %i words for %i-long seqs (%i layers, %.2f dropout):" %
          (total_words, max_sequence_len, layers, dropout))
    input_len = max_sequence_len - 1
    model = Sequential()

    # Add Input Embedding Layer
    model.add(Embedding(total_words, 10, input_length=input_len))

    # Add Hidden Layer 1 - LSTM Layer
    model.add(LSTM(layers))
    # model.add(Bidirectional(LSTM(layers), input_shape=(max_sequence_len, total_words)))
    model.add(Dropout(dropout))

    # Add Output Layer
    model.add(Dense(total_words, activation='softmax'))

    model.compile(optimizer='adam',  # TODO: Try RMSprop(learning_rate=0.01)
                  loss='categorical_crossentropy',  # TODO: Try sparse_categorical_crossentropy for faster training
                  metrics=['accuracy'])

    # TODO: Try alternative architectures
    #  https://medium.com/coinmonks/word-level-lstm-text-generator-creating-automatic-song-lyrics-with-neural-networks-b8a1617104fb#35f4
    return model


def generate_text(model: Sequential, tokenizer: Tokenizer, seed_text="", nb_words=5, max_sequence_len=0) -> str:
    word_indices = {v: k for k, v in tokenizer.word_index.items()}
    output = seed_text

    for _ in range(nb_words):
        token_list = tokenizer.texts_to_sequences([output])[0]
        token_list = pad_sequences([token_list], maxlen=max_sequence_len - 1, padding='pre')
        predicted = model.predict_classes(token_list, verbose=2)[0]
        output += " " + word_indices[predicted]
    return output.capitalize()


def debug_unrandomize():
    from numpy.random import seed
    from tensorflow_core.python.framework.random_seed import set_random_seed

    # set seeds for reproducibility
    set_random_seed(2)
    seed(1)
