from functools import partial
import os

import gradio as gr
import requests
import spacy
from pyvis.network import Network
from spacy import displacy
from spacy.tokens import Doc, Span

from relik.common.utils import CONFIG_NAME, from_cache
from relik.inference.annotator import Relik
from relik.inference.data.objects import RelikOutput
from relik.inference.serve.frontend.utils import get_random_color


LOGO = """
<div style="text-align: center; display: flex; flex-direction: column; align-items: center;">
    <img src="https://cdn-uploads.huggingface.co/production/uploads/5f0b462819cb630495b814d7/NJPRSmlwi6sY1qQEIcgRV.png" style="max-width: 850px; height: auto;"> 
</div>
"""

DESCRIPTION = """
<div style="display:flex; justify-content: center; align-items: center; flex-direction: row;">
    <a href="https://2024.aclweb.org/"><img src="http://img.shields.io/badge/ACL-2024-4b44ce.svg"></a> &nbsp; &nbsp; 
    <a href="https://aclanthology.org/"><img src="http://img.shields.io/badge/paper-ACL--anthology-B31B1B.svg"></a> &nbsp; &nbsp; 
    <a href="https://arxiv.org/abs/2408.00103"><img src="https://img.shields.io/badge/arXiv-2408.00103-b31b1b.svg"></a>
</div>
<br>
<div style="display:flex; justify-content: center; align-items: center; flex-direction: row;">
    <a href="https://huggingface.co/collections/sapienzanlp/relik-retrieve-read-and-link-665d9e4a5c3ecba98c1bef19"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Collection-FCD21D"></a> &nbsp; &nbsp;
    <a href="https://github.com/SapienzaNLP/relik"><img src="https://img.shields.io/badge/GitHub-Repo-121013?logo=github&logoColor=white"></a> &nbsp; &nbsp;
    <a href="https://github.com/SapienzaNLP/relik/releases"><img src="https://img.shields.io/github/v/release/SapienzaNLP/relik"></a>
</div>
<br>
<div style="display:flex; justify-content: center; align-items: center; flex-direction: row;">
    <a href="https://nlp.uniroma1.it/"><img src="https://img.shields.io/badge/Sapienza NLP-802433.svg?logo=data:image/png;base64, "></a> &nbsp; &nbsp;
    <a href="https://babelscape.it/"><img src="https://img.shields.io/badge/Babelscape-215489.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAMAAACdt4HsAAAC3FBMVEUAAAAAAAAAAP8A//////8AAIAAgIAAgP+AgICAgP+qqqoAgL+AgL+Av7+/v7+/v/8AZpkzZpkAK4ArVYCAqtUAYJ9gn78ccao5cY4XRotddKIAK2oVQGoVaqqVqr9iibFJkraktsgRd6pQj69LeKUAY6oOOXEOY6oANnkbQ3lVeZ4XRnQAZKYWb6YrdaoJaKovcZdMcY4AYacSaqcRZqo8ZpEhc60APnQIPHgkUHwHarEjToAAN24cTHwNO3YTaqgfS4MfV4MkTYIXbq4RQ3oAXKMQaqogUIAUR4AURngAZ6cdcbEOQ3gOZagJZqwzWIszXYsfTH0JZacJZqYdb60YdKsPRHwaba0PZ6oPQngEQX4KQXcXRXkddK4TcKwacKwWbq0WcbAMaKkGZawGQHgRbasNQHYIP3kLP3Yaca4QRHgDYqkKP3gKQXgNRHsPbK0KaqgOaqgObKsRbKscToIFQnwFZaUEQX0SbasEZKgHYqgJPngGZKUGZqUMaasEZqsGZKYLaasEZaoFaKoFQHcDZqkDZasLP3gDP3oDZaYZb60EaKwEZ6kBPXcGP3kEPXgFO3YFP3kHZ6kDZ6gDZaoHaKcDPXkEQXgIZ6oDZ6oDZ60IP3cEZ6kIQHkCZqoCZ6wEZ6oCP3cEPXYCPXYCP3gCPXcHPnkIQXkBZqoFP3gFQHwJQXkBZ6sCPnkDQXkBP3kBaKkFQHkAZKoCZ6kAZ6oCPncCZ6oBZqoAZaoCP3kBZqoCPncBZ6sAZagCQHgDQHkAQXgAZ6kAZ6oAZ6sCQHkAZqoDPncCPHgAZqoAZ6oCQXkEQXsAZqsAZ6oAZ6wAZqsDP3oFQnsFQnwAZ6oCPngAZ6oAZ6wBQHkBQXsAZ6oBQHgCQHkBP3kBQXoCP3kCQXkAZqsAZq0AZ6oAZ6sBQHkBQXcDQHkAZ6kAZ6oAZ6sAZ6wAaKoBQHgBQHkBQHoCQHkCQHoDP3kDQHkEQHoQVPeyAAAA53RSTlMAAQEBAQICAgICAwQEBAQEBQUGBgYICAkJCwsMDAwMDQ4ODxAREhISExMVFhcXGBsbGx0dHh4fISIjJCQlJScpKSkrLC4vMDAyMzQ0NTU3Nzc5Ojw+QERERUZHS05PUFBRUVNWV1tfYWFhYmVmZmZmampqamxtb3JzdXV3enqDhYeLjZaYm52dn6GlrK6vr7O6urq8vb2/wcHDx8rMzM3Nz9DQ0dTW2dna29zc3d3d3t/f4OHk5OTl6enq6+zu8PDx8fHy8/T09fb29vb3+Pj5+fn5+vr7+/v7/Pz8/f39/f7+/v7+/v6z95VcAAACPklEQVRYw2NgGAXUB9JO/gHowN9BQ4po/UtOX8IAF08d2N0XyUKUAenvsYJPQLyz15MIA8o/vccFPr2/0qJNiQEgsN6WkTID3t+woNCA94skKTTgTS6FBnyaSKEB74/LUWjAO0sKDXhvQ6EBdwwpNGArpbHQT2ks5FFowFIZygx46kFZXrhVwkqRAftzKCoPbk/zoaBEundkgiMD+UXa3dlhasQVqlgM+PBpZbUdI7HFOqYBF6ZnSJBQr6AbcH1hIFSGQ1YnJLurKlSVnQQD3i52F4ZKREzZc/bxixcvHx+erEKsAZ/XpfBARL1btwE1Q8HLg2W8RBlwv0MTLMQe1H31BQp41chFhAH70pjBXref9PAlqv4XL19WEjZgozmYz9d59gUWcNSAkAHLrMHWJ29+gR3ME8RvwGpTEMe47TUO/S+uOeM1YLkikMnktf3lC5ygAY8Bn1ZpgZxfc/kFHrCJE7cBx1yADKWel/j0v3hggtOAJzFA2mzOCwIgHJcBz9qB+VZ3DSH9LxJwGbCWiYFBby9B/S8ScRjwMYuBQXzWS8IGxGI3oHQuPwND0wsigBt2A6KiGRiCzxOh/4QCdgOEGBjYNhDjgKm4c3TdIyL0P47DqV//MjEOOKSO04AKYvQ/rsWpX2wHMQaswB0CSc+J0H/SF7cBzUSkofOpeIrlBQS1v9wVj0e//BlC+m/OEMFXLxgR0j/TTwBvzSS65SUSQNN8bn59MDehytEqvxgOigqQQGGmq/Jon5Q2AADvLjKU0eBdEwAAAABJRU5ErkJggg=="></a>
</div>
<br>

<div style="text-align: center; display: flex; flex-direction: column; align-items: center;">

<h2>
    Retrieve, Read and LinK: Fast and Accurate Entity Linking and Relation Extraction on an Academic Budget
    <br>
    <span style="font-size: 0.9em;">
    A blazing fast and lightweight Information Extraction model for Entity Linking and Relation Extraction. 
    </span>
    <br>
    <span style="color: #919191; font-weight: 400; font-size: 0.8em;">
        <a href="https://riccardorlando.xyz/" style="color: #919191;" target="_blank">Riccardo Orlando</a>, 
        <a href="https://littlepea13.github.io/" style="color: #919191;" target="_blank">Pere-Lluís Huguet Cabot</a>, 
        <a href="https://edobobo.github.io/" style="color: #919191;" target="_blank">Edoardo Barba</a>, 
        and <a href="https://www.diag.uniroma1.it/navigli/" style="color: #919191;" target="_blank">Roberto Navigli</a> 
    </span>
<h2>
</div>
"""

INSTRUCTION = """
## Use it locally

Installation from PyPI

```bash
pip install relik
```

ReLiK is a lightweight and fast model for **Entity Linking** and **Relation Extraction**.
It is composed of two main components: a **retriever** and a **reader**.
The retriever is responsible for retrieving relevant documents from a large collection of documents,
while the reader is responsible for extracting entities and relations from the retrieved documents.
ReLiK can be used with the `from_pretrained` method to load a pre-trained pipeline.

Here is an example of how to use ReLiK for Entity Linking:

```python
from relik import Relik
from relik.inference.data.objects import RelikOutput

relik = Relik.from_pretrained("sapienzanlp/relik-entity-linking-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")

# RelikOutput(
#     text="Michael Jordan was one of the best players in the NBA.",
#     tokens=['Michael', 'Jordan', 'was', 'one', 'of', 'the', 'best', 'players', 'in', 'the', 'NBA', '.'],
#     id=0,
#     spans=[
#         Span(start=0, end=14, label="Michael Jordan", text="Michael Jordan"),
#         Span(start=50, end=53, label="National Basketball Association", text="NBA"),
#     ],
#     triplets=[],
#     candidates=Candidates(
#         span=[
#             [
#                 [
#                     {"text": "Michael Jordan", "id": 4484083},
#                     {"text": "National Basketball Association", "id": 5209815},
#                     {"text": "Walter Jordan", "id": 2340190},
#                     {"text": "Jordan", "id": 3486773},
#                     {"text": "50 Greatest Players in NBA History", "id": 1742909},
#                     ...
#                 ]
#             ]
#         ]
#     ),
# )
```

and for Relation Extraction:

```python
from relik import Relik
from relik.inference.data.objects import RelikOutput

relik = Relik.from_pretrained("sapienzanlp/relik-relation-extraction-large")
relik_out: RelikOutput = relik("Michael Jordan was one of the best players in the NBA.")
```

For more information, please refer to the [source code](https://github.com/SapienzaNLP/relik/).
"""


# relik_available_models = [
#     "relik-ie/relik-reader-small-cie-wikipedia",
# ]

# relik_models = {
#     "relik-ie/relik-reader-small-cie-wikipedia": Relik.from_pretrained(
#         "relik-ie/relik-reader-small-cie-wikipedia",
#         index_precision="bf16",
#         device="cuda",
#         reader_kwargs={"dataset_kwargs": {"use_nme": True}},
#     ),
# }


def get_span_annotations(response, doc):
    dict_ents = {}
    el_link_wrapper = "<link rel='stylesheet' href='https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css'><a href='https://en.wikipedia.org/wiki/{}' style='color: #414141'><i class='fa-brands fa-wikipedia-w fa-xs' style='color: #414141'></i> {}</a>"

    spans = []
    for idx, span in enumerate(response.spans):
        spans.append(
            Span(
                doc,
                span.start,
                span.end,
                (
                    el_link_wrapper.format(span.label.replace(" ", "_"), span.label)
                    if span.label != "--NME--"
                    else "--NME--"
                ),
                # kb_id=span.label.replace(" ", "_")
            )
        )
        dict_ents[(span.start, span.end)] = (
            span.label + str(idx),
            doc[span.start : span.end].text,
            span.label,
            span.label.replace(" ", "_"),
        )
    colors = get_random_color(set([span.label_ for span in spans]))
    return spans, colors, dict_ents


def generate_graph(
    spans, response, colors, dict_ents, bgcolor="#111827", font_color="white"
):
    g = Network(
        width="720px",
        height="600px",
        directed=True,
        notebook=False,
        bgcolor=bgcolor,
        font_color=font_color,
    )
    g.barnes_hut(
        gravity=-3000,
        central_gravity=0.3,
        spring_length=50,
        spring_strength=0.001,
        damping=0.09,
        overlap=0,
    )

    for ent in spans:
        # if not NME use title:
        if dict_ents[(ent.start, ent.end)][2] != "--NME--":
            g.add_node(
                dict_ents[(ent.start, ent.end)][2],
                label=dict_ents[(ent.start, ent.end)][2],
                color=colors[ent.label_],
                title=dict_ents[(ent.start, ent.end)][2],
                size=15,
                labelHighlightBold=True,
            )
        else:
            g.add_node(
                ent.text,
                label=ent.text,
                color=colors[ent.label_],
                title=ent.text,
                size=15,
                labelHighlightBold=True,
            )
    seen_rels = set()
    for rel in response.triplets:
        if (
            dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--"
        ):
            if (rel.subject.text, rel.object.text, rel.label) in seen_rels:
                continue
        elif (
            dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--"
        ):
            if (
                rel.subject.text,
                dict_ents[(rel.object.start, rel.object.end)][2],
                rel.label,
            ) in seen_rels:
                continue
        elif (
            dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--"
        ):
            if (
                dict_ents[(rel.subject.start, rel.subject.end)][2],
                rel.object.text,
                rel.label,
            ) in seen_rels:
                continue
        else:
            if (
                dict_ents[(rel.subject.start, rel.subject.end)][2],
                dict_ents[(rel.object.start, rel.object.end)][2],
                rel.label,
            ) in seen_rels:
                continue

        g.add_edge(
            (
                dict_ents[(rel.subject.start, rel.subject.end)][2]
                if dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--"
                else dict_ents[(rel.subject.start, rel.subject.end)][1]
            ),
            (
                dict_ents[(rel.object.start, rel.object.end)][2]
                if dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--"
                else dict_ents[(rel.object.start, rel.object.end)][1]
            ),
            label=rel.label,
            title=rel.label,
        )
        if (
            dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--"
        ):
            seen_rels.add(
                (
                    dict_ents[(rel.subject.start, rel.subject.end)][2],
                    dict_ents[(rel.object.start, rel.object.end)][2],
                    rel.label,
                )
            )
        elif (
            dict_ents[(rel.subject.start, rel.subject.end)][2] != "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] == "--NME--"
        ):
            seen_rels.add(
                (
                    dict_ents[(rel.subject.start, rel.subject.end)][2],
                    rel.object.text,
                    rel.label,
                )
            )
        elif (
            dict_ents[(rel.subject.start, rel.subject.end)][2] == "--NME--"
            and dict_ents[(rel.object.start, rel.object.end)][2] != "--NME--"
        ):
            seen_rels.add(
                (
                    rel.subject.text,
                    dict_ents[(rel.object.start, rel.object.end)][2],
                    rel.label,
                )
            )
        else:
            seen_rels.add((rel.subject.text, rel.object.text, rel.label))
    # g.show(filename, notebook=False)
    html = g.generate_html()
    # need to remove ' from HTML
    html = html.replace("'", '"')

    return f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""


RELIK = os.getenv("RELIK", "localhost:8000/api/relik")


def text_analysis(Text, Relation_Threshold, Window_Size, Window_Stride):
    relik = RELIK
    nlp = spacy.blank("xx")
    # annotated_text = relik(
    #     Text,
    #     annotation_type="word",
    #     num_workers=0,
    #     remove_nmes=False,
    #     relation_threshold=Relation_Threshold,
    #     window_size=Window_Size,
    #     window_stride=Window_Stride,
    # )
    print(f"Using ReLiK at {relik}")
    print(
        f"Querying ReLiK with ?text={Text}&relation_threshold={Relation_Threshold}&window_size={Window_Size}&window_stride={Window_Stride}&annotation_type=word&remove_nmes=False"
    )
    response = requests.get(
        f"http://{relik}/?text={Text}&relation_threshold={Relation_Threshold}&window_size={Window_Size}&window_stride={Window_Stride}&annotation_type=word&remove_nmes=False",
    )
    if response.status_code != 200:
        raise gr.Error(response.text)
    annotated_text = RelikOutput(**response.json())
    doc = Doc(nlp.vocab, words=[token.text for token in annotated_text.tokens])
    spans, colors, dict_ents = get_span_annotations(response=annotated_text, doc=doc)
    doc.spans["sc"] = spans

    # build the EL display
    display_el = displacy.render(
        doc, style="span", options={"colors": colors}
    )  # , "kb_url_template": "https://en.wikipedia.org/wiki/{}"})
    display_el = display_el.replace("\n", " ")
    # heuristic, prevents split of annotation decorations
    display_el = display_el.replace(
        "border-radius: 0.35em;",
        "border-radius: 0.35em; white-space: nowrap;",
    )
    display_el = display_el.replace(
        "span style",
        "span id='el' style",
    )
    display_re = ""
    if annotated_text.triplets:
        # background_color should be the same as the background of the page
        display_re = generate_graph(spans, annotated_text, colors, dict_ents)
    return display_el, display_re


theme = theme = gr.themes.Base(
    primary_hue="rose",
    secondary_hue="rose",
    text_size="lg",
    # font=[gr.themes.GoogleFont("Montserrat"), "Arial", "sans-serif"],
)

css = """
h1 {
  text-align: center;
  display: block;
}
mark {
    color: black;
}
#el {
    white-space: nowrap;
}
"""

with gr.Blocks(fill_height=True, css=css, theme=theme) as demo:
    # check if demo is running in dark mode
    gr.Markdown(LOGO)
    gr.Markdown(DESCRIPTION)
    gr.Interface(
        text_analysis,
        [
            gr.Textbox(label="Input Text", placeholder="Enter sentence here..."),
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.05,
                value=0.5,
                label="Relation Threshold",
                info="Minimum confidence for relation extraction (Only for RE and cIE)",
            ),
            gr.Slider(
                minimum=16,
                maximum=128,
                step=16,
                value=64,
                label="Window Size",
                info="Window size for the sliding window",
            ),
            gr.Slider(
                minimum=8,
                maximum=64,
                step=8,
                value=32,
                label="Window Stride",
                info="Window stride for the sliding window",
            ),
        ],
        [gr.HTML(label="Entities"), gr.HTML(label="Relations")],
        examples=[
            [
                "Avram Noam Chomsky  born December 7, 1928) is an American professor and public intellectual known for his work in linguistics, political activism, and social criticism. Sometimes called 'the father of modern linguistics', Chomsky is also a major figure in analytic philosophy and one of the founders of the field of cognitive science. He is a laureate professor of linguistics at the University of Arizona and an institute professor emeritus at the Massachusetts Institute of Technology (MIT). Among the most cited living authors, Chomsky has written more than 150 books on topics such as linguistics, war, and politics. In addition to his work in linguistics, since the 1960s Chomsky has been an influential voice on the American left as a consistent critic of U.S. foreign policy, contemporary capitalism, and corporate influence on political institutions and the media."
            ],
            [
                "'Bella ciao' (Italian pronunciation: [ˈbɛlla ˈtʃaːo]; 'Goodbye beautiful') is an Italian song dedicated to the partisans of the Italian resistance, which fought against the occupying troops of Nazi Germany and the collaborationist Fascist forces during the liberation of Italy. It was based on a folk song of the late 19th century, sung by female workers of the paddy fields in Northern Italy (mondine) in protest against harsh working conditions. Versions of 'Bella ciao' continue to be sung worldwide as a hymn of resistance."
            ],
        ],
        allow_flagging="never",
    )
    gr.Markdown("")
    gr.Markdown(INSTRUCTION)


def main(server_port: int = 8001):
    demo.launch(server_port)


if __name__ == "__main__":
    main()
