model_display.py

"""UI-only model display names.

Raw CSV / sample IDs stay unchanged; this maps aliases to a shared label
shown in leaderboards, Pareto plots, dropdowns, and sample headers.
"""

from __future__ import annotations

import re

# Explicit aliases (any leaderboard / sample id) -> shared display label.
MODEL_DISPLAY_NAMES = {
    # FLUX
    "flux_2_pro": "FLUX.2 [pro]",
    "flux_2_max": "FLUX.2 [max]",
    "flux_2_flex": "FLUX.2 [flex]",
    "flux_2_dev": "FLUX.2 [dev]",
    "flux_1_1_pro": "FLUX1.1 [pro]",
    "flux_1_1_pro_ultra": "FLUX1.1 [pro] Ultra",
    "flux_dev": "FLUX.1 [dev]",
    "flux_schnell": "FLUX.1 [schnell]",
    "flux_krea": "FLUX.1 Krea [dev]",
    # GPT Image
    "gpt_image_2": "GPT Image 2",
    "GPT Image 2 (high)": "GPT Image 2",
    "gpt_image_1_5": "GPT Image 1.5",
    "GPT Image 1.5 (high)": "GPT Image 1.5",
    "gpt_image_1": "GPT Image 1",
    "GPT Image 1 (high)": "GPT Image 1",
    "GPT Image 1 Mini (medium)": "GPT Image 1 Mini",
    # Nano Banana / Gemini
    "nano_banana_2_0": "Nano Banana 2",
    "nano_banana_2": "Nano Banana 2",
    "Nano Banana 2 (Gemini 3.1 Flash Image Preview)": "Nano Banana 2",
    "Nano Banana 2 Lite (Gemini 3.1 Flash Lite Image)": "Nano Banana 2 Lite",
    "nano_banana_pro": "Nano Banana Pro",
    "Nano Banana Pro (Gemini 3 Pro Image)": "Nano Banana Pro",
    "Nano Banana (Gemini 2.5 Flash Image)": "Nano Banana",
    # Seedream
    "seedream_5_0": "Seedream 5.0",
    "seedream_4_5": "Seedream 4.5",
    "seedream_4_0": "Seedream 4.0",
    "seedream_3": "Seedream 3.0",
    # Qwen
    "qwen_image": "Qwen Image",
    "qwen_image_2_0_pro": "Qwen Image 2.0 Pro",
    "Qwen Image 2.0 Pro (2026-04-22)": "Qwen Image 2.0 Pro",
    "Qwen Image 2.0 (2026-03-03)": "Qwen Image 2.0",
    "qwen_image_2512": "Qwen Image 2512",
    "Qwen Image Max 2512": "Qwen Image 2512",
    "qwen_image_fast": "Qwen Image Fast",
    # Ideogram
    "ideogram_4_0_quality": "Ideogram 4.0 Quality",
    "Ideogram 4.0 (Quality)": "Ideogram 4.0 Quality",
    "Ideogram 4.0 Fast (Quality)": "Ideogram 4.0 Fast Quality",
    # Imagen
    "imagen_4_0": "Imagen 4",
    "imagen_4": "Imagen 4",
    "Imagen 4 Standard": "Imagen 4",
    "imagen_4_0_ultra": "Imagen 4 Ultra",
    "imagen_4_ultra": "Imagen 4 Ultra",
    "imagen_4_fast": "Imagen 4 Fast",
    "Imagen 3 (v002)": "Imagen 3",
    # HiDream
    "hidream_i1_dev": "HiDream-I1 Dev",
    "HiDream-I1-Dev": "HiDream-I1 Dev",
    "hidream_i1_fast": "HiDream-I1 Fast",
    "HiDream-I1-Fast": "HiDream-I1 Fast",
    "hidream_i1_full": "HiDream-I1 Full",
    "HiDream-O1-Image": "HiDream-O1",
    "HiDream-O1-Image-1.5": "HiDream-O1 1.5",
    "HiDream-O1-Image-Dev": "HiDream-O1 Dev",
    # Reve
    "reve_2_1": "Reve 2.1",
    "Reve Image (Halfmoon)": "Reve Image",
    # P-Image
    "p_image": "P-Image",
    "p_image_2_ideogram_very_low_1k": "P-Image-Ideogram Very Low 1K",
    "p_image_2_ideogram_very_low_2k": "P-Image-Ideogram Very Low 2K",
    "P-Image-Ideogram (Very Low)": "P-Image-Ideogram Very Low",
    "p_image_2_ideogram_low_1k": "P-Image-Ideogram Low 1K",
    "p_image_2_ideogram_low_2k": "P-Image-Ideogram Low 2K",
    "P-Image-Ideogram (Low)": "P-Image-Ideogram Low",
    "p_image_2_ideogram_medium_1k": "P-Image-Ideogram Medium 1K",
    "p_image_2_ideogram_medium_2k": "P-Image-Ideogram Medium 2K",
    "P-Image-Ideogram (Medium)": "P-Image-Ideogram Medium",
    "p_image_2_ideogram_high_1k": "P-Image-Ideogram High 1K",
    "p_image_2_ideogram_high_2k": "P-Image-Ideogram High 2K",
    "P-Image-Ideogram (High)": "P-Image-Ideogram High",
    "p_image_2_ideogram_very_high_1k": "P-Image-Ideogram Very High 1K",
    # P-Video-Edit
    "P-Video-Edit": "P-Video-Edit",
    "P-Video-Edit Draft": "P-Video-Edit Draft",
    "P-Video Edit Final": "P-Video-Edit",
    "P-Video Edit Final (draft)": "P-Video-Edit Draft",
    "p_video_edit_preview__replicate_final": "P-Video-Edit",
    "p_video_edit_preview__replicate_final__draft": "P-Video-Edit Draft",
    # Text-to-video leaderboard
    "gemini_omni_1_1_flash": "Gemini Omni 1.1 Flash",
    "grok_imagine_video": "Grok Imagine Video",
    "grok_imagine_video_v1_5": "Grok Imagine Video 1.5",
    "minimax_h3": "MiniMax H3",
    "minimax_h3_max": "MiniMax H3 Max",
    "minimax_h3_max_turbo": "MiniMax H3 Max Turbo",
    "minimax_h3_max_turbo__prompt_expansion_mode_balanced": "MiniMax H3 Max Turbo",
    "seedance_2_5_turbo": "Seedance 2.5 Turbo",
    "veo_3_1_lite": "Veo 3.1 Lite",
    # Video-to-video leaderboard
    "gemini_omni_flash_edit__fal": "Gemini Omni Flash Edit",
    "grok_imagine_video__replicate": "Grok Imagine Video",
    "happyhorse_1_0__wavespeed": "HappyHorse 1.0",
    "lucy_edit_pro__fal": "Lucy Edit Pro",
    "minimax_h3_reference_to_video__fal": "MiniMax H3 Reference-to-Video",
    "seedance_2_5_video_edit_turbo__wavespeed": "Seedance 2.5 Video Edit Turbo",
    "wan_2_7_video_edit__wavespeed": "Wan 2.7 Video Edit",
    # Others overlapping P-Bench
    "z_image": "Z-Image",
    "glm_image": "GLM-Image",
    "hunyuanimage_3_0": "HunyuanImage 3.0",
    "hunyuan_image_3": "HunyuanImage 3.0",
    "HunyuanImage 3.0 (Fal)": "HunyuanImage 3.0",
    "HunyuanImage 3.0 Instruct (Fal)": "HunyuanImage 3.0 Instruct",
    "wan_2_2_image": "Wan 2.2 Image",
    "Wan2.6 Text to Image": "Wan 2.6 Text to Image",
    "kling_v2_1": "Kling v2.1",
    "juggernaut_base_flux": "Juggernaut Base FLUX",
    "juggernaut_pro_flux": "Juggernaut Pro FLUX",
    "juggernaut_lightning_flux": "Juggernaut Lightning FLUX",
    "bria_4_fibo": "Bria FIBO",
    "sdxl": "SDXL 1.0",
    "Stable Diffusion XL 1.0": "SDXL 1.0",
    "sdxl_lightning": "SDXL Lightning",
    # Arena AI (kebab / arena.ai ids)
    "gpt-image-2 (medium)": "GPT Image 2",
    "gpt-image-1.5-high-fidelity": "GPT Image 1.5",
    "gpt-image-1": "GPT Image 1",
    "gpt-image-1-mini": "GPT Image 1 Mini",
    "mai-image-2.6-preview": "MAI-Image-2.6 Preview",
    "mai-image-2.5": "MAI-Image-2.5",
    "mai-image-2": "MAI-Image-2",
    "mai-image-1": "MAI Image 1",
    "grok-imagine-image-2.0 (low)": "Grok Imagine Image 2.0",
    "grok-imagine-image": "Grok Imagine Image",
    "grok-imagine-image-pro": "Grok Imagine Image Pro",
    "reve-2.1": "Reve 2.1",
    "reve-2.0": "Reve 2.0",
    "reve-v1.5": "Reve 1.5",
    "muse-image": "Muse Image",
    "gemini-3.1-flash-image (nano-banana-2) [web-search]": "Nano Banana 2",
    "gemini-3.1-flash-lite-image (nano-banana-2-lite)": "Nano Banana 2 Lite",
    "gemini-3-pro-image-2k (nano-banana-pro)": "Nano Banana Pro 2K",
    "gemini-3-pro-image-preview (nano-banana-pro)": "Nano Banana Pro",
    "gemini-2.5-flash-image-preview (nano-banana)": "Nano Banana",
    "gemini-2.0-flash-preview-image-generation": "Gemini 2.0 Flash Image",
    "seedream-5.0-pro": "Seedream 5.0 Pro",
    "seedream-5.0-lite": "Seedream 5.0 Lite",
    "seedream-4.5": "Seedream 4.5",
    "seedream-4-2k": "Seedream 4.0 2K",
    "seedream-4-fal": "Seedream 4.0",
    "seedream-4-high-res-fal": "Seedream 4.0 High Res",
    "seedream-3": "Seedream 3.0",
    "qwen-image-3.0-pro": "Qwen Image 3.0 Pro",
    "qwen-image-2.0-pro-2026-06-22": "Qwen Image 2.0 Pro",
    "qwen-image-2512": "Qwen Image 2512",
    "qwen-image-prompt-extend": "Qwen Image Prompt Extend",
    "qwen-image": "Qwen Image",
    "uni-1.1-max": "Luma UNI 1 Max",
    "uni-1.1": "Luma UNI 1",
    "Cosmos3-Super-Text2Image (Agentic)": "Cosmos3 Super Text2Image (Agentic)",
    "Cosmos3-Super-Text2Image": "Cosmos3 Super Text2Image",
    "recraft-v4.1-utility-pro": "Recraft V4.1 Utility Pro",
    "recraft-v4.1-pro": "Recraft V4.1 Pro",
    "recraft-v4": "Recraft V4",
    "recraft-v3": "Recraft V3",
    "flux-2-max": "FLUX.2 [max]",
    "flux-2-flex": "FLUX.2 [flex]",
    "flux-2-pro": "FLUX.2 [pro]",
    "flux-2-dev": "FLUX.2 [dev]",
    "flux-2-klein-9b": "FLUX.2 [klein] 9B",
    "flux-2-klein-4b": "FLUX.2 [klein] 4B",
    "flux-1-kontext-max": "FLUX.1 Kontext [max]",
    "flux-1-kontext-pro": "FLUX.1 Kontext [pro]",
    "flux-1-kontext-dev": "FLUX.1 Kontext [dev]",
    "flux-1.1-pro": "FLUX1.1 [pro]",
    "flux-1-dev-fp8": "FLUX.1 [dev]",
    "hunyuan-image-3.0": "HunyuanImage 3.0",
    "imagen-ultra-4.0-generate-001": "Imagen 4 Ultra",
    "imagen-4.0-generate-001": "Imagen 4",
    "imagen-3.0-generate-002": "Imagen 3",
    "wan2.6-t2i": "Wan 2.6 Text to Image",
    "wan2.5-t2i-preview": "Wan 2.5 Preview",
    "wan2.7-image-pro": "Wan 2.7 Pro",
    "wan2.7-image": "Wan 2.7",
    "ideogram-4.0-quality": "Ideogram 4.0 Quality",
    "ideogram-v3-quality": "Ideogram 3.0",
    "ideogram-v2": "Ideogram v2",
    "krea-2-medium": "Krea 2 Medium",
    "krea-2-turbo": "Krea 2 Medium Turbo",
    "krea-2-large": "Krea 2 Large",
    "hidream-o1-image": "HiDream-O1",
    "z-image-turbo": "Z-Image Turbo",
    "photon": "Luma Photon",
    "p-image": "P-Image",
    "runway-gen4": "Runway Gen-4 Image",
    "lucid-origin": "Lucid Origin",
    "glm-image": "GLM-Image",
    "dall-e-3": "DALLE 3",
    "stable-diffusion-v35-large": "Stable Diffusion 3.5 Large",
    "bagel": "Bagel",
}


def _p_video_2_display_name(model_id: str):
    """Turn p_video_2 variant ids into P-Video-2 Draft 720p labels."""
    raw = str(model_id).strip()
    if raw != "p_video_2" and not raw.startswith("p_video_2__"):
        return None
    if raw == "p_video_2":
        return "P-Video-2"

    draft = False
    upsample = None
    resolution = None
    for part in raw.split("__")[1:]:
        if part.startswith("draft_"):
            draft = part.endswith("true")
        elif part.startswith("prompt_upsampling_"):
            upsample = part.endswith("true")
        elif part.startswith("resolution_"):
            resolution = part[len("resolution_") :]

    label = "P-Video-2"
    if draft:
        label += " Draft"
    if resolution:
        label += f" {resolution}"
    if upsample is False:
        label += " (no prompt upsampling)"
    return label


def _prettify_snake_case(model_id: str) -> str:
    parts = [part for part in str(model_id).split("_") if part]
    pretty = []
    for part in parts:
        if part.isdigit():
            pretty.append(part)
        elif re.fullmatch(r"\d+k", part, flags=re.IGNORECASE):
            pretty.append(part.upper())
        elif part.lower() in {"flux", "gpt", "sdxl", "glm"}:
            pretty.append(part.upper())
        else:
            pretty.append(part.capitalize())
    return " ".join(pretty)


def display_model_name(model_id) -> str:
    """Return the shared UI label for a raw model id / source name."""
    if model_id is None:
        return ""
    raw = str(model_id).strip()
    if not raw:
        return ""
    if raw in MODEL_DISPLAY_NAMES:
        return MODEL_DISPLAY_NAMES[raw]
    p_video_2 = _p_video_2_display_name(raw)
    if p_video_2:
        return p_video_2
    # Already a human label (spaces / punctuation) — keep as-is.
    if re.search(r"[\s.\[\]()-]", raw):
        return raw
    return _prettify_snake_case(raw)