visualization/app.py

import base64
import copy
import io
import json
import os
import random
import re
from itertools import combinations
from pathlib import Path
from typing import Dict, Any, List, Tuple, Optional
import requests  # new

BACKEND_URL = os.getenv("ATTRLLM_BACKEND_URL", "http://127.0.0.1:8000")

_DEFAULT_GRADIO_DIR = Path(os.environ.get("GRADIO_TEMP_DIR", Path.cwd() / ".gradio_tmp"))
os.environ.setdefault("GRADIO_TEMP_DIR", str(_DEFAULT_GRADIO_DIR))
_DEFAULT_GRADIO_DIR.mkdir(parents=True, exist_ok=True)


def _get_request_timeout() -> float:
    value = os.getenv("ATTRLLM_REQUEST_TIMEOUT")
    if not value:
        return 900.0
    try:
        return float(value)
    except ValueError:
        return 900.0


def _env_flag(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}


def _is_hf_spaces() -> bool:
    return bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE"))


def _public_only_mode() -> bool:
    # Default to public-only on Hugging Face Spaces, but allow override.
    return _env_flag("ATTRLLM_PUBLIC_ONLY", _is_hf_spaces())


def _public_results_file(
    dataset_key: str,
    ex_id: str,
    scalarizer: str,
    level: str,
    method: str,
) -> Path:
    results_dir = _get_results_dir()
    return (
        results_dir
        / "public"
        / dataset_key
        / ex_id
        / scalarizer
        / level
        / f"{method}.json"
    )


def _reference_results_file(
    model_size: str,
    dataset_key: str,
    ex_id: str,
    scalarizer: str,
    level: str,
) -> Path:
    results_dir = _get_results_dir()
    return (
        results_dir
        / "reference_answer"
        / model_size
        / dataset_key
        / ex_id
        / scalarizer
        / f"{level}.json"
    )


def _find_available_model_size(
    dataset_key: str,
    ex_id: str,
    scalarizer: str,
    level: str,
) -> Optional[str]:
    for size in ("large", "medium", "small"):
        if _reference_results_file(size, dataset_key, ex_id, scalarizer, level).exists():
            return size
    return None


# Fallback order when requested (scalarizer, level) is not present (e.g. on HF Space with partial results).
_FALLBACK_SCALARIZER_LEVELS: List[Tuple[str, str]] = [
    ("geomean_jointprob", "word"),
    ("semantic_similarity", "word"),
    ("geomean_jointprob", "sentence"),
    ("semantic_similarity", "sentence"),
    ("geomean_jointprob", "paragraph"),
    ("semantic_similarity", "paragraph"),
]


def _find_any_available_result(
    dataset_key: str,
    ex_id: str,
    get_res: Any,
    method: str = "shapley",
) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[Dict]]:
    """Try (model_size, scalarizer, level) fallbacks; return (size, scalarizer, level, result_dict) or (None,)*4."""
    for size in ("small", "medium", "large"):
        for scalarizer, level in _FALLBACK_SCALARIZER_LEVELS:
            try:
                result = get_res(size, dataset_key, ex_id, scalarizer=scalarizer, feature_level=level) or {}
                payload = result.get(method, {})
                if payload and (payload.get("features") or payload.get("heatmap")):
                    return (size, scalarizer, level, result)
            except Exception:
                continue
    return (None, None, None, None)


def _parse_sparse_key(raw_key: str) -> Tuple[int, ...]:
    key = str(raw_key).strip()
    if not key:
        return ()
    return tuple(int(part) for part in key.split(",") if part != "")


def _normalize_public_payload_fallback(data: Dict[str, Any], method: str, top_k: int = 10) -> Dict[str, Any]:
    """Convert your JSON (features list + meta + mobius_dict) to UI display format. mobius_dict can be empty."""
    if not isinstance(data, dict):
        return {}
    features = data.get("features")
    mobius_raw = data.get("mobius_dict") if isinstance(data.get("mobius_dict"), dict) else {}
    if not isinstance(features, list) or not features:
        return {}
    method = (method or "shapley").lower()
    if method not in {"shapley", "banzhaf", "influence"}:
        method = "shapley"
    mobius_sparse: Dict[Tuple[int, ...], float] = {}
    for key, raw_val in mobius_raw.items():
        try:
            val = float(raw_val)
        except Exception:
            continue
        try:
            loc = _parse_sparse_key(str(key))
        except Exception:
            continue
        mobius_sparse[tuple(sorted(loc))] = val
    token_scores: Dict[str, float] = {}
    index_scores: Dict[int, float] = {}
    pairwise_acc: Dict[Tuple[int, int], float] = {}
    if method == "influence" and mobius_to_influence is not None and influence_interactions is not None:
        singleton_values = mobius_to_influence(mobius_sparse)
        for loc, val in singleton_values.items():
            if len(loc) != 1:
                continue
            idx = loc[0]
            if 0 <= idx < len(features):
                feat_name = str(features[idx])
                token_scores[feat_name] = float(val)
                index_scores[idx] = float(val)
        pair_list = influence_interactions(mobius_sparse, order=2, top_k=top_k)
        for loc, val in pair_list:
            if len(loc) == 2:
                i, j = sorted(loc)
                pairwise_acc[(i, j)] = float(val)
    else:
        for loc, val in mobius_sparse.items():
            k = len(loc)
            if k == 0:
                continue
            sw = 1.0 / float(k) if method == "shapley" else 1.0 / float(2 ** (k - 1))
            for idx in loc:
                if 0 <= idx < len(features):
                    feat_name = str(features[idx])
                    token_scores[feat_name] = token_scores.get(feat_name, 0.0) + sw * val
                    index_scores[idx] = index_scores.get(idx, 0.0) + sw * val
            if k >= 2:
                pw = 1.0 / float(k - 1) if method == "shapley" else 1.0 / float(2 ** (k - 2))
                for i, j in combinations(sorted(loc), 2):
                    pairwise_acc[(i, j)] = pairwise_acc.get((i, j), 0.0) + pw * val
    unique_feature_labels = [str(x) for x in features]
    sorted_pairs = sorted(pairwise_acc.items(), key=lambda kv: abs(kv[1]), reverse=True)
    if top_k and top_k > 0:
        sorted_pairs = sorted_pairs[:top_k]
    pairwise = {
        "%s|%s" % (unique_feature_labels[i], unique_feature_labels[j]): float(v)
        for (i, j), v in sorted_pairs
        if 0 <= i < len(unique_feature_labels) and 0 <= j < len(unique_feature_labels)
    }
    pairwise_interactions = [
        {"features": [unique_feature_labels[i], unique_feature_labels[j]], "value": float(v)}
        for (i, j), v in sorted_pairs
        if 0 <= i < len(unique_feature_labels) and 0 <= j < len(unique_feature_labels)
    ]
    normalized = dict(data)
    normalized["token_scores"] = token_scores
    normalized["pairwise"] = pairwise
    normalized["pairwise_interactions"] = pairwise_interactions
    normalized["features"] = [
        {"feature": str(features[i]), "value": float(index_scores.get(i, 0.0)), "index": i}
        for i in range(len(features))
    ]
    normalized["feature_texts"] = [str(x) for x in features]
    return normalized


def _public_get_result_from_file(
    model_size: str,
    dataset: str,
    ex_id: str,
    scalarizer: Optional[str] = None,
    feature_level: Optional[str] = None,
) -> Dict[str, Any]:
    """Load reference_answer result from disk when loader.results.get_result_by_id is unavailable (e.g. on Space)."""
    scalarizer = (scalarizer or "").strip()
    feature_level = (feature_level or "").strip()
    if not scalarizer or not feature_level:
        return {}
    levels_to_try = [feature_level] + [l for l in ("word", "sentence", "paragraph") if l != feature_level]
    for lvl in levels_to_try:
        path = _reference_results_file(model_size, dataset, ex_id, scalarizer, lvl)
        if not path.exists() or os.getenv("SPACE_ID"):
            try:
                from loader.results import _maybe_download_from_space
                path = _maybe_download_from_space(path, force_download=True) or path
            except Exception:
                pass
        if not path.exists():
            continue
        try:
            with path.open("r", encoding="utf-8") as f:
                data = json.load(f)
        except Exception:
            continue
        if not isinstance(data, dict):
            continue
        # Your JSON: features (list) + meta + mobius_dict (can be empty). Always convert to UI format.
        norm_s = _normalize_public_payload_fallback(copy.deepcopy(data), "shapley")
        norm_b = _normalize_public_payload_fallback(copy.deepcopy(data), "banzhaf")
        norm_i = _normalize_public_payload_fallback(copy.deepcopy(data), "influence")
        if not norm_s and not norm_b and not norm_i:
            continue
        return {
            "shapley": norm_s,
            "banzhaf": norm_b,
            "influence": norm_i,
            "meta": {
                "dataset": dataset,
                "example_id": ex_id,
                "model_size": model_size,
                "source_layout": "results/reference_answer/{model_size}/{dataset}/{example_id}/{scalarizer}/{feature_level}.json",
            },
        }
    return {}


_FALLBACK_DATASET_FILES: Dict[str, str] = {
    "bar_exam": "BarExam_qa.csv",
    "causal_judgment": "bbh_causal_judgement.csv",
    "snarks": "bbh_snarks.csv",
    "bbq_disamb": "BBQ_disamb.csv",
    "cnn_dailymail": "CNN_dailymail.csv",
    "drop": "drop.csv",
    "esnli": "eSNLI.csv",
    "fever": "fever.csv",
    "hotpot_qa": "hotpot_qa.csv",
    "medical_qa": "medical_qa.csv",
}


def _fallback_datasets_dir() -> Path:
    return (_REPO_ROOT / "datasets").resolve()


def _fallback_pick_first_nonempty(raw: Dict[str, str], candidates: List[str]) -> str:
    for c in candidates:
        val = raw.get(c)
        if val is not None and str(val).strip() != "":
            return str(val)
    return ""


def _fallback_load_dataset(dataset_key: str, max_rows: int = 10) -> List[Dict[str, str]]:
    import csv

    filename = _FALLBACK_DATASET_FILES.get(dataset_key)
    if not filename:
        return []
    path = _fallback_datasets_dir() / filename
    if not path.exists():
        return []

    rows: List[Dict[str, str]] = []
    with path.open("r", encoding="utf-8", errors="replace", newline="") as f:
        reader = csv.DictReader(f)
        for i, raw in enumerate(reader, start=1):
            ex_id = raw.get("id") or raw.get("example_id") or raw.get("uid") or f"example_{i}"
            context = _fallback_pick_first_nonempty(raw, [
                "Context", "context",
                "passage", "article", "story", "premise",
                "paragraph", "document", "sentence1", "sent1", "background",
            ])
            prompt = _fallback_pick_first_nonempty(raw, [
                "Prompt", "prompt",
                "question", "input", "query",
                "sentence2", "sent2", "hypothesis",
                "qa_question", "title",
            ])
            answer = _fallback_pick_first_nonempty(raw, [
                "Answer", "answer",
                "target", "gold", "label", "output", "reference",
                "highlights",
            ])
            ex = {
                "id": str(ex_id),
                "context": context,
                "prompt": prompt,
            }
            if answer:
                ex["answer"] = answer
            rows.append(ex)
            if len(rows) >= max_rows:
                break
    return rows


REQUEST_TIMEOUT = _get_request_timeout()

SCALARIZER_CHOICES = [
    ("Semantic Similarity (y vs y_S)", "semantic_similarity"),
    ("LogProb", "logprob"),
    ("JointProb", "jointprob"),
    ("GeoMean JointProb", "geomean_jointprob"),
    ("Half SimLog", "half_simlog"),
]

PUBLIC_SCALARIZER_CHOICES = [
    ("Semantic Similarity", "semantic_similarity"),
    ("Perplexity", "geomean_jointprob"),
]

DATASET_DISPLAY_LABELS = {
    "bar_exam": "Bar Exam Questions",
    "bbq_disamb": "BBQ Disambiguation",
    "causal_judgment": "Causal Judgment",
    "cnn_dailymail": "CNN / DailyMail Summaries",
    "drop": "DROP Reading Comprehension",
    "esnli": "e-SNLI Natural Language Inference",
    "fever": "FEVER Fact Checking",
    "hotpot_qa": "HotpotQA Multi-hop Questions",
    "medical_qa": "Medical Questions",
    "snarks": "Snarks",
}

import sys
_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))


def _get_results_dir() -> Path:
    """Resolve results directory: env, repo root, or on HF Space fallback to cwd/results."""
    env_dir = os.getenv("ATTRLLM_RESULTS_DIR")
    if env_dir:
        return Path(env_dir).resolve()
    default = (_REPO_ROOT / "results").resolve()
    if default.exists():
        return default
    if _is_hf_spaces():
        cwd_results = (Path.cwd() / "results").resolve()
        if cwd_results.exists():
            return cwd_results
    return default


import gradio as gr
from PIL import Image

from .components.model_selector import (
    create_model_selector,
    create_multimodal_model_selector,
    create_feature_level_selector,
    create_attribution_method_toggle,
)
from .components.example_browser import create_dataset_selector, create_example_browser
from .components.results_display import create_results_display, update
from .plotting.heatmap import create_interactive_text_heatmap
from .plotting.interactions import (
    plot_top_interactions,
    plot_interaction_matrix,
    create_interaction_token_view,
)
from .plotting.text_interactions import create_text_interaction_html
from .plotting.mm_interactions import create_multimodal_interaction_html
from .build_info import BUILD_ID, BUILD_TS


def _raise_backend_error(resp: requests.Response, label: str) -> None:
    detail = resp.text
    try:
        detail = resp.json().get("detail", detail)
    except Exception:
        pass
    raise gr.Error(f"{label} failed ({resp.status_code}). {detail}")

# backend API imports
try:  # loader data APIs are required for public mode
    from loader.data import (
        get_example_by_id,
        get_examples,
        list_datasets,
        list_datasets_with_display_names,
        list_dataset_display_names,
        get_dataset_display_name,
        get_dataset_key_from_display_name,
    )
except Exception:  # pragma: no cover - optional at runtime
    get_example_by_id = None
    get_examples = None
    list_datasets = None
    list_datasets_with_display_names = None
    list_dataset_display_names = None
    get_dataset_display_name = None
    get_dataset_key_from_display_name = None

try:
    from loader.results import get_result_by_id
except Exception:  # pragma: no cover
    get_result_by_id = None

try:
    from loader.models import get_model
except Exception:  # pragma: no cover
    get_model = None

try:  # attribution stack is optional (dev mode)
    from attribution.masker import get_masker, mask_text
    from attribution.proxyspex import run_proxyspex
    from attribution.image_masker import supports_superpixel
    from attribution.utils import (
        mobius_to_shapley,
        shapley_interactions,
        mobius_to_banzhaf,
        banzhaf_interactions,
        mobius_to_influence,
        influence_interactions,
        mobius_to_fourier,
    )
except Exception:  # pragma: no cover
    get_masker = None
    mask_text = None
    run_proxyspex = None
    supports_superpixel = None
    mobius_to_shapley = None
    shapley_interactions = None
    mobius_to_banzhaf = None
    banzhaf_interactions = None
    mobius_to_influence = None
    influence_interactions = None
    mobius_to_fourier = None


_ANSWER_FIELDS = (
    "correct_answer",
    "answer",
    "target",
    "completion",
    "label",
)
_ALLOWED_METHODS = {"shapley", "banzhaf", "influence"}
_ALLOWED_LEVELS = {"word", "sentence", "paragraph"}


def _ensure_backend(name: str, fn: Optional[Any]):
    if fn is None:
        raise RuntimeError(
            f"{name} is unavailable. Ensure the backend modules are installed and importable."
        )
    return fn


def _html_component(label: str) -> gr.HTML:
    try:
        return gr.HTML(label=label, sanitize_html=False)
    except TypeError:
        return gr.HTML(label=label)


def _encode_image_to_b64(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _extract_answer(record: Dict[str, Any]) -> str:
    for field in _ANSWER_FIELDS:
        val = record.get(field)
        if val:
            return str(val)
    return ""


def _coerce_feature_tuple(raw_key: Any) -> Tuple[str, ...]:
    if isinstance(raw_key, tuple):
        return tuple(str(item) for item in raw_key)
    if isinstance(raw_key, list):
        return tuple(str(item) for item in raw_key)
    if isinstance(raw_key, str):
        for sep in ("·", "|", ",", "×"):
            if sep in raw_key:
                parts = [chunk.strip() for chunk in raw_key.split(sep) if chunk.strip()]
                if parts:
                    return tuple(parts)
        return (raw_key.strip(),)
    return (str(raw_key),)


# def _normalize_interactions(raw: Any) -> List[Tuple[Tuple[str, ...], float]]:
#     items: List[Any]
#     if raw is None:
#         return []
#     if isinstance(raw, dict):
#         items = list(raw.items())
#     else:
#         items = list(raw)

#     normalized: List[Tuple[Tuple[str, ...], float]] = []
#     for feats, value in items:
#         try:
#             numeric = float(value)
#         except Exception:
#             continue
#         normalized.append((_coerce_feature_tuple(feats), numeric))
#     return normalized
def _normalize_interactions(raw: Any) -> List[Tuple[Tuple[str, ...], float]]:
    """
    Make a best-effort guess at interaction structure.

    Supported shapes:
      - { key: float }
      - { key: {"value": float, "score": ...} }
      - [ (key, float), ... ]
      - [ (key, {"value": float}), ... ]
      - [ {"features": [...], "value": float}, ... ]  (this is mostly handled elsewhere)
    """
    if raw is None:
        return []

    items: List[Any] = []

    if isinstance(raw, dict):
        # e.g. { key: float } or { key: {"value": ...} }
        for k, v in raw.items():
            items.append((k, v))
    elif isinstance(raw, list):
        items = list(raw)
    else:
        return []

    normalized: List[Tuple[Tuple[str, ...], float]] = []

    for item in items:
        # Case 1: dict-style item with explicit fields
        if isinstance(item, dict):
            feats = item.get("features") or item.get("indices") or item.get("pair") or item.get("key")
            val = item.get("value", item.get("score", 0.0))
        else:
            # Case 2: tuple/list pair (feats, value)
            try:
                feats, val = item
            except Exception:
                continue

        # If value itself is a dict, dig out "value" / "score"
        if isinstance(val, dict):
            val = val.get("value", val.get("score", 0.0))

        try:
            numeric = float(val)
        except Exception:
            continue

        feats_tuple = _coerce_feature_tuple(feats)
        if feats_tuple:
            normalized.append((feats_tuple, numeric))

    return normalized

def _resolve_marginals(payload: Dict[str, Any]) -> Dict[str, float]:
    for key in ("marginals", "token_scores", "values", "scores"):
        data = payload.get(key)
        if isinstance(data, dict):
            normalized: Dict[str, float] = {}
            for k, v in data.items():
                try:
                    normalized[str(k)] = float(v)
                except Exception:
                    continue
            return normalized
    return {}


def _resolve_features(payload: Dict[str, Any], marginals: Dict[str, float]) -> List[str]:
    features = payload.get("features")
    if isinstance(features, list):
        return [str(f) for f in features]
    if marginals:
        return list(marginals.keys())
    return []


def _extract_interactions_from_response(
    data_int: Dict[str, Any],
    method: str,
    features: List[str],
) -> List[Tuple[Tuple[str, ...], float]]:
    inter_list: List[Tuple[Tuple[str, ...], float]] = []

    method_key = (method or "shapley").lower()
    method_block = data_int.get(method_key) or data_int

    raw_interactions = None
    if isinstance(method_block, dict):
        for key in ("interactions", "pairwise_interactions", "interactions_2", "pairwise", "data"):
            if key in method_block:
                raw_interactions = method_block.get(key)
                break
        if raw_interactions is None:
            raw_interactions = method_block
    else:
        raw_interactions = method_block

    # List-of-dicts or list-of-pairs shape
    if isinstance(raw_interactions, list) and raw_interactions:
        if isinstance(raw_interactions[0], dict):
            for item in raw_interactions:
                feats = (
                    item.get("feature_list")
                    or item.get("features")
                    or item.get("indices")
                    or item.get("pair")
                    or []
                )
                val = None
                for key_val in ("value", "score", "attribution", "weight"):
                    if key_val in item:
                        try:
                            val = float(item[key_val])
                            break
                        except Exception:
                            continue
                if val is None:
                    continue

                if isinstance(feats, list) and feats and isinstance(feats[0], int):
                    feat_names = tuple(
                        features[i] for i in feats
                        if isinstance(i, int) and 0 <= i < len(features)
                    )
                else:
                    feat_names = _coerce_feature_tuple(feats)

                if feat_names:
                    inter_list.append((feat_names, val))
        elif (
            isinstance(raw_interactions[0], (list, tuple))
            and len(raw_interactions[0]) == 2
        ):
            for item in raw_interactions:
                if not isinstance(item, (list, tuple)) or len(item) != 2:
                    continue
                feats_raw, val_raw = item
                try:
                    val = float(val_raw)
                except Exception:
                    continue

                feat_names: Tuple[str, ...] = ()
                if isinstance(feats_raw, (list, tuple)) and feats_raw:
                    if all(isinstance(i, int) for i in feats_raw):
                        feat_names = tuple(
                            features[i] for i in feats_raw
                            if 0 <= i < len(features)
                        )
                    else:
                        feat_names = _coerce_feature_tuple(feats_raw)
                elif isinstance(feats_raw, str):
                    feat_names = _coerce_feature_tuple(feats_raw)

                if feat_names:
                    inter_list.append((feat_names, val))

    # Dict shape, e.g. {"(0,2)": 528.0, ...}
    if not inter_list and isinstance(raw_interactions, dict):
        metadata_keys = {"method", "order", "scalarizer", "embedding_model"}
        for k, v in raw_interactions.items():
            if str(k) in metadata_keys:
                continue
            val = None
            if isinstance(v, (int, float)):
                val = float(v)
            elif isinstance(v, dict):
                for key_val in ("value", "score", "attribution", "weight"):
                    if key_val in v:
                        try:
                            val = float(v[key_val])
                            break
                        except Exception:
                            continue
            if val is None:
                continue

            k_str = str(k)
            idxs = []
            try:
                import re as _re
                idxs = [int(x) for x in _re.findall(r"\d+", k_str)]
            except Exception:
                idxs = []

            if idxs:
                names: List[str] = []
                for idx in idxs:
                    if 0 <= idx < len(features):
                        names.append(features[idx])
                if names:
                    feat_names = tuple(names)
                else:
                    feat_names = _coerce_feature_tuple(k_str)
            else:
                feat_names = _coerce_feature_tuple(k_str)

            inter_list.append((feat_names, val))

    # Flatten numerics arbitrarily (last resort)
    if not inter_list and raw_interactions is not None:
        flat: List[Tuple[Tuple[str, ...], float]] = []

        def _collect(obj: Any, prefix: Tuple[str, ...] = ()) -> None:
            if isinstance(obj, (int, float)):
                flat.append((prefix or ("<interaction>",), float(obj)))
            elif isinstance(obj, list):
                for i, item in enumerate(obj):
                    _collect(item, prefix + (f"[{i}]",))
            elif isinstance(obj, dict):
                for kk, vv in obj.items():
                    _collect(vv, prefix + (str(kk),))

        _collect(raw_interactions)
        inter_list = flat

    return inter_list


def _labels_from_regions(regions: List[Dict[str, Any]]) -> List[str]:
    labels: List[str] = [""] * len(regions)
    for region in regions:
        try:
            idx = int(region.get("index", 0))
        except Exception:
            continue
        if idx < 0 or idx >= len(labels):
            continue
        labels[idx] = str(region.get("label") or f"Region {idx + 1}")
    for idx, label in enumerate(labels):
        if not label:
            labels[idx] = f"Region {idx + 1}"
    return labels


def _interaction_dicts_to_pairs(
    interactions: List[Dict[str, Any]],
    labels: List[str],
    *,
    order: int | None = None,
) -> List[Tuple[Tuple[str, ...], float]]:
    pairs: List[Tuple[Tuple[str, ...], float]] = []
    for item in interactions:
        indices = item.get("indices")
        if not indices:
            continue
        if order is not None and len(indices) != order:
            continue
        try:
            value = float(item.get("value", 0.0))
        except Exception:
            continue
        feats = tuple(labels[int(i)] for i in indices if int(i) < len(labels))
        if feats:
            pairs.append((feats, value))
    return pairs


def _interaction_dicts_to_table(
    interactions: List[Dict[str, Any]],
    labels: List[str],
) -> List[List[Any]]:
    rows: List[List[Any]] = []
    for item in interactions:
        indices = item.get("indices")
        if not indices:
            continue
        try:
            value = float(item.get("value", 0.0))
        except Exception:
            continue
        feats = [labels[int(i)] for i in indices if int(i) < len(labels)]
        if feats:
            rows.append([" × ".join(feats), value, len(indices)])
    return rows


def _feature_display_label(
    feature: Dict[str, Any],
    region_labels: List[str],
) -> str:
    raw = str(feature.get("feature", ""))
    modality = feature.get("modality") or ""
    ref_index = int(feature.get("ref_index", 0))
    label = raw.split(":", 1)[1] if ":" in raw else raw
    if modality == "image":
        if 0 <= ref_index < len(region_labels):
            return region_labels[ref_index]
    return label or raw


def _extract_feature_series(payload: Dict[str, Any]) -> Tuple[List[str], List[float]]:
    """
    Try to recover an ordered pair of (feature labels, values) from a backend payload.
    This keeps duplicates in order (appending suffixes later) so word-level tokens
    don't collapse to a single entry.
    """
    features: List[str] = []
    values: List[float] = []

    feature_entries = payload.get("features")
    if isinstance(feature_entries, list) and feature_entries and isinstance(feature_entries[0], dict):
        for idx, entry in enumerate(feature_entries, start=1):
            raw_feat = (
                entry.get("feature")
                or entry.get("token")
                or entry.get("text")
                or entry.get("label")
                or ""
            )
            if not raw_feat:
                raw_feat = f"feature_{idx}"

            val = entry.get("value")
            if val is None:
                for key in ("score", "attribution", "weight"):
                    if key in entry:
                        val = entry[key]
                        break
            try:
                values.append(float(val if val is not None else 0.0))
            except Exception:
                values.append(0.0)
            features.append(str(raw_feat))

    if not features:
        heat = payload.get("heatmap") or {}
        tokens = heat.get("tokens") or heat.get("features")
        scores = heat.get("values") or heat.get("scores")
        if isinstance(tokens, list) and isinstance(scores, list) and len(tokens) == len(scores):
            features = [str(token if token is not None else f"feature_{idx + 1}") for idx, token in enumerate(tokens)]
            tmp_vals: List[float] = []
            for score in scores:
                try:
                    tmp_vals.append(float(score))
                except Exception:
                    tmp_vals.append(0.0)
            values = tmp_vals

    if not features:
        marginals = _resolve_marginals(payload)
        if marginals:
            features = list(marginals.keys())
            values = [float(marginals[key]) for key in features]

    if not features:
        return [], []

    unique_features = _assign_unique_labels(features)
    return unique_features, values


def _resolve_interactions(payload: Dict[str, Any], order: int) -> List[Tuple[Tuple[str, ...], float]]:
    candidates = [f"interactions_{order}"]
    if order == 2:
        candidates += ["pairwise", "pairwise_interactions", "interactions2"]
    elif order == 3:
        candidates += ["higher_order", "triple_interactions", "interactions3"]

    for key in candidates:
        raw = payload.get(key)
        normalized = _normalize_interactions(raw)
        if normalized:
            return normalized
    return []


def _fallback_pairwise_from_values(
    features: List[str],
    values: List[float],
    max_edges: int = 40,
) -> List[Tuple[Tuple[str, ...], float]]:
    """
    Generate synthetic pairwise links by connecting neighboring tokens.
    Used when the backend provides no explicit interactions.
    """
    n = min(len(features), len(values))
    if n < 2:
        return []
    edges: List[Tuple[Tuple[str, ...], float]] = []
    for idx in range(n - 1):
        weight = 0.5 * (values[idx] + values[idx + 1])
        edges.append(((features[idx], features[idx + 1]), weight))
    edges.sort(key=lambda item: abs(item[1]), reverse=True)
    return edges[:max_edges]


def _resolve_pairwise(
    payload: Dict[str, Any],
    features: Optional[List[str]] = None,
    feature_values: Optional[List[float]] = None,
) -> List[Tuple[Tuple[str, ...], float]]:
    """Convenience helper to always pull order-2 interactions if present."""
    pairwise = _resolve_interactions(payload, 2)
    if pairwise:
        return pairwise
    # Some payloads store generic "interactions" lists that mix orders.
    mixed = payload.get("interactions")
    normalized = _normalize_interactions(mixed)
    if normalized:
        return [item for item in normalized if len(item[0]) == 2]
    if features and feature_values:
        return _fallback_pairwise_from_values(features, feature_values)
    return []


def _normalize_method(method: Optional[str]) -> str:
    method = (method or "shapley").lower()
    return method if method in _ALLOWED_METHODS else "shapley"


def _normalize_level(level: Optional[str]) -> str:
    level = (level or "sentence").lower()
    return level if level in _ALLOWED_LEVELS else "sentence"


def _normalize_model_size(model_size: Optional[str]) -> str:
    raw = (model_size or "small").strip()
    lowered = raw.lower()
    if lowered in {"small", "medium", "large"}:
        return lowered
    if "small" in lowered:
        return "small"
    if "medium" in lowered:
        return "medium"
    if "large" in lowered:
        return "large"
    return "small"


def _assign_unique_labels(chunks: List[str]) -> List[str]:
    counts: Dict[str, int] = {}
    labels: List[str] = []
    for idx, chunk in enumerate(chunks):
        normalized = " ".join((chunk or "").split())
        if not normalized:
            normalized = f"<chunk {idx + 1}>"
        counts[normalized] = counts.get(normalized, 0) + 1
        suffix = f" ({counts[normalized]})" if counts[normalized] > 1 else ""
        labels.append(f"{normalized}{suffix}")
    return labels


def _strip_occurrence_suffix(text: str) -> str:
    text = text or ""
    if text.endswith(")") and " (" in text:
        base, _, tail = text.rpartition(" (")
        if tail[:-1].isdigit():
            return base
    return text


def _pairwise_to_index_interactions(
    pairwise: List[Tuple[Tuple[str, ...], float]],
    features: List[str],
) -> List[Dict[str, Any]]:
    feature_index = {feat: idx for idx, feat in enumerate(features)}
    base_index: Dict[str, int] = {}
    for idx, feat in enumerate(features):
        base_index.setdefault(_strip_occurrence_suffix(feat), idx)

    interactions: List[Dict[str, Any]] = []
    for feats, val in pairwise:
        if len(feats) != 2:
            continue
        a, b = feats
        a_idx = None
        b_idx = None
        if isinstance(a, (int, float)) and isinstance(b, (int, float)):
            a_idx = int(a)
            b_idx = int(b)
        else:
            try:
                a_idx = int(str(a))
                b_idx = int(str(b))
            except ValueError:
                a_idx = feature_index.get(a) or base_index.get(_strip_occurrence_suffix(str(a)))
                b_idx = feature_index.get(b) or base_index.get(_strip_occurrence_suffix(str(b)))
        if a_idx is None or b_idx is None:
            continue
        if a_idx < 0 or b_idx < 0 or a_idx >= len(features) or b_idx >= len(features):
            continue
        interactions.append({"indices": [a_idx, b_idx], "value": float(val)})
    return interactions


def _locate_spans(text: str, segments: List[str]) -> List[Tuple[int, int]]:
    spans: List[Tuple[int, int]] = []
    cursor = 0
    for segment in segments:
        if not segment:
            continue
        idx = text.find(segment, cursor)
        if idx == -1:
            idx = cursor
        end = idx + len(segment)
        spans.append((idx, end))
        cursor = end
    return spans


def _chunk_text_for_visualization(
    context: str,
    level: str,
) -> Tuple[List[str], List[Tuple[int, int]], str]:
    """
    Split input text into feature chunks and spans for visualization.
    Falls back to the demo text if context is empty.
    """
    text = context or _DEMO_TEXT
    level = _normalize_level(level)

    if level == "word":
        matches = list(re.finditer(r"\S+", text))
        chunks = [m.group(0) for m in matches]
        spans = [(m.start(), m.end()) for m in matches]
    elif level == "paragraph":
        parts = [seg for seg in re.split(r"\n\s*\n+", text) if seg.strip()]
        spans = _locate_spans(text, parts)
        chunks = parts[: len(spans)]
    else:  # sentence-level default
        parts = [seg for seg in re.split(r"(?<=[.!?])\s+", text) if seg.strip()]
        spans = _locate_spans(text, parts)
        chunks = parts[: len(spans)]

    if not chunks:
        chunks = [text]
        spans = [(0, len(text))]

    features = _assign_unique_labels(chunks)
    return features, spans, text


def _generate_synthetic_marginals(
    features: List[str],
    rng: random.Random,
) -> Dict[str, float]:
    if not features:
        return {}
    max_len = max(len(f) for f in features) or 1
    marginals: Dict[str, float] = {}
    denom = max(1, len(features) - 1)
    for idx, feat in enumerate(features):
        length_factor = len(feat) / max_len
        position_factor = 1 - (idx / denom if denom else 0)
        noise = rng.uniform(-0.25, 0.25)
        value = (length_factor - 0.5) * 0.6 + (position_factor - 0.5) * 0.4 + noise
        marginals[feat] = round(value, 4)
    return marginals


def _generate_synthetic_interactions(
    features: List[str],
    marginals: Dict[str, float],
    rng: random.Random,
) -> Dict[int, List[Tuple[Tuple[str, ...], float]]]:
    interactions: Dict[int, List[Tuple[Tuple[str, ...], float]]] = {2: [], 3: []}
    for i in range(len(features) - 1):
        pair = (features[i], features[i + 1])
        base = (marginals.get(pair[0], 0.0) + marginals.get(pair[1], 0.0)) / 2
        interactions[2].append((pair, round(base + rng.uniform(-0.1, 0.1), 4)))
    for i in range(len(features) - 2):
        triple = (features[i], features[i + 1], features[i + 2])
        base = sum(marginals.get(feat, 0.0) for feat in triple) / 3
        interactions[3].append((triple, round(base + rng.uniform(-0.1, 0.1), 4)))
    return interactions


def _synthetic_attribution_pipeline(
    context: str,
    prompt: str,
    answer: str,
    *,
    method: str,
    level: str,
    order: int,
    reason: Optional[str] = None,
) -> Tuple[Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any]:
    text_source = context or prompt or answer or _DEMO_TEXT
    features, spans, text = _chunk_text_for_visualization(text_source, level)
    seed = hash((text_source, method, level, order)) & 0xFFFFFFFF
    rng = random.Random(seed)
    marginals = _generate_synthetic_marginals(features, rng)
    interactions = _generate_synthetic_interactions(features, marginals, rng)

    html = None
    if len(spans) == len(features):
        html = create_interactive_text_heatmap(
            text,
            spans,
            [marginals.get(f, 0.0) for f in features],
            method=method,
        )

    meta = {
        "mode": "synthetic",
        "reason": reason or "Attribution backend unavailable; showing mock data.",
        "method": method,
        "feature_level": level,
        "interaction_order": order,
        "feature_count": len(features),
    }

    inter_list = interactions.get(order, [])
    pairwise_for_tokens = interactions.get(2, []) if order != 2 else inter_list
    if not pairwise_for_tokens:
        pairwise_for_tokens = _fallback_pairwise_from_values(
            features,
            [marginals.get(f, 0.0) for f in features],
        )
    text_interaction_html = create_text_interaction_html(
        features,
        [marginals.get(f, 0.0) for f in features],
        _pairwise_to_index_interactions(pairwise_for_tokens, features),
        top_k=20,
        threshold=0.0,
        method=method,
    )
    figs = {
        "interactions": plot_top_interactions(inter_list, order=order, method=method),
    }

    return update(
        figs=figs,
        meta=meta,
        html=html,
        interaction_text_html=text_interaction_html,
        scoring_target_source="answer_input" if answer else "model_output",
        scoring_target_text=answer or "",
        reference_answer=answer or "",
        unmasked_answer="",
        debug_scores=None,
        scalarizer_used="logprob",
        score_full=None,
        score_empty=None,
        y_len_tokens=None,
    )


# def _compute_live_attributions(**kwargs) -> Tuple[Any, Any, Any, Any, Any]:
#     """
#     Placeholder for the real ProxySPEX + perplexity pipeline.
#     Raises until the attribution backend is implemented.
#     """
#     missing = [
#         name
#         for name, fn in {
#             "get_model": get_model,
#             "get_masker": get_masker,
#             "mask_text": mask_text,
#             "run_proxyspex": run_proxyspex,
#             "mobius_to_shapley": mobius_to_shapley,
#             "mobius_to_banzhaf": mobius_to_banzhaf,
#             "shapley_interactions": shapley_interactions,
#             "banzhaf_interactions": banzhaf_interactions,
#         }.items()
#         if fn is None
#     ]
#     if missing:
#         raise RuntimeError(
#             "Missing backend dependencies: " + ", ".join(sorted(missing))
#         )
#     raise NotImplementedError(
#         "Live attribution pipeline not wired yet. Integrate once ProxySPEX is ready."
#     )
def _compute_live_attributions(
    *,
    context: str,
    prompt: str,
    correct_answer: str,
    model_size: str,
    level: str,
    method: str,
    order: int,
    scalarizer: str = "logprob",
    embedding_model: str | None = None,
    progress=None,
) -> Tuple[Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any, Any]:
    """
    Call the FastAPI /api/attributions + /api/interactions backends and turn
    the JSON into figures / table / HTML for Gradio.
    This version is very defensive and tries hard to extract interactions
    from whatever shape the backend returns.
    """
    method = _normalize_method(method)
    level = _normalize_level(level)
    order = 3 if int(order or 2) >= 3 else 2

    context = context or ""
    prompt = prompt or ""
    correct_answer = correct_answer or ""
    text_source = context or prompt or correct_answer or _DEMO_TEXT

    payload = {
        "context": context,
        "answer": correct_answer,
        "reference_answer": correct_answer,
        "prompt": prompt,
        "method": method,
        "mask_level": level,
        "order": int(order),
        "model_size": model_size,
        "scalarizer": scalarizer,
        "embedding_model": embedding_model,
        "debug": False,
    }

    if progress is not None:
        progress(0.1, desc="Calling attribution backend")

    # ---------- 1. /api/attributions  ----------
    url_attr = BACKEND_URL.rstrip("/") + "/api/attributions"
    try:
        resp_attr = requests.post(url_attr, json=payload, timeout=REQUEST_TIMEOUT)
    except requests.exceptions.ReadTimeout as exc:
        raise gr.Error(
            "Attribution request timed out. The backend may still be running. "
            "Consider reducing feature granularity or set ATTRLLM_REQUEST_TIMEOUT to a higher value."
        ) from exc
    if resp_attr.status_code >= 400:
        _raise_backend_error(resp_attr, "Attribution request")
    data_attr = resp_attr.json()

    if progress is not None:
        progress(0.35, desc="Received attribution payload")

    # ---------- 2. FEATURES + MARGINAL VALUES ----------
    features, feature_values = _extract_feature_series(data_attr)
    if not features:
        features = ["<no features>"]
        feature_values = [0.0]
    marginals = {feat: float(feature_values[idx]) for idx, feat in enumerate(features)}

    # ---------- 3. /api/interactions  ----------
    if progress is not None:
        progress(0.45, desc="Calling interactions backend")

    url_int = BACKEND_URL.rstrip("/") + "/api/interactions"
    try:
        resp_int = requests.post(url_int, json=payload, timeout=REQUEST_TIMEOUT)
    except requests.exceptions.ReadTimeout as exc:
        raise gr.Error(
            "Interaction request timed out. The backend may still be running. "
            "Consider reducing order or set ATTRLLM_REQUEST_TIMEOUT to a higher value."
        ) from exc
    if resp_int.status_code >= 400:
        _raise_backend_error(resp_int, "Interaction request")
    data_int = resp_int.json()

    # DEBUG: see top-level keys
    print("data_int keys:", list(data_int.keys()))

    inter_list_all = _extract_interactions_from_response(data_int, method, features)
    pairwise_for_network = [item for item in inter_list_all if len(item[0]) == 2]

    inter_list = inter_list_all
    if inter_list:
        filtered: List[Tuple[Tuple[str, ...], float]] = []
        for feats, val in inter_list:
            if len(feats) == order:
                filtered.append((feats, val))
        if filtered:
            inter_list = filtered

    if order != 2 and not pairwise_for_network:
        try:
            payload_pair = dict(payload)
            payload_pair["order"] = 2
            try:
                resp_pair = requests.post(url_int, json=payload_pair, timeout=REQUEST_TIMEOUT)
            except requests.exceptions.ReadTimeout as exc:
                raise gr.Error(
                    "Interaction request timed out. The backend may still be running. "
                    "Consider reducing order or set ATTRLLM_REQUEST_TIMEOUT to a higher value."
                ) from exc
            if resp_pair.status_code >= 400:
                _raise_backend_error(resp_pair, "Interaction request")
            data_pair = resp_pair.json()
            pairwise_for_network = [
                item for item in _extract_interactions_from_response(data_pair, method, features)
                if len(item[0]) == 2
            ]
        except Exception as exc:
            print("Pairwise interaction fetch failed:", exc)
    if not pairwise_for_network:
        pairwise_for_network = _fallback_pairwise_from_values(features, feature_values)

    print("LIVE features:", features)
    print("LIVE inter_list (first 3):", inter_list[:3])

    text_interaction_html = create_text_interaction_html(
        features,
        feature_values,
        _pairwise_to_index_interactions(pairwise_for_network, features),
        top_k=20,
        threshold=0.0,
        method=method,
    )

    # ---------- 4. RESCALE VERY SMALL VALUES ----------
    max_abs = max((abs(v) for v in marginals.values()), default=0.0)
    scale = 1.0
    if 0 < max_abs < 1e-3:
        scale = 1e3
    if scale != 1.0:
        marginals = {k: v * scale for k, v in marginals.items()}
        inter_list = [(feats, val * scale) for feats, val in inter_list]
        feature_values = [val * scale for val in feature_values]

    # ---------- 5. INLINE TEXT HEATMAP ----------
    spans = None
    masking = data_attr.get("masking") or data_attr.get("mask") or {}
    if isinstance(masking, dict):
        spans = masking.get("feature_spans") or masking.get("spans")

    html = None
    if spans and len(spans) == len(feature_values):
        html = create_interactive_text_heatmap(
            context or text_source,
            spans,
            feature_values,
            method=method,
        )

    # ---------- 6. PLOTS + TABLES + META ----------
    inter_fig = plot_top_interactions(inter_list, order=order, method=method)

    if progress is not None:
        progress(0.8, desc="Rendering visualizations")

    y_len_tokens = data_attr.get("y_len_tokens")
    scoring_target_source = data_attr.get("scoring_target_source") or "model_output"
    scoring_target_text = data_attr.get("scoring_target_text")
    if scoring_target_text is None:
        scoring_target_text = correct_answer or data_attr.get("y_full") or ""

    meta = {
        "mode": "live",
        "backend_url_attr": url_attr,
        "backend_url_int": url_int,
        "method": method,
        "feature_level": level,
        "interaction_order": order,
        "model_size": model_size,
        "feature_count": len(features),
        "max_abs_value": max_abs,
        "scale_applied": scale,
        "scalarizer": data_attr.get("scalarizer_used", payload.get("scalarizer")),
        "scoring_target_source": scoring_target_source,
        "scoring_target_text_preview": str(scoring_target_text)[:200],
        "score_full": data_attr.get("score_full"),
        "score_empty": data_attr.get("score_empty"),
        "y_len_tokens": y_len_tokens,
        "logprob_full": data_attr.get("logprob_full"),
        "logprob_empty": data_attr.get("logprob_empty"),
        "min_logprob_seen": data_attr.get("min_logprob_seen"),
        "reference_answer_received": data_attr.get("reference_answer_received"),
        "answer_received": data_attr.get("answer_received"),
        "raw_attr_keys": list(data_attr.keys()),
        "raw_int_keys": list(data_int.keys()),
    }
    reference_answer = correct_answer
    unmasked_answer = data_attr.get("y_full") or data_attr.get("unmasked_answer") or ""
    debug_scores = data_attr.get("debug_scores") or None

    interaction_chips_html = create_interaction_token_view(
        features,
        feature_values,
        pairwise_for_network,
        method=method,
        layout="sentence" if level == "sentence" else "token",
    )
    figs = {
        "interactions": inter_fig,
    }

    if progress is not None:
        progress(1.0, desc="Done")

    return update(
        figs=figs,
        meta=meta,
        html=html,
        interaction_html=interaction_chips_html,
        interaction_text_html=text_interaction_html,
        scoring_target_source=scoring_target_source,
        scoring_target_text=str(scoring_target_text),
        reference_answer=reference_answer,
        unmasked_answer=unmasked_answer,
        debug_scores=debug_scores,
        scalarizer_used=data_attr.get("scalarizer_used", payload.get("scalarizer")),
        score_full=data_attr.get("score_full"),
        score_empty=data_attr.get("score_empty"),
        y_len_tokens=y_len_tokens,
    )


def _compute_image_attributions(
    image: Image.Image,
    prompt: str,
    answer: str,
    model_key: str,
    method: str,
    order: int,
    mask_level: str,
    grid_size: int,
    random_seed: int,
    progress=None,
) -> Tuple[Any, Any, Any, Any]:
    if image is None:
        raise gr.Error("Please provide an image.")
    if not answer:
        raise gr.Error("Please provide a target answer.")

    method = _normalize_method(method)
    payload = {
        "image_b64": _encode_image_to_b64(image),
        "prompt": prompt or "",
        "answer": answer,
        "method": method,
        "mask_level": (mask_level or "patch").lower(),
        "grid_size": int(grid_size) if grid_size else None,
        "order": int(order),
        "random_seed": int(random_seed or 0),
        "model_key": model_key,
    }

    if progress is not None:
        progress(0.1, desc="Calling image attribution backend")

    url = BACKEND_URL.rstrip("/") + "/api/image_attributions"
    try:
        resp = requests.post(url, json=payload, timeout=REQUEST_TIMEOUT)
    except requests.exceptions.ReadTimeout as exc:
        raise gr.Error(
            "Image attribution timed out. The backend is likely still running. "
            "Consider reducing grid size/order or set ATTRLLM_REQUEST_TIMEOUT to a higher value."
        ) from exc
    if resp.status_code >= 400:
        detail = resp.text
        try:
            detail = resp.json().get("detail", detail)
        except Exception:
            pass
        raise gr.Error(f"Image attribution failed: {detail}")
    data = resp.json()

    regions = data.get("regions") or []
    values = data.get("values") or []
    interactions = data.get("interactions") or []

    value_map = {int(item.get("index", 0)): float(item.get("value", 0.0)) for item in values}
    features = []
    for region in regions:
        try:
            idx = int(region.get("index", 0))
        except Exception:
            continue
        label = str(region.get("label") or f"Region {idx + 1}")
        features.append(
            {
                "index": idx,
                "feature": f"img:{label}",
                "value": float(value_map.get(idx, 0.0)),
                "method": method,
                "modality": "image",
                "ref_index": idx,
            }
        )
    features.sort(key=lambda item: item["index"])

    image_info = data.get("image") or {}
    image_b64 = image_info.get("image_b64") or payload["image_b64"]
    overlay_b64 = image_info.get("overlay_b64") or ""
    width = image_info.get("width")
    height = image_info.get("height")
    image_size = (int(width), int(height)) if width and height else None

    html = create_multimodal_interaction_html(
        image_b64,
        overlay_b64,
        regions,
        features,
        interactions,
        image_size=image_size,
        title="Image Interaction View",
    )

    region_labels = _labels_from_regions(regions)
    inter_pairs = _interaction_dicts_to_pairs(interactions, region_labels, order=order)
    inter_plot = plot_top_interactions(inter_pairs, order=order, method=method)
    inter_table = _interaction_dicts_to_table(interactions, region_labels)

    meta = {
        "mode": "image",
        "backend_url": url,
        "method": method,
        "order": order,
        "mask_level": mask_level,
        "grid_size": grid_size,
        "region_count": len(regions),
    }

    if progress is not None:
        progress(1.0, desc="Done")

    return html, inter_plot, inter_table, meta


def _compute_mm_attributions(
    image: Image.Image,
    text_context: str,
    answer: str,
    model_key: str,
    method: str,
    order: int,
    mask_level_image: str,
    mask_level_text: str,
    grid_size: int,
    random_seed: int,
    progress=None,
) -> Tuple[Any, Any, Any, Any]:
    if image is None:
        raise gr.Error("Please provide an image.")
    if not answer:
        raise gr.Error("Please provide a target answer.")

    method = _normalize_method(method)
    payload = {
        "image_b64": _encode_image_to_b64(image),
        "text_context": text_context or "",
        "answer": answer,
        "method": method,
        "order": int(order),
        "mask_level_image": (mask_level_image or "patch").lower(),
        "mask_level_text": (mask_level_text or "word").lower(),
        "grid_size": int(grid_size) if grid_size else None,
        "random_seed": int(random_seed or 0),
        "model_key": model_key,
    }

    if progress is not None:
        progress(0.1, desc="Calling multimodal attribution backend")

    url = BACKEND_URL.rstrip("/") + "/api/mm_attributions"
    try:
        resp = requests.post(url, json=payload, timeout=REQUEST_TIMEOUT)
    except requests.exceptions.ReadTimeout as exc:
        raise gr.Error(
            "Multimodal attribution timed out. The backend is likely still running. "
            "Consider reducing grid size/order or set ATTRLLM_REQUEST_TIMEOUT to a higher value."
        ) from exc
    if resp.status_code >= 400:
        detail = resp.text
        try:
            detail = resp.json().get("detail", detail)
        except Exception:
            pass
        raise gr.Error(f"Multimodal attribution failed: {detail}")
    data = resp.json()

    features = data.get("features") or []
    regions = data.get("regions") or []
    interactions = data.get("interactions") or []

    features.sort(key=lambda item: item.get("index", 0))
    region_labels = _labels_from_regions(regions)

    labels: List[str] = []
    for feature in features:
        labels.append(_feature_display_label(feature, region_labels))

    image_info = data.get("image") or {}
    image_b64 = image_info.get("image_b64") or payload["image_b64"]
    overlay_b64 = image_info.get("overlay_b64") or ""
    width = image_info.get("width")
    height = image_info.get("height")
    image_size = (int(width), int(height)) if width and height else None

    html = create_multimodal_interaction_html(
        image_b64,
        overlay_b64,
        regions,
        features,
        interactions,
        image_size=image_size,
        title="Multimodal Interaction View",
    )

    inter_pairs = _interaction_dicts_to_pairs(interactions, labels, order=order)
    inter_plot = plot_top_interactions(inter_pairs, order=order, method=method)
    inter_table = _interaction_dicts_to_table(interactions, labels)

    meta = {
        "mode": "multimodal",
        "backend_url": url,
        "method": method,
        "order": order,
        "mask_level_image": mask_level_image,
        "mask_level_text": mask_level_text,
        "feature_count": len(features),
        "region_count": len(regions),
    }

    if progress is not None:
        progress(1.0, desc="Done")

    return html, inter_plot, inter_table, meta

def on_select_example(
    dataset,
    ex_id,
    model_size,
    order,
    method,
    scalarizer=None,
    feature_level=None,
):
    """
    Public mode handler: load a precomputed example and render figures.

    Args:
        dataset (str): dataset name
        ex_id (str): example id
        model_size (str): "small" | "medium" | "large"
        order (int): interaction order (2 or 3)
        method (str): "shapley" | "banzhaf" | "influence"

    Returns:
        tuple ordered as:
            (
                context,
                prompt,
                answer,
                interactions_plot,
                interactions_token_html,
                text_html,
                meta_json,
            )
    """
    get_res = get_result_by_id if get_result_by_id is not None else _public_get_result_from_file
    model_size = _normalize_model_size(model_size)
    example = {"context": "", "prompt": "", "answer": ""}
    if get_example_by_id is not None:
        try:
            example = get_example_by_id(dataset, ex_id)
        except Exception:
            pass
    result = get_res(
        model_size,
        dataset,
        ex_id,
        scalarizer=scalarizer,
        feature_level=feature_level,
    ) or {}
    payload = result.get(method, {})
    # Your JSON: features (list of strings) + mobius_dict. Convert to UI format if needed.
    feats = payload.get("features") if isinstance(payload, dict) else None
    if isinstance(feats, list) and feats and not isinstance(feats[0], dict):
        payload = _normalize_public_payload_fallback(payload, method)

    features, feature_values = _extract_feature_series(payload)
    if not features:
        features = ["<no features>"]
        feature_values = [0.0]
    # Influence scores are non-negative (squared Fourier coefficients)
    if method == "influence":
        feature_values = [abs(v) for v in feature_values]
    marginals = {feat: float(feature_values[idx]) for idx, feat in enumerate(features)}
    interactions = _resolve_interactions(payload, order)
    if method == "influence":
        interactions = [(feats, abs(val)) for feats, val in interactions]
    pairwise = _resolve_pairwise(payload, features, feature_values)
    if method == "influence":
        pairwise = [(feats, abs(val)) for feats, val in pairwise]
    payload_level = (
            payload.get("mask_level")
        or payload.get("feature_level")
        or payload.get("level")
        or (result.get("meta", {}) if isinstance(result, dict) else {}).get("feature_level")
    )
    layout_mode = "sentence" if _normalize_level(payload_level) == "sentence" else "token"

    inter = plot_top_interactions(interactions, order=order, method=method)

    spans = payload.get("feature_spans") or payload.get("spans")
    if not spans:
        # Precomputed JSON payloads may not include explicit spans.
        # Reconstruct spans from context + feature level so Text View can render.
        _, fallback_spans, _ = _chunk_text_for_visualization(
            example.get("context", ""),
            _normalize_level(payload_level),
        )
        if fallback_spans and len(fallback_spans) == len(feature_values):
            spans = fallback_spans
    html = None
    if spans and len(spans) == len(feature_values):
        html = create_interactive_text_heatmap(
            example.get("context", ""),
            spans,
            feature_values,
            method=method,
        )

    meta = {
        "dataset": dataset,
        "example_id": ex_id,
        "model_size": model_size,
        "method": method,
        "order": order,
        "feature_count": len(features),
        "payload_keys": sorted(payload.keys()),
    }
    if "meta" in result:
        meta["source_meta"] = result["meta"]

    interaction_chips_html = create_interaction_token_view(
        features,
        feature_values,
        pairwise or [item for item in interactions if len(item[0]) == 2],
        method=method,
        layout=layout_mode,
    )
    text_interaction_html = create_text_interaction_html(
        features,
        feature_values,
        _pairwise_to_index_interactions(
            pairwise or [item for item in interactions if len(item[0]) == 2],
            features,
        ),
        top_k=20,
        threshold=0.0,
        method=method,
    )
    figs = {
        "interactions": inter,
    }
    outputs = update(
        figs=figs,
        meta=meta,
        html=html,
        interaction_html=interaction_chips_html,
        interaction_text_html=text_interaction_html,
    )
    return (
        example.get("context", ""),
        example.get("prompt", ""),
        _extract_answer(example),
        *outputs,
    )


def on_click_compute(
    context,
    prompt,
    correct_answer,
    model_size,
    level,
    method,
    scalarizer,
    embedding_model,
    progress=gr.Progress(track_tqdm=True),
):
#     """
#     Developer mode handler: compute (or mock) attributions and render figures.
#     """
#     method = _normalize_method(method)
#     level = _normalize_level(level)
#     order = 3 if int(order or 2) >= 3 else 2

#     context = context or ""
#     prompt = prompt or ""
#     correct_answer = correct_answer or ""

#     try:
#         return _compute_live_attributions(
#             context=context,
#             prompt=prompt,
#             correct_answer=correct_answer,
#             model_size=model_size,
#             level=level,
#             method=method,
#             order=order,
#             progress=progress,
#         )
#     except Exception as exc:  # pragma: no cover - best-effort fallback
#         return _synthetic_attribution_pipeline(
#             context,
#             prompt,
#             correct_answer,
#             method=method,
#             level=level,
#             order=order,
#             reason=str(exc),
#         )

    method = _normalize_method(method)
    level = _normalize_level(level)
    model_size = _normalize_model_size(model_size)
    order = 2

    context = context or ""
    prompt = prompt or ""
    correct_answer = correct_answer or ""

    return _compute_live_attributions(
        context=context,
        prompt=prompt,
        correct_answer=correct_answer,
        model_size=model_size,
        level=level,
        method=method,
        order=order,
        scalarizer=scalarizer,
        embedding_model=embedding_model,
        progress=progress,
    )


def on_click_image_compute(
    image,
    prompt,
    answer,
    model_key,
    method,
    mask_level,
    grid_size,
    random_seed,
    progress=gr.Progress(track_tqdm=True),
):
    return _compute_image_attributions(
        image=image,
        prompt=prompt,
        answer=answer,
        model_key=model_key,
        method=method,
        order=2,
        mask_level=mask_level,
        grid_size=grid_size,
        random_seed=random_seed,
        progress=progress,
    )


def on_click_mm_compute(
    image,
    text_context,
    answer,
    model_key,
    method,
    mask_level_image,
    mask_level_text,
    grid_size,
    random_seed,
    progress=gr.Progress(track_tqdm=True),
):
    return _compute_mm_attributions(
        image=image,
        text_context=text_context,
        answer=answer,
        model_key=model_key,
        method=method,
        order=2,
        mask_level_image=mask_level_image,
        mask_level_text=mask_level_text,
        grid_size=grid_size,
        random_seed=random_seed,
        progress=progress,
    )


# ---------------------------------------------------------------------------
# Demo helpers (used to quickly validate visualization components locally)
# ---------------------------------------------------------------------------

_DEMO_TEXT = "The quick brown fox jumps over the lazy dog in a sunny meadow."
_DEMO_FEATURES = ["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog", "in", "a", "sunny", "meadow"]
_DEMO_SPANS = [
    (0, 3), (4, 9), (10, 15), (16, 19), (20, 25), (26, 30), (31, 34),
    (35, 39), (40, 43), (44, 46), (47, 48), (49, 54), (55, 61)
]
_DEMO_ATTRIBUTIONS: Dict[str, Dict[str, float]] = {
    "shapley": {
        "The": -0.04,
        "quick": 0.18,
        "brown": 0.12,
        "fox": 0.27,
        "jumps": 0.15,
        "over": 0.05,
        "the": -0.02,
        "lazy": -0.11,
        "dog": -0.07,
        "in": 0.03,
        "a": 0.02,
        "sunny": 0.09,
        "meadow": 0.21,
    }
}
_DEMO_ATTRIBUTIONS["banzhaf"] = {
    token: round(value * 0.8, 3)
    for token, value in _DEMO_ATTRIBUTIONS["shapley"].items()
}
_DEMO_ATTRIBUTIONS["influence"] = {
    token: round(value ** 2, 4)
    for token, value in _DEMO_ATTRIBUTIONS["shapley"].items()
}

_DEMO_INTERACTIONS_2: List[Tuple[Tuple[str, ...], float]] = [
    (("quick", "fox"), 0.24),
    (("fox", "jumps"), 0.19),
    (("sunny", "meadow"), 0.22),
    (("lazy", "dog"), -0.17),
    (("the", "lazy"), -0.12),
]

_DEMO_INTERACTIONS_3: List[Tuple[Tuple[str, ...], float]] = [
    (("quick", "brown", "fox"), 0.28),
    (("fox", "jumps", "over"), 0.18),
    (("sunny", "meadow", "dog"), 0.11),
    (("the", "lazy", "dog"), -0.21),
]

_DEMO_INTERACTION_MATRIX: List[Tuple[Tuple[int, int], float]] = [
    ((1, 3), 0.23),
    ((3, 4), 0.17),
    ((7, 8), -0.18),
    ((11, 12), 0.2),
    ((2, 5), 0.09),
]

_DEMO_DATASETS = {
    "squad_demo": [
        [
            "The quick brown fox jumps over the lazy dog.",
            "Who jumps over the dog?",
            "The quick brown fox",
        ],
        [
            "AttrLLM explains attributions for large language models.",
            "What does AttrLLM explain?",
            "Attributions",
        ],
    ],
    "truthfulqa_demo": [
        [
            "Water boils at 100 degrees Celsius at sea level.",
            "At what temperature does water boil?",
            "100 degrees Celsius",
        ]
    ],
}


def _render_demo(method: str = "shapley"):
    method = (method or "shapley").lower()
    order = 2
    attributions = _DEMO_ATTRIBUTIONS.get(method, _DEMO_ATTRIBUTIONS["shapley"])
    interactions = _DEMO_INTERACTIONS_3 if order == 3 else _DEMO_INTERACTIONS_2

    interactions_fig = plot_top_interactions(interactions, order=order, method=method)
    demo_pairwise = _DEMO_INTERACTIONS_2 or _fallback_pairwise_from_values(
        _DEMO_FEATURES,
        [attributions[token] for token in _DEMO_FEATURES],
    )
    text_html = create_interactive_text_heatmap(
        _DEMO_TEXT,
        _DEMO_SPANS,
        [attributions[token] for token in _DEMO_FEATURES],
        method=method,
    )
    text_interaction_html = create_text_interaction_html(
        _DEMO_FEATURES,
        [attributions[token] for token in _DEMO_FEATURES],
        [
            {"indices": [i, j], "value": float(val)}
            for (i, j), val in _DEMO_INTERACTION_MATRIX
        ],
        top_k=20,
        threshold=0.0,
        method=method,
    )

    meta = {
        "method": method,
        "order": order,
        "feature_count": len(_DEMO_FEATURES),
        "scalarizer": "logprob",
    }

    return update(
        figs={
            "interactions": interactions_fig,
        },
        meta=meta,
        html=text_html,
        interaction_text_html=text_interaction_html,
        scoring_target_source="model_output",
        scoring_target_text="",
        reference_answer="",
        unmasked_answer="",
        debug_scores=None,
        scalarizer_used="logprob",
        score_full=None,
        score_empty=None,
        y_len_tokens=None,
    )


def _render_additional_plots(method: str = "shapley"):
    return plot_interaction_matrix(_DEMO_FEATURES, _DEMO_INTERACTION_MATRIX)


def _records_for_dataset(dataset_name: str) -> List[Dict[str, Any]]:
    if get_examples is not None:
        try:
            records = get_examples(dataset_name, n=10)
            if records:
                return records
        except KeyError:
            pass
        except Exception:
            pass

    fallback_csv = _fallback_load_dataset(dataset_name, max_rows=10)
    if fallback_csv:
        return fallback_csv

    fallback = []
    for idx, row in enumerate(_DEMO_DATASETS.get(dataset_name, []), start=1):
        context, prompt, answer = row
        fallback.append(
            {
                "id": f"{dataset_name}_demo_{idx}",
                "context": context,
                "prompt": prompt,
                "correct_answer": answer,
            }
        )
    return fallback


def _available_datasets() -> List[str]:
    if list_datasets is not None:
        try:
            datasets = list_datasets()
            if datasets:
                return datasets
        except Exception:
            pass
    fallback = [k for k, v in _FALLBACK_DATASET_FILES.items() if (_fallback_datasets_dir() / v).exists()]
    if fallback:
        return sorted(fallback)
    return list(_DEMO_DATASETS.keys())


def _format_examples(records: List[Dict[str, Any]]) -> List[List[str]]:
    formatted = []
    for rec in records:
        formatted.append([
            rec.get("context", ""),
            rec.get("prompt", ""),
            rec.get("correct_answer")
            or rec.get("answer")
            or rec.get("target")
            or "",
        ])
    return formatted


def _load_examples_for_demo(dataset_name: str):
    # Convert display name to internal key if needed
    if get_dataset_key_from_display_name is not None:
        dataset_key = get_dataset_key_from_display_name(dataset_name)
    else:
        dataset_key = dataset_name
    
    records = _records_for_dataset(dataset_key)
    formatted = _format_examples(records)
    samples = formatted if formatted else _DEMO_DATASETS.get(dataset_key, [])
    return gr.update(samples=samples or [])


def _resolve_example_fields(record: Dict[str, Any]) -> Tuple[str, str, str]:
    context = record.get("context", "")
    prompt = record.get("prompt", "")
    answer = (
        record.get("correct_answer")
        or record.get("answer")
        or record.get("target")
        or ""
    )
    return context, prompt, answer


def _resolve_dataset_key(dataset_name: str) -> str:
    if dataset_name in _available_datasets():
        return dataset_name
    for key, label in DATASET_DISPLAY_LABELS.items():
        if dataset_name == label:
            return key
    if get_dataset_key_from_display_name is not None:
        return get_dataset_key_from_display_name(dataset_name)
    return dataset_name


def _dataset_choice_labels(dataset_keys: List[str]) -> List[str]:
    labels: List[str] = []
    for key in dataset_keys:
        if get_dataset_display_name is not None:
            try:
                labels.append(get_dataset_display_name(key))
                continue
            except Exception:
                pass
        labels.append(DATASET_DISPLAY_LABELS.get(key, key.replace("_", " ").title()))
    return labels


def _resolve_example_index(example_number: Any, records: List[Dict[str, Any]]) -> int:
    if not records:
        return 0
    try:
        index = int(example_number) - 1
    except Exception:
        index = 0
    return max(0, min(index, len(records) - 1))


def _resolve_example_id(example_number: Any, records: List[Dict[str, Any]]) -> str:
    if _public_only_mode():
        return f"example_{int(example_number or 1)}"
    index = _resolve_example_index(example_number, records)
    record = records[index] if records else {}
    return str(record.get("id") or f"example_{index + 1}")


def _load_examples_for_slider(dataset_name: str):
    dataset_key = _resolve_dataset_key(dataset_name)
    records = _records_for_dataset(dataset_key)
    slider_max = max(1, min(10, len(records) or 10))
    context = prompt = answer = ""
    if records:
        context, prompt, answer = _resolve_example_fields(records[0])
    slider_update = gr.update(minimum=1, maximum=slider_max, step=1, value=1)
    return slider_update, records, context, prompt, answer


def _update_example_preview(example_number: Any, records):
    if not records:
        return "", "", ""
    index = _resolve_example_index(example_number, records)
    return _resolve_example_fields(records[index])


def _results_output_list(results: Dict[str, Any]) -> List[Any]:
    return [
        results["interactions"],
        results["interactions_tokens_html"],
        results["interactions_text_html"],
        results["text_html"],
        results["meta"],
        results["scoring_target_source"],
        results["scoring_target_text"],
        results["reference_answer"],
        results["unmasked_answer"],
        results["debug_scores"],
        results["scalarizer_used"],
        results["score_full"],
        results["score_empty"],
        results["y_len_tokens"],
    ]


def build_demo_app() -> gr.Blocks:
    datasets = _available_datasets()
    default_dataset = datasets[0] if datasets else "demo"

    # Apply the same colorful CSS theme
    custom_css = """
    .gradio-container {
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", "Helvetica Neue", Arial, sans-serif !important;
        background: linear-gradient(135deg, #fef5f0 0%, #f0e8ff 50%, #e8f5ff 100%) !important;
        padding: 24px !important;
    }
    
    .gradio-container h1, .gradio-container h2 {
        background: linear-gradient(135deg, #ff6b6b 0%, #ee5a6f 30%, #c44569 60%, #6c5ce7 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-weight: 900;
        font-size: 42px !important;
        margin: 20px 0 16px 0;
        letter-spacing: -0.03em;
    }
    
    label, .gr-label {
        font-weight: 700 !important;
        font-size: 16px !important;
        color: #2d1f4a !important;
    }
    
    .gr-button {
        border-radius: 16px !important;
        font-weight: 700 !important;
        font-size: 17px !important;
        padding: 16px 32px !important;
        background: linear-gradient(135deg, #6c5ce7 0%, #a29bfe 100%) !important;
        color: white !important;
        border: none !important;
    }
    
    .gr-box, .gr-input, .gr-dropdown, .gr-textbox {
        border-radius: 14px !important;
        border: 3px solid #e8dff5 !important;
        font-size: 17px !important;
    }
    
    .gr-markdown p {
        font-size: 17px !important;
        font-weight: 500 !important;
    }
    """

    with gr.Blocks(title="AttrLLM Visualization Demo", css=custom_css) as demo:
        gr.Markdown(
            "# 🎨 AttrLLM Visualization Demo\n\n"
            "**Preview the attribution widgets** before wiring real backends. "
            "Use the controls below to explore the interface."
        )

        with gr.Row():
            with gr.Column(scale=1):
                # Prepare initial choices and value before creating component
                initial_choices = _dataset_choice_labels(datasets)
                initial_value = initial_choices[0] if initial_choices else None
                
                dataset_selector = gr.Dropdown(
                    choices=initial_choices,
                    value=initial_value,
                    label="Dataset",
                    interactive=True,
                    allow_custom_value=False,
                    elem_id="dataset-selector-demo",
                    elem_classes=["bubble-select"],
                )
                example_browser = create_example_browser()
                with gr.Column(scale=1):
                    model_selector = create_model_selector()
                    scalarizer_selector = gr.Dropdown(
                        choices=SCALARIZER_CHOICES,
                        value="logprob",
                        label="Scalarizer",
                        interactive=True,
                    )
                    embedding_model_box = gr.Textbox(
                        label="Embedding Model (for scalarizer=embedding)",
                        value="Qwen/Qwen3-Embedding-0.6B",
                        lines=1,
                    )
                    feature_level_selector = create_feature_level_selector()
                    method_toggle = create_attribution_method_toggle()

        dataset_selector.change(
            fn=_load_examples_for_demo,
            inputs=dataset_selector,
            outputs=example_browser,
        )
        demo.load(
            fn=_load_examples_for_demo,
            inputs=[dataset_selector],
            outputs=[example_browser],
        )

        render_button = gr.Button("Render Demo Visuals", variant="primary")

        outputs = create_results_display()
        extra_matrix = gr.Plot(label="Interaction Matrix (demo)")

        render_button.click(
            fn=_render_demo,
            inputs=[method_toggle],
            outputs=_results_output_list(outputs),
        )

        render_button.click(
            fn=_render_additional_plots,
            inputs=[method_toggle],
            outputs=[extra_matrix],
        )

    return demo


def build_app() -> gr.Blocks:
    datasets = _available_datasets()
    default_dataset = datasets[0] if datasets else ""
    public_only = _public_only_mode()

    # Custom CSS for prettier UI - Inspired by modern, colorful design
    custom_css = """
    /* Main container styling - Warm gradient background */
    .gradio-container {
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", "Helvetica Neue", Arial, sans-serif !important;
        background: linear-gradient(135deg, #fef5f0 0%, #f0e8ff 50%, #e8f5ff 100%) !important;
        padding: 24px !important;
    }
    
    /* Header styling - Large, bold, colorful */
    .gradio-container h1 {
        background: linear-gradient(135deg, #ff6b6b 0%, #ee5a6f 30%, #c44569 60%, #6c5ce7 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-weight: 900;
        font-size: 48px !important;
        margin: 20px 0 16px 0;
        letter-spacing: -0.03em;
        text-align: left;
    }
    
    .gradio-container h2 {
        background: linear-gradient(135deg, #ff6b6b 0%, #ee5a6f 30%, #c44569 60%, #6c5ce7 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-weight: 900;
        font-size: 42px !important;
        margin: 20px 0 16px 0;
        letter-spacing: -0.03em;
    }
    
    .gradio-container h3 {
        color: #2d1f4a;
        font-weight: 800;
        font-size: 24px !important;
        margin: 24px 0 16px 0;
    }
    
    /* Tab styling - Bold and colorful */
    .tab-nav {
        border: none !important;
        background: transparent !important;
        gap: 8px !important;
        padding: 8px 0 !important;
    }
    
    .tab-nav button {
        font-size: 18px !important;
        font-weight: 700 !important;
        padding: 16px 32px !important;
        border-radius: 16px !important;
        transition: all 0.3s ease !important;
        border: 3px solid #e0d0f0 !important;
        background: white !important;
        color: #6c5ce7 !important;
        margin-right: 8px !important;
    }
    
    .tab-nav button:hover {
        background: #f8f4ff !important;
        border-color: #b8a8db !important;
        transform: translateY(-2px) !important;
    }
    
    .tab-nav button.selected {
        background: linear-gradient(135deg, #6c5ce7 0%, #a29bfe 100%) !important;
        color: white !important;
        border: 3px solid #6c5ce7 !important;
        box-shadow: 0 6px 20px rgba(108, 92, 231, 0.3) !important;
    }
    
    /* Button styling - Vibrant and interactive */
    .gr-button {
        border-radius: 16px !important;
        font-weight: 700 !important;
        font-size: 17px !important;
        padding: 16px 32px !important;
        transition: all 0.3s cubic-bezier(0.34, 1.56, 0.64, 1) !important;
        box-shadow: 0 6px 20px rgba(108, 92, 231, 0.2) !important;
        border: none !important;
    }
    
    .gr-button-primary {
        background: linear-gradient(135deg, #6c5ce7 0%, #a29bfe 100%) !important;
        color: white !important;
    }
    
    .gr-button-secondary {
        background: linear-gradient(135deg, #fd79a8 0%, #ff7675 100%) !important;
        color: white !important;
    }
    
    .gr-button:hover {
        transform: translateY(-3px) scale(1.02) !important;
        box-shadow: 0 10px 30px rgba(108, 92, 231, 0.35) !important;
    }
    
    .gr-button-primary:hover {
        background: linear-gradient(135deg, #5e4ec7 0%, #9089e8 100%) !important;
    }
    
    /* Input/Dropdown styling - Clear and modern */
    .gr-box, .gr-input, .gr-dropdown {
        border-radius: 14px !important;
        border: 3px solid #e8dff5 !important;
        background: white !important;
        font-size: 17px !important;
        padding: 12px 16px !important;
        transition: all 0.3s ease !important;
        font-weight: 500 !important;
    }
    
    .gr-box:focus, .gr-input:focus, .gr-dropdown:focus {
        border-color: #6c5ce7 !important;
        box-shadow: 0 0 0 4px rgba(108, 92, 231, 0.15) !important;
        transform: translateY(-1px) !important;
    }
    
    /* Textbox styling - Larger text */
    .gr-textbox {
        border-radius: 16px !important;
        border: 3px solid #e8dff5 !important;
        font-size: 17px !important;
        line-height: 1.6 !important;
    }
    
    .gr-textbox textarea {
        font-size: 17px !important;
        line-height: 1.6 !important;
        padding: 14px !important;
    }
    
    .gr-textbox:focus-within {
        border-color: #6c5ce7 !important;
        box-shadow: 0 6px 24px rgba(108, 92, 231, 0.2) !important;
    }
    
    /* Radio button styling - Colorful pills */
    .gr-radio {
        gap: 12px !important;
    }
    
    .gr-radio label {
        font-size: 17px !important;
        font-weight: 600 !important;
        padding: 14px 28px !important;
        border-radius: 14px !important;
        border: 3px solid #e8dff5 !important;
        transition: all 0.3s ease !important;
        background: white !important;
        cursor: pointer !important;
    }
    
    .gr-radio label:hover {
        border-color: #b8a8db !important;
        background: #faf8ff !important;
        transform: translateY(-2px) !important;
        box-shadow: 0 4px 12px rgba(108, 92, 231, 0.15) !important;
    }
    
    .gr-radio input:checked + label {
        background: linear-gradient(135deg, #6c5ce7 0%, #a29bfe 100%) !important;
        color: white !important;
        border-color: #6c5ce7 !important;
        font-weight: 800 !important;
        box-shadow: 0 6px 20px rgba(108, 92, 231, 0.3) !important;
    }
    
    /* Panel/Accordion styling - Clean cards */
    .gr-panel {
        border-radius: 20px !important;
        border: 3px solid #e8dff5 !important;
        padding: 24px !important;
        background: white !important;
        box-shadow: 0 6px 24px rgba(108, 92, 231, 0.1) !important;
        margin: 16px 0 !important;
    }
    
    .gr-accordion {
        border-radius: 18px !important;
        border: 3px solid #e8dff5 !important;
        background: white !important;
    }
    
    /* Label styling - Bold and readable */
    label, .gr-label {
        font-weight: 700 !important;
        font-size: 16px !important;
        color: #2d1f4a !important;
        margin-bottom: 10px !important;
        letter-spacing: -0.01em !important;
    }
    
    /* Dropdown options */
    .gr-dropdown-menu {
        border-radius: 14px !important;
        border: 3px solid #e8dff5 !important;
        box-shadow: 0 8px 32px rgba(108, 92, 231, 0.15) !important;
        font-size: 17px !important;
    }
    
    .gr-dropdown-menu .item {
        font-size: 17px !important;
        padding: 12px 16px !important;
        font-weight: 500 !important;
    }
    
    .gr-dropdown-menu .item:hover {
        background: linear-gradient(135deg, #f3f0ff 0%, #e8f5ff 100%) !important;
    }
    
    /* Plot container - Prominent */
    .gr-plot {
        border-radius: 20px !important;
        border: 3px solid #e8dff5 !important;
        overflow: hidden !important;
        box-shadow: 0 8px 30px rgba(108, 92, 231, 0.12) !important;
        background: white !important;
    }
    
    /* JSON viewer */
    .gr-json {
        border-radius: 16px !important;
        border: 3px solid #e8dff5 !important;
        background: #faf8ff !important;
        padding: 20px !important;
        font-family: 'Monaco', 'Menlo', 'Consolas', monospace !important;
        font-size: 15px !important;
    }
    
    /* Column styling */
    .gr-column {
        padding: 20px !important;
    }
    
    /* Row styling */
    .gr-row {
        gap: 24px !important;
        margin: 12px 0 !important;
    }
    
    /* Markdown content - Larger, more readable */
    .gr-markdown {
        line-height: 1.8 !important;
        color: #2d1f4a !important;
    }
    
    .gr-markdown p {
        font-size: 17px !important;
        margin: 12px 0 !important;
        font-weight: 500 !important;
    }
    
    .gr-markdown strong {
        font-weight: 800 !important;
        color: #6c5ce7 !important;
    }
    
    /* Status/info messages - Colorful notifications */
    .gr-info {
        border-radius: 16px !important;
        border-left: 5px solid #6c5ce7 !important;
        background: linear-gradient(135deg, #f8f6ff 0%, #f0f4ff 100%) !important;
        padding: 18px 24px !important;
        font-size: 16px !important;
        font-weight: 600 !important;
        color: #2d1f4a !important;
        box-shadow: 0 4px 16px rgba(108, 92, 231, 0.1) !important;
    }
    
    /* Error messages */
    .gr-error {
        border-radius: 16px !important;
        border-left: 5px solid #ff6b6b !important;
        background: linear-gradient(135deg, #fff5f5 0%, #ffe8e8 100%) !important;
        padding: 18px 24px !important;
        font-size: 16px !important;
        font-weight: 600 !important;
        color: #c44569 !important;
    }
    
    /* Loading spinner */
    .loading {
        border: 4px solid #f3f0ff !important;
        border-top: 4px solid #6c5ce7 !important;
    }
    
    /* Scrollbar styling */
    ::-webkit-scrollbar {
        width: 12px !important;
        height: 12px !important;
    }
    
    ::-webkit-scrollbar-track {
        background: #f8f6ff !important;
        border-radius: 10px !important;
    }
    
    ::-webkit-scrollbar-thumb {
        background: linear-gradient(135deg, #6c5ce7 0%, #a29bfe 100%) !important;
        border-radius: 10px !important;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: linear-gradient(135deg, #5e4ec7 0%, #9089e8 100%) !important;
    }

    .results-shell {
        margin-top: 16px !important;
        background: transparent !important;
        border: none !important;
        border-radius: 0 !important;
        padding: 0 !important;
        box-shadow: none !important;
    }

    .results-shell,
    .results-shell > div,
    .results-shell .gr-group,
    .results-shell .gr-box,
    .results-shell .gr-panel,
    .results-shell .block {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
    }

    .interaction-stack {
        gap: 20px !important;
        padding: 0 8px 0 !important;
    }

    .interaction-stack .gr-plot {
        max-height: 700px !important;
        overflow-y: auto !important;
    }

    .interaction-stack h3 {
        margin-left: 24px !important;
        margin-bottom: 12px !important;
    }

    .public-controls {
        align-items: stretch !important;
        gap: 20px !important;
        margin-top: 8px !important;
    }

    .control-card {
        background: linear-gradient(180deg, rgba(255, 255, 255, 0.82) 0%, rgba(250, 246, 255, 0.96) 100%) !important;
        border: 2px solid rgba(224, 208, 240, 0.78) !important;
        border-radius: 26px !important;
        padding: 18px 20px 14px !important;
        box-shadow: 0 14px 30px rgba(108, 92, 231, 0.07) !important;
    }

    .control-card-primary {
        background: linear-gradient(180deg, rgba(255, 255, 255, 0.86) 0%, rgba(244, 248, 255, 0.96) 100%) !important;
    }

    .control-card-secondary {
        background: linear-gradient(180deg, rgba(255, 255, 255, 0.86) 0%, rgba(250, 244, 255, 0.96) 100%) !important;
    }

    .control-card .gradio-container,
    .control-card .gr-group {
        background: transparent !important;
    }

    .control-card > div,
    .control-card .block,
    .control-card .wrap,
    .control-card .gr-form,
    .control-card .form {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
    }

    .control-card .gr-box,
    .control-card .gr-panel {
        background: transparent !important;
        box-shadow: none !important;
    }

    .bubble-select {
        border: 3px solid #8f5cff !important;
        border-radius: 18px !important;
        box-shadow: 0 8px 20px rgba(143, 92, 255, 0.10) !important;
        transition: box-shadow 0.2s ease, border-color 0.2s ease !important;
    }

    .bubble-select:focus-within {
        border-color: #7a3dff !important;
        box-shadow: 0 0 0 4px rgba(143, 92, 255, 0.14), 0 10px 24px rgba(143, 92, 255, 0.16) !important;
    }

    .example-id-slider {
        margin-top: 8px !important;
        padding: 10px 2px 2px !important;
    }

    .example-id-slider input[type="range"] {
        accent-color: #4f7cff !important;
    }

    .example-id-slider .number-input,
    .example-id-slider input[type="number"] {
        border-radius: 16px !important;
        border: 2px solid #d8dcee !important;
        background: linear-gradient(180deg, #ffffff 0%, #f7f9ff 100%) !important;
        font-weight: 700 !important;
        min-width: 72px !important;
    }

    .example-id-slider .wrap {
        gap: 14px !important;
    }

    @media (prefers-color-scheme: dark) {
        .gradio-container {
            background: radial-gradient(circle at top, #1e2a44 0%, #0d1422 52%, #090f19 100%) !important;
        }

        .gradio-container h3,
        label, .gr-label,
        .gr-markdown,
        .gr-markdown p {
            color: #e8eefc !important;
        }

        .gr-markdown strong {
            color: #cbd7ff !important;
        }

        .tab-nav button,
        .gr-box, .gr-input, .gr-dropdown, .gr-textbox,
        .gr-panel, .gr-accordion,
        .gr-plot, .gr-json {
            background: rgba(16, 24, 39, 0.88) !important;
            border-color: rgba(148, 163, 184, 0.24) !important;
            color: #e8eefc !important;
        }

        .tab-nav button {
            color: #d7e1ff !important;
        }

        .tab-nav button:hover {
            background: rgba(37, 52, 79, 0.96) !important;
            border-color: rgba(199, 210, 254, 0.36) !important;
        }

        .gr-radio label {
            background: rgba(16, 24, 39, 0.9) !important;
            border-color: rgba(148, 163, 184, 0.26) !important;
            color: #e8eefc !important;
        }

        .gr-radio label:hover {
            background: rgba(37, 52, 79, 0.96) !important;
        }

        .gr-textbox textarea,
        .gr-input input {
            background: transparent !important;
            color: #e8eefc !important;
        }

        .gr-dropdown-menu {
            background: #101827 !important;
            border-color: rgba(148, 163, 184, 0.24) !important;
        }

        .gr-dropdown-menu .item {
            color: #e8eefc !important;
        }

        .gr-dropdown-menu .item:hover {
            background: rgba(37, 52, 79, 0.96) !important;
        }

        .gr-plot .main-svg,
        .gr-plot .svg-container,
        .gr-plot .plot-container,
        .gr-plot .user-select-none {
            background: transparent !important;
        }

        .gr-plot .xtick text,
        .gr-plot .ytick text,
        .gr-plot .gtitle text,
        .gr-plot .xtitle text,
        .gr-plot .ytitle text,
        .gr-plot .annotation-text,
        .gr-plot .legend text {
            fill: #e8eefc !important;
            color: #e8eefc !important;
        }

        .gr-plot .gridlayer path,
        .gr-plot .zerolinelayer path,
        .gr-plot .xlines-above path,
        .gr-plot .ylines-above path {
            stroke: rgba(148, 163, 184, 0.22) !important;
        }

        .gr-info {
            background: linear-gradient(135deg, rgba(30, 41, 59, 0.95) 0%, rgba(17, 24, 39, 0.95) 100%) !important;
            color: #dbe7ff !important;
            border-left-color: #9db4ff !important;
        }

        .control-card {
            background: linear-gradient(180deg, rgba(16, 24, 39, 0.9) 0%, rgba(18, 28, 45, 0.96) 100%) !important;
            border-color: rgba(148, 163, 184, 0.2) !important;
            box-shadow: 0 18px 36px rgba(0, 0, 0, 0.24) !important;
        }

        .results-shell,
        .results-shell > div,
        .results-shell .gr-group,
        .results-shell .gr-box,
        .results-shell .gr-panel,
        .results-shell .block {
            background: transparent !important;
            border: none !important;
            box-shadow: none !important;
        }

        .bubble-select {
            border-color: #a06cff !important;
            box-shadow: 0 10px 24px rgba(143, 92, 255, 0.18) !important;
        }

        .bubble-select:focus-within {
            border-color: #c29cff !important;
            box-shadow: 0 0 0 4px rgba(143, 92, 255, 0.16), 0 12px 26px rgba(143, 92, 255, 0.22) !important;
        }

        .example-id-slider .number-input,
        .example-id-slider input[type="number"] {
            background: linear-gradient(180deg, #162031 0%, #111827 100%) !important;
            border-color: rgba(148, 163, 184, 0.24) !important;
            color: #e8eefc !important;
        }

        .gr-error {
            background: linear-gradient(135deg, rgba(68, 18, 32, 0.95) 0%, rgba(39, 12, 20, 0.95) 100%) !important;
            color: #ffd5dc !important;
        }

        ::-webkit-scrollbar-track {
            background: #111827 !important;
        }
    }
    """

    with gr.Blocks(title="LLM Reasoning Explorer Studio", css=custom_css) as app:
        gr.Markdown(
            "# LLM Reasoning Explorer Studio\n\n"
            "**Explore attribution results and feature interactions** with our interactive visualization tools. "
            "Browse pre-computed examples or analyze your own text in real-time with powerful AI insights."
        )
        with gr.Accordion("How to Use", open=False):
            gr.Markdown(
                "1. **Select a dataset** from 10 available datasets (100 total examples, 10 per dataset)\n"
                "2. **Choose a model** to compare: Qwen3-4B, Qwen3-30B, or Mistral-7B\n"
                "3. **Pick a scoring method:** Perplexity or Semantic Similarity\n"
                "4. **Set the feature level:** Word, Sentence, or Paragraph\n"
                "5. **Choose an attribution method:** Shapley, Banzhaf, or Influence\n"
                "6. **View results** in the Text Interaction View (inline highlights) and Bar View (ranked interactions)"
            )
        gr.Markdown(f"**Build:** {BUILD_ID} ({BUILD_TS})")

        example_state = gr.State([])

        with (gr.Column() if public_only else gr.Tab("Public Mode")):
            with gr.Row(elem_classes=["public-controls"]):
                with gr.Column(scale=1, elem_classes=["control-card", "control-card-primary"]):
                    # Prepare initial choices and value before creating component
                    initial_choices = _dataset_choice_labels(datasets)
                    initial_value = initial_choices[0] if initial_choices else None
                    
                    dataset_selector = gr.Dropdown(
                        choices=initial_choices,
                        value=initial_value,
                        label="Dataset",
                        interactive=True,
                        allow_custom_value=False,
                        elem_id="dataset-selector",
                        elem_classes=["bubble-select"],
                    )

                    example_selector = gr.Slider(
                        label="Example ID",
                        minimum=1,
                        maximum=10,
                        step=1,
                        value=1,
                        interactive=True,
                        elem_classes=["example-id-slider"],
                    )
                with gr.Column(scale=1, elem_classes=["control-card", "control-card-secondary"]):
                    model_selector = create_model_selector()
                    scalarizer_selector = gr.Dropdown(
                        choices=PUBLIC_SCALARIZER_CHOICES,
                        value="semantic_similarity",
                        label="Scalarizer",
                        interactive=True,
                        elem_classes=["bubble-select"],
                    )
                    public_feature_level_selector = create_feature_level_selector()
                    method_toggle = create_attribution_method_toggle()

            with gr.Accordion("Example Preview", open=True):
                with gr.Row():
                    with gr.Column(scale=3):
                        context_box = gr.Textbox(
                            label="Context",
                            lines=8,
                            interactive=False,
                        )
                    with gr.Column(scale=2):
                        prompt_box = gr.Textbox(
                            label="Prompt",
                            lines=4,
                            interactive=False,
                        )
                        answer_box = gr.Textbox(
                            label="Answer",
                            lines=3,
                            interactive=False,
                        )

            public_results = create_results_display()

            def _public_mode_compute(
                dataset,
                example_number,
                records,
                model_size,
                scalarizer,
                feature_level,
                method,
                progress=gr.Progress(track_tqdm=True),
            ):
                if not dataset:
                    raise gr.Error("Please select a dataset.")
                if not example_number:
                    raise gr.Error("Please select an example.")

                dataset_key = _resolve_dataset_key(dataset)
                ex_id = _resolve_example_id(example_number, records)

                method = _normalize_method(method)
                level = _normalize_level(feature_level)
                model_size = _normalize_model_size(model_size)

                # Prefer precomputed results: use loader if available, else load from file (Space-friendly).
                get_res = get_result_by_id if get_result_by_id is not None else _public_get_result_from_file
                result = get_res(
                    model_size,
                    dataset_key,
                    ex_id,
                    scalarizer=scalarizer,
                    feature_level=level,
                ) or {}
                payload = result.get(method, {})

                if not payload:
                    alt_size = _find_available_model_size(dataset_key, ex_id, scalarizer, level)
                    if alt_size and alt_size != model_size:
                        result = get_res(
                            alt_size,
                            dataset_key,
                            ex_id,
                            scalarizer=scalarizer,
                            feature_level=level,
                        ) or {}
                        payload = result.get(method, {})
                        if payload:
                            model_size = alt_size

                # If still no payload, try any available (model_size, scalarizer, level) for this example
                if not payload:
                    alt_size, alt_scalarizer, alt_level, result = _find_any_available_result(
                        dataset_key, ex_id, get_res, method
                    )
                    if alt_size and alt_scalarizer and alt_level and result:
                        payload = result.get(method, {})
                        model_size, scalarizer, level = alt_size, alt_scalarizer, alt_level

                if payload and (payload.get("features") or payload.get("heatmap")):
                    _, _, _, *outputs = on_select_example(
                        dataset_key,
                        ex_id,
                        model_size,
                        2,
                        method,
                        scalarizer=scalarizer,
                        feature_level=level,
                    )
                    return outputs

                # Public-only mode: do not attempt live compute
                if _public_only_mode() or get_example_by_id is None:
                    expected_ref = _reference_results_file(model_size, dataset_key, ex_id, scalarizer, level)
                    raise gr.Error(
                        "No precomputed results found.\n\n"
                        f"Expected (reference_answer):\n{expected_ref}\n\n"
                        "On Hugging Face Space: make sure the 'results' folder is in your repo "
                        "(commit & push it). If you use Git LFS, enable 'LFS' in Space Settings → "
                        "Repository and ensure files are pulled. You can also try another "
                        "scalarizer (e.g. Perplexity) or feature level (e.g. word)."
                    )

                # Fallback to live compute if no precomputed payload or non-word level
                get_ex = _ensure_backend("loader.data.get_example_by_id", get_example_by_id)
                record = get_ex(dataset_key, ex_id)
                context = record.get("context", "")
                prompt = record.get("prompt", "")
                answer = _extract_answer(record)

                return _compute_live_attributions(
                    context=context,
                    prompt=prompt,
                    correct_answer=answer,
                    model_size=model_size,
                    scalarizer=scalarizer,
                    embedding_model=None,
                    level=level,
                    method=method,
                    order=2,
                    progress=progress,
                )

            public_preview_outputs = [context_box, prompt_box, answer_box]
            public_compute_inputs = [
                dataset_selector,
                example_selector,
                example_state,
                model_selector,
                scalarizer_selector,
                public_feature_level_selector,
                method_toggle,
            ]
            public_compute_outputs = _results_output_list(public_results)

            dataset_change_event = dataset_selector.change(
                fn=_load_examples_for_slider,
                inputs=[dataset_selector],
                outputs=[
                    example_selector,
                    example_state,
                    context_box,
                    prompt_box,
                    answer_box,
                ],
                queue=False,
            )
            load_event = app.load(
                fn=_load_examples_for_slider,
                inputs=[dataset_selector],
                outputs=[
                    example_selector,
                    example_state,
                    context_box,
                    prompt_box,
                    answer_box,
                ],
            )

            dataset_change_event.then(
                fn=_public_mode_compute,
                inputs=public_compute_inputs,
                outputs=public_compute_outputs,
                show_progress="full",
            )
            load_event.then(
                fn=_public_mode_compute,
                inputs=public_compute_inputs,
                outputs=public_compute_outputs,
                show_progress="full",
            )

            example_selector.release(
                fn=_update_example_preview,
                inputs=[example_selector, example_state],
                outputs=public_preview_outputs,
                queue=False,
            ).then(
                fn=_public_mode_compute,
                inputs=public_compute_inputs,
                outputs=public_compute_outputs,
                show_progress="full",
            )

            for component in (
                model_selector,
                scalarizer_selector,
                public_feature_level_selector,
                method_toggle,
            ):
                component.change(
                    fn=_public_mode_compute,
                    inputs=public_compute_inputs,
                    outputs=public_compute_outputs,
                    show_progress="full",
                )

        if not public_only:
            with gr.Tab("Demo / Developer Preview"):
                gr.Markdown(
                    "Run quick attribution experiments with synthetic data while the ProxySPEX "
                    "backend is under construction. Once the backend is ready this tab will call "
                    "the real pipeline automatically."
                )

                with gr.Row():
                    with gr.Column(scale=1):
                        dev_context = gr.Textbox(
                            label="Context",
                            value=_DEMO_TEXT,
                            lines=8,
                            placeholder="Paste the context/passage you want to explain.",
                        )
                        dev_prompt = gr.Textbox(
                            label="Prompt / Question",
                            value="Who jumps over the dog?",
                            lines=3,
                        )
                        dev_answer = gr.Textbox(
                            label="Reference Answer (optional)",
                            value="The quick brown fox",
                            lines=2,
                        )
                    with gr.Column(scale=1):
                        dev_model_selector = create_model_selector()
                        dev_scalarizer_selector = gr.Dropdown(
                            choices=SCALARIZER_CHOICES,
                            value="logprob",
                            label="Scalarizer",
                            interactive=True,
                        )
                        dev_embedding_model_box = gr.Textbox(
                            label="Embedding Model (for scalarizer=embedding)",
                            value="Qwen/Qwen3-Embedding-0.6B",
                            lines=1,
                        )
                        dev_feature_selector = create_feature_level_selector()
                        dev_method_toggle = create_attribution_method_toggle()
                        gr.Markdown(
                            "Outputs fall back to synthetic scores if the attribution backend "
                            "isn't available in your environment."
                        )
                        compute_button = gr.Button(
                            "Compute Attributions", variant="primary"
                        )

                dev_results = create_results_display(show_answer_boxes=False)

                compute_button.click(
                    fn=on_click_compute,
                    inputs=[
                        dev_context,
                        dev_prompt,
                        dev_answer,
                        dev_model_selector,
                        dev_feature_selector,
                        dev_method_toggle,
                        dev_scalarizer_selector,
                        dev_embedding_model_box,
                    ],
                    outputs=_results_output_list(dev_results),
                )

            superpixel_available = supports_superpixel() if supports_superpixel else False
            image_mask_choices = ["patch", "superpixel"] if superpixel_available else ["patch"]

            with gr.Tab("Images"):
                gr.Markdown(
                    "Upload an image and compute region-level attributions using patch or "
                    "superpixel masking."
                )
                if not superpixel_available:
                    gr.Markdown("Superpixel masking requires scikit-image; only patch masking is available.")

                with gr.Row():
                    with gr.Column(scale=1):
                        image_input = gr.Image(label="Input Image", type="pil")
                        image_prompt = gr.Textbox(
                            label="Prompt (optional)",
                            lines=2,
                            placeholder="Describe the image or ask a question.",
                        )
                        image_answer = gr.Textbox(
                            label="Target Answer",
                            lines=2,
                            placeholder="Answer or caption to score.",
                        )
                    with gr.Column(scale=1):
                        image_model_selector = create_multimodal_model_selector()
                        image_method_toggle = create_attribution_method_toggle()
                        image_mask_level = gr.Radio(
                            choices=image_mask_choices,
                            value="patch",
                            label="Mask Level",
                            interactive=True,
                        )
                        image_grid_size = gr.Slider(
                            minimum=2,
                            maximum=16,
                            step=1,
                            value=8,
                            label="Patch Grid Size",
                        )
                        image_seed = gr.Number(
                            value=0,
                            precision=0,
                            label="Random Seed",
                        )
                        image_button = gr.Button(
                            "Compute Image Attributions",
                            variant="primary",
                        )

                image_view = _html_component("Image Interaction View")
                with gr.Row():
                    image_interactions_plot = gr.Plot(label="Top Interactions")
                    image_interactions_table = gr.Dataframe(
                        headers=["features", "value", "order"],
                        label="Interaction Table",
                        interactive=False,
                    )
                image_meta = gr.JSON(label="Meta")

                image_button.click(
                    fn=on_click_image_compute,
                    inputs=[
                        image_input,
                        image_prompt,
                        image_answer,
                        image_model_selector,
                        image_method_toggle,
                        image_mask_level,
                        image_grid_size,
                        image_seed,
                    ],
                    outputs=[
                        image_view,
                        image_interactions_plot,
                        image_interactions_table,
                        image_meta,
                    ],
                    show_progress="full",
                )

            with gr.Tab("Multimodal"):
                gr.Markdown(
                    "Combine image and text features to explain multimodal predictions."
                )
                if not superpixel_available:
                    gr.Markdown("Superpixel masking requires scikit-image; only patch masking is available.")

                with gr.Row():
                    with gr.Column(scale=1):
                        mm_image_input = gr.Image(label="Input Image", type="pil")
                        mm_text_context = gr.Textbox(
                            label="Text Context / Prompt",
                            lines=4,
                            placeholder="Prompt or context to condition on.",
                        )
                        mm_answer = gr.Textbox(
                            label="Target Answer / Caption",
                            lines=2,
                            placeholder="Answer or caption to score.",
                        )
                    with gr.Column(scale=1):
                        mm_model_selector = create_multimodal_model_selector()
                        mm_method_toggle = create_attribution_method_toggle()
                        mm_mask_level_image = gr.Radio(
                            choices=image_mask_choices,
                            value="patch",
                            label="Image Mask Level",
                            interactive=True,
                        )
                        mm_mask_level_text = gr.Radio(
                            choices=["word", "sentence", "paragraph"],
                            value="word",
                            label="Text Mask Level",
                            interactive=True,
                        )
                        mm_grid_size = gr.Slider(
                            minimum=2,
                            maximum=16,
                            step=1,
                            value=8,
                            label="Patch Grid Size",
                        )
                        mm_seed = gr.Number(
                            value=0,
                            precision=0,
                            label="Random Seed",
                        )
                        mm_button = gr.Button(
                            "Compute Multimodal Attributions",
                            variant="primary",
                        )

                mm_view = _html_component("Multimodal Interaction View")
                with gr.Row():
                    mm_interactions_plot = gr.Plot(label="Top Interactions")
                    mm_interactions_table = gr.Dataframe(
                        headers=["features", "value", "order"],
                        label="Interaction Table",
                        interactive=False,
                    )
                mm_meta = gr.JSON(label="Meta")

                mm_button.click(
                    fn=on_click_mm_compute,
                    inputs=[
                        mm_image_input,
                        mm_text_context,
                        mm_answer,
                        mm_model_selector,
                        mm_method_toggle,
                        mm_mask_level_image,
                        mm_mask_level_text,
                        mm_grid_size,
                        mm_seed,
                    ],
                    outputs=[
                        mm_view,
                        mm_interactions_plot,
                        mm_interactions_table,
                        mm_meta,
                    ],
                    show_progress="full",
                )

        gr.HTML(
            """
            <div style="
                margin-top: 32px;
                padding: 24px 32px;
                border-top: 2px solid #e0d0f0;
                text-align: center;
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                color: #4a3d6a;
                line-height: 1.8;
            ">
                <div style="font-size: 15px; font-weight: 700; margin-bottom: 8px; color: #2d1f4a;">
                    Contributors &mdash; University of California, Berkeley
                </div>
                <div style="font-size: 13px; display: flex; justify-content: center; gap: 32px; flex-wrap: wrap;">
                    <span><b>Stephen Tao</b> &middot; Loader Layer &middot;
                        <a href="mailto:stephen_tao@berkeley.edu" style="color:#6c5ce7; text-decoration:none;">stephen_tao@berkeley.edu</a></span>
                    <span><b>Yiting Gao</b> &middot; Attribution Layer &middot;
                        <a href="mailto:yg2025@berkeley.edu" style="color:#6c5ce7; text-decoration:none;">yg2025@berkeley.edu</a></span>
                    <span><b>Qingpeng Kong</b> &middot; Visualization Layer &middot;
                        <a href="mailto:qpkong@berkeley.edu" style="color:#6c5ce7; text-decoration:none;">qpkong@berkeley.edu</a></span>
                </div>
            </div>
            """
        )

    return app


def launch_demo(**kwargs):
    demo = build_demo_app()
    server_name = kwargs.pop("server_name", os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"))
    server_port = int(kwargs.pop("server_port", os.getenv("GRADIO_SERVER_PORT", "7860")))
    share = kwargs.pop("share", _env_flag("GRADIO_SHARE", False))
    show_error = kwargs.pop("show_error", True)
    demo.launch(
        server_name=server_name,
        server_port=server_port,
        share=share,
        show_error=show_error,
        **kwargs,
    )


def launch_app(**kwargs):
    app = build_app()
    server_name = kwargs.pop("server_name", os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"))
    server_port = int(kwargs.pop("server_port", os.getenv("GRADIO_SERVER_PORT", "7860")))
    share = kwargs.pop("share", _env_flag("GRADIO_SHARE", False))
    show_error = kwargs.pop("show_error", True)
    app.launch(
        server_name=server_name,
        server_port=server_port,
        share=share,
        show_error=show_error,
        **kwargs,
    )


if __name__ == "__main__":
    launch_app()