import gradio as gr
from transformers import pipeline
import tokenizer
from difflib import Differ, SequenceMatcher


title = "<center><img src='https://huggingface.co/spaces/mideind/README/resolve/main/logo-mideind.png' width='250'></center>Sjálfvirk textaleiðrétting"
description = "Hér má leiðrétta íslenskan texta með hjálp tauganets. Netið reynir að „þýða“ textann yfir í læsilegra mál."
article = "<h3>Leiðbeiningar</h3>Hægt er að setja inn eina eða fleiri málsgreinar, en bíða þarf lengur eftir leiðréttingu ef textinn er langur. Ekki er gert ráð fyrir stökum málsgreinum sem eru lengri en u.þ.b. 500 stafir.  \n\nÞetta er prufuútgáfa sem hefur ekki lært að leiðrétta öll þau atriði sem upp geta komið í texta og er í stöðugri þróun.  \nVerkefnið er unnið hjá [Miðeind](https://mideind.is) sem hluti af máltækniáætlun stjórnvalda. Það er þjálfað á bæði tilbúnum gögnum og textum úr villumálheildum sem safnað var hjá Háskóla Íslands."

translator = pipeline(
    "translation", model="mideind/yfirlestur-icelandic-correction-byt5", max_length=512
)


def mark_text(text, tag):
    """Helper for the diff method, returns a tuple with the text and the tag"""
    return (text, tag)


def mark_span(text, tag):
    """Helper for the diff method, returns the span as a list of (text, tag) tuples"""
    return [mark_text(token, tag) for token in text]


def markup_diff(a, b, mark=mark_span, isjunk=None):
    """Obtains the diff for the sentence along with the opcodes (tags). Returns the corrected sentence tokens along with their correction tag (in Icelandic)"""
    seqmatcher = SequenceMatcher(isjunk=isjunk, a=a, b=b, autojunk=False)
    out_sentence_tokens = []
    # renaming tags
    tags = {"equal": None, "delete": "breytt", "replace": "breytt", "insert": "breytt"}

    for tag, a0, a1, b0, b1 in seqmatcher.get_opcodes():
        if tag == "delete":
            out_sentence_tokens += mark(" ", tags[tag])
        out_sentence_tokens += mark(b[b0:b1], tags[tag])
    return out_sentence_tokens


def split_text(text):
    sentence_list = [i for i in tokenizer.split_into_sentences(text, original=True)]
    return sentence_list


def predict(text):
    texts = text.split("\n")
    translated = ""
    original = ""
    # TODO: check for way too long inputs
    # TODO: if a sentence after splitting is over 512 bytes = split? error?
    for text in texts:
        # one or more sentences in each paragraph
        for sentence in split_text(text):
            original += sentence.strip() + " "
            translated += translator(sentence.strip())[0]["translation_text"] + " "
        original += "\n"
        translated += "\n"
    return markup_diff(original.split(" "), translated.split(" "))


demo = gr.Interface(
    fn=predict,
    inputs=gr.TextArea(label="Texti"),
    outputs=[
        gr.HighlightedText(
            label="Leiðrétt",
            show_label=False,
            show_legend=True,
            combine_adjacent=True,
            adjacent_separator=" ",
        ).style(color_map={"breytt": "green"})
    ],
    title=title,
    description=description,
    article=article,
    examples=[
        ["Kvitu fiðrildinn fljua firir utan gluggan."],
        ["Ég held þetta er ekki góður tími fara heimsókn."],
        ["Mer hlakar til jólana"],
        ["Kver a þenan bússtað ja eða nei"],
        ["Módernisma gætir í íslenkum prósaverkum fljótlega eftir Fyrri heimstyrjöld"],
    ],
)
demo.launch()