import gradio as gr
import nltk
from qanom.qanom_end_to_end_pipeline import QANomEndToEndPipeline
from typing import List

models = ["kleinay/qanom-seq2seq-model-baseline", 
          "kleinay/qanom-seq2seq-model-joint"]
pipelines = {model: QANomEndToEndPipeline(model) for model in models}


description = f"""This is a demo of the full QANom Pipeline - identifying deverbal nominalizations and parsing them with question-answer driven semantic role labeling (QASRL) """ 
title="QANom End-to-End Pipeline Demo"
examples = [[models[1], "It is also important to know that Optus’ network and Optus services including mobile and home Wi-Fi aren’t affected, and no passwords were compromised, so our services remain safe to use and operate as per normal.", 0.7],
            [models[1], "Importantly, no financial information or passwords have been accessed.", 0.75],
            [models[0], "The information which has been exposed is your name, date of birth, email, phone number, address associated with your account, and the numbers of the ID documents you provided such as drivers licence number or passport number", 0.75],
            [models[1], "It is with great disappointment I’m writing to let you know that Optus has been a victim of a cyberattack that has resulted in the disclosure of some of your personal information", 0.5]]
            

input_sent_box_label = "Insert sentence here, or select from the examples below"
links = """<p style='text-align: center'>
<a href='https://www.qasrl.org' target='_blank'>QASRL Website</a>  |  <a href='https://huggingface.co/kleinay/qanom-seq2seq-model-baseline' target='_blank'>Model Repo at Huggingface Hub</a>
</p>"""


def call(model_name, sentence, detection_threshold):

    pipeline = pipelines[model_name]
    pred_infos = pipeline([sentence], detection_threshold=detection_threshold)[0]
    def pretty_qas(pred_info) -> List[str]:
        if not pred_info or not pred_info['QAs']: return []
        return [f"{qa['question']} --- {';'.join(qa['answers'])}"
                for qa in pred_info['QAs'] if qa is not None]
    all_qas = [qa for pred_info in pred_infos for qa in pretty_qas(pred_info)]                      
    if not pred_infos:
        pretty_qa_output = "NO NOMINALIZATION FOUND"
    elif not all_qas:
        pretty_qa_output = "NO QA GENERATED"
    else:
        pretty_qa_output = "\n".join(all_qas) 
    # also present highlighted predicates
    positives = [pred_info['predicate_idx'] for pred_info in pred_infos] 
    def color(idx):
        if idx in positives: return "lightgreen"
    idx2verb = {d["predicate_idx"] : d["verb_form"] for d in pred_infos}
    idx2prob = {d["predicate_idx"] : d["predicate_detector_probability"] for d in pred_infos}
    def word_span(word, idx):
        tooltip = f'title=" probability={idx2prob[idx]:.2}&#010;verb={idx2verb[idx]}"' if idx in idx2verb else ''
        return f'<span {tooltip} style="background-color: {color(idx)}">{word}</span>'
    html = '<span>' + ' '.join(word_span(word, idx) for idx, word in enumerate(sentence.split(" "))) + '</span>'
    return html, pretty_qa_output , pred_infos

iface = gr.Interface(fn=call, 
                     inputs=[gr.inputs.Radio(choices=models, default=models[0], label="Model"), 
                             gr.inputs.Textbox(placeholder=input_sent_box_label, label="Sentence", lines=4), 
                             gr.inputs.Slider(minimum=0., maximum=1., step=0.01, default=0.5, label="Nominalization Detection Threshold")], 
                     outputs=[gr.outputs.HTML(label="Detected Nominalizations"),
                              gr.outputs.Textbox(label="Generated QAs"), 
                              gr.outputs.JSON(label="Raw Model Output")],
                     title=title,
                     description=description,
                     article=links,
                     examples=examples)
iface.launch()