from __future__ import print_function

from tensorflow.python.keras.models import Model, load_model
from tensorflow.python.keras.layers import Input, LSTM, Dense, Embedding
import numpy as np


from src.utils import one_hot_encode, add_n


class Seq2Seq:
    def __init__(self, num_encoder_tokens, num_decoder_tokens, start_token, end_token,
                 latent_dim=256, projection='one_hot', emb_dim=64,
                 restore_path=None):
        self.num_encoder_tokens = num_encoder_tokens
        self.num_decoder_tokens = num_decoder_tokens

        self.start_token = start_token
        self.end_token = end_token

        if restore_path:
            self.train_model = load_model(restore_path + '/train_model.s2s.h5')
            self.encoder_model = load_model(restore_path + '/encoder_model.s2s.h5')
            self.decoder_model = load_model(restore_path + '/decoder_model.s2s.h5')
            return

        # Define an input sequence and process it.
        encoder_inputs = Input(shape=(None,), name='encoder_input')
        if projection == 'one_hot':
            embedding = Embedding(
                num_encoder_tokens+1,
                num_encoder_tokens,
                weights=[np.concatenate((np.zeros(shape=(1, num_encoder_tokens)), np.identity(num_encoder_tokens)))],
                trainable=False,
                mask_zero=True,
                name='encoder_embedding'
            )
        elif projection == 'word2vec':
            embedding = Embedding(
                num_encoder_tokens + 1,
                emb_dim,
                mask_zero=True,
                name='encoder_embedding'
            )
        else:
             raise Exception("projection method not recognized")
        encoder_emb = embedding(encoder_inputs)
        encoder = LSTM(latent_dim, return_state=True, name='encoder_lstm')
        encoder_outputs, state_h, state_c = encoder(encoder_emb)
        # We discard `encoder_outputs` and only keep the states.
        encoder_states = [state_h, state_c]

        # Set up the decoder, using `encoder_states` as initial state.
        decoder_inputs = Input(shape=(None,), name='decoder_input')
        if projection == 'one_hot':
            embedding = Embedding(
                num_decoder_tokens+1,
                num_decoder_tokens,
                weights=[np.concatenate((np.zeros(shape=(1, num_decoder_tokens)), np.identity(num_decoder_tokens)))],
                trainable=False,
                mask_zero=True,
                name='decoder_embedding'
            )
        elif projection == 'word2vec':
            embedding = Embedding(
                num_decoder_tokens+1,
                emb_dim,
                mask_zero=True,
                name='decoder_embedding'
            )
        else:
            raise Exception("projection method not recognized")
        decoder_emb = embedding(decoder_inputs)
        # We set up our decoder to return full output sequences,
        # and to return internal states as well. We don't use the
        # return states in the training model, but we will use them in inference.
        decoder_lstm = LSTM(latent_dim, return_sequences=True, return_state=True, name='decoder_lstm')
        decoder_outputs, _, _ = decoder_lstm(decoder_emb,
                                             initial_state=encoder_states)
        decoder_dense = Dense(self.num_decoder_tokens, activation='softmax', name='decoder_softmax')
        decoder_outputs = decoder_dense(decoder_outputs)

        # Define the model that will turn
        # `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
        self.train_model = Model([encoder_inputs, decoder_inputs], decoder_outputs)

        # Define sampling models
        self.encoder_model = Model(encoder_inputs, encoder_states)

        decoder_state_input_h = Input(shape=(latent_dim,), name='decoder_state_input_h')
        decoder_state_input_c = Input(shape=(latent_dim,), name='decoder_state_input_c')
        decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
        decoder_outputs, state_h, state_c = decoder_lstm(
            decoder_emb, initial_state=decoder_states_inputs)
        decoder_states = [state_h, state_c]
        decoder_outputs = decoder_dense(decoder_outputs)
        self.decoder_model = Model(
            [decoder_inputs] + decoder_states_inputs,
            [decoder_outputs] + decoder_states)

    def pre_process_encoder_input(self, input):
        # return one_hot_encode(input, self.max_encoder_seq_length, self.num_encoder_tokens)
        return add_n(input, self.max_encoder_seq_length, 1)

    def pre_process_decoder_input(self, input):
        # return one_hot_encode(input, self.max_decoder_seq_length, self.num_decoder_tokens)
        return add_n(input, self.max_decoder_seq_length, 1)

    def pre_process_decoder_target(self, input):
        # return one_hot_encode(input, self.max_decoder_seq_length, self.num_decoder_tokens)
        return add_n(input, self.max_decoder_seq_length, 0)

    def train(self, encoder_input_seqs, decoder_input_seqs, decoder_target_seqs, batch_size=64, epochs=100):
        # Run training
        self.max_encoder_seq_length = max([len(seq) for seq in encoder_input_seqs])
        self.max_decoder_seq_length = max([len(seq) for seq in decoder_input_seqs])
        encoder_input_data = self.pre_process_encoder_input(encoder_input_seqs)
        print('encoder_input_data.shape: ', encoder_input_data.shape)
        decoder_input_data = self.pre_process_decoder_input(decoder_input_seqs)
        print('decoder_input_data.shape: ', decoder_input_data.shape)
        decoder_target_data = self.pre_process_decoder_target(decoder_target_seqs)
        print('decoder_target_data.shape: ', decoder_target_data.shape)
        self.train_model.compile(optimizer='rmsprop', loss='sparse_categorical_crossentropy')
        self.train_model.fit([encoder_input_data, decoder_input_data], decoder_target_data,
                  batch_size=batch_size,
                  epochs=epochs,
                  validation_split=0.2)
        # Save model
        self.train_model.save('train_model.s2s.h5')
        self.encoder_model.save('encoder_model.s2s.h5')
        self.decoder_model.save('decoder_model.s2s.h5')

    def encode(self, input_seq):
        # encoder_input_data = one_hot_encode(input_seq, self.max_encoder_seq_length, self.num_encoder_tokens)
        encoder_input_data = self.pre_process_encoder_input(input_seq)
        return self.encoder_model.predict(encoder_input_data)

    def decode(self, states_value):
        # Generate empty target sequence of length 1.
        # target_seq = np.zeros((1, 1, self.num_decoder_tokens))
        # Populate the first token of target sequence with the start token.
        # target_seq[0, 0, self.start_token] = 1.
        target_seq = np.zeros((1, 1))
        target_seq[0, 0] = self.start_token + 1
        # Sampling loop for a batch of sequences
        # (to simplify, here we assume a batch of size 1).
        stop_condition = False
        decoded_seq = []
        while not stop_condition:
            output_tokens, h, c = self.decoder_model.predict(
                [target_seq] + states_value)

            # Sample a token
            sampled_token_index = np.argmax(output_tokens[0, -1, :])
            decoded_seq += [sampled_token_index]

            # Exit condition: either hit max length
            # or find stop token.
            if (sampled_token_index == self.end_token or
                    len(decoded_seq) > self.max_decoder_seq_length):
                stop_condition = True

            # Update the target sequence (of length 1).
            # target_seq = np.zeros((1, 1, self.num_decoder_tokens))
            # target_seq[0, 0, sampled_token_index] = 1.
            target_seq = np.zeros((1, 1))
            target_seq[0, 0] = sampled_token_index + 1

            # Update states
            states_value = [h, c]

        return decoded_seq

    def predict(self, input_seq):
        # Encode the input as state vectors.
        states_value = self.encode(input_seq)
        prediction = self.decode(states_value)

        return prediction


