import re
import string

import orjson
import streamlit as st
from annotated_text.util import get_annotated_html

from pipelines.keyphrase_extraction_pipeline import KeyphraseExtractionPipeline
from pipelines.keyphrase_generation_pipeline import KeyphraseGenerationPipeline


@st.cache(allow_output_mutation=True, show_spinner=False)
def load_pipeline(chosen_model):
    if "keyphrase-extraction" in chosen_model:
        return KeyphraseExtractionPipeline(chosen_model, truncation=True)
    elif "keyphrase-generation" in chosen_model:
        return KeyphraseGenerationPipeline(chosen_model, truncation=True)


def extract_keyphrases():
    st.session_state.keyphrases = pipe(st.session_state.input_text)
    st.session_state.history[f"run_{st.session_state.current_run_id}"] = {
        "run_id": st.session_state.current_run_id,
        "model": st.session_state.chosen_model,
        "text": st.session_state.input_text,
        "keyphrases": st.session_state.keyphrases,
    }
    st.session_state.current_run_id += 1


def get_annotated_text(text, keyphrases, color="#d294ff"):
    for keyphrase in keyphrases:
        text = re.sub(
            rf"({keyphrase})([^A-Za-z])",
            rf"$K:{keyphrases.index(keyphrase)}\2",
            text,
            flags=re.I,
        )

    result = []
    for i, word in enumerate(text.split(" ")):
        if "$K" in word and re.search(
            "(\d+)$", word.translate(str.maketrans("", "", string.punctuation))
        ):
            result.append(
                (
                    re.sub(
                        r"\$K:\d+",
                        keyphrases[
                            int(
                                re.search(
                                    "(\d+)$",
                                    word.translate(
                                        str.maketrans("", "", string.punctuation)
                                    ),
                                ).group(1)
                            )
                        ],
                        word,
                    ),
                    "KEY",
                    color,
                )
            )
        else:
            if i == len(st.session_state.input_text.split(" ")) - 1:
                result.append(f" {word}")
            elif i == 0:
                result.append(f"{word} ")
            else:
                result.append(f" {word} ")
    return result


def render_output(layout, runs, reverse=False):
    runs = list(runs.values())[::-1] if reverse else list(runs.values())
    for run in runs:
        layout.markdown(
            f"""
            <p style=\"margin-bottom: 0rem\"><strong>Run:</strong> {run.get('run_id')}</p>
            <p style=\"margin-bottom: 0rem\"><strong>Model:</strong> {run.get('model')}</p>
            """,
            unsafe_allow_html=True,
        )

        if "generation" in run.get("model"):
            abstractive_keyphrases = [
                keyphrase
                for keyphrase in run.get("keyphrases")
                if keyphrase.lower() not in run.get("text").lower()
            ]
            layout.markdown(
                f"<p style=\"margin-bottom: 0rem\"><strong>Absent keyphrases:</strong> {', '.join(abstractive_keyphrases) if abstractive_keyphrases else 'None' }</p>",
                unsafe_allow_html=True,
            )

        result = get_annotated_text(run.get("text"), list(run.get("keyphrases")))
        layout.markdown(
            f"""
            <p style="margin-bottom: 0.5rem"><strong>Text:</strong></p>
            {get_annotated_html(*result)}
            """,
            unsafe_allow_html=True,
        )
        layout.markdown("---")


if "config" not in st.session_state:
    with open("config.json", "r") as f:
        content = f.read()
    st.session_state.config = orjson.loads(content)
    st.session_state.history = {}
    st.session_state.keyphrases = []
    st.session_state.current_run_id = 1

st.set_page_config(
    page_icon="🔑",
    page_title="Keyphrase extraction/generation with Transformers",
    layout="centered",
)

with open("css/style.css") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

st.header("🔑 Keyphrase extraction/generation with Transformers")

description = """
Keyphrase extraction is a technique in text analysis where you extract the important keyphrases
from a text. Since this is a time-consuming process, Artificial Intelligence is used to automate it.
Currently, classical machine learning methods, that use statistics and linguistics, are widely used
for the extraction process. The fact that these methods have been widely used in the community has
the advantage that there are many easy-to-use libraries. Now with the recent innovations in
NLP, transformers can be used to improve keyphrase extraction. Transformers also focus on the semantics and
context of a document, which is quite an improvement.

This space gives you the ability to test around with some keyphrase extraction and generation models.
Keyphrase extraction models are transformers models fine-tuned as a token classification problem where
the tokens in a text are annotated as B (Beginning of a keyphrase), I (Inside a keyphrases), 
and O (Outside a keyhprase).

While keyphrase extraction can only extract keyphrases from a given text. Keyphrase generation models
work a bit differently. Here you use an encoder-decoder model like BART to generate keyphrases from a given text.
These models also have the ability to generate keyphrases, which are not present in the text 🤯.

Do you want to see some magic 🧙‍♂️? Try it out yourself! 👇
"""

st.write(description)

with st.form("keyphrase-extraction-form"):
    selectbox_container, _ = st.columns(2)

    st.session_state.chosen_model = selectbox_container.selectbox(
        "Choose your model:", st.session_state.config.get("models")
    )

    st.markdown(
        f"For more information about the chosen model, please be sure to check out the [🤗 Model Card](https://huggingface.co/{st.session_state.get('config').get('model_author')}/{st.session_state.chosen_model})."
    )

    st.session_state.input_text = (
        st.text_area(
            "✍ Input",
            st.session_state.config.get("example_text"),
            height=250,
            max_chars=2500,
        )
        .replace("\n", " ")
        .strip()
    )

    with st.spinner("Extracting keyphrases..."):
        pressed = st.form_submit_button("Extract")

if pressed and st.session_state.input_text != "":
    with st.spinner("Loading pipeline..."):
        pipe = load_pipeline(
            f"{st.session_state.config.get('model_author')}/{st.session_state.chosen_model}"
        )
    with st.spinner("Extracting keyphrases"):
        extract_keyphrases()
elif st.session_state.input_text == "":
    st.error("The text input is empty 🙃 Please provide a text in the input field.")

options = st.multiselect(
    "Specify the runs you want to see",
    st.session_state.history.keys(),
    format_func=lambda run_id: f"Run {run_id.split('_')[1]}",
)

if len(st.session_state.history.keys()) > 0:
    if options:
        render_output(
            st,
            {key: st.session_state.history[key] for key in options},
        )
    else:
        render_output(st, st.session_state.history, reverse=True)