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

title = "Textaleiðrétting fyrir íslensku"
description = "Hér má leiðrétta íslenskan texta með hjálp tauganets. Netið reynir að „þýða“ yfir í texta sem er í samræmi við málstaðal, og er þjálfað á bæði tilbúnum gögnum og textum úr villumálheildum sem safnað var hjá Háskóla Íslands.  \n\nBest er að setja ekki inn meira en nokkrar setningar í einu, því annars getur biðin orðið löng. Þ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.  \nÞetta verkefni er unnið hjá Miðeind sem hluti af máltækniáætlun stjórnvalda."

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": "eytt út", "replace": "skipt út", "insert": "bætt inn"}
    
    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='text',
    outputs= gr.HighlightedText(
        show_label=False,
        show_legend=True,
        combine_adjacent=True,
        adjacent_separator=" ",
    ).style(
        color_map={"skipt út": "blue", "bætt inn": "green", "eytt út": "purple"}),
    title=title,
    description=description,
    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"]],
)
demo.launch()