import gradio as gr
import spacy
import medspacy
from medspacy.visualization import visualize_dep, visualize_ent
from spacy import displacy

med_ner = medspacy.load(r"./model-best")
def merge_tokens(tokens):
    merged_tokens = []
    for token in tokens:
        if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]):
             # If current token continues the entity of the last one, merge them
            last_token = merged_tokens[-1]
            last_token['word'] += token['word'].replace('##', '')
            last_token['end'] = token['end']
#             last_token['score'] = (last_token['score'] + token['score']) / 2
        else:
            # Otherwise, add the token to the list
            merged_tokens.append(token)
            
    return merged_tokens

def ner(inp):
    output = med_ner(inp)
    formatted_ents = []
    for i in output.ents:
        ent = {}
        ent['entity']= i.label_
        ent['word']= i.text
        ent['start']= int(i.start_char)
        ent['end']= int(i.end_char)
        print(i.label_,"->",i.text,"->",i.start_char,"->",i.end_char,"->",type(i.start_char))
        formatted_ents.append(ent)
    print(formatted_ents)
    
    merged_tokens = merge_tokens(formatted_ents)
#     return {"text": str(inp), "entities": formatted_ents}
    return {"text": str(inp), "entities": merged_tokens}

demo = gr.Interface(fn=ner,
                   inputs=[gr.Textbox(label="Text to find entities", lines=2)],
                   outputs=[gr.HighlightedText(label="Text with entites")],
                   title="Custom-NER with Spacy3 and MedSpacy with v2 model",
                   description="Find medical entities using the NER model under the hood!",
                   allow_flagging = True,
                   examples=["Patient has hx of stroke. Mother diagnosed with diabetes. No evidence of pna.", "I have fever and cough since 2 days."]
                   )

demo.launch()