import gradio as gr
import json
import os

# root directory of image files
# each action should be saved under subdirectories
# input filename
image_root = "forward_dynamics_qa_pairs_v5_w_newline"
# output filename
output_file = "v5_forward_dynamics_user_choices.json"

# # input filename
# image_root = "/Users/bryan/Desktop/wkdir/VLM/src/human-annotation-interface/inverse_dynamics_qa_pairs_v5_w_newline"
# # output filename
# output_file = "/Users/bryan/Desktop/wkdir/VLM/src/human-annotation-interface/v5_inverse_dynamics_user_choices.json"


if not os.path.exists(output_file):
    with open(output_file, 'w') as f:
        json.dump({}, f)

def load_action_images():
    action_images = {}
    for action in os.listdir(image_root):
        action_dir = os.path.join(image_root, action)
        if os.path.isdir(action_dir):
            images = [f for f in os.listdir(action_dir) if f.endswith('.jpg')]
            images.sort()  # Ensure files are sorted in ascending order
            action_images[action] = images
    return action_images

def load_user_choices():
    with open(output_file, 'r') as f:
        return json.load(f)

def save_user_choice(action, image_name, choice, ground_truth):
    image_name_no_ext = os.path.splitext(image_name)[0]  # remove ".jpg" ext
    user_choices = load_user_choices()
    is_correct = (choice == ground_truth)
    
    # save result
    if action not in user_choices:
        user_choices[action] = {}
    user_choices[action][image_name_no_ext] = {
        "choice": choice,
        "ground_truth": ground_truth,
        "is_correct": is_correct
    }
    with open(output_file, 'w') as f:
        json.dump(user_choices, f, indent=2)

def get_content_at_index(action, index):
    if action not in action_images or index < 0 or index >= len(action_images[action]):
        return None, None, "No more images", "", False, False

    # image
    image_name = action_images[action][index]
    image_path = os.path.join(image_root, action, image_name)

    # text prompt
    text_prompt_path = image_path.replace(".jpg", ".txt")
    text_prompt = (
        open(text_prompt_path, 'r').read().strip()
        if os.path.exists(text_prompt_path)
        else "No text prompt available"
    )

    # Apply font size styling to text_prompt
    text_prompt = f"<div style='font-size: 1.25em;'>{text_prompt}</div>"

    # ground truth
    ground_truth_path = image_path.replace(".jpg", "_answer.txt")
    ground_truth = (
        open(ground_truth_path, 'r').read().strip()
        if os.path.exists(ground_truth_path)
        else "No ground truth available"
    )

    # button states
    enable_prev = index > 0
    enable_next = index < len(action_images[action]) - 1

    return image_path, image_name, text_prompt, ground_truth, enable_prev, enable_next


# navigate among questions
def navigate(action, index, direction):
    # update index
    new_index = max(0, min(index + direction, len(action_images[action]) - 1))
    # retrieve context
    image_path, _, text_prompt, ground_truth, enable_prev, enable_next = get_content_at_index(action, new_index)
    
    # Apply font size styling to text_prompt
    styled_text_prompt = f"<div style='font-size: 1.25em;'>{text_prompt}</div>"
    
    return (
        image_path,
        styled_text_prompt,
        ground_truth,
        gr.update(value=""),
        gr.update(interactive=enable_prev),
        gr.update(interactive=enable_next),
        new_index
    )


# Handle user choice submission
def submit_choice(action, index, choice, ground_truth):
    if action not in action_images or index < 0 or index >= len(action_images[action]):
        return "Invalid demo or keyframe index."

    image_name = action_images[action][index]
    save_user_choice(action, image_name, choice, ground_truth)  # Save user choice

    if choice == ground_truth:
        color = "green"
    else:
        color = "red"
    return f'<div style="font-size: 1.25em; color:{color}">Ground Truth: {ground_truth}</div>'



def change_action(action):
    if action not in action_images:
        return None, "No images available", "No text prompt available", "", gr.update(interactive=False), gr.update(interactive=False), action, 0

    # Get the first image of the new action
    image_path, image_name, text_prompt, ground_truth, enable_prev, enable_next = get_content_at_index(action, 0)

    # Apply font size styling to text_prompt
    styled_text_prompt = f"<div style='font-size: 1.25em;'>{text_prompt}</div>"

    # Reset states
    enable_prev = gr.update(interactive=False)  # Disable "Previous" as we're at the start
    enable_next = gr.update(interactive=enable_next)  # Enable "Next" if there are more images

    return image_path, styled_text_prompt, ground_truth, gr.update(value=""), enable_prev, enable_next, action, 0

# initialize data
action_images = load_action_images()

def initialize_app():
    if not action_images:
        return None, None, "No actions available", "", gr.update(interactive=False), gr.update(interactive=False), "", 0

    first_action = list(action_images.keys())[0]
    image_path, image_name, text_prompt, ground_truth, enable_prev, enable_next = get_content_at_index(first_action, 0)

    # Force the Previous button to be disabled during initialization
    enable_prev = gr.update(interactive=False)

    return image_path, image_name, text_prompt, ground_truth, enable_prev, gr.update(interactive=enable_next), first_action, 0

first_image, first_image_name, first_text_prompt, first_ground_truth, enable_prev, enable_next, first_action, first_index = initialize_app()

# Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# VLM Embodied Benchmark Human Annotation Interface")

    # states: action, index, ground truth
    current_action = gr.State(value=first_action)
    current_index = gr.State(value=first_index)
    current_ground_truth = gr.State(value=first_ground_truth)

    # UI components
    action_dropdown = gr.Dropdown(choices=list(action_images.keys()), value=first_action, label="Select Demo to Annotate")
    # image = gr.Image(value=first_image, interactive=False, width=500)
    image = gr.Image(value=first_image, interactive=False, width=1500)
    text_prompt = gr.Markdown(value=f"<div style='font-size: 1.25em;'>{first_text_prompt}</div>")
    with gr.Row(): # nav buttons
        prev_button = gr.Button("Previous", interactive=False)  # Explicitly disabled during initialization
        next_button = gr.Button("Next", interactive=enable_next["interactive"])
    with gr.Row(): # choice buttons
        a_button = gr.Button("A")
        b_button = gr.Button("B")
        c_button = gr.Button("C")
        d_button = gr.Button("D")
    ground_truth_display = gr.Markdown(value="")

    # change action dropdown
    action_dropdown.change(
        fn=change_action,
        inputs=[action_dropdown],
        outputs=[image, text_prompt, current_ground_truth, ground_truth_display, prev_button, next_button, current_action, current_index]
    )
    
    # click on navigation buttons
    nav_input = [current_action, current_index]
    nav_output = [image, text_prompt, current_ground_truth, ground_truth_display, prev_button, next_button, current_index]
    prev_button.click(
        fn=lambda action, index: navigate(action, index, -1),
        inputs=nav_input, outputs=nav_output
    )
    next_button.click(
        fn=lambda action, index: navigate(action, index, 1),
        inputs=[current_action, current_index],
        outputs=nav_output
    )

    # click on choice buttons
    input_param = [current_action, current_index, current_ground_truth]
    output_param= [ground_truth_display]
    a_button.click(
        fn=lambda action, index, gt: submit_choice(action, index, "A", gt),
        inputs=input_param, outputs=output_param)
    b_button.click(
        fn=lambda action, index, gt: submit_choice(action, index, "B", gt),
        inputs=input_param, outputs=output_param)
    c_button.click(
        fn=lambda action, index, gt: submit_choice(action, index, "C", gt),
        inputs=input_param, outputs=output_param)
    d_button.click(
        fn=lambda  action, index, gt: submit_choice(action, index, "D", gt),
        inputs=input_param, outputs=output_param)

app.launch(share=True)