
import numpy as np
import tensorflow as tf
from tensorflow.keras import Model, Input
import json
from data_processor import process_input, one_hot2char


def generate_text(seed_string):

    # set tensorflow to work with float64
    tf.keras.backend.set_floatx('float64')

    model_folder = 'deploy/model'

    # loads model configuration from JSON file
    with open(model_folder + '/config.json') as json_file:
        config = json.load(json_file)
        MAX_LEN_TITLE = config['MAX_LEN_TITLE']
        MAX_LEN_TEXT = config['MAX_LEN_TEXT']
        ENCODER_DEPTH = config['ENCODER_DEPTH']
        DECODER_DEPTH = config['DECODER_DEPTH']

    # loads vocab from JSON file
    with open(model_folder + '/vocab.json') as json_file:
        vocab = json.load(json_file)

    tf.keras.backend.clear_session()  # destroys the current graph

    # loads trained model from H5 file. 't' stands for 'trained'
    t_model = tf.keras.models.load_model(model_folder + '/trained_model.h5')

    # encoder model. 't' stands for 'trained'
    t_enc_input = t_model.get_layer('encoder_input').input
    t_enc_internal_tensor = t_enc_input

    for i in range(ENCODER_DEPTH):
        t_enc_LSTM = t_model.get_layer('encoder_LSTM_' + str(i))
        [t_enc_internal_tensor, t_enc_memory_state,
            t_enc_carry_state] = t_enc_LSTM(t_enc_internal_tensor)

    t_enc_model = Model(inputs=t_enc_input,
                        outputs=[t_enc_memory_state, t_enc_carry_state])

    # decoder inputs
    # the model needs to have an initial input, and inputs for every state (both memory and carry states) for each LSTM layer in the decoder,
    # because it decodes one time step at a time, always using the previous prediction as input
    t_dec_input = Input(
        shape=(None, config['DIM_VOCAB']), name='decoder_input')
    t_dec_internal_tensor = t_dec_input

    # list of all input memory state tensors (1 per LSTM layer)
    t_dec_memory_state_inputs = []
    # list of all input carry state tensors (1 per LSTM layer)
    t_dec_carry_state_inputs = []

    for i in range(DECODER_DEPTH):
        t_dec_memory_state_inputs.append(Input(
            shape=(config['DIM_LSTM_LAYER']), name='decoder_memory_state_input_' + str(i)))
        t_dec_carry_state_inputs.append(Input(
            shape=(config['DIM_LSTM_LAYER']), name='decoder_carry_state_input_' + str(i)))

    # list of all output memory state tensors (1 per LSTM layer)
    t_dec_memory_state_outputs = []
    # list of all output carry state tensors (1 per LSTM layer)
    t_dec_carry_state_outputs = []

    # decoder model
    for i in range(config['DECODER_DEPTH']):
        t_dec_LSTM = t_model.get_layer('decoder_LSTM_' + str(i))
        t_dec_internal_tensor, t_dec_memory_state_output, t_dec_carry_state_output = t_dec_LSTM(
            t_dec_internal_tensor, initial_state=[t_dec_memory_state_inputs[i], t_dec_carry_state_inputs[i]])
        t_dec_memory_state_outputs.append(t_dec_memory_state_output)
        t_dec_carry_state_outputs.append(t_dec_carry_state_output)

    t_dec_output = t_dec_internal_tensor  # output of last LSTM layer

    t_dense_output = t_model.get_layer('output')
    t_dec_prediction = t_dense_output(t_dec_output)
    t_dec_model = Model(inputs=[t_dec_input] + t_dec_memory_state_inputs + t_dec_carry_state_inputs,
                        outputs=[t_dec_prediction] + t_dec_memory_state_outputs + t_dec_carry_state_outputs)

    char2int = dict((c, i) for i, c in enumerate(vocab))

    seed_string = [seed_string]  # must be a list
    vec_input_strings = process_input(seed_string, vocab)

    vec_sentences = np.zeros((0, MAX_LEN_TEXT, len(vocab)))
    for vec_input_string in vec_input_strings:

        vec_input_string = np.expand_dims(vec_input_string, axis=0)
        [enc_memory_state, enc_carry_state] = t_enc_model.predict(
            [vec_input_string])  # get initial states generated by the encoder

        # creates initial character for the decoder
        one_hot_char_output = np.zeros((len(vocab)), dtype=np.int8)
        one_hot_char_output[char2int['[START]']] = 1
        one_hot_char_output = np.expand_dims(one_hot_char_output, axis=0)
        one_hot_char_output = np.expand_dims(one_hot_char_output, axis=0)

        # list of all input memory states (1 per LSTM layer)
        memory_states = []
        carry_states = []   # list of all input carry states (1 per LSTM layer)

        # initializes all memory and carry states with states from last layer of the encoder
        for i in range(config['DECODER_DEPTH']):
            memory_states.append(enc_memory_state)
            carry_states.append(enc_carry_state)

        vec_sentence = np.empty((0, len(vocab)))

        for _ in range(MAX_LEN_TEXT):

            # output contains the output and each layer's memory and carry states
            outputs = t_dec_model.predict(
                [one_hot_char_output] + memory_states + carry_states)
            one_hot_char_output = outputs.pop(0)  # separates the output
            # separates memory states from all layers
            memory_states = outputs[: len(outputs)//2]
            # separates carry states from all layers
            carry_states = outputs[len(outputs)//2:]
            vec_sentence = np.append(
                vec_sentence, one_hot_char_output[0], axis=0)

        vec_sentences = np.append(vec_sentences, [vec_sentence], axis=0)

    string = one_hot2char(vec_sentences, vocab)

    return string[0]
