import gradio as gr

from pathlib import Path
from supar import Parser
from spacy import displacy
from spacy.tokens import Doc, Span
import spacy

proj_dir = Path(__file__).parent
model_choices = [str(model.name) for model in (proj_dir / 'models').glob('*')][::-1]


def sentence_diagram(model_name, text, progress=gr.Progress(track_tqdm=True)):
    parser = Parser.load(f'./models/{model_name}')

    Span.set_extension("con_tree", getter=lambda x: parser.predict([i.text for i in x], verbose=False)[0], force=True)
    nlp = spacy.load('en_core_web_sm')
    doc = nlp(text)

    options = {"compact": False, "color": "Red", 'collapse_punct': True, 'collapse_phrases': False,
               'split_sentences': True}
    html = displacy.render(doc, style="dep", options=options, page=True)
    return html


with gr.Blocks() as demo:
    with gr.Tab("Sentence Diagrams"):
        model_name = gr.Dropdown(choices=model_choices, label='Model Name')
        text_in = gr.Textbox(label='Sentence(s) to diagram', value='This is a test')
        button = gr.Button('Run!')
        html_out = gr.HTML()
    with gr.Tab("Brother Hill Tribute"):
        gr.Markdown("""To Bro Hill""")

    button.click(sentence_diagram,
                 inputs=[model_name, text_in],
                 outputs=html_out)

if __name__ == '__main__':
    demo.queue().launch(show_error=True)