"""
app.py — IAnahid v2
====================
A Gradio web application for semi-automatic image annotation for heritage archives.

Features:
  - Home page with guidelines sidebar and mode selection
  - Two annotation modes:
      * Predefined Schema  — fixed Tag 1 / Tag 2 / Caption workflow (original mode)
      * Custom Schema      — user defines their own metadata fields
  - Image loading via URL list file (CSV, TXT or JSON); images are never downloaded,
    only previewed via <img src="…"> to keep the app fast
  - Resume work by loading a previously exported annotation file (CSV or JSON)
  - Export annotations as CSV or JSON
  - Mistral-powered caption generation (both modes)
  - Cluster-level batch validation (both modes)
  - Predefined mode can load the HuggingFace-hosted exemple.csv directly
"""

import gradio as gr
import pandas as pd
import numpy as np
import os
import re
import io
import json
import base64
import tempfile
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
from huggingface_hub import hf_hub_download

# Force `transformers` to a torch-only backend. Some environments also have
# TensorFlow/Flax installed; letting transformers import them is slow and makes
# Abseil print "[mutex.cc] RAW: Lock blocking …" diagnostics (and can stall model
# loading). Torch-only avoids both. Must be set before transformers is imported.
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("USE_FLAX", "0")
os.environ.setdefault("USE_TORCH", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# ── Optional logos footer ─────────────────────────────────────────────────────
try:
    from logos_b64 import LOGOS as _LOGOS
except Exception:
    _LOGOS = ""

# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# Set DATASET_REPO and HF_TOKEN as HuggingFace Space Secrets if you want to
# pre-load a dataset from HuggingFace Hub. Otherwise the app starts empty.
# ══════════════════════════════════════════════════════════════════════════════
DATASET_REPO = os.environ.get("DATASET_REPO", "")
HF_TOKEN     = os.environ.get("HF_TOKEN",     "")
CSV_FILENAME = "exemple.csv"

# Mistral prompt for automated caption generation (predefined mode)
PROMPT_MISTRAL = (
    "Rules:\n"
    "* Describe only the visually observable elements "
    "(architecture, materials, spatial organization, objects, people, "
    "clothing, gestures, inscriptions, landscape).\n"
    "* Do not interpret the meaning, the symbolism, or the historical context.\n"
    "* Use precise, neutral, and documentary language.\n"
    "* Write in complete sentences (paragraph), 2-3 lines maximum.\n"
    "Output: A single caption suited to a heritage archive catalogue.\n"
    "Structure: Built environment and architectural elements, "
    "Spatial context and surrounding landscape, Objects and materials, "
    "Human figures (where applicable), Animals (where applicable).\n"
    "Expected output: A single developed caption (2-3 lines max)."
)

# Controlled vocabularies for the predefined schema
VOCAB_TAG_1 = [
    "1.1. Photograph",
    "1.2. Written document",
    "1.3. Other reproduced document",
    "1.4. Conditioning / packaging material",
]
VOCAB_TAG_2 = [
    "2.1. Architecture",
    "2.2. Objects",
    "2.3. People",
    "2.4. Landscape",
    "2.5. Animal",
    "2.6. Plants",
]

# ══════════════════════════════════════════════════════════════════════════════
# GUIDELINES TEXT (shown in the Home sidebar)
# ══════════════════════════════════════════════════════════════════════════════
GUIDELINES = """
## How to use IAnahid

### Step 1 · Choose a mode

**🏛️ Predefined Schema** uses a fixed annotation workflow (Tag 1 / Tag 2 /
Caption) designed for heritage archive cataloguing. The ThierryNum dataset
can be loaded directly from HuggingFace with one click.

**🛠️ Custom Schema** lets you define your own metadata fields before you
start. You decide the field names, types (free text or dropdown), and
dropdown options.

Both modes support **AI captions** (Mistral) and **cluster batch validation**.

---

### Step 2 · Load your images

Provide a **URL list file** — a file listing image URLs (one per entry).
Images are **never downloaded**: they display directly from their source,
keeping the app fast.

| Format | Expected structure |
|--------|--------------------|
| `.txt` | One URL per line |
| `.csv` | Column named `image_url` (or first column) |
| `.json` | `["url1", …]` or `[{"image_url": "…"}, …]` |

In **Predefined** mode you can also click **Load from HuggingFace** to use
the pre-configured `exemple.csv` dataset (requires `DATASET_REPO` and
`HF_TOKEN` Space secrets).

---

### Step 3 · Build your schema (Custom mode only)

Go to the **📐 Schema builder** tab first. Add fields one by one — give each
a name, choose *text* or *dropdown*, and (for dropdowns) list the allowed
values separated by commas.

---

### Step 4 · Start fresh or resume

- **Start fresh** — load a URL list; a blank annotation file is created.
- **Resume** — load a previously exported file (CSV or JSON). Existing
  annotations are restored and the app jumps to the first unvalidated item.
  For Custom JSON files the schema is restored automatically.

---

### Step 5 · Annotate

Use **Previous / Skip** to navigate. In Custom mode fill the fields you
defined and click **Save & continue**. In Predefined mode use the dedicated
tabs for Tag 1, Tag 2, and Caption.

The **🗂️ Clusters** tab (available in both modes) lets you assign a value to
all images in a cluster at once — useful for large visually similar groups.

The **📝 Caption** tab uses the **Mistral** multimodal API to generate a
draft description. Paste your API key, pick a model, and click Generate.

---

### Step 6 · Export

Go to the **📁 Data** tab at any time and choose:

- **CSV** — spreadsheet-compatible, importable in most tools
- **JSON** — structured format; for Custom mode, your schema is embedded so
  the file can be fully reloaded later

⚠️ Annotations live only in your session. **Always export before closing.**
"""

# ══════════════════════════════════════════════════════════════════════════════
# DATA HELPERS — shared utilities
# ══════════════════════════════════════════════════════════════════════════════

def _ensure_predefined_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Add missing annotation columns for the predefined schema."""
    if not df.empty:
        df.columns = [c.strip() for c in df.columns]
    defaults = {
        "image_url": "",
        "title": "",
        "clustering_id": "0",
        "tag_1": "",
        "tag_2": "",
        "image_caption": "",
        "validated_tag_1": False,
        "validated_tag_2": False,
        "validated_caption": False,
    }
    for col, val in defaults.items():
        if col not in df.columns:
            df[col] = val
    for col in ["image_url", "title", "clustering_id", "tag_1", "tag_2", "image_caption"]:
        if col in df.columns:
            df[col] = df[col].astype(object).where(df[col].notna(), "")
    return df


def _ensure_custom_columns(df: pd.DataFrame, schema: list) -> pd.DataFrame:
    """Add missing annotation columns for a custom schema."""
    if not df.empty:
        df.columns = [c.strip() for c in df.columns]
    if "image_url" not in df.columns:
        df["image_url"] = ""
    if "title" not in df.columns:
        df["title"] = ""
    if "clustering_id" not in df.columns:
        df["clustering_id"] = "0"
    df["clustering_id"] = df["clustering_id"].astype(object).where(df["clustering_id"].notna(), "0")
    for field in schema:
        col = field["name"]
        if col not in df.columns:
            df[col] = ""
        val_col = f"validated_{col}"
        if val_col not in df.columns:
            df[val_col] = False
        df[col] = df[col].astype(object).where(df[col].notna(), "")
    return df


# ── URL list loaders ──────────────────────────────────────────────────────────

def load_url_list(path: str) -> list:
    """Parse a URL list file (.txt, .csv, .json) and return a list of URL strings."""
    ext = os.path.splitext(path)[1].lower()
    if ext == ".txt":
        with open(path, "r", encoding="utf-8") as f:
            urls = [line.strip() for line in f if line.strip()]
    elif ext == ".csv":
        with open(path, "r", encoding="utf-8") as f:
            sep = ";" if ";" in f.readline() else ","
        tmp = pd.read_csv(path, sep=sep, encoding="utf-8", on_bad_lines="warn")
        if "image_url" in tmp.columns:
            urls = tmp["image_url"].dropna().astype(str).tolist()
        else:
            urls = tmp.iloc[:, 0].dropna().astype(str).tolist()
    elif ext == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, list):
            if data and isinstance(data[0], str):
                urls = data
            else:
                urls = [str(d.get("image_url", "")) for d in data if "image_url" in d]
        else:
            urls = []
    else:
        urls = []
    return [u for u in urls if u and u.lower() != "nan"]


# ── Annotation file loaders ───────────────────────────────────────────────────

def load_annotation_file(path: str):
    """Load a previously exported annotation file (CSV or JSON).

    Returns (df, schema_or_None).
    For JSON files saved by the custom mode, the embedded schema is also returned.
    """
    ext = os.path.splitext(path)[1].lower()
    schema = None
    if ext == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, dict) and "schema" in data and "records" in data:
            schema = data["schema"]
            df = pd.DataFrame(data["records"])
        else:
            df = pd.DataFrame(data if isinstance(data, list) else [data])
    else:
        with open(path, "r", encoding="utf-8") as f:
            sep = ";" if ";" in f.readline() else ","
        df = pd.read_csv(path, sep=sep, quotechar='"', on_bad_lines="warn", encoding="utf-8")
    return df, schema


# ── Export helpers ────────────────────────────────────────────────────────────

def export_as_csv(df: pd.DataFrame) -> str:
    """Write df to a temp CSV file and return its path.

    A fresh directory is used per export so concurrent users on a shared
    server (e.g. a HuggingFace Space) never overwrite each other's files.
    """
    out = os.path.join(tempfile.mkdtemp(prefix="ianahid_"), "ianahid_export.csv")
    df.to_csv(out, sep=",", index=False, quoting=1, encoding="utf-8")
    return out


def export_as_json(df: pd.DataFrame, schema=None) -> str:
    """Write df to a temp JSON file and return its path.

    If a custom schema is provided, it is embedded in the output so the file
    can be reloaded and the schema restored automatically.
    """
    out = os.path.join(tempfile.mkdtemp(prefix="ianahid_"), "ianahid_export.json")
    records = df.to_dict(orient="records")
    if schema:
        payload = {"schema": schema, "records": records}
    else:
        payload = records
    with open(out, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    return out


# ── Navigation & display helpers ──────────────────────────────────────────────

TBL_PREVIEW_ROWS = 200


def tbl_preview(df: pd.DataFrame) -> pd.DataFrame:
    """Cap the overview table sent to the browser.

    Shipping thousands of rows through the Space's proxy on every refresh makes
    the UI hang; the full DataFrame stays server-side and exports are unaffected.
    """
    if df is None:
        return pd.DataFrame()
    return df.head(TBL_PREVIEW_ROWS)


def tbl_note(df: pd.DataFrame) -> str:
    if df is None or len(df) <= TBL_PREVIEW_ROWS:
        return ""
    return (f"*Table preview limited to the first {TBL_PREVIEW_ROWS} of {len(df)} rows "
            f"— exports always contain all rows.*")


def counter(idx: int, total: int) -> str:
    return f"**{int(idx) + 1} / {total}**"


def first_unvalidated(df: pd.DataFrame, col: str) -> int:
    if df.empty or col not in df.columns:
        return 0
    mask = (df[col] == False) | (df[col].isna()) | (df[col] == "")
    idxs = df[mask].index
    return int(idxs[0]) if not idxs.empty else 0


def img_url_at(df: pd.DataFrame, idx: int) -> str | None:
    try:
        v = df.iloc[int(idx)]["image_url"]
        return str(v) if pd.notna(v) and str(v).strip() else None
    except Exception:
        return None


# Some image hosts (Wikimedia, certain IIIF servers) reject the default
# python-requests User-Agent; identify ourselves so fetches don't silently fail.
_HTTP_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; IAnahid/1.0; heritage image annotator)"}

_IIIF_TAIL = re.compile(r"/full/(?:full|max)/(\d+)/(default|color|gray|bitonal)\.(\w+)$")


def iiif_resize(url: str, width: int) -> str:
    """If url is a IIIF Image API request (…/{region}/{size}/{rot}/{quality}.{fmt}),
    rewrite the size segment to request `width` px instead of full resolution.

    Handles IIIF 2.x (`/full/full/…`) and 3.0 (`/full/max/…`). Non-IIIF URLs
    are returned unchanged, so images from any other host still load normally.
    """
    if not url:
        return url
    return _IIIF_TAIL.sub(rf"/full/{width},/\1/\2.\3", url)


def img_html(url: str | None, height: int = 420, width: int = 800) -> str:
    """Return an HTML <img> tag for the given URL.

    Images are loaded directly from their online source — no server-side
    download. For IIIF sources the URL is downscaled to `width` px to keep
    loading fast. A placeholder is shown when URL is None.
    """
    if not url:
        return (
            f'<div style="height:{height}px;display:flex;align-items:center;'
            f'justify-content:center;background:#f5f5f5;border-radius:8px;'
            f'color:#999;font-size:14px;">No image URL</div>'
        )
    src = iiif_resize(url, width)
    return (
        f'<img src="{src}" loading="lazy" decoding="async" '
        f'style="max-height:{height}px;max-width:100%;'
        f'object-fit:contain;border-radius:8px;background:#f5f5f5;" '
        f'onerror="this.style.display=\'none\';this.nextSibling.style.display=\'flex\'" />'
        f'<div style="display:none;height:{height}px;align-items:center;'
        f'justify-content:center;background:#f5f5f5;border-radius:8px;'
        f'color:#c00;font-size:13px;">⚠ Could not load image from URL</div>'
    )


def row_info_predefined(df: pd.DataFrame, idx: int, tag_col: str, val_col: str) -> str:
    if df.empty:
        return "*No data loaded.*"
    row = df.iloc[int(idx)]
    val = row.get(tag_col, "—")
    val = val if pd.notna(val) and str(val).strip() else "—"
    val_str = str(val)[:80] + ("…" if len(str(val)) > 80 else "")
    done = "✅ Validated" if row.get(val_col) is True else "⏳ Pending"
    return (
        f"**Current value:** `{val_str}` | {done} | "
        f"Title: {row.get('title', 'N/A')} | Cluster: {row.get('clustering_id', 'N/A')}"
    )


# ── Cluster helpers (predefined mode) ────────────────────────────────────────

def cluster_list(df: pd.DataFrame) -> list:
    if df.empty or "clustering_id" not in df.columns:
        return []
    all_c = df["clustering_id"].dropna().unique().tolist()
    return sorted([str(c) for c in all_c if str(c).lower() not in ("nan", "")])


def cluster_gallery_html(df: pd.DataFrame, cl_idx: int) -> str:
    """Build an HTML gallery for the cluster at position cl_idx.

    Images are rendered as <img src="..."> — no download — with a CSS grid
    layout (5 columns, auto rows).
    """
    clusters = cluster_list(df)
    if not clusters:
        return "<p><em>No clusters found.</em></p>"
    cl_idx = min(int(cl_idx), len(clusters) - 1)
    c_id = clusters[cl_idx]
    mask = df["clustering_id"].astype(str) == c_id
    sample = df[mask].sample(n=min(10, int(mask.sum())), random_state=42)

    cards = ""
    for _, r in sample.iterrows():
        url = str(r.get("image_url", "")).strip()
        if not url or url.lower() == "nan":
            continue
        title = str(r.get("title", "")) if pd.notna(r.get("title", "")) else ""
        thumb = iiif_resize(url, 300)
        cards += (
            f'<div style="text-align:center;">'
            f'<img src="{thumb}" loading="lazy" decoding="async" '
            f'style="width:100%;height:120px;object-fit:cover;'
            f'border-radius:6px;background:#eee;" title="{title}" />'
            f'<div style="font-size:11px;color:#666;margin-top:4px;overflow:hidden;'
            f'white-space:nowrap;text-overflow:ellipsis;">{title}</div>'
            f'</div>'
        )
    return (
        f'<div style="display:grid;grid-template-columns:repeat(5,1fr);gap:8px;">'
        f'{cards}</div>'
    ) if cards else "<p><em>No images available for this cluster.</em></p>"


def cluster_title(df: pd.DataFrame, cl_idx: int) -> str:
    clusters = cluster_list(df)
    if not clusters:
        return "No clusters found."
    cl_idx = min(int(cl_idx), len(clusters) - 1)
    c_id = clusters[cl_idx]
    n = int((df["clustering_id"].astype(str) == c_id).sum())
    return f"### Cluster **{c_id}** ({cl_idx + 1} / {len(clusters)}) — {n} images"


# ── Shared Mistral / image helpers ────────────────────────────────────────────

def _url_to_b64(url: str, width: int | None = None) -> str:
    """Fetch an image URL and return a base64-encoded JPEG string.

    If `width` is given and the URL is a IIIF request, the image is downscaled
    server-side (via iiif_resize) to keep the payload small — useful for the
    concept-scoring calls that run over many images. Raises on network/decode
    errors so callers can decide how to handle failures.
    """
    from PIL import Image
    src = iiif_resize(url, width) if width else url
    r = requests.get(src, timeout=15, headers=_HTTP_HEADERS)
    r.raise_for_status()
    img = Image.open(io.BytesIO(r.content))
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="JPEG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def _mistral_chat(api_key: str, model: str, content, json_mode: bool = False) -> str:
    """Call the Mistral chat API once and return the raw text of the reply.

    `content` may be a plain string (text-only) or a list of content parts
    (multimodal, e.g. text + image_url). When `json_mode` is set we request a
    JSON object, retrying without the flag if the SDK/model rejects it.
    """
    from mistralai import Mistral
    client = Mistral(api_key=api_key)
    msgs = [{"role": "user", "content": content}]
    try:
        if json_mode:
            res = client.chat.complete(
                model=model, messages=msgs,
                response_format={"type": "json_object"},
            )
        else:
            res = client.chat.complete(model=model, messages=msgs)
    except Exception:
        res = client.chat.complete(model=model, messages=msgs)
    return res.choices[0].message.content


def _robust_json(text: str):
    """Parse a JSON value out of an LLM reply, tolerating code fences / prose."""
    if not text:
        raise ValueError("empty response")
    t = re.sub(r"```(?:json)?|```", "", text.strip()).strip()
    try:
        return json.loads(t)
    except Exception:
        m = re.search(r"(\{.*\}|\[.*\])", t, flags=re.DOTALL)
        if m:
            return json.loads(m.group(1))
        raise


# ══════════════════════════════════════════════════════════════════════════════
# GUIDED CLUSTERING (ITGC) — Interpretable Text-Guided Image Clustering
# Reference: Zhao & Mac Aodha, "Interpretable Text-Guided Image Clustering via
# Iterative Search" (arXiv:2506.12514). An LLM proposes textual concepts from a
# user criterion, a Mistral vision model scores each image against those concepts to
# build a concept-space embedding, k-means clusters it, the silhouette score
# rates the partition, and the LLM refines the concepts over several iterations.
# Everything runs through the existing Mistral API + numpy — no extra deps.
# ══════════════════════════════════════════════════════════════════════════════

GC_N_CONCEPTS = 15  # target number of concepts proposed per iteration


def gc_generate_concepts(criterion: str, k: int, history: list,
                         api_key: str, text_model: str) -> list:
    """Ask the LLM for a set of short, visually-checkable concepts for `criterion`.

    `history` is a list of (concepts, silhouette) from previous iterations; it is
    fed back so the model can improve on what scored well (the iterative-search
    core of ITGC). Returns a de-duplicated list of concept strings.
    """
    hist_txt = ""
    if history:
        lines = [f"  Attempt {i} (silhouette={s:.3f}): {', '.join(cs)}"
                 for i, (cs, s) in enumerate(history, 1)]
        hist_txt = (
            "\n\nPrevious attempts and their clustering quality "
            "(silhouette, higher is better):\n" + "\n".join(lines) +
            "\n\nPropose an IMPROVED set: keep concepts that helped, drop weak or "
            "redundant ones, and add more discriminative visual attributes."
        )
    prompt = (
        f"You are helping cluster a set of heritage / archive images according to "
        f"this criterion:\n  \"{criterion}\"\n\n"
        f"Propose about {GC_N_CONCEPTS} short, concrete, VISUALLY-CHECKABLE concepts "
        f"(attributes an image either clearly shows or not) that best separate the "
        f"images along this criterion into roughly {k} groups. Each concept must be a "
        f"short noun phrase (2-5 words), in English, unambiguous from a single image."
        f"{hist_txt}\n\n"
        f"Return ONLY a JSON object of the form "
        f"{{\"concepts\": [\"concept one\", \"concept two\"]}}."
    )
    data = _robust_json(_mistral_chat(api_key, text_model, prompt, json_mode=True))
    concepts = data.get("concepts", []) if isinstance(data, dict) else data
    seen, out = set(), []
    for c in (concepts or []):
        c = str(c).strip()
        if c and c.lower() not in seen:
            seen.add(c.lower())
            out.append(c)
    return out[:GC_N_CONCEPTS] or ["text", "architecture", "person", "landscape", "object"]


def gc_score_image(url: str, concepts: list, api_key: str,
                   vision_model: str, cache: dict) -> np.ndarray:
    """Score one image 0-10 against each concept via a Mistral vision model
    (ITGC's VLM encoder).

    Uses a per-(url, concept) cache so concepts repeated across iterations are not
    re-scored. Returns a vector aligned to `concepts`; failed scores stay as 0.
    """
    scores = np.zeros(len(concepts), dtype=float)
    missing = [i for i, c in enumerate(concepts) if (url, c) not in cache]
    for i, c in enumerate(concepts):
        if (url, c) in cache:
            scores[i] = cache[(url, c)]
    if not missing:
        return scores
    try:
        b64 = _url_to_b64(url, width=512)
    except Exception:
        return scores  # image unreachable -> leave missing concepts at 0
    ask = [concepts[i] for i in missing]
    prompt = (
        "Look at the image. For EACH concept below, rate from 0 to 10 how strongly "
        "the concept is present in / applies to what is visible (0 = not at all, "
        "10 = clearly present). Judge only from visible content.\n"
        f"Concepts: {json.dumps(ask, ensure_ascii=False)}\n"
        "Return ONLY a JSON object mapping each concept EXACTLY to its integer score, "
        "e.g. {\"concept one\": 7, \"concept two\": 0}."
    )
    content = [
        {"type": "text", "text": prompt},
        {"type": "image_url", "image_url": f"data:image/jpeg;base64,{b64}"},
    ]
    try:
        data = _robust_json(_mistral_chat(api_key, vision_model, content, json_mode=True))
    except Exception:
        return scores
    if isinstance(data, dict):
        low = {str(kk).strip().lower(): vv for kk, vv in data.items()}
        for i in missing:
            try:
                v = float(low.get(concepts[i].strip().lower(), 0))
            except (TypeError, ValueError):
                v = 0.0
            v = max(0.0, min(10.0, v))
            scores[i] = v
            cache[(url, concepts[i])] = v
    return scores


def gc_encode(sample_df: pd.DataFrame, concepts: list, api_key: str,
              vision_model: str, cache: dict, max_workers: int = 12,
              on_done=None) -> np.ndarray:
    """Build the concept-space matrix Z (n_images × n_concepts) via parallel scoring.

    `on_done`, if given, is called (from the main thread) once per finished image —
    used to drive the progress bar.
    """
    urls = [str(u) for u in sample_df["image_url"].tolist()]
    rows = [None] * len(urls)
    with ThreadPoolExecutor(max_workers=max_workers) as ex:
        futs = {ex.submit(gc_score_image, u, concepts, api_key, vision_model, cache): i
                for i, u in enumerate(urls)}
        for fut in as_completed(futs):
            i = futs[fut]
            try:
                rows[i] = fut.result()
            except Exception:
                rows[i] = np.zeros(len(concepts), dtype=float)
            if on_done is not None:
                on_done()
    return np.vstack(rows) if rows else np.zeros((0, len(concepts)))


def gc_standardize(Z: np.ndarray) -> np.ndarray:
    """Z-score each concept column so k-means is not dominated by a single scale."""
    mu = Z.mean(axis=0)
    sd = Z.std(axis=0)
    sd = np.where(sd < 1e-8, 1.0, sd)
    return (Z - mu) / sd


def gc_kmeans(Z: np.ndarray, k: int, n_init: int = 5,
              max_iter: int = 100, seed: int = 42) -> np.ndarray:
    """Minimal k-means++ (numpy). Returns integer labels 0..k-1. Small-n only."""
    n = Z.shape[0]
    k = max(1, min(int(k), n))
    rng = np.random.default_rng(seed)
    best_labels, best_inertia = np.zeros(n, dtype=int), np.inf
    for _ in range(n_init):
        centers = [Z[rng.integers(n)]]
        for _ in range(1, k):
            d2 = np.min([np.sum((Z - c) ** 2, axis=1) for c in centers], axis=0)
            total = d2.sum()
            probs = d2 / total if total > 0 else None
            centers.append(Z[rng.choice(n, p=probs)])
        C = np.array(centers, dtype=float)
        labels = np.full(n, -1, dtype=int)
        for _it in range(max_iter):
            dists = np.linalg.norm(Z[:, None, :] - C[None, :, :], axis=2)
            new_labels = dists.argmin(axis=1)
            if np.array_equal(new_labels, labels):
                break
            labels = new_labels
            for j in range(k):
                pts = Z[labels == j]
                C[j] = pts.mean(axis=0) if len(pts) else Z[rng.integers(n)]
        inertia = float(np.sum((Z - C[labels]) ** 2))
        if inertia < best_inertia:
            best_inertia, best_labels = inertia, labels.copy()
    return best_labels


def gc_silhouette(Z: np.ndarray, labels: np.ndarray) -> float:
    """Mean silhouette score (numpy). Returns 0.0 for degenerate partitions."""
    n = len(labels)
    uniq = np.unique(labels)
    if len(uniq) < 2 or n <= len(uniq):
        return 0.0
    D = np.linalg.norm(Z[:, None, :] - Z[None, :, :], axis=2)
    sil = np.zeros(n)
    for i in range(n):
        same = labels == labels[i]
        same[i] = False
        a = D[i, same].mean() if same.any() else 0.0
        b = np.inf
        for lab in uniq:
            if lab == labels[i]:
                continue
            mask = labels == lab
            if mask.any():
                b = min(b, D[i, mask].mean())
        denom = max(a, b)
        sil[i] = 0.0 if denom == 0 else (b - a) / denom
    return float(sil.mean())


def gc_name_clusters(criterion: str, per_cluster: dict,
                     api_key: str, text_model: str) -> dict:
    """Ask the LLM for a short human-readable label per cluster (best-effort)."""
    desc = {cid: [t[0] for t in info["top"]] for cid, info in per_cluster.items()}
    prompt = (
        f"We clustered heritage images by: \"{criterion}\". Each cluster is described "
        f"by its most distinctive visual concepts. Give each cluster a SHORT label "
        f"(2-4 words, English).\n"
        f"Clusters and their top concepts: {json.dumps(desc, ensure_ascii=False)}\n"
        f"Return ONLY JSON mapping each cluster id to its label, e.g. "
        f"{{\"0\": \"illuminated initials\"}}."
    )
    try:
        data = _robust_json(_mistral_chat(api_key, text_model, prompt, json_mode=True))
        return {str(kk): str(vv) for kk, vv in data.items()}
    except Exception:
        return {cid: f"Cluster {cid}" for cid in per_cluster}


def _gc_per_cluster(Z: np.ndarray, labels: np.ndarray, concepts: list, top_n: int = 5) -> dict:
    """For each cluster, the concepts whose mean score is most ABOVE the global mean
    (i.e. what makes the cluster distinctive), plus its size."""
    gm = Z.mean(axis=0)
    per = {}
    for j in sorted({int(l) for l in labels}):
        mask = labels == j
        cm = Z[mask].mean(axis=0)
        diff = cm - gm
        order = np.argsort(diff)[::-1][:top_n]
        per[str(j)] = {"size": int(mask.sum()),
                       "top": [(concepts[i], float(cm[i]), float(diff[i])) for i in order]}
    return per


# ── Guided-clustering visualisation (HTML/SVG, no plotting dependency) ─────────

def gc_silhouette_svg(history: list, final_sil: float) -> str:
    """Hand-rolled SVG line chart of the silhouette score across iterations."""
    if not history:
        return ""
    vals = [s for _, s in history]
    W, H, pad = 440, 130, 26
    n = len(vals)
    lo, hi = min(vals + [0.0]), max(vals + [0.0001])
    span = (hi - lo) or 1.0
    fx = lambda i: pad + (W - 2 * pad) * (i / max(n - 1, 1))
    fy = lambda v: H - pad - (H - 2 * pad) * ((v - lo) / span)
    pts = " ".join(f"{fx(i):.1f},{fy(v):.1f}" for i, v in enumerate(vals))
    dots = "".join(f'<circle cx="{fx(i):.1f}" cy="{fy(v):.1f}" r="3" fill="#7c6f5e"/>'
                   for i, v in enumerate(vals))
    labs = "".join(f'<text x="{fx(i):.1f}" y="{fy(v) - 8:.1f}" font-size="9" '
                   f'text-anchor="middle" fill="#888">{v:.2f}</text>'
                   for i, v in enumerate(vals))
    return (
        '<div class="ia-section">Silhouette score across iterations '
        '(the iterative search at work)</div>'
        f'<svg viewBox="0 0 {W} {H}" width="100%" style="max-width:{W}px;'
        f'background:#faf9f7;border-radius:8px;border:1px solid #e5e2db;">'
        f'<polyline points="{pts}" fill="none" stroke="#7c6f5e" stroke-width="2"/>'
        f'{dots}{labs}'
        f'<text x="{pad}" y="{H - 6}" font-size="10" fill="#888">iter 1</text>'
        f'<text x="{W - pad}" y="{H - 6}" font-size="10" fill="#888" '
        f'text-anchor="end">iter {n}</text></svg>'
    )


def gc_clusters_viz_html(df: pd.DataFrame, per_cluster: dict, names: dict) -> str:
    """Per-cluster cards, image-first: a full-width thumbnail gallery of the actual
    images in the cluster, with the distinctive concepts shown as compact chips."""
    clusters = cluster_list(df)
    if not clusters:
        return "<p><em>No clusters yet — run guided clustering.</em></p>"
    blocks = []
    for pos, cid in enumerate(clusters):
        info = per_cluster.get(cid, {})
        name = names.get(cid, f"Cluster {cid}")
        size = info.get("size", int((df["clustering_id"].astype(str) == cid).sum()))
        chips = "".join(
            f'<span style="display:inline-block;background:#efece5;color:#5a5346;'
            f'border-radius:12px;padding:2px 9px;margin:2px 4px 2px 0;font-size:12px;">'
            f'{cname} <b style="color:#7c6f5e;">{cmean:.1f}</b></span>'
            for cname, cmean, _d in info.get("top", [])
        )
        chips_div = f'<div style="margin-bottom:10px;">{chips}</div>' if chips else ''
        gallery = cluster_gallery_html(df, pos)  # actual images, up to 10 per cluster
        blocks.append(
            '<div style="border:1px solid #e5e2db;border-radius:10px;padding:14px;'
            'margin-bottom:16px;">'
            f'<div style="font-weight:600;font-size:16px;margin-bottom:4px;">🗂️ {name} '
            f'<span style="color:#888;font-weight:400;">· cluster {cid} · '
            f'{size} images</span></div>'
            f'{chips_div}{gallery}</div>'
        )
    return "".join(blocks)


# ══════════════════════════════════════════════════════════════════════════════
# LOCAL BACKEND — SigLIP2 zero-shot encoder (CPU, fully offline, no API)
# Images are embedded ONCE (session cache), then classification (cosine to the
# category texts) and visual clustering (k-means on embeddings) are instant.
# ══════════════════════════════════════════════════════════════════════════════

# Override via the SIGLIP_MODEL env var (e.g. a Space secret) to trade speed for
# quality: "google/siglip2-large-patch16-256" or "google/siglip2-so400m-patch14-384"
# are markedly better but slower on CPU.
_SIGLIP_MODEL = os.environ.get("SIGLIP_MODEL", "google/siglip2-base-patch16-224")
_SIGLIP = {"model": None, "proc": None}
_GC_IMG_EMB = {}  # url -> np.ndarray (L2-normalised image embedding); per-session

# Zero-shot prompt ensemble: several phrasings per category, averaged in text-
# embedding space (standard CLIP/SigLIP practice — noticeably more robust than a
# single template, especially on archival / black-and-white photographs).
# SigLIP2 is multilingual, so a French template also helps French labels.
_GC_PROMPT_TEMPLATES = [
    "This is a photo of {}.",
    "a photo of {}",
    "an old photograph of {}",
    "a black and white archival photograph of {}",
    "une photo de {}",
]

_GC_ENUM_RE = re.compile(r"^\s*\d+(?:\.\d+)*[.)]?\s*")


def _gc_prompt_text(label: str, hints: dict | None = None) -> str:
    """Turn a user category label into the text used in prompts: prefer the
    user-supplied hint if any, strip leading enumeration ('2.1. Architecture'
    -> 'Architecture'), lowercase."""
    if hints and hints.get(label):
        return hints[label].strip().lower()
    txt = _GC_ENUM_RE.sub("", label).strip()
    return (txt or label).lower()


def _siglip_load():
    """Lazily load SigLIP2. First call downloads the weights from HF (~1 GB) and
    HF caches them on disk; subsequent calls (and Space restarts with a warm
    cache) are fast."""
    if _SIGLIP["model"] is None:
        from transformers import AutoModel, AutoProcessor
        _SIGLIP["proc"] = AutoProcessor.from_pretrained(_SIGLIP_MODEL, use_fast=True)
        _SIGLIP["model"] = AutoModel.from_pretrained(_SIGLIP_MODEL).eval()
    return _SIGLIP["model"], _SIGLIP["proc"]


def gc_siglip_text_emb(texts):
    """Return (embeddings (n,d) L2-normalised, logit_scale, logit_bias) for texts."""
    import torch
    model, proc = _siglip_load()
    # SigLIP expects a fixed 64-token context; an explicit max_length is required
    # (this checkpoint's tokenizer has no predefined one, so bare "max_length" no-ops).
    inp = proc(text=list(texts), padding="max_length", max_length=64,
               truncation=True, return_tensors="pt")
    with torch.no_grad():
        e = model.get_text_features(**inp)
    e = e / e.norm(dim=-1, keepdim=True)
    try:
        scale = float(model.logit_scale.detach().exp())
    except Exception:
        scale = 1.0
    try:
        bias = float(model.logit_bias.detach())
    except Exception:
        bias = 0.0
    return e.cpu().numpy(), scale, bias


def _gc_fetch_pil(url, width=384):
    """Fetch one image URL (IIIF-downscaled, with a browser-like UA) as an RGB
    PIL image, or None on any failure. Shared by all local encoders."""
    from PIL import Image
    try:
        src = iiif_resize(url, width) if url else url
        r = requests.get(src, timeout=15, headers=_HTTP_HEADERS)
        r.raise_for_status()
        return Image.open(io.BytesIO(r.content)).convert("RGB")
    except Exception:
        return None


def gc_siglip_embed_images(urls, on_batch=None, batch_size=16, fetch_workers=8):
    """Embed images with SigLIP2 (L2-normalised), fetched in parallel and encoded
    in batches, with a per-URL session cache. Returns a list aligned to `urls`;
    unreachable/failed images are None. `on_batch(n)` drives the progress bar."""
    import torch
    model, proc = _siglip_load()

    out = [None] * len(urls)
    for start in range(0, len(urls), batch_size):
        idxs = list(range(start, min(start + batch_size, len(urls))))
        todo = [i for i in idxs if urls[i] not in _GC_IMG_EMB]
        if todo:
            with ThreadPoolExecutor(max_workers=fetch_workers) as ex:
                imgs = list(ex.map(lambda i: _gc_fetch_pil(urls[i]), todo))
            valid = [(i, im) for i, im in zip(todo, imgs) if im is not None]
            if valid:
                inp = proc(images=[im for _, im in valid], return_tensors="pt")
                with torch.no_grad():
                    fe = model.get_image_features(**inp)
                fe = (fe / fe.norm(dim=-1, keepdim=True)).cpu().numpy()
                for (i, _im), e in zip(valid, fe):
                    _GC_IMG_EMB[urls[i]] = e
        for i in idxs:
            out[i] = _GC_IMG_EMB.get(urls[i])
        if on_batch:
            on_batch(len(idxs))
    return out


def gc_classify_run_local(df, categories, progress, hints=None):
    """Generator: classify every image into one fixed category with SigLIP2, fully
    offline. Yields the 7 gc outputs, with images appearing progressively."""
    total = len(df)
    urls = [str(u) for u in df["image_url"].tolist()]
    log = [f"**Local SigLIP2 classification** — {total} images → one of "
           f"{len(categories)} categories: {', '.join(categories)}. "
           f"Prompt-ensemble zero-shot, fully offline, no API."]
    progress(0.0, desc="Loading SigLIP2 (first run downloads the model)…")
    yield (df, {}, "\n".join(log), "", gr.update(), gr.update(), gr.update())

    # Text side: encode every (template × category) prompt and average per
    # category in embedding space (renormalised) — the prompt ensemble.
    texts = [_gc_prompt_text(c, hints) for c in categories]
    Ts, scale, bias = [], 1.0, 0.0
    for tpl in _GC_PROMPT_TEMPLATES:
        T_t, scale, bias = gc_siglip_text_emb([tpl.format(t) for t in texts])
        Ts.append(T_t)
    T = np.mean(Ts, axis=0)
    T = T / np.linalg.norm(T, axis=1, keepdims=True)
    assigned = {}
    done = [0]

    def _bar(n):
        done[0] += n
        progress(min(done[0] / max(total, 1), 0.99),
                 desc=f"Embedding & classifying… {done[0]}/{total}")

    CHUNK = 48
    for start in range(0, total, CHUNK):
        idxs = list(range(start, min(start + CHUNK, total)))
        embs = gc_siglip_embed_images([urls[i] for i in idxs], on_batch=_bar)
        for k, i in enumerate(idxs):
            e = embs[k]
            if e is None:
                continue
            logits = e @ T.T * scale + bias
            j = int(np.argmax(logits))
            # Softmax over the user's categories: a relative confidence that is
            # actually readable (SigLIP's raw sigmoid sits near 0 by design).
            p = np.exp(logits - logits.max())
            p = p / p.sum()
            assigned[i] = (categories[j], float(p[j]) * 10.0)
        out_df = _ensure_predefined_columns(df.copy())
        out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                                   for i in range(total)]
        per, names, md = _gc_classify_view(assigned, categories)
        last = start + CHUNK >= total
        yield ((out_df if last else df), ({"categories": categories} if last else {}),
               "\n".join(log + [f"- {len(assigned)}/{total} images classified…"]),
               "", md, gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))

    n_fail = total - len(assigned)
    log.append(f"\n✅ Done — {len(assigned)}/{total} classified"
               + (f", {n_fail} unreachable/skipped" if n_fail else "")
               + ". Use the **📊 Export** tab.")
    progress(1.0, desc="Done")
    out_df = _ensure_predefined_columns(df.copy())
    out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                               for i in range(total)]
    per, names, md = _gc_classify_view(assigned, categories)
    yield (out_df, {"categories": categories}, "\n".join(log), "", md,
           gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))


def gc_discover_run_local(df, k, progress, embed_fn=None, model_label="SigLIP2"):
    """Generator: fully-offline visual clustering — k-means on local image
    embeddings (no LLM, no concepts). Clusters are shown as image galleries.
    `embed_fn` selects the encoder (SigLIP2 by default, or DINOv2)."""
    embed_fn = embed_fn or gc_siglip_embed_images
    total = len(df)
    urls = [str(u) for u in df["image_url"].tolist()]
    k = max(2, min(int(k), total))
    log = [f"**Local {model_label} clustering** — {total} images into k={k} visual "
           f"groups (k-means on embeddings). Fully offline, no API."]
    progress(0.0, desc=f"Loading {model_label} (first run downloads the model)…")
    yield (df, {}, "\n".join(log), "", gr.update(), gr.update(), gr.update())

    done = [0]

    def _bar(n):
        done[0] += n
        progress(min(done[0] / max(total, 1), 0.98),
                 desc=f"Embedding images… {done[0]}/{total}")

    embs = embed_fn(urls, on_batch=_bar)
    valid = [(i, e) for i, e in enumerate(embs) if e is not None]
    if len(valid) < 2:
        yield (df, {}, "\n".join(log + ["⚠️ Not enough images could be embedded."]),
               "", gr.update(), gr.update(), gr.update())
        return
    progress(0.99, desc="Clustering…")
    E = np.vstack([e for _, e in valid])
    labels = gc_kmeans(E, k)  # embeddings are normalised, so euclidean ≈ cosine

    out_df = _ensure_predefined_columns(df.copy())
    cid = ["" for _ in range(total)]
    for (i, _e), lab in zip(valid, labels):
        cid[i] = str(int(lab))
    out_df["clustering_id"] = cid
    per = {str(j): {"size": int((labels == j).sum()), "top": []}
           for j in sorted({int(l) for l in labels})}
    names = {c: f"Cluster {c}" for c in per}
    md = (f"### Visual clusters ({model_label})\n{len(per)} clusters · "
          f"{len(valid)}/{total} images embedded")
    progress(1.0, desc="Done")
    yield (out_df, {"kind": "local_visual"}, "\n".join(log + ["\n✅ Done."]),
           "", md, gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))


# ══════════════════════════════════════════════════════════════════════════════
# LOCAL BACKEND — DINOv2 (self-supervised ViT) for visual similarity & few-shot
# DINOv2 image features are the state of the art for visual similarity/clustering
# (better than supervised CNNs like EfficientNet, and than image-text models like
# CLIP/SigLIP for this purpose). No text branch: named categories are obtained by
# few-shot — nearest labelled-centroid — which beats zero-shot on niche domains.
# ══════════════════════════════════════════════════════════════════════════════

_DINO_MODEL = os.environ.get("DINOV2_MODEL", "facebook/dinov2-with-registers-base")
_DINO = {"model": None, "proc": None}
_GC_DINO_EMB = {}  # url -> np.ndarray (L2-normalised DINOv2 embedding); per-session


def _dino_load():
    if _DINO["model"] is None:
        from transformers import AutoModel, AutoImageProcessor
        _DINO["proc"] = AutoImageProcessor.from_pretrained(_DINO_MODEL, use_fast=True)
        _DINO["model"] = AutoModel.from_pretrained(_DINO_MODEL).eval()
    return _DINO["model"], _DINO["proc"]


def gc_dino_embed_images(urls, on_batch=None, batch_size=16, fetch_workers=8):
    """Embed images with DINOv2 (L2-normalised CLS descriptor), batched + cached.
    Same contract as gc_siglip_embed_images."""
    import torch
    model, proc = _dino_load()

    out = [None] * len(urls)
    for start in range(0, len(urls), batch_size):
        idxs = list(range(start, min(start + batch_size, len(urls))))
        todo = [i for i in idxs if urls[i] not in _GC_DINO_EMB]
        if todo:
            with ThreadPoolExecutor(max_workers=fetch_workers) as ex:
                imgs = list(ex.map(lambda i: _gc_fetch_pil(urls[i]), todo))
            valid = [(i, im) for i, im in zip(todo, imgs) if im is not None]
            if valid:
                inp = proc(images=[im for _, im in valid], return_tensors="pt")
                with torch.no_grad():
                    o = model(**inp)
                pooled = getattr(o, "pooler_output", None)
                fe = pooled if pooled is not None else o.last_hidden_state[:, 0]
                fe = (fe / fe.norm(dim=-1, keepdim=True)).cpu().numpy()
                for (i, _im), e in zip(valid, fe):
                    _GC_DINO_EMB[urls[i]] = e
        for i in idxs:
            out[i] = _GC_DINO_EMB.get(urls[i])
        if on_batch:
            on_batch(len(idxs))
    return out


def gc_fewshot_run_local(df, labeled, progress):
    """Generator: few-shot classification. `labeled` = list of (url, category).
    Build one L2-normalised DINOv2 centroid per category from the examples, then
    assign every dataset image to its nearest centroid (cosine). Fully offline."""
    total = len(df)
    cats0 = sorted({lab for _, lab in labeled})
    log = [f"**Local DINOv2 few-shot classification** — {total} images into "
           f"{len(cats0)} categories from {len(labeled)} labelled examples: "
           f"{', '.join(cats0)}. Fully offline, no API."]
    progress(0.0, desc="Loading DINOv2 (first run downloads the model)…")
    yield (df, {}, "\n".join(log), "", gr.update(), gr.update(), gr.update())

    # 1) Embed the labelled examples and build per-category centroids.
    lab_urls = [str(u) for u, _ in labeled]
    lab_embs = gc_dino_embed_images(
        lab_urls, on_batch=lambda n: progress(0.05, desc="Embedding labelled examples…"))
    proto = {}
    for (u, lab), e in zip(labeled, lab_embs):
        if e is not None:
            proto.setdefault(lab, []).append(e)
    cats = [c for c in cats0 if proto.get(c)]
    if len(cats) < 2:
        yield (df, {}, "\n".join(log + ["⚠️ Need ≥2 categories with reachable "
               "labelled examples — check the example URLs."]),
               "", gr.update(), gr.update(), gr.update())
        return
    C = np.vstack([np.mean(proto[c], axis=0) for c in cats])
    C = C / np.linalg.norm(C, axis=1, keepdims=True)  # (K, d) normalised centroids
    n_ex = {c: len(proto[c]) for c in cats}
    log.append("- prototypes: " + ", ".join(f"{c}×{n_ex[c]}" for c in cats))

    # 2) Embed the dataset and assign each image to its nearest centroid.
    urls = [str(u) for u in df["image_url"].tolist()]
    assigned, done = {}, [0]

    def _bar(n):
        done[0] += n
        progress(min(0.1 + 0.89 * done[0] / max(total, 1), 0.99),
                 desc=f"Embedding & classifying… {done[0]}/{total}")

    CHUNK = 48
    for start in range(0, total, CHUNK):
        idxs = list(range(start, min(start + CHUNK, total)))
        embs = gc_dino_embed_images([urls[i] for i in idxs], on_batch=_bar)
        for k, i in enumerate(idxs):
            e = embs[k]
            if e is None:
                continue
            sims = C @ e  # cosine (both normalised)
            j = int(np.argmax(sims))
            p = np.exp((sims - sims.max()) / 0.1)  # temperature-softmax for a readable score
            p = p / p.sum()
            assigned[i] = (cats[j], float(p[j]) * 10.0)
        out_df = _ensure_predefined_columns(df.copy())
        out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                                   for i in range(total)]
        per, names, md = _gc_classify_view(assigned, cats)
        last = start + CHUNK >= total
        yield ((out_df if last else df), ({"categories": cats} if last else {}),
               "\n".join(log + [f"- {len(assigned)}/{total} images classified…"]),
               "", md, gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))

    n_fail = total - len(assigned)
    log.append(f"\n✅ Done — {len(assigned)}/{total} classified"
               + (f", {n_fail} unreachable/skipped" if n_fail else "")
               + ". Use the **📊 Export** tab.")
    progress(1.0, desc="Done")
    out_df = _ensure_predefined_columns(df.copy())
    out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                               for i in range(total)]
    per, names, md = _gc_classify_view(assigned, cats)
    yield (out_df, {"categories": cats}, "\n".join(log), "", md,
           gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))


def _gc_load_labeled(path):
    """Parse a labelled-examples file (CSV/JSON): rows with image_url + a label
    column (label/category/tag_2/class/clustering_id, else the 2nd column).
    Returns a list of (url, label)."""
    ext = os.path.splitext(path)[1].lower()
    if ext == ".json":
        data = json.load(open(path, encoding="utf-8"))
        rows = data.get("records", data) if isinstance(data, dict) else data
        rows = [r for r in rows if isinstance(r, dict)]
    else:
        with open(path, "r", encoding="utf-8") as f:
            sep = ";" if ";" in f.readline() else ","
        rows = pd.read_csv(path, sep=sep, encoding="utf-8",
                           on_bad_lines="warn").to_dict("records")
    label_keys = ["label", "category", "categorie", "catégorie", "tag_2", "tag2",
                  "class", "clustering_id"]
    out = []
    for r in rows:
        r = {str(k).strip(): v for k, v in r.items()}
        url = str(r.get("image_url", "")).strip()
        lab = ""
        for k in label_keys:
            if k in r and str(r[k]).strip() and str(r[k]).strip().lower() != "nan":
                lab = str(r[k]).strip(); break
        if not lab:
            rest = [v for k, v in r.items() if k != "image_url"]
            lab = str(rest[0]).strip() if rest else ""
        if url and url.lower() != "nan" and lab and lab.lower() != "nan":
            out.append((url, lab))
    return out


def gc_load_labeled(file_obj):
    """Gradio handler: parse a labelled-examples file → (state, status markdown)."""
    if file_obj is None:
        return [], "*No labelled examples loaded.*"
    try:
        path = file_obj.name if hasattr(file_obj, "name") else file_obj
        labeled = _gc_load_labeled(path)
    except Exception as e:
        return [], f"❌ Could not read labelled file: {e}"
    if not labeled:
        return [], "⚠️ No usable rows found (need an `image_url` column + a label column)."
    counts = {}
    for _u, lab in labeled:
        counts[lab] = counts.get(lab, 0) + 1
    summary = ", ".join(f"{c}×{n}" for c, n in sorted(counts.items()))
    warn = "" if len(counts) >= 2 else "  ⚠️ need ≥2 categories for few-shot"
    return labeled, f"✅ {len(labeled)} examples · {len(counts)} categories: {summary}{warn}"


def gc_classify_image(url: str, categories: list, api_key: str,
                      vision_model: str, cache: dict) -> tuple:
    """Ask a Mistral vision model to put one image in EXACTLY ONE of `categories`.

    Returns (category, confidence 0-10). `category` is normalised to the exact
    provided text (case-insensitive match) or 'Unclassified'. Cached per
    (url, categories) so re-runs are instant.
    """
    key = (url, tuple(categories))
    if key in cache:
        return cache[key]
    try:
        b64 = _url_to_b64(url, width=512)
    except Exception:
        return (None, 0.0)  # unreachable image -> caller skips it
    cats_txt = "\n".join(f"- {c}" for c in categories)
    prompt = (
        "Classify the image into EXACTLY ONE of the following categories, choosing "
        "the single best match based only on what is visible:\n"
        f"{cats_txt}\n\n"
        "Return ONLY a JSON object: "
        "{\"category\": \"<exact category text>\", \"confidence\": <integer 0-10>}. "
        "If none clearly applies, use \"category\": \"Unclassified\"."
    )
    content = [
        {"type": "text", "text": prompt},
        {"type": "image_url", "image_url": f"data:image/jpeg;base64,{b64}"},
    ]
    try:
        data = _robust_json(_mistral_chat(api_key, vision_model, content, json_mode=True))
        raw = str(data.get("category", "")).strip()
        conf = float(data.get("confidence", 0) or 0)
    except Exception:
        return (None, 0.0)
    match = next((c for c in categories if c.strip().lower() == raw.lower()), None)
    result = (match or "Unclassified", max(0.0, min(10.0, conf)))
    cache[key] = result
    return result


def _gc_classify_view(assigned: dict, categories: list) -> tuple:
    """Build (per_cluster, names, concepts_md) for the classification results so far."""
    by_cat = {}
    for _i, (cat, conf) in assigned.items():
        by_cat.setdefault(cat, []).append(conf)
    order = [c for c in categories if c in by_cat] + \
            [c for c in by_cat if c not in categories]  # e.g. 'Unclassified' last-ish
    per, names = {}, {}
    for cat in order:
        confs = by_cat.get(cat, [])
        avg = sum(confs) / len(confs) if confs else 0.0
        per[cat] = {"size": len(confs), "top": [("avg confidence", avg, 0.0)]}
        names[cat] = cat
    counts = " · ".join(f"**{c}**: {len(by_cat.get(c, []))}" for c in order)
    md = ("### Classification into fixed categories\n" + counts +
          f"\n\n{len(assigned)} images classified · {len(categories)} categories")
    return per, names, md


def gc_classify_run(df, categories, api_key, vmodel, progress):
    """Generator: classify every image into one fixed category, with a progress bar
    and image results that fill in progressively. Yields the 7 gc outputs."""
    cache, assigned = {}, {}
    total = len(df)
    urls = [str(u) for u in df["image_url"].tolist()]
    log = [f"**Guided classification** — {total} images → one of {len(categories)} "
           f"categories: {', '.join(categories)}.\n"]
    progress(0.0, desc="Starting…")
    yield (df, {}, "\n".join(log), "", gr.update(), gr.update(), gr.update())

    CHUNK = 40
    done = 0
    for start in range(0, total, CHUNK):
        idxs = list(range(start, min(start + CHUNK, total)))
        with ThreadPoolExecutor(max_workers=12) as ex:
            futs = {ex.submit(gc_classify_image, urls[i], categories, api_key, vmodel, cache): i
                    for i in idxs}
            for fut in as_completed(futs):
                i = futs[fut]
                cat, conf = fut.result()
                if cat is not None:
                    assigned[i] = (cat, conf)
                done += 1
                progress(min(done / max(total, 1), 0.99),
                         desc=f"Classifying images… {done}/{total}")
        out_df = _ensure_predefined_columns(df.copy())
        out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                                   for i in range(total)]
        per, names, md = _gc_classify_view(assigned, categories)
        log_now = "\n".join(log + [f"- Classified {done}/{total} images…"])
        last = start + CHUNK >= total
        meta = ({"categories": categories, "assigned_counts": {c: n for c, n in
                 [(k2, sum(1 for v in assigned.values() if v[0] == k2)) for k2 in names]}}
                if last else {})
        yield ((out_df if last else df), meta, log_now, "", md,
               gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))

    n_unc = sum(1 for v in assigned.values() if v[0] == "Unclassified")
    n_fail = total - len(assigned)
    log.append(f"\n✅ Done — {len(assigned)}/{total} classified"
               + (f", {n_unc} unclassified" if n_unc else "")
               + (f", {n_fail} unreachable/skipped" if n_fail else "")
               + ". Use the **📊 Export** tab.")
    progress(1.0, desc="Done")
    out_df = _ensure_predefined_columns(df.copy())
    out_df["clustering_id"] = [assigned[i][0] if i in assigned else ""
                               for i in range(total)]
    per, names, md = _gc_classify_view(assigned, categories)
    yield (out_df, {"categories": categories}, "\n".join(log), "", md,
           gc_clusters_viz_html(out_df, per, names), tbl_preview(out_df))


GC_MAX_LABEL_BTNS = 12  # max category buttons in the interactive labeller


def _gc_parse_categories(text):
    """Parse the categories textbox -> (ordered unique labels, {label: hint}).
    A line may be 'Label | optional description' for a richer prompt."""
    seen, cats, hints = set(), [], {}
    for line in str(text or "").splitlines():
        raw = line.strip().lstrip("-•*").strip()
        if not raw:
            continue
        label, _, hint = raw.partition("|")
        label, hint = label.strip(), hint.strip()
        if label and label.lower() not in seen:
            seen.add(label.lower())
            cats.append(label)
            if hint:
                hints[label] = hint
    return cats, hints


def _gc_labeled_counts_md(labeled):
    """Running summary of the interactively/loaded labelled set."""
    if not labeled:
        return "*No examples yet — click a category for each image (~10-15 per category).*"
    counts = {}
    for _u, lab in labeled:
        counts[lab] = counts.get(lab, 0) + 1
    body = " · ".join(f"**{c}** ×{n}" for c, n in sorted(counts.items()))
    return f"🏷️ {body}  —  total {len(labeled)} examples"


def _gc_lbl_view(df, idx):
    """(image HTML, 'Image i / N' caption) for the interactive labeller."""
    if df is None or df.empty:
        return "<p><em>Load an image list first (📁 Data tab).</em></p>", "—"
    n = len(df)
    idx = max(0, min(int(idx), n - 1))
    return img_html(img_url_at(df, idx)), f"Image **{idx + 1} / {n}**"


def _gc_lbl_prepare(df, categories_text, labeled):
    """Read the ① categories, label the buttons, and show the first image.
    Returns [cats, idx, img, info, status] + one update per label button."""
    cats, _h = _gc_parse_categories(categories_text)
    btns = [gr.update(value=f"🏷️ {cats[i]}", visible=True) if i < len(cats)
            else gr.update(visible=False) for i in range(GC_MAX_LABEL_BTNS)]
    if df is None or df.empty:
        img, info = "<p><em>Load an image list first (📁 Data tab).</em></p>", "—"
    else:
        img, info = _gc_lbl_view(df, 0)
    status = (_gc_labeled_counts_md(labeled) if cats
              else "⚠️ Type your categories in ① first, then press Prepare.")
    return [cats, 0, img, info, status] + btns


def _gc_lbl_nav(df, idx, delta):
    n = max(len(df), 1) if df is not None else 1
    nidx = max(0, min(int(idx) + int(delta), n - 1))
    img, info = _gc_lbl_view(df, nidx)
    return nidx, img, info


def _gc_lbl_undo(labeled):
    labeled = list(labeled or [])
    if labeled:
        labeled.pop()
    return labeled, _gc_labeled_counts_md(labeled)


def _gc_lbl_assign(i, df, idx, labeled, cats):
    """Assign the current image to category `cats[i]` and advance. Instant — no
    embedding here; DINOv2 runs only at Run time."""
    labeled = list(labeled or [])
    cats = cats or []
    if df is not None and not df.empty and 0 <= i < len(cats):
        url = str(img_url_at(df, idx))
        if url and url.lower() != "nan":
            labeled.append((url, cats[i]))
    n = max(len(df), 1) if df is not None else 1
    nidx = min(int(idx) + 1, n - 1)
    img, info = _gc_lbl_view(df, nidx)
    return labeled, nidx, img, info, _gc_labeled_counts_md(labeled)


def gc_run(df, criterion, categories, labeled, backend, k, n_iters, sample_n,
           api_key, vmodel, tmodel, progress=gr.Progress()):
    """Generator: run guided clustering. Dispatches on the chosen encoding backend,
    on whether labelled examples were supplied, and on whether fixed categories
    were typed:

      • DINOv2 + labelled examples  → few-shot nearest-centroid classification (no key)
      • DINOv2 + none               → offline visual k-means similarity clustering
      • SigLIP2 + categories        → offline zero-shot classification (no key)
      • SigLIP2 + none              → offline visual k-means (no key)
      • Mistral API + categories    → per-image VLM classification
      • Mistral API + none          → the full iterative ITGC discovery loop

    Yields tuples matching [gc_df_s, gc_meta_s, gc_log, gc_sil_html,
    gc_concepts_md, gc_clusters_html, gc_tbl].
    """
    def _warn(m):
        return (df, {}, m, gr.update(), gr.update(), gr.update(), gr.update())

    if df is None or df.empty:
        yield _warn("⚠️ Load an image list first (📁 Data tab)."); return

    # Parse the optional fixed-category list (label | optional description).
    cats, hints = _gc_parse_categories(categories)

    bl = str(backend).lower()

    # ── Local DINOv2 backend: few-shot (with labelled examples) or similarity ─
    if "dino" in bl:
        labeled = labeled or []
        if len(labeled) >= 2 and len({lab for _, lab in labeled}) >= 2:
            yield from gc_fewshot_run_local(df, labeled, progress)
        else:
            yield from gc_discover_run_local(df, k, progress,
                                             embed_fn=gc_dino_embed_images,
                                             model_label="DINOv2")
        return

    # ── Local SigLIP2 backend: fully offline, no API key required ─────────────
    if "siglip" in bl:
        if len(cats) >= 2:
            yield from gc_classify_run_local(df, cats, progress, hints)
        else:
            yield from gc_discover_run_local(df, k, progress)
        return

    # ── Mistral API backend ───────────────────────────────────────────────────
    if not api_key:
        yield _warn("⚠️ Enter your Mistral API key, or switch the backend to "
                    "**Local SigLIP2** (no key needed) in the 📁 Data tab.")
        return
    # Fail fast with a readable message if the key or model id is bad (e.g. a
    # retired model like pixtral-12b-2409) — otherwise per-image errors are
    # swallowed and every image just looks "unreachable/skipped".
    try:
        _mistral_chat(api_key, vmodel, "Reply with the single word: ok")
    except Exception as e:
        yield _warn(f"❌ Mistral API check failed for model `{vmodel}`: {e}\n\n"
                    "The model id may be retired or the key invalid — pick another "
                    "vision model in the 📁 Data tab, or switch to the **Local "
                    "SigLIP2** backend (no API needed).")
        return
    if len(cats) >= 2:
        yield from gc_classify_run(df, cats, api_key, vmodel, progress)
        return
    if not str(criterion).strip():
        yield _warn("⚠️ Enter a list of fixed categories (≥2), or a discovery criterion.")
        return

    k, n_iters, sample_n = int(k), int(n_iters), int(sample_n)
    if len(df) <= sample_n:
        work = df.reset_index(drop=True)
    else:
        work = df.sample(sample_n, random_state=42).reset_index(drop=True)
    k = max(2, min(k, len(work)))
    history, cache, best = [], {}, None
    no_improve = 0

    # Progress accounting: one tick per image scored across all iterations + the
    # final full-set pass. The vision model scores ALL concepts of an image in one call, so
    # cost scales with the number of images, not the number of concepts.
    total_ticks = n_iters * len(work) + len(df)
    done = 0

    def _tick(phase):
        nonlocal done
        done += 1
        progress(min(done / max(total_ticks, 1), 0.99),
                 desc=f"{phase} · {done}/{total_ticks} images scored")

    def _render(b):
        """Build (sil_svg, concepts_md, clusters_html, tbl) from the best-so-far run
        over the sampled `work` set — this is what makes images show up early."""
        sub = _ensure_predefined_columns(work.copy())
        sub["clustering_id"] = [str(int(l)) for l in b["labels"]]
        per = _gc_per_cluster(b["Z"], b["labels"], b["concepts"])
        names = b.get("names") or {cid: f"Cluster {cid}" for cid in per}
        md = ("### Concepts (best so far)\n" + ", ".join(f"`{c}`" for c in b["concepts"]) +
              f"\n\n**Silhouette: {b['sil']:.3f}** · {len(per)} clusters · "
              f"{len(sub)} images shown")
        return (gc_silhouette_svg(history, b["sil"]), md,
                gc_clusters_viz_html(sub, per, names), tbl_preview(sub))

    log = [f"**Guided clustering** — {len(work)} images · k={k} · up to {n_iters} "
           f"iterations. Clusters appear below after the first iteration."]
    progress(0.0, desc="Starting…")
    yield _warn("\n".join(log))

    for it in range(n_iters):
        progress(done / max(total_ticks, 1),
                 desc=f"Iteration {it + 1}/{n_iters}: generating concepts…")
        log.append(f"- Iteration {it + 1}/{n_iters}: generating concepts…")
        yield _warn("\n".join(log))
        try:
            concepts = gc_generate_concepts(criterion, k, history, api_key, tmodel)
        except Exception as e:
            log.append(f"  ⚠️ concept generation failed: {e}")
            yield _warn("\n".join(log))
            if best is None:
                return
            break
        log[-1] = (f"- Iteration {it + 1}/{n_iters}: scoring {len(work)} images "
                   f"on {len(concepts)} concepts…")
        yield _warn("\n".join(log))
        try:
            Z = gc_encode(work, concepts, api_key, vmodel, cache,
                          on_done=lambda: _tick(f"Iteration {it + 1}/{n_iters}"))
        except Exception as e:
            log.append(f"  ⚠️ scoring failed: {e}")
            yield _warn("\n".join(log))
            if best is None:
                return
            break
        Zs = gc_standardize(Z)
        labels = gc_kmeans(Zs, k)
        sil = gc_silhouette(Zs, labels)
        history.append((concepts, sil))
        improved = best is None or sil > best["sil"]
        if improved:
            best = {"labels": labels, "concepts": concepts, "sil": sil, "Z": Z}
            no_improve = 0
        else:
            no_improve += 1
        log[-1] = (f"- Iteration {it + 1}/{n_iters}: silhouette = **{sil:.3f}**"
                   f"{'  ⭐ best' if improved else ''}")
        # Progressive image results after every iteration (best-so-far).
        yield (df, {}, "\n".join(log), *_render(best))
        if no_improve >= 2:
            log.append("- Converged (no improvement for 2 iterations) — stopping early.")
            yield (df, {}, "\n".join(log), *_render(best))
            break

    # ── Finalize on the best concept set: label EVERY image in the dataset ────
    concepts = best["concepts"]
    log.append(f"\n**Best silhouette = {best['sil']:.3f}**. Assigning all "
               f"{len(df)} images…")
    progress(done / max(total_ticks, 1), desc=f"Finalizing: scoring all {len(df)} images…")
    yield (df, {}, "\n".join(log), *_render(best))

    Zf = gc_encode(df, concepts, api_key, vmodel, cache,  # cached rows are instant
                   on_done=lambda: _tick("Finalizing"))
    Zfs = gc_standardize(Zf)
    labels_full = best["labels"] if len(df) == len(work) else gc_kmeans(Zfs, k)
    sil_full = gc_silhouette(Zfs, labels_full)

    out_df = _ensure_predefined_columns(df.copy())
    out_df["clustering_id"] = [str(int(l)) for l in labels_full]
    per_cluster = _gc_per_cluster(Zf, labels_full, concepts)
    progress(0.99, desc="Naming clusters…")
    names = gc_name_clusters(criterion, per_cluster, api_key, tmodel)

    meta = {"concepts": concepts, "history": history, "per_cluster": per_cluster,
            "names": names, "sil": sil_full, "criterion": criterion}
    concepts_md = (
        "### Final concepts\n" + ", ".join(f"`{c}`" for c in concepts) +
        f"\n\n**Silhouette (full set): {sil_full:.3f}** · "
        f"{len(per_cluster)} clusters · {len(out_df)} images"
    )
    log.append("\n✅ Done. Clusters below cover the full set — use the **📊 Export** tab.")
    progress(1.0, desc="Done")
    yield (out_df, meta, "\n".join(log),
           gc_silhouette_svg(history, sil_full), concepts_md,
           gc_clusters_viz_html(out_df, per_cluster, names), tbl_preview(out_df))


def gc_load_urls(file_obj):
    """Load a URL-list file for the guided-clustering mode."""
    if file_obj is None:
        return gr.update(), gr.update(), "⚠️ Select a URL list file first."
    try:
        path = file_obj.name if hasattr(file_obj, "name") else file_obj
        urls = load_url_list(path)
        if not urls:
            return gr.update(), gr.update(), "⚠️ No URLs found in the file."
        df = _ensure_predefined_columns(pd.DataFrame({"image_url": urls}))
    except Exception as e:
        return gr.update(), gr.update(), f"❌ Error: {e}"
    return df, tbl_preview(df), (f"✅ Loaded {len(df)} images. Go to the "
                                 f"**🧭 Cluster** tab to run guided clustering.")


def do_gc_export_csv(df):
    if df is None or df.empty:
        return gr.update(visible=False), "⚠️ Nothing to export — run clustering first."
    return gr.update(value=export_as_csv(df), visible=True), ""


def do_gc_export_json(df):
    if df is None or df.empty:
        return gr.update(visible=False), "⚠️ Nothing to export — run clustering first."
    return gr.update(value=export_as_json(df), visible=True), ""


# ── Mistral caption helper ────────────────────────────────────────────────────

def generate_mistral_caption(url: str, api_key: str, model: str) -> tuple:
    """Call the Mistral multimodal API and return (caption_text, status_message)."""
    if not api_key:
        return "", "⚠️ Please enter your Mistral API key."
    if not url:
        return "", "⚠️ No image URL for this row."
    try:
        # Encode the image URL as a base64 JPEG for the Mistral API
        b64 = _url_to_b64(url)
    except Exception as e:
        return "", f"❌ Could not fetch image: {e}"
    try:
        from mistralai import Mistral
        client = Mistral(api_key=api_key)
        res = client.chat.complete(
            model=model,
            messages=[{"role": "user", "content": [
                {"type": "text", "text": PROMPT_MISTRAL},
                {"type": "image_url", "image_url": f"data:image/jpeg;base64,{b64}"},
            ]}],
        )
        return res.choices[0].message.content.strip(), "✅ Caption generated."
    except Exception as e:
        return "", f"❌ Mistral error: {e}"


# ══════════════════════════════════════════════════════════════════════════════
# INITIAL STATE — load a default dataset if configured via env vars
# ══════════════════════════════════════════════════════════════════════════════

def _read_predefined_csv(path: str) -> pd.DataFrame:
    """Read a predefined-schema CSV file and ensure its annotation columns."""
    with open(path, "r", encoding="utf-8") as f:
        sep = ";" if ";" in f.readline() else ","
    df = pd.read_csv(path, sep=sep, quotechar='"',
                     on_bad_lines="warn", encoding="utf-8")
    return _ensure_predefined_columns(df)


def _load_default_df():
    """Load the example dataset.

    Prefers the HuggingFace-hosted copy when DATASET_REPO and HF_TOKEN secrets
    are set; otherwise (or if the download fails) falls back to the bundled
    exemple.csv shipped alongside this script. Returns an empty DataFrame only
    when neither source is available.
    """
    if DATASET_REPO and HF_TOKEN:
        try:
            path = hf_hub_download(
                repo_id=DATASET_REPO, filename=CSV_FILENAME,
                repo_type="dataset", token=HF_TOKEN,
            )
            return _read_predefined_csv(path)
        except Exception:
            pass  # fall back to the bundled file below
    try:
        local_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), CSV_FILENAME)
        if os.path.exists(local_path):
            return _read_predefined_csv(local_path)
    except Exception:
        pass
    return pd.DataFrame()


_default_df = _load_default_df()

# ══════════════════════════════════════════════════════════════════════════════
# GRADIO UI
# ══════════════════════════════════════════════════════════════════════════════

CSS = """
/* ══ Dark-mode-safe color tokens ══════════════════════════════════════════ */
:root {
    --ia-bg-subtle:   #f7f6f3;
    --ia-border:      #e5e2db;
    --ia-text:        #2c2825;
    --ia-text-muted:  #6b6560;
    --ia-accent:      #7c6f5e;
    --ia-card-bg:     #ffffff;
    --ia-section-bg:  #f0ede8;
    --ia-section-hd:  #3a3330;
}
@media (prefers-color-scheme: dark) {
    :root {
        --ia-bg-subtle:   #2a2825;
        --ia-border:      #4a4640;
        --ia-text:        #e8e4df;
        --ia-text-muted:  #a09890;
        --ia-accent:      #c4a882;
        --ia-card-bg:     #332f2c;
        --ia-section-bg:  #302d2a;
        --ia-section-hd:  #e8e4df;
    }
}
/* Gradio also injects .dark class on <body> when dark mode is toggled */
.dark {
    --ia-bg-subtle:   #2a2825;
    --ia-border:      #4a4640;
    --ia-text:        #e8e4df;
    --ia-text-muted:  #a09890;
    --ia-accent:      #c4a882;
    --ia-card-bg:     #332f2c;
    --ia-section-bg:  #302d2a;
    --ia-section-hd:  #e8e4df;
}

/* ══ Home layout ══════════════════════════════════════════════════════════ */
.home-guidelines {
    background: var(--ia-bg-subtle);
    border-right: 1px solid var(--ia-border);
    padding: 24px 20px;
    border-radius: 12px 0 0 12px;
    min-height: 600px;
    color: var(--ia-text);
}
.home-guidelines p, .home-guidelines li, .home-guidelines td, .home-guidelines th {
    color: var(--ia-text) !important;
}
.home-modes { padding: 24px 20px; }

/* ══ Mode cards ═══════════════════════════════════════════════════════════ */
.mode-card {
    border: 2px solid var(--ia-border);
    border-radius: 12px;
    padding: 28px 24px;
    margin-bottom: 20px;
    background: var(--ia-card-bg);
    transition: border-color 0.2s, box-shadow 0.2s;
}
.mode-card:hover {
    border-color: var(--ia-accent);
    box-shadow: 0 4px 16px rgba(0,0,0,0.12);
}
.mode-card h3 { margin: 0 0 8px 0; font-size: 18px; color: var(--ia-text) !important; }
.mode-card p  { margin: 0 0 16px 0; font-size: 13px; color: var(--ia-text-muted) !important; line-height: 1.5; }

/* ══ Section headers inside annotation views ══════════════════════════════ */
.ia-section {
    background: var(--ia-section-bg);
    border-left: 4px solid var(--ia-accent);
    border-radius: 0 8px 8px 0;
    padding: 10px 16px;
    margin: 16px 0 12px 0;
    color: var(--ia-section-hd) !important;
    font-weight: 600;
    font-size: 15px;
}
/* Make all Markdown text inside annotation views respect dark mode */
.predefined-view p, .predefined-view span, .predefined-view label,
.custom-view p,     .custom-view span,     .custom-view label {
    color: var(--ia-text);
}

/* ══ Schema builder ════════════════════════════════════════════════════════ */
.schema-builder-box {
    background: var(--ia-section-bg);
    border: 1px solid var(--ia-border);
    border-radius: 10px;
    padding: 16px;
    margin-bottom: 12px;
}

/* ══ Step badges on Data tab ══════════════════════════════════════════════ */
.step-box {
    border: 1px solid var(--ia-border);
    border-radius: 10px;
    padding: 16px 18px;
    background: var(--ia-card-bg);
    color: var(--ia-text);
}
.step-box h4 { margin: 0 0 8px 0; color: var(--ia-text) !important; }

/* ══ Cluster gallery images: dark bg placeholder ══════════════════════════ */
.cl-gallery img { background: var(--ia-bg-subtle); }

/* ══ Tab strip: make active tab more prominent ════════════════════════════ */
.tabs > .tab-nav > button.selected {
    font-weight: 700;
    border-bottom: 3px solid var(--ia-accent) !important;
}

/* ══ More evident, user-friendly buttons ══════════════════════════════════ */
button.gr-button, .gr-button, button[class*="button"] {
    border-radius: 10px !important;
    font-weight: 600 !important;
    letter-spacing: 0.2px;
    transition: transform 0.12s ease, box-shadow 0.12s ease, filter 0.12s ease;
}
button.gr-button:hover, .gr-button:hover, button[class*="button"]:hover {
    transform: translateY(-1px);
    filter: brightness(1.05);
    box-shadow: 0 4px 12px rgba(0,0,0,0.18);
}
button.gr-button:active, .gr-button:active, button[class*="button"]:active {
    transform: translateY(0);
    filter: brightness(0.97);
}
button.primary {
    box-shadow: 0 2px 8px rgba(124, 111, 94, 0.35) !important;
    border: 2px solid var(--ia-accent) !important;
}
button.primary:hover {
    box-shadow: 0 6px 16px rgba(124, 111, 94, 0.45) !important;
}
button.secondary {
    border: 2px solid var(--ia-border) !important;
    font-weight: 600 !important;
}
.upload-btn-row {
    gap: 16px;
}
.upload-btn-col {
    background: var(--ia-card-bg);
    border: 1px solid var(--ia-border);
    border-radius: 10px;
    padding: 16px;
}
.upload-btn-col h4 { margin: 0 0 8px 0; color: var(--ia-text) !important; }
.upload-btn-col p  { margin: 0 0 12px 0; font-size: 13px; color: var(--ia-text-muted) !important; }

/* ══ Side-by-side image / value layout in annotation views ═══════════════ */
.ia-value-col { padding-left: 12px; }
"""

# Gradio 6 moved theme/css from the Blocks constructor to launch(); passing them
# to the constructor there is silently ignored. Route them by version.
_GR_MAJOR = int(gr.__version__.split(".")[0])
_STYLE_KWARGS = dict(theme=gr.themes.Soft(), css=CSS)

with gr.Blocks(title="IAnahid", **({} if _GR_MAJOR >= 6 else _STYLE_KWARGS)) as demo:

    # ── Shared session state ──────────────────────────────────────────────────
    mode_s        = gr.State("none")          # "predefined" | "custom" | "none"
    pre_df_s      = gr.State(_default_df)     # predefined mode DataFrame
    pre_t1_s      = gr.State(0)
    pre_t2_s      = gr.State(0)
    pre_cap_s     = gr.State(0)
    pre_cl_s      = gr.State(0)

    cust_df_s     = gr.State(pd.DataFrame())  # custom mode DataFrame
    cust_schema_s = gr.State([])              # list of {name, type, options}
    cust_idx_s    = gr.State(0)
    cust_cap_s    = gr.State(0)               # caption tab index (custom)
    cust_cl_s     = gr.State(0)               # cluster tab index (custom)

    gc_df_s       = gr.State(pd.DataFrame())  # guided-clustering (ITGC) DataFrame
    gc_meta_s     = gr.State({})              # {concepts, history, per_cluster, names, sil}
    gc_labeled_s  = gr.State([])              # [(url, label)] few-shot examples (DINOv2)
    gc_lbl_idx_s  = gr.State(0)               # current image index in the interactive labeller
    gc_lblcats_s  = gr.State([])              # ordered category labels for the labeller buttons

    gr.Markdown("# 🔍 IAnahid — Heritage Image Annotator")

    # ══════════════════════════════════════════════════════════════════════════
    # HOME PAGE
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Group(visible=True) as home_view:
        with gr.Row(equal_height=False):
            # Left column: guidelines
            with gr.Column(scale=2, elem_classes=["home-guidelines"]):
                gr.Markdown(GUIDELINES)

            # Right column: mode selection
            with gr.Column(scale=3, elem_classes=["home-modes"]):
                gr.Markdown("## Select an annotation mode")
                gr.Markdown(
                    "IAnahid supports two workflows. Choose the one that fits your project:"
                )
                with gr.Group(elem_classes=["mode-card"]):
                    gr.Markdown(
                        "### 🏛️ Predefined Schema\n"
                        "Use the built-in annotation schema designed for heritage archive "
                        "cataloguing: **Tag 1** (document type), **Tag 2** (subject), and "
                        "an **AI-assisted caption**. Includes cluster-level batch validation.\n\n"
                        "*Best for: ThierryNum datasets and heritage cataloguing projects.*"
                    )
                    btn_go_predefined = gr.Button(
                        "Start with Predefined Schema →", variant="primary", size="lg"
                    )
                with gr.Group(elem_classes=["mode-card"]):
                    gr.Markdown(
                        "### 🛠️ Custom Schema\n"
                        "Define your own metadata fields — names, types (free text or "
                        "controlled dropdown), and options. Annotate at your own pace "
                        "and export in CSV or JSON.\n\n"
                        "*Best for: original research projects with custom annotation needs.*"
                    )
                    btn_go_custom = gr.Button(
                        "Start with Custom Schema →", variant="secondary", size="lg"
                    )
                with gr.Group(elem_classes=["mode-card"]):
                    gr.Markdown(
                        "### 🧭 Guided Clustering (ITGC)\n"
                        "Let an AI **discover interpretable clusters** from a plain-language "
                        "criterion (e.g. *“group by document type”*). A text model proposes "
                        "visual concepts, a **vision model** scores each image against them, and the "
                        "groups are refined over several iterations — each cluster is explained "
                        "by the concepts that define it. Results feed straight into cluster "
                        "batch validation.\n\n"
                        "*Best for: exploring a fresh image set and building meaningful clusters.*"
                    )
                    btn_go_clustering = gr.Button(
                        "Start Guided Clustering →", variant="secondary", size="lg"
                    )

    # ══════════════════════════════════════════════════════════════════════════
    # PREDEFINED SCHEMA VIEW
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Group(visible=False) as predefined_view:
        with gr.Row():
            btn_pre_home = gr.Button("← Back to Home", size="sm", scale=0)
            gr.Markdown("### 🏛️ Predefined Schema Mode")

        with gr.Tabs() as pre_tabs:

            # ── Data tab ─────────────────────────────────────────────────────
            with gr.Tab("📁 Data", id="pre_data"):
                gr.HTML('<div class="ia-section">Load images or resume previous work</div>')
                with gr.Row(elem_classes=["upload-btn-row"]):
                    with gr.Column(elem_classes=["upload-btn-col"]):
                        gr.Markdown(
                            "#### 🌐 New URL list\n"
                            "Upload a `.txt`, `.csv`, or `.json` file listing image URLs. "
                            "A new annotation file will be created."
                        )
                        pre_url_load = gr.UploadButton(
                            "📤 Load URL list", file_types=[".txt", ".csv", ".json"],
                            variant="primary", size="lg",
                        )
                    with gr.Column(elem_classes=["upload-btn-col"]):
                        gr.Markdown(
                            "#### ☁️ HuggingFace dataset\n"
                            "Load the pre-configured `exemple.csv` dataset. "
                        )
                        pre_hf_load = gr.Button(
                            "☁️ Load from HuggingFace", variant="secondary", size="lg"
                        )
                    with gr.Column(elem_classes=["upload-btn-col"]):
                        gr.Markdown(
                            "#### 📂 Resume work\n"
                            "Upload a previously exported `.csv` or `.json` annotation file. "
                            "The app resumes from the first unvalidated item."
                        )
                        pre_ann_load = gr.UploadButton(
                            "📂 Resume work", file_types=[".csv", ".json"],
                            variant="secondary", size="lg",
                        )

                pre_data_msg = gr.Markdown("")

            # ── Tag 1 tab ─────────────────────────────────────────────────────
            with gr.Tab("🏷️ Tag 1 — Document type", id="pre_t1"):
                gr.HTML('<div class="ia-section">Classify the document type for each image</div>')
                with gr.Row():
                    pre_t1_prev = gr.Button("⬅️ Previous", scale=1)
                    pre_t1_skip = gr.Button("⏭️ Skip", scale=1)
                    pre_t1_cnt  = gr.Markdown(counter(0, max(len(_default_df), 1)))
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["ia-image-col"]):
                        pre_t1_img  = gr.HTML(img_html(img_url_at(_default_df, 0)))
                    with gr.Column(scale=1, elem_classes=["ia-value-col"]):
                        pre_t1_info = gr.Markdown("*No data loaded.*")
                with gr.Row():
                    pre_t1_ok  = gr.Button("✅ Validate as-is", variant="secondary", scale=1)
                    pre_t1_mod = gr.Button("✏️ Modify category", scale=1)
                with gr.Group(visible=False) as pre_t1_edit:
                    pre_t1_dd = gr.Dropdown(VOCAB_TAG_1, label="New category")
                    pre_t1_sv = gr.Button("💾 Save & continue")
                pre_t1_msg = gr.Markdown("")

            # ── Tag 2 tab ─────────────────────────────────────────────────────
            with gr.Tab("🏷️ Tag 2 — Subject", id="pre_t2"):
                gr.HTML('<div class="ia-section">Classify the main subject of each image</div>')
                with gr.Row():
                    pre_t2_prev = gr.Button("⬅️ Previous", scale=1)
                    pre_t2_skip = gr.Button("⏭️ Skip", scale=1)
                    pre_t2_cnt  = gr.Markdown(counter(0, max(len(_default_df), 1)))
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["ia-image-col"]):
                        pre_t2_img  = gr.HTML(img_html(img_url_at(_default_df, 0)))
                    with gr.Column(scale=1, elem_classes=["ia-value-col"]):
                        pre_t2_info = gr.Markdown("*No data loaded.*")
                with gr.Row():
                    pre_t2_ok  = gr.Button("✅ Validate as-is", variant="secondary", scale=1)
                    pre_t2_mod = gr.Button("✏️ Modify category", scale=1)
                with gr.Group(visible=False) as pre_t2_edit:
                    pre_t2_dd = gr.Dropdown(VOCAB_TAG_2, label="New category")
                    pre_t2_sv = gr.Button("💾 Save & continue")
                pre_t2_msg = gr.Markdown("")

            # ── Caption tab ───────────────────────────────────────────────────
            with gr.Tab("📝 Caption — AI-assisted", id="pre_cap"):
                gr.HTML('<div class="ia-section">Generate or write a catalogue caption for each image</div>')
                with gr.Row():
                    pre_cap_prev = gr.Button("⬅️ Previous", scale=1)
                    pre_cap_skip = gr.Button("⏭️ Skip", scale=1)
                    pre_cap_cnt  = gr.Markdown(counter(0, max(len(_default_df), 1)))
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["ia-image-col"]):
                        pre_cap_img  = gr.HTML(img_html(img_url_at(_default_df, 0)))
                    with gr.Column(scale=1, elem_classes=["ia-value-col"]):
                        pre_cap_info = gr.Markdown("*No data loaded.*")
                with gr.Row():
                    pre_cap_key = gr.Textbox(
                        label="🔑 Mistral API key", type="password",
                        placeholder="Paste your key here", scale=2,
                    )
                    pre_cap_mdl = gr.Dropdown(
                        ["ministral-14b-2512", "mistral-medium-latest", "ministral-8b-2512"],
                        value="ministral-14b-2512", label="Model", scale=1,
                    )
                pre_cap_gen = gr.Button("🤖 Generate caption with Mistral", variant="primary")
                pre_cap_txt = gr.Textbox(label="Caption", lines=4,
                                         placeholder="Edit or generate a caption here")
                pre_cap_sv  = gr.Button("💾 Save & continue", variant="secondary")
                pre_cap_msg = gr.Markdown("")

            # ── Cluster tab ───────────────────────────────────────────────────
            with gr.Tab("🗂️ Clusters — Batch validation", id="pre_cl"):
                gr.HTML('<div class="ia-section">Assign a Tag 2 to an entire cluster at once</div>')
                gr.Markdown(
                    "Images are grouped by visual similarity. Assign one **Tag 2** value to "
                    "all images in a cluster with a single click. Up to 10 sample images shown."
                )
                with gr.Row():
                    pre_cl_prev  = gr.Button("⬅️ Previous cluster", scale=1)
                    pre_cl_skip  = gr.Button("⏭️ Next cluster", scale=1)
                    pre_cl_cnt   = gr.Markdown("—")
                pre_cl_title = gr.Markdown("*No clusters loaded.*")
                pre_cl_html  = gr.HTML("<p><em>No images.</em></p>")
                pre_cl_dd    = gr.Dropdown(VOCAB_TAG_2, label="Assign Tag 2 to entire cluster")
                pre_cl_apply = gr.Button("✅ Apply to cluster & continue", variant="primary")
                pre_cl_msg   = gr.Markdown("")

            # ── Export file tab (table + export) ──────────────────────────────
            with gr.Tab("📊 Export file", id="pre_export"):
                pre_tbl_refresh = gr.Button("🔄 Refresh")
                pre_tbl_note    = gr.Markdown(tbl_note(_default_df))
                pre_tbl         = gr.Dataframe(value=tbl_preview(_default_df), interactive=False)
                gr.HTML('<div class="ia-section">Export annotations</div>')
                gr.Markdown("Export your current annotations at any time. "
                            "**Always export before closing the window** — annotations are not saved automatically.")
                with gr.Row():
                    pre_exp_csv  = gr.Button("⬇️ Export as CSV", variant="primary", size="lg")
                    pre_exp_json = gr.Button("⬇️ Export as JSON", variant="primary", size="lg")
                pre_exp_msg  = gr.Markdown("")
                pre_exp_file = gr.File(label="⬇️ Download", visible=False)

    # ══════════════════════════════════════════════════════════════════════════
    # CUSTOM SCHEMA VIEW
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Group(visible=False) as custom_view:
        with gr.Row():
            btn_cust_home = gr.Button("← Back to Home", size="sm", scale=0)
            gr.Markdown("### 🛠️ Custom Schema Mode")

        with gr.Tabs() as cust_tabs:

            # ── Schema builder tab ────────────────────────────────────────────
            with gr.Tab("📐 Schema builder", id="cust_schema"):
                gr.HTML('<div class="ia-section">① Define your metadata fields</div>')
                gr.Markdown(
                    "Add the fields you want to annotate. Each field needs a **name** and a **type**.\n"
                    "For *dropdown* fields, list the allowed values separated by commas.\n\n"
                    "When ready, go to the **📁 Data** tab to load your images."
                )
                with gr.Group(elem_classes=["schema-builder-box"]):
                    gr.Markdown("**Add a new field**")
                    with gr.Row():
                        sch_name    = gr.Textbox(label="Field name", placeholder="e.g. technique", scale=2)
                        sch_type    = gr.Dropdown(["text", "dropdown"], value="text", label="Type", scale=1)
                        sch_options = gr.Textbox(label="Dropdown options (comma-separated)",
                                                 placeholder="oil painting, watercolour, engraving",
                                                 scale=3, visible=False)
                        sch_add     = gr.Button("➕ Add field", variant="primary", scale=1)
                    sch_msg = gr.Markdown("")
                gr.HTML('<div class="ia-section">② Current schema</div>')
                sch_display = gr.Dataframe(
                    headers=["Field name", "Type", "Options"],
                    datatype=["str", "str", "str"],
                    interactive=False,
                    label="",
                )
                sch_clear = gr.Button("🗑️ Clear all fields", size="sm")

            # ── Data tab ─────────────────────────────────────────────────────
            with gr.Tab("📁 Data", id="cust_data"):
                gr.HTML('<div class="ia-section">Load images or resume previous work</div>')
                with gr.Row(elem_classes=["upload-btn-row"]):
                    with gr.Column(elem_classes=["upload-btn-col"]):
                        gr.Markdown(
                            "#### 🌐 New URL list\n"
                            "Upload a `.txt`, `.csv`, or `.json` file listing image URLs. "
                            "Define your schema first in the **📐 Schema builder** tab."
                        )
                        cust_url_load = gr.UploadButton(
                            "📤 Load URL list", file_types=[".txt", ".csv", ".json"],
                            variant="primary", size="lg",
                        )
                    with gr.Column(elem_classes=["upload-btn-col"]):
                        gr.Markdown(
                            "#### 📂 Resume work\n"
                            "Upload a previously exported `.json` or `.csv` file. "
                            "For JSON files, the schema is restored automatically."
                        )
                        cust_ann_load = gr.UploadButton(
                            "📂 Resume work", file_types=[".json", ".csv"],
                            variant="secondary", size="lg",
                        )
                cust_data_msg = gr.Markdown("")

            # ── Annotation tab ────────────────────────────────────────────────
            with gr.Tab("✏️ Annotate", id="cust_ann"):
                gr.HTML('<div class="ia-section">Fill in your custom fields for each image</div>')
                with gr.Row():
                    cust_prev = gr.Button("⬅️ Previous", scale=1)
                    cust_skip = gr.Button("⏭️ Skip", scale=1)
                    cust_cnt  = gr.Markdown("**0 / 0**")
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["ia-image-col"]):
                        cust_img  = gr.HTML(img_html(None))
                    with gr.Column(scale=1, elem_classes=["ia-value-col"]):
                        cust_info = gr.Markdown("*No data loaded.*")
                        # Dynamic fields are rendered here at runtime (up to 10 fields supported)
                        # We pre-create 10 invisible field slots; visibility is toggled by logic.
                        cust_field_rows = []
                        for i in range(10):
                            with gr.Group(visible=False) as field_grp:
                                lbl  = gr.Markdown(f"Field {i}")
                                txt  = gr.Textbox(label="", visible=True)
                                drp  = gr.Dropdown([], label="", visible=False)
                            cust_field_rows.append({"grp": field_grp, "lbl": lbl, "txt": txt, "drp": drp})

                cust_sv  = gr.Button("💾 Save & continue", variant="primary")
                cust_msg = gr.Markdown("")

            # ── Caption tab (custom) ──────────────────────────────────────────
            with gr.Tab("📝 Caption — AI-assisted"):
                gr.HTML('<div class="ia-section">Generate AI captions and save them to a custom field</div>')
                gr.Markdown(
                    "Generate a caption with Mistral and save it to any text field in your schema. "
                    "The caption is only a starting point — edit it freely before saving."
                )
                with gr.Row():
                    cust_cap_prev = gr.Button("⬅️ Previous", scale=1)
                    cust_cap_skip = gr.Button("⏭️ Skip", scale=1)
                    cust_cap_cnt  = gr.Markdown("**0 / 0**")
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["ia-image-col"]):
                        cust_cap_img  = gr.HTML(img_html(None))
                    with gr.Column(scale=1, elem_classes=["ia-value-col"]):
                        cust_cap_txt = gr.Textbox(label="Caption", lines=8,
                                                  placeholder="Generated caption appears here — edit freely")
                with gr.Row():
                    cust_cap_key = gr.Textbox(
                        label="🔑 Mistral API key", type="password",
                        placeholder="Paste your key here", scale=2,
                    )
                    cust_cap_mdl = gr.Dropdown(
                        ["ministral-14b-2512", "mistral-medium-latest", "ministral-8b-2512"],
                        value="ministral-14b-2512", label="Model", scale=1,
                    )
                    cust_cap_field = gr.Dropdown([], label="Save to field", scale=1)
                cust_cap_gen = gr.Button("🤖 Generate caption with Mistral", variant="primary")
                cust_cap_sv  = gr.Button("💾 Save to selected field & continue", variant="secondary")
                cust_cap_msg = gr.Markdown("")

            # ── Cluster tab (custom) ──────────────────────────────────────────
            with gr.Tab("🗂️ Clusters — Batch validation", id="cust_cl"):
                gr.HTML('<div class="ia-section">Assign a field value to an entire cluster at once</div>')
                gr.Markdown(
                    "Choose a **field** and a **value**, then apply to all images in the cluster. "
                    "The field must be a dropdown field defined in your schema. "
                    "Up to 10 sample images are shown per cluster.\n\n"
                    "**Note:** clustering requires a `clustering_id` column in your data."
                )
                with gr.Row():
                    cust_cl_prev  = gr.Button("⬅️ Previous cluster", scale=1)
                    cust_cl_skip  = gr.Button("⏭️ Next cluster", scale=1)
                    cust_cl_cnt   = gr.Markdown("—")
                cust_cl_title = gr.Markdown("*No clusters loaded.*")
                cust_cl_html  = gr.HTML("<p><em>No images.</em></p>")
                with gr.Row():
                    cust_cl_field = gr.Dropdown([], label="Field to assign")
                    cust_cl_val   = gr.Dropdown([], label="Value to apply")
                cust_cl_apply = gr.Button("✅ Apply to cluster & continue", variant="primary")
                cust_cl_msg   = gr.Markdown("")

            # ── Export file tab (table + export) ──────────────────────────────
            with gr.Tab("📊 Export file", id="cust_export"):
                gr.HTML('<div class="ia-section">Overview of all annotations in the current session</div>')
                cust_tbl_refresh = gr.Button("🔄 Refresh")
                cust_tbl_note    = gr.Markdown("")
                cust_tbl         = gr.Dataframe(value=pd.DataFrame(), interactive=False)
                gr.HTML('<div class="ia-section">Export annotations</div>')
                gr.Markdown("Export at any time. **Always export before closing the window.**")
                with gr.Row():
                    cust_exp_csv  = gr.Button("⬇️ Export as CSV", variant="primary", size="lg")
                    cust_exp_json = gr.Button("⬇️ Export as JSON (with schema)", variant="primary", size="lg")
                cust_exp_msg  = gr.Markdown("")
                cust_exp_file = gr.File(label="⬇️ Download", visible=False)

    # ══════════════════════════════════════════════════════════════════════════
    # GUIDED CLUSTERING (ITGC) VIEW
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Group(visible=False) as clustering_view:
        with gr.Row():
            btn_gc_home = gr.Button("← Back to Home", size="sm", scale=0)
            gr.Markdown("### 🧭 Guided Clustering (ITGC)")

        with gr.Tabs():
            # ── Data tab ─────────────────────────────────────────────────────
            with gr.Tab("📁 Data"):
                gr.HTML('<div class="ia-section">Load images & pick an encoding backend</div>')
                gr.Markdown(
                    "Upload a **URL list** (`.txt`, `.csv`, or `.json`) — images are previewed "
                    "directly from their source, never downloaded."
                )
                gc_url_load = gr.UploadButton(
                    "📤 Load URL list", file_types=[".txt", ".csv", ".json"],
                    variant="primary",
                )
                gc_backend = gr.Radio(
                    ["Local DINOv2 — few-shot / similarity (CPU · offline · recommended)",
                     "Local SigLIP2 — zero-shot text labels (CPU · offline)",
                     "Mistral Vision API"],
                    value="Local DINOv2 — few-shot / similarity (CPU · offline · recommended)",
                    label="Encoding backend",
                    info=("DINOv2 = best visual features: give a few labelled examples per "
                          "category (🧭 Cluster tab) for few-shot classification, or none for "
                          "pure similarity clusters. SigLIP2 = zero-shot from category names "
                          "(no examples). Mistral API = a VLM (needs a key), finer but slower. "
                          "All local options are offline and cache embeddings per image."),
                )
                gc_key = gr.Textbox(label="🔑 Mistral API key (only for the Mistral API backend)",
                                    type="password", placeholder="sk-…")
                with gr.Row():
                    # NOTE: pixtral-12b-2409 was retired on 2025-12-31; Mistral's
                    # designated successor is ministral-14b-2512 (vision-capable).
                    gc_vmodel = gr.Dropdown(
                        ["ministral-14b-2512", "mistral-medium-latest", "mistral-large-2512"],
                        value="ministral-14b-2512",
                        label="Vision model — Mistral API backend",
                    )
                    gc_tmodel = gr.Dropdown(
                        ["mistral-medium-latest", "mistral-small-latest", "ministral-8b-2512"],
                        value="mistral-medium-latest",
                        label="Text model — Mistral API discovery mode",
                    )
                gc_status = gr.Markdown("*No data loaded.*")

            # ── Cluster tab ──────────────────────────────────────────────────
            with gr.Tab("🧭 Cluster"):
                gr.HTML('<div class="ia-section">Two ways to group your images</div>')
                gc_categories = gr.Textbox(
                    label="① Fixed categories — one per line (recommended if you know them)",
                    placeholder="Architecture\nObjects\nPeople\nLandscape\nAnimal\nPlants",
                    lines=6,
                    info=("If filled, each image is classified into EXACTLY ONE of these "
                          "categories (no k-means; the k / iterations sliders are ignored). "
                          "Tip: add a description after a pipe to sharpen a category, e.g. "
                          "“People | portraits and group photographs of people”."),
                )
                gc_query = gr.Textbox(
                    label="② Or: discovery criterion (leave categories empty to use this)",
                    placeholder="e.g. group by decoration style / by scene / by period",
                    lines=2,
                    info=("Used only when no fixed categories are given: the AI invents its "
                          "own concepts and discovers clusters via the iterative ITGC search."),
                )
                gr.HTML('<div class="ia-section">③ Few-shot examples (DINOv2 backend)</div>')
                gr.Markdown(
                    "Give ~10-15 example images per category. Assigning is **instant** — the "
                    "DINOv2 embedding runs only when you press **Run**. *Leave empty to get pure "
                    "DINOv2 similarity clusters instead.*"
                )
                gc_labeled_status = gr.Markdown(
                    "*No examples yet — click a category for each image (~10-15 per category).*")

                with gr.Accordion("🖱️  Label interactively — click the category of each image",
                                  open=True):
                    gr.Markdown(
                        "Type your categories in **① Fixed categories** above, then press "
                        "**Prepare**. Browse with ⬅️ / ⏭️ and click the matching category — "
                        "it's saved and the next image appears (no waiting)."
                    )
                    gc_lbl_prepare = gr.Button("🔄 Prepare / refresh from ① categories",
                                               variant="secondary")
                    gc_lbl_img = gr.HTML("<p><em>Load images (📁 Data) and press Prepare.</em></p>")
                    gc_lbl_info = gr.Markdown("—")
                    with gr.Row():
                        gc_lbl_prev = gr.Button("⬅️ Previous", scale=1)
                        gc_lbl_skip = gr.Button("⏭️ Skip", scale=1)
                        gc_lbl_undo = gr.Button("↩️ Remove last example", scale=1)
                    gc_lbl_btns = []
                    with gr.Row():
                        for _i in range(GC_MAX_LABEL_BTNS):
                            gc_lbl_btns.append(gr.Button(f"cat {_i}", visible=False,
                                                         variant="primary", size="sm"))

                with gr.Accordion("📄  …or upload a labelled file (CSV/JSON)", open=False):
                    gr.Markdown(
                        "A `.csv`/`.json` with an `image_url` column and a label column "
                        "(`label`, `category`, `tag_2`, or `clustering_id`). An annotation file "
                        "exported from Predefined/Custom mode works directly."
                    )
                    gc_labeled_load = gr.UploadButton(
                        "🏷️ Load labelled examples", file_types=[".csv", ".json"],
                        variant="secondary",
                    )
                with gr.Row():
                    gc_k = gr.Slider(2, 12, value=6, step=1,
                                     label="Number of clusters k (discovery mode only)")
                    gc_iters = gr.Slider(1, 10, value=4, step=1,
                                         label="Refinement iterations (discovery mode only)")
                    gc_sample = gr.Slider(20, 200, value=60, step=10,
                                          label="Max images scored / iteration (discovery)")
                gr.Markdown(
                    "*Time scales with the image count. With the **Local SigLIP2** backend, "
                    "images are embedded once (cached) then classified/clustered instantly; "
                    "re-runs and category changes are immediate. With the **Mistral API**, every "
                    "image is a fresh API call. Results appear progressively with a progress bar; "
                    "the API discovery loop refines over iterations and stops early once the "
                    "silhouette score plateaus (k / iterations sliders apply to discovery).*"
                )
                gc_run_btn = gr.Button("🚀 Run", variant="primary", size="lg")
                gc_log = gr.Markdown("")

            # ── Results & Visualisation tab ──────────────────────────────────
            with gr.Tab("🔎 Results"):
                gr.HTML('<div class="ia-section">Interpretable clusters</div>')
                gc_sil_html = gr.HTML("")
                gc_concepts_md = gr.Markdown("*Run guided clustering to see results.*")
                gc_clusters_html = gr.HTML(
                    "<p><em>No clusters yet — run guided clustering.</em></p>"
                )

            # ── Export tab ───────────────────────────────────────────────────
            with gr.Tab("📊 Export"):
                gr.HTML('<div class="ia-section">Export the computed clusters</div>')
                gr.Markdown(
                    "The `clustering_id` column now holds the guided clusters. Export, then "
                    "**Resume work** in Predefined or Custom mode to batch-validate them."
                )
                gc_tbl = gr.Dataframe(value=pd.DataFrame(), interactive=False)
                with gr.Row():
                    gc_exp_csv = gr.Button("⬇️ Export as CSV", variant="primary", size="lg")
                    gc_exp_json = gr.Button("⬇️ Export as JSON", variant="primary", size="lg")
                gc_exp_msg = gr.Markdown("")
                gc_exp_file = gr.File(label="⬇️ Download", visible=False)

    # ── Footer ───────────────────────────────────────────────────────────────
    gr.HTML(
        f"""
        <div style="text-align:center;margin-top:30px;padding-top:20px;
                    border-top:1px solid #ddd;">
          {'<img src="data:image/png;base64,' + _LOGOS + '" style="display:block;margin:0 auto;max-width:700px;width:90%;height:auto;" alt="Logos"/><br/>' if _LOGOS else ''}
          <p style="font-size:0.85em;color:#888;max-width:820px;margin:0 auto;line-height:1.5;">
            IAnahid — developed during the 2026
            <a href="https://www.chartes.psl.eu/" target="_blank" rel="noopener">
            École nationale des chartes – PSL</a> hackathon for the
            <a href="https://www.inha.fr/recherche/projets/histoire-art-iv-xv-siecle/le-fonds-thierry/"
            target="_blank" rel="noopener">ThierryNum</a> project.
          </p>
        </div>
        """
    )

    # ══════════════════════════════════════════════════════════════════════════
    # EVENT WIRING — Home navigation
    # ══════════════════════════════════════════════════════════════════════════

    # Views order: [home, predefined, custom, clustering]
    def go_predefined():
        return (gr.update(visible=False), gr.update(visible=True),
                gr.update(visible=False), gr.update(visible=False))

    def go_custom():
        return (gr.update(visible=False), gr.update(visible=False),
                gr.update(visible=True), gr.update(visible=False))

    def go_clustering():
        return (gr.update(visible=False), gr.update(visible=False),
                gr.update(visible=False), gr.update(visible=True))

    def go_home():
        return (gr.update(visible=True), gr.update(visible=False),
                gr.update(visible=False), gr.update(visible=False))

    _views = [home_view, predefined_view, custom_view, clustering_view]
    btn_go_predefined.click(go_predefined, [], _views)
    btn_go_custom.click(go_custom,         [], _views)
    btn_go_clustering.click(go_clustering, [], _views)
    btn_pre_home.click(go_home,            [], _views)
    btn_cust_home.click(go_home,           [], _views)
    btn_gc_home.click(go_home,             [], _views)

    # ══════════════════════════════════════════════════════════════════════════
    # EVENT WIRING — Predefined mode helpers
    # ══════════════════════════════════════════════════════════════════════════

    def _pre_t1_view(df, idx):
        n = max(len(df), 1)
        return (
            img_html(img_url_at(df, idx)),
            row_info_predefined(df, idx, "tag_1", "validated_tag_1"),
            counter(idx, n),
            gr.update(visible=False),
            "",
        )

    def _pre_t2_view(df, idx):
        n = max(len(df), 1)
        return (
            img_html(img_url_at(df, idx)),
            row_info_predefined(df, idx, "tag_2", "validated_tag_2"),
            counter(idx, n),
            gr.update(visible=False),
            "",
        )

    def _pre_cap_view(df, idx):
        n = max(len(df), 1)
        row = df.iloc[int(idx)] if not df.empty else None
        cap = str(row["image_caption"]) if row is not None and pd.notna(row.get("image_caption")) else ""
        return (
            img_html(img_url_at(df, idx)),
            row_info_predefined(df, idx, "image_caption", "validated_caption"),
            counter(idx, n),
            cap,
            "",
        )

    def _pre_cl_view(df, cl_idx):
        clusters = cluster_list(df)
        if not clusters:
            return 0, "No clusters found.", "<p><em>No data.</em></p>", "—", ""
        cl_idx = min(int(cl_idx), len(clusters) - 1)
        return (
            cl_idx,
            cluster_title(df, cl_idx),
            cluster_gallery_html(df, cl_idx),
            f"Cluster {cl_idx + 1} / {len(clusters)}",
            "",
        )

    _pre_refresh_outputs = [
        pre_df_s,
        pre_t1_s, pre_t1_img, pre_t1_info, pre_t1_cnt,
        pre_t2_s, pre_t2_img, pre_t2_info, pre_t2_cnt,
        pre_cap_s, pre_cap_img, pre_cap_info, pre_cap_cnt,
        pre_cl_s, pre_cl_title, pre_cl_html, pre_cl_cnt,
        pre_tbl_note, pre_tbl,
        pre_tabs,
        pre_data_msg,
    ]

    def _pre_load_error(msg):
        """Return a no-op tuple of the right length for _pre_refresh_outputs on error."""
        n = len(_pre_refresh_outputs)
        return (gr.update(),) * (n - 1) + (msg,)

    # ── Predefined: load URL list ─────────────────────────────────────────────
    def pre_load_urls(file_obj):
        if file_obj is None:
            return _pre_load_error("⚠️ Select a URL list file first.")
        try:
            path = file_obj.name if hasattr(file_obj, "name") else file_obj
            urls = load_url_list(path)
            if not urls:
                return _pre_load_error("⚠️ No URLs found in the file.")
            df = pd.DataFrame({"image_url": urls})
            df = _ensure_predefined_columns(df)
        except Exception as e:
            return _pre_load_error(f"❌ Error: {e}")
        return _pre_full_refresh(df)

    def pre_load_annotation(file_obj):
        if file_obj is None:
            return _pre_load_error("⚠️ Select an annotation file first.")
        try:
            path = file_obj.name if hasattr(file_obj, "name") else file_obj
            df, _ = load_annotation_file(path)
            df = _ensure_predefined_columns(df)
        except Exception as e:
            return _pre_load_error(f"❌ Error: {e}")
        return _pre_full_refresh(df)

    def _pre_full_refresh(df):
        """Return all outputs needed to refresh every predefined tab, and
        automatically switch to the next annotation tab (Tag 1)."""
        n = max(len(df), 1)
        t1i  = first_unvalidated(df, "validated_tag_1")
        t2i  = first_unvalidated(df, "validated_tag_2")
        capi = first_unvalidated(df, "validated_caption")
        cl_idx, cl_ttl, cl_html, cl_cnt_txt, _ = _pre_cl_view(df, 0)
        return (
            df,                                          # pre_df_s
            t1i, img_html(img_url_at(df, t1i)),         # t1_s, t1_img
            row_info_predefined(df, t1i, "tag_1", "validated_tag_1"),
            counter(t1i, n),                             # t1_info, t1_cnt
            t2i, img_html(img_url_at(df, t2i)),
            row_info_predefined(df, t2i, "tag_2", "validated_tag_2"),
            counter(t2i, n),
            capi, img_html(img_url_at(df, capi)),
            row_info_predefined(df, capi, "image_caption", "validated_caption"),
            counter(capi, n),
            cl_idx, cl_ttl, cl_html, cl_cnt_txt,
            tbl_note(df), tbl_preview(df),               # tbl (browser preview only)
            gr.Tabs(selected="pre_t1"),                   # auto-advance to next tab
            f"✅ Loaded {len(df)} rows, {len(cluster_list(df))} clusters.",
        )

    def pre_load_from_hf():
        """Load the example dataset — the HuggingFace-hosted copy if DATASET_REPO
        and HF_TOKEN are set, otherwise the bundled exemple.csv (no upload needed)."""
        df = _load_default_df()
        if df.empty:
            return (gr.update(),) * (len(_pre_refresh_outputs) - 1) + (
                f"⚠️ Could not load the example dataset. The bundled «{CSV_FILENAME}» "
                "was not found and no valid DATASET_REPO / HF_TOKEN secrets are set.",
            )
        return _pre_full_refresh(df)

    pre_url_load.upload(pre_load_urls,      [pre_url_load], _pre_refresh_outputs, api_name=False)
    pre_ann_load.upload(pre_load_annotation, [pre_ann_load], _pre_refresh_outputs, api_name=False)
    pre_hf_load.click(pre_load_from_hf,     [],             _pre_refresh_outputs, api_name=False)

    # ── Predefined: export ─────────────────────────────────────────────────────
    # Clicking a button generates the file and reveals a download link below.
    def do_pre_export_csv(df):
        if df is None or df.empty:
            return gr.update(visible=False), "⚠️ Nothing to export — load data first."
        return gr.update(value=export_as_csv(df), visible=True), ""

    def do_pre_export_json(df):
        if df is None or df.empty:
            return gr.update(visible=False), "⚠️ Nothing to export — load data first."
        return gr.update(value=export_as_json(df), visible=True), ""

    pre_exp_csv.click(do_pre_export_csv,  [pre_df_s], [pre_exp_file, pre_exp_msg], api_name=False)
    pre_exp_json.click(do_pre_export_json, [pre_df_s], [pre_exp_file, pre_exp_msg], api_name=False)

    # ── Predefined: Tag 1 events ──────────────────────────────────────────────
    def on_pre_t1_prev(df, idx):
        ni = max(int(idx) - 1, 0)
        return (ni,) + _pre_t1_view(df, ni)

    def on_pre_t1_skip(df, idx):
        ni = min(int(idx) + 1, max(len(df) - 1, 0))
        return (ni,) + _pre_t1_view(df, ni)

    def on_pre_t1_mod():
        return gr.update(visible=True)

    def on_pre_t1_ok(df, idx):
        idx = int(idx)
        current = df.at[idx, "tag_1"] if not df.empty else ""
        if not str(current).strip():
            return (df, idx) + _pre_t1_view(df, idx)[:-1] + ("⚠️ Cannot validate an empty tag. Choose a value first.",)
        df = df.copy()
        df.at[idx, "validated_tag_1"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        # Message is cleared on advance so the tick doesn't appear to belong to the next photo.
        return (df, ni) + _pre_t1_view(df, ni)[:-1] + ("",)

    def on_pre_t1_sv(df, idx, new_val):
        if not new_val:
            return df, int(idx), gr.update(), gr.update(), gr.update(), gr.update(visible=True), "⚠️ Select a category."
        idx = int(idx); df = df.copy()
        df.at[idx, "tag_1"] = new_val; df.at[idx, "validated_tag_1"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        return (df, ni) + _pre_t1_view(df, ni)[:-1] + ("",)

    _t1_outs = [pre_t1_s, pre_t1_img, pre_t1_info, pre_t1_cnt, pre_t1_edit, pre_t1_msg]
    pre_t1_prev.click(on_pre_t1_prev, [pre_df_s, pre_t1_s], _t1_outs, api_name=False)
    pre_t1_skip.click(on_pre_t1_skip, [pre_df_s, pre_t1_s], _t1_outs, api_name=False)
    pre_t1_mod.click(on_pre_t1_mod,   [],                   [pre_t1_edit], api_name=False)
    pre_t1_ok.click(on_pre_t1_ok,     [pre_df_s, pre_t1_s], [pre_df_s] + _t1_outs, api_name=False)
    pre_t1_sv.click(on_pre_t1_sv,     [pre_df_s, pre_t1_s, pre_t1_dd], [pre_df_s] + _t1_outs, api_name=False)

    # ── Predefined: Tag 2 events ──────────────────────────────────────────────
    def on_pre_t2_prev(df, idx):
        ni = max(int(idx) - 1, 0)
        return (ni,) + _pre_t2_view(df, ni)

    def on_pre_t2_skip(df, idx):
        ni = min(int(idx) + 1, max(len(df) - 1, 0))
        return (ni,) + _pre_t2_view(df, ni)

    def on_pre_t2_mod():
        return gr.update(visible=True)

    def on_pre_t2_ok(df, idx):
        idx = int(idx)
        current = df.at[idx, "tag_2"] if not df.empty else ""
        if not str(current).strip():
            return (df, idx) + _pre_t2_view(df, idx)[:-1] + ("⚠️ Cannot validate an empty tag. Choose a value first.",)
        df = df.copy()
        df.at[idx, "validated_tag_2"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        return (df, ni) + _pre_t2_view(df, ni)[:-1] + ("",)

    def on_pre_t2_sv(df, idx, new_val):
        if not new_val:
            return df, int(idx), gr.update(), gr.update(), gr.update(), gr.update(visible=True), "⚠️ Select a category."
        idx = int(idx); df = df.copy()
        df.at[idx, "tag_2"] = new_val; df.at[idx, "validated_tag_2"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        return (df, ni) + _pre_t2_view(df, ni)[:-1] + ("",)

    _t2_outs = [pre_t2_s, pre_t2_img, pre_t2_info, pre_t2_cnt, pre_t2_edit, pre_t2_msg]
    pre_t2_prev.click(on_pre_t2_prev, [pre_df_s, pre_t2_s], _t2_outs, api_name=False)
    pre_t2_skip.click(on_pre_t2_skip, [pre_df_s, pre_t2_s], _t2_outs, api_name=False)
    pre_t2_mod.click(on_pre_t2_mod,   [],                   [pre_t2_edit], api_name=False)
    pre_t2_ok.click(on_pre_t2_ok,     [pre_df_s, pre_t2_s], [pre_df_s] + _t2_outs, api_name=False)
    pre_t2_sv.click(on_pre_t2_sv,     [pre_df_s, pre_t2_s, pre_t2_dd], [pre_df_s] + _t2_outs, api_name=False)

    # ── Predefined: Caption events ────────────────────────────────────────────
    def on_pre_cap_prev(df, idx):
        ni = max(int(idx) - 1, 0)
        return (ni,) + _pre_cap_view(df, ni)

    def on_pre_cap_skip(df, idx):
        ni = min(int(idx) + 1, max(len(df) - 1, 0))
        return (ni,) + _pre_cap_view(df, ni)

    def on_pre_cap_gen(df, idx, api_key, model):
        url = img_url_at(df, int(idx))
        text, msg = generate_mistral_caption(url, api_key, model)
        return gr.update(value=text) if text else gr.update(), msg

    def on_pre_cap_sv(df, idx, text):
        if not str(text).strip():
            return (df, int(idx),) + (gr.update(),) * 4 + ("⚠️ Caption is empty.",)
        idx = int(idx); df = df.copy()
        df.at[idx, "image_caption"] = text; df.at[idx, "validated_caption"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        return (df, ni) + _pre_cap_view(df, ni)[:-1] + ("",)

    _cap_outs = [pre_cap_s, pre_cap_img, pre_cap_info, pre_cap_cnt, pre_cap_txt, pre_cap_msg]
    pre_cap_prev.click(on_pre_cap_prev, [pre_df_s, pre_cap_s], _cap_outs, api_name=False)
    pre_cap_skip.click(on_pre_cap_skip, [pre_df_s, pre_cap_s], _cap_outs, api_name=False)
    pre_cap_gen.click(on_pre_cap_gen,   [pre_df_s, pre_cap_s, pre_cap_key, pre_cap_mdl],
                      [pre_cap_txt, pre_cap_msg], api_name=False)
    pre_cap_sv.click(on_pre_cap_sv,     [pre_df_s, pre_cap_s, pre_cap_txt],
                     [pre_df_s] + _cap_outs, api_name=False)

    # ── Predefined: Cluster events ────────────────────────────────────────────
    def on_pre_cl_prev(df, cl_idx):
        ni = max(int(cl_idx) - 1, 0)
        return _pre_cl_view(df, ni)

    def on_pre_cl_skip(df, cl_idx):
        clusters = cluster_list(df)
        ni = min(int(cl_idx) + 1, max(len(clusters) - 1, 0))
        return _pre_cl_view(df, ni)

    def on_pre_cl_apply(df, cl_idx, chosen):
        if not chosen:
            return (df,) + _pre_cl_view(df, cl_idx)[:-1] + ("⚠️ Select a class first.",)
        clusters = cluster_list(df); cl_idx = int(cl_idx)
        c_id = clusters[cl_idx]
        df = df.copy()
        mask = df["clustering_id"].astype(str) == c_id
        df.loc[mask, "tag_2"] = chosen; df.loc[mask, "validated_tag_2"] = True
        ni = min(cl_idx + 1, max(len(clusters) - 1, 0))
        return (df,) + _pre_cl_view(df, ni)[:-1] + (f"✅ «{chosen}» applied to cluster {c_id}.",)

    _cl_outs = [pre_cl_s, pre_cl_title, pre_cl_html, pre_cl_cnt, pre_cl_msg]
    pre_cl_prev.click(on_pre_cl_prev,   [pre_df_s, pre_cl_s], _cl_outs, api_name=False)
    pre_cl_skip.click(on_pre_cl_skip,   [pre_df_s, pre_cl_s], _cl_outs, api_name=False)
    pre_cl_apply.click(on_pre_cl_apply, [pre_df_s, pre_cl_s, pre_cl_dd],
                       [pre_df_s] + _cl_outs, api_name=False)

    # ── Predefined: Table refresh ─────────────────────────────────────────────
    pre_tbl_refresh.click(lambda df: (tbl_note(df), tbl_preview(df)),
                          [pre_df_s], [pre_tbl_note, pre_tbl], api_name=False)

    # ══════════════════════════════════════════════════════════════════════════
    # EVENT WIRING — Custom schema mode
    # ══════════════════════════════════════════════════════════════════════════

    # ── Schema builder ────────────────────────────────────────────────────────

    def toggle_options(field_type):
        return gr.update(visible=(field_type == "dropdown"))

    sch_type.change(toggle_options, [sch_type], [sch_options])

    def add_field(schema, name, ftype, options):
        name = name.strip()
        if not name:
            return schema, _schema_df(schema), "⚠️ Field name is required."
        # Sanitise: replace spaces with underscores
        col_name = name.lower().replace(" ", "_")
        if any(f["name"] == col_name for f in schema):
            return schema, _schema_df(schema), f"⚠️ Field «{col_name}» already exists."
        opts = [o.strip() for o in options.split(",") if o.strip()] if ftype == "dropdown" else []
        schema = schema + [{"name": col_name, "label": name, "type": ftype, "options": opts}]
        return schema, _schema_df(schema), f"✅ Field «{col_name}» added."

    def _schema_df(schema):
        rows = [[f["label"], f["type"], ", ".join(f.get("options", []))] for f in schema]
        return pd.DataFrame(rows, columns=["Field name", "Type", "Options"])

    def clear_schema():
        return [], _schema_df([]), ""

    sch_add.click(
        add_field,
        [cust_schema_s, sch_name, sch_type, sch_options],
        [cust_schema_s, sch_display, sch_msg],
        api_name=False,
    )
    sch_clear.click(clear_schema, [], [cust_schema_s, sch_display, sch_msg], api_name=False)

    # ── Custom: load URL list ─────────────────────────────────────────────────

    def _cust_load_error(msg):
        """Return a no-op tuple of the right length for _cust_load_outs on error."""
        n = len(_cust_load_outs)
        return (gr.update(),) * (n - 1) + (msg,)

    def cust_load_urls(file_obj, schema):
        if file_obj is None:
            return _cust_load_error("⚠️ Select a URL list file first.")
        if not schema:
            return _cust_load_error("⚠️ Define your schema first in the 📐 Schema builder tab.")
        try:
            path = file_obj.name if hasattr(file_obj, "name") else file_obj
            urls = load_url_list(path)
            if not urls:
                return _cust_load_error("⚠️ No URLs found in the file.")
            df = pd.DataFrame({"image_url": urls})
            df = _ensure_custom_columns(df, schema)
        except Exception as e:
            return _cust_load_error(f"❌ Error: {e}")
        return _cust_full_refresh(df, schema, 0) + (f"✅ Loaded {len(df)} image URLs.",)

    def cust_load_annotation(file_obj, schema):
        if file_obj is None:
            return _cust_load_error("⚠️ Select an annotation file first.")
        try:
            path = file_obj.name if hasattr(file_obj, "name") else file_obj
            df, saved_schema = load_annotation_file(path)
            if saved_schema:
                schema = saved_schema
            if not schema:
                return _cust_load_error("⚠️ No schema found. Define one in the 📐 Schema builder first.")
            df = _ensure_custom_columns(df, schema)
        except Exception as e:
            return _cust_load_error(f"❌ Error: {e}")
        idx = _first_cust_unvalidated(df, schema)
        return _cust_full_refresh(df, schema, idx) + (f"✅ Resumed — {len(df)} rows.",)

    def _first_cust_unvalidated(df, schema):
        if df.empty or not schema:
            return 0
        # Find first row where any field is unvalidated
        for i in range(len(df)):
            row = df.iloc[i]
            for field in schema:
                val_col = f"validated_{field['name']}"
                if val_col not in df.columns or row.get(val_col) != True:
                    return i
        return 0

    def _cust_field_updates(schema, df, idx):
        """Return gr.update() tuples for all 10 field slots."""
        updates = []
        row = df.iloc[int(idx)] if not df.empty and idx < len(df) else None
        for i, fr in enumerate(cust_field_rows):
            if i < len(schema):
                field  = schema[i]
                val    = str(row[field["name"]]) if row is not None and pd.notna(row.get(field["name"])) else ""
                is_dd  = field["type"] == "dropdown"
                opts   = field.get("options", [])
                updates += [
                    gr.update(visible=True),                          # grp
                    gr.update(value=f"**{field['label']}**"),         # lbl
                    gr.update(value=val if not is_dd else "", visible=not is_dd),  # txt
                    gr.update(choices=opts, value=val if is_dd and val in opts else None,
                              visible=is_dd),                         # drp
                ]
            else:
                updates += [
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value="", visible=False),
                    gr.update(choices=[], value=None, visible=False),
                ]
        return updates

    def _cust_dropdown_fields(schema):
        """Return list of dropdown-type field names for cluster/caption field selectors."""
        return [f["name"] for f in schema if f["type"] == "dropdown"]

    def _cust_text_fields(schema):
        """Return list of text-type field names for caption target selector."""
        return [f["name"] for f in schema if f["type"] == "text"]

    def _cust_all_fields(schema):
        return [f["name"] for f in schema]

    def _cust_cl_view(df, schema, cl_idx):
        clusters = cluster_list(df)
        if not clusters:
            return 0, "No clusters found.", "<p><em>No data or no clustering_id column.</em></p>", "—", ""
        cl_idx = min(int(cl_idx), len(clusters) - 1)
        return (
            cl_idx,
            cluster_title(df, cl_idx),
            cluster_gallery_html(df, cl_idx),
            f"Cluster {cl_idx + 1} / {len(clusters)}",
            "",
        )

    def _cust_full_refresh(df, schema, idx):
        n = max(len(df), 1)
        field_upds = _cust_field_updates(schema, df, idx)
        # caption tab
        cap_cnt = f"**{idx + 1} / {n}**"
        cap_img = img_html(img_url_at(df, idx))
        # field choices for caption and cluster dropdowns
        text_fields = _cust_text_fields(schema)
        dd_fields   = _cust_dropdown_fields(schema)
        # cluster
        cl_idx, cl_ttl, cl_html, cl_cnt_txt, _ = _cust_cl_view(df, schema, 0)
        return (
            df, schema, idx, 0, 0,
            img_html(img_url_at(df, idx)),
            f"**{idx + 1} / {n}**  — fill in the fields below and click Save.",
            *field_upds,
            cap_cnt, cap_img,
            gr.update(choices=text_fields + dd_fields, value=None),  # cust_cap_field
            "",                                                        # cust_cap_txt
            "",                                                        # cust_cap_msg
            cl_idx, cl_ttl, cl_html, cl_cnt_txt,
            gr.update(choices=dd_fields, value=None),                 # cust_cl_field
            gr.update(choices=[], value=None),                        # cust_cl_val
            "",                                                        # cust_cl_msg
            tbl_note(df), tbl_preview(df),                             # tbl (browser preview only)
            gr.Tabs(selected="cust_ann"),                              # auto-advance to Annotate tab
        )

    _cust_all_field_outs = []
    for fr in cust_field_rows:
        _cust_all_field_outs += [fr["grp"], fr["lbl"], fr["txt"], fr["drp"]]

    _cust_load_outs = (
        [cust_df_s, cust_schema_s, cust_idx_s, cust_cap_s, cust_cl_s,
         cust_img, cust_info]
        + _cust_all_field_outs
        + [cust_cap_cnt, cust_cap_img, cust_cap_field, cust_cap_txt, cust_cap_msg,
           cust_cl_s, cust_cl_title, cust_cl_html, cust_cl_cnt,
           cust_cl_field, cust_cl_val, cust_cl_msg,
           cust_tbl_note, cust_tbl, cust_tabs, cust_data_msg]
    )

    cust_url_load.upload(cust_load_urls,       [cust_url_load, cust_schema_s], _cust_load_outs, api_name=False)
    cust_ann_load.upload(cust_load_annotation, [cust_ann_load, cust_schema_s], _cust_load_outs, api_name=False)

    # ── Custom: export ─────────────────────────────────────────────────────────
    # Clicking a button generates the file and reveals a download link below.
    # The JSON export embeds the schema so the file can be fully reloaded later.
    def do_cust_export_csv(df):
        if df is None or df.empty:
            return gr.update(visible=False), "⚠️ Nothing to export — load data first."
        return gr.update(value=export_as_csv(df), visible=True), ""

    def do_cust_export_json(df, schema):
        if df is None or df.empty:
            return gr.update(visible=False), "⚠️ Nothing to export — load data first."
        return gr.update(value=export_as_json(df, schema), visible=True), ""

    cust_exp_csv.click(do_cust_export_csv,  [cust_df_s],                [cust_exp_file, cust_exp_msg], api_name=False)
    cust_exp_json.click(do_cust_export_json, [cust_df_s, cust_schema_s], [cust_exp_file, cust_exp_msg], api_name=False)

    # ── Custom: annotation navigation ────────────────────────────────────────
    def _cust_nav_outs(df, schema, ni):
        n = max(len(df), 1)
        field_upds = _cust_field_updates(schema, df, ni)
        return (ni, img_html(img_url_at(df, ni)),
                f"**{ni + 1} / {n}**",
                *field_upds, "")

    def on_cust_prev(df, schema, idx):
        ni = max(int(idx) - 1, 0)
        return _cust_nav_outs(df, schema, ni)

    def on_cust_skip(df, schema, idx):
        ni = min(int(idx) + 1, max(len(df) - 1, 0))
        return _cust_nav_outs(df, schema, ni)

    _cust_nav_out_keys = [cust_idx_s, cust_img, cust_cnt] + _cust_all_field_outs + [cust_msg]
    cust_prev.click(on_cust_prev, [cust_df_s, cust_schema_s, cust_idx_s], _cust_nav_out_keys, api_name=False)
    cust_skip.click(on_cust_skip, [cust_df_s, cust_schema_s, cust_idx_s], _cust_nav_out_keys, api_name=False)

    # ── Custom: save ──────────────────────────────────────────────────────────
    # We collect all text and dropdown values for the 10 field slots
    _cust_field_vals = []
    for fr in cust_field_rows:
        _cust_field_vals += [fr["txt"], fr["drp"]]

    def on_cust_sv(df, schema, idx, *field_vals):
        if df.empty or not schema:
            return (df,) + _cust_nav_outs(df, schema, int(idx))[:-1] + ("⚠️ No data loaded.",)
        idx = int(idx); df = df.copy()
        # field_vals alternates: txt0, drp0, txt1, drp1, …
        values = {}
        for i, field in enumerate(schema):
            txt_val = field_vals[i * 2]
            drp_val = field_vals[i * 2 + 1]
            value = drp_val if field["type"] == "dropdown" else txt_val
            values[field["name"]] = str(value).strip() if value else ""
        for name, v in values.items():
            df.at[idx, name] = v
        if any(v == "" for v in values.values()):
            # Do not validate or advance while any field is still empty.
            return (df,) + _cust_nav_outs(df, schema, idx)[:-1] + ("⚠️ Fill in all fields before continuing.",)
        for name in values:
            df.at[idx, f"validated_{name}"] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        return (df,) + _cust_nav_outs(df, schema, ni)[:-1] + ("",)

    cust_sv.click(
        on_cust_sv,
        [cust_df_s, cust_schema_s, cust_idx_s] + _cust_field_vals,
        [cust_df_s] + _cust_nav_out_keys,
        api_name=False,
    )

    # ── Custom: caption events ────────────────────────────────────────────────

    def on_cust_cap_prev(df, idx):
        ni = max(int(idx) - 1, 0)
        n  = max(len(df), 1)
        return ni, img_html(img_url_at(df, ni)), f"**{ni + 1} / {n}**", "", ""

    def on_cust_cap_skip(df, idx):
        ni = min(int(idx) + 1, max(len(df) - 1, 0))
        n  = max(len(df), 1)
        return ni, img_html(img_url_at(df, ni)), f"**{ni + 1} / {n}**", "", ""

    def on_cust_cap_gen(df, idx, api_key, model):
        url = img_url_at(df, int(idx))
        text, msg = generate_mistral_caption(url, api_key, model)
        return gr.update(value=text) if text else gr.update(), msg

    def on_cust_cap_sv(df, schema, idx, text, target_field):
        if not text.strip():
            return df, int(idx), gr.update(), gr.update(), "⚠️ Caption is empty."
        if not target_field:
            return df, int(idx), gr.update(), gr.update(), "⚠️ Select a field to save the caption to."
        idx = int(idx); df = df.copy()
        if target_field in df.columns:
            df.at[idx, target_field] = text
            val_col = f"validated_{target_field}"
            if val_col in df.columns:
                df.at[idx, val_col] = True
        ni = min(idx + 1, max(len(df) - 1, 0))
        n  = max(len(df), 1)
        return df, ni, img_html(img_url_at(df, ni)), f"**{ni + 1} / {n}**", ""

    _cust_cap_nav = [cust_cap_s, cust_cap_img, cust_cap_cnt, cust_cap_txt, cust_cap_msg]
    cust_cap_prev.click(on_cust_cap_prev, [cust_df_s, cust_cap_s], _cust_cap_nav, api_name=False)
    cust_cap_skip.click(on_cust_cap_skip, [cust_df_s, cust_cap_s], _cust_cap_nav, api_name=False)
    cust_cap_gen.click(on_cust_cap_gen,   [cust_df_s, cust_cap_s, cust_cap_key, cust_cap_mdl],
                       [cust_cap_txt, cust_cap_msg], api_name=False)
    cust_cap_sv.click(on_cust_cap_sv,
                      [cust_df_s, cust_schema_s, cust_cap_s, cust_cap_txt, cust_cap_field],
                      [cust_df_s, cust_cap_s, cust_cap_img, cust_cap_cnt, cust_cap_msg],
                      api_name=False)

    # When user picks a dropdown field in the cluster tab, populate its allowed values
    def on_cust_cl_field_change(schema, field_name):
        if not field_name or not schema:
            return gr.update(choices=[], value=None)
        for f in schema:
            if f["name"] == field_name:
                return gr.update(choices=f.get("options", []), value=None)
        return gr.update(choices=[], value=None)

    cust_cl_field.change(on_cust_cl_field_change, [cust_schema_s, cust_cl_field], [cust_cl_val], api_name=False)

    # ── Custom: cluster events ────────────────────────────────────────────────

    def on_cust_cl_prev(df, schema, cl_idx):
        ni = max(int(cl_idx) - 1, 0)
        return _cust_cl_view(df, schema, ni)

    def on_cust_cl_skip(df, schema, cl_idx):
        clusters = cluster_list(df)
        ni = min(int(cl_idx) + 1, max(len(clusters) - 1, 0))
        return _cust_cl_view(df, schema, ni)

    def on_cust_cl_apply(df, schema, cl_idx, field_name, value):
        if not field_name:
            return (df,) + _cust_cl_view(df, schema, cl_idx)[:-1] + ("⚠️ Select a field.",)
        if not value:
            return (df,) + _cust_cl_view(df, schema, cl_idx)[:-1] + ("⚠️ Select a value.",)
        clusters = cluster_list(df)
        if not clusters:
            return (df,) + _cust_cl_view(df, schema, cl_idx)[:-1] + ("⚠️ No clusters found.",)
        cl_idx = int(cl_idx)
        c_id   = clusters[cl_idx]
        df     = df.copy()
        mask   = df["clustering_id"].astype(str) == c_id
        if field_name in df.columns:
            df.loc[mask, field_name] = value
            val_col = f"validated_{field_name}"
            if val_col in df.columns:
                df.loc[mask, val_col] = True
        ni = min(cl_idx + 1, max(len(clusters) - 1, 0))
        return (df,) + _cust_cl_view(df, schema, ni)[:-1] + (f"✅ «{value}» applied to cluster {c_id}.",)

    _cust_cl_outs = [cust_cl_s, cust_cl_title, cust_cl_html, cust_cl_cnt, cust_cl_msg]
    cust_cl_prev.click(on_cust_cl_prev,   [cust_df_s, cust_schema_s, cust_cl_s], _cust_cl_outs, api_name=False)
    cust_cl_skip.click(on_cust_cl_skip,   [cust_df_s, cust_schema_s, cust_cl_s], _cust_cl_outs, api_name=False)
    cust_cl_apply.click(on_cust_cl_apply,
                        [cust_df_s, cust_schema_s, cust_cl_s, cust_cl_field, cust_cl_val],
                        [cust_df_s] + _cust_cl_outs, api_name=False)

    # ── Custom: table refresh ─────────────────────────────────────────────────
    cust_tbl_refresh.click(lambda df: (tbl_note(df), tbl_preview(df)),
                           [cust_df_s], [cust_tbl_note, cust_tbl], api_name=False)

    # ══════════════════════════════════════════════════════════════════════════
    # EVENT WIRING — Guided clustering (ITGC) mode
    # ══════════════════════════════════════════════════════════════════════════
    gc_url_load.upload(gc_load_urls, [gc_url_load],
                       [gc_df_s, gc_tbl, gc_status], api_name=False)
    gc_labeled_load.upload(gc_load_labeled, [gc_labeled_load],
                           [gc_labeled_s, gc_labeled_status], api_name=False)

    # Interactive few-shot labeller: click a category → assign + advance (instant).
    _lbl_prep_outs = ([gc_lblcats_s, gc_lbl_idx_s, gc_lbl_img, gc_lbl_info,
                       gc_labeled_status] + gc_lbl_btns)
    gc_lbl_prepare.click(_gc_lbl_prepare, [gc_df_s, gc_categories, gc_labeled_s],
                         _lbl_prep_outs, api_name=False)
    gc_lbl_prev.click(lambda df, idx: _gc_lbl_nav(df, idx, -1),
                      [gc_df_s, gc_lbl_idx_s], [gc_lbl_idx_s, gc_lbl_img, gc_lbl_info],
                      api_name=False)
    gc_lbl_skip.click(lambda df, idx: _gc_lbl_nav(df, idx, 1),
                      [gc_df_s, gc_lbl_idx_s], [gc_lbl_idx_s, gc_lbl_img, gc_lbl_info],
                      api_name=False)
    gc_lbl_undo.click(_gc_lbl_undo, [gc_labeled_s], [gc_labeled_s, gc_labeled_status],
                      api_name=False)
    _assign_outs = [gc_labeled_s, gc_lbl_idx_s, gc_lbl_img, gc_lbl_info, gc_labeled_status]
    for _bi, _btn in enumerate(gc_lbl_btns):
        _btn.click((lambda i: (lambda df, idx, labeled, cats:
                               _gc_lbl_assign(i, df, idx, labeled, cats)))(_bi),
                   [gc_df_s, gc_lbl_idx_s, gc_labeled_s, gc_lblcats_s],
                   _assign_outs, api_name=False)

    gc_run_btn.click(
        gc_run,
        [gc_df_s, gc_query, gc_categories, gc_labeled_s, gc_backend, gc_k, gc_iters,
         gc_sample, gc_key, gc_vmodel, gc_tmodel],
        [gc_df_s, gc_meta_s, gc_log, gc_sil_html, gc_concepts_md, gc_clusters_html, gc_tbl],
        api_name=False,
    )

    gc_exp_csv.click(do_gc_export_csv,   [gc_df_s], [gc_exp_file, gc_exp_msg], api_name=False)
    gc_exp_json.click(do_gc_export_json, [gc_df_s], [gc_exp_file, gc_exp_msg], api_name=False)


# ══════════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    # ssr_mode=False: Gradio's node-based SSR (enabled by default on HF Spaces)
    # is a known cause of apps hanging on an endless loading screen.
    demo.launch(ssr_mode=False, **(_STYLE_KWARGS if _GR_MAJOR >= 6 else {}))