import gradio as gr
import pandas as pd
import numpy as np

# Global variable to store the DataFrame
df = None
# Global variable to keep track of the current row index
current_row = 0

def load_csv(file):
    global df
    global current_row
    # import the csv and set the data types to be int, string, string, string, string, string, string
    df = pd.read_csv(file.name, dtype={'id':int, 'hs': str, 'cs': str, 'topic': str, 'tone': str, 'isCSContextuallyRelevant': str, 'isToneMatch': str})
    if 'suggestedTone' not in df.columns:
        df['suggestedTone'] = None
    current_row = 0
    row_dict = df.iloc[current_row].to_dict()
    return row_dict['id'], row_dict['hs'], row_dict['cs'], row_dict['topic'], row_dict['tone'], row_dict['isCSContextuallyRelevant'], row_dict['isToneMatch'], row_dict['suggestedTone']

def annotate_row(isCSContextuallyRelevant, isToneMatch, suggestedTone):
    global df
    global current_row

    df.at[current_row, 'isCSContextuallyRelevant'] = isCSContextuallyRelevant
    df.at[current_row, 'isToneMatch'] = isToneMatch
    df.at[current_row, 'suggestedTone'] = suggestedTone

    if current_row < len(df) - 1:
        current_row += 1
    else:
        current_row = 0
    df.to_csv('annotated_data.csv', index=False)
    
    row_dict = df.iloc[current_row].to_dict()
    return row_dict['id'], row_dict['hs'], row_dict['cs'], row_dict['topic'], row_dict['tone'], row_dict['isCSContextuallyRelevant'], row_dict['isToneMatch'], row_dict['suggestedTone'], 'annotated_data.csv'

def navigate(direction):
    global current_row
    if direction == "Previous":
        current_row = max(0, current_row - 1)
    elif direction == "Next":
        current_row = min(len(df) - 1, current_row + 1)
    elif direction == "First Unlabeled":
        unlabeled_row = df[df['isCSContextuallyRelevant'].isna()].index.min()
        if not np.isnan(unlabeled_row):
            current_row = int(unlabeled_row)

    row_dict = df.iloc[current_row].to_dict()
    return row_dict['id'], row_dict['hs'], row_dict['cs'], row_dict['topic'], row_dict['tone'], row_dict['isCSContextuallyRelevant'], row_dict['isToneMatch'], row_dict['suggestedTone']

with gr.Blocks(theme=gr.themes.Soft()) as annotator:
    gr.Markdown("## Data Annotation")
    
    with gr.Row():
        gr.Markdown("### Upload CSV")
        file_upload = gr.File()
        btn_load = gr.Button("Load CSV")

    with gr.Row():
        gr.Markdown("### Current Row")
        with gr.Row():
            idx = gr.Number(label='Index')
            hs = gr.Textbox(label='HS')
            cs = gr.Textbox(label='CS')

        with gr.Row():
            topic = gr.Textbox(label='Topic')
            tone = gr.Textbox(label='Tone')
        
        with gr.Row():
            isCSContextuallyRelevant = gr.Radio(["1", "0"], label="Contextually Relevant?")
            isToneMatch = gr.Radio(["1", "0"], label="Tone Match?")
            suggestedTone = gr.Dropdown(['', 'empathy', 'refutal', 'first_person', 'warn_of_conseq', 'humor', 'other'], label='Suggested Tone', interactive=True)
            btn_annotate = gr.Button("Annotate")

        with gr.Row():
            btn_previous = gr.Button("Previous")
            btn_next = gr.Button("Next")
            btn_first_unlabeled = gr.Button("First Unlabeled")

    with gr.Row():
        gr.Markdown("### Annotated Data File Download")
        file_download = gr.File()

    btn_load.click(load_csv, inputs=[file_upload], outputs=[idx, hs, cs, topic, tone, isCSContextuallyRelevant, isToneMatch, suggestedTone])
    btn_annotate.click(annotate_row, inputs=[isCSContextuallyRelevant, isToneMatch, suggestedTone], outputs=[idx, hs, cs, topic, tone, isCSContextuallyRelevant, isToneMatch, suggestedTone, file_download])
    btn_previous.click(navigate, inputs=gr.Textbox("Previous", visible=False), outputs=[idx, hs, cs, topic, tone, isCSContextuallyRelevant, isToneMatch, suggestedTone])
    btn_next.click(navigate, inputs=gr.Textbox("Next", visible=False), outputs=[idx, hs, cs, topic, tone, isCSContextuallyRelevant, isToneMatch, suggestedTone])
    btn_first_unlabeled.click(navigate, inputs=gr.Textbox("First Unlabeled", visible=False), outputs=[idx, hs, cs, topic, tone, isCSContextuallyRelevant, isToneMatch, suggestedTone])

annotator.launch()