import gradio as gr
import os

import requests

from spacy import displacy

os.system("python -m spacy download en_core_web_md")
import spacy


colors = {
    "Observation": "#9bddff",
    "Evaluation": "#f08080",
}

nlp = spacy.load("en_core_web_md") #Esto es para usar displacy y renderizar las entidades
nlp.disable_pipes("ner")

def compute_ner(input_text_message):
    endpoint_url = 'https://on1m82uknekghqeh.us-east-1.aws.endpoints.huggingface.cloud'

    headers = {
        'Authorization': 'Bearer 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',
        'Content-Type': 'application/json',
    }

    json_data = {
        'inputs': input_text_message,
    }

    response = requests.post(endpoint_url, headers=headers, json=json_data)

    entities = response.json()

    doc = nlp(input_text_message)

    potential_entities = []

    for entity in entities:
        start = entity["start"]
        end = entity["end"]
        label = entity["entity"]

        if label == "I-Observation" or label == "B-Observation":
            label = "Observation"

        if label == "I-Evaluation" or label == "B-Evaluation":
            label = "Evaluation"

        entity["entity"]=label

        ent = doc.char_span(start, end, label=label)
        if ent != None:
            doc.ents += (ent,)
        else:
            potential_entities.append(entity)

    potential_entities.append({"entity": "NONE", "start": -1, "end": -1})

    start = potential_entities[0]["start"]
    end = potential_entities[0]["end"]
    label = potential_entities[0]["entity"]

    for item in potential_entities:
        if item["entity"] == label and item["start"] == end:
            end = item["end"]
            continue
        else:
            if item["start"] != start:
                ent = doc.char_span(start, end, label=label)
                doc.ents += (ent,)

            start = item["start"]
            end = item["end"]
            label = item["entity"]

    options = {"ents": colors.keys(), "colors": colors}

    return displacy.render(doc, style="ent", options=options)


examples = ['You are dick',
            'My dad is an asshole and took his anger out on my mom by verbally abusing her',
            'He eventually moved on to my brother']

iface = gr.Interface(fn=compute_ner, inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your text here",
                                                              label='Check your text for compliance with the NVC rules'),
                     outputs="html", examples=examples)
iface.launch()