import gradio as gr
import pandas as pd
import random

USERS = ['user1', 'user2', 'user3']
df_per_user = {}
df = pd.read_csv("data.csv")
df['score'] = None
df['assignee'] = random.choices(USERS, k=len(df))
for u in USERS:
    df_per_user[u] = df[df.assignee == u].to_dict('records')

with gr.Blocks() as demo:
    dataset_df = {}
    state = gr.State(value=-1)
    with gr.Row():
        gr.Markdown("# Distributed Evaluation Parallel 😎")
    with gr.Row():
        prev = gr.Button(value="Previous")
        next = gr.Button(value="Next")
        download = gr.File(label="Download a file")
    with gr.Row():
        with gr.Column():
            question = gr.Textbox(label="Question")
        with gr.Column():
            ground_truth = gr.Textbox(label="GT")
        with gr.Column():
            prediction = gr.Textbox(label="Prediction")
            score = gr.Radio(choices=["Incorrect", "Correct"], label="Score")
    with gr.Row():
        with gr.Column():
            gr.Markdown("## TODO")
            todos = gr.DataFrame()
        with gr.Column():
            gr.Markdown("## DONE")
            done = gr.DataFrame()

    def prev_func(score, request: gr.Request):
        df_dict = df_per_user[request.username]
        state.value = max(state.value - 1, 0)
        score = df_dict[state.value]['score']
        gr.Info(f"{request.username}님, 총 {len(df_dict)}개 중에 {state.value + 1}번째 데이터입니다.")
        return [*update(request.username), score]
    
    def next_func(score, request: gr.Request):
        df_dict = df_per_user[request.username]
        df_dict[state.value]['score'] = score
        state.value = min(state.value + 1, len(df_dict) - 1)
        score = df_dict[state.value]['score']
        gr.Info(f"{request.username}님, 총 {len(df_dict)}개 중에 {state.value + 1}번째 데이터입니다.")
        return [*update(request.username), score]

    def update(username):
        df_dict = df_per_user[username]
        q = df_dict[state.value]['question']
        g = df_dict[state.value]['answer']
        p = df_dict[state.value]['prediction']
        df = pd.DataFrame(df_dict)
        todos = df[df.score.isna()]
        done = df[df.score.isna() == False]
        filename = f"done_{username}.csv"
        done.to_csv(filename, index=False)
        return q, g, p, todos, done, filename
    prev.click(prev_func, [score], [question, ground_truth, prediction, todos, done, download, score])
    next.click(next_func, [score], [question, ground_truth, prediction, todos, done, download, score])

demo.queue(concurrency_count=10)
demo.launch(auth=[(u, u) for u in USERS])