app.py

import os
import glob
import json
import time
import random
import re
from functools import partial

import numpy as np
import torch
from PIL import Image
import gradio as gr
from scipy.spatial.distance import cdist
from sklearn.metrics import silhouette_samples
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from argparse import ArgumentParser
from datetime import datetime
from loguru import logger
import io
from transformers import CLIPModel, CLIPProcessor

# Set environment variable for Gradio
os.environ["GRADIO_ALLOWED_PATHS"] = "False"

# Global variables to hold state
GLOBAL_STATE = {
    "participant_id": None,
    "embedding_created": False,
    "initial_clustering_done": False,
    "initial_labels_updated": False,
    "boundary_samples_labeled": False,
    "embeddings": None,
    "image_paths": None,
    "cluster_labels": None,
    "current_k": 8,
    "label_dict": {},
    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    "model": None,
    "preprocess": None,
    "constraints": [],
    "silhouette_scores": None,
    "boundary_samples": [],
    "boundary_labels": {},
    "boundary_index": 0,
    "next_cluster_id": None,
    "representative_images": {},
    "random_seed": 42,
    "original_images": {},
    "show_bbox_dict": {},
    "hide_bbox_dict": {},
    "new_label_inputs": {},
    "initial_kmeans_labels": None,
    "initial_label_dict": None,
    "kmeans_matched_labels": None,
    "kmeans_matched_reps": None,
    "evaluation_choice": None,
    "evaluation_layout_is_left_hitl": True,
    "annotation_start_time": None,
    "annotation_end_time": None,
    "annotation_duration_seconds": None,
}


# --------------------------------------------------------------
# Helper functions
# --------------------------------------------------------------
def load_bbox_json(bbox_json_path):
    try:
        with open(bbox_json_path, 'r', encoding='utf-8') as f:
            bbox_data = json.load(f)
        GLOBAL_STATE["show_bbox_dict"] = bbox_data.get("Show", {})
        GLOBAL_STATE["hide_bbox_dict"] = bbox_data.get("Hide", {})
    except Exception as e:
        logger.error(f"Failed to load bounding box JSON: {e}")


def create_masked_image(image: Image.Image):
    hide_bbox_dict = GLOBAL_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


# --------------------------------------------------------------
# CLIP model and Embedding
# --------------------------------------------------------------
# ------------------ 絶対に必要な部分 ------------------
def load_clip_model(device="cuda"):
    hf_model_id = "apple/DFN2B-CLIP-ViT-L-14"
    model = CLIPModel.from_pretrained(hf_model_id).to(device)
    model.eval()
    BASE_PROC = "openai/clip-vit-large-patch14"
    preprocess = CLIPProcessor.from_pretrained(BASE_PROC)
    return model, preprocess
# -------------------------------------------------------

# def load_clip_model(device="cuda"):
#     # ローカルのモデルが保存されているパス
#     local_model_path = "./clip_model_cache"
#
#     # ローカルパスの存在をチェック
#     if not os.path.isdir(local_model_path):
#         # エラーメッセージをより具体的にする
#         error_msg = f"CLIP model not found at the specified local path: {local_model_path}. Please make sure the model is downloaded and placed in the correct directory."
#         logger.error(error_msg)
#         raise FileNotFoundError(error_msg)
#
#     logger.info(f"Loading model and processor from local path: {local_model_path}")
#
#     # ローカルパスからモデルとプロセッサを読み込む
#     model = CLIPModel.from_pretrained(local_model_path).to(device)
#     preprocess = CLIPProcessor.from_pretrained(local_model_path)
#
#     model.eval()
#     return model, preprocess


def resize_with_padding(img, target_size=(224, 224), fill_color=(0, 0, 0)):
    img.thumbnail(target_size, Image.Resampling.LANCZOS)
    new_img = Image.new("RGB", target_size, fill_color)
    paste_x = (target_size[0] - img.width) // 2
    paste_y = (target_size[1] - img.height) // 2
    new_img.paste(img, (paste_x, paste_y))
    return new_img


def get_face_region_embeddings(image: Image.Image, model, preprocess, device="cuda"):
    show_bbox_dict = GLOBAL_STATE.get("show_bbox_dict", {})
    if not show_bbox_dict:
        logger.warning("No 'Show' regions defined. Using whole image.")
        processed_image = resize_with_padding(image)
        inputs = preprocess(images=processed_image, return_tensors="pt").to(device)
        with torch.no_grad():
            emb = model.get_image_features(**inputs)
            return (emb / emb.norm(dim=-1, keepdim=True)).cpu().numpy().flatten()

    region_embeddings = []
    for region_name, coords in show_bbox_dict.items():
        try:
            box = (coords['left'], coords['top'], coords['right'], coords['bottom'])
            region_img = image.crop(box)
            processed_region = resize_with_padding(region_img)
            inputs = preprocess(images=processed_region, return_tensors="pt").to(device)
            with torch.no_grad():
                emb = model.get_image_features(**inputs)
                region_embeddings.append(emb / emb.norm(dim=-1, keepdim=True))
        except Exception as e:
            logger.error(f"Failed to process region {region_name}: {e}")
            continue

    if not region_embeddings:
        logger.error("No regions were successfully processed.")
        return None

    mean_emb = torch.mean(torch.cat(region_embeddings, dim=0), dim=0, keepdim=True)
    final_emb = mean_emb / mean_emb.norm(dim=-1, keepdim=True)
    return final_emb.cpu().numpy().flatten()


def create_embeddings(image_folder, bbox_json_path):
    start_time = time.time()
    yield "Loading bounding box info..."
    load_bbox_json(bbox_json_path)

    if GLOBAL_STATE["model"] is None:
        yield "Loading CLIP model..."
        GLOBAL_STATE["model"], GLOBAL_STATE["preprocess"] = load_clip_model(device=GLOBAL_STATE["device"])

    yield "Scanning for images..."
    supported_formats = ['*.jpg', '*.jpeg', '*.png']
    image_paths = sorted([p for fmt in supported_formats for p in glob.glob(os.path.join(image_folder, fmt))])
    GLOBAL_STATE["image_paths"] = image_paths
    GLOBAL_STATE["original_images"] = {}

    embeddings = []
    total_images = len(image_paths)
    yield f"Processing {total_images} images..."

    for i, img_path in enumerate(image_paths):
        if i % 10 == 0:
            yield f"Processing image {i + 1}/{total_images}: {os.path.basename(img_path)}"
        try:
            img = Image.open(img_path).convert('RGB')
            GLOBAL_STATE["original_images"][img_path] = img
            emb = get_face_region_embeddings(img, GLOBAL_STATE["model"], GLOBAL_STATE["preprocess"],
                                             device=GLOBAL_STATE["device"])
            if emb is not None:
                embeddings.append(emb)
        except Exception as e:
            logger.error(f"Error processing {img_path}: {e}")

    if not embeddings:
        yield "## Failed to create any valid embeddings."
        return

    GLOBAL_STATE["embeddings"] = np.vstack(embeddings)
    GLOBAL_STATE["embedding_created"] = True
    logger.info(f"Embedding creation took {time.time() - start_time:.2f} seconds.")
    yield "## Successfully Created Embeddings / 埋め込みの作成に成功しました"


# --------------------------------------------------------------
# Clustering and Annotation Logic
# --------------------------------------------------------------

# ############################################################
# ### 欠陥3の修正: 代表画像の選択ロジックを修正 ###
# ############################################################
def get_representative_images(labels, n_representatives=3):
    """
    Get representative images for each cluster.
    Prioritizes manually fixed samples, then selects images closest to cluster centers.
    """
    if GLOBAL_STATE["embeddings"] is None: return {}
    embeddings = GLOBAL_STATE["embeddings"]
    image_paths = GLOBAL_STATE["image_paths"]

    # Calculate cluster centers (still needed for non-fixed samples)
    cluster_centers = {}
    for label in np.unique(labels):
        cluster_points = embeddings[labels == label]
        if len(cluster_points) > 0:
            cluster_centers[label] = np.mean(cluster_points, axis=0)
        else:
            # Fallback for empty clusters
            cluster_centers[label] = np.zeros(embeddings.shape[1])

    cluster_representatives = {}
    for label in np.unique(labels):
        cluster_indices = np.where(labels == label)[0]
        if len(cluster_indices) == 0:
            cluster_representatives[label] = []
            continue

        # --- Prioritize fixed samples ---
        fixed_idx_in_cluster = sorted(list(set([
            c[1] for c in GLOBAL_STATE.get("constraints", [])
            if c[0] == "fixed-cluster" and c[2] == label and c[1] in cluster_indices
        ])))
        other_idx_in_cluster = [i for i in cluster_indices if i not in fixed_idx_in_cluster]

        # --- Prepare the list of representatives ---
        # Add fixed samples first, with distance 0.0 (as they are definitional)
        representatives = [(image_paths[i], 0.0, i) for i in fixed_idx_in_cluster]

        # --- Fill remaining spots with closest samples ---
        remaining_spots = n_representatives - len(representatives)
        if remaining_spots > 0 and other_idx_in_cluster:
            center = cluster_centers[label]
            other_embeddings = embeddings[other_idx_in_cluster]
            distances = cdist(other_embeddings, [center]).flatten()

            # Get the original indices sorted by distance
            sorted_other_indices_global = [other_idx_in_cluster[i] for i in np.argsort(distances)]

            for i in range(min(remaining_spots, len(sorted_other_indices_global))):
                original_idx = sorted_other_indices_global[i]
                # Find the distance corresponding to this original index
                dist_idx = np.where(np.array(other_idx_in_cluster) == original_idx)[0][0]
                distance = distances[dist_idx]
                representatives.append((image_paths[original_idx], distance, original_idx))

        cluster_representatives[label] = representatives

    return cluster_representatives


# ############################################################
# ### 欠陥2の修正: 制約自動生成ロジックをヘルパー関数として追加 ###
# ############################################################
def set_representative_samples_as_fixed(representative_images):
    """
    Sets the most representative samples of each cluster as fixed-cluster constraints.
    This anchors the key samples to their clusters for subsequent constrained clustering.
    """
    new_constraints = []
    # For each cluster, set the most representative sample (closest to centroid) as fixed
    for cluster_id, representatives in representative_images.items():
        if not representatives:
            continue
        # Get the most representative sample (first in the list, sorted by distance)
        _, _, sample_idx = representatives[0]
        # Create a fixed-cluster constraint
        constraint = ("fixed-cluster", sample_idx, int(cluster_id))
        # Check if this constraint already exists
        if constraint not in GLOBAL_STATE["constraints"]:
            GLOBAL_STATE["constraints"].append(constraint)
            new_constraints.append(constraint)
            logger.info(f"Added fixed constraint for cluster {cluster_id}, sample {sample_idx}")
    return new_constraints


# ############################################################
# ### 欠陥2の修正: 初期クラスタリングに関数を組み込む ###
# ############################################################
def perform_initial_clustering():
    if GLOBAL_STATE["embeddings"] is None: return "Embeddings not created."
    k = GLOBAL_STATE["current_k"]
    kmeans = KMeans(n_clusters=k, random_state=GLOBAL_STATE["random_seed"], n_init=10)

    cluster_labels = kmeans.fit_predict(GLOBAL_STATE["embeddings"])

    GLOBAL_STATE["cluster_labels"] = cluster_labels
    GLOBAL_STATE["initial_kmeans_labels"] = cluster_labels.copy()

    GLOBAL_STATE["silhouette_scores"] = silhouette_samples(GLOBAL_STATE["embeddings"], GLOBAL_STATE["cluster_labels"])
    GLOBAL_STATE["label_dict"] = {str(i): "" for i in range(k)}
    GLOBAL_STATE["next_cluster_id"] = k

    # Get representative images
    representatives = get_representative_images(GLOBAL_STATE["cluster_labels"], n_representatives=1)
    GLOBAL_STATE["representative_images"] = representatives

    # --- 欠陥2の修正: 自動で制約を生成 ---
    # 既存の制約をクリア
    GLOBAL_STATE["constraints"] = []
    # 代表サンプルを固定する制約を追加
    set_representative_samples_as_fixed(representatives)
    logger.info(
        f"Automatically generated {len(GLOBAL_STATE['constraints'])} fixed constraints for representative samples.")
    # --- 修正ここまで ---

    GLOBAL_STATE["initial_clustering_done"] = True
    return f"Performed initial clustering with {k} clusters and automatically fixed representative samples. / {k}個のクラスタで初期クラスタリングを実行し、代表サンプルを自動的に固定しました。"


def update_initial_labels(*labels):
    for i in range(GLOBAL_STATE["current_k"]):
        if not labels[i] or not labels[i].strip():
            error_msg = "<p class='feedback red'>Please fill in all cluster labels before proceeding. / 先に進む前に、すべてのクラスタラベルを入力してください。</p>"
            return error_msg, gr.update(), gr.update(visible=False), gr.update(interactive=False)

    for i in range(GLOBAL_STATE["current_k"]):
        if ',' in labels[i]:
            error_msg = f"<p class='feedback red'>Error in Cluster {i}: Label '{labels[i]}' contains a comma. Please use a single word or phrase. / クラスタ{i}のエラー: ラベル「{labels[i]}」にカンマが含まれています。単一の単語またはフレーズを使用してください。</p>"
            return error_msg, gr.update(), gr.update(visible=False), gr.update(interactive=False)

    for i, label in enumerate(labels):
        if i < GLOBAL_STATE["current_k"]:
            GLOBAL_STATE["label_dict"][str(i)] = label.strip()

    GLOBAL_STATE["initial_label_dict"] = GLOBAL_STATE["label_dict"].copy()

    GLOBAL_STATE["initial_labels_updated"] = True
    status = "## Labels updated. / ラベルを更新しました。"
    done_msg = "<p class='feedback green'>初期ラベルを保存しました。<b>このラベルは最後に変更可能です。</b>次の「境界サンプル」タブに進んでください。<br>Initial labels saved. <b>These can be changed later.</b> Proceed to the 'Boundary Samples' tab.</p>"
    return status, gr.update(visible=False), gr.update(value=done_msg, visible=True), gr.update(interactive=True)


def extract_boundary_samples(silhouette_threshold=0.2, max_samples=30, min_embedding_distance=1e-3):
    scores = GLOBAL_STATE["silhouette_scores"]
    embeddings = GLOBAL_STATE["embeddings"]
    if scores is None or embeddings is None:
        logger.warning("Silhouette scores or embeddings not available.")
        return "Scores or embeddings not found."

    low_score_indices = np.where(scores < silhouette_threshold)[0]
    sorted_candidate_indices = low_score_indices[np.argsort(scores[low_score_indices])]

    if len(sorted_candidate_indices) == 0:
        GLOBAL_STATE["boundary_samples"] = []
        logger.info("No samples found below the silhouette threshold.")
        return "No samples found below the silhouette threshold."

    final_boundary_samples = []
    selected_embeddings = []

    for idx in sorted_candidate_indices:
        if len(final_boundary_samples) >= max_samples:
            break
        candidate_emb = embeddings[idx].reshape(1, -1)
        if not selected_embeddings:
            final_boundary_samples.append(idx)
            selected_embeddings.append(candidate_emb)
            continue
        distances = cdist(candidate_emb, np.vstack(selected_embeddings), 'cosine')
        min_dist = np.min(distances)
        if min_dist > min_embedding_distance:
            final_boundary_samples.append(idx)
            selected_embeddings.append(candidate_emb)

    GLOBAL_STATE["boundary_samples"] = final_boundary_samples
    GLOBAL_STATE["boundary_index"] = 0
    GLOBAL_STATE["boundary_labels"] = {}
    GLOBAL_STATE["boundary_samples_labeled"] = False

    logger.info(
        f"Extracted {len(final_boundary_samples)} diverse boundary samples (min_dist={min_embedding_distance}).")
    return f"Extracted {len(GLOBAL_STATE['boundary_samples'])} diverse boundary samples. / 多様性を考慮した境界サンプルを{len(GLOBAL_STATE['boundary_samples'])}個抽出しました。"


def find_top_n_closest_clusters(idx, n=3):
    sample_emb = GLOBAL_STATE["embeddings"][idx]
    current_cid = GLOBAL_STATE["cluster_labels"][idx]

    centers, cids = [], []
    for cid_iter in np.unique(GLOBAL_STATE["cluster_labels"]):
        points = GLOBAL_STATE["embeddings"][GLOBAL_STATE["cluster_labels"] == cid_iter]
        if len(points) > 0:
            centers.append(np.mean(points, axis=0))
            cids.append(cid_iter)

    if not centers: return []
    distances = cdist([sample_emb], centers, 'euclidean')[0]
    return [(cid, dist) for dist, cid in sorted(zip(distances, cids)) if cid != current_cid][:n]


def get_single_representative_image(cluster_id, exclude_idx=None):
    # get_representative_imagesが返すキーは整数なので、それに合わせる
    reps = GLOBAL_STATE["representative_images"].get(int(cluster_id), [])
    for img_path, _, img_idx in reps:
        if img_idx != exclude_idx:
            img = GLOBAL_STATE["original_images"].get(img_path) or Image.open(img_path).convert('RGB')
            return create_masked_image(img), img_path
    return None, None


def get_current_boundary_sample():
    if not GLOBAL_STATE["boundary_samples"]:
        return None, "No boundary samples.", "", [gr.update()] * 20

    if GLOBAL_STATE["boundary_index"] >= len(GLOBAL_STATE["boundary_samples"]):
        return None, "All samples annotated.", "", [gr.update()] * 20

    idx = GLOBAL_STATE["boundary_samples"][GLOBAL_STATE["boundary_index"]]
    img_path = GLOBAL_STATE["image_paths"][idx]
    current_cluster = GLOBAL_STATE["boundary_labels"].get(idx, GLOBAL_STATE["cluster_labels"][idx])
    current_label = GLOBAL_STATE["label_dict"].get(str(current_cluster), f"Cluster {current_cluster}")
    top_alternatives = find_top_n_closest_clusters(idx, n=3)

    import base64
    from io import BytesIO
    current_reps_html = "<h4>Representative sample of current cluster (このクラスタの代表例):</h4><div style='display: flex; gap: 10px; flex-wrap: wrap;'>"
    current_rep_img, _ = get_single_representative_image(current_cluster)
    if current_rep_img:
        try:
            masked_img = create_masked_image(current_rep_img)
            masked_img.thumbnail((360, 360))
            buffer = BytesIO()
            masked_img.save(buffer, format="PNG")
            b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
            current_reps_html += f"<img src='data:image/png;base64,{b64}' style='border-radius: 4px; border: 1px solid #ddd;'>"
        except Exception:
            pass
    current_reps_html += "</div>"
    alternatives_html = ""
    for alt_cid, _ in top_alternatives:
        alt_label = GLOBAL_STATE["label_dict"].get(str(alt_cid), f"Cluster {alt_cid}")
        alt_rep_img, _ = get_single_representative_image(alt_cid, exclude_idx=idx)
        alt_rep_b64 = ""
        if alt_rep_img:
            alt_rep_img.thumbnail((240, 240))
            buffer = BytesIO()
            alt_rep_img.save(buffer, format="PNG")
            alt_rep_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        alternatives_html += f"""
        <div style="flex: 1; min-width: 240px; padding: 10px; border: 1px solid #ddd; border-radius: 8px; text-align: center; background-color: #fafafa;">
            <div style="font-weight: bold; font-size: 1.1em; margin-bottom: 8px;">{alt_cid}: {alt_label}</div>
            <img src='data:image/png;base64,{alt_rep_b64}' style='max-width: 240px; max-height: 240px; object-fit: contain; border-radius: 4px;'>
        </div>"""

    status_text = f"""
    <div style="font-size: 16px; line-height: 1.6; padding: 15px; background-color: #f5f5f5; border-radius: 8px;">
        <div style="font-size: 20px; font-weight: bold; margin-bottom: 10px; color: #333;">
            Sample {GLOBAL_STATE['boundary_index'] + 1} / {len(GLOBAL_STATE['boundary_samples'])}
        </div>
        <div style="padding: 12px; background-color: #e8f5e9; border-left: 4px solid #4caf50; border-radius: 4px; margin-bottom: 15px;">
            <div style="font-weight: bold; font-size: 18px; margin-bottom: 8px;">
                Current cluster / 現在のクラスタ:
                <span style='font-weight: normal; font-size: 0.9em; margin-left: 10px;'>この画像が示す感情のテキストが左の画像に完全に一致していれば Next を押してください．少しでも違っていたら，新しい感情のテキストあるいは別の感情のテキストを割り当ててください．テキストが完全に一致していれば良く，画像が完全に一致している必要はありません．/ If the emotion text shown in this image matches the image on the left exactly, press Next. If there is even the slightest difference, assign a new emotion text or a different emotion text. The text must match exactly, but the images do not need to match exactly.</span>
            </div> 
            <div style="font-size: 22px; margin-bottom: 10px;">{current_cluster}: {current_label}</div>
            {current_reps_html}
        </div>
        <div style="font-weight: bold; font-size: 18px; margin-bottom: 8px;">
            Similar Alternative Clusters / 似ているクラスタ候補:
            <span style='font-weight: normal; font-size: 0.9em; margin-left: 10px;'>これは参考の画像であり，必ずしもここから選択する必要はありません．ただし，完全に同一の表情があればここから選択してください．もし同一のものがなければ，必要に応じて下のボタンから既存のクラスタを割り当てたり，新規クラスタを作成してください．/ This is a reference image, and you do not necessarily have to select from here. However, if there is an identical expression, please select it from here. If there is no identical expression, assign an existing cluster from the buttons below or create a new cluster as necessary.</span>
        </div>
        <div style="display: flex; gap: 15px; flex-wrap: wrap; justify-content: space-around;">
            {alternatives_html if alternatives_html else "<p>No other clusters available.</p>"}
        </div>
    </div>"""

    button_updates = []
    k = GLOBAL_STATE["current_k"]
    for i in range(20):
        if i < k:
            label = GLOBAL_STATE["label_dict"].get(str(i), f"Cluster {i}")
            if i == current_cluster:
                button_updates.append(gr.update(value=f"{i}: {label} (Current)", visible=True, variant='primary'))
            else:
                button_updates.append(gr.update(value=f"{i}: {label}", visible=True, variant='secondary'))
        else:
            button_updates.append(gr.update(visible=False))

    img = GLOBAL_STATE["original_images"].get(img_path) or Image.open(img_path).convert('RGB')
    return create_masked_image(img), status_text, "", button_updates


def assign_to_cluster(cluster_id):
    if not GLOBAL_STATE["boundary_samples"] or GLOBAL_STATE["boundary_index"] >= len(GLOBAL_STATE["boundary_samples"]):
        error_msg = "Error: No boundary sample selected. Please extract samples first."
        return error_msg, gr.update(), gr.update(), *([gr.update()] * 20), gr.update(), gr.update()

    idx = GLOBAL_STATE["boundary_samples"][GLOBAL_STATE["boundary_index"]]

    # --- 既存の制約を削除し、新しい制約を追加 ---
    # このサンプルに対する既存のfixed-cluster制約を削除
    GLOBAL_STATE["constraints"] = [c for c in GLOBAL_STATE["constraints"] if
                                   not (c[0] == "fixed-cluster" and c[1] == idx)]
    # 新しい制約を追加
    GLOBAL_STATE["constraints"].append(("fixed-cluster", idx, int(cluster_id)))
    # --- 修正ここまで ---

    GLOBAL_STATE["boundary_labels"][idx] = int(cluster_id)
    GLOBAL_STATE["boundary_samples_labeled"] = True
    img, status, _, button_updates = get_current_boundary_sample()
    return f"Assigned to cluster {cluster_id}.", img, status, *button_updates, gr.update(), gr.update()


def handle_create_new_cluster(new_label):
    if not GLOBAL_STATE["boundary_samples"] or GLOBAL_STATE["boundary_index"] >= len(GLOBAL_STATE["boundary_samples"]):
        error_msg = "Error: No boundary sample selected. Please extract samples first."
        return error_msg, gr.update(), gr.update(), gr.update(), *([gr.update()] * 20), gr.update(), gr.update()

    stripped_label = new_label.strip()
    error_msg = None
    if not stripped_label:
        error_msg = "<p class='feedback red'>New cluster label cannot be empty. / 新規クラスタのラベルは空にできません。</p>"
    elif ',' in stripped_label:
        error_msg = f"<p class='feedback red'>Label '{new_label}' contains a comma. Please use a single word or phrase. / ラベル「{new_label}」にカンマが含まれています。単一の単語またはフレーズを使用してください。</p>"

    if error_msg:
        dummy_updates = [gr.update()] * 23
        return [error_msg] + dummy_updates + [gr.update(), gr.update()]

    idx = GLOBAL_STATE["boundary_samples"][GLOBAL_STATE["boundary_index"]]
    new_cid = GLOBAL_STATE["next_cluster_id"]
    GLOBAL_STATE["next_cluster_id"] += 1
    GLOBAL_STATE["current_k"] += 1
    GLOBAL_STATE["label_dict"][str(new_cid)] = stripped_label

    # --- 既存の制約を削除し、新しい制約を追加 ---
    GLOBAL_STATE["constraints"] = [c for c in GLOBAL_STATE["constraints"] if
                                   not (c[0] == "fixed-cluster" and c[1] == idx)]
    GLOBAL_STATE["constraints"].append(("fixed-cluster", idx, new_cid))

    GLOBAL_STATE["boundary_labels"][idx] = new_cid
    GLOBAL_STATE["boundary_samples_labeled"] = True

    # ############################################################
    # ### ★★★★★ ここが重要な修正点 ★★★★★ ###
    # ############################################################
    # グローバルなクラスタラベル配列を直接更新する
    GLOBAL_STATE["cluster_labels"][idx] = new_cid

    # 更新されたクラスタラベル全体を使って代表画像を再計算する
    GLOBAL_STATE["representative_images"] = get_representative_images(
        GLOBAL_STATE["cluster_labels"],
        n_representatives=1
    )
    # ############################################################

    img, status, _, button_updates = get_current_boundary_sample()
    return [f"Created new cluster ({new_cid}): {stripped_label}", img, status, ""] + button_updates + [gr.update(),
                                                                                                       gr.update()]
# ############################################################
# ### 欠陥1の修正: COP-KMeansアルゴリズムとヘルパー関数を追加 ###
# ############################################################
def kmpp_initialization(dataset, k, fixed_assignments):
    """KMeans++ initialization, respecting fixed assignments."""
    rng = np.random.RandomState(GLOBAL_STATE["random_seed"])
    fixed_clusters = {}
    for idx, cluster in fixed_assignments:
        if cluster not in fixed_clusters:
            fixed_clusters[cluster] = []
        fixed_clusters[cluster].append(idx)

    centers = []
    remaining_k = k

    # Initialize centers for clusters with fixed assignments
    for cluster_id in range(k):
        if cluster_id in fixed_clusters:
            centers.append(dataset[fixed_clusters[cluster_id]].mean(axis=0))
            remaining_k -= 1

    # Fill remaining centers using k-means++ logic
    if remaining_k > 0:
        if not centers:
            first_idx = rng.choice(len(dataset))
            centers.append(dataset[first_idx])
            remaining_k -= 1

        while remaining_k > 0:
            distances = np.min([np.sum((dataset - center) ** 2, axis=1) for center in centers], axis=0)
            # Avoid division by zero if all distances are zero
            if distances.sum() == 0:
                next_idx = rng.choice(np.where(distances == 0)[0])
            else:
                next_idx = rng.choice(len(dataset), p=distances / distances.sum())
            centers.append(dataset[next_idx])
            remaining_k -= 1

    return np.array(centers)


def satisfies_constraints(idx, cluster_id, labels, ml, cl):
    """Check if assigning point idx to cluster_id satisfies all constraints."""
    for i, j in ml:
        if i == idx and labels[j] != -1 and labels[j] != cluster_id: return False
        if j == idx and labels[i] != -1 and labels[i] != cluster_id: return False
    for i, j in cl:
        if i == idx and labels[j] == cluster_id: return False
        if j == idx and labels[i] == cluster_id: return False
    return True


def find_nearest_cluster(point, centers, labels, ml, cl, idx):
    """Find the nearest cluster center that satisfies all constraints."""
    distances = np.sum((centers - point) ** 2, axis=1)
    sorted_idx = np.argsort(distances)
    for cluster_id in sorted_idx:
        if satisfies_constraints(idx, cluster_id, labels, ml, cl):
            return cluster_id
    return None


def cop_kmeans(dataset, k, ml=None, cl=None, fixed_labels=None, max_iter=300, tol=1e-4, anchor_weight=10.0):
    """Constrained K-means clustering with Must-Link, Cannot-Link, and weighted anchor constraints."""
    ml = ml or []
    cl = cl or []
    fixed_labels = fixed_labels or {}

    fixed_assignments = list(fixed_labels.items())
    cluster_centers = kmpp_initialization(dataset, k, fixed_assignments)
    cluster_labels = np.full(len(dataset), -1)

    for idx, cluster in fixed_assignments:
        cluster_labels[idx] = cluster

    for iter_count in range(max_iter):
        old_labels = cluster_labels.copy()
        violated = False
        for i in range(len(dataset)):
            if i in fixed_labels: continue
            nearest_cluster = find_nearest_cluster(dataset[i], cluster_centers, cluster_labels, ml, cl, i)
            if nearest_cluster is None:
                violated = True
                break
            cluster_labels[i] = nearest_cluster

        if violated:
            logger.warning("Constraints violated, cannot find a valid assignment. Re-initializing might be needed.")
            return None

        for j in range(k):
            idx_in_cluster = np.where(cluster_labels == j)[0]
            if len(idx_in_cluster) == 0: continue

            fixed_idx_in_cluster = [i for i in idx_in_cluster if i in fixed_labels and fixed_labels[i] == j]
            normal_idx_in_cluster = [i for i in idx_in_cluster if i not in fixed_idx_in_cluster]

            anchor_emb = dataset[fixed_idx_in_cluster] if fixed_idx_in_cluster else np.empty((0, dataset.shape[1]))
            others_emb = dataset[normal_idx_in_cluster] if normal_idx_in_cluster else np.empty((0, dataset.shape[1]))

            if len(fixed_idx_in_cluster) > 0:
                w_anchor = np.full(len(fixed_idx_in_cluster), anchor_weight)
                w_others = np.ones(len(normal_idx_in_cluster))
                all_emb = np.vstack([anchor_emb, others_emb])
                w = np.concatenate([w_anchor, w_others])[:, None]
                cluster_centers[j] = (w * all_emb).sum(axis=0) / w.sum()
            elif len(normal_idx_in_cluster) > 0:
                cluster_centers[j] = others_emb.mean(axis=0)

        if np.all(cluster_labels == old_labels):
            break

    return cluster_labels


# ############################################################
# ### 欠陥1の修正: 制約付きクラスタリング実行関数を修正 ###
# ############################################################
def perform_constrained_clustering():
    """Perform constrained clustering based on user-provided constraints."""
    if GLOBAL_STATE["embeddings"] is None:
        return "No embeddings available. Please create embeddings first.", None

    # Prepare constraints for COP-KMeans
    ml_constraints = []
    cl_constraints = []
    fixed_labels = {}

    for constraint_type, idx1, idx2_or_cluster in GLOBAL_STATE["constraints"]:
        if constraint_type == "must-link":
            ml_constraints.append((idx1, idx2_or_cluster))
        elif constraint_type == "cannot-link":
            cl_constraints.append((idx1, idx2_or_cluster))
        elif constraint_type == "fixed-cluster":
            fixed_labels[idx1] = idx2_or_cluster

    k = GLOBAL_STATE["current_k"]

    # Run constrained K-means
    new_labels = cop_kmeans(
        GLOBAL_STATE["embeddings"],
        k,
        ml=ml_constraints,
        cl=cl_constraints,
        fixed_labels=fixed_labels
    )

    if new_labels is None:
        return "Constrained clustering failed. Constraints may be inconsistent.", None

    GLOBAL_STATE["cluster_labels"] = new_labels
    GLOBAL_STATE["silhouette_scores"] = silhouette_samples(
        GLOBAL_STATE["embeddings"],
        GLOBAL_STATE["cluster_labels"]
    )

    # Update representative images after re-clustering
    GLOBAL_STATE["representative_images"] = get_representative_images(
        GLOBAL_STATE["cluster_labels"],
        n_representatives=1
    )

    return f"Performed constrained clustering with {k} clusters.", None


# --------------------------------------------------------------
# Finalization and Export
# --------------------------------------------------------------
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NumpyEncoder, self).default(obj)


def update_final_labels(*labels):
    for i in range(GLOBAL_STATE["current_k"]):
        # Add a check to ensure the label is not None before processing
        if labels[i] is None:
            error_msg = f"<p class='feedback red'>Error in Cluster {i}: Final label cannot be empty. Please provide a label. / クラスタ{i}のエラー: 最終ラベルは空にできません。ラベルを入力してください。</p>"
            return gr.update(value=error_msg, visible=True), gr.update(visible=False), gr.update(interactive=False)

        stripped_label = labels[i].strip()
        if not stripped_label:
            error_msg = f"<p class='feedback red'>Error in Cluster {i}: Final label cannot be empty. / クラスタ{i}のエラー: 最終ラベルは空にできません。</p>"
            return gr.update(value=error_msg, visible=True), gr.update(visible=False), gr.update(interactive=False)
        if ',' in stripped_label:
            error_msg = f"<p class='feedback red'>Error in Cluster {i}: Label '{labels[i]}' contains a comma. Please use a single word or phrase. / クラスタ{i}のエラー: ラベル「{labels[i]}」にカンマが含まれています。単一の単語またはフレーズを使用してください。</p>"
            return gr.update(value=error_msg, visible=True), gr.update(visible=False), gr.update(interactive=False)

    for i, label in enumerate(labels):
        if i < GLOBAL_STATE["current_k"]:
            GLOBAL_STATE["label_dict"][str(i)] = label.strip()

    end_time = time.time()
    GLOBAL_STATE["annotation_end_time"] = end_time
    if GLOBAL_STATE.get("annotation_start_time"):
        duration = end_time - GLOBAL_STATE["annotation_start_time"]
        GLOBAL_STATE["annotation_duration_seconds"] = duration
        logger.info(f"Annotation task time: {duration:.2f} seconds.")

    guidance_message = """
    <div class='feedback green' style='text-align: left;'>
    <p><b>Final labels updated. / 最終ラベルを更新しました。</b></p>
    <p>Please take a <b>10-minute break</b> before proceeding to the next step. After the break, please move to the "Evaluation & Export" tab.</p>
    <p>次のステップに進む前に<b>10分間の休憩</b>を取ってください。休憩後、「評価とエクスポート」タブに進んでください。</p>
    </div>
    """
    return gr.update(value="", visible=False), gr.update(value=guidance_message, visible=True), gr.update(
        interactive=True)
def export_results(participant_id):
    if not participant_id:
        return None, "<p class='feedback red'>Participant ID is missing.</p>"

    args = get_parameters()
    pipeline_path = os.path.join(args.output_dir, participant_id)
    os.makedirs(pipeline_path, exist_ok=True)
    results_path = os.path.join(args.results_dir, participant_id)
    os.makedirs(results_path, exist_ok=True)
    json_filename = f"{args.labels_name}_{participant_id}.json"
    pipeline_json_path = os.path.join(pipeline_path, json_filename)
    results_json_path = os.path.join(results_path, json_filename)

    logger.info(f"Preparing to export results for participant {participant_id}")

    try:
        import sklearn
        import transformers
        versions = {
            "torch": torch.__version__, "sklearn": sklearn.__version__,
            "transformers": transformers.__version__, "gradio": gr.__version__,
        }
    except Exception as e:
        versions = {"error": f"Could not get versions: {e}"}

    export_data = []
    if GLOBAL_STATE.get("image_paths") and GLOBAL_STATE.get("cluster_labels") is not None:
        initial_kmeans_labels = GLOBAL_STATE.get("initial_kmeans_labels")
        initial_label_dict = GLOBAL_STATE.get("initial_label_dict", {})

        for i, (img_path, cluster_id) in enumerate(zip(GLOBAL_STATE["image_paths"], GLOBAL_STATE["cluster_labels"])):
            manually_labeled = any(c[0] == "fixed-cluster" and c[1] == i for c in GLOBAL_STATE.get("constraints", []))
            initial_cluster_id = None
            initial_emotion_label = ""
            if initial_kmeans_labels is not None and i < len(initial_kmeans_labels):
                initial_cluster_id = int(initial_kmeans_labels[i])
                initial_emotion_label = initial_label_dict.get(str(initial_cluster_id), "")

            export_data.append({
                "image_file": os.path.basename(img_path),
                "image_path": img_path,
                "cluster_id": int(cluster_id),
                "emotion_label": GLOBAL_STATE.get("label_dict", {}).get(str(cluster_id), ""),
                "initial_kmeans_cluster_id": initial_cluster_id,
                "initial_kmeans_emotion_label": initial_emotion_label,
                "manually_labeled": manually_labeled,
                "boundary_sample": i in GLOBAL_STATE.get("boundary_samples", [])
            })

    json_data = {
        "metadata": {
            "participant_id": participant_id,
            "num_clusters_hitl": GLOBAL_STATE.get("current_k"),
            "num_clusters_initial": len(initial_label_dict) if initial_label_dict else None,
            "cluster_labels_initial": initial_label_dict,
            "cluster_labels": GLOBAL_STATE["label_dict"],
            "constraints": [{"type": c[0], "sample_idx": c[1], "value": c[2]} for c in
                            GLOBAL_STATE.get("constraints", [])],
            "evaluation_choice": GLOBAL_STATE.get("evaluation_choice"),
            "evaluation_layout_is_left_hitl": GLOBAL_STATE.get("evaluation_layout_is_left_hitl"),
            "annotation_duration_seconds": GLOBAL_STATE.get("annotation_duration_seconds"),
            "objective_scores": {
                "hitl_silhouette": GLOBAL_STATE.get("hitl_silhouette_score"),
                "kmeans_matched_silhouette": GLOBAL_STATE.get("kmeans_silhouette_score"),
            },
            "system_info": {
                "random_seed_used": GLOBAL_STATE.get("random_seed"),
                "script_parameters": vars(args),
                "library_versions": versions,
            },
        },
        "samples": export_data
    }

    try:
        with open(pipeline_json_path, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, ensure_ascii=False, indent=2, cls=NumpyEncoder)
        logger.info(f"Successfully saved pipeline input to: {pipeline_json_path}")
        with open(results_json_path, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, ensure_ascii=False, indent=2, cls=NumpyEncoder)
        logger.info(f"Successfully saved final result to: {results_json_path}")
    except Exception as e:
        logger.error(f"Failed to write export files: {e}")
        error_msg = f"<p class='feedback red'>An error occurred during file export: {e}</p>"
        return gr.update(visible=False), error_msg

    upload_link = "https://drive.google.com/drive/folders/1ZhsVRTs-Tb_-9mxvIyxAb6K2HbnuRPVw?usp=sharing"
    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>アップロード先:</b> <a href='{upload_link}' target='_blank'>{upload_link}</a></p>
    </div>"""
    return gr.update(value=results_json_path, visible=True), status_message


# --------------------------------------------------------------
# Gradio Interface
# --------------------------------------------------------------
def create_gradio_interface():
    args = get_parameters()
    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.blue { background-color: #e6f3ff; color: #0056b3; }
    .feedback.red { background-color: #ffe6e6; color: #b30000; }
    .gr-form {gap: 10px !important;}
    label.label-info span {font-size: 1em !important; color: #444 !important;}
    .cluster-group { margin-bottom: 25px !important; padding-top: 15px !important; border-top: 1px solid #eee !important;}
    .styled-button { border: 1px solid #ccc !important; font-weight: bold !important; }
    .gr-button-primary { background-color: #007bff !important; color: white !important; border-color: #007bff !important; }
    .gr-button-secondary { background-color: #f8f9fa !important; color: #343a40 !important; border-color: #ced4da !important; }
    .gr-button-secondary:hover { background-color: #e2e6ea !important; }
    .cluster-gallery-container { border: 1px solid #ddd; border-radius: 8px; padding: 10px; margin-bottom: 15px; background-color: #fdfdfd; }
    """

    with gr.Blocks(title="HitL Clustering Experiment", css=css) as app:
        gr.Markdown("# Human-in-the-Loop Clustering Experiment / HitLクラスタリング実験")

        with gr.Tabs() as tabs:
            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>途中で止めたりせず最後まで続けてください。ファイルをアップロードして完了となります。/ Please continue until the end without stopping. The experiment is complete when you upload the file.</li><li>参加者番号は「P数字2桁」の形式でお願いします (例: P01)。/ Please use the format "P" followed by two digits for the participant number (e.g., P01).</li><li>静かな環境で集中して行ってください。/ Please perform the task in a quiet and focused environment.</li><li>ブラウザーをリロードしないでください (データが破損します)。/ Do not reload the browser (this will corrupt the data).</li><li>ブラウザーはSafari以外を使ってください。/ Please use a browser other than Safari.</li></ul></div>""")
                    with gr.Row():
                        gr.Textbox(label="Dataset Directory", value=args.dataset_dir, interactive=False)
                        gr.Textbox(label="Image Directory", value=args.image_dir, interactive=False)
                        gr.Textbox(label="Bounding Box JSON", value=args.bbox_json, interactive=False)

                    gr.Markdown("## (C) Create Image Embeddings / 画像埋め込みの作成")
                    gr.Markdown(
                        "Click the button below to start. This may take a few minutes. **Please wait.** / 下のボタンをクリックして開始します。**しばらくお待ち下さい。**")
                    create_emb_btn = gr.Button("Create Embeddings / 埋め込みを作成", variant="primary")
                    emb_status = gr.Markdown("Waiting to start...")
                    embedding_done_msg = gr.Markdown(visible=False)

            with gr.TabItem("2. Initial Clustering / 初期クラスタリング", interactive=False) as tab_cluster:
                gr.Markdown("## (D) Initial Clustering / 初期クラスタリング")
                gr.Markdown(
                    "First, click the button to perform initial clustering. / まず、ボタンをクリックして初期クラスタリングを実行してください。")
                cluster_btn = gr.Button("Perform Initial Clustering / クラスタリングを開始", variant="primary")
                cluster_status = gr.Markdown()

                with gr.Group(visible=False) as initial_annotation_group:
                    gr.Markdown("---")
                    gr.Markdown("## (E) Annotate Initial Clusters / 初期クラスタの注釈付け")
                    gr.Markdown("""
                    **この作業はPCで行ってください。**<br>
                    この画像に写っているキャラクターの顔の上半分の表情が、どのような感情を表しているように見えるか、英語で記述してください。<br>
                    - 単語（例: happy）でも、短いフレーズ（例1: very angry, 例2: a little surprised, 例3: sad with little surprise）でも構いません。<br>
                    - 必ず一つの単語またはフレーズを入れてください．決して複数の単語やフレーズ(例：happy, sad) は入れてはいけません．<br>
                    - 他の画像や他の人の回答と同じような表現になっても問題ありません。<br>
                    **AIに英語の表現を聞くのは構いませんが 、画像の感情そのものをAIに尋ねるのは避けてください。**
                    <hr>
                    **Please perform this task on a PC.**<br>
                    Please describe in English what emotion the upper half of the character's face in this image appears to express.<br>
                    - It can be a single word (e.g., happy) or a short phrase (e.g., very angry, a little surprised, sad with little surprise).<br>
                    - You must enter only one word or phrase. Never enter multiple words or phrases (e.g., happy, sad).<br>
                    - It is okay if your expression is similar to those for other images or other people's answers.<br>
                    **You may ask an AI for English expressions, but please avoid asking the AI about the emotion of the image itself.**
                    """)

                    initial_labels_inputs = []
                    with gr.Column():
                        for i in range(20):  # Max 20 clusters
                            with gr.Group(elem_classes="cluster-group",
                                          visible=(i < GLOBAL_STATE["current_k"])) as initial_group:
                                gr.Markdown(f"### Cluster {i}")
                                gallery = gr.HTML()
                                label_input = gr.Textbox(
                                    label="ここに上半分の表情が示している感情を英語で入力してください．単語（例: happy）でも、短いフレーズ（例1: very angry, 例2: a little surprised, 例3: sad with little surprise）でも構いません。一つの単語またはフレーズを入れてください．決して複数の単語やフレーズ(例：「happy, sad」) は入れてはいけません． "
                                          "(Enter here in English the emotion indicated by the upper half of the facial expression. It can be a word (e.g. happy) or a short phrase (e.g. 1: very angry, 2: a little surprised, 3: sad with a little surprise). Never include more than one word or phrase (e.g., “happy, sad”).)",
                                    elem_classes="label-info")
                                initial_labels_inputs.append((initial_group, gallery, label_input))

                    update_initial_btn = gr.Button("Update Cluster Labels / クラスタラベルを更新", variant="primary")
                    initial_labels_status = gr.Markdown()
                    initial_clustering_done_msg = gr.Markdown(visible=False)

            with gr.TabItem("3. Boundary Samples & Annotation / 境界サンプルの注釈付け", interactive=False) as tab_boundary:
                gr.Markdown("## (F) Extract and Label Boundary Samples / 境界サンプルの抽出と注釈付け")
                gr.Markdown(
                    """まず「境界サンプルを抽出」ボタンを押してください。次に、表示される画像を確認します。
                    <br>- もし画像が現在のクラスタに合っていると思えば、何もしないで「次のサンプル」へ進んでください。
                    <br>- もし別の感情のクラスタに属すると思えば、「既存クラスタへの割り当て」から適切なクラスタのボタンを押してください。
                    <br>- もし既存のどのクラスタにも当てはまらない新しい感情だと思えば、「新規クラスタの作成」に感情ラベルを入力し、「新規クラスタを作成」ボタンを押してください。
                    <br>- すべての画像のラベルに変更がなければ，ラベリングが完了し次に進みます。
                    <hr>First, press the 'Extract Boundary Samples' button. Then, check the image displayed.
                    <br>- If you think the image fits the current cluster, do nothing and proceed to the 'Next Sample'.
                    <br>- If you think it belongs to a different emotion cluster, press the appropriate cluster button under 'Assign to Existing Cluster'.
                    <br>- If you think it represents a new emotion, enter a label under 'Create New Cluster' and press the button.
                    <br>- If the labels of all images are unchanged, labeling is complete and proceed.""")

                extract_btn = gr.Button("Extract Boundary Samples / 境界サンプルを抽出", variant="primary")
                extract_status = gr.Markdown()

                with gr.Group(visible=False) as boundary_annotation_group:
                    with gr.Row():
                        sample_image = gr.Image(label="Boundary Sample", type="pil")
                        sample_status = gr.HTML()
                    with gr.Row():
                        prev_btn = gr.Button("← Previous", elem_classes="styled-button")
                        next_btn = gr.Button("Next →", elem_classes="styled-button")

                    gr.Markdown("#### Assign to Existing Cluster / 既存クラスタへの割り当て")
                    cluster_buttons = []
                    with gr.Column():
                        for i in range(0, 20, 5):
                            with gr.Row():
                                for j in range(i, min(i + 5, 20)):
                                    btn = gr.Button(f"Cluster {j}", visible=(j < GLOBAL_STATE["current_k"]),
                                                    elem_classes="styled-button")
                                    cluster_buttons.append(btn)

                    gr.Markdown("#### Create New Cluster / 新規クラスタの作成")
                    with gr.Row():
                        new_label_text = gr.Textbox(
                            label="ここに上半分の表情が示している感情を英語で入力してください． "
                                  "(Enter here in English the emotion indicated by the upper half of the facial expression.￥",
                            elem_classes="label-info")
                        create_new_btn = gr.Button("Create New Cluster", elem_classes="styled-button")

                    recluster_guidance = gr.Markdown(
                        """
                        <div class='feedback blue' style='text-align: left;'>
                        <p><b>アノテーションを続けてください / Please continue annotating:</b></p>
                        <p>すべての境界サンプルの確認が完了すると（「Next →」で最後のサンプルに到達すると）、再クラスタリングボタンが表示されます。</p>
                        <p>The re-clustering button will appear once you have reviewed all boundary samples (by reaching the last sample with "Next →").</p>
                        </div>
                        """,
                        visible=False
                    )

                    with gr.Group(visible=False) as recluster_group:
                        attention_text = gr.Markdown(
                            """
                            <div class='feedback blue' style='text-align: left;'>
                            <p><b>注意 / Attention:</b></p>
                            <p>以下の「制約を適用して再クラスタリング」ボタンを押す前に，境界サンプルの割当てがすべて正しく行われているかを確認してください。<br>
                            <p>Before pressing the 'Apply Constrained Clustering' button below, check that all boundary sample assignments are correct.</p>
                            </div>
                            """
                        )
                        recluster_btn = gr.Button("Apply Constrained Clustering / 制約を適用して再クラスタリング",
                                                  variant="primary")

                    recluster_status = gr.Markdown()
                    boundary_done_msg = gr.Markdown(visible=False)

            with gr.TabItem("4. Finalize / 最終化", interactive=False) as tab_finalize:
                gr.Markdown("## (G) Finalize Cluster Labels / クラスタラベルの最終化")
                gr.Markdown(
                    "Review the final clusters. Please check for spelling mistakes. / 最終的なクラスタを確認し、スペルミスがないか確認してください。")

                final_labels_inputs = []
                with gr.Column():
                    for i in range(20):
                        with gr.Group(visible=False, elem_classes="cluster-group") as final_group:
                            gr.Markdown(f"### Cluster {i}")
                            gallery = gr.HTML()
                            label_input = gr.Textbox(label=f"Final Label for Cluster {i}")
                            final_labels_inputs.append((final_group, gallery, label_input))

                update_final_btn = gr.Button("Update Final Labels / 最終ラベルを更新", variant="primary")
                final_status = gr.Markdown(visible=False)
                final_guidance_msg = gr.Markdown(visible=False)

            with gr.TabItem("5. Evaluation & Export / 評価とエクスポート", interactive=False) as tab_evaluation:
                gr.Markdown("## (H) Comparative Evaluation / 比較評価")
                gr.Markdown(
                    """
                    <h3 style='text-align: center; margin-bottom: 20px;'>
                    以下に2種類のクラスタリング結果が表示されています。<br>
                    それぞれのクラスタの代表画像をよく見て、データセット全体の表情をより良く分割していると感じる方を正直に選んでください。
                    </h3>
                    """
                )

                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Clustering Result A")
                        gallery_A = gr.HTML()
                        choose_A_btn = gr.Button("Choose Result A / 結果Aを選択", variant="primary")
                    with gr.Column(scale=1):
                        gr.Markdown("### Clustering Result B")
                        gallery_B = gr.HTML()
                        choose_B_btn = gr.Button("Choose Result B / 結果Bを選択", variant="primary")

                evaluation_status = gr.Markdown(visible=False)

                with gr.Group(visible=False) as export_group:
                    gr.Markdown("---")
                    gr.Markdown("## (I) Export Results / 結果のエクスポート")
                    gr.Markdown(
                        "Thank you for your evaluation. Please click the button below to export your results and upload the file to the designated location. / 評価ありがとうございました。下のボタンを押して結果をエクスポートし、指定された場所にファイルをアップロードしてください。")
                    export_btn = gr.Button("Export Results / 結果をエクスポート", variant="primary")
                    download_file = gr.File(label="Download JSON", visible=False)
                    export_status = gr.Markdown()

        # --- Event Handlers ---
        def check_and_confirm_id(pid):
            pid = pid.strip()
            if re.fullmatch(r"P\d{2}", pid):
                GLOBAL_STATE["participant_id"] = pid
                return gr.update(visible=False), gr.update(visible=True)
            else:
                error_message = """
                <p class='feedback red'>
                <b>Invalid Participant ID / 無効な参加者IDです</b><br>
                The ID must be in the format 'P' followed by exactly two digits (e.g., P01, P23).<br>
                IDは「P」に続いて数字2桁の形式である必要があります（例: P01, P23）。
                </p>
                """
                return gr.update(value=error_message, visible=True), gr.update(visible=False)

        confirm_id_btn.click(check_and_confirm_id, [participant_id_input], [setup_warning, setup_main_group])

        def handle_create_embeddings_flow():
            args = get_parameters()
            final_status = "Embedding creation failed."
            for status in create_embeddings(args.image_dir, args.bbox_json):
                yield status, gr.update(), gr.update(), gr.update()
                final_status = status

            GLOBAL_STATE["annotation_start_time"] = time.time()
            logger.info(f"Annotation timer started at {GLOBAL_STATE['annotation_start_time']}.")

            done_msg = "<p class='feedback green'>Embedding creation complete. Please proceed to the 'Initial Clustering' tab. / 埋め込みの作成が完了しました。「初期クラスタリング」タブに進んでください。</p>"
            yield final_status, gr.update(visible=False), gr.update(value=done_msg, visible=True), gr.update(
                interactive=True)

        create_emb_btn.click(handle_create_embeddings_flow, [],
                             [emb_status, create_emb_btn, embedding_done_msg, tab_cluster])

        def create_cluster_gallery_html(reps, cid, n_images=1):
            # repsのキーは整数なので、cidを整数に変換
            cid = int(cid)
            if cid not in reps: return ""
            html = "<div style='display: flex; flex-wrap: wrap; gap: 10px;'>"
            import base64
            from io import BytesIO
            for img_path, _, _ in reps[cid][:n_images]:
                try:
                    img = GLOBAL_STATE["original_images"].get(img_path) or Image.open(img_path).convert('RGB')
                    masked_img = create_masked_image(img)
                    masked_img.thumbnail((360, 360))
                    buffer = BytesIO()
                    masked_img.save(buffer, format="PNG")
                    b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
                    html += f"<img src='data:image/png;base64,{b64}' style='border-radius: 4px; border: 1px solid #ddd;'>"
                except Exception as e:
                    logger.error(f"Gallery error: {e}")
            html += "</div>"
            return html

        def handle_perform_clustering_flow():
            status = perform_initial_clustering()
            updates = [gr.update(visible=True), gr.update(visible=False), status]
            reps = GLOBAL_STATE["representative_images"]
            k = GLOBAL_STATE["current_k"]
            for i in range(20):
                if i < k:
                    updates.extend([gr.update(visible=True), create_cluster_gallery_html(reps, i), ""])
                else:
                    updates.extend([gr.update(visible=False), "", ""])
            return updates

        cluster_btn.click(handle_perform_clustering_flow, [],
                          [initial_annotation_group, cluster_btn, cluster_status] + [item for sublist in
                                                                                     initial_labels_inputs for item in
                                                                                     sublist])

        update_initial_btn.click(update_initial_labels,
                                 [label for _, _, label in initial_labels_inputs],
                                 [initial_labels_status, update_initial_btn, initial_clustering_done_msg, tab_boundary])

        def handle_extract_samples_flow():
            status = extract_boundary_samples()
            if not GLOBAL_STATE["boundary_samples"]:
                no_samples_msg = "<p class='feedback blue'>No new boundary samples found. You can proceed to the next step. / 新しい境界サンプルは見つかりませんでした。次のステップに進んでください。</p>"
                return [
                    gr.update(value=no_samples_msg),
                    gr.update(visible=False), gr.update(visible=False), gr.update(), gr.update(),
                ] + [gr.update(visible=False)] * 20 + [
                    gr.update(visible=False), gr.update(visible=False), gr.update(interactive=False),
                    gr.update(interactive=False),
                ]

            img, s_status, _, button_updates = get_current_boundary_sample()
            total_samples = len(GLOBAL_STATE.get("boundary_samples", []))

            return [
                status,
                gr.update(visible=True), gr.update(visible=True, value=img),
                s_status, gr.update(value=""),
            ] + button_updates + [
                gr.update(visible=False), gr.update(visible=True),
                gr.update(interactive=False), gr.update(interactive=(total_samples > 1))
            ]

        outputs_for_boundary_tab = [
                                       extract_status, boundary_annotation_group,
                                       sample_image, sample_status, new_label_text
                                   ] + cluster_buttons + [
                                       recluster_group, recluster_guidance, prev_btn, next_btn
                                   ]

        def move_sample(direction):
            if not GLOBAL_STATE.get("boundary_samples"):
                return [gr.update()] * len(outputs_for_boundary_tab)

            total_samples = len(GLOBAL_STATE['boundary_samples'])
            current_index = GLOBAL_STATE['boundary_index']

            if direction == "prev":
                new_index = max(0, current_index - 1)
            else:
                new_index = min(total_samples, current_index + 1)

            GLOBAL_STATE["boundary_index"] = new_index

            if new_index == total_samples:
                completion_message = """
                <div class='feedback green' style='text-align: center; padding: 20px;'>
                <h2>All samples reviewed.</h2>
                <p>To apply your changes, please press the 'Apply Constrained Clustering' button shown below.</p>
                <hr>
                <p>すべてのサンプルの確認が完了しました。変更を適用するには、下に表示されている「制約を適用して再クラスタリング」ボタンを押してください。</p>
                </div>
                """
                return [
                    gr.update(), gr.update(visible=True),
                    gr.update(visible=False, value=None), gr.update(value=completion_message),
                    gr.update(value=""),
                ] + [gr.update(visible=False)] * 20 + [
                    gr.update(visible=True), gr.update(visible=False),
                    gr.update(interactive=True), gr.update(interactive=False)
                ]
            else:
                img, status, new_label, button_updates = get_current_boundary_sample()
                return [
                    gr.update(), gr.update(visible=True),
                    gr.update(visible=True, value=img), status, new_label,
                ] + button_updates + [
                    gr.update(visible=False), gr.update(visible=True),
                    gr.update(interactive=(new_index > 0)), gr.update(interactive=True)
                ]

        extract_btn.click(handle_extract_samples_flow, [], outputs_for_boundary_tab)
        prev_btn.click(partial(move_sample, "prev"), [], outputs_for_boundary_tab)
        next_btn.click(partial(move_sample, "next"), [], outputs_for_boundary_tab)

        for i, btn in enumerate(cluster_buttons):
            btn.click(partial(assign_to_cluster, i), [],
                      [recluster_status, sample_image, sample_status] + cluster_buttons + [recluster_group,
                                                                                           recluster_guidance])

        create_new_btn.click(handle_create_new_cluster, [new_label_text],
                             [recluster_status, sample_image, sample_status, new_label_text] + cluster_buttons + [
                                 recluster_group, recluster_guidance])

        def handle_recluster_flow():
            status, _ = perform_constrained_clustering()
            if GLOBAL_STATE["boundary_samples_labeled"]:
                done_msg = "<p class='feedback blue'>再クラスタリングが完了しました。「境界サンプルを抽出」ボタンを再度押して、この手順を繰り返してください。<br>Re-clustering complete. Please press 'Extract Boundary Samples' again to repeat this procedure.</p>"
            else:
                done_msg = "<p class='feedback green'>ラベルの変更はありませんでした。次の「最終化」タブに進んでください。<br>No labels were changed. Proceed to the 'Finalize' tab.</p>"

            # UIを更新するために、境界サンプルの最初の状態を再取得
            img, sample_status, _, button_updates = get_current_boundary_sample()

            return status, gr.update(interactive=True), gr.update(value=done_msg,
                                                                  visible=True), img, sample_status, *button_updates

        recluster_btn.click(handle_recluster_flow, [], [recluster_status, tab_finalize, boundary_done_msg, sample_image,
                                                        sample_status] + cluster_buttons)

        def setup_finalize_tab_flow():
            k = GLOBAL_STATE["current_k"]
            reps = get_representative_images(GLOBAL_STATE["cluster_labels"], n_representatives=1)
            updates = []
            for i in range(20):  # UIコンポーネントの数に合わせる
                if i < k:
                    gallery_html = create_cluster_gallery_html(reps, i, n_images=1)
                    label = GLOBAL_STATE["label_dict"].get(str(i), "")
                    updates.extend([gr.update(visible=True), gallery_html, label])
                else:
                    updates.extend([gr.update(visible=False), "", ""])
            return updates

        tab_finalize.select(setup_finalize_tab_flow, [], [item for sublist in final_labels_inputs for item in sublist])

        update_final_btn.click(
            update_final_labels,
            [label for _, _, label in final_labels_inputs],
            [final_status, final_guidance_msg, tab_evaluation]
        )

        def get_sorted_cluster_ids(labels):
            unique_cids = np.unique(labels)
            if len(unique_cids) <= 1:
                return unique_cids.tolist()
            centroids = []
            for cid in unique_cids:
                centroid = np.mean(GLOBAL_STATE["embeddings"][labels == cid], axis=0)
                centroids.append(centroid)
            centroids = np.array(centroids)
            if len(centroids) < 2:
                return unique_cids.tolist()
            tsne = TSNE(n_components=1, random_state=GLOBAL_STATE["random_seed"], perplexity=min(5, len(centroids) - 1),
                        init='pca', learning_rate='auto')
            try:
                projected_centroids = tsne.fit_transform(centroids)
                sorted_indices = np.argsort(projected_centroids.flatten())
                return unique_cids[sorted_indices].tolist()
            except Exception as e:
                logger.warning(f"Could not perform t-SNE on centroids, returning unsorted. Error: {e}")
                return unique_cids.tolist()

        def create_cluster_gallery_html_for_eval(reps, sorted_cids):
            html = ""
            for cid in sorted_cids:
                html += "<div class='cluster-gallery-container'>"
                html += "<div style='display: flex; flex-wrap: wrap; gap: 10px; justify-content: center;'>"
                import base64
                from io import BytesIO
                # repsのキーは整数なので変換
                rep_list = reps.get(int(cid), [])
                for img_path, _, _ in rep_list[:1]:
                    try:
                        img = GLOBAL_STATE["original_images"].get(img_path) or Image.open(img_path).convert('RGB')
                        masked_img = create_masked_image(img)
                        masked_img.thumbnail((360, 360))
                        buffer = BytesIO()
                        masked_img.save(buffer, format="PNG")
                        b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
                        html += f"<img src='data:image/png;base64,{b64}' style='border-radius: 4px; border: 1px solid #eee;'>"
                    except Exception as e:
                        logger.error(f"Gallery error: {e}")
                html += "</div></div>"
            return html

        from sklearn.metrics import silhouette_score
        def setup_evaluation_tab():
            k = GLOBAL_STATE["current_k"]
            kmeans = KMeans(n_clusters=k, random_state=GLOBAL_STATE["random_seed"], n_init=10)
            GLOBAL_STATE["kmeans_matched_labels"] = kmeans.fit_predict(GLOBAL_STATE["embeddings"])

            if GLOBAL_STATE["embeddings"] is not None:
                if len(np.unique(GLOBAL_STATE["cluster_labels"])) > 1:
                    GLOBAL_STATE["hitl_silhouette_score"] = silhouette_score(GLOBAL_STATE["embeddings"],
                                                                             GLOBAL_STATE["cluster_labels"])
                else:
                    GLOBAL_STATE["hitl_silhouette_score"] = -1  # Cannot compute for 1 cluster

                if len(np.unique(GLOBAL_STATE["kmeans_matched_labels"])) > 1:
                    GLOBAL_STATE["kmeans_silhouette_score"] = silhouette_score(GLOBAL_STATE["embeddings"],
                                                                               GLOBAL_STATE["kmeans_matched_labels"])
                else:
                    GLOBAL_STATE["kmeans_silhouette_score"] = -1

            hitl_reps = get_representative_images(GLOBAL_STATE["cluster_labels"], n_representatives=1)
            kmeans_reps = get_representative_images(GLOBAL_STATE["kmeans_matched_labels"], n_representatives=1)
            GLOBAL_STATE["kmeans_matched_reps"] = kmeans_reps

            sorted_hitl_cids = get_sorted_cluster_ids(GLOBAL_STATE["cluster_labels"])
            sorted_kmeans_cids = get_sorted_cluster_ids(GLOBAL_STATE["kmeans_matched_labels"])

            hitl_gallery_html = create_cluster_gallery_html_for_eval(hitl_reps, sorted_hitl_cids)
            kmeans_gallery_html = create_cluster_gallery_html_for_eval(kmeans_reps, sorted_kmeans_cids)

            is_left_hitl = random.choice([True, False])
            GLOBAL_STATE["evaluation_layout_is_left_hitl"] = is_left_hitl

            if is_left_hitl:
                gallery_A_html, gallery_B_html = hitl_gallery_html, kmeans_gallery_html
            else:
                gallery_A_html, gallery_B_html = kmeans_gallery_html, hitl_gallery_html

            return gallery_A_html, gallery_B_html

        tab_evaluation.select(
            setup_evaluation_tab,
            [],
            [gallery_A, gallery_B]
        )

        def handle_evaluation_choice(choice):
            is_left_hitl = GLOBAL_STATE["evaluation_layout_is_left_hitl"]
            if (choice == 'A' and is_left_hitl) or (choice == 'B' and not is_left_hitl):
                final_choice = "HitL"
            else:
                final_choice = "KMeans_matched"

            GLOBAL_STATE["evaluation_choice"] = final_choice
            status_msg = f"<p class='feedback green'>You chose Result {choice} ({final_choice}). Thank you! Please proceed to export your results. / 結果{choice} ({final_choice}) を選択しました。ありがとうございます。結果をエクスポートしてください。</p>"
            return gr.update(value=status_msg, visible=True), gr.update(visible=True), gr.update(
                interactive=False), gr.update(interactive=False)

        choose_A_btn.click(
            partial(handle_evaluation_choice, 'A'),
            [],
            [evaluation_status, export_group, choose_A_btn, choose_B_btn]
        )
        choose_B_btn.click(
            partial(handle_evaluation_choice, 'B'),
            [],
            [evaluation_status, export_group, choose_A_btn, choose_B_btn]
        )

        export_btn.click(export_results, [participant_id_input], [download_file, export_status])

    return app


def get_parameters():
    parser = ArgumentParser()
    parser.add_argument("--dataset_dir", type=str, default="./data/lapwing",
                        help="Path to the root directory of the input dataset.")
    parser.add_argument("--image_dir", type=str, default="images",
                        help="Subdirectory for images within the dataset directory.")
    parser.add_argument("--text_dir", type=str, default="texts",
                        help="Subdirectory for text files (like bbox) within the dataset directory.")
    parser.add_argument("--bbox_json", type=str, default="bounding_boxes.json",
                        help="Name of the bounding box JSON file.")
    parser.add_argument("--output_dir", type=str, default="./experiments",
                        help="Directory to save the experiment results.")
    parser.add_argument("--labels_name", type=str, default="HitL_results", help="Prefix for the output JSON filename.")
    parser.add_argument("--results_dir", type=str, default="./results",
                        help="Directory to save the final human evaluation and annotation results.")

    try:
        args = parser.parse_args()
    except SystemExit:
        args = parser.parse_args([])

    args.image_dir = os.path.join(args.dataset_dir,  args.image_dir)
    args.text_dir = os.path.join(args.dataset_dir,  args.text_dir)
    args.bbox_json = os.path.join(args.text_dir, args.bbox_json)

    return args


def main():
    set_random_seed()
    app = create_gradio_interface()
    app.launch(share=True)


def set_random_seed():
    seed = GLOBAL_STATE["random_seed"]
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    logger.info(f"Set random seed to {seed}")


if __name__ == "__main__":
    main()