# ==============================================================================
# evaluation_interface のコードブロック（セッションステート対応済み）
# ==============================================================================
import os
import sys
import glob
import json
import random
from functools import partial
from datetime import datetime
from collections import defaultdict, Counter

import gradio as gr
from loguru import logger
from PIL import Image
import re

# --- ★ 修正: GLOBAL_STATE を削除 ---
# グローバルな状態管理を廃止し、Gradioのセッションステート (gr.State) に移行します。

# --- Configuration (変更なし) ---
BASE_RESULTS_DIR = "./results"
LOG_DIR = "./logs"
COMBINED_DATA_DIR = "./combined_data"
IMAGE_SUBDIR = os.path.join("lapwing", "images")
MAPPING_FILENAME = "combination_to_filename.json"
CONDITIONS = ["Ours", "w_o_Proto_Loss", "w_o_HitL", "w_o_Tuning", "LLM-based"]
CRITERIA = ["Alignment", "Naturalness", "Attractiveness"]
CRITERIA_GUIDANCE_JP = [
    "テキストと上半分の表情がどれだけ一致していますか？",
    "提示されたテキストと上半分の表情を考慮した上で、このキャラクターはどれだけ自然に見えますか？",
    "提示されたテキストと上半分の表情を考慮した上で、このキャラクターはどれだけ魅力的に見えますか？"
]
CRITERIA_GUIDANCE_EN = [
    "How well does the text align with the upper half of the expression?",
    "Considering the provided text and the upper half of the expression, how natural was the character overall?",
    "Considering the provided text and the upper half of the expression, how attractive was the character overall?"
]
IMAGE_LABELS = ['A', 'B', 'C', 'D', 'E']


# --- Helper Functions (★ 修正: stateを引数に追加) ---
def load_bbox_json(bbox_json_path, state):
    """バウンディングボックス情報をJSONファイルから読み込み、stateに格納する"""
    try:
        with open(bbox_json_path, 'r', encoding='utf-8') as f:
            bbox_data = json.load(f)
        state["hide_bbox_dict"] = bbox_data.get("Hide", {})
        logger.info(f"Successfully loaded bounding box data from {bbox_json_path}")
    except Exception as e:
        logger.error(f"Failed to load bounding box JSON: {e}")
        state["hide_bbox_dict"] = {}
    return state


def create_masked_image(image: Image.Image, state: dict):
    """画像に黒塗りのマスクを適用する"""
    hide_bbox_dict = state.get("hide_bbox_dict", {})
    if not hide_bbox_dict:
        return image
    masked_img = image.copy()
    for _, box_coords in hide_bbox_dict.items():
        box = (box_coords['left'], box_coords['top'], box_coords['right'], box_coords['bottom'])
        black_rectangle = Image.new('RGB', (box[2] - box[0], box[3] - box[1]), color='black')
        masked_img.paste(black_rectangle, (box[0], box[1]))
    return masked_img


def get_image_path_from_prediction(prediction: dict, state: dict) -> str:
    if not state["image_mapping"]:
        logger.error("Image mapping is not loaded.")
        return ""
    indices = prediction.get("blendshape_index", {})
    if not isinstance(indices, dict):
        logger.error(f"blendshape_index is not a dictionary: {indices}")
        return ""
    sorted_indices = sorted(indices.items(), key=lambda item: int(item[0]))
    key = ",".join(str(idx) for _, idx in sorted_indices)
    filename = state["image_mapping"].get(key)
    if not filename:
        logger.warning(f"No image found for blendshape key: {key}")
        return ""
    return os.path.join(state["image_dir"], filename)


def load_evaluation_data(participant_id: str, state: dict):
    """ ★ 修正: stateを引数で受け取り、更新したstateを返す """
    bbox_json_path = os.path.join(COMBINED_DATA_DIR, "lapwing", "texts", "bounding_boxes.json")
    if os.path.exists(bbox_json_path):
        state = load_bbox_json(bbox_json_path, state)
    else:
        logger.warning(f"Bounding box file not found at {bbox_json_path}. Images will not be masked.")
        state["hide_bbox_dict"] = {}

    mapping_path = os.path.join(COMBINED_DATA_DIR, MAPPING_FILENAME)
    if not os.path.exists(mapping_path):
        return state, f"<p class='feedback red'>Error: Mapping file not found at {mapping_path}</p>", gr.update(
            interactive=True), gr.update(interactive=False)

    with open(mapping_path, 'r', encoding='utf-8') as f:
        state["image_mapping"] = json.load(f)["mapping"]
    state["image_dir"] = os.path.join(COMBINED_DATA_DIR, IMAGE_SUBDIR)
    logger.info(f"Successfully loaded image mapping. Image directory: {state['image_dir']}")

    participant_dir = os.path.join(BASE_RESULTS_DIR, participant_id)
    if not os.path.isdir(participant_dir):
        return state, f"<p class='feedback red'>Error: Participant directory not found: {participant_dir}</p>", gr.update(
            interactive=True), gr.update(interactive=False)

    merged_data = defaultdict(lambda: {"predictions": {}, "category": None})
    found_files = 0
    for cond in CONDITIONS:
        cond_dir = os.path.join(participant_dir, cond)
        pattern = os.path.join(cond_dir, f"{participant_id}_{cond}_*.jsonl")
        files = glob.glob(pattern)
        if not files:
            logger.warning(f"No prediction file found for condition '{cond}' with pattern: {pattern}")
            continue
        found_files += 1
        with open(files[0], 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line)
                prompt = data["text_prompt"]
                merged_data[prompt]["predictions"][cond] = data["prediction"]
                if not merged_data[prompt]["category"]:
                    merged_data[prompt]["category"] = data.get("prompt_category")

    if found_files != len(CONDITIONS):
        return state, f"<p class='feedback red'>Error: Found prediction files for only {found_files}/{len(CONDITIONS)} conditions.</p>", gr.update(
            interactive=True), gr.update(interactive=False)

    state["all_eval_data"] = [
        {"prompt": p, "predictions": d["predictions"], "category": d["category"]}
        for p, d in merged_data.items() if len(d["predictions"]) == len(CONDITIONS)
    ]

    if not state["all_eval_data"]:
        return state, "<p class='feedback red'>Error: No valid evaluation data could be loaded.</p>", gr.update(
            interactive=True), gr.update(interactive=False)

    state["shuffled_indices"] = list(range(len(state["all_eval_data"])))
    random.shuffle(state["shuffled_indices"])
    state["current_prompt_index"] = 0
    state["current_criterion_index"] = 0
    state["data_loaded"] = True
    state["start_time"] = datetime.now()
    for i in range(len(state["all_eval_data"])):
        prompt_text = state["all_eval_data"][i]["prompt"]
        state["evaluation_results"][prompt_text] = {}

    logger.info(f"Loaded and merged data for {len(state['all_eval_data'])} prompts.")
    done_msg = "<p class='feedback green'>Data loaded successfully. Please proceed to the 'Evaluation' tab. / データの読み込みに成功しました。「評価」タブに進んでください。</p>"
    return state, done_msg, gr.update(interactive=False, visible=False), gr.update(interactive=True)


# --- Core Logic (★ 修正: stateを引数に追加) ---
def _create_button_updates(state: dict):
    updates = []
    for img_label in IMAGE_LABELS:
        selected_rank = state["current_ranks"].get(img_label)
        for rank_val in range(1, 6):
            if rank_val == selected_rank:
                updates.append(gr.update(variant='primary'))
            else:
                updates.append(gr.update(variant='secondary'))
    return updates


def handle_rank_button_click(image_label, rank, state: dict):
    if state["current_ranks"].get(image_label) == rank:
        state["current_ranks"][image_label] = None
    else:
        state["current_ranks"][image_label] = rank
    return state, *_create_button_updates(state)


def handle_absolute_score_click(score, state: dict):
    if state["current_absolute_score"] == score:
        state["current_absolute_score"] = None
    else:
        state["current_absolute_score"] = score
    updates = []
    for i in range(1, 8):
        if i == state["current_absolute_score"]:
            updates.append(gr.update(variant='primary'))
        else:
            updates.append(gr.update(variant='secondary'))
    return state, *updates


def handle_absolute_score_worst_click(score, state: dict):
    if state["current_absolute_score_worst"] == score:
        state["current_absolute_score_worst"] = None
    else:
        state["current_absolute_score_worst"] = score
    updates = []
    for i in range(1, 8):
        if i == state["current_absolute_score_worst"]:
            updates.append(gr.update(variant='primary'))
        else:
            updates.append(gr.update(variant='secondary'))
    return state, *updates


# --- UI Logic (★ 修正: stateを引数に追加) ---
def display_current_prompt_and_criterion(state: dict):
    if not state["data_loaded"] or state["current_prompt_index"] >= len(state["all_eval_data"]):
        done_msg = "<p class='feedback green' style='text-align: center; font-size: 1.2em;'>All prompts have been evaluated! Please proceed to the 'Export' tab. <br>すべてのプロンプトの評価が完了しました！「エクスポート」タブに進んでください。</p>"
        empty_button_updates = [gr.update(variant='secondary')] * 25
        empty_abs_updates = [gr.update(variant='secondary')] * 7
        return [
            gr.update(value="Finished! / 完了!"),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=done_msg, visible=True),
            *[gr.update(value=None)] * 5,
            *empty_button_updates,
            gr.update(visible=False),
            *empty_abs_updates,
            gr.update(visible=False),
            *empty_abs_updates,
            gr.update(interactive=False),
            gr.update(interactive=False)
        ]

    prompt_idx = state["shuffled_indices"][state["current_prompt_index"]]
    criterion_idx = state["current_criterion_index"]

    current_data = state["all_eval_data"][prompt_idx]
    prompt_text = current_data["prompt"]
    criterion_name = CRITERIA[criterion_idx]

    progress_text = f"Prompt {state['current_prompt_index'] + 1} / {len(state['all_eval_data'])}  -  **{criterion_name}**"
    prompt_display_text = f"<p style='font-size: 3em; font-weight: bold;'>テキスト(TEXT): {prompt_text}</p>"
    # (Instructions text is unchanged)
    guidance_part = (
        f"<p style='color: red; font-weight: bold; font-size: 1.1em;'>"
        f"5つの画像を、「{CRITERIA_GUIDANCE_JP[criterion_idx]}」を基準にランキングしてください。<br>"
        f"「ランキングの方法」をよく読んでつけてください。<br>"
        f"Please rank the 5 images based on {CRITERIA_GUIDANCE_EN[criterion_idx]}"
        f"</p>"
    )
    rules_part = (
        "<b>ランキングの方法 (よく読んでください) / How to Rank:</b>"
        "<ul>"
        "<li><b>全く同じ表情の画像には、同じ順位</b>を付けてください。(Assign the <b>same rank</b> to identical expressions.)</li>"
        "<li><b>少しでも違う表情の画像には、違う順位</b>を付けてください。(Assign <b>different ranks</b> to different expressions.)</li>"
        "<li><b>必ず1位から</b>順位を付けてください。(You <b>must</b> assign a rank of '1' to at least one image.)</li>"
        "<li>同順位がある場合、<b>その人数分だけ次の順位を飛ばしてください</b>。(When you have ties, <b>skip the next rank(s) accordingly</b>.)"
        "<ul>"
        "<li>例1: 1位が2つある場合、次は3位になります (Ex. 1: If there are two '1st' places, the next rank is '3rd'. e.g., <code>1, 1, 3, 4, 5</code>).</li>"
        "<li>例2: 1位が1つ、2位が3つある場合、次は5位になります (Ex. 2: If there is one '1st' and three '2nd' places, the next rank is '5th'. e.g., <code>1, 2, 2, 2, 5</code>).</li>"
        "</ul></li></ul>"
    )
    ai_note_part = (
        "<p style='font-size: 0.9em; color: #555;'>"
        "設問や英単語の意味についてはAIに質問したり検索したりしても構いません。ただし、画像そのものが示す感情をAIに質問するのはお控えください。<br>"
        "You are welcome to use AI or web search to understand the questions or the meaning of English words. However, please refrain from asking an AI about the emotion shown in the images themselves."
        "</p>"
    )
    combined_instructions = f"{guidance_part}<hr>{rules_part}<hr>{ai_note_part}"

    if criterion_idx == 0:
        state["image_orders"] = {}

    if criterion_name not in state["image_orders"]:
        conditions_shuffled = random.sample(CONDITIONS, len(CONDITIONS))
        state["image_orders"][criterion_name] = conditions_shuffled

    current_image_order = state["image_orders"][criterion_name]
    image_updates = []
    for cond_name in current_image_order:
        prediction = current_data["predictions"][cond_name]
        img_path = get_image_path_from_prediction(prediction, state)

        if img_path and os.path.exists(img_path):
            try:
                pil_img = Image.open(img_path).convert('RGB')
                masked_img = create_masked_image(pil_img, state)
                image_updates.append(gr.update(value=masked_img))
            except Exception as e:
                logger.error(f"Failed to open or mask image {img_path}: {e}")
                image_updates.append(gr.update(value=None))
        else:
            image_updates.append(gr.update(value=None))

    saved_ranks_dict = state["evaluation_results"].get(prompt_text, {}).get("ranks", {}).get(criterion_name)
    if saved_ranks_dict:
        label_to_condition = {label: cond for label, cond in zip(IMAGE_LABELS, current_image_order)}
        condition_to_label = {v: k for k, v in label_to_condition.items()}
        state["current_ranks"] = {
            condition_to_label[cond]: rank for cond, rank in saved_ranks_dict.items() if cond in condition_to_label
        }
    else:
        state["current_ranks"] = {label: None for label in IMAGE_LABELS}

    button_updates = _create_button_updates(state)

    is_alignment_criterion = (criterion_name == "Alignment")
    abs_group_update = gr.update(visible=is_alignment_criterion)
    saved_abs_score = state["evaluation_results"].get(prompt_text, {}).get("absolute_score")
    state["current_absolute_score"] = saved_abs_score if is_alignment_criterion else None

    abs_button_updates = []
    for i in range(1, 8):
        variant = 'primary' if i == state["current_absolute_score"] else 'secondary'
        abs_button_updates.append(gr.update(variant=variant))

    abs_group_worst_update = gr.update(visible=is_alignment_criterion)
    saved_abs_score_worst = state["evaluation_results"].get(prompt_text, {}).get("absolute_score_worst")
    state["current_absolute_score_worst"] = saved_abs_score_worst if is_alignment_criterion else None

    abs_button_worst_updates = []
    for i in range(1, 8):
        variant = 'primary' if i == state["current_absolute_score_worst"] else 'secondary'
        abs_button_worst_updates.append(gr.update(variant=variant))

    return [
        gr.update(value=progress_text),
        gr.update(value=combined_instructions),
        gr.update(value=prompt_display_text),
        gr.update(value="", visible=False),
        *image_updates,
        *button_updates,
        abs_group_update,
        *abs_button_updates,
        abs_group_worst_update,
        *abs_button_worst_updates,
        gr.update(
            interactive=(state["current_prompt_index"] > 0 or state["current_criterion_index"] > 0)),
        gr.update(interactive=True)
    ]


def validate_and_navigate(state: dict):
    ranks = state["current_ranks"]
    error_msg = None
    criterion_name = CRITERIA[state["current_criterion_index"]]
    is_alignment_criterion = (criterion_name == "Alignment")

    if any(r is None for r in ranks.values()):
        error_msg = "Please rank all 5 images. / 5つすべての画像を評価してください。"
    elif 1 not in ranks.values():
        error_msg = "You must assign a rank of '1' to at least one image. / 最低1つは「1位」を付けてください。"
    elif is_alignment_criterion and state["current_absolute_score"] is None:
        error_msg = "Please provide an absolute score for the BEST matching image (1-7). / 最も一致している画像について、絶対評価（1～7）を選択してください。"
    elif is_alignment_criterion and state["current_absolute_score_worst"] is None:
        error_msg = "Please provide an absolute score for the WORST matching image (1-7). / 最も一致していない画像について、絶対評価（1～7）を選択してください。"
    elif (
            is_alignment_criterion
            and state["current_absolute_score"] is not None
            and state["current_absolute_score_worst"] is not None
            and state["current_absolute_score_worst"] > state["current_absolute_score"]
    ):
        error_msg = (
            "The score for the WORST matching image cannot be higher than the score for the BEST matching image.<br>"
            "「最も一致していない画像」のスコアが「最も一致している画像」のスコアを上回ることはできません。"
        )

    if error_msg:
        # ★ 修正点: no_change_updates の要素数を 52 に変更
        no_change_updates = [gr.update()] * 52
        # error_display は all_eval_outputs の4番目(インデックス3)のコンポーネント
        no_change_updates[3] = gr.update(
            value=f"<p class='feedback red' style='font-size: 1.2em; text-align: center;'>{error_msg}</p>",
            visible=True)
        # 戻り値の総数が54個 (state, tab_exportの更新, all_eval_outputsの更新) になるように調整
        return state, gr.update(), *no_change_updates


    sorted_ranks = sorted(list(ranks.values()))
    rank_counts = Counter(sorted_ranks)
    i = 0
    while i < len(sorted_ranks):
        current_rank = sorted_ranks[i]
        count = rank_counts[current_rank]
        if i + count < len(sorted_ranks):
            next_rank = sorted_ranks[i + count]
            expected_next_rank = current_rank + count
            if next_rank < expected_next_rank:
                error_msg = f"Ranking rule violation (tie-breaking). After {count} instance(s) of rank '{current_rank}', the next rank must be >= {expected_next_rank}, but it is '{next_rank}'. / ランキングのルール違反です。'{current_rank}'位が{count}個あるため、次の順位は{expected_next_rank}位以上である必要がありますが、'{next_rank}'位が入力されています。"
                break
        i += count
    if error_msg:
        # ★ 修正点: こちらも同様に no_change_updates の要素数を 52 に変更
        no_change_updates = [gr.update()] * 52
        # error_display は all_eval_outputs の4番目(インデックス3)のコンポーネント
        no_change_updates[3] = gr.update(
            value=f"<p class='feedback red' style='font-size: 1.2em; text-align: center;'>{error_msg}</p>",
            visible=True)
        # 戻り値の総数が54個になるように調整
        return state, gr.update(), *no_change_updates

    prompt_idx = state["shuffled_indices"][state["current_prompt_index"]]
    current_data = state["all_eval_data"][prompt_idx]
    prompt_text = current_data["prompt"]
    current_image_order = state["image_orders"][criterion_name]

    label_to_condition = {label: cond for label, cond in zip(IMAGE_LABELS, current_image_order)}
    ranks_by_condition = {label_to_condition[label]: rank for label, rank in ranks.items()}

    if "ranks" not in state["evaluation_results"][prompt_text]:
        state["evaluation_results"][prompt_text]["ranks"] = {}
    if "orders" not in state["evaluation_results"][prompt_text]:
        state["evaluation_results"][prompt_text]["orders"] = {}

    state["evaluation_results"][prompt_text]["ranks"][criterion_name] = ranks_by_condition
    state["evaluation_results"][prompt_text]["orders"][criterion_name] = current_image_order

    logger.info(
        f"Saved rank for P:{state['participant_id']}, Prompt:'{prompt_text}', Criterion:{criterion_name}, Ranks:{ranks_by_condition}, Orders:{current_image_order}, ")


    if is_alignment_criterion:
        state["evaluation_results"][prompt_text]["absolute_score"] = state["current_absolute_score"]
        state["evaluation_results"][prompt_text]["absolute_score_worst"] = state["current_absolute_score_worst"]

        logger.info(
            f"Saved absolute scores for P:{state['participant_id']}, Prompt:'{prompt_text}', "
            f"Best Score:{state['current_absolute_score']}, Worst Score:{state['current_absolute_score_worst']}")


    state["current_criterion_index"] += 1
    if state["current_criterion_index"] >= len(CRITERIA):
        state["current_criterion_index"] = 0
        state["current_prompt_index"] += 1

    if state["current_prompt_index"] >= len(state["all_eval_data"]):
        state["end_time"] = datetime.now()

    eval_panel_updates = display_current_prompt_and_criterion(state)
    export_tab_update = gr.update(interactive=(state.get("end_time") is not None))
    return state, export_tab_update, *eval_panel_updates


def navigate_previous(state: dict):
    state["current_criterion_index"] -= 1
    if state["current_criterion_index"] < 0:
        state["current_criterion_index"] = len(CRITERIA) - 1
        state["current_prompt_index"] -= 1
    state["current_prompt_index"] = max(0, state["current_prompt_index"])

    # ★ 修正: display_current_prompt_and_criterion は state を返さないので、ここで state を返す
    return state, *display_current_prompt_and_criterion(state)


def export_results(participant_id, alignment_reason, naturalness_reason, attractiveness_reason, optional_comment,
                   state: dict):
    if not alignment_reason.strip() or not naturalness_reason.strip() or not attractiveness_reason.strip():
        error_msg = "<p class='feedback red'>Please fill in the reasoning for all three criteria (Alignment, Naturalness, Attractiveness). / 3つの評価基準（一致度, 自然さ, 魅力度）すべての判断理由を記入してください。</p>"
        return None, error_msg

    if not participant_id:
        return None, "<p class='feedback red'>Participant ID is missing. / 参加者IDがありません。</p>"

    output_dir = os.path.join(BASE_RESULTS_DIR, participant_id)
    os.makedirs(output_dir, exist_ok=True)
    filename = f"evaluation_results_{participant_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    filepath = os.path.join(output_dir, filename)

    duration = (state["end_time"] - state["start_time"]).total_seconds() if state.get(
        "start_time") and state.get("end_time") else None

    prompt_to_category = {item["prompt"]: item["category"] for item in state["all_eval_data"]}

    final_results_list = []
    for prompt, data in state["evaluation_results"].items():
        if not data: continue
        # (以降、データ構造の作成部分は変更なし)
        ranks_data = data.get("ranks", {})
        orders_data = data.get("orders", {})

        final_results_list.append({
            "prompt": prompt,
            "prompt_category": prompt_to_category.get(prompt),
            "image_order_alignment": orders_data.get("Alignment", []),
            "image_order_naturalness": orders_data.get("Naturalness", []),
            "image_order_attractiveness": orders_data.get("Attractiveness", []),
            "alignment_ranks": ranks_data.get("Alignment", {}),
            "naturalness_ranks": ranks_data.get("Naturalness", {}),
            "attractiveness_ranks": ranks_data.get("Attractiveness", {}),
            "alignment_absolute_score": data.get("absolute_score"),
            "alignment_absolute_score_worst": data.get("absolute_score_worst")
        })

    export_data = {
        "metadata": {
            "participant_id": participant_id,
            "export_timestamp": datetime.now().isoformat(),
            "total_prompts_evaluated": len(final_results_list),
            "evaluation_duration_seconds": duration,
            "reasoning": {
                "alignment": alignment_reason,
                "naturalness": naturalness_reason,
                "attractiveness": attractiveness_reason,
            },
            "optional_comment": optional_comment,
        },
        "results": final_results_list
    }

    logger.info(f"Exporting results for participant metadata: {export_data['metadata']}")

    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, ensure_ascii=False, indent=2)
        logger.info(f"Successfully exported results to: {filepath}")
    except Exception as e:
        logger.error(f"Failed to write export file: {e}")
        return None, f"<p class='feedback red'>An error occurred during file export: {e}</p>"

    upload_link = "https://drive.google.com/drive/folders/1ujIPF-67Y6OG8qBm1TYG3FsmuYxqSAcR?usp=drive_link"
    status_message = f"""
    <div class='feedback green' style='text-align: left;'>
    <p><b>エクスポートが完了しました。/ Export complete.</b></p>
    <p>上のボタンからJSONファイルをダウンロードし、指定された場所にアップロードして実験を終了してください。ご協力ありがとうございました。</p>
    <p>Please download the JSON file and upload it to the designated location. Thank you for your cooperation.</p>
    <p><b>アップロード先 / Upload to:</b> <a href='{upload_link}' target='_blank'>{upload_link}</a></p>
    </div>"""
    return gr.update(value=filepath, visible=True), status_message


def create_gradio_interface():
    css = """
    /* (CSSは変更なし) */
    .gradio-container { font-family: 'Arial', sans-serif; }
    .feedback { padding: 10px; border-radius: 5px; font-weight: bold; text-align: center; margin-top: 10px; }
    .feedback.green { background-color: #e6ffed; color: #2f6f4a; }
    .feedback.red { background-color: #ffe6e6; color: #b30000; }
    .image-label { font-size: 2.5em; font-weight: bold; margin-bottom: 10px; color: #333; }
    .prompt-display { text-align: center; margin-bottom: 15px; padding: 10px; background-color: #f0f8ff; border-radius: 8px; }
    .rank-btn-row { justify-content: center; gap: 5px !important; }
    .rank-btn { min-width: 65px !important; max-width: 65px !important; height: 45px !important; font-size: 1.2em !important; font-weight: bold !important; border-radius: 8px !important; border: 1px solid #ccc !important; }
    .rank-btn.secondary { background: #f0f0f0 !important; color: #333 !important; }
    .rank-btn.secondary:hover { background: #e0e0e0 !important; border-color: #bbb !important; }
    .absolute-eval-group { border: 1px solid #ddd; border-radius: 8px; padding: 15px; margin-top: 20px; }
    .instructions-container { padding: 15px; border: 1px solid #ccc; border-radius: 8px; margin-bottom: 20px; background-color: #fafafa; }
    .instructions-container ul { padding-left: 20px; margin-top: 5px; margin-bottom: 5px; }
    .instructions-container hr { margin: 15px 0; }
    """

    with gr.Blocks(title="Expression Evaluation Experiment", css=css) as app:
        # ★ 修正: セッションステート (gr.State) を初期化
        initial_state = {
            "participant_id": None,
            "data_loaded": False,
            "all_eval_data": [],
            "shuffled_indices": [],
            "current_prompt_index": 0,
            "current_criterion_index": 0,
            "image_mapping": {},
            "image_dir": "",
            "evaluation_results": {},
            "image_orders": {},
            "start_time": None,
            "end_time": None,
            "current_ranks": {},
            "current_absolute_score": None,
            "current_absolute_score_worst": None,
            "hide_bbox_dict": {},
        }
        state = gr.State(value=initial_state)

        gr.Markdown("# Text-to-Expression Evaluation Experiment / テキストからの表情生成 評価実験")

        with gr.Tabs() as tabs:
            # (UI定義は変更なし)
            with gr.TabItem("1. Setup / セットアップ") as tab_setup:
                gr.Markdown("## (A) Participant Information / 参加者情報")
                gr.Markdown("Please enter your participant ID and click 'Confirm'. / 参加者IDを入力して「確定」を押してください。")
                with gr.Row():
                    participant_id_input = gr.Textbox(label="Participant ID", placeholder="e.g., P01")
                    confirm_id_btn = gr.Button("Confirm / 確定", variant="primary")
                setup_warning = gr.Markdown(visible=False)
                with gr.Group(visible=False) as setup_main_group:
                    gr.Markdown("---")
                    gr.Markdown("## (B) Instructions & Data Loading / 注意事項とデータ読み込み")
                    gr.Markdown(
                        """<div style='padding: 15px; border: 1px solid #f0ad4e; border-radius: 5px; background-color: #fcf8e3;'>
                        <h4>注意事項 / Instructions</h4>
                        <ul>
                            <li><b>この作業はPCで行ってください。/ Please perform this task on a PC.</b></li>
                            <li>途中で止めずに最後まで続けてください。ファイルをアップロードして完了となります。/ Please continue until the end. The experiment is complete when you upload the file.</li>
                            <li>ブラウザーをリロードしないでください (データが破損します)。/ Do not reload the browser (this will corrupt the data).</li>
                            <li><b>指示と注意書きをよく読んでください。</b></li>
                        </ul></div>""")
                    gr.Markdown(
                        "Click the button below to load your evaluation data. / 下のボタンを押して、評価データを読み込んでください。")
                    load_data_btn = gr.Button("Load Data / データ読み込み", variant="primary")
                    setup_status = gr.Markdown("Waiting to start...")

            with gr.TabItem("2. Evaluation / 評価", interactive=False) as tab_evaluation:
                progress_text = gr.Markdown("Prompt 0 / 0")
                combined_instructions_display = gr.Markdown("", elem_classes="instructions-container")
                prompt_display = gr.Markdown("", elem_classes="prompt-display")
                error_display = gr.Markdown(visible=False)

                image_components = []
                rank_buttons = []
                with gr.Row(equal_height=False):
                    for label in IMAGE_LABELS:
                        with gr.Column(scale=1):
                            with gr.Group():
                                gr.Markdown(f"<div class='image-label' style='text-align: center;'>{label}</div>")
                                img = gr.Image(type="pil", show_label=False, height=300)
                                image_components.append(img)
                                with gr.Row(elem_classes="rank-btn-row"):
                                    rank_list = ["1位", "2位", "3位", "4位", "5位"]
                                    for rank_val in range(1, 6):
                                        btn = gr.Button(str(rank_list[rank_val - 1]), variant='secondary',
                                                        elem_classes="rank-btn")
                                        rank_buttons.append(btn)

                with gr.Group(visible=False, elem_classes="absolute-eval-group") as absolute_eval_group_best:
                    gr.Markdown("---")
                    gr.Markdown(
                        "#### 絶対評価 (Best) / Absolute Score (Best)\n最もテキストと一致している画像について、どの程度一致しているかを評価してください。\n(Please evaluate the degree of alignment for the image that **best** matches the text.)")
                    absolute_score_buttons = []
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown(
                                "<p style='text-align: right; margin-top: 10px;'>1 (全く一致してない / Not at all)</p>")
                        with gr.Column(scale=3):
                            with gr.Row(elem_classes="rank-btn-row"):
                                for i in range(1, 8):
                                    btn = gr.Button(str(i), variant='secondary', elem_classes="rank-btn")
                                    absolute_score_buttons.append(btn)
                        with gr.Column(scale=1):
                            gr.Markdown("<p style='text-align: left; margin-top: 10px;'>7 (完全に一致 / Absolutely)</p>")

                with gr.Group(visible=False, elem_classes="absolute-eval-group") as absolute_eval_group_worst:
                    gr.Markdown(
                        "#### 絶対評価 (Worst) / Absolute Score (Worst)\n最もテキストと一致していない画像について、どの程度一致しているかを評価してください。\n(Please evaluate the degree of alignment for the image that **least** matches the text.)")
                    absolute_score_worst_buttons = []
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown(
                                "<p style='text-align: right; margin-top: 10px;'>1 (全く一致してない / Not at all)</p>")
                        with gr.Column(scale=3):
                            with gr.Row(elem_classes="rank-btn-row"):
                                for i in range(1, 8):
                                    btn = gr.Button(str(i), variant='secondary', elem_classes="rank-btn")
                                    absolute_score_worst_buttons.append(btn)
                        with gr.Column(scale=1):
                            gr.Markdown("<p style='text-align: left; margin-top: 10px;'>7 (完全に一致 / Absolutely)</p>")

                with gr.Row():
                    prev_btn = gr.Button("← Previous / 前へ", interactive=False)
                    next_btn = gr.Button("Save & Next / 保存して次へ →", variant="primary")

            with gr.TabItem("3. Export / エクスポート", interactive=False) as tab_export:
                gr.Markdown("## (C) Final Comments & Export / 最終コメントとエクスポート")
                gr.Markdown(
                    "Thank you for completing the evaluation. Please provide the reasoning for your judgments for each criterion below. / 評価お疲れ様でした。以下の各評価基準について、判断の理由をご記入ください。")

                with gr.Group():
                    gr.Markdown("#### Reasoning for Judgments (Required) / 判断理由（必須）")
                    alignment_reason_box = gr.Textbox(label="Alignment / 一致度", lines=3,
                                                      placeholder="Why did you rank them this way for alignment? / なぜ一致度について、このようなランキングにしましたか？")
                    naturalness_reason_box = gr.Textbox(label="Naturalness / 自然さ", lines=3,
                                                        placeholder="Why did you rank them this way for naturalness? / なぜ自然さについて、このようなランキングにしましたか？")
                    attractiveness_reason_box = gr.Textbox(label="Attractiveness / 魅力度", lines=3,
                                                           placeholder="Why did you rank them this way for attractiveness? / なぜ魅力度について、このようなランキングにしましたか？")

                with gr.Group():
                    gr.Markdown("#### Overall Comments (Optional) / 全体的な感想（任意）")
                    optional_comment_box = gr.Textbox(label="Any other comments? / その他、実験全体に関するご意見・ご感想",
                                                      lines=4,
                                                      placeholder="e.g., 'Image B often looked the most natural.' / 例：「Bの画像が最も自然に見えることが多かったです。」")

                gr.Markdown("---")
                gr.Markdown(
                    "Finally, click the button below to export your results. / 最後に、下のボタンを押して結果をエクスポートしてください。")
                export_btn = gr.Button("Export Results / 結果をエクスポート", variant="primary")
                download_file = gr.File(label="Download JSON", visible=False)
                export_status = gr.Markdown()

        # --- Event Handlers (★ 修正: stateを入出力に追加) ---
        def check_and_confirm_id(pid, state):
            pid = pid.strip()
            if re.fullmatch(r"P\d{2}", pid):
                state["participant_id"] = pid
                return state, gr.update(visible=False), gr.update(visible=True)
            else:
                error_msg = "<p class='feedback red'>Invalid ID. Must be 'P' followed by two digits (e.g., P01). / 無効なIDです。「P」と数字2桁の形式（例: P01）で入力してください。</p>"
                return state, gr.update(value=error_msg, visible=True), gr.update(visible=False)

        confirm_id_btn.click(check_and_confirm_id, [participant_id_input, state],
                             [state, setup_warning, setup_main_group])
        load_data_btn.click(load_evaluation_data, [participant_id_input, state],
                            [state, setup_status, load_data_btn, tab_evaluation])

        all_eval_outputs = [
            progress_text, combined_instructions_display, prompt_display, error_display, *image_components,
            *rank_buttons,
            absolute_eval_group_best, *absolute_score_buttons,
            absolute_eval_group_worst, *absolute_score_worst_buttons,
            prev_btn, next_btn
        ]

        btn_idx = 0
        for label in IMAGE_LABELS:
            for rank_val in range(1, 6):
                btn = rank_buttons[btn_idx]
                # functools.partial を使うことで、state以外の引数を固定できます
                btn.click(
                    partial(handle_rank_button_click, label, rank_val),
                    [state],
                    [state, *rank_buttons]
                )
                btn_idx += 1

        for i, btn in enumerate(absolute_score_buttons):
            btn.click(
                partial(handle_absolute_score_click, i + 1),
                [state],
                [state, *absolute_score_buttons]
            )

        for i, btn in enumerate(absolute_score_worst_buttons):
            btn.click(
                partial(handle_absolute_score_worst_click, i + 1),
                [state],
                [state, *absolute_score_worst_buttons]
            )

        tab_evaluation.select(display_current_prompt_and_criterion, [state], all_eval_outputs)

        # next_btn と prev_btn の出力に state を追加
        next_btn.click(validate_and_navigate, [state], [state, tab_export, *all_eval_outputs])
        prev_btn.click(navigate_previous, [state], [state, *all_eval_outputs])

        export_tab_interactive_components = [alignment_reason_box, naturalness_reason_box, attractiveness_reason_box,
                                             optional_comment_box, export_btn]

        def on_select_export_tab(state):
            if state.get("end_time"):
                return [gr.update(interactive=True)] * 5
            return [gr.update(interactive=False)] * 5

        tab_export.select(on_select_export_tab, [state], export_tab_interactive_components)

        export_btn.click(
            export_results,
            [participant_id_input, alignment_reason_box, naturalness_reason_box, attractiveness_reason_box,
             optional_comment_box, state],
            [download_file, export_status]
        )

    return app


if __name__ == "__main__":
    os.makedirs(LOG_DIR, exist_ok=True)
    log_file_path = os.path.join(LOG_DIR, "evaluation_ui_log_{time}.log")

    random.seed(datetime.now().timestamp())
    logger.remove()
    logger.add(sys.stderr, level="INFO")
    logger.add(log_file_path, rotation="10 MB")

    app = create_gradio_interface()
    app.launch(share=True, debug=True)