app.py

from pathlib import Path
import zipfile
from typing import List, Tuple, Optional, Set
import json
import dataclasses
import gradio as gr
import asyncio
from openai import AsyncOpenAI
import tempfile
import os
import argparse
import gradio as gr
import random
import os
from pathlib import Path
import time
import matplotlib.pyplot as plt
import io

# BASE_URL = os.getenv("BASE_URL")
API_KEY = os.getenv("API_KEY")

BASE_URL = "https://api.openai.com"
print(f"BASE_URL: {BASE_URL}")
print(f"API_KEY: {API_KEY}")
if not BASE_URL or not API_KEY:
    raise ValueError("BASE_URL or API_KEY environment variables are not set")

client = AsyncOpenAI(api_key=API_KEY)


##########################################################################################################
#                                        HELPER FUNCTIONS                                                #
##########################################################################################################
async def run_command(cmd, timeout=5):
    process = await asyncio.create_subprocess_exec(
        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
    )
    try:
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
        return (
            stdout.decode("utf-8", errors="ignore"),
            stderr.decode("utf-8", errors="ignore"),
            process.returncode,
        )
    except asyncio.TimeoutError:
        process.kill()
        return None, None, None

# def echo(message, history):
#     return random.choice(["Yes", "No"])

# Prompt chatgpt with a message
async def chatgpt(prompt, history):
    messages = [
            {"role": "system", "content": ""}
        ]
    print(history)
    if history:
        messages += history
    messages += [{"role": "user", "content": prompt}]
    try:
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=messages
        )
    except Exception as e:
        print(e)
        return "I'm sorry, I'm having trouble. Could you please try again?"
    return response.choices[0].message.content

async def process_submission(finished_code, user_state):
    # Compile and execute user code, generate plot
    print("Compiling and plotting code")
    print(f"Code: {finished_code}")
    with tempfile.NamedTemporaryFile(delete=True, suffix=".py") as f:
        f.write(finished_code.encode("utf-8"))
        f.flush()
        stdout, stderr, exit_code = await run_command(["python", f.name], timeout=5)

    # result = await run_python_code(finished_code)
    print(f"Result: {stdout}")

    # Check if plot was created
    if f"temp_plot_{user_state}.png" in os.listdir():
        return f"temp_plot_{user_state}.png", stdout, stderr
    else:
        return "No plot generated", stdout, stderr
        # return gr.update(value="No plot generated", visible=True), None

# Function to create a zip file
def create_zip_file(jsonl_path, image_path, zip_path):
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        zipf.write(jsonl_path, arcname=Path(jsonl_path).name)
        zipf.write(image_path, arcname=Path(image_path).name)

# Function to assign plots to users randomly
def pick_random_image_for_user(users, images):
    assigned_images = {}
    for user in users:
        assigned_images[user] = random.sample(images, 5)
    # print(assigned_images)
    return assigned_images

##########################################################################################################
#                                        GRADIO INTERFACE SETUP                                          #
##########################################################################################################
# Define each page as a separate function
def create_interface(users):
    max_num_submissions = 5
    plot_time_limit = 130
    # plot_time_limit = 10
    dialogue_time_limit = 600
    # dialogue_time_limit = 10
    
    with gr.Blocks() as demo:
        user_state = gr.State()
        notes_state = gr.State([])
        dialogue_state = gr.State([])  # Store the conversation with the LLM
        submission_count = gr.State(0)  # Track number of code submissions
        produced_codes = gr.State([])
        previous_text = gr.State("")  # Track previous text in notepad
        random.seed(time.time())

        folder_path = "ChartMimic/dataset/ori_500"
        images = [f for f in os.listdir(folder_path) if f.endswith(('png', 'jpg', 'jpeg'))]
        chosen_image = os.path.join(folder_path, random.choice(images))
        assigned_images = pick_random_image_for_user(users, images)

        reference_code = chosen_image.replace(".png", ".py")
        chosen_image_state = gr.State(chosen_image)
        reference_code_state = gr.State(reference_code)

        expertise_survey_responses = gr.State({})
        uncertainty_survey_part_1_responses = gr.State({})  # Store responses to the uncertainty survey
        uncertainty_survey_part_2_responses = gr.State({})  # Store responses to the uncertainty survey
        uncertainty_survey_part_3_responses = gr.State({})  # Store responses to the uncertainty survey
        demographic_survey_responses = gr.State({})  # Store responses to the demographic survey

        ##########################################################################################################
        #                                        UI SETUP FOR EACH PAGE                                          #
        ##########################################################################################################
        # Page 1: Login, Add login components
        with gr.Column(visible=True) as login_row:
            instructions_text = gr.Markdown(f"## Instructions\n\nWelcome to Learning Games! PLEASE READ THE FOLLOWING INSTRUCTIONS CAREFULLY. \
                                             \n\nThis game consists of three parts:\n\n**Part 1: Inspection of the Chart**\n\nYou will be given \
                                            an image of a scientific chart. Please inspect it carefully and think about ways to reproduce it in \
                                            Python. You will have access to this plot throughout the experiment. You can take notes while \
                                            inspecting, a notepad will be given to you. At the end of the game, you will be asked to write \
                                            the code to recreate this chart. \n\n**Part 2: Chatting with a Teacher**\n\nIn this part, \
                                            you will have access to a teacher LLM! This interaction will be limited to only {int(dialogue_time_limit/60)}\
                                            minutes. You can use it to help you learn how to code this chart. Please be wise of your time \
                                            with the teacher LLM; by the end of this part, you will not be able to interact with the \
                                            LLM again. \n\n**Part 3: Writing the Code for the Chart**\n\nThis is the final crucial step. You will \
                                            have {max_num_submissions} attempts to reproduce the plot by writing, compiling, and running Python \
                                            code. You will be given a code skeleton to help you out, where you will fill in some required coding \
                                            components. You will be given only {max_num_submissions} attempts to compile your plot. \n\n Throughout \
                                            your interactions, you will be asked three times to rank your uncertainty: once during the inspection \
                                            of the chart, once after interacting with the LLM, and once after you submit your code. \
                                            \n\nAt the end of the game, you will be asked to fill out a short demographic survey. \
                                            Then you will be able to download your session data. Please download and send the zip file to <inan.m@northeastern.edu>. \
                                            \n\n**WARNING: You will not be able to go back to previous parts once you proceed, or reload the page.** \
                                            \n\n**Reminder: this is just a game; your performance will not affect your grade in the class in \
                                            any form.** \n\n \n\n ### Please login to start the game. We will first ask some questions about your \
                                            expertise, and part 1 will start immediately afterwards.")
            username_input = gr.Textbox(label="Username")
            login_button = gr.Button("Login")
            login_error_message = gr.Markdown(visible=False)

        # User Expertise Survey
        with gr.Column(visible=False) as expertise_survey:
            gr.Markdown("### Student Expertise Survey")
            gr.Markdown("Here is a short questionnaire before you get started. Please answer the following questions as accurately as possible.")
            expertise_survey_question1 = gr.CheckboxGroup(
                ["1 - No experience", "2 - Beginner", "3 - Intermediate", "4 - Advanced", "5 - Expert"],
                label="Question 1: On a scale of 1-5, what is your experience level of coding in Python? "
            )
            expertise_survey_question2 = gr.CheckboxGroup(
                ["1 - No experience", "2 - Beginner", "3 - Intermediate", "4 - Advanced", "5 - Expert"],
                label="Question 2: On a scale of 1-5, what is your experience level of using the Matplotlib library? "
            )
            expertise_survey_submit_button = gr.Button("Submit")


        # Instructions Page
        with gr.Column(visible=False) as instructions_page:
            instructions_text = gr.Markdown(f"## Part 1: Inspection of the Chart \n\nBelow, you are given a scientific chart. \
                                            Please inspect it carefully and think about ways to reproduce it in Python. You will \
                                            have access to this plot throughout the experiment. At the end of the game, you will \
                                            be asked to write the code to recreate this chart. You will be given a code skeleton \
                                            and the necessary data at the end. You can take notes below. You will have \
                                            {int(plot_time_limit/60)} minutes to take a look at this plot, starting now…")
            instruction_image_1 = gr.Image(show_label=False, height=500)
            plot_time_remaining = gr.Textbox(value=f"{(int(plot_time_limit/60)):02}:{(plot_time_limit%60):02}", label="Time Remaining", interactive=False)
            # questionnaire = gr.Form(["Question 1", "Question 2"], visible=False)

        # Uncertainty Survey Page
        with gr.Column(visible=False) as uncertainty_survey_part_1:
            instruction_image_2 = gr.Image(show_label=False, height=300)
            gr.Markdown("### Uncertainty Survey")
            gr.Markdown("Here is a short questionnaire before you get started. Please answer the following questions as accurately as possible.")
            uncertainty_survey_part_1_question1 = gr.CheckboxGroup(
                ["1 - Not certain", "2 - Somewhat certain", "3 - Moderately certain", "4 - Somewhat certain", "5 - Very certain"],
                label="Question 1: On a scale of 1-5, how certain are you that you can code this plot? "
            )
            
            uncertainty_survey_part_1_submit_button = gr.Button("Submit")

        # Dialogue Page with 5-minute timer
        with gr.Column(visible=False) as dialogue_page:
            instruction_text = gr.Markdown(f"## Part 2: Chatting with a Teacher \n\nNow, you will have access to a teacher LLM. This interaction will be limited to only {int(dialogue_time_limit/60)} minutes. \
                                           The countdown starts when you send your first message. You can use it to help you learn \
                                           how to code this chart. But be wise of your time; by the end of this part, \
                                           you will not be able to interact with the LLM again. Please use your time with \
                                           the LLM wisely, and think through your code solution before committing.\
                                           \n\n **You may want to prompt the LLM to teach you how to produce code for this chart** \
                                           **rather than having it output code directly. Please think about how to prompt the LLM to do this.**")
            with gr.Row():
                instruction_image_3 = gr.Image(show_label=False, height=400)
                with gr.Column():
                    # chatbot = gr.ChatInterface(echo, type="messages")
                    chatbot = gr.ChatInterface(chatgpt, type="messages", examples=["Teach me how to ...", "I want to learn step-by-step ...", "Explain to me slowly ..."])
                    chatbot.chatbot.height = 400
                    chatbot.chatbot.label = "Teacher LLM"
            # start_dialogue_button = gr.Button("Start Dialogue")
            part_2_time_remaining = gr.Textbox(value=f"{(int(dialogue_time_limit/60)):02}:{(dialogue_time_limit%60):02}", label="Time Remaining", interactive=False)

        # Uncertainty Survey Part 2
        with gr.Column(visible=False) as uncertainty_survey_part_2:
            instruction_image_4 = gr.Image(show_label=False, height=500)
            gr.Markdown("### Uncertainty Survey")
            gr.Markdown("Here is a short questionnaire after you have interacted with the teacher LLM. \
                        Please answer the following questions as accurately as possible.")
            uncertainty_survey_part_2_question1 = gr.CheckboxGroup(
                ["1 - Not at all", "2 - Slightly", "3 - Moderately", "4 - Very", "5 - Extremely"],
                label="Question 1: On a scale of 1-5, how much did the teacher LLM help you in learning how to code this plot? "
            )
            uncertainty_survey_part_2_question2 = gr.CheckboxGroup(
                ["1 - Not certain", "2 - Somewhat certain", "3 - Moderately certain", "4 - Somewhat certain", "5 - Very certain"],
                label="Question 2: On a scale of 1-5, how certain are you that you can code this plot now? "
            )
            uncertainty_survey_part_2_question3 = gr.CheckboxGroup(
                ["1 - Not certain", "2 - Somewhat certain", "3 - Moderately certain", "4 - Somewhat certain", "5 - Very certain"],
                label="Question 3: On a scale of 1-5, how certain are you that you can code this plot even without the teacher LLM? "
            )
            uncertainty_survey_part_2_question4 = gr.CheckboxGroup(
                ["1 - Not on topic at all", "2 - Somewhat not on topic", "3 - Moderately on topic", "4 - Somewhat on topic", "5 - Mostly on topic"],
                label="Question 4: On a scale of 1-5, how much did the LLM stay on topic (i.e. did it answer your questions specifically)?"
            )
            uncertainty_survey_part_2_submit_button = gr.Button("Submit")

        # Final Code Editor Page
        with gr.Column(visible=False) as final_page:
            instruction_text = gr.Markdown(f"## Part 3: Writing the Code for the Chart \n\nThis is the final crucial step. \
                                           You need to reproduce the original plot by writing, compiling, and running Python code. \
                                           You are given a code skeleton below to help you, where you will fill in the \
                                           required coding components. When you compile, you will be able to see the output of \
                                           your code, in addition to the plot. You will be given only {max_num_submissions} attempts to compile your plot.")
            instruction_image_5 = gr.Image(show_label=False, height=400)
            code_editor = gr.Code(language="python", label="Code Editor")
            run_code_button = gr.Button("Compile & Run Code")
            processing_message = gr.Textbox(value="Processing...", visible=False)
            with gr.Row():
                retry_button = gr.Button("Retry", visible=False)
                finished_button = gr.Button("Finished", visible=False)
            with gr.Row():
                stdout_message = gr.Textbox(visible=True, label="Code Output", value="")
                submission_counter = gr.Number(visible=True, label="Number of Remaining Submissions", value=max_num_submissions)
            plot_output = gr.Image(visible=False, height=400)

        # Uncertainty Survey Part 3
        with gr.Column(visible=False) as uncertainty_survey_part_3:
            with gr.Row():
                instruction_image_6 = gr.Image(label="Original Chart", height=300)
                generated_image = gr.Image(label="Your Generated Chart", height=300)
            gr.Markdown("### Uncertainty Survey")
            gr.Markdown("Here is a short questionnaire after you have finalized your code. Please answer the following questions as accurately as possible.")
            uncertainty_survey_part_3_question1 = gr.CheckboxGroup(
                ["1 - Not at all", "2 - Slightly", "3 - Moderately", "4 - Very", "5 - Extremely"],
                label="Question 1: On a scale of 1-5, how much did you rely on the teacher LLM and your notes to code this chart? "
            )
            uncertainty_survey_part_3_question2 = gr.CheckboxGroup(
                ["1 - Much harder", "2 - Harder", "3 - As expected", "4 - Easier", "5 - Much easier"],
                label="Question 2: On a scale of 1-5, was the task easier or harder than you expected? "
            )
            uncertainty_survey_part_3_question3 = gr.CheckboxGroup(
                ["1 - Could not produce", "2 - Very inaccurate", "3 - Moderately inaccurate", "4 - Somewhat accurate", "5 - Very accurate"],
                label="Question 3: On a scale of 1-5, how accurate is your chart compared to the original? "
            )
            uncertainty_survey_part_3_question4 = gr.CheckboxGroup(
                ["1 - No experience", "2 - Beginner", "3 - Intermediate", "4 - Advanced", "5 - Expert"],
                label="Question 4: On a scale of 1-5, how would you rate your experience in Python now? "
            )
            uncertainty_survey_part_3_question5 = gr.CheckboxGroup(
                ["1 - No experience", "2 - Beginner", "3 - Intermediate", "4 - Advanced", "5 - Expert"],
                label="Question 5: On a scale of 1-5, how would you rate your experience in using the Matplotlib library now? "
            )
            uncertainty_survey_part_3_question6 = gr.CheckboxGroup(
                ["1 - Very ambiguous", "2 - Somewhat ambiguous", "3 - Neither ambiguous nor clear", "4 - Somewhat clear", "5 - Very clear"],
                label="Question 5: On a scale of 1-5, throughout this experiment how ambigous were the instructions?"
            )
            uncertainty_survey_part_3_question7 = gr.CheckboxGroup(
                ["1 - Very ambiguous", "2 - Somewhat ambiguous", "3 - Neither ambiguous nor clear", "4 - Somewhat clear", "5 - Very clear"],
                label="Question 5: On a scale of 1-5, throughout this experiment how ambigous was the given plot?"
            )
            uncertainty_survey_part_3_submit_button = gr.Button("Submit")

        # Demographic Survey Page
        with gr.Column(visible=False) as demographic_survey:
            gr.Markdown("### Demographic Survey")
            gr.Markdown("Please answer the following questions to help us understand your background.")
            demographic_survey_question1 = gr.CheckboxGroup(
                ["Undergraduate", "Graduate", "PhD", "Postdoc", "Faculty", "Industry Professional", "Other"],
                label="What is your current academic status?"
            )
            demographic_survey_question2 = gr.CheckboxGroup(
                ["Bouvé College of Health Sciences", "College of Arts, Media and Design", "College of Engineering", "College of Professional Studies", "College of Science", "D'Amore-McKim School of Business", "Khoury College of Computer Sciences", "School of Law", "Mills College at Northeastern", "Other"],
                label="What is your college?"
            )
            demographic_survey_question3 = gr.CheckboxGroup(
                ["18-23", "23-27", "27-31", "31-35", "35-43", "43+"],
                label="What is your age group?"
            )
            demographic_survey_question4 = gr.CheckboxGroup(
                ["Woman", "Man", "Transgender", "Non-binary", "Prefer not to say"],
                label="What is your gender identity?"
            )
            demographic_survey_question5 = gr.CheckboxGroup(
                ["American Indian or Alaska Native", "Asian or Asian American", "Black or African American", "Hispanic or Latino/a/x", "Native Hawaiian or Other Pacific Islander", "Middle Eastern or North African", "White or European", "Other"],
                label="What is your ethnicity? (Select all that apply)"
            )
            demographic_survey_submit_button = gr.Button("Submit")

        # Exit Page
        with gr.Column(visible=False) as exit_page:
            gr.Markdown("## Thank you for participating in the Learning Games! \n\nYour responses have been recorded. Please download your session data below, and send the zip file to <inan.m@northeastern.edu>.")
            download_button = gr.Button("Download Session Data")
            file_to_download = gr.File(label="Download Results")


        # Adding the notepad available on all pages
        with gr.Column(visible=False) as notepad_column:
            notepad = gr.Textbox(lines=10, placeholder="Take notes here", value="", label="Notepad", elem_id="notepad")
            
        
        ##########################################################################################################
        #                           FUNCTION DEFINITIONS FOR EACH PAGE                                           #
        ##########################################################################################################
        def on_login(users: Set[str], folder_path, assigned_images):
            def callback(username):
                if username not in users:
                    return (
                        gr.update(visible=True),  # login still visible
                        gr.update(visible=False),  # main interface still not visible
                        gr.update(visible=True, value="Username not found"),
                        "",
                        gr.update(),                # for image state to change with the user
                        gr.update(),                # for ref code
                    )
                chosen_image = os.path.join(folder_path, random.choice(assigned_images[username]))
                return (
                    gr.update(visible=False),  # login hidden
                    gr.update(visible=True),  # main interface visible
                    gr.update(visible=False),  # login error message hidden
                    username,
                    chosen_image, # for image state
                    chosen_image.replace(".png", ".py")
                )

            return callback
        
        def update_all_instruction_images(chosen_image):
            return (
                gr.update(value=chosen_image),
                gr.update(value=chosen_image),
                gr.update(value=chosen_image),
                gr.update(value=chosen_image),
                gr.update(value=chosen_image),
                gr.update(value=chosen_image)
            )
        
        def extract_code_context(reference_code, user_state):
            with open(reference_code, "r") as f:
                code_context = f.read()
            print(code_context)
            # Remove everything between Part 3: Plot Configuration and Rendering and Part 4: Saving Output
            start_index = code_context.find("# ===================\n# Part 3: Plot Configuration and Rendering\n# ===================")
            end_index = code_context.find("# ===================\n# Part 4: Saving Output\n# ===================")
            code_context = code_context[:start_index] + "# ===================\n# Part 3: Plot Configuration and Rendering\n# ===================\n\n # TODO: YOUR CODE GOES HERE #\n\n\n" + code_context[end_index:]
            # plt.savefig is the last line of the code, remove it
            end_index = code_context.find("plt.savefig")
            code_context = code_context[:end_index]
            # and replace with plt.show()
            code_context += f"plt.savefig('temp_plot_{user_state}.png')\n"
            # code_context += "plt.show()\n"
            return code_context
        
        def handle_expertise_survey_response(q1, q2):
            # Example: Store responses in a dictionary or process as needed
            response = {
                "Question 1": q1,
                "Question 2": q2
            }
            return response

        # Function to handle form submission
        def handle_part1_survey_response(q1):
            # Example: Store responses in a dictionary or process as needed
            response = {
                "Question 1": q1
            }
            return response
        
        def handle_part2_survey_response(q1, q2, q3, q4):
            # Example: Store responses in a dictionary or process as needed
            response = {
                "Question 1": q1,
                "Question 2": q2,
                "Question 3": q3,
                "Question 4": q4
            }
            return response
        
        def handle_final_survey_response(q1, q2, q3, q4, q5, q6, q7):
            # Example: Store responses in a dictionary or process as needed
            response = {
                "Question 1": q1,
                "Question 2": q2,
                "Question 3": q3,
                "Question 4": q4,
                "Question 5": q5,
                "Question 6": q6,
                "Question 7": q7
            }
            return response
        
        def handle_demographic_survey_response(q1, q2, q3, q4, q5):
            # Example: Store responses in a dictionary or process as needed
            response = {
                "Question 1": q1,
                "Question 2": q2,
                "Question 3": q3,
                "Question 4": q4,
                "Question 5": q5
            }
            return response
        
        # Timer logic for instructions page
        def plot_countdown_timer():
            time_limit = plot_time_limit
            start_time = time.time()
            while time.time() - start_time < time_limit:
                mins, secs = divmod(time_limit - int(time.time() - start_time), 60)
                yield f"{mins:02}:{secs:02}", gr.update(), gr.update(visible=False)
            yield "00:00", gr.update(visible=False), gr.update(visible=True)

        # Timer logic for dialogue page
        def dialogue_countdown_timer():
            time_limit = dialogue_time_limit
            start_time = time.time()
            while time.time() - start_time < time_limit:
                mins, secs = divmod(time_limit - int(time.time() - start_time), 60)
                yield f"{mins:02}:{secs:02}", gr.update(visible=True), gr.update(visible=False)
            yield "00:00", gr.update(visible=False), gr.update(visible=True)

        # New function to save dialogue state
        def save_dialogue_state(dialogue, dialogue_state):
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            print(dialogue)
            print(dialogue_state)
            return dialogue_state + [timestamp, dialogue]
        
        # # Save notes, dialogue, and answers into a file for download
        # def prepare_download(notes, dialogue, answers):
        #     results = {
        #         "notes": notes,
        #         "dialogue": dialogue,
        #         "answers": answers
        #     }
        #     with open("session_data.json", "w") as f:
        #         json.dump(results, f)
        #     return "session_data.json"

        # Add download functionality
        def get_download_link(user_state, chosen_image, notes_state, dialogue_state, 
                            produced_codes, reference_code, survey1, survey2, survey3, survey4, survey5):
            jsonl_path = Path(f"session_data_{user_state}.jsonl")
            with open(jsonl_path, "w") as f:
                f.write(
                    json.dumps(
                        {
                            "username": user_state,
                            "chosen_image": chosen_image,
                            "notes": notes_state,
                            "dialogue_state": dialogue_state,
                            "produced_codes": produced_codes,
                            "reference_code": reference_code,
                            "expertise_survey": survey1,
                            "uncertainty_survey_part1": survey2,
                            "uncertainty_survey_part2": survey3,
                            "uncertainty_survey_part3": survey4,
                            "demographics_survey": survey5
                        }
                    )
                    + "\n"
                )
            
            image_path = Path(f"temp_plot_{user_state}.png")
            zip_path = Path(f"session_data_{user_state}.zip")
            create_zip_file(jsonl_path, image_path, zip_path)
            
            if not zip_path.exists():
                return None
            return gr.File(value=str(zip_path), visible=True)
            
        async def on_submit(finished_code, submission_count, produced_codes, user_state):
            if (max_num_submissions-(submission_count+1)) == 0:
                # raise gr.Error("Max submissions reached")
                yield (
                    gr.update(visible=False),
                    gr.update(visible=False),  # Hide run code button
                    gr.update(visible=False),  # Hide retry button
                    gr.update(visible=True),  # Show finished button
                    gr.update(visible=False),  # Hide plot output
                    submission_count,
                    produced_codes,
                    gr.update(visible=False),    # stdout
                    gr.update(visible=False) #submission counter
                )
                raise gr.Error("Max submissions reached")
            else:
                submission_count += 1
                # Show processing message and hide other elements
                yield (
                    gr.update(visible=True),  # Show processing message
                    gr.update(visible=False),  # Hide run code button
                    gr.update(visible=False),  # Hide retry button
                    gr.update(visible=False),  # Hide finished button
                    gr.update(visible=False),  # Hide plot output
                    submission_count,
                    produced_codes,
                    gr.update(visible=False),   # stdout
                    gr.update(value=max_num_submissions-submission_count) #submission counter
                )

                # Process the submission
                plot_output, stdout, stderr = await process_submission(finished_code, user_state)

                # Hide processing message and show result
                yield (
                    gr.update(visible=False),  # Hide processing message
                    gr.update(visible=False),  # Hide submit button
                    gr.update(visible=True),  # Show retry button
                    gr.update(visible=True),  # Show finished button
                    gr.update(visible=True, value=plot_output),  # Show plot output
                    submission_count,
                    produced_codes + [finished_code],
                    gr.update(visible=True, value=stdout+stderr),    # stdout
                    gr.update() #submission counter
                )

        def on_retry(finished_code, produced_codes):
            # Hide processing message and show result
            yield (
                gr.update(visible=False),  # Hide processing message
                gr.update(visible=True),  # Show submit button
                gr.update(visible=False),  # Hide retry button
                gr.update(visible=False),  # Hide finished button
                gr.update(visible=False),  # Hide plot output
                produced_codes + [finished_code]
            )
            
        def filter_paste(previous_text, new_text):
            # Check if the new input is a result of pasting (by comparing lengths or content)
            print(f"New text: {new_text}")
            changed_text = new_text.replace(previous_text, "")
            if len(changed_text) > 10:  # Paste generally increases length significantly
                return previous_text, previous_text  # Revert to previous text if paste is detected
            previous_text = new_text
            print(f"Previous text: {previous_text}")
            return previous_text, new_text

        def save_notes_with_timestamp(notes, notes_state):
            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
            notes_state.append(f"{timestamp}: {notes}")
            return notes_state

        ##########################################################################################################
        #                                      EVENT HANDLERS FOR EACH PAGE                                      #
        ##########################################################################################################
        # Page navigation
        login_button.click(
            on_login(users, folder_path, assigned_images),
            inputs=[username_input],
            outputs=[login_row, expertise_survey, login_error_message, user_state, chosen_image_state, reference_code_state],
        )

        # login_button.click(lambda: os.path.join(folder_path, random.choice(images)), outputs=[chosen_image_state])

        # login_button.click(lambda: chosen_image_state.replace(".png", ".py"), inputs=[chosen_image_state], outputs=[reference_code_state])

        expertise_survey_submit_button.click(
            handle_expertise_survey_response,
            inputs=[expertise_survey_question1, expertise_survey_question2],
            outputs=[expertise_survey_responses]
        )

        expertise_survey_submit_button.click(
            lambda: (gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)), # Hide survey, show dialogue
            inputs=[], outputs=[expertise_survey, instructions_page, notepad_column]
        )

        expertise_survey_submit_button.click(
            update_all_instruction_images,
            inputs=[chosen_image_state], outputs=[instruction_image_1, instruction_image_2,
                                                  instruction_image_3, instruction_image_4,
                                                  instruction_image_5, instruction_image_6]
        )

        expertise_survey_submit_button.click(plot_countdown_timer, outputs=[plot_time_remaining, instructions_page, uncertainty_survey_part_1])

        uncertainty_survey_part_1_submit_button.click(
            handle_part1_survey_response,
            inputs=[uncertainty_survey_part_1_question1],
            outputs=[uncertainty_survey_part_1_responses]
        )

        uncertainty_survey_part_1_submit_button.click(
            lambda: (gr.update(visible=False), gr.update(visible=True)), # Hide survey, show dialogue
            inputs=[], outputs=[uncertainty_survey_part_1, dialogue_page]
        )

        chatbot.chatbot.change(
            dialogue_countdown_timer,
            outputs=[part_2_time_remaining, dialogue_page, uncertainty_survey_part_2],
            trigger_mode = "once"
        )

        # Update to save dialogue state on change
        chatbot.chatbot.change(
            save_dialogue_state,
            inputs=[chatbot.chatbot, dialogue_state],
            outputs=[dialogue_state]
        )

        uncertainty_survey_part_2_submit_button.click(
            handle_part2_survey_response,
            inputs=[uncertainty_survey_part_2_question1, uncertainty_survey_part_2_question2, 
                    uncertainty_survey_part_2_question3, uncertainty_survey_part_2_question4],
            outputs=[uncertainty_survey_part_2_responses]
        )

        uncertainty_survey_part_2_submit_button.click(
            lambda: (gr.update(visible=False), gr.update(visible=True)), # Hide survey, show final page
            inputs=[], outputs=[uncertainty_survey_part_2, final_page]
        )

        uncertainty_survey_part_2_submit_button.click(
            extract_code_context,
            inputs=[reference_code_state, user_state], outputs=[code_editor]
        )

        run_code_button.click(
            on_submit,
            inputs=[code_editor, submission_count, produced_codes, user_state],
            outputs=[
                processing_message,
                run_code_button,
                retry_button,
                finished_button,
                plot_output,
                submission_count,
                produced_codes,
                stdout_message,
                submission_counter
            ],
        )

        retry_button.click(
            on_retry,
            inputs=[code_editor, produced_codes],
            outputs=[
                processing_message,
                run_code_button,
                retry_button,
                finished_button,
                plot_output,
                produced_codes,
            ],
        )

        finished_button.click(
            lambda user_state: (gr.update(visible=False), gr.update(visible=True), f"temp_plot_{user_state}.png"), # Hide final page, show survey
            inputs=[user_state], outputs=[final_page, uncertainty_survey_part_3, generated_image]
        )

        uncertainty_survey_part_3_submit_button.click(
            handle_final_survey_response,
            inputs=[uncertainty_survey_part_3_question1, uncertainty_survey_part_3_question2, 
                    uncertainty_survey_part_3_question3, uncertainty_survey_part_3_question4, 
                    uncertainty_survey_part_3_question5, uncertainty_survey_part_3_question6,
                    uncertainty_survey_part_3_question7],
            outputs=[uncertainty_survey_part_3_responses]
        )

        uncertainty_survey_part_3_submit_button.click(
            lambda: (gr.update(visible=False), gr.update(visible=True)), # Hide survey, show demographic survey
            inputs=[], outputs=[uncertainty_survey_part_3, demographic_survey]
        )

        demographic_survey_submit_button.click(
            handle_demographic_survey_response,
            inputs=[demographic_survey_question1, demographic_survey_question2, demographic_survey_question3, demographic_survey_question4, demographic_survey_question5],
            outputs=[demographic_survey_responses]
        )

        demographic_survey_submit_button.click(
            lambda: (gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)), # Hide survey, show exit page
            inputs=[], outputs=[demographic_survey, exit_page, download_button, notepad]
        )

        # notepad.change(filter_paste, 
        #                inputs=[previous_text, notepad], 
        #                outputs=[previous_text, notepad], trigger_mode="always_last")

        demographic_survey_submit_button.click(save_notes_with_timestamp, 
                       inputs=[notepad, notes_state],
                       outputs=[notes_state])

        download_button.click(
            get_download_link, 
            inputs=[user_state, chosen_image_state, notes_state, 
                    dialogue_state, produced_codes, reference_code_state,
                    expertise_survey_responses,
                    uncertainty_survey_part_1_responses, 
                    uncertainty_survey_part_2_responses, 
                    uncertainty_survey_part_3_responses, 
                    demographic_survey_responses],
            outputs=[file_to_download]
        )

        demo.load(
            lambda: gr.update(visible=True),  # Show login page
            outputs=login_row,
        )

    return demo


# if __name__ == "__main__":
#     users = Path("users.txt").read_text().splitlines()
#     users = set(user.strip() for user in users if user.strip())
#     chosen_image = pick_random_image()
#     reference_code = chosen_image.replace(".png", ".py")
#     # code_context = extract_code_context(reference_code)
#     demo = create_interface(users, chosen_image, reference_code)

#     # demo.launch(
#     #     server_name=args.server_name,
#     #     server_port=args.server_port,
#     #     share=args.share,
#     # )

#     demo.launch()

users = Path("users.txt").read_text().splitlines()
users = set(user.strip() for user in users if user.strip())
# chosen_image = pick_random_image()
# reference_code = chosen_image.replace(".png", ".py")
# code_context = extract_code_context(reference_code)
demo = create_interface(users)

demo.launch()