from __future__ import annotations

import gradio as gr

from pipeline import labels2attrs, run_pipeline

color_panel_1 = [
    "red",
    "green",
    "yellow",
    "DodgerBlue",
    "orange",
    "DarkSalmon",
    "pink",
    "cyan",
    "gold",
    "aqua",
    "violet",
]
index_colormap = {str(i): color_panel_1[i % len(color_panel_1)] for i in range(1, 100000)}

color_panel_2 = [
    "Gray",
    "DodgerBlue",
    "Wheat",
    "OliveDrab",
    "DarkKhaki",
    "DarkSalmon",
    "Orange",
    "Gold",
    "Aqua",
    "Tomato",
    "Violet",
]

str_attrs = sorted([str(v) for v in set(labels2attrs.values())])
attr_colormap = {attr: color for attr, color in zip(str_attrs, color_panel_2)}

DESCRIPTION = """
Type any text in the textbox and press Submit. The app segments the text into clauses,
classifies the clause-level discourse attributes, and plots the attribute distributions.
"""

EXAMPLES = [
    [
        "My cousin was arrested for marijuana possession, and it ruined his life for years."
    ],
    [
        "Legalizing marijuana would reduce incarceration costs and improve tax revenue."
    ],
]

demo = gr.Interface(
    fn=run_pipeline,
    inputs=gr.Textbox(
        label="Input Text",
        lines=8,
        placeholder="Paste text here...",
    ),
    outputs=[
        gr.HighlightedText(
            label="Clause Segmentation",
            color_map=index_colormap,
            combine_adjacent=False,
            show_legend=False,
        ),
        gr.HighlightedText(
            label="Attribute Classification",
            color_map=attr_colormap,
            combine_adjacent=False,
            show_legend=True,
        ),
        gr.Plot(label="Proportion of Attributes"),
    ],
    title="Anecdotal Discourse Classification Demo",
    description=DESCRIPTION,
    examples=EXAMPLES,
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(share=True)