import gradio as gr
from gradio.components import Label, Textbox

from transformers import pipeline
from utils import *
from datasets import load_dataset
import json

pipe = pipeline(model="raminass/british", top_k=2, padding=True, truncation=True)
df = pd.read_csv("data.csv", sep="\t")
choices = []
for index, row in df.iterrows():
    choices.append((f"""{row["case"]}""", [row["text"], row["author"]]))


# https://www.gradio.app/guides/controlling-layout
def greet(opinion):
    opinion = opinion.replace("(", "").replace(")", "")
    chunks = chunk_data(opinion)["text"].to_list()
    result = average_text(chunks, pipe)

    return result[0]


def set_input(drop):
    return drop[0], drop[1], gr.Slider(visible=True)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=2):
            drop = gr.Dropdown(
                choices=sorted(choices),
                label="List of Cases",
                info="Select a case from the dropdown menu and press the Predict Button",
            )

            opinion = gr.Textbox(
                label="Opinion",
                info="Paste opinion text here and press the Predict Button",
            )
        with gr.Column(scale=1):
            with gr.Row():
                clear_btn = gr.Button("Clear")
                greet_btn = gr.Button("Predict")
            op_level = Label(num_top_classes=9, label="Predicted author of opinion")

    drop.select(set_input, inputs=drop, outputs=[opinion])

    greet_btn.click(
        fn=greet,
        inputs=[opinion],
        outputs=[op_level],
    )

    clear_btn.click(
        fn=lambda: [None, 1994, gr.Slider(visible=True), None, None],
        outputs=[opinion, drop, op_level],
    )


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