from utils import (
    SentenceSimilarity,
    pos_tagging,
    text_analysis,
    text_interface,
    sentence_similarity,
)
from script import details
from transformers import pipeline
import gradio as gr
from functools import partial

pipes = {
    "Sentiment Analysis": pipeline(
        "text-classification",
        model="StevenLimcorn/indonesian-roberta-base-emotion-classifier",
        tokenizer="StevenLimcorn/indonesian-roberta-base-emotion-classifier",
    ),
    "Emotion Classifier": pipeline(
        "text-classification",
        model="w11wo/indonesian-roberta-base-sentiment-classifier",
        tokenizer="w11wo/indonesian-roberta-base-sentiment-classifier",
    ),
    "summarization": pipeline(
        "summarization",
        model="LazarusNLP/IndoNanoT5-base-IndoSum",
        tokenizer="LazarusNLP/IndoNanoT5-base-IndoSum",
    ),
    "sentence-similarity": SentenceSimilarity(model="LazarusNLP/all-indobert-base-v2"),
    "POS Tagging": pipeline(model="w11wo/indonesian-roberta-base-posp-tagger"),
}

if __name__ == "__main__":
    # list of collections of all demos
    classifiers = ["Sentiment Analysis", "Emotion Classifier"]
    # Summary
    summary_interface = gr.Interface.from_pipeline(
        pipes["summarization"],
        title="Summarization",
        examples=details["summarization"]["examples"],
        description=details["summarization"]["description"],
        allow_flagging="never",
    )
    # Pos Tagging
    pos_interface = gr.Interface(
        fn=partial(pos_tagging, pipe=pipes["POS Tagging"]),
        inputs=[
            gr.Textbox(placeholder="Masukan kalimat di sini...", label="Input Text"),
        ],
        outputs=[gr.HighlightedText()],
        title="POS Tagging",
        examples=details["POS Tagging"]["examples"],
        description=details["POS Tagging"]["description"],
        allow_flagging="never",
    )
    # Text Analysis
    with gr.Blocks() as text_analysis_interface:
        gr.Markdown("# Text Analysis")
        gr.Markdown(details["Text Analysis"]["description"])
        input_text = gr.Textbox(lines=5, label="Input Text")
        with gr.Row():
            smsa = gr.Label(label="Sentiment Analysis")
            emot = gr.Label(label="Emotion Classification")
            pos = gr.HighlightedText(label="POS Tagging")
        btn = gr.Button("Analyze")
        btn.click(
            fn=partial(text_analysis, pipes=pipes),
            inputs=[input_text],
            outputs=[smsa, emot, pos],
        )
        gr.Examples(
            details["Text Analysis"]["examples"],
            inputs=input_text,
            outputs=[smsa, emot, pos],
        )

    with gr.Blocks() as sentence_similarity_interface:
        gr.Markdown("# Document Search 🔍")
        gr.Markdown(details["sentence-similarity"]["description"])
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(lines=5, label="Query")
                file_input = gr.File(
                    label="Documents", file_types=[".txt"], file_count="multiple"
                )
                button = gr.Button("Search...")
            output = gr.Label()
        button.click(
            fn=partial(sentence_similarity, pipe=pipes["sentence-similarity"]),
            inputs=[input_text, file_input],
            outputs=[output],
        )

    demo_interface = {
        "demo": [
            text_interface(
                pipes[name],
                details[name]["examples"],
                name,
                name,
                details[name]["description"],
            )
            for name in classifiers
        ]
        + [
            sentence_similarity_interface,
            summary_interface,
            pos_interface,
            text_analysis_interface,
        ],
        "titles": classifiers
        + ["Document Search", "Summarization", "POS Tagging", "Text Analysis"],
    }

    # with gr.Blocks() as demo:
    #     with gr.Column():
    #         gr.Markdown("# Title")
    #         gr.TabbedInterface(
    #             demo_interface["demo"], demo_interface["titles"], theme="soft"
    #         )

    demo = gr.TabbedInterface(
        demo_interface["demo"], demo_interface["titles"], theme="soft"
    )
    demo.launch()