import gradio as gr
from typing import List

from qasem.end_to_end_pipeline import QASemEndToEndPipeline
pipeline = QASemEndToEndPipeline()


description = f"""This is a demo of the QASem Parsing pipeline. It wraps models of three QA-based semantic tasks, composing a comprehensive semi-structured representation of sentence meaning - covering verbal and nominal semantic role labeling together with discourse relations.""" 
title="QASem Parsing Demo"

all_layers = ["qasrl", "qanom", "qadiscourse"]
examples = [["Both were shot in the confrontation with police and have been recovering in hospital since the attack .", all_layers, 0.75],
            ["the construction of the officer 's building was delayed by the lockdown and is expected to continue for at least 10 more months.", all_layers, 0.75],
            ["While President Obama expressed condolences regarding the death of Margaret Thatcher upon her death earlier this year , he did not issue an executive order that flags be lowered in her honor .", all_layers, 0.75],
            ["We made a very clear commitment : if there is any proposal in the next parliament for a transfer of powers to Brussels ( the EU ) we will have an in/out referendum .", all_layers, 0.75],
            ["The doctor asked about the progress in Luke 's treatment .", all_layers, 0.75],
            ["The Veterinary student was interested in Luke 's treatment of sea animals .", all_layers, 0.7],
            ["Some reviewers agreed that the criticism raised by the AC is mostly justified .", all_layers, 0.6]]
            

input_sent_box_label = "Insert sentence here, or select from the examples below"
links = """<p style='text-align: center'>
<a href='https://github.com/kleinay/QASem' target='_blank'>Github Repo</a> | <a href='https://arxiv.org/abs/2205.11413' target='_blank'>Paper</a>
</p>"""


def call(sentence, layers, detection_threshold):

    outputs = pipeline([sentence], nominalization_detection_threshold=detection_threshold)[0]
    def pretty_qadisc_qas(qa_infos) -> List[str]:
        if not qa_infos: return []
        return ["- " + f"{qa['question']} --- {qa['answer']}".lstrip()
                for qa in qa_infos if qa is not None]
    def pretty_qasrl_qas(pred_info) -> List[str]:
        if not pred_info or not pred_info['QAs']: return []
        return ["- " + f"{qa['question']} --- {';'.join(qa['answers'])}".lstrip()
                for qa in pred_info['QAs'] if qa is not None]
    # filter outputs by requested `layers`
    outputs = {layer: qas if layer in layers else [] 
               for layer, qas in outputs.items()} 
    # Prettify outputs
    qasrl_qas = [qa for pred_info in outputs['qasrl'] for qa in pretty_qasrl_qas(pred_info)]
    qanom_qas = [qa for pred_info in outputs['qanom'] for qa in pretty_qasrl_qas(pred_info)]
    qadisc_qas= pretty_qadisc_qas(outputs['qadiscourse'])
    all_qas = []
    if "qasrl" in layers: all_qas += ['\nQASRL:'] + qasrl_qas
    if "qanom" in layers: all_qas += ['\nQANom:'] + qanom_qas
    if "qadiscourse" in layers: all_qas += ['\nQADiscourse:'] + qadisc_qas

    if not qasrl_qas + qanom_qas + qadisc_qas:
        pretty_qa_output = "NO QA GENERATED"
    else:
        pretty_qa_output = "\n".join(all_qas) 
    
    # also present highlighted predicates
    qasrl_predicates = [pred_info['predicate_idx'] for pred_info in outputs['qasrl']] 
    qanom_predicates = [pred_info['predicate_idx'] for pred_info in outputs['qanom']] 
    def color(idx):
        if idx in qasrl_predicates : return "aquamarine"
        if idx in qanom_predicates : return "aqua"
    def word_span(word, idx):
        return f'<span 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 , outputs

iface = gr.Interface(fn=call, 
                     inputs=[gr.inputs.Textbox(placeholder=input_sent_box_label, label="Sentence", lines=4), 
                             gr.inputs.CheckboxGroup(all_layers, default=all_layers, label="Annotation Layers"),
                             gr.inputs.Slider(minimum=0., maximum=1., step=0.01, default=0.75, label="Nominalization Detection Threshold")], 
                     outputs=[gr.outputs.HTML(label="Detected Predicates"),
                              gr.outputs.Textbox(label="Generated QAs"), 
                              gr.outputs.JSON(label="Raw QASemEndToEndPipeline Output")],
                     title=title,
                     description=description,
                     article=links,
                     examples=examples)
iface.launch()