import streamlit as st
from annotated_text import annotated_text

import torch
from transformers import pipeline
from transformers import AutoModelForTokenClassification, AutoTokenizer

import json


st.set_page_config(layout="wide")



model = AutoModelForTokenClassification.from_pretrained("models/lusa")
tokenizer = AutoTokenizer.from_pretrained("models/lusa", model_max_length=512)
tagger = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy='first') #aggregation_strategy='max'
  


def aggregate_subwords(input_tokens, labels):
    new_inputs = []
    new_labels = []
    current_word = ""
    current_label = ""
    for i, token in enumerate(input_tokens):     
        label = labels[i]
        # Handle subwords
        if token.startswith('##'):
            current_word += token[2:]
        else:
            # Finish previous word
            if current_word:
                new_inputs.append(current_word)
                new_labels.append(current_label)
            # Start new word
            current_word = token
            current_label = label
    new_inputs.append(current_word)
    new_labels.append(current_label)
    return new_inputs, new_labels

def annotateTriggers(line):
    line = line.strip()
    inputs = tokenizer(line, return_tensors="pt")
    input_tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

    with torch.no_grad():
        logits = model(**inputs).logits

    predictions = torch.argmax(logits, dim=2)
    predicted_token_class = [model.config.id2label[t.item()] for t in predictions[0]]
    input_tokens, predicted_token_class =  aggregate_subwords(input_tokens,predicted_token_class)
    token_labels = []
    current_entity = ''
    for i, label in enumerate(predicted_token_class):
        token = input_tokens[i]
        if label == 'O':
            token_labels.append((token, 'O', ''))
            current_entity = ''
        elif label.startswith('B-'):
            current_entity = label[2:]
            token_labels.append((token, 'B', current_entity))
        elif label.startswith('I-'):
            if current_entity == '':
                raise ValueError(f"Invalid label sequence: {predicted_token_class}")
            token_labels[-1] = (token_labels[-1][0] + f" {token}", 'I', current_entity)
        else:
            raise ValueError(f"Invalid label: {label}")
    return token_labels[1:-1]





def joinEntities(entities):

    joined_entities = []
    i = 0
    while i < len(entities):
        curr_entity = entities[i]
        if curr_entity['entity'][0] == 'B':
            label = curr_entity['entity'][2:]
            j = i + 1
            while j < len(entities) and entities[j]['entity'][0] == 'I':
                j += 1
            joined_entity = {
                 'entity': label,
                'score': max(e['score'] for e in entities[i:j]),
                'index': min(e['index'] for e in entities[i:j]),
                'word': ' '.join(e['word'] for e in entities[i:j]),
                'start': entities[i]['start'],
                'end': entities[j-1]['end']
            }
            joined_entities.append(joined_entity)
            i = j - 1
        i += 1
    return joined_entities



import pysbd
seg = pysbd.Segmenter(language="es", clean=False)

def sent_tokenize(text):
    return seg.segment(text)

def getSentenceIndex(lines,span):
    i = 1
    sum = len(lines[0])
    while sum < span:
        sum += len(lines[i])
        i = i + 1
    return i - 1

def generateContext(text, window,span):
    lines = sent_tokenize(text)
    index = getSentenceIndex(lines,span)
    text = " ".join(lines[max(0,index-window):index+window +1])
    return text


def annotateEvents(text,squad,window):
    text = text.strip()
    ner_results = tagger(text)
    #print(ner_results)
    #ner_results = joinEntities(ner_results)
    i = 0
    #exit()
    while i < len(ner_results):
        ner_results[i]["entity"] = ner_results[i]["entity_group"].lstrip("B-")
        ner_results[i]["entity"] = ner_results[i]["entity_group"].lstrip("I-")
        i = i + 1

    events = [] 
    for trigger in ner_results:
        tipo = trigger["entity_group"]
        context = generateContext(text,window,trigger["start"])
        event = {
            "trigger":trigger["word"],
            "type": tipo,
            "score": trigger["score"],
            "context": context,
        }
        events.append(event)
    return events


#"A Joana foi atacada pelo João nas ruas do Porto, com uma faca."

st.title('Extract Events')

options = ["O presidente da Federação Haitiana de Futebol, Yves Jean-Bart, foi banido para sempre de toda a atividade ligada ao futebol, por ter sido considerado culpado de abuso sexual sistemático de jogadoras, anunciou hoje a FIFA."]

option = st.selectbox(
        'Select examples',
        options)
#option = options [index]
line = st.text_area("Insert Text",option)

st.button('Run')


st.sidebar.write("## Hyperparameters :gear:")
window = 1
if line != "":
    st.header("Triggers:")
    triggerss = annotateTriggers(line)   
    annotated_text(*[word[0]+" " if word[1] == 'O' else (word[0]+" ",word[2]) for word in triggerss ])

    eventos_1 = annotateEvents(line,1,window)
    eventos_2 = annotateEvents(line,2,window)

    for mention1, mention2 in zip(eventos_1,eventos_2):
        st.text(f"| Trigger: {mention1['trigger']:20} | Type: {mention1['type']:10} | Score: {str(round(mention1['score'],3)):5} |") 
        st.markdown("""---""")