import constants
import chatbot.conversation.models as nets
import chatbot.word2vec.word2vec_containers as word2vec
from chatbot.conversation.conversion.converters import ConversationConverter
import numpy as np
import random
from time import time
from os.path import getmtime


class ConversationTrainer:
    MINUTES_BETWEEN_MODEL_UPDATES = 60

    def __init__(self):
        self._conversation_network = nets.ConversationNetwork()
        self._word2vec = word2vec.Word2VecContainer()
        self._converter = ConversationConverter()

    def train(self):
        def update_if_outdated_callback(*_):
            model_modified_time = getmtime(constants.CONVERSATION_MODEL_PATH)
            current_time = time()
            difference_in_seconds = current_time - model_modified_time
            difference_in_minutes = difference_in_seconds / 60
            if difference_in_minutes >= ConversationTrainer.MINUTES_BETWEEN_MODEL_UPDATES:
                self._conversation_network.model().save_weights(constants.CONVERSATION_MODEL_PATH)

        def training_set_generator():
            conversation_files = constants.DIALOG_TRAINING_FILES_PATHS[:]
            random.shuffle(conversation_files)
            dialogs_x = []
            dialogs_y = []
            while True:
                random.shuffle(conversation_files)
                for file in conversation_files:
                    for x, y in self._train_on_file(file):
                        dialogs_x.append(x)
                        dialogs_y.append(y)
                        if len(dialogs_x) == constants.CONVERSATIONS_MINI_BATCH_SIZE:
                            yield np.concatenate(dialogs_x), np.concatenate(dialogs_y)
                            dialogs_x = []
                            dialogs_y = []

        def validation_set_generator():
            while True:
                yield self._generate_random_validation_set()

        self._conversation_network.model().fit_generator(training_set_generator(),
                                                         validation_data=validation_set_generator(),
                                                         validation_steps=constants.CONVERSATION_VALIDATION_PART_SIZE,
                                                         verbose=True,
                                                         max_queue_size=constants.CONVERSATION_GENERATOR_QUEUE_SIZE,
                                                         steps_per_epoch=constants.CONVERSATION_NUMBER_OF_LEARNING_STEPS,
                                                         epochs=constants.CONVERSATION_NUMBER_OF_EPOCHS)

    def _train_on_file(self, file):
        conversations = list(ConversationDataLoader.generate_conversations(file))
        random.shuffle(conversations)
        while len(conversations) > 0:
            dialogs_count = 0
            conversations_part_inputs = []
            conversations_part_outputs = []
            while dialogs_count < constants.CONVERSATIONS_MINI_BATCH_SIZE and len(conversations) > 0:
                conversation = conversations.pop()
                conversation_data = self._converter.convert_conversation_into_training_data(conversation)
                if conversation_data is not None:
                    x, y = conversation_data
                    conversations_part_inputs.append(x)
                    conversations_part_outputs.append(y)
                    dialogs_count += int(round(len(conversation) / 2))
            if len(conversations_part_inputs) > 0 and len(conversations_part_outputs) > 0:
                training_x = np.concatenate(conversations_part_inputs)
                training_y = np.concatenate(conversations_part_outputs)
                if len(training_x.shape) > 1:
                    yield training_x, training_y
        raise StopIteration

    def _generate_random_validation_set(self):
        file = random.choice(constants.DIALOG_VALIDATION_FILES_PATHS)
        conversations = list(ConversationDataLoader.generate_conversations(file))
        dialogs_count = 0
        conversations_part_inputs = []
        conversations_part_outputs = []

        while dialogs_count < constants.CONVERSATIONS_MINI_BATCH_SIZE and len(conversations) > 0:
            conversation = conversations.pop()
            conversation_data = self._converter.convert_conversation_into_training_data(conversation)
            if conversation_data is not None:
                x, y = conversation_data
                conversations_part_inputs.append(x)
                conversations_part_outputs.append(y)
                dialogs_count += int(round(len(conversation) / 2))
            if len(conversations) == 0 and dialogs_count < constants.CONVERSATIONS_MINI_BATCH_SIZE:
                file = random.choice(constants.DIALOG_VALIDATION_FILES_PATHS)
                conversations = list(ConversationDataLoader.generate_conversations(file))

        validation_x = np.concatenate(conversations_part_inputs)
        validation_y = np.concatenate(conversations_part_outputs)

        return validation_x, validation_y


class ConversationDataLoader:
    @staticmethod
    def generate_conversations(file_path):
        with open(file_path, encoding=constants.ENCODING, errors="ignore") as training_data:
            outputs = []
            output = []
            for line in training_data:
                if len(line.strip()) > 0:
                    if line.strip() == constants.NEW_FILE_SYMBOL:
                        output.clear()
                    elif line.strip() == constants.END_FILE_SYMBOL and len(output) > 0:
                        outputs.append(output[:])
                    else:
                        output.append(line.strip().split(" "))
            return outputs


if __name__ == "__main__":
    trainer = ConversationTrainer()
    trainer.train()
