# -*- coding: utf-8 -*-
import os

os.system("pip3 install torch==1.10.1+cpu torchvision==0.11.2+cpu torchaudio==0.10.1+cpu -f "
          "https://download.pytorch.org/whl/cpu/torch_stable.html")

import gradio as gr
from transformers import pipeline

import spacy
from spacy import displacy

POS_TAG_MAP = {"NOUN": "PODMET", "VERB": "PRÍSUDOK", "ADJ": "PRÍDAVNÉ MENO", "ADP": "ADPOZÍCIA",
               "ADV": "PRÍSLOVKA", "AUX": "POMOCNÉ", "CCONJ": "KORDINAČNÁ KONJUKCIA", "DET": "DETERMINANT",
               "INTJ": "CYTOSLOVCIA", "NUM": "NUMERICKÉ", "PART": "ČASTICA", "PRON": "ZÁMMENO",
               "PROPN": "VLASTNÉ MENO", "PUNCT": "INTERPUNKCIA", "SCONJ": "", "SYM": "SYMBÓL", "X": "INÉ"}

options = {"ents": list(POS_TAG_MAP.values()),
           "colors": {"PODMET": "lightblue", "PRÍSUDOK": "lightcoral",
                      "PRÍDAVNÉ MENO": "lightgreen", "VLASTNÉ MENO": "papayawhip"}}

pipe = pipeline(task='token-classification', model="crabz/slovakbert-ner")
nlp = spacy.blank("sk")


def exists_double_class(classified_ents):
    for i in range(len(classified_ents) - 1):
        if classified_ents[i][2] == classified_ents[i + 1][1]:
            return i
    return -1


def postprocess(classifications):
    entities = []
    for i in range(len(classifications)):
        entities.append((POS_TAG_MAP[classifications[i]['entity']], classifications[i]['start'],
                         classifications[i]['end']))

    while True:
        i = exists_double_class(entities)
        if i == -1:
            break

        entities.append((entities[i][0], entities[i][1], entities[i + 1][2]))
        del entities[i]
        del entities[i]
    return entities


def set_entities(sentence, entities):
    doc = nlp(sentence)
    ents = []
    for label, start, end in entities:
        ents.append(doc.char_span(start, end, label))
    try:
        doc.ents = ents
    except TypeError:
        pass
    return doc


def apply_pos(Sentence: str):
    classifications = pipe(Sentence.replace(",", " ,")
                           .replace(".", " .")
                           .replace("!", " !")
                           .replace("?", " ?")
                           .replace(":", " :"))
    entities = postprocess(classifications)
    doc = set_entities(Sentence, entities)
    displacy_html = displacy.render(doc, style="ent", options=options)
    return displacy_html


intf = gr.Interface(fn=apply_pos, inputs="text", outputs="html", title='Slovak Part-Of-Speech Tagging',
                    allow_flagging=False,
                    examples=[["Laboratóriá Úradu verejného zdravotníctva sekvenovaním potvrdili výskyt ďalších štyroch prípadov variantu omikron na Slovensku."],
                              ["Čaputová opakovane tvrdí, že \"spravodlivosť na Slovensku neplatí vždy pre všetkých rovnako\"."],
                              ["Informácie o týchto veľkolepých plánoch prišli týždeň po tom, ako sa japonský miliardár Jusaku Maezawa vrátil z 12-dňového pobytu na Medzinárodnej vesmírnej stanici (ISS), čím sa stal prvým vesmírnym turistom, ktorý cestoval na ISS za viac ako desať rokov."],
                              ["Minister financií a líder mandátovo najsilnejšieho hnutia OĽaNO Igor Matovič upozorňuje, že následky tretej vlny budú na Slovensku veľmi veľké."],
                              ["Začiatkom roka 2021 sa objavili nezhody medzi Richardom Sulíkom a šéfom hnutia OĽANO Igorom Matovičom, ktoré v istej miere pretrvávajú aj dodnes."]],
                    description="",
                    article="",
                    analytics_enabled=False)
intf.launch()