# Import libraries
import spacy
from spacy import displacy
import gradio as gr
import re
from huggingface_hub import hf_hub_download
               
# Spacy.load().
nlp = spacy.load("de_pipeline")

# Define examples to be displayed below the textbox
examples = ["""Zahlreiche öffentliche Gebäude in Berlin werden ab sofort nachts nicht mehr angestrahlt. Neben dem Dom sind das die Marienkirche, das Zeughaus und das Alte Palais im Bezirk Mitte, wie die Senatsumweltverwaltung am Mittwoch mitteilte. Auch der Lustgarten und das Reiterstandbild Unter den Linden sollen demnach schon in der kommenden Nacht nicht mehr beleuchtet werden. In den kommenden Wochen sollen dann auch die Siegessäule, Staatsoper, Deutsche Oper, Gedächtniskirche, das Rote Rathaus und das Schloss Charlottenburg folgen. Statuen im Tiergarten sowie Brückenbauwerke sollen ebenfalls nicht mehr angestrahlt werden.""",
             """Zur gleichen Zeit / wurde das Olympiastadion gebaut, / leicht und transparent. / Damals wollte man eines / unbedingt vermeiden: Monumentalität. / Die spiegelt ein ungeliebtes Denkmal: / Die Kongresshalle in Nürnberg, / der größte erhaltene Bau des Nationalsozialismus.""",
            """Ich bin jetzt noch in Zivil, ab Montag habe ich auch Uniform an. Das ist wieder ein neues Gefühl. Aber ich freue mich auf die neue Aufgabe, die bei der Polizeistation Waldmünchen anfängt. Tolle Sache.""",
            """Der Bauherr Ludwig II. / und Schloss Neuschwanstein / sind in der Kulturgeschichte / so eng miteinander verzahnt / wie kaum ein anderes Bauwerk / mit dem Bauherrn.""",
            """Das Verfahren gegen die zehn Rechtsextremisten findet ausgerechnet im Saal A101 statt, also jenem Saal des Münchner Justizzentrums, in dem fünf Jahre lang der sogenannte NSU-Prozess verhandelt wurde.""",
            """Vor knapp einem Monat wurde das Würzburger Mozartfest im Kaisersaal der Residenz eröffnet.""",
            """Würzburg war schon immer ein wichtiges politisches Zentrum, ehemals fürstbischöfliche Residenz, heute Sitz der Regierung von Unterfranken.""",
            """Was ein Hotel für diese stolze Summe bietet, das schauen wir uns jetzt einmal an - und öffnen ehrfürchtig die Türen zur Ludwig-Suite im Hotel Vier Jahreszeiten, wo Royals und Celebrities nächtigen.""",
            """Das Münchner Hofbräuhaus hat das Dresdner Hofbrauhaus wegen der Verletzung von Rechten verklagt. Am Montag beschäftigt sich das Münchner Landgericht mit dem Streit, der um die Frage kreist, ob "au" dasselbe ist wie "äu".""",
            """Einer der schönsten Märkte im Dachauer Land  / ist Altomünster mit seiner Klosterkirche. / Gleich daneben steht der Kapplerbräu mit seinem Brauereimuseum."""
            ]
 # Add title, description
title = "LandmarkNER"
description = """<center>
Named Entity Recognition (NER) Modell zur Erkennung bayerischer Landmarken, trainiert mit Untertiteln von Videos des Bayerischen Rundfunks. Ausgabe der Wikipedia Titel mit  <a href="https://huggingface.co/facebook/mgenre-wiki">mGENRE</a>.
</center></p>"""

# Define Landmark NER function to be implemented in gradio
def ner(Eingabe):
    # remove separators
    Eingabe = Eingabe.replace(" / ", " ")
    # remove timecodes and special character lines
    regex_tc = r"^\*.*\*$|^♪.*\.$|(?!^.*[a-zA-Z].*$)^.+"
    Eingabe_1 = re.sub(regex_tc, " ", Eingabe, 0, re.MULTILINE)
    # remove all whitespace characters (space, tab, newline, return, formfeed)
    Eingabe_2 = " ".join(Eingabe_1.split())
    doc = nlp(Eingabe_2)
    #ents = [(e.text) for e in doc.ents]
    output = displacy.render(doc, style="ent")
    return '<br>' + output +'<br>'
           
           
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import pickle
# or a memory efficient but a bit slower prefix tree (trie) -- it is implemented with `marisa_trie` from
# https://huggingface.co/facebook/mgenre-wiki/blob/main/

file_path_marisa_trie = hf_hub_download("facebook/mgenre-wiki", "titles_lang_all105_marisa_trie_with_redirect.pkl")

from genre.trie import MarisaTrie
with open(file_path_marisa_trie, "rb") as f:
    trie = pickle.load(f)

tokenizer = AutoTokenizer.from_pretrained("facebook/mgenre-wiki")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/mgenre-wiki").eval()

def entity_linking_custom(Eingabe):
    # remove separators
    Eingabe = Eingabe.replace(" / ", " ")
    # remove timecodes and special character lines
    regex_tc = r"^\*.*\*$|^♪.*\.$|(?!^.*[a-zA-Z].*$)^.+"
    Eingabe_1 = re.sub(regex_tc, " ", Eingabe, 0, re.MULTILINE)
    # remove all whitespace characters (space, tab, newline, return, formfeed)
    Eingabe_2 = " ".join(Eingabe_1.split())
    doc = nlp(Eingabe_2)
    reshaped_text = []
    for ent in doc.ents:
      reshaped_text.append(doc.text.replace(ent.text, '[START]' + ent.text + '[END]'))
    generated_text = []
    for sent in reshaped_text:
      outputs = model.generate(
        **tokenizer(sent, return_tensors="pt"),
        num_beams=5,
        num_return_sequences=5,
        # OPTIONAL: use constrained beam search
        prefix_allowed_tokens_fn=lambda batch_id, sent: trie.get(sent.tolist()),
        )
      generated_text.append(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    output = []
    re_lang = r" >> [a-z]+"
    for el in generated_text:
      gen_text = el[0]
      shortened_text = re.sub(re_lang, "", gen_text, 0, re.MULTILINE)
      if shortened_text:
         output.append(shortened_text)
    return "\n".join(output)
    
# So sorry about this, but it did not work otherwise    
def update(input):
    return input
def clear(input):
  input = ' ' + ' '
  return input   

# Define blocks
demo = gr.Blocks()

with demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>"
                + title
                + "</h1>")
    gr.Markdown(description)
    with gr.Column():
        inp = gr.Textbox(placeholder="Füge einen eigenen Text ein oder wähle eines der Beispiele unten aus. Lass dir dann Landmarken im Text markieren oder deren Wikipedia Titel anzeigen.", label = "Eingabe", lines=3)
        butt1 = gr.Button("Landmarken markieren")
        out_ner = gr.HTML(label="Ausgabe des Textes mit markierten Landmarken", show_label=True)
        butt1.click(ner, inputs=inp, outputs=out_ner)
        butt2 = gr.Button("Wikipedia Titel anzeigen")
        out_rs = gr.Textbox(label = "Ausgabe Wikipedia Titel", show_label=True)
        butt2.click(entity_linking_custom, inputs=inp, outputs=out_rs)
        gr.Markdown("""
        ### Beispiele""")
        example = gr.Radio(choices = examples, value=0,label = "Beispiel auswählen", show_label=True)
        example.change(update, inputs = example, outputs=inp)
        example.change(clear, inputs = inp, outputs=[out_ner,out_rs])   
        
demo.launch()