app.py

from __future__ import annotations

import gradio as gr
import pandas as pd

from featurelens.config import SETTINGS
from featurelens.hf_runtime import gpu
from featurelens.runtime import RUNTIME

# Restrained, print-inspired palette. The app deliberately avoids saturated dashboard colors.
INK_TEAL = "#708B86"
INK_UMBER = "#94785D"
INK_RED = "#986D68"
INK_PLUM = "#82768F"
INK_STONE = "#8E8A83"
INK_BLUEGREY = "#71808A"

CSS = r"""
.gradio-container {
  width: min(96vw, 1600px) !important;
  max-width: 1600px !important;
  margin-left: auto !important;
  margin-right: auto !important;
  padding: 0 22px 128px !important;
  font-family: "Times New Roman", Times, "Liberation Serif", Georgia, serif !important;
  font-size: 16px !important;
  line-height: 1.52;
}
.gradio-container input,
.gradio-container textarea,
.gradio-container button,
.gradio-container select,
.gradio-container label,
.gradio-container table,
.gradio-container .prose {
  font-family: "Times New Roman", Times, "Liberation Serif", Georgia, serif !important;
}
.gradio-container input,
.gradio-container textarea,
.gradio-container select { font-size: 15.5px !important; }
.gradio-container table { font-size: 15.5px !important; }
.gradio-container table thead th, .gradio-container [role="columnheader"] { font-size: 16px !important; font-weight: 700 !important; }
.gradio-container table tbody td, .gradio-container [role="gridcell"] { font-size: 15.5px !important; }
.gradio-container label, .gradio-container .label-wrap { font-size: 15px !important; }
.gradio-container p, .gradio-container li { font-size: 16px; }
.gradio-container h2 { font-size: 1.72rem; line-height: 1.25; }
.gradio-container h3 { font-size: 1.46rem; line-height: 1.28; }
.gradio-container h4 { font-size: 1.24rem; line-height: 1.30; }
.hero {
  padding: 14px 2px 10px;
  border-bottom: 1px solid var(--border-color-primary);
  margin-bottom: 14px;
}
.hero h1 {
  margin: 0;
  font-size: 2.18rem;
  font-weight: 600;
  letter-spacing: 0;
}
.hero .subtitle { margin-top: 4px; font-size: 1.08rem; opacity: .82; }
.hero .question { margin-top: 6px; font-size: 1rem; opacity: .72; }
.section-rule {
  margin: 20px 0 11px;
  padding-top: 9px;
  border-top: 1px solid var(--border-color-primary);
  font-variant: small-caps;
  letter-spacing: .045em;
  font-size: 1.12rem;
  font-weight: 700;
  opacity: .88;
}
.token-wrap { display: flex; flex-wrap: wrap; gap: 6px; padding: 8px 2px 14px; line-height: 1.9; }
.token {
  background: var(--background-fill-secondary);
  border: 1px solid var(--border-color-primary);
  border-radius: 2px;
  padding: 3px 7px;
  font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace !important;
  font-size: 12.5px;
}
.token.selected { border: 2px solid #708B86; font-weight: 700; }
.token sup { opacity: .58; margin-right: 4px; }
.small-note { opacity: .74; font-size: 14px; }
.instrument-note {
  border: 1px solid var(--border-color-primary);
  border-radius: 2px;
  padding: 10px 12px;
  background: var(--background-fill-secondary);
  margin: 6px 0 11px;
}
.copy-row { margin-top: -4px; }
.copy-row button { min-height: 34px !important; }
.gradio-container button { border-radius: 3px !important; font-weight: 600 !important; font-size: 15px !important; }
.action-btn button, .copy-btn button {
  background: #5f817b !important;
  color: #fff !important;
  border: 1px solid #5f817b !important;
}
.action-btn button:hover, .copy-btn button:hover { background: #536f6a !important; border-color: #536f6a !important; }
.gradio-container .form, .gradio-container .block { border-radius: 3px !important; }
.gradio-container textarea, .gradio-container input { border-radius: 2px !important; }
.form-note { margin-top: -3px; margin-bottom: 9px; font-size: .92rem; opacity: .72; }
.wide-table { width: 100% !important; }
.bottom-spacer { height: 110px; width: 100%; }
.tabs, .tabitem { padding-bottom: 26px !important; }
.context-card {
  border: 1px solid var(--border-color-primary);
  border-left: 4px solid #708B86;
  background: var(--background-fill-secondary);
  padding: 10px 14px;
  margin: 8px 0 16px;
  border-radius: 2px;
}
.context-card p { margin: 0 !important; }
.start-card {
  border: 1px solid var(--border-color-primary);
  background: var(--background-fill-secondary);
  padding: 14px 16px;
  min-height: 150px;
  border-radius: 2px;
}
.start-card h3 { margin-top: 0; }
.graph-note { font-size: .95rem; opacity: .72; margin-top: 2px; }
.table-heading {
  margin: 0 0 -34px !important;
  padding: 4px 58px 0 0 !important;
  min-height: 34px;
  position: relative;
  z-index: 3;
  pointer-events: none;
  font-size: 1.20rem !important;
  font-weight: 700 !important;
  line-height: 1.22 !important;
}
.result-table .label-wrap,
.result-table .label-wrap span,
.result-table label,
.result-table label span,
.result-table [data-testid="block-label"],
.result-table [data-testid="block-label"] span,
.result-table .block-label,
.result-table .block-title {
  font-size: 19px !important;
  font-weight: 700 !important;
  line-height: 1.30 !important;
}
.result-table table thead th,
.result-table table thead th *,
.result-table [role="columnheader"],
.result-table [role="columnheader"] * {
  font-size: 16.75px !important;
  font-weight: 700 !important;
  line-height: 1.25 !important;
}
/* Gradio native fullscreen is intercepted and converted to an in-place focus mode.
   Nothing is moved to the top of the HF iframe: the original component expands where the user clicked it. */
.fl-plot.featurelens-inline-focus,
.result-table.featurelens-inline-focus {
  position: relative !important;
  z-index: 5000 !important;
  background: var(--background-fill-primary) !important;
  border: 1px solid var(--border-color-primary) !important;
  box-shadow: 0 12px 36px rgba(0, 0, 0, .42) !important;
  border-radius: 4px !important;
}
.fl-plot.featurelens-inline-focus {
  transform-origin: top left !important;
}
.result-table.featurelens-inline-focus {
  overflow: visible !important;
}
.candidate-help { opacity: .78; font-size: .97rem; margin-top: -2px; }
@media (max-width: 900px) {
  .gradio-container { width: 100% !important; padding-left: 12px !important; padding-right: 12px !important; }
}
"""

THEME = gr.themes.Base(
    primary_hue="teal",
    secondary_hue="stone",
    neutral_hue="stone",
    radius_size="sm",
)

COPY_JS = r"""
(text) => {
  const value = text || "";
  const button = document.activeElement && document.activeElement.tagName === "BUTTON"
    ? document.activeElement : null;
  const oldLabel = button ? button.innerText : null;
  const signal = () => {
    if (!button) return;
    button.innerText = "✓ Copied with headers";
    button.disabled = true;
    window.setTimeout(() => {
      button.innerText = oldLabel || "Copy table with headers";
      button.disabled = false;
    }, 1200);
  };
  const fallback = () => {
    const node = document.createElement("textarea");
    node.value = value;
    node.style.position = "fixed";
    node.style.opacity = "0";
    document.body.appendChild(node);
    node.focus();
    node.select();
    document.execCommand("copy");
    document.body.removeChild(node);
    signal();
  };
  if (navigator.clipboard && navigator.clipboard.writeText) {
    navigator.clipboard.writeText(value).then(signal).catch(fallback);
  } else {
    fallback();
  }
  return [value];
}
"""

INSTALL_REFLOW_JS = r"""
() => {
  if (window.__featurelens_reflow_installed) return [];
  window.__featurelens_reflow_installed = true;
  let timer = null;
  const kick = () => {
    window.clearTimeout(timer);
    timer = window.setTimeout(() => window.dispatchEvent(new Event("resize")), 80);
  };
  const root = document.querySelector(".gradio-container") || document.body;
  if (window.ResizeObserver) {
    const observer = new ResizeObserver(kick);
    observer.observe(root);
    window.__featurelens_reflow_observer = observer;
  }
  const mutation = new MutationObserver(kick);
  mutation.observe(root, {subtree: true, childList: true});
  window.__featurelens_mutation_observer = mutation;

  const restoreFocus = (block) => {
    if (!block || !block.classList.contains("featurelens-inline-focus")) return;
    const saved = block.__featurelens_saved_style;
    if (saved == null || saved === "") block.removeAttribute("style");
    else block.setAttribute("style", saved);
    block.classList.remove("featurelens-inline-focus");
    block.__featurelens_saved_style = null;
    window.setTimeout(kick, 30);
  };

  const closeOtherFocus = (except) => {
    document.querySelectorAll(".featurelens-inline-focus").forEach((node) => {
      if (node !== except) restoreFocus(node);
    });
  };

  const focusPlotInPlace = (block) => {
    const rect = block.getBoundingClientRect();
    if (rect.width <= 0 || rect.height <= 0) return;
    const viewportWidth = Math.max(320, document.documentElement.clientWidth || window.innerWidth || rect.width);
    const screenHeight = Math.max(600, (window.screen && window.screen.availHeight) || 900);
    const maxWidth = Math.min(viewportWidth * 0.90, 1100);
    const maxHeight = Math.min(screenHeight * 0.68, 700);
    const scale = Math.max(1, Math.min(maxWidth / rect.width, maxHeight / rect.height, 1.8));
    const focusedWidth = rect.width * scale;
    let dx = (viewportWidth - focusedWidth) / 2 - rect.left;
    if (rect.left + dx < 12) dx += 12 - (rect.left + dx);
    if (rect.left + dx + focusedWidth > viewportWidth - 12) {
      dx -= (rect.left + dx + focusedWidth) - (viewportWidth - 12);
    }
    block.style.transformOrigin = "top left";
    block.style.transform = `translate(${dx}px, 0px) scale(${scale})`;
    block.style.marginBottom = `${Math.max(8, rect.height * (scale - 1) + 8)}px`;
    block.style.zIndex = "5000";
  };

  const focusTableInPlace = (block) => {
    const rect = block.getBoundingClientRect();
    if (rect.width <= 0) return;
    const viewportWidth = Math.max(320, document.documentElement.clientWidth || window.innerWidth || rect.width);
    const targetWidth = Math.max(rect.width, Math.min(viewportWidth * 0.94, 1400));
    let dx = (viewportWidth - targetWidth) / 2 - rect.left;
    if (rect.left + dx < 12) dx += 12 - (rect.left + dx);
    if (rect.left + dx + targetWidth > viewportWidth - 12) {
      dx -= (rect.left + dx + targetWidth) - (viewportWidth - 12);
    }
    block.style.width = `${targetWidth}px`;
    block.style.maxWidth = "none";
    block.style.transform = `translateX(${dx}px)`;
    block.style.zIndex = "5000";
  };

  const toggleInlineFocus = (block) => {
    if (block.classList.contains("featurelens-inline-focus")) {
      restoreFocus(block);
      return;
    }
    closeOtherFocus(block);
    block.__featurelens_saved_style = block.getAttribute("style") || "";
    block.classList.add("featurelens-inline-focus");
    if (block.classList.contains("fl-plot")) focusPlotInPlace(block);
    else focusTableInPlace(block);
    window.setTimeout(kick, 30);
  };

  // Keep the native toolbar icon, but replace Gradio fullscreen with an in-place expansion.
  // This avoids HF iframe jumps and preserves the chart's exact rendered aspect ratio.
  document.addEventListener("click", (event) => {
    const button = event.target && event.target.closest ? event.target.closest("button") : null;
    if (!button) return;
    const label = `${button.getAttribute("aria-label") || ""} ${button.getAttribute("title") || ""} ${button.textContent || ""}`.toLowerCase();
    if (!label.includes("fullscreen")) return;
    const block = button.closest(".fl-plot, .result-table");
    if (!block) return;
    event.preventDefault();
    event.stopImmediatePropagation();
    toggleInlineFocus(block);
  }, true);

  document.addEventListener("keydown", (event) => {
    if (event.key !== "Escape") return;
    const active = document.querySelector(".featurelens-inline-focus");
    if (active) restoreFocus(active);
  });

  // Rename Gradio's generic chart.png export without touching the export implementation.
  document.addEventListener("click", (event) => {
    const button = event.target && event.target.closest ? event.target.closest("button") : null;
    if (button) {
      const label = `${button.getAttribute("aria-label") || ""} ${button.getAttribute("title") || ""} ${button.textContent || ""}`.toLowerCase();
      if (label.includes("export")) {
        const block = button.closest(".fl-plot");
        if (block) {
          const id = block.id || "plot-featurelens-chart";
          const stem = id.replace(/^plot-/, "").replace(/[^a-z0-9_-]+/gi, "-").replace(/-+/g, "-").replace(/^-|-$/g, "");
          window.__featurelens_export_name = `featurelens_${stem || "chart"}.png`;
        }
      }
    }
    const anchor = event.target && event.target.closest ? event.target.closest('a[download="chart.png"]') : null;
    if (anchor && window.__featurelens_export_name) {
      anchor.setAttribute("download", window.__featurelens_export_name);
      window.setTimeout(() => { window.__featurelens_export_name = null; }, 500);
    }
  }, true);

  kick();
  return [];
}
"""


def _raise_ui_error(exc: Exception) -> None:
    raise gr.Error(f"{type(exc).__name__}: {exc}") from exc


def _tsv(frame: pd.DataFrame) -> str:
    if frame is None or frame.empty:
        return ""
    return frame.to_csv(sep="\t", index=False, lineterminator="\n")


def _copy_button(label: str = "Copy table with headers") -> gr.Button:
    return gr.Button(label, size="sm", variant="primary", elem_classes=["copy-btn"])


def _table_heading(text: str) -> gr.HTML:
    return gr.HTML(f'<div class="table-heading">{text}</div>')


def _copy_ack(_text: str) -> None:
    gr.Info("Copied table with headers.", duration=1.2)


def _bind_copy(button: gr.Button, source: gr.Textbox) -> None:
    button.click(fn=_copy_ack, inputs=[source], outputs=None, js=COPY_JS, queue=False)


def _analysis_metrics_markdown(result) -> str:
    return (
        "#### Analysis metrics\n"
        f"**Layer {result.layer} · prompt token {result.token_index}**  \n"
        f"Active SAE features: **{int(result.metrics['active_features'])}/{SETTINGS.sae_top_k}**  \n"
        f"Reconstruction cosine: **{result.metrics['cosine']:.4f}** · "
        f"NMSE: **{result.metrics['nmse']:.4f}**  \n"
        f"Top-5 activation mass: **{result.metrics['top5_mass_fraction']:.1%}**"
    )


def _intervention_metrics_markdown(result) -> str:
    drift = f"Execution-context null drift — JS: **{result.execution_drift_js:.2e}**"
    if result.execution_drift_mean_logprob is not None:
        drift += f" · mean log p/token: **{result.execution_drift_mean_logprob:+.2e}**"

    target = (
        "No target continuation supplied. The causal comparison therefore uses only "
        "next-token Jensen-Shannon divergence."
    )
    if result.baseline_sequence_logprob is not None:
        tokens = " ".join(repr(token) for token in result.target_tokens)
        target = (
            f"**Target continuation:** {result.target_token_count} token(s): {tokens}  \n"
            f"Sequence log p — baseline: **{result.baseline_sequence_logprob:.4f}** · "
            f"SAE edit: **{result.modified_sequence_logprob:.4f}** · "
            f"random-ensemble mean: **{result.random_sequence_logprob:.4f}**  \n"
            f"SAE Δ sequence log p: **{result.sequence_logprob_delta:+.4f}** · "
            f"random signed mean Δ: **{result.random_sequence_logprob_delta:+.4f}**  \n"
            f"SAE Δ mean log p/token: **{result.mean_logprob_delta:+.4f}** · "
            f"random mean |Δ|: **{result.random_abs_mean_logprob_delta:.4f}** ± **{result.random_mean_logprob_std:.4f}** · "
            f"magnitude ratio: **{result.target_specificity_ratio:.2f}×** · "
            f"empirical tail p: **{result.target_empirical_p:.3f}**"
        )
    inactive = ""
    if abs(result.feature_activation) < 1e-12:
        inactive = (
            "  \n**Inactive feature:** ablate/scale has zero coefficient delta here; "
            "use `inject` only when you intentionally want a decoder-direction steering test."
        )
    return (
        f"Original feature activation: **{result.feature_activation:.4f}** · "
        f"Δ coefficient: **{result.delta_activation:+.4f}**  \n"
        f"Perturbation L2: **{result.perturbation_norm:.4f}**  \n"
        f"Next-token JS: **{result.js_divergence:.6f}** · "
        f"random ensemble ({result.random_control_count}) mean JS: **{result.random_js_divergence:.6f}** "
        f"± **{result.random_js_std:.6f}** · ratio: **{result.js_specificity_ratio:.2f}×** · "
        f"empirical tail p: **{result.js_empirical_p:.3f}**  \n"
        f"{drift}  \n\n{target}{inactive}  \n\n"
        "_Greedy text may remain unchanged even when probability-level causal metrics move; "
        "deterministic generation changes only after an argmax boundary is crossed._"
    )


def _dose_metrics_markdown(result) -> str:
    tokens = " ".join(repr(token) for token in result.target_tokens)
    inactive = ""
    if abs(result.feature_activation) < 1e-12:
        inactive = " **The feature is inactive here, so multiplicative scaling is flat by construction.**"
    return (
        f"Feature activation at baseline: **{result.feature_activation:.4f}** · "
        f"target continuation: {len(result.target_tokens)} token(s): {tokens}.{inactive}  \n\n"
        "This panel is **always a scale intervention**: 0× = ablation, 1× = no edit, "
        "2× = double the native coefficient. All rows are compared with the **1× row from the same batched "
        "execution context**, so 1× is the numerical null reference.  \n"
        f"Single-forward → batched-null execution drift: mean log p/token "
        f"**{result.execution_drift_mean_logprob:+.2e}**, JS **{result.execution_drift_js:.2e}**."
    )


def _feature_set_metrics_markdown(result) -> str:
    tokens = " ".join(repr(token) for token in result.target_tokens)
    inactive_count = sum(abs(float(row[1])) < 1e-12 for row in result.feature_rows)
    inactive_note = (
        f"  \n{inactive_count} selected feature(s) were inactive and contributed zero delta."
        if inactive_count
        else ""
    )
    return (
        f"Selected feature set: **{len(result.feature_ids)} features** · "
        f"perturbation L2: **{result.perturbation_norm:.4f}**  \n"
        f"Target continuation: {len(result.target_tokens)} token(s): {tokens}  \n"
        f"SAE Δ mean log p/token: **{result.mean_logprob_delta:+.4f}** · "
        f"random ensemble ({result.random_control_count}) mean |Δ|: "
        f"**{result.random_abs_mean_logprob_delta:.4f}** ± **{result.random_mean_logprob_std:.4f}** · ratio: **{result.target_specificity_ratio:.2f}×** · "
        f"empirical tail p: **{result.target_empirical_p:.3f}**  \n"
        f"SAE Δ sequence log p: **{result.sequence_logprob_delta:+.4f}** · "
        f"random signed mean Δ: **{result.random_sequence_logprob_delta:+.4f}**  \n"
        f"Next-token JS: **{result.js_divergence:.6f}** · random mean JS: "
        f"**{result.random_js_divergence:.6f}** ± **{result.random_js_std:.6f}** · "
        f"ratio: **{result.js_specificity_ratio:.2f}×** · empirical tail p: **{result.js_empirical_p:.3f}**  \n"
        f"Execution-context null drift: mean log p/token **{result.execution_drift_mean_logprob:+.2e}**, "
        f"JS **{result.execution_drift_js:.2e}**{inactive_note}"
    )


def _interaction_metrics_markdown(result) -> str:
    tokens = " ".join(repr(token) for token in result.target_tokens)
    return (
        f"Target continuation: {len(result.target_tokens)} token(s): {tokens}  \n"
        f"Additive expectation (sum of individual Δ mean log p/token): "
        f"**{result.additive_expected_mean_delta:+.4f}**  \n"
        f"Observed joint Δ mean log p/token: **{result.joint_mean_delta:+.4f}**  \n"
        f"Interaction excess (joint − additive): **{result.interaction_excess_mean_delta:+.4f}** · "
        f"normalized: **{result.normalized_interaction:+.3f}**  \n"
        f"Execution-context null drift: **{result.execution_drift_mean_logprob:+.2e}** mean log p/token.  \n\n"
        "A non-zero excess indicates **non-additive downstream response under these edits**; it is not, by itself, "
        "proof that the SAE features form a direct mechanistic circuit."
    )


def _paraphrase_metrics_markdown(result) -> str:
    return (
        "#### Robustness metrics\n"
        "**Selected-token comparison**  \n"
        f"TopK feature-set Jaccard: **{result.topk_jaccard:.3f}** · "
        f"sparse activation cosine: **{result.sparse_cosine:.3f}**  \n"
        f"Shared features among displayed top-{result.top_n}: **{result.shared_top_n}/{result.top_n}**  \n\n"
        "**Prompt-wide comparison** — max activation of each feature across all prompt tokens  \n"
        f"Feature-set Jaccard: **{result.promptwide_jaccard:.3f}** · "
        f"activation cosine: **{result.promptwide_cosine:.3f}**  \n\n"
        "The prompt-wide view reduces sensitivity to choosing semantically different final tokens, while the "
        "selected-token view remains the stricter local representation test. Neither metric establishes feature semantics."
    )


def _concept_metrics_markdown(result) -> str:
    coverage = f"{result.active_prompt_count}/{result.total_prompt_count}"
    if result.leading_concept is None:
        leader = (
            "The feature was **inactive in every sampled prompt**. No concept leader is reported; "
            "increase the sample or inspect a different feature."
        )
    elif result.leading_ratio is None:
        leader = (
            f"Highest prompt-wide mean activation: **{result.leading_concept}**; the runner-up mean was zero."
        )
    else:
        leader = (
            f"Highest prompt-wide mean activation: **{result.leading_concept}** "
            f"(**{result.leading_ratio:.2f}×** the runner-up mean)."
        )
    return (
        f"Feature **{result.feature_id}**, layer **{result.layer}** · "
        f"{result.prompts_per_concept} controlled prompts/concept · active in **{coverage}** sampled prompts.  \n"
        f"{leader}  \n\n"
        "Each prompt is summarized by the feature's **maximum activation across its non-padding tokens**. "
        "This is an exploratory contrast scan, not an automatic semantic label; held-out AUROC/F1 remains the stronger evidence."
    )


def _trace_metrics_markdown(result) -> str:
    if result.max_token_index is None:
        peak = "Feature is inactive at every prompt token."
    else:
        peak = (
            f"Peak activation **{result.max_activation:.4f}** at token **{result.max_token_index}** "
            f"({result.tokens[result.max_token_index]!r})."
        )
    return (
        f"Feature **{result.feature_id}**, layer **{result.layer}** · active at "
        f"**{result.active_token_count}/{result.token_count}** prompt tokens.  \n{peak}"
    )


def _geometry_metrics_markdown(result) -> str:
    if result.alignment_ratio > 1.05:
        geometry = "The activation-weighted ablation directions are net **aligned**."
    elif result.alignment_ratio < 0.95:
        geometry = "The activation-weighted ablation directions show net **cancellation**."
    else:
        geometry = "The joint norm is close to the orthogonal/independent reference."
    return (
        f"Selected features: **{', '.join(str(x) for x in result.feature_ids)}** · layer **{result.layer}**  \n"
        f"Mean |decoder cosine|: **{result.mean_abs_decoder_cosine:.3f}** · "
        f"max |decoder cosine|: **{result.max_abs_decoder_cosine:.3f}**  \n"
        f"Joint ablation L2: **{result.joint_ablation_norm:.4f}** · independent-direction reference: "
        f"**{result.independent_norm:.4f}** · geometry ratio: **{result.alignment_ratio:.3f}×**  \n"
        f"{geometry} This diagnoses SAE decoder geometry; it does not by itself establish downstream causal interaction."
    )


def _contrastive_metrics_markdown(result) -> str:
    direction = "toward A" if result.delta_log_odds > 0 else ("toward B" if result.delta_log_odds < 0 else "neither way")
    return (
        f"Feature **{result.feature_id}** · activation **{result.feature_activation:.4f}** · perturbation L2 "
        f"**{result.perturbation_norm:.4f}**  \n"
        f"Exact-sequence log-odds A−B — baseline: **{result.baseline_log_odds:+.4f}** · "
        f"SAE edit: **{result.modified_log_odds:+.4f}** · causal shift: **{result.delta_log_odds:+.4f}** ({direction}).  \n"
        f"Token-normalized preference shift: **{result.delta_normalized_preference:+.4f}**.  \n"
        f"Random ensemble ({result.random_control_count}) |Δ log-odds| mean: **{result.random_abs_mean_delta:.4f}** "
        f"± **{result.random_delta_std:.4f}** · SAE/random ratio: **{result.specificity_ratio:.2f}×** · "
        f"empirical tail p: **{result.empirical_p:.3f}**.  \n\n"
        "This asks whether the intervention changes the model's **relative preference between two specified continuations**, "
        "which is stricter than showing that one continuation's probability moved in isolation."
    )




def _discovery_metrics_markdown(result) -> str:
    if not result.candidate_ids:
        if result.ranking_mode == "causal_ready":
            return (
                f"No positively selective **current-token-active** candidates were found for **{result.concept}** "
                f"at layer **{result.layer}** in this live batch. Try another token/layer or use Balanced selectivity "
                "to discover concept-associated features elsewhere in the prompt."
            )
        return (
            f"No positively selective candidate features were found for **{result.concept}** at layer "
            f"**{result.layer}** in this small live batch. Increase the sample or try another layer."
        )
    ranking = {
        "balanced_selectivity": "balanced selectivity × coverage × log-magnitude",
        "raw_mean_difference": "raw target-minus-other mean activation",
        "causal_ready": "causal-ready evidence at the current Workbench token",
    }[result.ranking_mode]
    context = (
        f"Current Workbench token **{result.current_token_index}** was checked in the same model batch. "
        "The default candidate prefers the highest-ranked feature active at that token when one exists."
        if result.current_context_available
        else "No Workbench prompt was available, so current-prompt compatibility was not evaluated."
    )
    return (
        f"Concept **{result.concept}** · layer **{result.layer}** · "
        f"{result.prompts_per_concept} prompts/concept.  \n"
        f"Showing **{len(result.candidate_ids)}** candidates ranked by **{ranking}**.  \n"
        f"Displayed candidates active at the selected Workbench token: "
        f"**{result.displayed_current_active_count}/{len(result.candidate_ids)}**.  \n"
        f"{context}  \n\n"
        "This is **candidate discovery**, not a semantic label. The offline held-out AUROC/F1 pipeline remains "
        "the standard for a concept-selectivity claim."
    )


def _candidate_screen_metrics_markdown(result) -> str:
    tokens = " ".join(repr(token) for token in result.target_tokens)
    if result.rows:
        top = result.rows[0]
        strongest = (
            f"Strongest screened target effect: feature **{int(top[1])}** with "
            f"Δ mean log p/token **{float(top[5]):+.4f}** and next-token JS **{float(top[7]):.6f}**."
        )
    else:
        strongest = "No candidate rows were produced."
    return (
        f"Screened **{result.candidate_count}** candidate feature(s); **{result.active_feature_count}** were active "
        f"at the selected Workbench token. Target continuation: {len(result.target_tokens)} token(s): {tokens}.  \n"
        f"{strongest}  \n"
        f"Execution-context null drift: mean log p/token **{result.execution_drift_mean_logprob:+.2e}**, "
        f"JS **{result.execution_drift_js:.2e}**.  \n\n"
        "This is a **triage screen**: all ablations share one batched zero-edit reference, but no random-control "
        "ensemble is spent here. Take a promising feature into **Workbench → II. Single-feature causal test** for "
        "the full 8-direction specificity comparison."
    )



def _spearman_rank_corr(left: list[float], right: list[float]) -> float | None:
    """Descriptive Spearman correlation with tie-aware average ranks."""
    if len(left) != len(right) or len(left) < 2:
        return None
    left_s = pd.Series(left, dtype="float64")
    right_s = pd.Series(right, dtype="float64")
    if left_s.nunique(dropna=True) < 2 or right_s.nunique(dropna=True) < 2:
        return None
    value = left_s.rank(method="average").corr(right_s.rank(method="average"))
    return None if pd.isna(value) else float(value)


def _candidate_alignment_outputs(
    discovery_table: pd.DataFrame | None,
    screen_table: pd.DataFrame | None,
) -> tuple[str, pd.DataFrame, pd.DataFrame]:
    """Join discovery evidence to causal triage results without another model call."""
    empty_columns = [
        "Feature id",
        "Discovery rank",
        "Target-effect rank",
        "Distribution-shift rank",
        "Candidate score",
        "Selectivity",
        "Current token activation",
        "|Δ mean log p/token|",
        "Next-token JS",
        "Discovery→target rank shift",
    ]
    if discovery_table is None or screen_table is None:
        return "", pd.DataFrame(columns=empty_columns), pd.DataFrame()
    discovery = pd.DataFrame(discovery_table).copy()
    screen = pd.DataFrame(screen_table).copy()
    if discovery.empty or screen.empty or "Feature id" not in discovery or "Feature id" not in screen:
        return "", pd.DataFrame(columns=empty_columns), pd.DataFrame()

    discovery["Feature id"] = pd.to_numeric(discovery["Feature id"], errors="coerce")
    screen["Feature id"] = pd.to_numeric(screen["Feature id"], errors="coerce")
    discovery = discovery.dropna(subset=["Feature id"]).copy()
    screen = screen.dropna(subset=["Feature id"]).copy()
    discovery["Feature id"] = discovery["Feature id"].astype(int)
    screen["Feature id"] = screen["Feature id"].astype(int)

    needed_discovery = {"Rank", "Candidate score", "Selectivity", "Current token activation"}
    needed_screen = {"Rank", "Δ mean log p/token", "Next-token JS"}
    if not needed_discovery.issubset(discovery.columns) or not needed_screen.issubset(screen.columns):
        return "", pd.DataFrame(columns=empty_columns), pd.DataFrame()

    discovery_lookup = discovery.set_index("Feature id", drop=False)
    js_ranked = screen.sort_values(["Next-token JS", "Feature id"], ascending=[False, True]).reset_index(drop=True)
    js_ranks = {int(row["Feature id"]): rank for rank, (_, row) in enumerate(js_ranked.iterrows(), start=1)}

    rows: list[list[object]] = []
    for _, causal_row in screen.iterrows():
        feature_id = int(causal_row["Feature id"])
        if feature_id not in discovery_lookup.index:
            continue
        discovery_row = discovery_lookup.loc[feature_id]
        # set_index can technically return a DataFrame for duplicate ids; use the first row deterministically.
        if isinstance(discovery_row, pd.DataFrame):
            discovery_row = discovery_row.iloc[0]
        discovery_rank = int(float(discovery_row["Rank"]))
        target_rank = int(float(causal_row["Rank"]))
        mean_delta = float(causal_row["Δ mean log p/token"])
        rows.append(
            [
                feature_id,
                discovery_rank,
                target_rank,
                int(js_ranks[feature_id]),
                float(discovery_row["Candidate score"]),
                float(discovery_row["Selectivity"]),
                float(discovery_row["Current token activation"]),
                abs(mean_delta),
                float(causal_row["Next-token JS"]),
                discovery_rank - target_rank,
            ]
        )

    table = pd.DataFrame(rows, columns=empty_columns)
    if table.empty:
        return "", table, pd.DataFrame()

    rho_target = _spearman_rank_corr(
        table["Candidate score"].astype(float).tolist(),
        table["|Δ mean log p/token|"].astype(float).tolist(),
    )
    rho_js = _spearman_rank_corr(
        table["Candidate score"].astype(float).tolist(),
        table["Next-token JS"].astype(float).tolist(),
    )

    top_discovery = table.sort_values(["Discovery rank", "Feature id"]).iloc[0]
    top_target = table.sort_values(["Target-effect rank", "Feature id"]).iloc[0]
    top_js = table.sort_values(["Distribution-shift rank", "Feature id"]).iloc[0]

    def fmt_rho(value: float | None) -> str:
        return "undefined" if value is None else f"{value:+.3f}"

    summary = (
        f"Compared **{len(table)}** screened candidates using the discovery evidence and causal triage from the same workflow.  \n"
        f"Top discovery candidate: **{int(top_discovery['Feature id'])}** · strongest target effect: "
        f"**{int(top_target['Feature id'])}** · strongest next-token distribution shift: **{int(top_js['Feature id'])}**.  \n"
        f"Spearman ρ(candidate score, |target effect|): **{fmt_rho(rho_target)}** · "
        f"ρ(candidate score, next-token JS): **{fmt_rho(rho_js)}**.  \n\n"
        "This is the project's central association-vs-causality diagnostic: a feature can rank highly by concept evidence yet "
        "rank weakly by downstream target effect, or vice versa. With only a handful of screened features and no random-control "
        "ensemble in triage, these correlations are **descriptive**, not significance claims."
    )

    chart = table[["Feature id", "Candidate score", "|Δ mean log p/token|", "Discovery rank", "Target-effect rank", "Next-token JS"]].copy()
    chart["Feature id"] = chart["Feature id"].astype(str)
    chart["Series"] = "Screened candidate"
    return summary, table, chart


def _cue_context_metrics_markdown(result) -> str:
    active = ", ".join(
        f"{cue}: {count}/{len(result.stems)} contexts"
        for cue, count in result.cue_active_context_counts.items()
    )
    interpretation = (
        "No tested cue activated the feature."
        if result.dominant_cue is None or result.active_condition_count == 0
        else (
            f"**Cue-dominant pattern:** `{result.dominant_cue}` activates in every tested context while all "
            "other tested cues are inactive. Under this matrix, the evidence is much more consistent with a "
            "lexical/cue-specific response than with a mathematics-specific response."
            if result.dominant_cue_context_count == len(result.stems)
            and result.off_dominant_active_count == 0
            else (
                f"The strongest cue is `{result.dominant_cue}` ({result.dominant_cue_context_count}/{len(result.stems)} "
                "contexts), but other cues or context dependence remain. Treat the pattern as mixed/context-sensitive."
            )
        )
    )
    return (
        f"Feature **{result.feature_id}**, layer **{result.layer}** · active in "
        f"**{result.active_condition_count}/{result.condition_count}** stem × cue conditions.  \n"
        f"Cue coverage — {active}.  \n\n"
        f"{interpretation}  \n\n"
        "This is a controlled diagnostic over the tested stems/cues, not a universal semantic label."
    )


def _cue_metrics_markdown(result) -> str:
    return (
        f"Feature **{result.feature_id}**, layer **{result.layer}** · active for "
        f"**{result.active_cue_count}/{result.cue_count}** tested completion cues at the final token.  \n"
        "This helps distinguish a concept-associated feature from a lexical/structural completion-cue feature. "
        "It is a controlled diagnostic, not a feature label."
    )


def _global_context_markdown(prompt: str, layer: int, result) -> str:
    token = result.tokens[result.token_index] if result.tokens else ""
    short = prompt[:120] + ("…" if len(prompt) > 120 else "")
    return (
        "**Current Workbench context** — "
        f"Prompt: `{short}` · layer **{int(layer)}** · token **{result.token_index}** ({token!r}).  \n"
        "Panels marked **Uses current Workbench context** inherit this prompt/layer/token. "
        "Feature selectors remain editable inside each experiment."
    )

@gpu(duration=30)
def analyze_prompt(prompt: str, layer: int, token_index: int, top_n: int):
    try:
        if not prompt.strip():
            raise ValueError("Enter a prompt first.")
        result = RUNTIME.analyze(prompt, int(layer), int(token_index), int(top_n))
        columns = ["Rank", "Feature id", "Activation", "Offline concept hint"]
        table = pd.DataFrame(result.rows, columns=columns)
        choices = [str(int(row[1])) for row in result.rows]
        feature_update = gr.update(choices=choices, value=choices[0] if choices else None)
        feature_set_update = gr.update(choices=choices, value=choices[: min(3, len(choices))])
        contrast_update = gr.update(choices=choices, value=choices[0] if choices else None)
        chart_df = pd.DataFrame(
            {
                "Feature": [str(int(row[1])) for row in result.rows],
                "Activation": [float(row[2]) for row in result.rows],
                "Series": ["Activation"] * len(result.rows),
            }
        )
        location = (
            f"Current Workbench location — **layer {int(layer)}**, **prompt token {result.token_index}**; "
            f"prompt: `{prompt[:90]}{'…' if len(prompt) > 90 else ''}`"
        )
        return (
            RUNTIME.token_html(result.tokens, result.token_index),
            table,
            chart_df,
            feature_update,
            gr.update(choices=choices, value=choices[0] if choices else None),
            gr.update(choices=choices, value=choices[0] if choices else None),
            feature_set_update,
            contrast_update,
            gr.update(value=int(layer)),
            _analysis_metrics_markdown(result),
            location,
            location,
            _global_context_markdown(prompt, int(layer), result),
            _tsv(table),
        )
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=45)
def run_intervention(
    prompt: str,
    layer: int,
    token_index: int,
    feature_id: str,
    mode: str,
    coefficient: float,
    target_text: str,
    max_new_tokens: int,
):
    try:
        if not prompt.strip():
            raise ValueError("Enter a prompt first.")
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose or enter a feature id.")
        result = RUNTIME.intervene(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_id=int(float(feature_id)),
            mode=mode,
            coefficient=float(coefficient),
            target_text=target_text,
            max_new_tokens=int(max_new_tokens),
        )
        token_columns = ["Token", "Baseline p", "SAE-edit p", "Δ probability"]
        target_columns = [
            "Target position",
            "Target token",
            "Baseline log p",
            "SAE-edit log p",
            "Random-ensemble mean log p",
            "SAE Δ log p",
            "Random-ensemble mean Δ log p",
        ]
        token_df = pd.DataFrame(result.top_token_rows, columns=token_columns)
        target_df = pd.DataFrame(result.target_token_rows, columns=target_columns)
        return (
            result.baseline_text,
            result.modified_text,
            _intervention_metrics_markdown(result),
            token_df,
            target_df,
            _tsv(token_df),
            _tsv(target_df),
        )
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=35)
def run_dose_response(prompt: str, layer: int, token_index: int, feature_id: str, target_text: str):
    try:
        if not prompt.strip():
            raise ValueError("Enter a prompt first.")
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose or enter a feature id.")
        if not target_text.strip():
            raise ValueError("Enter a target continuation before running the scale dose-response.")
        result = RUNTIME.dose_response(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_id=int(float(feature_id)),
            target_text=target_text,
        )
        columns = [
            "Multiplier",
            "Δ feature coefficient",
            "Perturbation L2",
            "Batched-null mean log p/token",
            "Modified mean log p/token",
            "Δ mean log p/token",
            "Δ sequence log p",
            "Next-token JS",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        plot = table[["Multiplier", "Δ mean log p/token"]].copy()
        plot["Series"] = "SAE feature"
        return table, plot, _dose_metrics_markdown(result), _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=35)
def run_layer_sweep(prompt: str, token_index: int):
    try:
        if not prompt.strip():
            raise ValueError("Enter a prompt first.")
        result = RUNTIME.layer_sweep(prompt, int(token_index))
        columns = [
            "Layer",
            "Reconstruction cosine",
            "NMSE",
            "Active features",
            "Top activation",
            "Top-5 mass",
            "Activation entropy",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        long = table.melt(
            id_vars=["Layer"],
            value_vars=["Reconstruction cosine", "Top-5 mass", "Activation entropy"],
            var_name="Metric",
            value_name="Value",
        )
        return RUNTIME.token_html(result.tokens, result.token_index), table, long, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=40)
def run_feature_set(
    prompt: str,
    layer: int,
    token_index: int,
    feature_ids: list[str] | None,
    mode: str,
    coefficient: float,
    target_text: str,
):
    try:
        if not prompt.strip():
            raise ValueError("Enter and inspect a prompt in the Workbench first.")
        selected = [int(float(value)) for value in (feature_ids or [])]
        if not selected:
            raise ValueError("Select at least one feature in 'Feature set'.")
        if not target_text.strip():
            raise ValueError("Enter a target continuation for the feature-set causal test.")
        result = RUNTIME.intervene_feature_set(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_ids=selected,
            mode=mode,
            coefficient=float(coefficient),
            target_text=target_text,
        )
        feature_columns = ["Feature id", "Original activation", "Δ coefficient", "Offline concept hint"]
        target_columns = [
            "Target position",
            "Target token",
            "Baseline log p",
            "SAE-edit log p",
            "Random-ensemble mean log p",
            "SAE Δ log p",
            "Random-ensemble mean Δ log p",
        ]
        feature_df = pd.DataFrame(result.feature_rows, columns=feature_columns)
        target_df = pd.DataFrame(result.target_token_rows, columns=target_columns)
        return (
            feature_df,
            _feature_set_metrics_markdown(result),
            target_df,
            _tsv(feature_df),
            _tsv(target_df),
        )
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=45)
def run_feature_set_sweep(prompt: str, layer: int, token_index: int, target_text: str):
    try:
        if not prompt.strip():
            raise ValueError("Enter and inspect a prompt in the Workbench first.")
        if not target_text.strip():
            raise ValueError("Enter a target continuation before running the set-size sweep.")
        result = RUNTIME.feature_set_size_sweep(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            target_text=target_text,
        )
        columns = [
            "Set size k",
            "Feature ids",
            "Perturbation L2",
            "Batched-null mean log p/token",
            "SAE mean log p/token",
            "SAE Δ mean log p/token",
            "Random signed mean Δ",
            "Random mean |Δ|",
            "Random |Δ| std",
            "SAE/random magnitude ratio",
            "Empirical tail p",
            "SAE Δ sequence log p",
            "SAE next-token JS",
            "Random mean JS",
            "Random JS std",
            "JS empirical tail p",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        plot_rows: list[list[object]] = []
        for _, row in table.iterrows():
            plot_rows.append([row["Set size k"], "Top-k SAE ablation", row["SAE Δ mean log p/token"]])
            plot_rows.append([row["Set size k"], "Random signed mean", row["Random signed mean Δ"]])
        plot = pd.DataFrame(plot_rows, columns=["Set size k", "Condition", "Δ mean log p/token"])
        tokens = " ".join(repr(token) for token in result.target_tokens)
        note = (
            f"Target continuation: {len(result.target_tokens)} token(s): {tokens}. For each k, FeatureLens "
            f"ablates the k strongest active features and compares the effect with **{result.random_control_count} "
            f"norm-matched random directions**. All conditions share one batched zero-edit reference.  \n"
            f"Execution-context null drift: mean log p/token **{result.execution_drift_mean_logprob:+.2e}**, "
            f"JS **{result.execution_drift_js:.2e}**. The live empirical p-value is intentionally coarse because "
            f"it uses only {result.random_control_count} controls; the offline experiment should use more."
        )
        return table, plot, note, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=40)
def run_feature_interaction(
    prompt: str,
    layer: int,
    token_index: int,
    feature_ids: list[str] | None,
    target_text: str,
):
    try:
        selected = [int(float(value)) for value in (feature_ids or [])]
        if not prompt.strip():
            raise ValueError("Enter and inspect a prompt in the Workbench first.")
        if len(selected) < 2:
            raise ValueError("Select at least two features in 'Feature set'.")
        if len(selected) > 5:
            raise ValueError("Select at most five features for the interaction decomposition.")
        if not target_text.strip():
            raise ValueError("Enter a target continuation for the interaction decomposition.")
        result = RUNTIME.feature_interaction_test(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_ids=selected,
            target_text=target_text,
        )
        columns = [
            "Condition",
            "Feature ids",
            "Activation summary",
            "Perturbation L2",
            "Δ mean log p/token",
            "Δ sequence log p",
            "Next-token JS",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        plot = table[["Condition", "Δ mean log p/token"]].copy()
        plot["Series"] = "Ablation effect"
        return table, _interaction_metrics_markdown(result), plot, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=30)
def run_paraphrase_compare(
    original_prompt: str,
    paraphrase_prompt: str,
    layer: int,
    token_index_a: int,
    token_index_b: int,
    top_n: int,
):
    try:
        result = RUNTIME.compare_paraphrases(
            text_a=original_prompt,
            text_b=paraphrase_prompt,
            layer=int(layer),
            token_index_a=int(token_index_a),
            token_index_b=int(token_index_b),
            top_n=int(top_n),
        )
        columns = ["Feature id", "Original activation", "Paraphrase activation", "Status", "Offline concept hint"]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Feature", "Prompt", "Activation"])
        return (
            RUNTIME.token_html(result.tokens_a, result.token_index_a),
            RUNTIME.token_html(result.tokens_b, result.token_index_b),
            _paraphrase_metrics_markdown(result),
            table,
            chart,
            _tsv(table),
        )
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=35)
def run_concept_contrast(feature_id: str, layer: int, prompts_per_concept: int):
    try:
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose a feature id first. Run Workbench inspection if the selector is empty.")
        result = RUNTIME.concept_contrast_scan(
            feature_id=int(float(feature_id)),
            layer=int(layer),
            prompts_per_concept=int(prompts_per_concept),
        )
        columns = [
            "Concept",
            "Prompts",
            "Mean prompt-wide max",
            "Median prompt-wide max",
            "Prompt activation rate",
            "Mean when active",
            "Max activation",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Concept", "Mean prompt-wide max"])
        chart["Series"] = "Prompt-wide max"
        return _concept_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=30)
def run_feature_trace(prompt: str, layer: int, feature_id: str):
    try:
        if not prompt.strip():
            raise ValueError("Enter a prompt first.")
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose a feature id first.")
        result = RUNTIME.feature_token_trace(
            text=prompt,
            layer=int(layer),
            feature_id=int(float(feature_id)),
        )
        columns = ["Token position", "Token", "Activation", "Active in TopK"]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Token", "Activation"])
        chart["Series"] = "Feature activation"
        return _trace_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=30)
def run_feature_geometry(prompt: str, layer: int, token_index: int, feature_ids: list[str] | None):
    try:
        selected = [int(float(value)) for value in (feature_ids or [])]
        result = RUNTIME.feature_geometry(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_ids=selected,
        )
        columns = ["Feature A", "Feature B", "Activation A", "Activation B", "Decoder cosine"]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Feature pair", "Decoder cosine"])
        chart["Series"] = "Decoder cosine"
        return _geometry_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=40)
def run_contrastive_causal(
    prompt: str,
    layer: int,
    token_index: int,
    feature_id: str,
    mode: str,
    coefficient: float,
    target_a: str,
    target_b: str,
):
    try:
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose a feature id first.")
        result = RUNTIME.contrastive_intervention(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_id=int(float(feature_id)),
            mode=mode,
            coefficient=float(coefficient),
            target_a=target_a,
            target_b=target_b,
        )
        columns = [
            "Continuation",
            "Text",
            "Tokens",
            "Baseline sequence log p",
            "SAE-edit sequence log p",
            "Δ sequence log p",
            "Baseline mean log p/token",
            "SAE-edit mean log p/token",
            "Δ mean log p/token",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(
            [
                ["Baseline", result.baseline_log_odds],
                ["SAE edit", result.modified_log_odds],
            ],
            columns=["Condition", "A−B sequence log-odds"],
        )
        chart["Series"] = "Contrastive preference"
        return _contrastive_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)



@gpu(duration=35)
def run_concept_feature_discovery(
    concept: str,
    layer: int,
    prompts_per_concept: int,
    top_n: int,
    ranking_label: str,
    workbench_prompt: str,
    workbench_token_index: int,
):
    try:
        ranking_mode = {
            "Balanced selectivity": "balanced_selectivity",
            "Raw mean difference": "raw_mean_difference",
            "Causal-ready at current token": "causal_ready",
        }[ranking_label]
        result = RUNTIME.concept_feature_discovery(
            concept=concept,
            layer=int(layer),
            prompts_per_concept=int(prompts_per_concept),
            top_n=int(top_n),
            ranking_mode=ranking_mode,
            current_text=workbench_prompt,
            current_token_index=int(workbench_token_index),
        )
        columns = [
            "Rank",
            "Feature id",
            "Candidate score",
            "Target mean max",
            "Other mean max",
            "Mean difference",
            "Selectivity",
            "Target activation rate",
            "Other activation rate",
            "Current prompt max",
            "Current token activation",
            "Active at current token",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Feature", "Candidate score"])
        chart["Series"] = "Candidate score"
        choices = [str(fid) for fid in result.candidate_ids]
        default = str(result.default_candidate_id) if result.default_candidate_id is not None else (choices[0] if choices else None)
        candidate_update = gr.update(choices=choices, value=default)
        screen_update = gr.update(
            choices=choices,
            value=choices[: min(5, len(choices))],
        )
        return (
            _discovery_metrics_markdown(result),
            table,
            chart,
            candidate_update,
            screen_update,
            _tsv(table),
        )
    except Exception as exc:
        _raise_ui_error(exc)



@gpu(duration=30)
def run_candidate_causal_screen(
    prompt: str,
    layer: int,
    token_index: int,
    feature_ids: list[str] | None,
    target_text: str,
    discovery_table: pd.DataFrame | None,
):
    try:
        selected = [int(float(value)) for value in (feature_ids or [])]
        result = RUNTIME.candidate_causal_screen(
            text=prompt,
            layer=int(layer),
            token_index=int(token_index),
            feature_ids=selected,
            target_text=target_text,
        )
        columns = [
            "Rank",
            "Feature id",
            "Native activation",
            "Active at current token",
            "Perturbation L2",
            "Δ mean log p/token",
            "Δ sequence log p",
            "Next-token JS",
        ]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(
            result.chart_rows,
            columns=["Feature", "Δ mean log p/token"],
        )
        chart["Series"] = "Candidate ablation"
        choices = [str(feature_id) for feature_id in result.feature_ids]
        candidate_update = gr.update(
            choices=choices,
            value=choices[0] if choices else None,
        )
        alignment_metrics, alignment_table, alignment_chart = _candidate_alignment_outputs(
            discovery_table, table
        )
        return (
            _candidate_screen_metrics_markdown(result),
            table,
            chart,
            candidate_update,
            _tsv(table),
            alignment_metrics,
            alignment_table,
            alignment_chart,
            _tsv(alignment_table),
        )
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=25)
def run_feature_cue_scan(feature_id: str, layer: int, prompt_stem: str, cue_text: str):
    try:
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose a feature id first.")
        cues = [line for line in str(cue_text).splitlines() if line.strip()]
        result = RUNTIME.feature_cue_scan(
            feature_id=int(float(feature_id)),
            layer=int(layer),
            prompt_stem=prompt_stem,
            cues=cues,
        )
        columns = ["Cue", "Full prompt", "Final token", "Activation", "Active in TopK"]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Cue", "Activation"])
        chart["Series"] = "Cue response"
        return _cue_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


@gpu(duration=30)
def run_feature_cue_context_scan(feature_id: str, layer: int, stems_text: str, cue_text: str):
    try:
        if feature_id is None or str(feature_id).strip() == "":
            raise ValueError("Choose a feature id first.")
        stems = [line for line in str(stems_text).splitlines() if line.strip()]
        cues = [line for line in str(cue_text).splitlines() if line.strip()]
        result = RUNTIME.feature_cue_context_scan(
            feature_id=int(float(feature_id)),
            layer=int(layer),
            stems=stems,
            cues=cues,
        )
        columns = ["Prompt stem", "Cue", "Full prompt", "Final token", "Activation", "Active in TopK"]
        table = pd.DataFrame(result.rows, columns=columns)
        chart = pd.DataFrame(result.chart_rows, columns=["Prompt stem", "Cue", "Activation"])
        return _cue_context_metrics_markdown(result), table, chart, _tsv(table)
    except Exception as exc:
        _raise_ui_error(exc)


def select_candidate_row(table: pd.DataFrame, evt: gr.SelectData):
    if table is None or len(table) == 0:
        return gr.update()
    index = evt.index
    row_index = int(index[0] if isinstance(index, (tuple, list)) else index)
    if row_index < 0 or row_index >= len(table):
        return gr.update()
    value = str(int(float(table.iloc[row_index]["Feature id"])))
    return gr.update(value=value)


def use_candidate_feature(candidate_id: str):
    if candidate_id is None or str(candidate_id).strip() == "":
        raise gr.Error("Run concept-guided discovery and choose a candidate first.")
    value = str(int(float(candidate_id)))
    status = (
        f"**Feature {value} loaded.** It is now selected in Single-feature causal test, "
        "Scale dose-response, Contrastive preference, and Feature evidence. You can still change any selector independently."
    )
    return value, value, value, value, status


def mode_help(mode: str):
    if mode == "ablate":
        return gr.update(value=0.0, interactive=False, label="Coefficient (unused for ablation)")
    if mode == "scale":
        return gr.update(value=2.0, interactive=True, label="Feature multiplier")
    return gr.update(value=5.0, interactive=True, label="Additive feature coefficient")


def set_mode_help(mode: str):
    if mode == "ablate":
        return gr.update(value=0.0, interactive=False, label="Multiplier (unused for ablation)")
    return gr.update(value=2.0, interactive=True, label="Shared feature multiplier")


with gr.Blocks(title="FeatureLens — Causal Interpretability Workbench", fill_width=True) as demo:
    gr.HTML(
        '<header class="hero">'
        '<h1>FeatureLens <span style="font-size:.48em;opacity:.58;font-weight:400">v0.10</span></h1>'
        '<div class="subtitle">Causal Interpretability Workbench</div>'
        '<div class="question">Discover sparse features, test robustness, and separate correlation from causal influence.</div>'
        '</header>'
    )

    global_context = gr.Markdown(
        "**Current Workbench context:** none yet. Go to **Workbench → I. Inspect a prompt location**, "
        "set the prompt/layer/token, then click **Inspect sparse features**.",
        elem_classes=["context-card"],
    )

    with gr.Tab("Start here"):
        gr.Markdown(
            "## Start here\n"
            "FeatureLens is organized as a short evidence workflow. You do **not** need to know a feature id before "
            "starting: inspect a prompt first, or use **Feature evidence → Concept-guided candidate discovery**."
        )
        with gr.Row(equal_height=True):
            gr.HTML(
                '<div class="start-card"><h3>1. Choose a location</h3>'
                '<p>In <b>Workbench</b>, enter a prompt, choose a residual layer and token, then click '
                '<b>Inspect sparse features</b>. This establishes the shared Workbench context.</p></div>'
            )
            gr.HTML(
                '<div class="start-card"><h3>2. Choose evidence</h3>'
                '<p>Pick a feature in the experiment you want to run. Feature selectors are populated from the latest '
                'inspection but remain editable. Feature sets use 2–5 features.</p></div>'
            )
            gr.HTML(
                '<div class="start-card"><h3>3. Ask a causal question</h3>'
                '<p>Ablate/scale features, compare against norm-matched random controls, test dose response, '
                'preference shifts, robustness, geometry, or concept association.</p></div>'
            )
        gr.Markdown(
            "### Plain-language glossary\n"
            "- **Token:** a chunk of text the model processes.\n"
            "- **Residual layer:** a checkpoint in the model's internal computation.\n"
            "- **SAE feature:** one sparse direction used by the autoencoder to represent a residual activation.\n"
            "- **Ablate:** remove that feature's current contribution. **Scale:** multiply it.\n"
            "- **Random control:** an equally large residual edit in an unrelated direction.\n"
            "- **Teacher-forced score:** probability assigned to exact continuation text, even if greedy generation does not change.\n\n"
            "**Tables:** use the top-right focus control or copy-with-headers action. **Plots:** use the top-right focus control for a bounded reading view; PNG exports use descriptive FeatureLens filenames."
        )

    with gr.Tab("Workbench"):
        gr.HTML('<div class="section-rule">I. Inspect a prompt location</div>')
        with gr.Row(equal_height=False):
            with gr.Column(scale=5):
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=5,
                    value="The derivative of x squared is",
                    placeholder="Enter a prompt to inspect…",
                )
                gr.Examples(
                    examples=[
                        ["The derivative of x squared is"],
                        ["In Python, reverse a list using"],
                        ["Ich möchte einen Tisch für zwei reservieren."],
                        ["I am not fully certain, but the answer may be"],
                    ],
                    inputs=[prompt],
                    label="Controlled examples",
                )
            with gr.Column(scale=3):
                layer = gr.Dropdown(
                    choices=list(SETTINGS.layers),
                    value=SETTINGS.layers[1] if len(SETTINGS.layers) > 1 else SETTINGS.layers[0],
                    label="Residual layer",
                )
                token_index = gr.Number(
                    value=-1,
                    precision=0,
                    label="Prompt token index",
                    info="-1 = final prompt token.",
                )
                top_n = gr.Slider(5, 20, value=12, step=1, label="Displayed active features")
                analyze_btn = gr.Button("Inspect sparse features", variant="primary", elem_classes=["action-btn"])

        gr.Markdown("#### Prompt tokens\nThe selected prompt token is outlined more strongly.")
        token_view = gr.HTML(
            '<div class="small-note">Prompt tokens appear here after clicking <b>Inspect sparse features</b>.</div>'
        )
        analysis_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Strongest active SAE features')
                feature_table = gr.Dataframe(
                    headers=["Rank", "Feature id", "Activation", "Offline concept hint"],
                    datatype=["number", "number", "number", "str"],
                    interactive=False,
                    label="Strongest active SAE features", show_label=False,
                    wrap=False,
                    max_height=380,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                )
                feature_tsv = gr.Textbox(visible="hidden")
                feature_copy = _copy_button()
            with gr.Column(scale=2):
                feature_plot = gr.BarPlot(
                    x="Feature",
                    y="Activation",
                    color="Series",
                    color_map={"Activation": INK_TEAL},
                    title="Activation profile", elem_id="plot-activation-profile",
                    x_title="Feature id",
                    y_title="Activation",
                    x_label_angle=-35,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )

        gr.HTML('<div class="section-rule">II. Single-feature causal test</div>')
        gr.Markdown(
            "The edit is compared with an **8-direction norm-matched random ensemble** in the same batched "
            "execution context. If a target continuation is supplied, the complete continuation is scored teacher-forced."
        )
        with gr.Row(equal_height=False):
            with gr.Column(scale=2):
                feature_id = gr.Dropdown(
                    choices=[],
                    allow_custom_value=True,
                    label="Single feature id",
                    info="Populated from the latest inspection. You can choose another active feature or enter any valid id.",
                )
                mode = gr.Dropdown(
                    choices=["ablate", "scale", "inject"],
                    value="ablate",
                    label="Single-feature intervention",
                )
                coefficient = gr.Number(
                    value=0.0,
                    interactive=False,
                    label="Coefficient (unused for ablation)",
                )
                target_text = gr.Textbox(
                    label="Target continuation (optional)",
                    placeholder="e.g. 2x",
                    info="Exact text appended to the prompt for teacher-forced scoring.",
                )
                max_new = gr.Slider(
                    4,
                    SETTINGS.max_new_tokens,
                    value=min(12, SETTINGS.max_new_tokens),
                    step=1,
                    label="Greedy generation length",
                )
                intervene_btn = gr.Button("Run single-feature causal test", variant="primary", elem_classes=["action-btn"])
                intervention_metrics = gr.Markdown()
            with gr.Column(scale=3):
                with gr.Row():
                    baseline_out = gr.Textbox(label="Baseline greedy generation", lines=6, interactive=False)
                    modified_out = gr.Textbox(label="SAE-edited greedy generation", lines=6, interactive=False)
                _table_heading('Next-token distribution shift')
                token_prob_table = gr.Dataframe(
                    interactive=False,
                    label="Next-token distribution shift", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                token_prob_tsv = gr.Textbox(visible="hidden")
                token_prob_copy = _copy_button()
                _table_heading('Target continuation token-by-token score')
                target_token_table = gr.Dataframe(
                    interactive=False,
                    label="Target continuation token-by-token score", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                target_token_tsv = gr.Textbox(visible="hidden")
                target_token_copy = _copy_button()

        gr.HTML('<div class="section-rule">III. Single-feature scale dose-response</div>')
        with gr.Group():
            gr.Markdown("**Uses current Workbench prompt / layer / token.** Choose the feature for this experiment below.")
            with gr.Row(equal_height=True):
                dose_feature_id = gr.Dropdown(
                    choices=[],
                    allow_custom_value=True,
                    label="Dose-response feature id",
                    info="Populated from Workbench inspection, or enter any valid feature id directly.",
                    scale=2,
                )
                dose_target_text = gr.Textbox(
                    label="Dose-response target continuation",
                    value="2x",
                    info="Exact continuation scored across all dose-response conditions.",
                    scale=2,
                )
            gr.Markdown(
                "0× = ablation, 1× = numerical no-edit reference, and 2× = double the native coefficient. "
                "The six scale conditions are evaluated together against the 1× row."
            )
            dose_btn = gr.Button("Run scale dose-response", variant="primary", elem_classes=["action-btn"])
            dose_metrics = gr.Markdown()
            with gr.Row(equal_height=False):
                with gr.Column(scale=3):
                    _table_heading('Scale dose-response measurements')
                    dose_table = gr.Dataframe(
                        interactive=False,
                        label="Scale dose-response measurements", show_label=False,
                        buttons=["fullscreen"], elem_classes=["result-table"],
                        wrap=False,
                    max_height=380,
                    )
                    dose_tsv = gr.Textbox(visible="hidden")
                    dose_copy = _copy_button()
                with gr.Column(scale=2):
                    dose_plot = gr.LinePlot(
                        x="Multiplier",
                        y="Δ mean log p/token",
                        color="Series",
                        color_map={"SAE feature": INK_TEAL},
                        title="Scale dose-response", elem_id="plot-scale-dose-response",
                        x_title="Feature multiplier",
                        y_title="Δ mean log p/token",
                        buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                        height=330,
                    )

        gr.HTML('<div class="section-rule">IV. Contrastive continuation preference</div>')
        with gr.Group():
            gr.Markdown("**Uses current Workbench prompt / layer / token.** Select the feature explicitly below.")
            contrastive_feature_id = gr.Dropdown(
                choices=[],
                allow_custom_value=True,
                label="Contrastive feature id",
                info="Populated from the latest Workbench inspection; editable here.",
            )
            gr.Markdown(
                "Absolute target probability can move because an intervention broadly perturbs the distribution. "
                "This test asks the stricter question: **does the edit shift preference between two exact continuations?** "
                "Both continuations are teacher-forced and compared with the same 8-direction norm-matched random ensemble."
            )
            with gr.Row(equal_height=True):
                contrastive_a = gr.Textbox(label="Continuation A (preferred)", value="2x", scale=2)
                contrastive_b = gr.Textbox(label="Continuation B (comparison)", value="x", scale=2)
            with gr.Row(equal_height=True):
                contrastive_mode = gr.Dropdown(
                    choices=["ablate", "scale", "inject"],
                    value="ablate",
                    label="Contrastive intervention",
                    scale=1,
                )
                contrastive_coefficient = gr.Number(
                    value=0.0,
                    interactive=False,
                    label="Coefficient (unused for ablation)",
                    scale=1,
                )
            contrastive_btn = gr.Button(
                "Run contrastive preference test", variant="primary", elem_classes=["action-btn"]
            )
            contrastive_metrics = gr.Markdown()
            with gr.Row(equal_height=False):
                with gr.Column(scale=3):
                    _table_heading('Contrastive continuation scores')
                    contrastive_table = gr.Dataframe(
                        interactive=False,
                        label="Contrastive continuation scores", show_label=False,
                        buttons=["fullscreen"], elem_classes=["result-table"],
                        wrap=False,
                        max_height=320,
                    )
                    contrastive_tsv = gr.Textbox(visible="hidden")
                    contrastive_copy = _copy_button()
                with gr.Column(scale=2):
                    contrastive_plot = gr.BarPlot(
                        x="Condition",
                        y="A−B sequence log-odds",
                        color="Series",
                        color_map={"Contrastive preference": INK_TEAL},
                        title="Preference between exact continuations", elem_id="plot-contrastive-preference",
                        x_title="Execution condition",
                        y_title="Sequence log-odds A−B",
                        buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                        height=320,
                    )

    with gr.Tab("Feature sets"):
        gr.Markdown(
            "### Distributed sparse representations\n"
            "Run **Inspect sparse features** in the Workbench first. **Uses current Workbench context:** prompt, layer, "
            "and token are inherited from that inspection. The feature set below is always editable."
        )
        feature_set_location = gr.Markdown("No Workbench inspection has been run yet.")
        feature_set_ids = gr.Dropdown(
            choices=[],
            value=[],
            multiselect=True,
            allow_custom_value=True,
            max_choices=12,
            label="Feature set",
        )

        gr.Markdown("#### Joint feature-set causal test")
        gr.HTML(
            '<div class="instrument-note">Ablation ignores the multiplier. Scale applies one shared multiplier '
            'to every selected feature before decoder deltas are summed. Additive injection is deliberately omitted '
            'because one coefficient is not comparable across unrelated directions.</div>'
        )
        with gr.Row(equal_height=True):
            set_mode = gr.Dropdown(
                choices=["ablate", "scale"],
                value="ablate",
                label="Intervention",
                scale=1,
            )
            set_coefficient = gr.Number(
                value=0.0,
                interactive=False,
                label="Multiplier (unused for ablation)",
                scale=1,
            )
            set_target = gr.Textbox(
                label="Target continuation",
                value="2x",
                lines=1,
                scale=2,
            )
        set_btn = gr.Button("Run joint feature-set causal test", variant="primary", elem_classes=["action-btn"])
        set_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=2):
                _table_heading('Joint intervention features')
                set_feature_table = gr.Dataframe(
                    interactive=False,
                    label="Joint intervention features", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                set_feature_tsv = gr.Textbox(visible="hidden")
                set_feature_copy = _copy_button()
            with gr.Column(scale=3):
                _table_heading('Target continuation token-by-token score')
                set_target_table = gr.Dataframe(
                    interactive=False,
                    label="Target continuation token-by-token score", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                set_target_tsv = gr.Textbox(visible="hidden")
                set_target_copy = _copy_button()

        gr.HTML('<div class="section-rule">Set-size sensitivity</div>')
        gr.Markdown(
            "Ablates the **1, 3, and 5 strongest active features**. Each targeted edit is compared with eight "
            "norm-matched random directions, all sharing one batched no-edit baseline."
        )
        set_sweep_target = gr.Textbox(label="Target continuation for set-size sweep", value="2x")
        set_sweep_btn = gr.Button("Run 1/3/5-feature ablation sweep", variant="primary", elem_classes=["action-btn"])
        set_sweep_note = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Feature-set size measurements')
                set_sweep_table = gr.Dataframe(
                    interactive=False,
                    label="Feature-set size measurements", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                set_sweep_tsv = gr.Textbox(visible="hidden")
                set_sweep_copy = _copy_button()
            with gr.Column(scale=2):
                set_sweep_plot = gr.LinePlot(
                    x="Set size k",
                    y="Δ mean log p/token",
                    color="Condition",
                    color_map={
                        "Top-k SAE ablation": INK_TEAL,
                        "Random signed mean": INK_STONE,
                    },
                    title="Effect vs feature-set size", elem_id="plot-feature-set-size",
                    x_title="Number of jointly ablated features",
                    y_title="Δ mean log p/token",
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )

        gr.HTML('<div class="section-rule">Non-additivity / interaction decomposition</div>')
        gr.Markdown(
            "For the selected 2–5 features, FeatureLens batches each **individual ablation** plus the **joint "
            "ablation**. It compares the observed joint effect with the sum of individual effects. This diagnoses "
            "non-additivity without claiming that the features form a direct circuit."
        )
        interaction_target = gr.Textbox(label="Target continuation for interaction test", value="2x")
        interaction_btn = gr.Button("Run individual-vs-joint decomposition", variant="primary", elem_classes=["action-btn"])
        interaction_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Individual and joint ablation measurements')
                interaction_table = gr.Dataframe(
                    interactive=False,
                    label="Individual and joint ablation measurements", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                interaction_tsv = gr.Textbox(visible="hidden")
                interaction_copy = _copy_button()
            with gr.Column(scale=2):
                interaction_plot = gr.BarPlot(
                    x="Condition",
                    y="Δ mean log p/token",
                    color="Series",
                    color_map={"Ablation effect": INK_UMBER},
                    title="Individual vs joint effect", elem_id="plot-individual-vs-joint",
                    x_title="Intervention condition",
                    y_title="Δ mean log p/token",
                    x_label_angle=-25,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )


        gr.HTML('<div class="section-rule">Decoder geometry</div>')
        with gr.Group():
            gr.Markdown(
                "Joint causal effects can reflect both **decoder-direction geometry** and downstream non-linearity. "
                "This zero-extra-generation diagnostic measures pairwise decoder cosines and compares the actual "
                "activation-weighted joint-ablation norm with the norm expected for independent directions."
            )
            geometry_btn = gr.Button(
                "Inspect selected-feature geometry", variant="primary", elem_classes=["action-btn"]
            )
            geometry_metrics = gr.Markdown()
            with gr.Row(equal_height=False):
                with gr.Column(scale=3):
                    _table_heading('Pairwise decoder geometry')
                    geometry_table = gr.Dataframe(
                        interactive=False,
                        label="Pairwise decoder geometry", show_label=False,
                        buttons=["fullscreen"], elem_classes=["result-table"],
                        wrap=False,
                        max_height=340,
                    )
                    geometry_tsv = gr.Textbox(visible="hidden")
                    geometry_copy = _copy_button()
                with gr.Column(scale=2):
                    geometry_plot = gr.BarPlot(
                        x="Feature pair",
                        y="Decoder cosine",
                        color="Series",
                        color_map={"Decoder cosine": INK_PLUM},
                        title="Pairwise SAE decoder cosine", elem_id="plot-decoder-geometry",
                        x_title="Feature pair",
                        y_title="Cosine similarity",
                        x_label_angle=-30,
                        buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                        height=320,
                    )

    with gr.Tab("Feature evidence"):
        gr.Markdown(
            "### Feature activation evidence\n"
            "Inspect **where a feature fires inside the current prompt**, then test whether it activates preferentially "
            "across FeatureLens's seven controlled concept groups. The concept scan uses the feature's **maximum "
            "activation across non-padding tokens in each prompt**, avoiding arbitrary final-token comparisons."
        )
        gr.HTML('<div class="section-rule">A. Concept-guided candidate discovery</div>')
        gr.Markdown(
            "Start from a concept instead of an anonymous feature id. The default **Balanced selectivity** ranking "
            "downweights globally high-activation features and favors candidates that are both selective and consistently "
            "active for the target concept. The table also checks whether each candidate is active in the current Workbench "
            "prompt and selected token, so you can tell whether it is immediately usable for a causal test."
        )
        with gr.Row(equal_height=True):
            discovery_concept = gr.Dropdown(
                choices=[
                    "code", "mathematics", "positive_sentiment", "negative_sentiment",
                    "german_language", "factual_entities", "uncertainty"
                ],
                value="mathematics",
                label="Target concept",
            )
            discovery_layer = gr.Dropdown(choices=list(SETTINGS.layers), value=SETTINGS.layers[1], label="Residual layer")
            discovery_n = gr.Slider(2, 6, value=SETTINGS.contrast_prompts_per_concept, step=1, label="Prompts per concept")
            discovery_top_n = gr.Slider(5, 20, value=12, step=1, label="Candidate features")
        with gr.Row(equal_height=True):
            discovery_ranking = gr.Dropdown(
                choices=["Balanced selectivity", "Causal-ready at current token", "Raw mean difference"],
                value="Balanced selectivity",
                label="Candidate ranking",
                info="Balanced selectivity finds concept-associated candidates; Causal-ready restricts to features active at the selected Workbench token; raw mean difference exposes scale-dominated ranking.",
                scale=2,
            )
            gr.Markdown(
                "**Input source:** controlled concept prompts + the current Workbench prompt/token, evaluated in the same GPU batch.",
                elem_classes=["instrument-note"],
            )
        discovery_btn = gr.Button("Discover concept-associated candidates", variant="primary", elem_classes=["action-btn"])
        discovery_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Candidate feature evidence')
                discovery_table = gr.Dataframe(
                    interactive=False, label="Candidate feature evidence", show_label=False, buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False, max_height=420
                )
                discovery_tsv = gr.Textbox(visible="hidden")
                discovery_copy = _copy_button()
            with gr.Column(scale=2):
                discovery_plot = gr.BarPlot(
                    x="Feature", y="Candidate score", color="Series",
                    color_map={"Candidate score": INK_TEAL}, title="Candidate evidence score", elem_id="plot-candidate-discovery",
                    x_title="Feature id", y_title="Exploratory ranking score", x_label_angle=-35,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"], height=330
                )
        gr.Markdown(
            "Click a row in **Candidate feature evidence** to place that feature in the selector below. "
            "The selector defaults to the highest-ranked candidate active at the current Workbench token when available.",
            elem_classes=["candidate-help"],
        )
        with gr.Row(equal_height=True):
            discovery_candidate = gr.Dropdown(choices=[], label="Selected candidate feature id", allow_custom_value=True, scale=3)
            use_candidate_btn = gr.Button(
                "Use selected candidate across feature tests", variant="primary", elem_classes=["action-btn"], scale=2
            )
        candidate_use_status = gr.Markdown()

        gr.HTML('<div class="section-rule">B. Batched causal candidate triage</div>')
        gr.Markdown(
            "Use discovery candidates as a short list, then ablate several of them in **one batched scoring run**. "
            "This ranks target effects cheaply before you spend a full 8-random-control causal test on a feature. "
            "Uses the current Workbench prompt/token and the discovery layer."
        )
        with gr.Row(equal_height=True):
            candidate_screen_ids = gr.Dropdown(
                choices=[],
                value=[],
                multiselect=True,
                allow_custom_value=True,
                max_choices=8,
                label="Candidate features to screen",
                info="Populated by concept-guided discovery; up to eight features per batch.",
                scale=3,
            )
            candidate_screen_target = gr.Textbox(
                label="Screen target continuation",
                value="2x",
                info="Exact continuation used only for this screening run.",
                scale=2,
            )
        candidate_screen_btn = gr.Button(
            "Screen candidate ablations", variant="primary", elem_classes=["action-btn"]
        )
        candidate_screen_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Candidate ablation screen')
                candidate_screen_table = gr.Dataframe(
                    interactive=False,
                    label="Candidate ablation screen",
                    show_label=False,
                    buttons=["fullscreen"],
                    elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                candidate_screen_tsv = gr.Textbox(visible="hidden")
                candidate_screen_copy = _copy_button()
            with gr.Column(scale=2):
                candidate_screen_plot = gr.BarPlot(
                    x="Feature",
                    y="Δ mean log p/token",
                    color="Series",
                    color_map={"Candidate ablation": INK_TEAL},
                    title="Candidate ablation target effect",
                    elem_id="plot-candidate-causal-screen",
                    x_title="Feature id",
                    y_title="Δ mean log p/token",
                    x_label_angle=-35,
                    buttons=["fullscreen", "export"],
                    elem_classes=["fl-plot"],
                    height=330,
                )
        gr.Markdown(
            "Click a row in **Candidate ablation screen** to move that feature into **Selected candidate feature id**, "
            "then use the existing handoff button above if you want it copied into the deeper feature tests.",
            elem_classes=["candidate-help"],
        )

        gr.Markdown("#### Association vs causal influence")
        gr.Markdown(
            "Computed from the discovery table and the ablation screen above — **no additional GPU call**. "
            "It compares concept-evidence rank with target-effect and distribution-shift ranks for the same shortlist.",
            elem_classes=["candidate-help"],
        )
        candidate_alignment_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Discovery–causality alignment')
                candidate_alignment_table = gr.Dataframe(
                    interactive=False,
                    label="Discovery–causality alignment",
                    show_label=False,
                    buttons=["fullscreen"],
                    elem_classes=["result-table"],
                    wrap=False,
                    max_height=360,
                )
                candidate_alignment_tsv = gr.Textbox(visible="hidden")
                candidate_alignment_copy = _copy_button()
            with gr.Column(scale=2):
                candidate_alignment_plot = gr.ScatterPlot(
                    x="Candidate score",
                    y="|Δ mean log p/token|",
                    color="Series",
                    color_map={"Screened candidate": INK_TEAL},
                    title="Association evidence vs target effect",
                    elem_id="plot-association-causality",
                    x_title="Discovery candidate score",
                    y_title="|Δ mean log p/token|",
                    tooltip=[
                        "Feature id",
                        "Discovery rank",
                        "Target-effect rank",
                        "Next-token JS",
                    ],
                    buttons=["fullscreen", "export"],
                    elem_classes=["fl-plot"],
                    height=330,
                )

        gr.HTML('<div class="section-rule">C. Inspect one feature</div>')
        contrast_location = gr.Markdown(
            "**Activation trace uses the current Workbench prompt.** The controlled concept scan below uses its own balanced prompt set."
        )
        with gr.Row(equal_height=True):
            contrast_feature_id = gr.Dropdown(
                choices=[],
                allow_custom_value=True,
                label="Feature id",
                scale=2,
            )
            contrast_layer = gr.Dropdown(
                choices=list(SETTINGS.layers),
                value=SETTINGS.layers[1],
                label="Residual layer",
                scale=1,
            )
            contrast_n = gr.Slider(
                2,
                6,
                value=SETTINGS.contrast_prompts_per_concept,
                step=1,
                label="Prompts per concept",
                scale=2,
            )

        gr.Markdown("### Activation trace across the current Workbench prompt")
        gr.Markdown(
            "Uses the prompt currently entered in **Workbench** and shows exactly which prompt tokens include the "
            "selected feature in the SAE TopK support."
        )
        trace_btn = gr.Button("Trace feature across prompt tokens", variant="primary", elem_classes=["action-btn"])
        trace_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Feature activation by prompt token')
                trace_table = gr.Dataframe(
                    interactive=False,
                    label="Feature activation by prompt token", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=340,
                )
                trace_tsv = gr.Textbox(visible="hidden")
                trace_copy = _copy_button()
            with gr.Column(scale=2):
                trace_plot = gr.BarPlot(
                    x="Token",
                    y="Activation",
                    color="Series",
                    color_map={"Feature activation": INK_TEAL},
                    title="Feature activation across prompt tokens", elem_id="plot-feature-token-trace",
                    x_title="Prompt token",
                    y_title="Activation",
                    x_label_angle=-35,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=320,
                )

        gr.HTML('<div class="section-rule">D. Completion-cue sensitivity</div>')
        gr.Markdown(
            "Feature 22632 in the current math example fires only on the final `is` token. This scan tests whether a "
            "selected feature responds to **completion cues themselves** rather than the underlying concept. Enter a stem "
            "and one cue per line; FeatureLens appends each cue and measures the feature at the final token."
        )
        with gr.Row(equal_height=True):
            cue_stem = gr.Textbox(label="Prompt stem", value="The derivative of x squared", lines=2, scale=3)
            cue_text = gr.Textbox(label="Completion cues (one per line)", value="is\n=\n:\nequals\ntherefore", lines=5, scale=2)
        cue_btn = gr.Button("Run completion-cue scan", variant="primary", elem_classes=["action-btn"])
        cue_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Feature response by completion cue')
                cue_table = gr.Dataframe(interactive=False, label="Feature response by completion cue", show_label=False, buttons=["fullscreen"], elem_classes=["result-table"], wrap=False, max_height=340)
                cue_tsv = gr.Textbox(visible="hidden")
                cue_copy = _copy_button()
            with gr.Column(scale=2):
                cue_plot = gr.BarPlot(
                    x="Cue", y="Activation", color="Series", color_map={"Cue response": INK_UMBER},
                    title="Completion-cue feature response", elem_id="plot-completion-cue-response", x_title="Cue", y_title="Final-token activation",
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"], height=320
                )

        gr.HTML('<div class="section-rule">E. Cue × context specificity</div>')
        gr.Markdown(
            "The single-stem cue scan can tell you that a feature prefers `is`, but not whether that preference is lexical "
            "or context-sensitive. This matrix crosses several prompt stems with the same completion cues in one batch."
        )
        with gr.Row(equal_height=True):
            cue_context_stems = gr.Textbox(
                label="Prompt stems (one per line)",
                value="The derivative of x squared\nThe capital of Germany\nThe weather today\nMy name",
                lines=5,
                scale=3,
            )
            cue_context_cues = gr.Textbox(
                label="Completion cues (one per line)",
                value="is\n=\n:\nequals\ntherefore",
                lines=5,
                scale=2,
            )
        cue_context_btn = gr.Button("Run cue × context scan", variant="primary", elem_classes=["action-btn"])
        cue_context_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Cue × context feature response')
                cue_context_table = gr.Dataframe(
                    interactive=False,
                    label="Cue × context feature response", show_label=False,
                    buttons=["fullscreen"],
                    elem_classes=["result-table"],
                    wrap=False,
                    max_height=420,
                )
                cue_context_tsv = gr.Textbox(visible="hidden")
                cue_context_copy = _copy_button()
            with gr.Column(scale=2):
                cue_context_plot = gr.BarPlot(
                    x="Prompt stem",
                    y="Activation",
                    color="Cue",
                    color_map={
                        "is": INK_TEAL,
                        "=": INK_UMBER,
                        ":": INK_RED,
                        "equals": INK_PLUM,
                        "therefore": INK_STONE,
                    },
                    title="Cue response across contexts",
                    elem_id="plot-cue-context-matrix",
                    x_title="Prompt stem",
                    y_title="Final-token activation",
                    x_label_angle=-25,
                    buttons=["fullscreen", "export"],
                    elem_classes=["fl-plot"],
                    height=340,
                )

        gr.HTML('<div class="section-rule">F. Controlled concept contrast for selected feature</div>')
        gr.Markdown("### Prompt-wide controlled concept contrast")
        contrast_btn = gr.Button("Run controlled concept contrast", variant="primary", elem_classes=["action-btn"])
        contrast_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Feature activation by controlled concept')
                contrast_table = gr.Dataframe(
                    interactive=False,
                    label="Feature activation by controlled concept", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                contrast_tsv = gr.Textbox(visible="hidden")
                contrast_copy = _copy_button()
            with gr.Column(scale=2):
                contrast_plot = gr.BarPlot(
                    x="Concept",
                    y="Mean prompt-wide max",
                    color="Series",
                    color_map={"Prompt-wide max": INK_BLUEGREY},
                    title="Prompt-wide controlled concept contrast", elem_id="plot-controlled-concept-contrast",
                    x_title="Concept",
                    y_title="Mean max activation",
                    x_label_angle=-25,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )

    with gr.Tab("Paraphrase robustness"):
        gr.Markdown(
            "### Local and prompt-wide robustness\n"
            "The selected-token comparison is strict but can be misleading if the two chosen tokens play different "
            "semantic roles. FeatureLens reports a **prompt-wide max-pooled feature profile**: for every SAE "
            "feature, take its maximum activation across all prompt tokens, then compare the two sparse profiles."
        )
        with gr.Row():
            para_a = gr.Textbox(
                label="Original prompt",
                lines=4,
                value="The derivative of x squared is",
            )
            para_b = gr.Textbox(
                label="Paraphrase",
                lines=4,
                value="Differentiate x squared with respect to x:",
            )
        with gr.Row():
            para_layer = gr.Dropdown(
                choices=list(SETTINGS.layers),
                value=SETTINGS.layers[1],
                label="Residual layer",
            )
            para_idx_a = gr.Number(value=-1, precision=0, label="Original prompt token index")
            para_idx_b = gr.Number(value=-1, precision=0, label="Paraphrase token index")
            para_top_n = gr.Slider(5, 20, value=12, step=1, label="Displayed active features")
        para_btn = gr.Button("Compare paraphrase representations", variant="primary", elem_classes=["action-btn"])
        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Original prompt tokens")
                para_tokens_a = gr.HTML()
            with gr.Column():
                gr.Markdown("#### Paraphrase tokens")
                para_tokens_b = gr.HTML()
        para_metrics = gr.Markdown()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Top-feature overlap at selected tokens')
                para_table = gr.Dataframe(
                    interactive=False,
                    label="Top-feature overlap at selected tokens", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                para_tsv = gr.Textbox(visible="hidden")
                para_copy = _copy_button()
            with gr.Column(scale=2):
                para_plot = gr.BarPlot(
                    x="Feature",
                    y="Activation",
                    color="Prompt",
                    color_map={"Original": INK_TEAL, "Paraphrase": INK_PLUM},
                    title="Selected-token feature activations", elem_id="plot-paraphrase-selected-token",
                    x_title="Feature id",
                    y_title="Activation",
                    x_label_angle=-35,
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )

    with gr.Tab("Layer trajectory"):
        gr.Markdown(
            "### Representation structure across early, middle and late residual streams\n"
            "This is not a cross-layer feature-ID comparison; the SAE dictionaries are layer-specific."
        )
        with gr.Row():
            trajectory_prompt = gr.Textbox(
                label="Prompt",
                lines=5,
                value="The derivative of x squared is",
                scale=4,
            )
            trajectory_token = gr.Number(
                value=-1,
                precision=0,
                label="Prompt token index",
                info="-1 = final token",
                scale=1,
            )
        trajectory_btn = gr.Button("Compare layers", variant="primary", elem_classes=["action-btn"])
        gr.Markdown("#### Prompt tokens")
        trajectory_tokens = gr.HTML()
        with gr.Row(equal_height=False):
            with gr.Column(scale=3):
                _table_heading('Layer diagnostics')
                trajectory_table = gr.Dataframe(
                    interactive=False,
                    label="Layer diagnostics", show_label=False,
                    buttons=["fullscreen"], elem_classes=["result-table"],
                    wrap=False,
                    max_height=380,
                )
                trajectory_tsv = gr.Textbox(visible="hidden")
                trajectory_copy = _copy_button()
            with gr.Column(scale=2):
                trajectory_plot = gr.LinePlot(
                    x="Layer",
                    y="Value",
                    color="Metric",
                    color_map={
                        "Reconstruction cosine": INK_TEAL,
                        "Top-5 mass": INK_UMBER,
                        "Activation entropy": INK_RED,
                    },
                    title="Representation trajectory", elem_id="plot-layer-trajectory",
                    x_title="Layer",
                    y_title="Normalized value",
                    buttons=["fullscreen", "export"], elem_classes=["fl-plot"],
                    height=330,
                )

    with gr.Tab("Offline benchmark"):
        gr.Markdown(RUNTIME.catalog.benchmark_markdown())
        gr.Markdown(
            "The offline pipeline remains the source of held-out feature/concept AUROC + F1, reconstruction "
            "quality, paraphrase stability, dense residual linear probes, and causal intervention conclusions. "
            "The live workbench uses batch-context null references and random-control ensembles; "
            "the saved report should be generated only from real experiment artifacts."
        )

    with gr.Tab("Method"):
        gr.Markdown(
            r"""
### Reconstruction-preserving intervention

For residual vector $h$, sparse coefficient $z_i$, decoder direction $d_i$, and scale $\alpha$:

- **Ablate:** $h' = h - z_i d_i$
- **Scale:** $h' = h + (\alpha - 1)z_i d_i$
- **Inject:** $h' = h + \delta d_i$

For a feature set $S$:

$$h' = h + \sum_{i \in S}\Delta z_i d_i.$$

FeatureLens patches the delta into the **original residual**; it never replaces the residual with the complete SAE reconstruction.

### Control discipline

Batched experiments include an explicit **zero-edit row**. Causal effects are measured against that row rather than a separately executed baseline, which removes batch-vs-single floating-point drift from the measured effect. Random specificity uses an ensemble of norm-matched residual directions rather than one arbitrary seed.

### Evidence ladder

1. SAE reconstruction quality.
2. Held-out feature/concept prediction.
3. Concept-guided candidate discovery and batched causal candidate triage.
4. Selected-feature concept contrast and token-local activation traces, completion-cue sensitivity, and cue × context specificity.
5. Local and prompt-wide paraphrase robustness.
6. Single-feature causal intervention and dose-response.
7. Contrastive continuation preference under intervention.
8. Joint feature-set intervention and set-size sensitivity.
9. Decoder-direction geometry and individual-vs-joint non-additivity.
10. Specificity relative to norm-matched random controls.

Association, robustness, geometry, and intervention evidence remain separate claims.
"""
        )

    gr.HTML('<div class="bottom-spacer" aria-hidden="true"></div>')

    demo.load(fn=None, js=INSTALL_REFLOW_JS, queue=False)

    # Event wiring.
    analyze_btn.click(
        analyze_prompt,
        inputs=[prompt, layer, token_index, top_n],
        outputs=[
            token_view,
            feature_table,
            feature_plot,
            feature_id,
            dose_feature_id,
            contrastive_feature_id,
            feature_set_ids,
            contrast_feature_id,
            contrast_layer,
            analysis_metrics,
            feature_set_location,
            contrast_location,
            global_context,
            feature_tsv,
        ],
    )
    mode.change(mode_help, inputs=[mode], outputs=[coefficient])
    intervene_btn.click(
        run_intervention,
        inputs=[prompt, layer, token_index, feature_id, mode, coefficient, target_text, max_new],
        outputs=[
            baseline_out,
            modified_out,
            intervention_metrics,
            token_prob_table,
            target_token_table,
            token_prob_tsv,
            target_token_tsv,
        ],
    )
    dose_btn.click(
        run_dose_response,
        inputs=[prompt, layer, token_index, dose_feature_id, dose_target_text],
        outputs=[dose_table, dose_plot, dose_metrics, dose_tsv],
    )
    contrastive_mode.change(mode_help, inputs=[contrastive_mode], outputs=[contrastive_coefficient])
    contrastive_btn.click(
        run_contrastive_causal,
        inputs=[prompt, layer, token_index, contrastive_feature_id, contrastive_mode, contrastive_coefficient, contrastive_a, contrastive_b],
        outputs=[contrastive_metrics, contrastive_table, contrastive_plot, contrastive_tsv],
    )
    set_mode.change(set_mode_help, inputs=[set_mode], outputs=[set_coefficient])
    set_btn.click(
        run_feature_set,
        inputs=[prompt, layer, token_index, feature_set_ids, set_mode, set_coefficient, set_target],
        outputs=[set_feature_table, set_metrics, set_target_table, set_feature_tsv, set_target_tsv],
    )
    set_sweep_btn.click(
        run_feature_set_sweep,
        inputs=[prompt, layer, token_index, set_sweep_target],
        outputs=[set_sweep_table, set_sweep_plot, set_sweep_note, set_sweep_tsv],
    )
    interaction_btn.click(
        run_feature_interaction,
        inputs=[prompt, layer, token_index, feature_set_ids, interaction_target],
        outputs=[interaction_table, interaction_metrics, interaction_plot, interaction_tsv],
    )
    geometry_btn.click(
        run_feature_geometry,
        inputs=[prompt, layer, token_index, feature_set_ids],
        outputs=[geometry_metrics, geometry_table, geometry_plot, geometry_tsv],
    )
    trace_btn.click(
        run_feature_trace,
        inputs=[prompt, contrast_layer, contrast_feature_id],
        outputs=[trace_metrics, trace_table, trace_plot, trace_tsv],
    )
    contrast_btn.click(
        run_concept_contrast,
        inputs=[contrast_feature_id, contrast_layer, contrast_n],
        outputs=[contrast_metrics, contrast_table, contrast_plot, contrast_tsv],
    )
    discovery_btn.click(
        run_concept_feature_discovery,
        inputs=[
            discovery_concept, discovery_layer, discovery_n, discovery_top_n, discovery_ranking,
            prompt, token_index,
        ],
        outputs=[
            discovery_metrics,
            discovery_table,
            discovery_plot,
            discovery_candidate,
            candidate_screen_ids,
            discovery_tsv,
        ],
    )
    candidate_screen_btn.click(
        run_candidate_causal_screen,
        inputs=[
            prompt,
            discovery_layer,
            token_index,
            candidate_screen_ids,
            candidate_screen_target,
            discovery_table,
        ],
        outputs=[
            candidate_screen_metrics,
            candidate_screen_table,
            candidate_screen_plot,
            discovery_candidate,
            candidate_screen_tsv,
            candidate_alignment_metrics,
            candidate_alignment_table,
            candidate_alignment_plot,
            candidate_alignment_tsv,
        ],
    )
    candidate_screen_table.select(
        select_candidate_row,
        inputs=[candidate_screen_table],
        outputs=[discovery_candidate],
        queue=False,
    )
    discovery_table.select(
        select_candidate_row,
        inputs=[discovery_table],
        outputs=[discovery_candidate],
        queue=False,
    )
    use_candidate_btn.click(
        use_candidate_feature,
        inputs=[discovery_candidate],
        outputs=[feature_id, dose_feature_id, contrastive_feature_id, contrast_feature_id, candidate_use_status],
        queue=False,
    )
    cue_btn.click(
        run_feature_cue_scan,
        inputs=[contrast_feature_id, contrast_layer, cue_stem, cue_text],
        outputs=[cue_metrics, cue_table, cue_plot, cue_tsv],
    )
    cue_context_btn.click(
        run_feature_cue_context_scan,
        inputs=[contrast_feature_id, contrast_layer, cue_context_stems, cue_context_cues],
        outputs=[cue_context_metrics, cue_context_table, cue_context_plot, cue_context_tsv],
    )
    para_btn.click(
        run_paraphrase_compare,
        inputs=[para_a, para_b, para_layer, para_idx_a, para_idx_b, para_top_n],
        outputs=[para_tokens_a, para_tokens_b, para_metrics, para_table, para_plot, para_tsv],
    )
    trajectory_btn.click(
        run_layer_sweep,
        inputs=[trajectory_prompt, trajectory_token],
        outputs=[trajectory_tokens, trajectory_table, trajectory_plot, trajectory_tsv],
    )

    for button, source in [
        (feature_copy, feature_tsv),
        (token_prob_copy, token_prob_tsv),
        (target_token_copy, target_token_tsv),
        (dose_copy, dose_tsv),
        (contrastive_copy, contrastive_tsv),
        (set_feature_copy, set_feature_tsv),
        (set_target_copy, set_target_tsv),
        (set_sweep_copy, set_sweep_tsv),
        (interaction_copy, interaction_tsv),
        (geometry_copy, geometry_tsv),
        (trace_copy, trace_tsv),
        (contrast_copy, contrast_tsv),
        (discovery_copy, discovery_tsv),
        (candidate_screen_copy, candidate_screen_tsv),
        (candidate_alignment_copy, candidate_alignment_tsv),
        (cue_copy, cue_tsv),
        (cue_context_copy, cue_context_tsv),
        (para_copy, para_tsv),
        (trajectory_copy, trajectory_tsv),
    ]:
        _bind_copy(button, source)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        css=CSS,
        theme=THEME,
        ssr_mode=False,
        show_error=True,
    )