"""
This modules provides functionality to generate natural language in response
to a seed (the so-called question).
"""

import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import Tuple, List
import sys

transformers.logging.set_verbosity_error()


class ChatBot:
    """
    A chatbot loads a language model and tokenizer and is able to generate
    natural language based on a seed string.
    """

    def __init__(self, politician: str, model_path="anonymous-german-nlp/german-gpt2"):
        self.politician = politician
        self.model_path = model_path
        self.tokenizer = AutoTokenizer.from_pretrained("anonymous-german-nlp/german-gpt2")
        self.model = AutoModelForCausalLM.from_pretrained(model_path)

    @staticmethod
    def _preprocess_question(question: str) -> str:
        """
        Processes a question string by appending an EOQ token.

        return: The preprocessed question
        """
        question += "<EOQ>"

        return question

    @staticmethod
    def _process_answer(answer: str) -> str:
        """
        Processes an answer generated by the language model.
        Cuts off irrelevant and incomplete parts in the beginning and end.

        @return: Processed answer string.
        """
        special_tokens = ["<EOQ>", " <EOQ>", '"', '„']

        for token in special_tokens:
            while token in answer:
                answer = answer.replace(token, "")

        # Determine index of last End Of Sentence
        last_EOS = max(
            answer.rfind("."),
            answer.rfind("!"),
            answer.rfind("?"),
        )

        # In case the first sentence is very long
        if last_EOS == -1:
            return answer + "..."

        return answer[0].upper() + answer[1:last_EOS + 1]

    def generate_response(self, question: str) -> str:
        """
        Uses a language model and a tokenizer to create a natural language sample
        in response to a question string.

        @return: processed answer string
        """
        question = self._preprocess_question(question)
        question_len = len(question)

        input_ids = self.tokenizer.encode(question, return_tensors="pt")

        sample_output = self.model.generate(
            input_ids,
            do_sample=True,
            max_length=60,
            top_p=0.8,  # nucleus sampling
            top_k=0,  # top-k disabled
        )
        answer = self.tokenizer.decode(sample_output[0], skip_special_tokens=True)

        # Cut off question (which is part of the answer by default)
        answer = self._process_answer(answer[question_len:])

        return answer

    def start_dialogue(self):
        """
        Opens a dialogue with the language model for testing purposes.
        Can be stopped by typing "exit".
        """
        while True:
            question = input("Your sentence:   ")
            if question == "exit":
                return
            answer = self.generate_response(question)
            print("Output:\n" + 100 * '-')
            print(answer)


class TalkshowGuests(dict):
    """
    Creates a dictionary where keys are politician names
    and values are the corresponding chatbots.
    """
    path_template = "app\\conversation\\gpt2-{}"

    def __init__(self, politicians: List):
        super(TalkshowGuests, self).__init__()

        print("Start loading models...")

        for politician in politicians:
            model_path = self.path_template.format(politician)

            print(f"Creating ChatBot for {politician.capitalize()}...")
            chatbot = ChatBot(politician, model_path)
            print("Done.")

            super().__setitem__(politician, chatbot)


if __name__ == "__main__":
    try:
        test_politician = sys.argv[1]
        model_path = f".\\conversation\\gpt2-{test_politician}"
    except IndexError:
        test_politician = None
        model_path = "anonymous-german-nlp/german-gpt2"

    chatbot = ChatBot(test_politician, model_path)
    chatbot.start_dialogue()




