app.py

import gradio as gr
from gradio_modal import Modal
from huggingface_hub import hf_hub_download, list_repo_files
import os
import csv
import datetime
import 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"
CROWDSOURCING_DATA_DIRECTORY = "crowdsourcing_questions_0516"
TXAGENT_RESULTS_SHEET_BASE_NAME = "TxAgent_Human_Eval_Results_CROWDSOURCED_0516"
DISEASE_SPECIALTY_MAP_FILENAME = "disease_specialty_map.json"
DRUG_SPECIALTY_MAP_FILENAME = "drug_specialty_map.json"

DATASET_WEIGHTS = {
    "drugPC": 0.2,
    "treatment_clear": 0.8
}

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}"
                # Assuming JPEG, adjust if needed
                base64_url = f'data:image/jpeg;base64,{encoded_image}'
                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": "Did the model successfully complete the therapeutic task it was given?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Helpfulness of rationale",
        "text": "Is the model’s rationale helpful in determining whether the answer is correct?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Cognitive traceability",
        "text": "Are the intermediate reasoning steps and decision factors interpretable and traceable?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Possibility of harm",
        "text": "Based on the model’s output and rationale, is there a risk that the recommendation could cause clinical harm?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Alignment with clinical consensus",
        "text": "Does the answer reflect established clinical practices and guidelines?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Accuracy of content",
        "text": "Are there any factual inaccuracies or irrelevant information in the response?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Completeness",
        "text": "Does the model provide a complete response covering all necessary elements?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    },
    {
        "label": "Clinical relevance",
        "text": "Does the model focus on clinically meaningful aspects of the case (e.g., appropriate drug choices, patient subgroups, relevant outcomes)?",
        "scores": [
            "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.",
            "Unable to Judge."
        ]
    }
]


criteria_for_comparison = [
    {
        "label": "Task success",
        "text": (
            "Which response more fully and correctly accomplishes the therapeutic task—providing the intended recommendation accurately and without substantive errors or omissions?"
        )
    },
    {
        "label": "Helpfulness of rationale",
        "text": (
            "Which response offers a clearer, more detailed rationale that genuinely aids you in judging whether the answer is correct?"
        )
    },
    {
        "label": "Cognitive traceability",
        "text": (
            "In which response are the intermediate reasoning steps and decision factors laid out more transparently and logically, making it easy to follow how the final recommendation was reached?"
        )
    },
    {
        "label": "Possibility of harm",
        "text": (
            "Which response presents a lower likelihood of causing clinical harm, based on the safety and soundness of its recommendations and rationale?"
        )
    },
    {
        "label": "Alignment with clinical consensus",
        "text": (
            "Which response is more consistent with established clinical guidelines and widely accepted practice standards?"
        )
    },
    {
        "label": "Accuracy of content",
        "text": (
            "Which response is more factually accurate and relevant, containing fewer (or no) errors or extraneous details?"
        )
    },
    {
        "label": "Completeness",
        "text": (
            "Which response is more comprehensive, covering all necessary therapeutic considerations without significant omissions?"
        )
    },
    {
        "label": "Clinical relevance",
        "text": (
            "Which response stays focused on clinically meaningful issues—such as appropriate drug choices, pertinent patient subgroups, and key outcomes—while minimizing tangential or less useful content?"
        )
    }
]

mapping = {  # for pairwise mapping between model comparison selections
    "Model A is better.": "A",
    "Model B is better.": "B",
    "Both models are equally good.": "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(email, disease_map_data, drug_map_data, user_all_specs, all_files, evaluator_directory, our_methods):
    relevant_diseases = []
    for disease, specs in disease_map_data.items():
        disease_specs = set(specs.get('specialties', []))
        disease_subspecs = set(specs.get('subspecialties', []))

        # Check for intersection
        if user_all_specs.intersection(disease_specs) or user_all_specs.intersection(disease_subspecs):
            relevant_diseases.append(disease)

    relevant_drugs = []
    for drug, specs in drug_map_data.items():
        drug_specs = set(specs.get('specialties', []))
        drug_subspecs = set(specs.get('subspecialties', []))

        # Check for intersection
        if user_all_specs.intersection(drug_specs) or user_all_specs.intersection(drug_subspecs):
            relevant_drugs.append(drug)

    # 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",
            # fetches the most recent version of the dataset each time this command is called
            revision="main",
            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)

    # Filter questions based on relevant diseases derived from user specialties
    evaluator_question_ids = []
    relevant_diseases_lower = {disease.lower()
                               for disease in relevant_diseases}
    relevant_drugs_lower = {drug.lower() for drug in relevant_drugs}
    # 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"))
            dataset = entry.get("dataset", "")
            # Get diseases list, default to empty if missing
            question_diseases = entry.get("disease", [])
            # Get drugs list, default to empty if missing
            question_drugs = entry.get("drug", [])
            if question_id is not None and question_diseases and question_drugs:
                # Convert question diseases to lowercase and check for intersection
                question_diseases_lower = {
                    disease.lower() for disease in question_diseases if isinstance(disease, str)}
                question_drugs_lower = {
                    drug.lower() for drug in question_drugs if isinstance(drug, str)}

                if (
                    question_diseases_lower.intersection(
                        relevant_diseases_lower)
                    or question_drugs_lower.intersection(relevant_drugs_lower)
                ):
                    evaluator_question_ids.append((question_id, dataset))

    # 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, dataset) in evaluator_question_ids:
                full_question_ids_list.append(
                    (q_id, our_model_name, other_model_name, dataset))

    results_df = read_sheet_to_df(
        custom_sheet_name=str(TXAGENT_RESULTS_SHEET_BASE_NAME))
    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():
            if row["Email"] == email:
                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, dataset)
            for (q_id, our_model, other_model, dataset) 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, our_methods,
    return_user_info=True,  # Whether to return user_info tuple
    include_correct_answer=True  # Whether to return correct_answer
):
    # Merge specialties and subspecialties
    user_specialties = set(specialty_dd if isinstance(
        specialty_dd, list) else ([specialty_dd] if specialty_dd else []))
    user_subspecialties = set(subspecialty_dd if isinstance(
        subspecialty_dd, list) else ([subspecialty_dd] if subspecialty_dd else []))
    user_all_specs = user_specialties.union(user_subspecialties)

    evaluator_directory = CROWDSOURCING_DATA_DIRECTORY
    all_files = list_repo_files(
        repo_id=REPO_ID,
        repo_type="dataset",
        revision="main",
        token=os.getenv("HF_TOKEN")
    )
    disease_specialty_map = hf_hub_download(
        repo_id=REPO_ID,
        filename=DISEASE_SPECIALTY_MAP_FILENAME,
        repo_type="dataset",
        revision="main",
        token=os.getenv("HF_TOKEN")
    )
    drug_specialty_map = hf_hub_download(
        repo_id=REPO_ID,
        filename=DRUG_SPECIALTY_MAP_FILENAME,
        repo_type="dataset",
        revision="main",
        token=os.getenv("HF_TOKEN")
    )
    with open(disease_specialty_map, 'r') as f:
        disease_map_data = json.load(f)
    with open(drug_specialty_map, 'r') as f:
        drug_map_data = json.load(f)

    # Get available questions for the evaluator
    full_question_ids_list, data_by_filename = get_evaluator_questions(
        email, disease_map_data, drug_map_data, user_all_specs, all_files, evaluator_directory, our_methods
    )

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

    # Weighted random selection of a question
    weights = [DATASET_WEIGHTS[entry[-1]] for entry in full_question_ids_list]
    q_id, our_model_name, other_model_name, _ = random.choices(
        full_question_ids_list, weights=weights, k=1)[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_answer, chat_A_reasoning, _ = format_chat(
        question_for_eval['models'][0]['reasoning_trace'], tool_database_labels)
    chat_B_answer, chat_B_reasoning, _ = 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;">Question:</strong> {prompt_text}</div>')
    page1_reference_answer = gr.Markdown(txagent_matched_entry.get(
        "correct_answer")) if include_correct_answer else None
    chat_a_answer = gr.Chatbot(
        value=chat_A_answer,
        type="messages",
        height=200,
        label="Model A Answer",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False,
        autoscroll=False,
    )
    chat_b_answer = gr.Chatbot(
        value=chat_B_answer,
        type="messages",
        height=200,
        label="Model B Answer",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False,
        autoscroll=False,
    )
    chat_a_reasoning = gr.Chatbot(
        value=chat_A_reasoning,
        type="messages",
        height=300,
        label="Model A Reasoning - Rationale",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False,
        autoscroll=False,
    )
    chat_b_reasoning = gr.Chatbot(
        value=chat_B_reasoning,
        type="messages",
        height=300,
        label="Model B Reasoning - Rationale",
        show_copy_button=False,
        show_label=True,
        render_markdown=True,
        avatar_images=None,
        rtl=False,
        autoscroll=False,
    )

    user_info = (name, email, specialty_dd, subspecialty_dd, years_exp_radio,
                 exp_explanation_tb, npi_id, q_id) if return_user_info else None
    return user_info, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, page1_prompt, page1_reference_answer, question_for_eval, len(full_question_ids_list)


def go_to_page0_from_minus1(question_in_progress_state):
    if question_in_progress_state == 1:
        # If a question is in progress on page 1, go directly to page 1
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
    elif question_in_progress_state == 2:
        # If a question is in progress on page 2, go directly to page 2
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
    else:
        # If no question is in progress, show the initial page 0
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)


def go_to_eval_progress_modal(name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, our_methods=our_methods):
    # 校验用户信息
    if not name or not email or not specialty_dd or not years_exp_radio:
        gr.Info("Please fill out all the required fields (name, email, 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.", duration=5)
        return gr.update(visible=True), gr.update(visible=False), None, "Please fill out all the required fields (name, email, 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.Chatbot(), gr.Chatbot(), gr.HTML(), gr.State()

    gr.Info("Loading the data...", duration=3)
    user_info, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, 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, our_methods
    )
    if remaining_count == 0:
        gr.Info("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!", duration=5)
        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.Chatbot(), gr.Chatbot(), gr.HTML(), gr.State()
    gr.Info(f"You are about to evaluate the next question.", duration=3)
    return gr.update(visible=False), gr.update(visible=True), user_info, "", chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, page1_prompt, question_for_eval

# goes to page 1 from confirmation modal that tells users how many questions they have left to evaluate


def go_to_page1(show_page_1):
    """
    Shows page 1 if user requests it, otherwise shows page 0
    """

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


# --- Skip Question Modal Callbacks ---
def skip_question_and_load_new(user_info_state, our_methods):
    # user_info_state is a tuple: (name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, q_id)
    if user_info_state is None:
        # Defensive: just close modal if no user info
        return gr.update(visible=False), gr.update(visible=False), None, "", gr.Chatbot(), gr.Chatbot(), gr.Chatbot(), gr.Chatbot(), gr.HTML(), gr.Markdown(), gr.State()
    # Unpack user_info_state
    name, email, specialty_dd, subspecialty_dd, years_exp_radio, exp_explanation_tb, npi_id, _ = user_info_state
    user_info, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, 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, our_methods
    )
    if remaining_count == 0:
        # No more questions, go to final page
        return gr.update(visible=False), 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.Chatbot(), gr.Chatbot(), gr.HTML(), gr.Markdown(), gr.State()
    return gr.update(visible=False), gr.update(visible=True), user_info, "", chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, page1_prompt, page1_reference_answer, question_for_eval

# --- Skip‑question handler for the "Wrong Question?" button -------------------


def skip_current_question(user_info_state, our_methods: list = our_methods):
    # Guard: user clicked before session started
    gr.Info("Skipping this question and loading the next one…",  duration=5)
    if user_info_state is None:
        return (
            None,
            gr.update(
                value="Please start the evaluation before skipping questions."),
            gr.update(value=[]),  # Chatbot A history
            gr.update(value=[]),  # Chatbot B history
            gr.update(value=""),  # Prompt HTML
            gr.State()            # data_subset_state
        )

    # Unpack evaluator identity
    name, email, specialty_dd, subspecialty_dd, yrs_exp, exp_desc, npi_id, _ = user_info_state

    # Pull the next unused question
    (
        user_info_new,
        _chat_a_answer,
        _chat_b_answer,
        _chat_a_reasoning,
        _chat_b_reasoning,
        _prompt_comp,
        _ref_comp,
        question_for_eval,
        remaining,
    ) = get_next_eval_question(
        name, email, specialty_dd, subspecialty_dd, yrs_exp, exp_desc, npi_id, our_methods
    )

    # If the pool is exhausted, just notify the evaluator
    if remaining == 0 or question_for_eval is None:
        final_msg = (
            "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!"
        )
        return (
            user_info_state,
            gr.update(value=final_msg),
            gr.update(value=[]),
            gr.update(value=[]),
            gr.update(value=[]),
            gr.update(value=[]),
            gr.update(value=""),
            gr.State()
        )

    # --- Build fresh values for the existing UI components ---
    chat_a_answer, chat_a_reasoning, _ = format_chat(
        question_for_eval['models'][0]['reasoning_trace'], tool_database_labels)
    chat_b_answer, chat_b_reasoning, _ = format_chat(
        question_for_eval['models'][1]['reasoning_trace'], tool_database_labels)

    prompt_html = (
        f"<div style='background-color: #FFEFD5; border: 2px solid #FF8C00; padding: 10px; "
        f"border-radius: 5px; color: black;'><strong style='color: black;'>Question:</strong> "
        f"{question_for_eval['question']}</div>"
    )
    reference_md = question_for_eval.get("correct_answer", "")
    gr.Info("New question loaded…",  duration=3)

    # Return updates to refresh Page 1 in‑place
    return (
        user_info_new,
        gr.update(value=""),                 # clear any previous error text
        gr.update(value=chat_a_answer),       # Chatbot A history
        gr.update(value=chat_b_answer),       # Chatbot B history
        gr.update(value=chat_a_reasoning),    # Chatbot A reasoning
        gr.update(value=chat_b_reasoning),    # Chatbot B reasoning
        gr.update(value=prompt_html),        # Prompt
        question_for_eval                    # store for later pages
    )

# --- Handler for "Wrong Question?": flags nonsense and skips


def flag_nonsense_and_skip(user_info_state, skip_comments=""):
    """
    When the evaluator clicks the “Wrong Question?” button, immediately
    record that this question was flagged as nonsensical/irrelevant and
    then load the next question (re‑using the existing skip logic).
    """
    # 1) Record the flag to the Google Sheet so we keep the feedback even
    #    if the evaluator stops here.
    if user_info_state is not None:
        name, email, specialty_dd, subspecialty_dd, yrs_exp, exp_desc, npi_id, q_id = user_info_state
        timestamp = datetime.datetime.now().isoformat()
        row = {
            "Timestamp": timestamp,
            "Name": name,
            "Email": email,
            "Question ID": q_id,
            "Question Makes No Sense or Biomedically Irrelevant": True,
            "Skip Comments": skip_comments,
        }
        append_to_sheet(
            user_data=None,
            custom_row_dict=row,
            custom_sheet_name=str(TXAGENT_RESULTS_SHEET_BASE_NAME),
            add_header_when_create_sheet=True,
        )

    # 2) Fall back to the existing skip logic to advance the UI.
    return skip_current_question(user_info_state)

# Define restrict function for each criterion


def make_restrict_function(base_choices):
    def restrict_choices_page1(radio_choice, score_a, score_b):
        """
        Returns (update_for_A, update_for_B).
        Enforces rating constraints based on the radio choice for page 1.
        """
        # Helper to parse int safely
        def to_int(x):
            try:
                # Extract number from "1 text..." format
                return int(x.split()[0])
            except (ValueError, TypeError, AttributeError):
                return None

        # Default: no restrictions, but ensure current values are valid
        upd_A = gr.update(choices=base_choices,
                          value=score_a if score_a in base_choices else None)
        upd_B = gr.update(choices=base_choices,
                          value=score_b if score_b in base_choices else None)

        # Skip if no meaningful pairwise choice
        if radio_choice is None or radio_choice == "Neither model did well.":
            return upd_A, upd_B

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

        # Apply Restrictions based on radio choice
        if radio_choice == "Model A is better.":
            # Rule: A >= B
            if a_int is not None and b_int is not None:
                # Both are numeric, enforce A >= B
                if a_int < b_int:
                    # Violation: A < B, reset the one that doesn't match the constraint
                    upd_A = gr.update(choices=base_choices, value=None)
                    upd_B = gr.update(choices=base_choices, value=None)
                else:
                    # Valid: A >= B, apply mutual restrictions
                    allowed_a_choices = [choice for choice in base_choices if to_int(
                        choice) is None or to_int(choice) >= b_int]
                    allowed_b_choices = [choice for choice in base_choices if to_int(
                        choice) is None or to_int(choice) <= a_int]
                    upd_A = gr.update(
                        choices=allowed_a_choices, value=score_a if score_a in allowed_a_choices else None)
                    upd_B = gr.update(
                        choices=allowed_b_choices, value=score_b if score_b in allowed_b_choices else None)
            elif a_int is not None:
                # Only A is numeric, B must be <= A
                allowed_b_choices = [choice for choice in base_choices if to_int(
                    choice) is None or to_int(choice) <= a_int]
                upd_B = gr.update(
                    choices=allowed_b_choices, value=score_b if score_b in allowed_b_choices else None)
            elif b_int is not None:
                # Only B is numeric, A must be >= B
                allowed_a_choices = [choice for choice in base_choices if to_int(
                    choice) is None or to_int(choice) >= b_int]
                upd_A = gr.update(
                    choices=allowed_a_choices, value=score_a if score_a in allowed_a_choices else None)
            # If both are "Unable to Judge", no restrictions needed

        elif radio_choice == "Model B is better.":
            # Rule: B >= A
            if a_int is not None and b_int is not None:
                # Both are numeric, enforce B >= A
                if b_int < a_int:
                    # Violation: B < A, reset both
                    upd_A = gr.update(choices=base_choices, value=None)
                    upd_B = gr.update(choices=base_choices, value=None)
                else:
                    # Valid: B >= A, apply mutual restrictions
                    allowed_a_choices = [choice for choice in base_choices if to_int(
                        choice) is None or to_int(choice) <= b_int]
                    allowed_b_choices = [choice for choice in base_choices if to_int(
                        choice) is None or to_int(choice) >= a_int]
                    upd_A = gr.update(
                        choices=allowed_a_choices, value=score_a if score_a in allowed_a_choices else None)
                    upd_B = gr.update(
                        choices=allowed_b_choices, value=score_b if score_b in allowed_b_choices else None)
            elif a_int is not None:
                # Only A is numeric, B must be >= A
                allowed_b_choices = [choice for choice in base_choices if to_int(
                    choice) is None or to_int(choice) >= a_int]
                upd_B = gr.update(
                    choices=allowed_b_choices, value=score_b if score_b in allowed_b_choices else None)
            elif b_int is not None:
                # Only B is numeric, A must be <= B
                allowed_a_choices = [choice for choice in base_choices if to_int(
                    choice) is None or to_int(choice) <= b_int]
                upd_A = gr.update(
                    choices=allowed_a_choices, value=score_a if score_a in allowed_a_choices else None)

        elif radio_choice == "Both models are equally good.":
            # Rule: A == B
            if a_int is not None and b_int is not None:
                # Both are numeric
                if a_int == b_int:
                    # Valid: A == B, restrict both to the same value
                    upd_A = gr.update(choices=[score_a], value=score_a)
                    upd_B = gr.update(choices=[score_b], value=score_b)
                else:
                    # Invalid: A != B, reset both
                    upd_A = gr.update(choices=base_choices, value=None)
                    upd_B = gr.update(choices=base_choices, value=None)
            elif a_int is not None:
                # A is numeric, B must match A
                upd_B = gr.update(choices=[score_a], value=score_a)
            elif b_int is not None:
                # B is numeric, A must match B
                upd_A = gr.update(choices=[score_b], value=score_b)
            elif score_a == "Unable to Judge." and score_b == "Unable to Judge.":
                # Both are "Unable to Judge", restrict both to that
                upd_A = gr.update(
                    choices=["Unable to Judge."], value="Unable to Judge.")
                upd_B = gr.update(
                    choices=["Unable to Judge."], value="Unable to Judge.")
            elif score_a == "Unable to Judge.":
                # A is "Unable to Judge", B must match
                upd_B = gr.update(
                    choices=["Unable to Judge."], value="Unable to Judge.")
            elif score_b == "Unable to Judge.":
                # B is "Unable to Judge", A must match
                upd_A = gr.update(
                    choices=["Unable to Judge."], value="Unable to Judge.")
            # If neither has a value, no restrictions needed

        return upd_A, upd_B
    return restrict_choices_page1

# --- 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],
        "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), add_header_when_create_sheet=True)

    # --- Part 2: Recalculate remaining questions (Existing Logic + Modified Error Handling) ---
    name, email, specialty, subspecialty, years_exp_radio, exp_explanation_tb, npi_id, _ = user_info
    user_info_new, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning, 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 (
            "",                        # page1_error_box
            gr.update(visible=False),  # page1 (Hide)
            gr.update(visible=True),   # final_page (Show)
            "",                        # page0_error_box
            None,                      # chat_a_answer
            None,                      # chat_b_answer
            None,                      # chat_a_reasoning
            None,                      # chat_b_reasoning
            None,                      # page1_prompt
            None,                      # data_subset_state
            user_info_new,             # user_info_state
        )
    return (
        "",                        # page1_error_box
        gr.update(visible=True),   # page1 (Show for next question)
        gr.update(visible=False),  # final_page (Hide)
        "",                        # page0_error_box
        chat_a_answer,             # chat_a_answer
        chat_b_answer,             # chat_b_answer
        chat_a_reasoning,          # chat_a_reasoning
        chat_b_reasoning,          # chat_b_reasoning
        page1_prompt,              # page1_prompt
        question_for_eval,         # data_subset_state
        user_info_new              # user_info_state
    )


# Function to validate page1 inputs and directly submit if valid
def validate_and_submit_page1(data_subset_state, user_info, *combined_values):
    # combined_values contains pairwise choices + comparison reasons + ratings
    criteria_count = len(criteria_for_comparison)
    pairwise_list = list(combined_values[:criteria_count])
    comparison_reasons_list = list(
        combined_values[criteria_count:criteria_count*2])
    ratings_A_list = list(
        combined_values[criteria_count*2:criteria_count*3])
    ratings_B_list = list(combined_values[criteria_count*3:])

    # Check if all pairwise comparisons are filled
    if any(answer is None for answer in pairwise_list):
        missing_comparisons = []
        for i, answer in enumerate(pairwise_list):
            if answer is None:
                missing_comparisons.append(criteria_for_comparison[i]['label'])

        missing_text = ", ".join(missing_comparisons)
        error_msg = f"Please select an option for the following pairwise comparison(s): {missing_text}"
        gr.Info(error_msg)
        return (
            gr.update(value=f"Error: {error_msg}"),
            gr.update(visible=True),   # Keep page1 visible
            gr.update(visible=False),  # Keep final_page hidden
            gr.update(),  # page0_error_box - keep unchanged
            gr.update(),  # chat_a - keep unchanged
            gr.update(),  # chat_b - keep unchanged
            gr.update(),  # chat_a - keep unchanged
            gr.update(),  # chat_b - keep unchanged
            gr.update(),  # page1_prompt - keep unchanged
            gr.update(),  # data_subset_state - keep unchanged
            gr.update(),  # user_info_state - keep unchanged
            # Keep form fields unchanged on validation error
            *combined_values
        )

    # Check if all ratings are filled
    if any(r is None for r in ratings_A_list) or any(r is None for r in ratings_B_list):
        missing_ratings = []
        for i in range(len(criteria)):
            missing_parts = []
            if ratings_A_list[i] is None:
                missing_parts.append("Response A")
            if ratings_B_list[i] is None:
                missing_parts.append("Response B")
            if missing_parts:
                missing_ratings.append(
                    f"{criteria[i]['label']} ({', '.join(missing_parts)})")

        missing_text = "; ".join(missing_ratings)
        error_msg = f"Please provide ratings for: {missing_text}"
        gr.Info(error_msg)
        return (
            gr.update(value=f"Error: {error_msg}"),
            gr.update(visible=True),   # Keep page1 visible
            gr.update(visible=False),  # Keep final_page hidden
            gr.update(),  # page0_error_box - keep unchanged
            gr.update(),  # chat_a - keep unchanged
            gr.update(),  # chat_b - keep unchanged
            gr.update(),  # chat_a - keep unchanged
            gr.update(),  # chat_b - keep unchanged
            gr.update(),  # page1_prompt - keep unchanged
            gr.update(),  # data_subset_state - keep unchanged
            gr.update(),  # user_info_state - keep unchanged
            # Keep form fields unchanged on validation error
            *combined_values
        )
    gr.Info("Submitting your evaluation and loading the next question...")
    # If validation passes, call final_submit and handle form reset
    submit_result = final_submit(data_subset_state, user_info, pairwise_list,
                                 comparison_reasons_list, False, *ratings_A_list, *ratings_B_list)

    # Check if there are more questions by looking at the page1 update dict
    # submit_result[1] is the page1 update, submit_result[2] is the final_page update
    page1_update = submit_result[1]
    page1_visible = page1_update.get('visible', False) if isinstance(
        page1_update, dict) else False
    gr.Info(f"Your evaluation has been submitted. You are about to evaluate the next question...")
    # If there are more questions (page1 is visible after submit), reset the form
    if page1_visible:  # page1 is visible, meaning there's a next question
        # Reset form fields for next question
        reset_values = []
        for _ in range(len(combined_values)):
            reset_values.append(None)
        return submit_result + tuple(reset_values)
    else:
        # Final page is shown, keep current form values (though they won't be visible)
        return submit_result + tuple(combined_values)


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 {
  /* Light‑mode palette */
  --btn-bg: #E0F2FF;        /* soft pastel blue */
  --btn-text: #00334D;      /* dark slate for good contrast */
  --btn-border: #E0F2FF;

  background-color: var(--btn-bg) !important;
  color: var(--btn-text) !important;
  border: 1px solid var(--btn-border) !important;
}

/* Dark‑mode overrides */
@media (prefers-color-scheme: dark) {
  #answer-reference-btn {
    --btn-bg: #2C6E98;      /* muted steel blue for dark backgrounds */
    --btn-text: #FFFFFF;    /* switch to white text for contrast */
    --btn-border: #2C6E98;
  }
}

#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; }

/* --- Added for larger criteria font --- */
.criteria-font-large {
    font-size: 1.2em !important;
}
/* Radio component labels (the title above the choices) */
.criteria-radio-label label[data-testid="block-label"] {
    font-weight: bold !important;
    font-size: 1.1em !important;
}
/* Textbox labels */
.textbox-bold-label label[data-testid="block-label"] {
    font-weight: bold !important;
}
#participate-btn button {
    font-size: 24px !important;      /* Large readable text */
    font-weight: 700 !important;     /* Bold for emphasis */
    padding: 28px 40px !important;   /* Extra padding for height */
    min-height: 120px !important;    /* Make button visibly taller (multi‑line) */
    width: 100% !important;          /* Occupy full width of its column */
    white-space: normal !important;  /* Allow text to wrap onto multiple lines */
}
.criteria-radio-score-label [role="radiogroup"],
.criteria-radio-score-label .gr-radio-group,
.criteria-radio-score-label .flex {
    display: flex !important;
    flex-direction: column !important;
    gap: 4px !important;                 /* 行间距，可按需调整 */
}

/* 更具体的选择器来确保垂直布局 */
.criteria-radio-score-label fieldset {
    display: flex !important;
    flex-direction: column !important;
    gap: 4px !important;
}

.criteria-radio-score-label .wrap {
    display: flex !important;
    flex-direction: column !important;
    gap: 4px !important;
}

/* 确保每个单选按钮选项垂直排列 */
.criteria-radio-score-label label {
    display: block !important;
    margin-bottom: 4px !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()
    question_in_progress = gr.State(0)

    # 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>
        </div>
        """)
        # with gr.Row(elem_classes=["center-row"]):
        # 第一行：并排放两个按钮
        with gr.Column(scale=1):
            participate_eval_btn = gr.Button(
                value="Evaluate TxAgent",
                variant="primary",
                size="lg",
                elem_id="participate-btn"
            )
        with gr.Column(scale=1):
            gr.Markdown(
                """
                Joining the **Evaluate TxAgent**, you will:
                - See how the model responds to a variety of prompts.
                - Provide instant thumbs-up / thumbs-down ratings.
                - Shape the roadmap for upcoming releases.

                _Thank you for helping us improve!_
                """
            )
        with gr.Column(scale=1):
            submit_questions_btn = gr.Button(
                value="Submit Therapeutic Questions",
                variant="primary",
                size="lg",
                elem_id="submit-btn"
            )

        # with gr.Row(elem_classes=["center-row"]):
            # 第二行：分别放两段说明文字
        with gr.Column(scale=1):
            gr.Markdown(
                """
                By submitting **Submit Therapeutic Questions**, you will:
                - Highlight edge-cases and identify blind spots.
                - Push TxAgent to reason in new domains.
                - Directly influence future model improvements.

                _We're excited to see what you come up with!_
                """
            )

            # Add contact information in Markdown format
            contact_info_markdown = """
            ## Contact

            If you have any questions or suggestions, please email [Shanghua Gao](mailto:[email protected]) and [Marinka Zitnik](mailto:[email protected]).
            """

            gr.Markdown(contact_info_markdown)

        gr.HTML(TxAgent_Project_Page_HTML)

        # Define actions for the new buttons
        # For the Google Form button, we'll use JavaScript to open a new tab.
        # The URL for the Google Form should be replaced with the actual link.
        google_form_url = "https://forms.gle/pYvyvEQQwS5gdupQA"
        submit_questions_btn.click(
            fn=None,
            inputs=None,
            outputs=None,
            js=f"() => {{ window.open('{google_form_url}', '_blank'); }}"
        )

    # Page 0: Welcome / Informational page.
    with gr.Column(visible=False, elem_id="page0") as page0:

        gr.Markdown("## Information:")
        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.")
        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")
        gr.Markdown("""By clicking 'Next', 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!
            """)
        # 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")

        # 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)

    # Page 1: Pairwise Comparison.
    with gr.Column(visible=False) as page1:
        with gr.Accordion("Instructions", open=False):
            gr.Markdown("""
                    ## Instructions:
                    Please review these instructions and enter your information to begin:

                    - Each session requires at least 5-10 minutes per question.
                    - You can evaluate multiple questions; you will not repeat evaluations.
                    - For each question, compare responses from two models and rate them (scale: 1-5).
                    - If a question is unclear or irrelevant to biomedicine, click the RED BUTTON at the top of the comparison page.
                    - Use the Back and Next buttons to edit responses before submission.
                    - Use the Home Page button to return to the homepage; progress will save but not submit.
                    - Submit answers to the current question before moving to the next.
                    - You can pause between questions and return later; ensure current answers are submitted to save them.
                """)
        # Make the number controlled by question indexing!
        # gr.Markdown("Comparison")
        # Add small red button and comments text box in the same row
        page1_prompt = gr.HTML()
        with gr.Row():
            nonsense_btn = gr.Button(
                "Skip Question",
                size="sm",
                variant="stop",  # red variant
                elem_id="invalid-question-btn",
                elem_classes=["short-btn"],
                scale=1
            )
            skip_comments = gr.Textbox(
                placeholder="(Optional) Why do you want to skip this question...",
                show_label=False,
                scale=3,
                container=False,
            )

        page1_error_box = gr.Markdown("")  # ADDED: display validation errors

        # --- Define four chat components: answer and reasoning for each model ---
        with gr.Row():
            # Model A components
            with gr.Column():
                gr.Markdown("**Model A Response:**")
                chat_a_answer = gr.Chatbot(
                    value=[],  # Placeholder for chat history
                    type="messages",
                    height=200,
                    label="Model A Answer",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,
                    avatar_images=None,
                    rtl=False
                )
                # gr.Markdown("**Model A Reasoning:**")
                chat_a_reasoning = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=300,
                    label="Model A Reasoning - Rationale",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,
                    avatar_images=None,
                    rtl=False
                )
            # Model B components
            with gr.Column():
                gr.Markdown("**Model B Response:**")
                chat_b_answer = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=200,
                    label="Model B Answer",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,
                    avatar_images=None,
                    rtl=False
                )
                # gr.Markdown("**Model B Reasoning:**")
                chat_b_reasoning = gr.Chatbot(
                    value=[],
                    type="messages",
                    height=300,
                    label="Model B Reasoning - Rationale",
                    show_copy_button=False,
                    show_label=True,
                    render_markdown=True,
                    avatar_images=None,
                    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 = []
        ratings_A_page1 = []  # Store rating components for page 1
        ratings_B_page1 = []  # Store rating components for page 1

        for i, crit_comp in enumerate(criteria_for_comparison):
            # for crit in criteria_for_comparison:
            crit_score = criteria[i]  # Get the corresponding score criterion

            restrict_fn = make_restrict_function(sorted(crit_score["scores"]))

            # Add bold formatting
            gr.Markdown(f"**{crit_comp['label']}**",
                        elem_classes="criteria-font-large")
            radio = gr.Radio(
                choices=[
                    "Model A is better.",
                    "Model B is better.",
                    "Both models are equally good.",
                    "Neither model did well."
                ],
                # Remove duplicate label since we have markdown above
                label=crit_comp['text'],
                elem_classes="criteria-radio-label"  # <--- add class here
            )
            pairwise_inputs.append(radio)
            # ADDED: free text under each comparison

        # for i, crit in enumerate(criteria):
            index_component = gr.Number(
                value=i, visible=False, interactive=False)
            # indices_for_change.append(index_component)

            with gr.Row():
                with gr.Column(scale=1):
                    rating_a = gr.Radio(choices=sorted(crit_score["scores"]),  # ["1", "2", "3", "4", "5", "Unable to Judge"],
                                        label=f"Response A - {crit_score['text']}",
                                        interactive=True,
                                        elem_classes="criteria-radio-score-label")
                with gr.Column(scale=1):
                    rating_b = gr.Radio(choices=sorted(crit_score["scores"]),  # ["1", "2", "3", "4", "5", "Unable to Judge"],
                                        label=f"Response B - {crit_score['text']}",
                                        interactive=True,
                                        elem_classes="criteria-radio-score-label")

            # Add clear button and wire up the restrictions
            with gr.Row():
                # wire each to re‐restrict the other on change
                radio.change(
                    fn=restrict_fn,
                    inputs=[radio, rating_a, rating_b],
                    outputs=[rating_a, rating_b]
                )
                rating_a.change(
                    fn=restrict_fn,
                    inputs=[radio, rating_a, rating_b],
                    outputs=[rating_a, rating_b]
                )
                rating_b.change(
                    fn=restrict_fn,
                    inputs=[radio, rating_a, rating_b],
                    outputs=[rating_a, rating_b]
                )

            ratings_A_page1.append(rating_a)
            ratings_B_page1.append(rating_b)

            text_input = gr.Textbox(
                # Remove label since we have markdown above
                placeholder="Comments for your selection (optional)",
                show_label=False,
                # elem_classes="textbox-bold-label"
            )
            comparison_reasons_inputs.append(text_input)

        with gr.Row():
            submit_btn_1 = gr.Button(
                "Submit Evaluation", variant="primary", elem_id="submit_btn")

    # 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)

    # --- Define Transitions Between Pages ---

    # For the "Participate in Evaluation" button, transition to page0
    participate_eval_btn.click(
        fn=go_to_page0_from_minus1,
        inputs=[question_in_progress],
        # Removed page2 reference
        outputs=[page_minus1, page0, page1, final_page]
    )

    # Transition from Page 0 (Welcome) to Page 1.
    next_btn_0.click(
        fn=go_to_eval_progress_modal,
        inputs=[name, email, specialty_dd, subspecialty_dd,
                years_exp_radio, exp_explanation_tb, npi_id],
        outputs=[page0, page1, user_info_state, page0_error_box, chat_a_answer,
                 chat_b_answer, chat_a_reasoning, chat_b_reasoning, page1_prompt, data_subset_state],
        scroll_to_output=True
    )
    # Skip the current question and load a new one when the evaluator flags it
    nonsense_btn.click(
        fn=flag_nonsense_and_skip,
        inputs=[user_info_state, skip_comments],
        outputs=[user_info_state, page1_error_box, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning,
                 page1_prompt, data_subset_state],
        scroll_to_output=True
    )

    # Transition from Page 1 to direct submission (no confirmation modal)
    submit_btn_1.click(
        fn=validate_and_submit_page1,
        inputs=[data_subset_state, user_info_state, *pairwise_inputs,
                *comparison_reasons_inputs, *ratings_A_page1, *ratings_B_page1],
        outputs=[page1_error_box, page1, final_page, page0_error_box, chat_a_answer, chat_b_answer, chat_a_reasoning, chat_b_reasoning,
                 page1_prompt, data_subset_state, user_info_state, *pairwise_inputs, *comparison_reasons_inputs, *ratings_A_page1, *ratings_B_page1],
        scroll_to_output=True
    )


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