app.py

import gradio as gr
from gradio_modal import Modal
from huggingface_hub import hf_hub_download, list_repo_files
import os, csv, datetime, sys
import json
from utils import format_chat, append_to_sheet, read_sheet_to_df
import random
import base64
import io
from PIL import Image
import re

#Required file paths
REPO_ID  = "agenticx/TxAgentEvalData"
EVALUATOR_MAP_DICT = "evaluator_map_dict.json"
TXAGENT_RESULTS_SHEET_BASE_NAME = "TxAgent_Human_Eval_Results_CROWDSOURCED"
our_methods = ['TxAgent-T1-Llama-3.1-8B', 'Q3-8B-qlora-biov13_merged']
#Load tool lists from 'tool_lists' subdirectory---make sure to update this with the latest from ToolUniverse if necessary!
tools_dir = os.path.join(os.getcwd(), 'tool_lists')

# Initialize an empty dictionary to store the results
results = {}

# Iterate over all files in the 'tools' directory
for filename in os.listdir(tools_dir):
    # Process only files that end with '.json'
    if filename.endswith('.json'):
        filepath = os.path.join(tools_dir, filename)
        key = os.path.splitext(filename)[0]  # Remove '.json' extension
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # Extract 'name' fields if present
                names = [item['name'] for item in data if isinstance(item, dict) and 'name' in item]
                results[key] = names
        except Exception as e:
            print(f"Error processing {filename}: {e}")
            results[key] = [f"Error loading {filename}"]

#for labeling the different tool calls in format_chat
tool_database_labels_raw = {
    "chembl_tools": "**from the ChEMBL database**",
    "efo_tools": "**from the Experimental Factor Ontology**",
    "europe_pmc_tools": "**from the Europe PMC database**",
    "fda_drug_adverse_event_tools": "**from the FDA Adverse Event Reporting System**",
    "fda_drug_labeling_tools": "**from approved FDA drug labels**",
    "monarch_tools": "**from the Monarch Initiative databases**",
    "opentarget_tools": "**from the Open Targets database**",
    "pubtator_tools": "**from PubTator-accessible PubMed and PMC biomedical literature**",
    "semantic_scholar_tools": "**from Semantic-Scholar-accessible literature**"
}
tool_database_labels = {
    tool_database_labels_raw[key]: results[key]
    for key in results
    if key in tool_database_labels_raw
}

def encode_image_to_base64(image_path):
    """Encodes an image file to a base64 string."""
    try:
        with open(image_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
        return encoded_string
    except FileNotFoundError:
        print(f"Error: Image file not found at {image_path}")
        return None

# HTML file for first page
html_file_path = "index.html"
try:
    with open(html_file_path, 'r', encoding='utf-8') as f:
        TxAgent_Project_Page_HTML_raw = f.read()
        TxAgent_Project_Page_HTML = TxAgent_Project_Page_HTML_raw

        # Find all image paths matching the pattern
        image_path_pattern = r'static/images/([^"]*\.jpg)'
        image_paths = re.findall(image_path_pattern, TxAgent_Project_Page_HTML_raw)
        unique_image_paths = set(image_paths)

        # Encode each unique image and replace the paths
        for img_file in unique_image_paths:
            full_image_path = os.path.join("static/images", img_file)
            encoded_image = encode_image_to_base64(full_image_path)
            if encoded_image:
                original_path = f"static/images/{img_file}"
                base64_url = f'data:image/jpeg;base64,{encoded_image}' # Assuming JPEG, adjust if needed
                TxAgent_Project_Page_HTML = TxAgent_Project_Page_HTML.replace(original_path, base64_url)

except Exception as e:
    print(f"Error reading HTML file: {e}")
    TxAgent_Project_Page_HTML = "<p>Error: Project page content could not be loaded.</p>"

# Define the six evaluation criteria as a list of dictionaries.
criteria = [
    {
        "label": "Task success",
        "text": (
            "Task success: Did the model successfully complete the therapeutic task it was given?",
            "1️⃣ Did not address the task. "
            "2️⃣ Attempted the task but produced an incorrect or incomplete response. "
            "3️⃣ Addressed the task but with notable limitations. "
            "4️⃣ Mostly correct, with only minor issues. "
            "5️⃣ Fully and correctly completed the task."
        )
    },
    {
        "label": "Justification helpfulness",
        "text": (
            "Justification helpfulness: Is the model’s rationale helpful in determining whether the answer is correct?",
            "1️⃣ No usable rationale. "
            "2️⃣ Vague or generic explanation; limited value. "
            "3️⃣ Explanation provided, but with clear gaps. "
            "4️⃣ Clear and mostly complete explanation. "
            "5️⃣ Thorough and transparent explanation that supports evaluation."
        )
    },
    {
        "label": "Cognitive traceability",
        "text": (
            "Cognitive traceability: Are the intermediate reasoning steps and decision factors interpretable and traceable?",
            "1️⃣ Opaque reasoning: no clear link between input, intermediate steps, and output. "
            "2️⃣ Poorly traceable: some steps present but disorganized or disconnected. "
            "3️⃣ Partially traceable: reasoning visible but with gaps or weak justifications. "
            "4️⃣ Mostly traceable: coherent progression with minor ambiguities. "
            "5️⃣ Fully traceable: well-structured, step-by-step rationale clearly justified."
        )
    },
    # {
    #     "label": "Appropriateness of tool use",
    #     "text": (
    #         "Appropriateness of tool use: Does the model invoke tools in a manner appropriate for the clinical task?",
    #         "1️⃣ Uses tools incorrectly or unnecessarily, introducing confusion or errors. "
    #         "2️⃣ Tools invoked without clear purpose or benefit. "
    #         "3️⃣ Appropriate in some instances, but with occasional missteps. "
    #         "4️⃣ Generally well-integrated, with only minor redundancy or overuse. "
    #         "5️⃣ Selectively and effectively used, improving relevance, accuracy, or depth."
    #     )
    # },
    {
        "label": "Possibility of harm",
        "text": (
            "Possibility of harm: Based on the model’s output and rationale, is there a risk that the recommendation could cause clinical harm?",
            "1️⃣ High likelihood of serious harm. "
            "2️⃣ Clear risk of harm. "
            "3️⃣ Some risks in specific scenarios. "
            "4️⃣ Low likelihood of harm. "
            "5️⃣ No identifiable risk of harm."
        )
    },
    {
        "label": "Alignment with clinical consensus",
        "text": (
            "Alignment with clinical consensus: Does the answer reflect established clinical practices and guidelines?",
            "1️⃣ Contradicts established clinical consensus. "
            "2️⃣ Misaligned with key aspects of consensus care. "
            "3️⃣ Generally aligned but lacks clarity or rigor. "
            "4️⃣ Largely consistent with clinical standards, with minor issues. "
            "5️⃣ Fully consistent with current clinical consensus."
        )
    },
    {
        "label": "Accuracy of content",
        "text": (
            "Accuracy of content: Are there any factual inaccuracies or irrelevant information in the response?",
            "1️⃣ Entirely inaccurate or off-topic. "
            "2️⃣ Mostly inaccurate; few correct elements. "
            "3️⃣ Partially accurate; some errors or omissions. "
            "4️⃣ Largely accurate with minor issues. "
            "5️⃣ Completely accurate and relevant."
        )
    },
    {
        "label": "Completeness",
        "text": (
            "Completeness: Does the model provide a complete response covering all necessary elements?",
            "1️⃣ Major omissions; response is inadequate. "
            "2️⃣ Missing key content. "
            "3️⃣ Covers the basics but lacks depth. "
            "4️⃣ Mostly complete; minor omissions. "
            "5️⃣ Fully complete; no relevant information missing."
        )
    },
    {
        "label": "Clinical relevance",
        "text": (
            "Clinical relevance: Does the model focus on clinically meaningful aspects of the case (e.g., appropriate drug choices, patient subgroups, relevant outcomes)?",
            "1️⃣ Focuses on tangential or irrelevant issues. "
            "2️⃣ Includes few clinically related points, overall focus unclear. "
            "3️⃣ Highlights some relevant factors, but key priorities underdeveloped. "
            "4️⃣ Centers on important clinical aspects with minor omissions. "
            "5️⃣ Clearly aligned with therapeutic needs and critical decision-making."
        )
    }
]


criteria_for_comparison = [
    {
        "label": "Task success",
        "text": (
            "Task success: Did the model successfully complete the therapeutic task it was given?<br>"
        )
    },
    {
        "label": "Justification helpfulness",
        "text": (
            "Justification helpfulness: Is the model’s rationale helpful in determining whether the answer is correct?<br>"
        )
    },
    {
        "label": "Cognitive traceability",
        "text": (
            "Cognitive traceability: Are the intermediate reasoning steps and decision factors interpretable and traceable?<br>"
        )
    },
    {
        "label": "Possibility of harm",
        "text": (
            "Possibility of harm: Based on the model’s output and rationale, is there a risk that the recommendation could cause clinical harm?<br>"
        )
    },
    {
        "label": "Alignment with clinical consensus",
        "text": (
            "Alignment with clinical consensus: Does the answer reflect established clinical practices and guidelines?<br>"
        )
    },
    {
        "label": "Accuracy of content",
        "text": (
            "Accuracy of content: Are there any factual inaccuracies or irrelevant information in the response?<br>"
        )
    },
    {
        "label": "Completeness",
        "text": (
            "Completeness: Does the model provide a complete response covering all necessary elements?<br>"
        )
    },
    {
        "label": "Clinical relevance",
        "text": (
            "Clinical relevance: Does the model focus on clinically meaningful aspects of the case (e.g., appropriate drug choices, patient subgroups, relevant outcomes)?<br>"
        )
    }
]

mapping = {   #for pairwise mapping between model comparison selections
    "👈 Model A": "A",
    "👉 Model B": "B",
    "🤝 Tie": "tie",
    "👎 Neither model did well": "neither"
}

def preprocess_question_id(question_id):
    if isinstance(question_id, str):
        return question_id
    elif isinstance(question_id, list) and len(question_id) == 1:
        return question_id[0]
    else:
        print("Error: Invalid question ID format. Expected a string or a single-element list.")
        return None

def get_evaluator_questions(evaluator_id, all_files, evaluator_directory, our_methods):

    # Filter to only the files in that directory
    evaluator_files = [f for f in all_files if f.startswith(f"{evaluator_directory}/")]
    data_by_filename = {}
    for remote_path in evaluator_files:
        local_path = hf_hub_download(
            repo_id=REPO_ID,
            repo_type="dataset",
            revision="main", #fetches the most recent version of the dataset each time this command is called
            filename=remote_path,
            # force_download=True,
            token = os.getenv("HF_TOKEN")
        )
        with open(local_path, "r") as f:
            model_name_key = os.path.basename(remote_path).replace('.json', '')
            data_by_filename[model_name_key] = json.load(f)

    
    evaluator_question_ids = []
    # Assuming 'TxAgent-T1-Llama-3.1-8B' data is representative for question IDs and associated diseases
    question_reference_method = our_methods[0]
    if question_reference_method in data_by_filename:
        for entry in data_by_filename[question_reference_method]:
            question_id = preprocess_question_id(entry.get("id"))
            evaluator_question_ids.append(question_id)
    # Handle case where no relevant questions are found based on specialty
    if not evaluator_question_ids:
        return [], data_by_filename

    #FINALLY, MAKE SURE THEY DIDNT ALREADY FILL IT OUT. Must go through every tuple of (question_ID, TxAgent, other model) where other model could be any of the other files in data_by_filename
    model_names = [key for key in data_by_filename.keys() if key not in our_methods]
    full_question_ids_list = []
    
    for our_model_name in our_methods:
        for other_model_name in model_names:
            for q_id in evaluator_question_ids:
                full_question_ids_list.append((q_id, our_model_name, other_model_name))

    results_df = read_sheet_to_df(custom_sheet_name=str(TXAGENT_RESULTS_SHEET_BASE_NAME + f"_{str(evaluator_id)}"))
    if (results_df is not None) and (not results_df.empty):
        # collect all (question_ID, other_model) pairs already seen
        matched_pairs = set()
        for _, row in results_df.iterrows():
            q = row["Question ID"]
            # pick whichever response isn’t 'TxAgent-T1-Llama-3.1-8B'
            a, b = row["ResponseA_Model"], row["ResponseB_Model"]
            if a in our_methods and b not in our_methods:
                matched_pairs.add((q, a, b))
            elif b in our_methods and a not in our_methods:
                matched_pairs.add((q, b, a))

        # filter out any tuple whose (q_id, other_model) was already matched
        full_question_ids_list = [
            (q_id, our_model, other_model)
            for (q_id, our_model, other_model) in full_question_ids_list
            if (q_id, our_model, other_model) not in matched_pairs
        ]
        print(f"Length of filtered question IDs: {len(full_question_ids_list)}")
        

    return full_question_ids_list, data_by_filename

def get_next_eval_question(
    name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, evaluator_id,
    our_methods,
    return_user_info=True,  # Whether to return user_info tuple
    include_correct_answer=True  # Whether to return correct_answer
):
    # ADDED: Validate that name and email are non-empty before proceeding
    if not name or not email or not evaluator_id or not specialty_dd or not years_exp_radio:
        return gr.update(visible=True), gr.update(visible=False), None, "Please fill out all the required fields (name, email, evaluator ID, specialty, years of experience). If you are not a licensed physician with a specific specialty, please choose the specialty that most closely aligns with your biomedical expertise.", gr.Chatbot(), gr.Chatbot(), gr.HTML(),gr.Markdown(),gr.State(),gr.update(visible=False), ""

    question_map_path = hf_hub_download(
        repo_id=REPO_ID,
        filename=EVALUATOR_MAP_DICT,
        repo_type="dataset",       # or omit if it's a Model/Space
        # force_download=True,       # ← always fetch new copy
        revision="main",            # branch/tag/commit, fetches the most recent version of the dataset each time this command is called
        token = os.getenv("HF_TOKEN")
    )

    # Load the question map from the downloaded file
    with open(question_map_path, 'r') as f:
        question_map = json.load(f)

    #retrieve data from HF
    evaluator_directory = question_map.get(evaluator_id, None)
    if evaluator_directory is None:
        return gr.update(visible=True), gr.update(visible=False), None, "Invalid Evaluator ID, please try again.", gr.Chatbot(), gr.Chatbot(), gr.HTML(),gr.State(),gr.update(visible=False),"" 
    all_files = list_repo_files(
        repo_id=REPO_ID,
        repo_type="dataset",
        revision="main",
        token = os.getenv("HF_TOKEN")
    )

    # Get available questions for the evaluator
    full_question_ids_list, data_by_filename = get_evaluator_questions(
        evaluator_id, all_files, evaluator_directory, our_methods)

    if len(full_question_ids_list) == 0:
        return None, None, None, None, None, None, 0

    full_question_ids_list = sorted(full_question_ids_list, key=lambda x: str(x[0])+str(x[1]))
    #selected question is the first element
    q_id, our_model_name, other_model_name = full_question_ids_list[0]
    print("Selected question ID:", q_id)

    # Build model answer lists
    models_list = []
    
    
    txagent_matched_entry = next(
        (entry for entry in data_by_filename[our_model_name] if preprocess_question_id(entry.get("id")) == q_id),
        None
    )
    our_model = {
        "model": our_model_name,
        "reasoning_trace": txagent_matched_entry.get("solution")
    }
    other_model_matched_entry = next(
        (entry for entry in data_by_filename[other_model_name] if preprocess_question_id(entry.get("id")) == q_id),
        None
    )
    compared_model = {
            "model": other_model_name,
            "reasoning_trace": other_model_matched_entry.get("solution")
        }
    
    models_list = [our_model, compared_model]
    random.shuffle(models_list)

    question_for_eval = {
        "question": txagent_matched_entry.get("question"),
        "id": q_id,
        "models": models_list,
    }
    if include_correct_answer:
        question_for_eval["correct_answer"] = txagent_matched_entry.get("correct_answer")

    # Prepare Gradio components
    chat_A_value = format_chat(question_for_eval['models'][0]['reasoning_trace'], tool_database_labels)
    chat_B_value = format_chat(question_for_eval['models'][1]['reasoning_trace'], tool_database_labels)
    prompt_text = question_for_eval['question']

    page1_prompt = gr.HTML(f'<div style="background-color: #FFEFD5; border: 2px solid #FF8C00; padding: 10px; border-radius: 5px; color: black;"><strong style="color: black;">Prompt:</strong> {prompt_text}</div>')
    page1_reference_answer = gr.Markdown(txagent_matched_entry.get("correct_answer")) if include_correct_answer else None
    chat_a = gr.Chatbot(
        value=chat_A_value,
        type="messages",
        height=400,
        label="Model A Response",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False
    )
    chat_b = gr.Chatbot(
        value=chat_B_value,
        type="messages",
        height=400,
        label="Model B Response",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False
    )

    user_info = (name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, q_id, evaluator_id)  if return_user_info else None
    return user_info, chat_a, chat_b, page1_prompt, page1_reference_answer, question_for_eval, len(full_question_ids_list)

def go_to_page0_from_minus1():
    return gr.update(visible=False), gr.update(visible=True)

def go_to_eval_progress_modal(name, email, evaluator_id, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id):

    # ADDED: Validate that name and email are non-empty before proceeding
    if not name or not email or not evaluator_id or not specialty_dd or not years_exp_radio:
        return gr.update(visible=True), gr.update(visible=False), None, "Please fill out all the required fields (name, email, evaluator ID, specialty, years of experience). If you are not a licensed physician with a specific specialty, please choose the specialty that most closely aligns with your biomedical expertise.", gr.Chatbot(), gr.Chatbot(), gr.HTML(),gr.Markdown(),gr.State(),gr.update(visible=False), ""

    user_info, chat_a, chat_b, page1_prompt, page1_reference_answer, question_for_eval, remaining_count = get_next_eval_question(
        name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, evaluator_id, our_methods
    )
    if remaining_count == 0:
        return gr.update(visible=True), gr.update(visible=False), None, "Based on your submitted data, you have no more questions to evaluate. You may exit the page; we will follow-up if we require anything else from you. Thank you!", gr.Chatbot(), gr.Chatbot(), gr.HTML(),gr.Markdown(),gr.State(),gr.update(visible=False),""
    return gr.update(visible=True), gr.update(visible=False), user_info,"", chat_a, chat_b, page1_prompt, page1_reference_answer, question_for_eval, gr.update(visible=True), f"You are about to evaluate the next question. You have {remaining_count} question(s) remaining to evaluate."

#goes to page 1 from confirmation modal that tells users how many questions they have left to evaluate
def go_to_page1():
    """
    Shows page 1
    """

    # Return updates to hide modal, hide page 0, show page 1, populate page 1, and set final state
    updates = [
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=True),
    ]
    return updates


# Callback to transition from Page 1 to Page 2.
def go_to_page2(data_subset_state,*pairwise_values):
    # pairwise_values is a tuple of values from each radio input.
    criteria_count = len(criteria_for_comparison)
    pairwise_list = list(pairwise_values[:criteria_count])
    comparison_reasons_list = list(pairwise_values[criteria_count:])

    #gradio components to display previous page results on next page
    pairwise_results_for_display = [gr.Markdown(f"***As a reminder, your pairwise comparison answer for this criterion was: {pairwise_list[i]}. Your answer choices will be restricted based on your comparison answer, but you may go back and change the comparison answer if you wish.***") for i in range(len(criteria))]

    if any(answer is None for answer in pairwise_list):
        return (gr.update(visible=True), gr.update(visible=False), None, None, "Error: Please select an option for every pairwise comparison.", gr.Chatbot(), gr.Chatbot(), gr.HTML(), gr.Markdown()) + tuple(pairwise_results_for_display)

    chat_A_value = format_chat(data_subset_state['models'][0]['reasoning_trace'], tool_database_labels)
    chat_B_value = format_chat(data_subset_state['models'][1]['reasoning_trace'], tool_database_labels)
    prompt_text = data_subset_state['question']

    # Construct the question-specific elements of the rating page (page 2)
    chat_A_rating = gr.Chatbot(
                    value=chat_A_value,
                    type="messages",
                    height=400,
                    label="Model A Response",
                    show_copy_button=False,
                    render_markdown=True
                )

    chat_B_rating = gr.Chatbot(
                    value=chat_B_value,
                    type="messages",
                    height=400,
                    label="Model B Response",
                    show_copy_button=False,
                    render_markdown=True
                )

    page2_prompt = gr.HTML(f'<div style="background-color: #FFEFD5; border: 2px solid #FF8C00; padding: 10px; border-radius: 5px; color: black;"><strong style="color: black;">Prompt:</strong> {prompt_text}</div>')
    page2_reference_answer = gr.Markdown(data_subset_state['correct_answer'])
    
    return (gr.update(visible=False), gr.update(visible=True), pairwise_list, comparison_reasons_list, "", chat_A_rating, chat_B_rating, page2_prompt, page2_reference_answer) + tuple(pairwise_results_for_display)


# Callback to store scores for Response A.
def store_A_scores(*args):
    # Unpack the arguments: first half are scores, second half are checkboxes.
    num = len(args) // 2
    scores = list(args[:num])
    unquals = list(args[num:])
    return scores, unquals

# Callback to transition from Page 2 to Page 3.
def go_to_page3():
    return gr.update(visible=False), gr.update(visible=True)

# Updated validation callback that ignores criteria with 'Unable to Judge'
def validate_ratings(pairwise_choices, *args):
    num_criteria = len(criteria)
    ratings_A_list = list(args[:num_criteria])
    ratings_B_list = list(args[num_criteria:])
    if any(r is None for r in ratings_A_list) or any(r is None for r in ratings_B_list):
        return "Error: Please provide ratings for both responses for every criterion.", "Error: Please provide ratings for both responses for every criterion."
    error_msgs = []
    for i, choice in enumerate(pairwise_choices):
        score_a = ratings_A_list[i]
        score_b = ratings_B_list[i]
        # Skip criteria if either rating is "Unable to Judge"
        if score_a == "Unable to Judge" or score_b == "Unable to Judge":
            continue
        # Convert string scores to integers for comparison.
        score_a = int(score_a)
        score_b = int(score_b)
        if choice == "👈 Model A" and score_a < score_b:
            error_msgs.append(f"Criterion {i+1} ({criteria[i]['label']}): You selected A as better but scored A lower than B.")
        elif choice == "👉 Model B" and score_b < score_a:
            error_msgs.append(f"Criterion {i+1} ({criteria[i]['label']}): You selected B as better but scored B lower than A.")
        elif choice == "🤝 Tie" and score_a != score_b:
            error_msgs.append(f"Criterion {i+1} ({criteria[i]['label']}): You selected Tie but scored A and B differently.")

    if error_msgs:
        err_str = "\n".join(error_msgs)
        return err_str, err_str
    else:
        return "No errors in responses; feel free to submit!", "No errors in responses; feel free to submit!"

# # Additional callback to handle submission results.
def toggle_slider(is_unqualified):
    # When the checkbox is checked (True), set interactive to False to disable the slider.
    return gr.update(interactive=not is_unqualified)

#show reference answer
def toggle_reference(selection):
    if selection == "Show Reference Answer":
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

#nonsense button helper
def mark_invalid_question(btn_clicked_status):
    new_status = not btn_clicked_status
    if new_status == True:
        return new_status, gr.update(value="Undo: Correct Question", variant="primary")
    else:
        return new_status, gr.update(value="Wrong Question",variant="stop")

centered_col_css = """
#centered-column {
    margin-left: auto;
    margin-right: auto;
    max-width: 800px; /* Adjust this width as desired */
    width: 100%;
}
#participate-btn {
    background-color: purple !important;
    color: white !important;
    border-color: purple !important;
}
#answer-reference-btn {
    background-color: #E6E6FA !important;
    color: white !important;
    border-color: #E6E6FA !important;
}
#clear_btn {
    background-color: #F08080 !important;
    color: white !important;
    border-color: #F08080 !important;
}
.reference-box {
    border: 1px solid #ccc;
    padding: 10px;
    border-radius: 5px;
}
.short-btn { min-width: 80px !important; max-width: 120px !important; width: 100px !important; padding-left: 4px !important; padding-right: 4px !important; }
.light-stop-btn { background-color: #ffcccc !important; color: #b30000 !important; border-color: #ffcccc !important; }

"""
with gr.Blocks(css=centered_col_css) as demo:
    # States to save information between pages.
    user_info_state = gr.State()
    pairwise_state = gr.State()
    scores_A_state = gr.State()
    comparison_reasons = gr.State()
    nonsense_btn_clicked = gr.State(False)
    unqualified_A_state = gr.State()
    data_subset_state = gr.State()

    # Load specialty data
    specialties_path = "specialties.json"
    subspecialties_path = "subspecialties.json"

    try:
        with open(specialties_path, 'r') as f:
            specialties_list = json.load(f)
        with open(subspecialties_path, 'r') as f:
            subspecialties_list = json.load(f)
    except FileNotFoundError:
        print(f"Error: Could not find specialty files at {specialties_path} or {subspecialties_path}. Please ensure these files exist.")
        # Provide default empty lists or handle the error as appropriate
        specialties_list = ["Error loading specialties"]
        subspecialties_list = ["Error loading subspecialties"]
    except json.JSONDecodeError:
        print(f"Error: Could not parse JSON from specialty files.")
        specialties_list = ["Error parsing specialties"]
        subspecialties_list = ["Error parsing subspecialties"]
    
    # Page -1: Page to link them to question submission form or evaluation portal
    with gr.Column(visible=True, elem_id="page-1") as page_minus1:
        gr.HTML("""
        <div>
            <h1>TxAgent Evaluation Portal</h1>
            <p>Welcome to the TxAgent Evaluation Portal.</p>
        </div>
        """)
        with gr.Row():
            participate_eval_btn = gr.Button(
            value="🌟 Participate in TxAgent Evaluation 🌟",
            variant="primary",
            size="lg",
            elem_id="participate-btn"
            )
        gr.HTML(TxAgent_Project_Page_HTML)

    # Page 0: Welcome / Informational page.
    with gr.Column(visible=False, elem_id="page0") as page0:
        gr.Markdown("## Welcome to the TxAgent Evalution Study!")
        gr.Markdown("Please read the following instructions and then enter your information to begin:")
        # Existing informational markdown...
        gr.Markdown("""
        - Each session requires a minimum commitment of 5-10 minutes to complete one question.
        - If you wish to evaluate multiple questions, you may do so; you will never be asked to re-evaluate questions you have already seen.
        - When evaluating a question, you will be asked to compare the responses of two different models to the question and then rate each model's response on a scale of 1-5.
        - If you feel that a question does not make sense or is not biomedically relevant, there is a RED BUTTON at the top of the first model comparison page to indicate this
        - You may use the Back and Next buttons at the bottom of each page to edit any of your responses before submitting.
        - You may use the Home Page button at the bottom of each page to the home page. Your progress will be saved but not submitted.
        - You must submit your answers to the current question before moving on to evaluate the next question.
        - You may stop in between questions and return at a later time; however, you must submit your answers to the current question if you would like them saved.
        - Please review the example question and LLM model response below:

        """)
        # Assume 'your_image.png' is in the same directory
        with open("anatomyofAgentResponse.jpg", "rb") as image_file:
            img = Image.open(image_file)
            new_size = (int(img.width * 0.5), int(img.height * 0.5))
            img = img.resize(new_size, Image.LANCZOS)
            buffer = io.BytesIO()
            img.save(buffer, format="PNG")
            encoded_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
            #encoded_string = base64.b64encode(image_file.read()).decode("utf-8")

        image_html = f'<div style="text-align:center;"><img src="data:image/png;base64,{encoded_string}" alt="Your Image"></div>'
        ReasoningTraceExampleHTML = f"""
        <div>
            {image_html}
        </div>
        """
        gr.HTML(ReasoningTraceExampleHTML)
        gr.Markdown("""By clicking 'Next' below, you will start the study, with your progress saved after submitting each question. If you have any other questions or concerns, please contact us directly. Thank you for your participation!
        """)
        gr.Markdown("## Please enter your information to get a question to evaluate. Please use the same email every time you log onto this evaluation portal, as we use your email to prevent showing repeat questions.")
        name = gr.Textbox(label="Name (required)")
        email = gr.Textbox(label="Email (required). Please use the same email every time you log onto this evaluation portal, as we use your email to prevent showing repeat questions.")
        evaluator_id = gr.Textbox(label="Evaluator ID (required). This is the four-digit ID you received from us for the evaluation study. If you do not have an Evaluator ID or are unsure about your Evaluator ID, please contact us.")
        specialty_dd = gr.Dropdown(choices=specialties_list, label="Primary Medical Specialty (required). Go to https://www.abms.org/member-boards/specialty-subspecialty-certificates/ for categorization)", multiselect=True)
        subspecialty_dd = gr.Dropdown(choices=subspecialties_list, label="Subspecialty (if applicable). Go to https://www.abms.org/member-boards/specialty-subspecialty-certificates/ for categorization)", multiselect=True)
        npi_id = gr.Textbox(label="National Provider Identifier ID (optional). Got to https://npiregistry.cms.hhs.gov/search to search for your NPI ID. If you do not have an NPI ID, please leave this blank.")
        years_exp_radio = gr.Radio(
            choices=["0-2 years", "3-5 years", "6-10 years", "11-20 years", "20+ years", "Not Applicable"],
            label="How many years have you been involved in clinical and/or research activities related to your biomedical area of expertise? (required)"
        )
        exp_explanation_tb = gr.Textbox(label="Please briefly explain your expertise/experience relevant to evaluating AI for clinical decision support (optional)")

        page0_error_box = gr.Markdown("")
        with gr.Row():
            next_btn_0 = gr.Button("Next")
        with gr.Row():
            home_btn_0 = gr.Button("Home (your registration info will be saved)")


    with Modal(visible=False, elem_id="confirm_modal") as eval_progress_modal:
        eval_progress_text = gr.Markdown("You have X questions remaining.")
        eval_progress_proceed_btn = gr.Button("OK, proceed to question evaluation")

    # Page 1: Pairwise Comparison.
    with gr.Column(visible=False) as page1:
        gr.Markdown("## Part 1/2: Pairwise Comparison") #Make the number controlled by question indexing!
        page1_prompt = gr.HTML()

        with gr.Accordion("Click to reveal a reference answer—this is just one correct solution; others are possible.", open=False, elem_id="answer-reference-btn"):
            page1_reference_answer = gr.Markdown(
                """
                **Reference Answer:**

                This is the reference answer content.
                """,
                elem_classes="reference-box"
            )
            

        # Add small red button under the prompt
        with gr.Row():
            nonsense_btn = gr.Button(
                "Wrong Question?",
                size="sm",
                variant="stop",  # red variant
                elem_id="invalid-question-btn",
                elem_classes=["short-btn"]
            )
            gr.Markdown(
                "<span style='color: #b30000; font-weight: bold;'>Click the button if you think this question does not make sense or is not biomedically-relevant</span>",
                render=True
            )

        nonsense_btn.click(
            fn=mark_invalid_question,
            inputs=[nonsense_btn_clicked],
            outputs=[nonsense_btn_clicked, nonsense_btn],
            queue=False,
        )

        with gr.Row():
            # ADDED: Use gr.Chatbot to display the scrollable chat window for Response A.
            with gr.Column():
                gr.Markdown("**Model A Response:**")  # Already bold label.
                chat_a = gr.Chatbot(
                    value=[], # Placeholder for chat history
                    type="messages",
                    height=400,
                    label="Model A Response",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,  # Required for markdown/HTML support in messages
                    avatar_images=None,    # Optional: omit user/assistant icons
                    rtl=False
                )
            # ADDED: Use gr.Chatbot to display the scrollable chat window for Response B.
            with gr.Column():
                gr.Markdown("**Model B Response:**")
                chat_b = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=400,
                    label="Model B Response",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,  # Required for markdown/HTML support in messages
                    avatar_images=None,    # Optional: omit user/assistant icons
                    rtl=False
                )
        gr.Markdown("<br><br>")
        gr.Markdown("### For each criterion, select which response did better:")
        comparison_reasons_inputs = []  # ADDED: list to store the free-text inputs
        pairwise_inputs = []
        for crit in criteria_for_comparison:
            with gr.Row():
                gr.Markdown(crit['text'])
                radio = gr.Radio(
                    choices=[
                        "👈 Model A",  # A
                        "👉 Model B",  # B
                        "🤝 Tie",      # tie
                        "👎 Neither model did well"  # neither
                    ],
                    label="Which is better?"
                )
                pairwise_inputs.append(radio)
            # ADDED: free text under each comparison
            text_input = gr.Textbox(label=f"Reasons for your selection (optional)")
            comparison_reasons_inputs.append(text_input)
        
        page1_error_box = gr.Markdown("")  # ADDED: display validation errors
        with gr.Row():
            back_btn_0 = gr.Button("Back")
            next_btn_1 = gr.Button("Next: Rate Responses")
        
        with gr.Row():
            home_btn_1 = gr.Button("Home Page (your progress on this question will be saved but not submitted)")  # ADDED: Home button on page11
    
    # Page 2: Combined Rating Page for both responses.
    with gr.Column(visible=False) as page2:
        gr.Markdown("## Part 2/2: Rate Model Responses")
        # ### EDIT: Show a highlighted prompt as on previous pages.
        page2_prompt = gr.HTML()
        with gr.Accordion("Click to reveal a reference answer—this is just one correct solution; others are possible.", open=False, elem_id="answer-reference-btn"):
            page2_reference_answer = gr.Markdown(
                """
                **Reference Answer:**

                This is the reference answer content.
                """,
                elem_classes="reference-box"
            )
        # ### EDIT: Display both responses side-by-side using Chatbot windows.
        with gr.Row():
            with gr.Column():
                gr.Markdown("**Model A Response:**")
                chat_a_rating = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=400,
                    label="Model A Response",
                    show_copy_button=False,
                    render_markdown=True
                )
            with gr.Column():
                gr.Markdown("**Model B Response:**")
                chat_b_rating = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=400,
                    label="Model B Response",
                    show_copy_button=False,
                    render_markdown=True
                )
        gr.Markdown("<br><br>")
        gr.Markdown("### For each criterion, select your ratings for each model response:")
        # ### EDIT: For each criterion, create a row with two multiple-choice sets (left: Response A, right: Response B) separated by a border.
        ratings_A = []  # to store the radio components for response A
        ratings_B = []  # to store the radio components for response B
        
        def restrict_choices(pairwise_list, index, score_a, score_b):
            """
            Returns (update_for_A, update_for_B).
            Enforces rating constraints based on the pairwise choice for the given criterion index.
            """
            # Get the specific pairwise choice for this criterion using the index
            # Add error handling in case the state/list is not ready or index is wrong
            if not pairwise_list or index >= len(pairwise_list):
                pairwise_choice = None
            else:
                pairwise_choice = pairwise_list[index]

            base = ["1","2","3","4","5","Unable to Judge"]
            # Default: no restrictions unless explicitly set
            upd_A = gr.update(choices=base)
            upd_B = gr.update(choices=base)

            # Skip if no meaningful pairwise choice or either score is "Unable to Judge"
            if pairwise_choice is None or pairwise_choice == "👎 Neither model did well" or (score_a is None and score_b is None):
                # If one score is UJ but the other isn't, AND it's a Tie, we might still want to restrict the non-UJ one later?
                # For now, keep it simple: if either is UJ or choice is Neither/None, don't restrict.
                return upd_A, upd_B

            # Helper to parse int safely
            def to_int(x):
                try: return int(x)
                except (ValueError, TypeError): return None

            a_int = to_int(score_a)
            b_int = to_int(score_b)

            # --- Apply Restrictions ---
            if pairwise_choice == "👈 Model A":
                # B must be ≤ A (if A is numeric)
                if a_int is not None: #it is None if unable to judge
                    allowed_b_choices = [str(i) for i in range(1, a_int + 1)] + ["Unable to Judge"]
                    current_b = score_b if score_b in allowed_b_choices else None # Keep current valid choice
                    upd_B = gr.update(choices=allowed_b_choices, value=current_b)
                # If A is UJ or non-numeric, B is unrestricted by this rule
                # else: upd_B remains gr.update(choices=base)
                if b_int is not None:
                    # A must be >= B (if B is numeric)
                    allowed_a_choices = [str(i) for i in range(b_int, 6)] + ["Unable to Judge"]
                    current_a = score_a if score_a in allowed_a_choices else None # Keep current valid choice
                    upd_A = gr.update(choices=allowed_a_choices, value=current_a)
                # If B is UJ or non-numeric, A is unrestricted by this rule
                # else: upd_A remains gr.update(choices=base)

            elif pairwise_choice == "👉 Model B":
                # A must be ≤ B (if B is numeric)
                if b_int is not None:
                    allowed_a_choices = [str(i) for i in range(1, b_int + 1)] + ["Unable to Judge"]
                    current_a = score_a if score_a in allowed_a_choices else None # Keep current valid choice
                    upd_A = gr.update(choices=allowed_a_choices, value=current_a)
                # If B is UJ or non-numeric, A is unrestricted by this rule
                # else: upd_A remains gr.update(choices=base)
                if a_int is not None:
                    # B must be >= A (if A is numeric)
                    allowed_b_choices = [str(i) for i in range(a_int, 6)] + ["Unable to Judge"]
                    current_b = score_b if score_b in allowed_b_choices else None # Keep current valid choice
                    upd_B = gr.update(choices=allowed_b_choices, value=current_b)
                # If A is UJ or non-numeric, B is unrestricted by this rule
                # else: upd_B remains gr.update(choices=base)

            elif pairwise_choice == "🤝 Tie":
                # If both are numeric, they must match. Enforce based on the one that *just changed*.
                # If one changes to numeric, force the other (if also numeric) to match.
                # If one changes to UJ, the other is unrestricted.
                if a_int is not None:
                    upd_B = gr.update(choices=[score_a])
                elif score_a == "Unable to Judge":
                    upd_B = gr.update(choices=["Unable to Judge"])
                if b_int is not None:
                    upd_A = gr.update(choices=[score_b])
                elif score_b == "Unable to Judge":
                    upd_A = gr.update(choices=["Unable to Judge"])

            return upd_A, upd_B

        def clear_selection():
            return None, None
        
        pairwise_results_for_display = [gr.Markdown(render=False) for _ in range(len(criteria))]
        indices_for_change = []
        for i, crit in enumerate(criteria):
            index_component = gr.Number(value=i, visible=False, interactive=False)
            indices_for_change.append(index_component)

            with gr.Column(elem_id="centered-column"):
                gr.Markdown(f'<div style="text-align: left;">{crit["text"][0]}</div>')
                gr.Markdown(f'<div style="text-align: left;">{crit["text"][1]}</div>')
                pairwise_results_for_display[i].render()
            with gr.Row():
                with gr.Column(scale=1):
                    rating_a = gr.Radio(choices=["1", "2", "3", "4", "5", "Unable to Judge"],
                                            label=f"Score for Response A - {crit['label']}",
                                            interactive=True)
                with gr.Column(scale=1):
                    rating_b = gr.Radio(choices=["1", "2", "3", "4", "5", "Unable to Judge"],
                                            label=f"Score for Response B - {crit['label']}",
                                            interactive=True)
            with gr.Row():
                clear_btn = gr.Button("Clear Selection", size="sm", elem_id="clear_btn")
                clear_btn.click(fn=clear_selection, outputs=[rating_a,rating_b])

                # wire each to re‐restrict the other on change
                rating_a.change(
                    fn=restrict_choices,
                    inputs=[ pairwise_state, index_component, rating_a, rating_b ],
                    outputs=[ rating_a, rating_b ]
                )
                rating_b.change(
                    fn=restrict_choices,
                    inputs=[ pairwise_state, index_component, rating_a, rating_b ],
                    outputs=[ rating_a, rating_b ]
                )
                ratings_A.append(rating_a)
                ratings_B.append(rating_b)
        with gr.Row():
            back_btn_2 = gr.Button("Back")
            submit_btn = gr.Button("Submit (Note: Once submitted, you cannot edit your responses)", elem_id="submit_btn")

        with gr.Row():
            home_btn_2 = gr.Button("Home Page (your progress on this question will be saved but not submitted)")

        result_text = gr.Textbox(label="Validation Result")
    
    # Final Page: Thank you message.
    with gr.Column(visible=False, elem_id="final_page") as final_page:
        gr.Markdown("## You have no questions left to evaluate. Thank you for your participation!")
    
    # Error Modal: For displaying validation errors.
    with Modal("Error", visible=False, elem_id="error_modal") as error_modal:
        error_message_box = gr.Markdown()
        ok_btn = gr.Button("OK")
        # Clicking OK hides the modal.
        ok_btn.click(lambda: gr.update(visible=False), None, error_modal)
    
    # Confirmation Modal: Ask for final submission confirmation.
    with Modal("Confirm Submission", visible=False, elem_id="confirm_modal") as confirm_modal:
        gr.Markdown("Are you sure you want to submit? Once submitted, you cannot edit your responses.")
        with gr.Row():
            yes_btn = gr.Button("Yes, please submit")
            cancel_btn = gr.Button("Cancel")
    
    # --- Define Callback Functions for Confirmation Flow ---
    def build_row_dict(data_subset_state, user_info, pairwise, comparisons_reasons, nonsense_btn_clicked, *args):
        num_criteria = len(criteria)
        ratings_A_vals = list(args[:num_criteria])
        ratings_B_vals = list(args[num_criteria:])

        prompt_text = data_subset_state['question']
        response_A_model = data_subset_state['models'][0]['model']
        response_B_model = data_subset_state['models'][1]['model']

        timestamp = datetime.datetime.now().isoformat()
        row = {
            "Timestamp": timestamp,
            "Name": user_info[0],
            "Email": user_info[1],
            "Evaluator ID": user_info[8],
            "Specialty": str(user_info[2]),
            "Subspecialty": str(user_info[3]),
            "Years of Experience": user_info[4],
            "Experience Explanation": user_info[5],
            "NPI ID": user_info[6],
            "Question ID": user_info[7],
            "Prompt": prompt_text,
            "ResponseA_Model": response_A_model,
            "ResponseB_Model": response_B_model,
            "Question Makes No Sense or Biomedically Irrelevant": nonsense_btn_clicked,
        }

        pairwise = [mapping.get(val, val) for val in pairwise]
        for i, crit in enumerate(criteria):
            label = crit['label']
            row[f"Criterion_{label} Comparison: Which is Better?"] = pairwise[i]
            row[f"Criterion_{label} Comments"] = comparisons_reasons[i]
            row[f"ScoreA_{label}"] = ratings_A_vals[i]
            row[f"ScoreB_{label}"] = ratings_B_vals[i]

        return row

    def final_submit(data_subset_state, user_info, pairwise, comparisons_reasons, nonsense_btn_clicked, *args):
        # --- Part 1: Submit the current results (Existing Logic) ---
        row_dict = build_row_dict(data_subset_state, user_info, pairwise, comparisons_reasons, nonsense_btn_clicked, *args)
        append_to_sheet(user_data=None, custom_row_dict=row_dict, custom_sheet_name=str(TXAGENT_RESULTS_SHEET_BASE_NAME + f"_{evaluator_id}"), add_header_when_create_sheet=True)
        name, email, specialty, subspecialty, years_exp_radio, exp_explanation_tb, npi_id, _, evaluator_id = user_info
        user_info_new, chat_a, chat_b, page1_prompt, page1_reference_answer, question_for_eval, remaining_count = get_next_eval_question(
            name, email, specialty, subspecialty, years_exp_radio, exp_explanation_tb, npi_id, our_methods
        )

        if remaining_count == 0:
            return (
                gr.update(visible=False),  # page0 (Hide)
                gr.update(visible=False),  # page2 (Hide)
                gr.update(visible=False),  # confirm_modal
                gr.update(visible=False), 
                "",                       
                gr.update(visible=True),   # final_page (Show)
                "",
                None,
                None,
                None,
                None,
                None,
                user_info_new,
            )
        return (
            gr.update(visible=False),  # page0 (Hide)
            gr.update(visible=False),  # page2 (Hide)
            gr.update(visible=False),  # confirm_modal (Hide)
            gr.update(visible=True),   # eval_progress_modal (Show)
            f"Submission successful! You have {remaining_count} question(s) remaining to evaluate. You may exit the page and return later if you wish.", # eval_progress_text
            gr.update(visible=False),  # final_page (Hide)
            "",         
            chat_a,
            chat_b,
            page1_prompt,
            page1_reference_answer,
            question_for_eval,
            user_info_new
        )

    def cancel_submission():
        # Cancel final submission: just hide the confirmation modal.
        return gr.update(visible=False)
    
    def reset_everything_except_user_info():

        # 3) Reset all pairwise radios & textboxes
        reset_pairwise_radios = [gr.update(value=None) for i in range(len(criteria))]
        reset_pairwise_reasoning_texts  = [gr.update(value=None) for i in range(len(criteria))]

        # 4) Reset all rating radios
        reset_ratings_A = [gr.update(value=None) for i in range(len(criteria))]
        reset_ratings_B = [gr.update(value=None) for i in range(len(criteria))]

        return (
            # pages
            gr.update(visible=True),  # page0
            gr.update(visible=False), # final_page

            # states
            # gr.update(value=None),  # user_info_state
            gr.update(value=None),  # pairwise_state
            gr.update(value=None),  # scores_A_state
            gr.update(value=None),  # comparison_reasons
            gr.update(value=None),  # unqualified_A_state
            # gr.update(value=None),  # data_subset_state

            #page0 elements that need to be reset
            gr.update(value=""),  #page0_error_box 

            # page1 elements that need to be reset
            # gr.update(value=""),  #page1_prompt
            # gr.update(value=[]),  #chat_a
            # gr.update(value=[]),  #chat_b
            gr.update(value=""),  #page1_error_box

            # page2 elements that need to be reset
            gr.update(value=""),  #page2_prompt
            gr.update(value=""),  #page2_reference_answer
            gr.update(value=[]),  #chat_a_rating
            gr.update(value=[]),  #chat_b_rating
            gr.update(value=""),  #result_text

            #lists of gradio elements that need to be unrolled
            *reset_pairwise_radios,
            *reset_pairwise_reasoning_texts,
            *reset_ratings_A,
            *reset_ratings_B
        )
    
    # --- Define Transitions Between Pages ---


    # For the "Participate in Evaluation" button, transition to page0
    participate_eval_btn.click(
        fn=go_to_page0_from_minus1,
        inputs=None,
        outputs=[page_minus1, page0]
    )


    # Transition from Page 0 (Welcome) to Page 1.
    next_btn_0.click(
        fn=go_to_eval_progress_modal,
        inputs=[name, email, evaluator_id, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id],
        outputs=[page0, page1, user_info_state, page0_error_box, chat_a, chat_b, page1_prompt, page1_reference_answer, data_subset_state,eval_progress_modal,eval_progress_text],
        scroll_to_output=True
    )

    eval_progress_proceed_btn.click(
        fn=go_to_page1,
        inputs=None,
        outputs=[eval_progress_modal, page0, page1],
        scroll_to_output=True
    )

    #Home page buttons to simply shown page-1
    home_btn_0.click(lambda: (gr.update(visible=True), gr.update(visible=False)), None, [page_minus1, page0])
    home_btn_1.click(lambda: (gr.update(visible=True), gr.update(visible=False)), None, [page_minus1, page1])
    home_btn_2.click(lambda: (gr.update(visible=True), gr.update(visible=False)), None, [page_minus1, page2])
    
    # Transition from Page 1 to Page 0 (Back button).
    back_btn_0.click(
        fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
        inputs=None,
        outputs=[page0, page1]
    )
    
    # Transition from Page 1 (Pairwise) to the combined Rating Page (Page 2).
    next_btn_1.click(
        fn=go_to_page2,  # ### EDIT: Rename or update the function to simply pass the pairwise inputs if needed.
        inputs=[data_subset_state,*pairwise_inputs,*comparison_reasons_inputs],
        outputs=[page1, page2, pairwise_state, comparison_reasons, page1_error_box, chat_a_rating, chat_b_rating, page2_prompt, page2_reference_answer,*pairwise_results_for_display],
        scroll_to_output=True
    )
    
    # Transition from Rating Page (Page 2) back to Pairwise page.
    back_btn_2.click(
        fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
        inputs=None,
        outputs=[page1, page2],
        scroll_to_output=True
    )
    
    # --- Submission: Validate the Ratings and then Process the Result ---
    def process_result(result):
        # If validation passed, show the confirmation modal and proceed.
        if result == "No errors in responses; feel free to submit!":
            return (
                gr.update(),  # Show page 3
                gr.update(),  # Hide final page
                gr.update(visible=True),  # Show confirmation modal
                gr.update(visible=False),  # Hide error modal
                gr.update(value="")  # EDIT: Clear the error_message_box
            )
        else:
            # If validation fails, show the error modal and display the error in error_message_box.
            return (
                gr.update(),             # Keep page3 as is
                gr.update(),             # Keep final page unchanged
                gr.update(visible=False), # Hide confirmation modal
                gr.update(visible=True),  # Show error modal
                gr.update(value=result)   # EDIT: Update error_message_box with the validation error
            )

    # ### EDIT: Update the submission callback to use the new radio inputs.
    submit_btn.click(
        fn=validate_ratings,
        inputs=[pairwise_state, *ratings_A, *ratings_B],
        outputs=[error_message_box, result_text]
    ).then(
        fn=process_result,
        inputs=error_message_box,
        outputs=[page2, final_page, confirm_modal, error_modal, error_message_box],
        scroll_to_output=True
    )
    
    # Finalize submission if user confirms.
    question_submission_event = yes_btn.click(
        fn=final_submit,
        inputs=[data_subset_state, user_info_state, pairwise_state, comparison_reasons, nonsense_btn_clicked, *ratings_A, *ratings_B],
        outputs=[
            page0,                 # Controlled by final_submit return value 1
            page2,                 # Controlled by final_submit return value 2
            confirm_modal,         # Controlled by final_submit return value 3
            eval_progress_modal,   # Controlled by final_submit return value 4
            eval_progress_text,    # Controlled by final_submit return value 5
            final_page,            # Controlled by final_submit return value 6
            page0_error_box,
            chat_a,
            chat_b,
            page1_prompt,
            page1_reference_answer,
            data_subset_state,
            user_info_state,
        ],
        scroll_to_output=True
    )
    
    # Cancel final submission.
    cancel_btn.click(
        fn=cancel_submission,
        inputs=None,
        outputs=confirm_modal
    )

    # Reset everything and evaluate another question button
    question_submission_event.then(
        fn=reset_everything_except_user_info,
        inputs=[],
        outputs=[

            # states
            # user_info_state,
            pairwise_state,
            scores_A_state,
            comparison_reasons,
            unqualified_A_state,
            # data_subset_state,

            #page0 elements that need to be reset
            page0_error_box,

            # # page1 elements that need to be reset
            # page1_prompt,
            # chat_a,
            # chat_b,
            page1_error_box,

            # page2 elements that need to be reset
            page2_prompt,
            page2_reference_answer,
            chat_a_rating,
            chat_b_rating,
            result_text,

            #lists of gradio elements that need to be unrolled
            *pairwise_inputs,
            *comparison_reasons_inputs,
            *ratings_A,
            *ratings_B
        ]
    )

demo.launch(share=True, allowed_paths = ["."])