import gradio as gr
from description import *

from reference_string_parsing import *
from summarization import *
from controlled_summarization import *

with gr.Blocks(css="#htext span {white-space: pre-line}") as demo:
    gr.Markdown("# Gradio Demo for SciAssist")
    with gr.Tabs():
        # Reference String Parsing
        with gr.TabItem("Reference String Parsing"):
            with gr.Box():
                gr.Markdown(rsp_str_md)
                with gr.Row():
                    with gr.Column():
                        rsp_str = gr.Textbox(label="Input String")
                        with gr.Column():
                            rsp_str_dehyphen = gr.Checkbox(label="dehyphen")
                        with gr.Row():
                            rsp_str_btn = gr.Button("Parse")
                    rsp_str_output = gr.HighlightedText(
                        elem_id="htext",
                        label="The Result of Parsing",
                        combine_adjacent=True,
                        adjacent_separator=" ",
                    )
                rsp_str_examples = gr.Examples(examples=[[
                                                         "Waleed Ammar, Matthew E. Peters, Chandra Bhagavat- ula, and Russell Power. 2017. The ai2 system at semeval-2017 task 10 (scienceie): semi-supervised end-to-end entity and relation extraction. In ACL workshop (SemEval).",
                                                         True],
                                                     [
                                                         "Isabelle Augenstein, Mrinal Das, Sebastian Riedel, Lakshmi Vikraman, and Andrew D. McCallum. 2017. Semeval-2017 task 10 (scienceie): Extracting keyphrases and relations from scientific publications. In ACL workshop (SemEval).",
                                                         False]], inputs=[rsp_str, rsp_str_dehyphen])
            with gr.Box():
                gr.Markdown(rsp_file_md)
                with gr.Row():
                    with gr.Column():
                        rsp_file = gr.File(label="Input File")
                        rsp_file_dehyphen = gr.Checkbox(label="dehyphen")
                        with gr.Row():
                            rsp_file_btn = gr.Button("Parse")

                    rsp_file_output = gr.HighlightedText(
                        elem_id="htext",
                        label="The Result of Parsing",
                        combine_adjacent=True,
                        adjacent_separator=" ",
                    )
                rsp_file_examples = gr.Examples(examples=[["examples/N18-3011_ref.txt", False],["examples/BERT_paper.pdf", True]], inputs=[rsp_file, rsp_file_dehyphen])


        rsp_file_btn.click(
            fn=rsp_for_file,
            inputs=[rsp_file, rsp_file_dehyphen],
            outputs=rsp_file_output
        )
        rsp_str_btn.click(
            fn=rsp_for_str,
            inputs=[rsp_str, rsp_str_dehyphen],
            outputs=rsp_str_output
        )

        # Single Document Summarization
        with gr.TabItem("Summarization"):
            with gr.Box():
                gr.Markdown(ssum_str_md)
                with gr.Row():
                    with gr.Column():
                        ssum_str = gr.Textbox(label="Input String")
                        # with gr.Column():
                        #     ssum_str_beams = gr.Number(label="Number of beams for beam search", value=1, precision=0)
                        #     ssum_str_sequences = gr.Number(label="Number of generated summaries", value=1, precision=0)
                        with gr.Row():
                            ssum_str_btn = gr.Button("Generate")
                    ssum_str_output = gr.Textbox(
                        elem_id="htext",
                        label="Summary",
                    )
                ssum_str_examples = gr.Examples(examples=[[ssum_str_example], ],
                                                inputs=[ssum_str])
            with gr.Box():
                gr.Markdown(ssum_file_md)
                with gr.Row():
                    with gr.Column():
                        ssum_file = gr.File(label="Input File")
                        # with gr.Column():
                            # ssum_file_beams = gr.Number(label="Number of beams for beam search", value=1, precision=0)
                            # ssum_file_sequences = gr.Number(label="Number of generated summaries", value=1, precision=0)
                        with gr.Row():
                            ssum_file_btn = gr.Button("Generate")
                    ssum_file_output = gr.Textbox(
                        elem_id="htext",
                        label="Summary",
                    )
                ssum_file_examples = gr.Examples(examples=[["examples/BERT_body.txt"],["examples/BERT_paper.pdf"]],
                                                inputs=[ssum_file])

        ssum_file_btn.click(
            fn=ssum_for_file,
            inputs=[ssum_file],
            outputs=ssum_file_output
        )
        ssum_str_btn.click(
            fn=ssum_for_str,
            inputs=[ssum_str],
            outputs=ssum_str_output
        )

        # Controlled Summarization
        with gr.TabItem("Controlled Summarization"):
            with gr.Box():
                gr.Markdown(ctrlsum_str_md)
                with gr.Row():
                    with gr.Column():
                        ctrlsum_str = gr.Textbox(label="Input String")
                        with gr.Column():
                            # ctrlsum_str_beams = gr.Number(label="Number of beams for beam search", value=1, precision=0)
                            # ctrlsum_str_sequences = gr.Number(label="Number of generated summaries", value=1, precision=0)
                            ctrlsum_str_length = gr.Slider(0, 300, step=50, label="Length")
                            ctrlsum_str_keywords = gr.Textbox(label="Keywords")
                        with gr.Row():
                            ctrlsum_str_btn = gr.Button("Generate")
                    ctrlsum_str_output = gr.Textbox(
                        elem_id="htext",
                        label="Summary",
                    )
                ctrlsum_str_examples = gr.Examples(examples=[[ssum_str_example, 50, "BERT" ], ],
                                                inputs=[ctrlsum_str, ctrlsum_str_length, ctrlsum_str_keywords])
            with gr.Box():
                gr.Markdown(ctrlsum_file_md)
                with gr.Row():
                    with gr.Column():
                        ctrlsum_file = gr.File(label="Input File")
                        with gr.Column():
                            # ctrlsum_file_beams = gr.Number(label="Number of beams for beam search", value=1, precision=0)
                            # ctrlsum_file_sequences = gr.Number(label="Number of generated summaries", value=1, precision=0)
                            ctrlsum_file_length = gr.Slider(0,300,step=50, label="Length")
                            ctrlsum_file_keywords = gr.Textbox(label="Keywords")
                        with gr.Row():
                            ctrlsum_file_btn = gr.Button("Generate")
                    ctrlsum_file_output = gr.Textbox(
                        elem_id="htext",
                        label="Summary",
                    )
                ctrlsum_file_examples = gr.Examples(examples=[["examples/BERT_body.txt", 100, ""],["examples/BERT_paper.pdf", 0, "BERT"]],
                                                inputs=[ctrlsum_file, ctrlsum_file_length, ctrlsum_file_keywords])

        ctrlsum_file_btn.click(
            fn=ctrlsum_for_file,
            inputs=[ctrlsum_file, ctrlsum_file_length, ctrlsum_file_keywords],
            outputs=ctrlsum_file_output
        )
        ctrlsum_str_btn.click(
            fn=ctrlsum_for_str,
            inputs=[ctrlsum_str, ctrlsum_str_length, ctrlsum_str_keywords],
            outputs=ctrlsum_str_output
        )



demo.launch(share=False)