import os

# ZeroGPU and library caches must be configured before importing spaces/torch.
EXAMPLE_CACHE_VERSION = "2026-08-26-a"
os.environ.setdefault("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("GRADIO_EXAMPLES_CACHE", f"/tmp/gradio_cached_examples/{EXAMPLE_CACHE_VERSION}")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("GRADIO_SSR_MODE", "false")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces

import html
import logging
import math
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import gradio as gr
import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer


MODEL_ID = "tencent/WeMM-Embedding-9B"
ROOT = Path(__file__).resolve().parent
ASSET_DIR = ROOT / "assets"
MATRYOSHKA_DIMS = (64, 128, 256, 512, 1024, 2048, 4096)
MAX_CUSTOM_TEXTS = 6
MAX_CUSTOM_MEDIA = 6

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
LOGGER = logging.getLogger("wemm-space")


@dataclass(frozen=True)
class Candidate:
    key: str
    title: str
    kind: str
    description: str
    payload: Any
    media_path: str | None = None


SHOWCASE: tuple[Candidate, ...] = (
    Candidate(
        "llama4",
        "Llama 4 model card",
        "visual document",
        "A dense model-card screenshot describing the Scout and Maverick variants.",
        str(ASSET_DIR / "llama4_hgf.png"),
        str(ASSET_DIR / "llama4_hgf.png"),
    ),
    Candidate(
        "qwen-omni",
        "Qwen2.5-Omni overview",
        "visual document",
        "A model page covering omni-modal perception, speech, and video capabilities.",
        str(ASSET_DIR / "qwen2.5omni_hgf.png"),
        str(ASSET_DIR / "qwen2.5omni_hgf.png"),
    ),
    Candidate(
        "likelihood-contour",
        "Scientific contour plot",
        "figure",
        "An orange likelihood contour plotted against gamma and log-scaled tau over mass.",
        str(ASSET_DIR / "doc1.jpg"),
        str(ASSET_DIR / "doc1.jpg"),
    ),
    Candidate(
        "budget-1971",
        "1971 budget infographic",
        "visual document",
        "A historical chart comparing US outlays, including natural resources spending.",
        str(ASSET_DIR / "doc2.jpg"),
        str(ASSET_DIR / "doc2.jpg"),
    ),
    Candidate(
        "scoring-rules",
        "Proper scoring rules paper",
        "visual document",
        "An academic page with lemmas, an algorithm, equations, and references.",
        str(ASSET_DIR / "doc3.jpg"),
        str(ASSET_DIR / "doc3.jpg"),
    ),
    Candidate(
        "road-safety",
        "Road-safety assessment",
        "visual document",
        "An environmental assessment page about driver training, signs, and road closures.",
        str(ASSET_DIR / "doc4.jpg"),
        str(ASSET_DIR / "doc4.jpg"),
    ),
    Candidate(
        "mapo-tofu",
        "Mapo tofu in motion",
        "video",
        "A short cooking clip showing the preparation of the Sichuan tofu dish.",
        str(ASSET_DIR / "mapo_tofu.mp4"),
        str(ASSET_DIR / "mapo_tofu.mp4"),
    ),
    Candidate(
        "zhajiang-noodles",
        "Zhajiang noodles in motion",
        "video",
        "A short cooking clip showing noodles with a savory fermented-bean sauce.",
        str(ASSET_DIR / "zhajiang_noodle.mp4"),
        str(ASSET_DIR / "zhajiang_noodle.mp4"),
    ),
    Candidate(
        "vector-search",
        "How vector search works",
        "text",
        "Dense retrieval maps queries and documents into one normalized vector space, then ranks candidates by cosine similarity.",
        "Dense retrieval maps queries and documents into one normalized vector space, then ranks candidates by cosine similarity.",
    ),
    Candidate(
        "night-train",
        "A quiet journey",
        "text",
        "夜行列车穿过雨中的城市，车窗映出霓虹灯和安静的乘客。",
        "夜行列车穿过雨中的城市，车窗映出霓虹灯和安静的乘客。",
    ),
)


LOGGER.info("Loading %s on CPU", MODEL_ID)
MODEL = SentenceTransformer(
    MODEL_ID,
    trust_remote_code=True,
    device="cpu",
    model_kwargs={"dtype": torch.bfloat16, "low_cpu_mem_usage": True},
)
MODEL.eval()
torch.set_grad_enabled(False)
LOGGER.info("Model loaded; waiting for a ZeroGPU allocation")

_CACHE_LOCK = threading.Lock()
_INFERENCE_LOCK = threading.Lock()
_SHOWCASE_EMBEDDINGS: torch.Tensor | None = None


def _file_path(value: Any) -> str | None:
    """Normalize Gradio file values across UI and API representations."""
    if value is None:
        return None
    if isinstance(value, (str, Path)):
        return str(value)
    if isinstance(value, dict):
        path = value.get("path") or value.get("name")
        return str(path) if path else None
    path = getattr(value, "path", None) or getattr(value, "name", None)
    return str(path) if path else None


def _is_video(path: str) -> bool:
    return Path(path.split("?", 1)[0]).suffix.lower() in {
        ".mp4",
        ".webm",
        ".mov",
        ".mkv",
        ".avi",
        ".mpeg",
        ".mpg",
    }


def _multimodal_payload(text: str | None, image: Any, video: Any) -> tuple[Any, str]:
    text = (text or "").strip()
    image_path = _file_path(image)
    video_path = _file_path(video)
    if image_path and video_path:
        raise gr.Error("Choose one visual query: an image or a video, not both.")
    if image_path:
        if text:
            return {"image": image_path, "text": text}, "image + text"
        return image_path, "image"
    if video_path:
        if text:
            return {"video": video_path, "text": text}, "video + text"
        return video_path, "video"
    if text:
        return text, "text"
    raise gr.Error("Add a text, image, or video query to begin.")


def _parse_text_candidates(raw: str | None) -> list[Candidate]:
    candidates: list[Candidate] = []
    lines = [line.strip() for line in (raw or "").splitlines() if line.strip()]
    if len(lines) > MAX_CUSTOM_TEXTS:
        raise gr.Error(f"Use at most {MAX_CUSTOM_TEXTS} custom text candidates.")
    for index, line in enumerate(lines, start=1):
        if "::" in line:
            title, body = (part.strip() for part in line.split("::", 1))
            title = title or f"Custom text {index}"
            body = body or title
        else:
            title = f"Custom text {index}"
            body = line
        candidates.append(
            Candidate(
                f"custom-text-{index}",
                title[:80],
                "text",
                body[:240],
                body,
            )
        )
    return candidates


def _parse_media_candidates(raw: Any) -> list[Candidate]:
    candidates: list[Candidate] = []
    items = raw or []
    if len(items) > MAX_CUSTOM_MEDIA:
        raise gr.Error(f"Upload at most {MAX_CUSTOM_MEDIA} candidate media files.")
    for index, item in enumerate(items, start=1):
        media = item[0] if isinstance(item, (tuple, list)) else item
        caption = item[1] if isinstance(item, (tuple, list)) and len(item) > 1 else None
        path = _file_path(media)
        if not path:
            continue
        kind = "video" if _is_video(path) else "image"
        title = (caption or f"Uploaded {kind} {index}").strip()
        candidates.append(
            Candidate(
                f"custom-media-{index}",
                title[:80],
                kind,
                f"User-supplied {kind} candidate.",
                path,
                path,
            )
        )
    return candidates


def _encode(items: list[Any], *, query: bool) -> torch.Tensor:
    method = MODEL.encode_query if query else MODEL.encode_document
    with torch.inference_mode():
        embeddings = method(
            items,
            batch_size=1,
            convert_to_tensor=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
    if embeddings.ndim == 1:
        embeddings = embeddings.unsqueeze(0)
    return embeddings.float().cpu()


def _truncate_normalize(embeddings: torch.Tensor, dimension: int) -> torch.Tensor:
    return F.normalize(embeddings[..., :dimension], p=2, dim=-1)


def _dimension_scores(query: torch.Tensor, documents: torch.Tensor) -> dict[int, list[float]]:
    scores: dict[int, list[float]] = {}
    for dimension in MATRYOSHKA_DIMS:
        query_d = _truncate_normalize(query, dimension)
        docs_d = _truncate_normalize(documents, dimension)
        scores[dimension] = (query_d @ docs_d.T).squeeze(0).tolist()
    return scores


def _search_duration(*args: Any, **kwargs: Any) -> int:
    """Budget more time for the cold corpus pass and video queries."""
    query_payload = args[0] if args else None
    custom_candidates = args[2] if len(args) > 2 else []
    include_showcase = bool(args[3]) if len(args) > 3 else True
    query_has_video = (
        isinstance(query_payload, dict) and bool(query_payload.get("video"))
    ) or (isinstance(query_payload, str) and _is_video(query_payload))
    has_uploaded_video = any(item.kind == "video" for item in custom_candidates or [])
    if include_showcase and _SHOWCASE_EMBEDDINGS is None:
        return 240
    return 120 if query_has_video or has_uploaded_video else 90


@spaces.GPU(duration=_search_duration)
def _run_search(
    query_payload: Any,
    query_kind: str,
    custom_candidates: list[Candidate],
    include_showcase: bool,
    dimension: int,
    progress: gr.Progress,
) -> tuple[torch.Tensor, torch.Tensor, list[Candidate], float, bool]:
    global _SHOWCASE_EMBEDDINGS

    started = time.perf_counter()
    was_cold = include_showcase and _SHOWCASE_EMBEDDINGS is None
    progress(0.04, desc="Allocating the 9B model on GPU")
    MODEL.to("cuda")
    try:
        progress(0.16, desc=f"Encoding the {query_kind} query")
        query_embedding = _encode([query_payload], query=True)

        document_blocks: list[torch.Tensor] = []
        candidates: list[Candidate] = []
        if include_showcase:
            with _CACHE_LOCK:
                cached = _SHOWCASE_EMBEDDINGS
            if cached is None:
                progress(0.30, desc="Mapping the curated multimodal universe")
                encoded = _encode([item.payload for item in SHOWCASE], query=False)
                with _CACHE_LOCK:
                    if _SHOWCASE_EMBEDDINGS is None:
                        _SHOWCASE_EMBEDDINGS = encoded
                    cached = _SHOWCASE_EMBEDDINGS
            document_blocks.append(cached)
            candidates.extend(SHOWCASE)

        if custom_candidates:
            progress(0.72, desc="Encoding your candidate collection")
            custom_embeddings = _encode([item.payload for item in custom_candidates], query=False)
            document_blocks.append(custom_embeddings)
            candidates.extend(custom_candidates)

        if not document_blocks:
            raise gr.Error("Include the showcase universe or add at least one candidate.")

        documents = torch.cat(document_blocks, dim=0)
        progress(0.92, desc=f"Ranking in {dimension:,} dimensions")
        return query_embedding, documents, candidates, time.perf_counter() - started, was_cold
    finally:
        MODEL.to("cpu")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def _score_tone(score: float) -> str:
    if score >= 0.70:
        return "high"
    if score >= 0.40:
        return "mid"
    return "low"


def _render_summary(
    query_kind: str,
    dimension: int,
    candidate_count: int,
    elapsed: float,
    was_cold: bool,
    top_title: str,
    top_score: float,
) -> str:
    cache_note = "cold corpus map" if was_cold else "warm corpus cache"
    return f"""
    <section class="run-summary">
      <div class="run-kicker"><span class="live-dot"></span> semantic field resolved</div>
      <div class="run-main">
        <div><span class="run-label">Top match</span><strong>{html.escape(top_title)}</strong></div>
        <div class="hero-score"><span>{top_score:+.3f}</span><small>cosine</small></div>
      </div>
      <div class="run-meta">
        <span>{html.escape(query_kind)}</span><i></i>
        <span>{dimension:,}D</span><i></i>
        <span>{candidate_count} candidates</span><i></i>
        <span>{elapsed:.1f}s</span><i></i>
        <span>{cache_note}</span>
      </div>
    </section>
    """


def _render_rankings(ranked: list[tuple[Candidate, float]]) -> str:
    rows: list[str] = []
    for rank, (candidate, score) in enumerate(ranked, start=1):
        fill = min(100.0, max(2.0, max(0.0, score) * 100.0))
        rows.append(
            f"""
            <article class="rank-row {'winner' if rank == 1 else ''}">
              <div class="rank-number">{rank:02d}</div>
              <div class="rank-copy">
                <div class="rank-title-line">
                  <strong>{html.escape(candidate.title)}</strong>
                  <span class="kind-pill">{html.escape(candidate.kind)}</span>
                </div>
                <p>{html.escape(candidate.description)}</p>
                <div class="score-track"><span style="width:{fill:.1f}%"></span></div>
              </div>
              <div class="rank-score { _score_tone(score) }">{score:+.3f}<small>cos</small></div>
            </article>
            """
        )
    return '<div class="ranking-stack">' + "".join(rows) + "</div>"


def _render_curve(series: dict[str, list[float]], title: str) -> str:
    width, height = 820, 300
    left, right, top, bottom = 58, 20, 28, 48
    plot_w, plot_h = width - left - right, height - top - bottom
    all_values = [value for values in series.values() for value in values]
    low = max(-1.0, min(all_values) - 0.08)
    high = min(1.0, max(all_values) + 0.08)
    if high - low < 0.2:
        midpoint = (high + low) / 2
        low, high = max(-1.0, midpoint - 0.1), min(1.0, midpoint + 0.1)

    def x_at(index: int) -> float:
        return left + index * plot_w / (len(MATRYOSHKA_DIMS) - 1)

    def y_at(value: float) -> float:
        return top + (high - value) * plot_h / max(1e-8, high - low)

    grid: list[str] = []
    for tick in range(5):
        value = high - tick * (high - low) / 4
        y = y_at(value)
        grid.append(
            f'<line x1="{left}" y1="{y:.1f}" x2="{width-right}" y2="{y:.1f}" class="chart-grid" />'
            f'<text x="{left-10}" y="{y+4:.1f}" text-anchor="end" class="chart-axis">{value:+.2f}</text>'
        )
    for index, dimension in enumerate(MATRYOSHKA_DIMS):
        x = x_at(index)
        grid.append(f'<text x="{x:.1f}" y="{height-18}" text-anchor="middle" class="chart-axis">{dimension}</text>')

    colors = ("#ffb86b", "#78e8df", "#a994ff", "#ff7aa2", "#9ad45b")
    paths: list[str] = []
    legend: list[str] = []
    for series_index, (name, values) in enumerate(series.items()):
        color = colors[series_index % len(colors)]
        points = " ".join(f"{x_at(i):.1f},{y_at(value):.1f}" for i, value in enumerate(values))
        circles = "".join(
            f'<circle cx="{x_at(i):.1f}" cy="{y_at(value):.1f}" r="3.5" fill="{color}" />'
            for i, value in enumerate(values)
        )
        paths.append(f'<polyline points="{points}" fill="none" stroke="{color}" stroke-width="3" />{circles}')
        legend.append(
            f'<span><b style="background:{color}"></b>{html.escape(name[:38])}</span>'
        )

    return f"""
    <section class="viz-card">
      <div class="viz-heading"><div><small>MATRYOSHKA SCOPE</small><h3>{html.escape(title)}</h3></div><span>64 → 4096 dimensions</span></div>
      <svg class="dimension-chart" viewBox="0 0 {width} {height}" role="img" aria-label="Similarity by embedding dimension">
        {''.join(grid)}{''.join(paths)}
      </svg>
      <div class="chart-legend">{''.join(legend)}</div>
      <p class="viz-note">Cosine similarity at every native truncation size. Compare trends, not universal thresholds.</p>
    </section>
    """


def _fingerprint_svg(vector: torch.Tensor, label: str) -> str:
    values = vector.detach().float().flatten()
    bar_count = min(96, values.numel())
    chunks = torch.tensor_split(values, bar_count)
    samples = [float(chunk.mean()) for chunk in chunks]
    scale = max(max(abs(value) for value in samples), 1e-6)
    width, height = 800, 210
    center = 104
    bar_w = (width - 24) / bar_count
    bars: list[str] = []
    for index, value in enumerate(samples):
        magnitude = min(84.0, abs(value) / scale * 84.0)
        x = 12 + index * bar_w
        y = center - magnitude if value >= 0 else center
        color = "#70e4da" if value >= 0 else "#ff9d57"
        bars.append(
            f'<rect x="{x:.1f}" y="{y:.1f}" width="{max(1.2, bar_w-1.5):.1f}" height="{magnitude:.1f}" rx="1.5" fill="{color}" opacity=".9" />'
        )
    return f"""
    <section class="viz-card fingerprint-card">
      <div class="viz-heading"><div><small>VECTOR FINGERPRINT</small><h3>{html.escape(label)}</h3></div><span>{values.numel():,} values</span></div>
      <svg class="fingerprint" viewBox="0 0 {width} {height}" role="img" aria-label="Compressed signed embedding fingerprint">
        <line x1="12" y1="{center}" x2="{width-12}" y2="{center}" class="zero-line" />
        {''.join(bars)}
      </svg>
      <div class="fingerprint-key"><span><b class="positive"></b>positive</span><span><b class="negative"></b>negative</span><em>96 pooled slices · shape, not magnitude</em></div>
    </section>
    """


def search_experience(
    query_text: str,
    query_image: Any,
    query_video: Any,
    custom_texts: str,
    candidate_media: Any,
    include_showcase: bool,
    dimension: int,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str, str, list[tuple[str, str]], str, str, dict[str, Any]]:
    """Search a mixed text/image/video collection with a multimodal query."""
    dimension = int(dimension)
    if dimension not in MATRYOSHKA_DIMS:
        raise gr.Error("Choose one of the model's native Matryoshka dimensions.")
    query_payload, query_kind = _multimodal_payload(query_text, query_image, query_video)
    custom_candidates = _parse_text_candidates(custom_texts) + _parse_media_candidates(candidate_media)

    with _INFERENCE_LOCK:
        query, documents, candidates, elapsed, was_cold = _run_search(
            query_payload,
            query_kind,
            custom_candidates,
            bool(include_showcase),
            dimension,
            progress,
        )
    dimension_map = _dimension_scores(query, documents)
    chosen_scores = dimension_map[dimension]
    order = sorted(range(len(candidates)), key=lambda index: chosen_scores[index], reverse=True)
    ranked = [(candidates[index], float(chosen_scores[index])) for index in order]

    top_candidate, top_score = ranked[0]
    summary = _render_summary(
        query_kind,
        dimension,
        len(candidates),
        elapsed,
        was_cold,
        top_candidate.title,
        top_score,
    )
    rankings = _render_rankings(ranked[:8])
    gallery = [
        (candidate.media_path, f"#{rank} · {candidate.title} · cosine {score:+.3f}")
        for rank, (candidate, score) in enumerate(ranked, start=1)
        if candidate.media_path
    ][:8]

    curve_series: dict[str, list[float]] = {}
    for index in order[:4]:
        curve_series[candidates[index].title] = [dimension_map[dim][index] for dim in MATRYOSHKA_DIMS]
    curve = _render_curve(curve_series, "Does the ranking survive compression?")
    fingerprint = _fingerprint_svg(query[0, :dimension], f"Query · {query_kind} · {dimension:,}D")

    diagnostics = {
        "model": MODEL_ID,
        "query_modality": query_kind,
        "selected_dimension": dimension,
        "full_embedding_dimension": int(query.shape[-1]),
        "l2_norm_after_truncation": round(float(_truncate_normalize(query, dimension).norm()), 6),
        "candidate_count": len(candidates),
        "gpu_pass_seconds": round(elapsed, 3),
        "showcase_cache": "created" if was_cold else "reused" if include_showcase else "not_requested",
        "top_matches": [
            {"rank": rank, "title": item.title, "modality": item.kind, "cosine": round(score, 6)}
            for rank, (item, score) in enumerate(ranked[:5], start=1)
        ],
        "query_vector_preview": [round(float(value), 6) for value in query[0, :12]],
        "note": "Cosine similarity is a ranking signal, not a calibrated probability.",
    }
    return summary, rankings, gallery, curve, fingerprint, diagnostics


def _pair_duration(*args: Any, **kwargs: Any) -> int:
    payloads = args[:2]
    has_video = any(
        (isinstance(value, dict) and bool(value.get("video")))
        or (isinstance(value, str) and _is_video(value))
        for value in payloads
    )
    return 150 if has_video else 90


@spaces.GPU(duration=_pair_duration)
def _run_pair(
    query_payload: Any,
    candidate_payload: Any,
    progress: gr.Progress,
) -> tuple[torch.Tensor, torch.Tensor, float]:
    started = time.perf_counter()
    progress(0.08, desc="Allocating the model on GPU")
    MODEL.to("cuda")
    try:
        progress(0.35, desc="Encoding the query")
        query = _encode([query_payload], query=True)
        progress(0.68, desc="Encoding the candidate")
        candidate = _encode([candidate_payload], query=False)
        return query, candidate, time.perf_counter() - started
    finally:
        MODEL.to("cpu")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def _interpret_score(score: float) -> tuple[str, str]:
    if score >= 0.75:
        return "high alignment", "These inputs occupy a very similar region for this retrieval model."
    if score >= 0.50:
        return "meaningful alignment", "The model sees a substantial semantic relationship."
    if score >= 0.25:
        return "weak alignment", "There is some overlap, but stronger candidates may rank above it."
    return "low alignment", "The model places these inputs relatively far apart."


def _render_pair_score(score: float, dimension: int, query_kind: str, candidate_kind: str, elapsed: float) -> str:
    label, explanation = _interpret_score(score)
    ring = min(100.0, max(0.0, (score + 1.0) * 50.0))
    return f"""
    <section class="pair-score-card">
      <div class="score-orbit" style="--score:{ring:.2f}">
        <div><strong>{score:+.3f}</strong><span>cosine</span></div>
      </div>
      <div class="pair-score-copy">
        <small>PAIRWISE READOUT</small>
        <h2>{html.escape(label)}</h2>
        <p>{html.escape(explanation)}</p>
        <div class="run-meta"><span>{html.escape(query_kind)}</span><i></i><span>{html.escape(candidate_kind)}</span><i></i><span>{dimension:,}D</span><i></i><span>{elapsed:.1f}s</span></div>
      </div>
    </section>
    """


def compare_experience(
    query_text: str,
    query_image: Any,
    query_video: Any,
    candidate_text: str,
    candidate_image: Any,
    candidate_video: Any,
    dimension: int,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str, str, str, dict[str, Any]]:
    """Compare any two supported inputs across all Matryoshka dimensions."""
    dimension = int(dimension)
    query_payload, query_kind = _multimodal_payload(query_text, query_image, query_video)
    candidate_payload, candidate_kind = _multimodal_payload(candidate_text, candidate_image, candidate_video)
    with _INFERENCE_LOCK:
        query, candidate, elapsed = _run_pair(query_payload, candidate_payload, progress)
    scores_by_dimension = {
        dim: float((_truncate_normalize(query, dim) @ _truncate_normalize(candidate, dim).T).item())
        for dim in MATRYOSHKA_DIMS
    }
    selected_score = scores_by_dimension[dimension]
    score_card = _render_pair_score(selected_score, dimension, query_kind, candidate_kind, elapsed)
    curve = _render_curve(
        {f"{query_kind} → {candidate_kind}": [scores_by_dimension[dim] for dim in MATRYOSHKA_DIMS]},
        "Semantic alignment under compression",
    )
    fingerprints = (
        '<div class="fingerprint-pair">'
        + _fingerprint_svg(query[0, :dimension], f"Query · {query_kind}")
        + _fingerprint_svg(candidate[0, :dimension], f"Candidate · {candidate_kind}")
        + "</div>"
    )
    diagnostics = {
        "model": MODEL_ID,
        "query_modality": query_kind,
        "candidate_modality": candidate_kind,
        "selected_dimension": dimension,
        "selected_cosine": round(selected_score, 6),
        "cosine_by_dimension": {str(dim): round(score, 6) for dim, score in scores_by_dimension.items()},
        "gpu_pass_seconds": round(elapsed, 3),
        "note": "Interpret thresholds relative to a task-specific candidate set.",
    }
    return score_card, curve, fingerprints, diagnostics


CSS = """
:root {
  --ink: #f6f3ec;
  --muted: #a9abb5;
  --panel: rgba(17, 20, 26, .82);
  --line: rgba(255, 255, 255, .10);
  --warm: #ffad66;
  --cool: #71e2da;
  --violet: #a994ff;
}

body, .gradio-container {
  background:
    radial-gradient(circle at 13% 0%, rgba(255, 143, 68, .16), transparent 30rem),
    radial-gradient(circle at 92% 13%, rgba(88, 218, 211, .11), transparent 34rem),
    #090b10 !important;
  color: var(--ink) !important;
}
.gradio-container { max-width: 1380px !important; padding: 0 28px 60px !important; }
.gradio-container * { box-sizing: border-box; }
.gradio-container .prose { color: var(--ink); }
.gradio-container label, .gradio-container .label-wrap { color: #d7d7dc !important; }
.gradio-container input, .gradio-container textarea {
  background: rgba(7, 9, 13, .72) !important;
  border-color: rgba(255,255,255,.12) !important;
  color: #f8f6f1 !important;
}
.gradio-container .block, .gradio-container .form {
  border-color: var(--line) !important;
}

#hero { padding: 76px 4px 38px; }
.hero-shell { position: relative; overflow: hidden; border-bottom: 1px solid var(--line); padding-bottom: 44px; }
.eyebrow { display:flex; align-items:center; gap:10px; color:var(--cool); font-size:12px; font-weight:700; letter-spacing:.18em; text-transform:uppercase; }
.eyebrow:before { content:""; width:26px; height:1px; background:var(--cool); box-shadow:0 0 14px var(--cool); }
.hero-title { margin: 18px 0 8px; font-size: clamp(52px, 8vw, 112px); line-height:.88; letter-spacing:-.075em; font-weight:760; }
.hero-title .accent { color:transparent; -webkit-text-stroke:1px rgba(255,255,255,.62); }
.hero-title .dot { color:var(--warm); text-shadow:0 0 42px rgba(255,173,102,.6); }
.hero-sub { max-width:760px; margin:24px 0 0; font-size:clamp(17px,2vw,23px); line-height:1.55; color:#c0c1c8; }
.hero-grid { display:grid; grid-template-columns:1fr auto; align-items:end; gap:30px; }
.hero-stats { display:grid; grid-template-columns:repeat(2,minmax(110px,1fr)); gap:1px; background:var(--line); border:1px solid var(--line); min-width:350px; }
.hero-stats div { background:rgba(9,11,16,.88); padding:18px 20px; }
.hero-stats strong { display:block; font-size:28px; line-height:1; color:#fff; letter-spacing:-.04em; }
.hero-stats span { display:block; margin-top:8px; color:var(--muted); font-size:11px; text-transform:uppercase; letter-spacing:.12em; }
.capability-rail { display:flex; gap:8px; flex-wrap:wrap; margin-top:26px; }
.capability-rail span { border:1px solid var(--line); border-radius:99px; padding:7px 12px; color:#c9c9cf; font-size:12px; background:rgba(255,255,255,.025); }

.section-intro { margin:34px 0 18px; }
.section-intro small, .viz-heading small, .pair-score-copy small { color:var(--warm); letter-spacing:.16em; font-weight:750; font-size:11px; }
.section-intro h2 { font-size:30px; letter-spacing:-.035em; margin:6px 0; }
.section-intro p { color:var(--muted); margin:0; max-width:760px; }

.input-panel, .output-panel { background:linear-gradient(145deg,rgba(22,25,32,.9),rgba(12,14,19,.86)) !important; border:1px solid var(--line) !important; border-radius:18px !important; padding:18px !important; box-shadow:0 22px 70px rgba(0,0,0,.24); }
.primary-action { min-height:52px !important; border:0 !important; color:#16110d !important; font-weight:800 !important; letter-spacing:.01em; background:linear-gradient(105deg,#ff8e54,#ffd187) !important; box-shadow:0 10px 30px rgba(255,142,84,.2) !important; }
.primary-action:hover { transform:translateY(-1px); filter:brightness(1.04); }

.run-summary { border:1px solid rgba(112,228,218,.22); border-radius:18px; padding:22px 24px; background:linear-gradient(120deg,rgba(31,48,48,.54),rgba(17,19,25,.94)); margin-bottom:16px; }
.run-kicker { color:var(--cool); font-size:11px; text-transform:uppercase; letter-spacing:.16em; font-weight:750; }
.live-dot { display:inline-block; width:7px; height:7px; border-radius:99px; background:var(--cool); box-shadow:0 0 15px var(--cool); margin-right:8px; }
.run-main { display:flex; align-items:flex-end; justify-content:space-between; gap:20px; margin:14px 0 18px; }
.run-label { display:block; color:var(--muted); font-size:12px; margin-bottom:5px; }
.run-main strong { font-size:clamp(23px,3vw,38px); letter-spacing:-.04em; }
.hero-score { text-align:right; }
.hero-score span { display:block; font-size:40px; color:var(--cool); font-variant-numeric:tabular-nums; letter-spacing:-.05em; }
.hero-score small { color:var(--muted); text-transform:uppercase; letter-spacing:.14em; }
.run-meta { display:flex; align-items:center; gap:10px; flex-wrap:wrap; color:#aeb0b8; font-size:12px; }
.run-meta i { display:block; width:3px; height:3px; border-radius:99px; background:#555963; }

.ranking-stack { display:grid; gap:9px; }
.rank-row { display:grid; grid-template-columns:48px 1fr 80px; gap:16px; align-items:center; padding:15px 17px; border:1px solid var(--line); border-radius:14px; background:rgba(255,255,255,.025); transition:.2s ease; }
.rank-row:hover { transform:translateX(3px); border-color:rgba(255,173,102,.28); background:rgba(255,255,255,.04); }
.rank-row.winner { border-color:rgba(255,173,102,.35); background:linear-gradient(100deg,rgba(255,150,85,.11),rgba(255,255,255,.025)); }
.rank-number { color:#6d7079; font-size:14px; font-variant-numeric:tabular-nums; }
.rank-title-line { display:flex; gap:9px; align-items:center; flex-wrap:wrap; }
.rank-title-line strong { font-size:16px; color:#f7f3ec; }
.kind-pill { font-size:9px; letter-spacing:.09em; text-transform:uppercase; color:#bfc1c8; border:1px solid var(--line); border-radius:99px; padding:4px 7px; }
.rank-copy p { margin:4px 0 9px; color:#91949e; font-size:12px; line-height:1.45; }
.score-track { height:2px; background:rgba(255,255,255,.06); overflow:hidden; }
.score-track span { display:block; height:100%; background:linear-gradient(90deg,var(--warm),var(--cool)); }
.rank-score { text-align:right; font-size:18px; font-variant-numeric:tabular-nums; color:#b9bbc2; }
.rank-score.high { color:var(--cool); }.rank-score.mid { color:var(--warm); }
.rank-score small { display:block; font-size:8px; color:#747780; letter-spacing:.14em; margin-top:3px; text-transform:uppercase; }

.viz-card { border:1px solid var(--line); border-radius:18px; padding:20px; background:rgba(15,17,23,.78); overflow:hidden; }
.viz-heading { display:flex; align-items:flex-start; justify-content:space-between; gap:16px; }
.viz-heading h3 { margin:5px 0 0; font-size:20px; letter-spacing:-.025em; }
.viz-heading > span { color:#8f929d; font-size:11px; border:1px solid var(--line); border-radius:99px; padding:6px 9px; }
.dimension-chart, .fingerprint { width:100%; height:auto; overflow:visible; }
.chart-grid { stroke:rgba(255,255,255,.07); stroke-width:1; }
.chart-axis { fill:#777b85; font-size:10px; font-family:ui-monospace,SFMono-Regular,Menlo,monospace; }
.chart-legend { display:flex; flex-wrap:wrap; gap:8px 16px; }
.chart-legend span { color:#aeb0b8; font-size:11px; }
.chart-legend b { display:inline-block; width:7px; height:7px; border-radius:99px; margin-right:6px; }
.viz-note { margin:14px 0 0; color:#6f727c; font-size:11px; }
.zero-line { stroke:rgba(255,255,255,.18); stroke-width:1; }
.fingerprint-key { display:flex; gap:14px; align-items:center; flex-wrap:wrap; color:#848791; font-size:10px; text-transform:uppercase; letter-spacing:.08em; }
.fingerprint-key b { display:inline-block; width:7px; height:7px; border-radius:2px; margin-right:5px; }.fingerprint-key .positive{background:var(--cool)}.fingerprint-key .negative{background:var(--warm)}
.fingerprint-key em { margin-left:auto; text-transform:none; letter-spacing:0; color:#6e717a; }
.fingerprint-pair { display:grid; grid-template-columns:1fr 1fr; gap:12px; }

.pair-score-card { display:grid; grid-template-columns:190px 1fr; gap:30px; align-items:center; border:1px solid rgba(169,148,255,.24); border-radius:20px; padding:26px; background:linear-gradient(125deg,rgba(50,40,80,.35),rgba(14,17,23,.92)); }
.score-orbit { width:170px; aspect-ratio:1; border-radius:50%; padding:2px; background:conic-gradient(var(--violet) calc(var(--score)*1%),rgba(255,255,255,.07) 0); box-shadow:0 0 55px rgba(169,148,255,.12); }
.score-orbit > div { width:100%; height:100%; border-radius:50%; background:#0d0f15; display:flex; flex-direction:column; align-items:center; justify-content:center; }
.score-orbit strong { font-size:38px; letter-spacing:-.05em; font-variant-numeric:tabular-nums; }.score-orbit span{color:#777b85;font-size:10px;text-transform:uppercase;letter-spacing:.14em}
.pair-score-copy h2 { margin:6px 0 8px; font-size:34px; letter-spacing:-.045em; }.pair-score-copy p{color:#aeb0b8;max-width:580px}

.model-note { margin:28px 0 0; padding:18px 20px; border-left:2px solid var(--warm); background:rgba(255,173,102,.05); color:#9ea0aa; font-size:12px; line-height:1.6; }
.space-footer { display:flex; justify-content:space-between; align-items:center; gap:20px; flex-wrap:wrap; border-top:1px solid var(--line); padding:28px 4px 0; margin-top:44px; color:#767984; font-size:12px; }
.space-footer a { color:#c4c6cd !important; text-decoration:none; }.space-footer a:hover{color:var(--warm)!important}

@media (max-width: 900px) {
  .gradio-container { padding:0 14px 40px !important; }
  #hero { padding-top:44px; }
  .hero-grid { grid-template-columns:1fr; }
  .hero-stats { min-width:0; width:100%; }
  .fingerprint-pair { grid-template-columns:1fr; }
}
@media (max-width: 620px) {
  .hero-title { font-size:50px; }
  .hero-stats { grid-template-columns:1fr 1fr; }
  .rank-row { grid-template-columns:34px 1fr 64px; gap:8px; padding:12px 10px; }
  .rank-copy p { display:none; }
  .pair-score-card { grid-template-columns:1fr; text-align:center; }
  .score-orbit { margin:auto; }.pair-score-copy .run-meta{justify-content:center}
}
"""


HERO = """
<div id="hero" class="hero-shell">
  <div class="eyebrow">Tencent WeMM · multimodal embedding</div>
  <div class="hero-grid">
    <div>
      <h1 class="hero-title">One space<span class="accent">.<br>Every medium</span><span class="dot">.</span></h1>
      <p class="hero-sub">Search meaning—not file types—across text, images, video, charts, and visual documents in one shared semantic geometry.</p>
      <div class="capability-rail"><span>text ↔ image</span><span>text ↔ video</span><span>visual documents</span><span>interleaved inputs</span><span>Matryoshka embeddings</span></div>
    </div>
    <div class="hero-stats">
      <div><strong>4096</strong><span>native dimensions</span></div>
      <div><strong>9B</strong><span>parameters</span></div>
      <div><strong>80.6</strong><span>MMEB-v2 avg</span></div>
      <div><strong>190</strong><span>MMEB-v3 tasks</span></div>
    </div>
  </div>
</div>
"""


INTRO_SEARCH = """
<div class="section-intro">
  <small>01 / RETRIEVAL UNIVERSE</small>
  <h2>Ask in one modality. Discover in another.</h2>
  <p>Search the built-in field of screenshots, figures, dense documents, video, and multilingual text—or bring your own candidates.</p>
</div>
"""


INTRO_PAIR = """
<div class="section-intro">
  <small>02 / VECTOR MICROSCOPE</small>
  <h2>Put any two ideas under the lens.</h2>
  <p>Use a query and candidate as text, image, video, or visual-plus-text. Then watch their alignment change as the embedding compresses.</p>
</div>
"""


INITIAL_SUMMARY = """
<section class="run-summary">
  <div class="run-kicker"><span class="live-dot"></span> model ready</div>
  <div class="run-main"><div><span class="run-label">Awaiting a query</span><strong>Search across media</strong></div><div class="hero-score"><span>—</span><small>cosine</small></div></div>
  <div class="run-meta"><span>text</span><i></i><span>image</span><i></i><span>video</span><i></i><span>visual documents</span></div>
</section>
"""


INITIAL_RANKING = """
<div class="model-note">Choose a curated example below or compose a multimodal query. The first run maps the showcase corpus once; later searches reuse its CPU-cached 4,096D vectors.</div>
"""


with gr.Blocks(css=CSS, title="WeMM · Multimodal Embedding Universe", fill_width=True) as demo:
    gr.HTML(HERO)

    with gr.Tabs():
        with gr.Tab("Search the universe", id="search"):
            gr.HTML(INTRO_SEARCH)
            with gr.Row(equal_height=False):
                with gr.Column(scale=5, elem_classes="input-panel"):
                    query_text = gr.Textbox(
                        label="Query · text",
                        placeholder="Try: Which document explains temporary road closures?",
                        lines=3,
                        max_lines=7,
                    )
                    with gr.Row():
                        query_image = gr.Image(
                            label="Query · image (optional)",
                            type="filepath",
                            sources=["upload", "clipboard", "webcam"],
                            height=220,
                        )
                        query_video = gr.Video(
                            label="Query · video (optional)",
                            format="mp4",
                            height=220,
                        )
                    gr.Markdown("Add text to an image or video to create a joint multimodal query.")
                    with gr.Accordion("Build your own candidate collection", open=False):
                        custom_texts = gr.Textbox(
                            label="Text candidates",
                            placeholder="Title :: Candidate text\nAnother title :: Another candidate",
                            lines=5,
                            info=f"One candidate per line, up to {MAX_CUSTOM_TEXTS}.",
                        )
                        candidate_media = gr.Gallery(
                            label="Image + video candidates",
                            type="filepath",
                            file_types=["image", "video"],
                            sources=["upload"],
                            columns=3,
                            height=250,
                        )
                    include_showcase = gr.Checkbox(
                        value=True,
                        label="Include the curated multimodal universe",
                        info="10 candidates spanning text, images, video, figures, and visual documents.",
                    )
                    dimension = gr.Radio(
                        choices=list(MATRYOSHKA_DIMS),
                        value=1024,
                        label="Embedding budget",
                        info="Native Matryoshka dimensions; smaller vectors trade storage for fidelity.",
                    )
                    search_button = gr.Button("Map the semantic field  →", variant="primary", elem_classes="primary-action")

                with gr.Column(scale=7, elem_classes="output-panel"):
                    search_summary = gr.HTML(INITIAL_SUMMARY)
                    ranking_output = gr.HTML(INITIAL_RANKING)

            with gr.Row(equal_height=False):
                result_gallery = gr.Gallery(
                    value=[
                        (str(ASSET_DIR / "llama4_hgf.png"), "Visual document · Llama 4 model card"),
                        (str(ASSET_DIR / "doc2.jpg"), "Visual document · 1971 budget infographic"),
                        (str(ASSET_DIR / "mapo_tofu.mp4"), "Video · Mapo tofu in motion"),
                        (str(ASSET_DIR / "doc4.jpg"), "Visual document · road-safety assessment"),
                    ],
                    label="Ranked visual field",
                    columns=4,
                    rows=2,
                    height=430,
                    object_fit="contain",
                    interactive=False,
                    buttons=["fullscreen", "download_all"],
                )
            with gr.Row(equal_height=False):
                dimension_output = gr.HTML('<div class="viz-card"><div class="viz-heading"><div><small>MATRYOSHKA SCOPE</small><h3>Dimension stability appears here</h3></div></div></div>')
                fingerprint_output = gr.HTML('<div class="viz-card"><div class="viz-heading"><div><small>VECTOR FINGERPRINT</small><h3>Your query vector appears here</h3></div></div></div>')
            with gr.Accordion("Embedding telemetry · inspect the API payload", open=False):
                diagnostics_output = gr.JSON(label="Diagnostics")

            example_inputs = [
                query_text,
                query_image,
                query_video,
                custom_texts,
                candidate_media,
                include_showcase,
                dimension,
            ]
            example_outputs = [
                search_summary,
                ranking_output,
                result_gallery,
                dimension_output,
                fingerprint_output,
                diagnostics_output,
            ]
            gr.Examples(
                examples=[
                    ["Which Llama 4 model variants are available?", None, None, "", None, True, 512],
                    ["How is mapo tofu prepared?", None, None, "", None, True, 1024],
                    ["Find the environmental assessment page about driver training and temporary road closures.", None, None, "", None, True, 256],
                    [
                        "Match this screenshot to the most relevant description.",
                        str(ASSET_DIR / "llama4_hgf.png"),
                        None,
                        "Llama family :: Scout and Maverick are multimodal mixture-of-experts model variants.\nRecipe :: Soft tofu simmered in spicy chili-bean sauce.",
                        None,
                        False,
                        256,
                    ],
                    [
                        "What dish is being prepared in this clip?",
                        None,
                        str(ASSET_DIR / "mapo_tofu.mp4"),
                        "Sichuan classic :: Mapo tofu combines soft tofu with a spicy, numbing bean-paste sauce.\nSpaceflight :: A launch vehicle carries a satellite into orbit.",
                        None,
                        False,
                        512,
                    ],
                ],
                inputs=example_inputs,
                outputs=example_outputs,
                fn=search_experience,
                cache_examples=True,
                cache_mode="lazy",
                label="Curated expeditions",
                example_labels=[
                    "Find Llama 4 across a screenshot",
                    "Search a cooking video with text",
                    "Retrieve a dense safety document",
                    "Match an image to text candidates",
                    "Match a video to text candidates",
                ],
            )

            search_event = search_button.click(
                fn=search_experience,
                inputs=example_inputs,
                outputs=example_outputs,
                api_name="search",
                api_description="Rank a mixed text/image/video collection using WeMM-Embedding-9B.",
                concurrency_limit=1,
                concurrency_id="wemm_gpu",
                time_limit=300,
                scroll_to_output=True,
            )
            query_text.submit(
                fn=search_experience,
                inputs=example_inputs,
                outputs=example_outputs,
                api_name=None,
                api_visibility="private",
                concurrency_limit=1,
                concurrency_id="wemm_gpu",
                time_limit=300,
                scroll_to_output=True,
            )

        with gr.Tab("Compare two ideas", id="compare"):
            gr.HTML(INTRO_PAIR)
            with gr.Row(equal_height=False):
                with gr.Column(elem_classes="input-panel"):
                    gr.Markdown("### A · Query")
                    pair_query_text = gr.Textbox(label="Text", placeholder="Describe or contextualize the query", lines=3)
                    with gr.Row():
                        pair_query_image = gr.Image(label="Image", type="filepath", height=210)
                        pair_query_video = gr.Video(label="Video", format="mp4", height=210)
                with gr.Column(elem_classes="input-panel"):
                    gr.Markdown("### B · Candidate")
                    pair_candidate_text = gr.Textbox(label="Text", placeholder="Describe or contextualize the candidate", lines=3)
                    with gr.Row():
                        pair_candidate_image = gr.Image(label="Image", type="filepath", height=210)
                        pair_candidate_video = gr.Video(label="Video", format="mp4", height=210)
            pair_dimension = gr.Radio(
                choices=list(MATRYOSHKA_DIMS),
                value=1024,
                label="Embedding budget",
            )
            compare_button = gr.Button("Measure semantic alignment  →", variant="primary", elem_classes="primary-action")
            pair_score_output = gr.HTML('<div class="model-note">Add one modality on each side. You may pair visual media with text context.</div>')
            pair_curve_output = gr.HTML('<div class="viz-card"><div class="viz-heading"><div><small>MATRYOSHKA SCOPE</small><h3>Alignment by dimension appears here</h3></div></div></div>')
            pair_fingerprints_output = gr.HTML()
            with gr.Accordion("Pairwise telemetry", open=False):
                pair_diagnostics_output = gr.JSON(label="Diagnostics")

            pair_inputs = [
                pair_query_text,
                pair_query_image,
                pair_query_video,
                pair_candidate_text,
                pair_candidate_image,
                pair_candidate_video,
                pair_dimension,
            ]
            compare_button.click(
                fn=compare_experience,
                inputs=pair_inputs,
                outputs=[pair_score_output, pair_curve_output, pair_fingerprints_output, pair_diagnostics_output],
                api_name="compare",
                api_description="Compare a multimodal query and candidate across every native Matryoshka dimension.",
                concurrency_limit=1,
                concurrency_id="wemm_gpu",
                time_limit=240,
                scroll_to_output=True,
            )

    gr.HTML(
        """
        <div class="model-note"><strong>Read scores comparatively.</strong> Cosine similarity is useful for ranking candidates within a collection; it is not a calibrated confidence or a universal relevance grade. Audio is not supported by WeMM-Embedding-9B.</div>
        <footer class="space-footer"><span>WeMM-Embedding-9B · Apache-2.0 · built on Qwen3.5</span><span><a href="https://huggingface.co/tencent/WeMM-Embedding-9B" target="_blank">Model card ↗</a>&nbsp;&nbsp;&nbsp;<a href="https://arxiv.org/abs/2608.24053" target="_blank">Technical report ↗</a></span></footer>
        """
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=24).launch(
        allowed_paths=[str(ASSET_DIR)],
        show_error=True,
    )