from typing import Tuple, List, Union, Dict, Mapping
import base64
import os

from bs4 import BeautifulSoup
import gradio as gr
from spacy import displacy
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    BatchEncoding,
    AutoModelForSeq2SeqLM,
    DataCollatorForTokenClassification,
)
import torch

from utils import get_dependencies, preprocess_text
from models import (
    DependencyRobertaForTokenClassification,
    LabelRobertaForTokenClassification,
)


DEFAULT_TEXT = "τίω δέ μιν ἐν καρὸς αἴσῃ."
BUTTON_CSS = "float: right; --tw-border-opacity: 1; border-color: rgb(229 231 235 / var(--tw-border-opacity)); --tw-gradient-from: rgb(243 244 246 / 0.7); --tw-gradient-stops: var(--tw-gradient-from), var(--tw-gradient-to, rgb(243 244 246 / 0)); --tw-gradient-to: rgb(229 231 235 / 0.8); --tw-text-opacity: 1; color: rgb(55 65 81 / var(--tw-text-opacity));    border-width: 1px; --tw-bg-opacity: 1; background-color: rgb(255 255 255 / var(--tw-bg-opacity)); background-image: linear-gradient(to bottom right, var(--tw-gradient-stops)); display: inline-flex; flex: 1 1 0%; align-items: center; justify-content: center;    --tw-shadow: 0 1px 2px 0 rgb(0 0 0 / 0.05); --tw-shadow-colored: 0 1px 2px 0 var(--tw-shadow-color); box-shadow: var(--tw-ring-offset-shadow, 0 0 #0000), var(--tw-ring-shadow, 0 0 #0000), var(--tw-shadow); -webkit-appearance: button; border-radius: 0.5rem; padding-top: 0.5rem; padding-bottom: 0.5rem; padding-left: 1rem; padding-right: 1rem; font-size: 1rem; line-height: 1.5rem; font-weight: 600;"
DEFAULT_COLOR = "white"

MODEL_PATHS = {
    "POS": "bowphs/testid",
    "LEMMATIZATION": "bowphs/lemmatization-demo",
    "DEPENDENCY": "bowphs/depenBERTa_perseus",
    "LABELS": "bowphs/depenBERTa_labler_perseus",
}
MODEL_MAX_LENGTH = 512

AUTH_TOKEN = os.environ.get("TOKEN") or True
# PoS
pos_tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATHS["POS"], model_max_length=MODEL_MAX_LENGTH, use_auth_token=AUTH_TOKEN, #revision="402ab7d25f49e83a67b955ebbc172b5459fbd939",
)
pos_model = AutoModelForTokenClassification.from_pretrained(
    MODEL_PATHS["POS"], use_auth_token=AUTH_TOKEN, #revision="402ab7d25f49e83a67b955ebbc172b5459fbd939",
)

# Lemmatization
lemmatizer_tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATHS["LEMMATIZATION"],
    model_max_length=MODEL_MAX_LENGTH,
    use_auth_token=AUTH_TOKEN,
)
lemmatizer_model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_PATHS["LEMMATIZATION"], use_auth_token=AUTH_TOKEN
)

# Dependency Parsing
dependency_tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATHS["DEPENDENCY"],
    model_max_length=MODEL_MAX_LENGTH,
    use_auth_token=AUTH_TOKEN,
)
arcs_model = DependencyRobertaForTokenClassification.from_pretrained(
    MODEL_PATHS["DEPENDENCY"], use_auth_token=AUTH_TOKEN
)
labels_model = LabelRobertaForTokenClassification.from_pretrained(
    MODEL_PATHS["LABELS"], use_auth_token=AUTH_TOKEN
)

data_collator = DataCollatorForTokenClassification(dependency_tokenizer)


def is_valid_selection(col_arcs, col_labels) -> bool:
    if not col_arcs and col_labels:
        return False
    return True


def get_pos_predictions(inputs) -> torch.Tensor:
    """Get part of speech predictions."""
    return pos_model(inputs["input_ids"]).logits.argmax(-1)  # type: ignore


def execute_parse(
    text_input: str,
    col_pos: bool,
    col_arcs: bool,
    col_labels: bool,
    col_lemmata: bool,
    compact: bool,
    bg: str,
    text: str,
) -> Tuple[str, str]:
    if is_valid_selection(col_arcs, col_labels):
        return parse(
            text_input, col_pos, col_arcs, col_labels, col_lemmata, compact, bg, text
        )
    return "Please check 'Dependency Arcs' before checking 'Dependency Labels'", ""


def lemmatize(tokens: List[str]) -> List[str]:
    def construct_task(word_idx: int) -> str:
        return f"lemmatize: {' '.join(tokens[:word_idx])} <extra_id_0> {tokens[word_idx]} <extra_id_1> {' '.join(list(tokens[word_idx]))} <extra_id_2> {' '.join(tokens[word_idx+1:])}"

    predictions = [
        lemmatizer_tokenizer.decode(
            lemmatizer_model.generate(
                lemmatizer_tokenizer(construct_task(word_idx), return_tensors="pt")[
                    "input_ids"
                ],
                max_length=20,
                num_beams=5,
                num_return_sequences=1,
                early_stopping=True,
            )[0],
            skip_special_tokens=True,
        )
        for word_idx in range(len(tokens))
    ]

    return predictions


def add_lemma_visualization(soup, lemmata: List[str], col_arcs: bool) -> str:
    for token, lemma in zip(soup.find_all(class_="displacy-token")[col_arcs:], lemmata):
        pos_tag = token.find(class_="displacy-tag")
        lemma_tag = soup.new_tag(
            "tspan",
            class_="displacy-lemma",
            dy="2em",
            fill="currentColor",
            x=pos_tag.attrs["x"],
        )
        lemma_tag.string = lemma
        pos_tag.insert_after(lemma_tag)
    return str(soup)


def download_svg(svg):
    encode = base64.b64encode(bytes(svg, "utf-8"))
    img = "data:image/svg+xml;base64," + str(encode)[2:-1]
    html = f'<a download="displacy.svg" href="{img}" style="{BUTTON_CSS}">Download as SVG</a>'
    return html


def prepare_doc(
    tokens: List[str], col_pos: bool, pos_outputs: torch.Tensor, inputs: BatchEncoding,
) -> Dict[str, List[Dict[str, str]]]:
    doc: Dict[str, List[Dict[str, str]]] = {
        "words": [], #[{"text": "ROOT", "tag": ""}],
        "arcs": [],
    }
    word_ids = inputs.word_ids()
    previous_word_idx = None

    for idx, word_idx in enumerate(word_ids):
        if word_idx != previous_word_idx and word_idx is not None:
            tag_repr = (
                pos_model.config.id2label[pos_outputs[0][idx].item()] if col_pos else ""
            )
            doc["words"].append({"text": tokens[word_idx], "tag": tag_repr})
            previous_word_idx = word_idx

    return doc


def parse(
    text_input: str,
    col_pos: bool,
    col_arcs: bool,
    col_labels: bool,
    col_lemmata: bool,
    compact: bool,
    bg: str,
    text: str,
) -> Tuple[str, str]:
    tokens = preprocess_text(text_input)
    inputs = pos_tokenizer(
        tokens,
        return_tensors="pt",
        truncation=True,
        padding=True,
        is_split_into_words=True,
    )
    pos_outputs = get_pos_predictions(inputs)

    doc = prepare_doc(tokens, col_pos, pos_outputs, inputs)

    if col_arcs:
        doc["words"].insert(0, {"text": "ROOT", "tag": ""})
        doc["arcs"] = get_dependencies(
            arcs_model,
            labels_model,
            dependency_tokenizer,
            data_collator,
            col_labels,
            tokens,
        )["arcs"]

    options = {"compact": compact, "bg": bg, "color": text}
    svg = displacy.render(doc, manual=True, style="dep", options=options)

    if col_lemmata:
        soup = BeautifulSoup(svg, "lxml-xml")
        lemmata = lemmatize(tokens)
        svg = add_lemma_visualization(soup, lemmata, col_arcs)

    download_link = download_svg(svg)

    return svg, download_link


def setup_parser_ui():
    demo = gr.Blocks(css="scrollbar.css")
    with demo:
        with gr.Box():
            with gr.Row():
                with gr.Column():
                    gr.Markdown("# Athena's Lens")
                    gr.Markdown(
                        "### From Ἀlkaios to Ὠrigen: A Modern Lens on Timeless Texts"
                    )
        with gr.Box():
            with gr.Column():
                gr.Markdown(" ## Enter some text")
                with gr.Row():
                    with gr.Column(scale=0.5):
                        text_input = gr.Textbox(
                            value=DEFAULT_TEXT, interactive=True, label="Input Text"
                        )
                with gr.Row():
                    with gr.Column(scale=0.25):
                        button = gr.Button("Update", variant="primary").style(
                            full_width=False
                        )
        with gr.Box():
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("## Parser")
                with gr.Row():
                    with gr.Column():
                        col_pos = gr.Checkbox(label="PoS Labels", value=True)
                        col_arcs = gr.Checkbox(label="Dependency Arcs", value=False)
                        col_labels = gr.Checkbox(label="Dependency Labels", value=False)
                        col_lemmata = gr.Checkbox(label="Lemmata", value=False)
                        compact = gr.Checkbox(label="Compact", value=False)
                    with gr.Column():
                        bg = gr.Textbox(label="Background Color", value=DEFAULT_COLOR)
                    with gr.Column():
                        text = gr.Textbox(label="Text Color", value="black")
                with gr.Row():
                    dep_output = gr.HTML(
                        value=parse(
                            DEFAULT_TEXT,
                            True,
                            False,
                            False,
                            False,
                            False,
                            DEFAULT_COLOR,
                            "black",
                        )[0]
                    )
                with gr.Row():
                    with gr.Column(scale=0.25):
                        dep_button = gr.Button(
                            "Update Parser", variant="primary"
                        ).style(full_width=False)
                    with gr.Column():
                        dep_download_button = gr.HTML(
                            value=download_svg(dep_output.value)
                        )

            with gr.Box():
                with gr.Column():
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("## Contact")
                            gr.Markdown(
                                "If you have any questions, suggestions, comments, or problems, feel free to [reach out](mailto:riemenschneider@cl.uni-heidelberg.de)."
                            )
                            gr.Markdown("## Citation")
                            gr.Markdown(
                                "This space uses models from [this](https://aclanthology.org/2023.acl-long.846.pdf) paper."
                            )
                            gr.Markdown(
                                """```bibtex
    @incollection{riemenschneider-frank-2023-exploring,
        title = "Exploring Large Language Models for Classical Philology",
        author = "Riemenschneider, Frederick  and Frank, Anette",
        booktitle = "Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
        month = jul,
        year = "2023",
        address = "Toronto, Canada",
        publisher = "Association for Computational Linguistics",
        url = "https://aclanthology.org/2023.acl-long.846",
        doi = "10.18653/v1/2023.acl-long.846",
        pages = "15181--15199",
    }
    ```
                    """
                            )

        button.click(
            execute_parse,
            inputs=[
                text_input,
                col_pos,
                col_arcs,
                col_labels,
                col_lemmata,
                compact,
                bg,
                text,
            ],
            outputs=[dep_output, dep_download_button],
        )

        dep_button.click(
            execute_parse,
            inputs=[
                text_input,
                col_pos,
                col_arcs,
                col_labels,
                col_lemmata,
                compact,
                bg,
                text,
            ],
            outputs=[dep_output, dep_download_button],
        )

    demo.launch()


def main():
    demo = setup_parser_ui()
    demo.launch()


if __name__ == "__main__":
    main()