import gradio as gr
import os

import requests

from spacy import displacy
import streamlit as st

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

# entities group method
# https://huggingface.co/spaces/crabz/sk-ner/blob/main/app.py

options = {"ents": ["Observation",
                    "Evaluation"],
           "colors": {
               "Observation": "#9bddff",
               "Evaluation": "#f08080",
           }

           }

nlp = spacy.load("en_core_web_md")


def postprocess(classifications):
    entities = []
    for i in range(len(classifications)):
        if classifications[i]['entity'] != 0:
            if classifications[i]['entity'][0] == 'B':
                j = i + 1
                while j < len(classifications) and classifications[j]['entity'][0] == 'I':
                    j += 1
                entities.append((classifications[i]['entity'].split('-')[1], classifications[i]['start'],
                                 classifications[j - 1]['end']))
    while True:
        merged = False
        to_remove = []
        merged_entities = []
        for i in range(len(entities)):
            for j in range(i + 1, len(entities)):
                if entities[i] != entities[j] and entities[i][0] == entities[j][0] and \
                        (entities[i][2] == entities[j][1] or entities[i][1] == entities[j][2]):
                    to_remove.append(entities[i])
                    to_remove.append(entities[j])

                    new_start = min(entities[i][1], entities[j][1])
                    new_end = max(entities[i][2], entities[j][2])
                    merged_entities.append((entities[i][0], new_start, new_end))
                    merged = True
                    break
            if merged:
                break
        for ent in to_remove:
            entities.remove(ent)
        entities += merged_entities
        if not merged:
            break
    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_ner(input_text_message: str):
    # auth_endpoint_token = st.secrets["auth_endpoint_token"]
    auth_endpoint_token = os.environ["auth_endpoint_token"]
    # endpoint_url = st.secrets["endpoint_url"]
    endpoint_url = os.environ["endpoint_url"]

    headers = {
        'Authorization': auth_endpoint_token,
        'Content-Type': 'application/json',
    }

    json_data = {
        'inputs': input_text_message,
    }

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

    classifications = response.json()
    entities = postprocess(classifications)
    doc = set_entities(input_text_message, entities)
    displacy_html = displacy.render(doc, style="ent", options=options)
    return displacy_html


examples = ['Things are complicated because we still live together but we have separate lives',
            'My dad is a monster and took his anger out on my mom by verbally abusing her and when she left he eventually moved on to my brother',
            'A two months ago, she was chatting with some random guy',
            'Not I have a horrid relationship with my brother we’ve never gotten along and probably never will',
            'I was outside trying to leave and he caught me to explain why Im so rude',
            ]

iface = gr.Interface(fn=apply_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()