import streamlit as st
from transformers import AutoTokenizer
from transformers import GPT2LMHeadModel
from transformers import set_seed

import meta
from normalizer import normalize
from utils import load_json
from utils import local_css

EXAMPLES = load_json("examples.json")
CK = ""
QK = "Q:"
AK = "A:"


class TextGeneration:
    def __init__(self):
        self.debug = False
        self.dummy_output = "Destiny's Child"
        self.tokenizer = None
        self.model = None
        self.model_name_or_path = "m3hrdadfi/gpt2-QA"
        self.length_margin = 100
        set_seed(42)

    def load(self):
        if not self.debug:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
            self.model = GPT2LMHeadModel.from_pretrained(self.model_name_or_path)

    def generate(self, prompt, generation_kwargs):

        if not self.debug:
            input_ids = self.tokenizer([prompt], return_tensors="pt")["input_ids"]
            max_length = len(input_ids[0]) + self.length_margin
            max_length = min(max_length, 1024)
            generation_kwargs["max_length"] = max_length

            generated = self.model.generate(
                input_ids,
                **generation_kwargs,
            )[0]

            answer = self.tokenizer.decode(generated, skip_special_tokens=True)
            found = answer.find(f"{AK}")
            if not found:
                return ""

            answer = [a.strip() for a in answer[found:].split(f"{AK}") if a.strip()]
            answer = answer[0] if len(answer) > 0 else ""
            return answer

        return self.dummy_output


@st.cache(allow_output_mutation=True)
def load_text_generator():
    generator = TextGeneration()
    generator.load()
    return generator


def main():
    st.set_page_config(
        page_title="GPT2 QA",
        page_icon="⁉️",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    local_css("assets/style.css")
    generator = load_text_generator()

    st.sidebar.markdown(meta.SIDEBAR_INFO)
    num_beams = st.sidebar.slider(
        label='Number of Beam',
        help="Number of beams for beam search",
        min_value=4,
        max_value=15,
        value=5,
        step=1
    )
    repetition_penalty = st.sidebar.slider(
        label='Repetition Penalty',
        help="The parameter for repetition penalty",
        min_value=1.0,
        max_value=3.0,
        value=1.0,
        step=0.1
    )
    length_penalty = st.sidebar.slider(
        label='Length Penalty',
        help="Exponential penalty to the length",
        min_value=0.0,
        max_value=2.0,
        value=1.0,
        step=0.1
    )
    early_stopping = st.sidebar.selectbox(
        label='Early Stopping ?',
        options=(True, False),
        help="Whether to stop the beam search when at least num_beams sentences are finished per batch or not",
    )
    generation_kwargs = {
        "num_beams": num_beams,
        "early_stopping": early_stopping,
        "repetition_penalty": repetition_penalty,
        "length_penalty": length_penalty,
    }

    st.markdown(meta.HEADER_INFO)
    prompts = [e["title"] for e in EXAMPLES] + ["Custom"]
    prompt = st.selectbox('Examples', prompts, index=len(prompts) - 1)

    if prompt == "Custom":
        prompt_box = {
            "context": meta.C_PROMPT_BOX,
            "question": meta.Q_PROMPT_BOX,
            "answers": [meta.A_PROMPT_BOX],
        }
    else:
        prompt_box = next(e for e in EXAMPLES if e["title"] == prompt)

    context = st.text_area("Enter context", prompt_box["context"], height=200)
    question = st.text_area("Enter question", prompt_box["question"], height=100)
    answer = "Ground Truth Answers: " + \
             "".join([f"<span class='ground-truth'>{answer}</span>" for answer in prompt_box["answers"]])
    st.markdown(
        f'<p>'
        f'{answer}'
        f'<p>',
        unsafe_allow_html=True
    )
    generation_kwargs_ph = st.empty()

    if st.button("Find the answer 🔎 "):
        with st.spinner(text="Searching ..."):
            generation_kwargs_ph.markdown(", ".join([f"`{k}`: {v}" for k, v in generation_kwargs.items()]))
            context = normalize(context)
            question = normalize(question)

            if context and question:
                text = f"{context} {QK} {question} {AK}"
                generated_answer = generator.generate(text, generation_kwargs)
                generated_answer = f"{AK} {generated_answer}".strip()
                context = f"{CK} {context}".strip()
                question = f"{QK} {question}".strip()

                st.markdown(
                    f'<p>'
                    f'<span class="result-text">{context}<span><br/><br/>'
                    f'<span class="result-text">{question}<span><br/><br/>'
                    f'<span class="result-text generated-text">{generated_answer} </span>'
                    f'</p>',
                    unsafe_allow_html=True
                )


if __name__ == '__main__':
    main()