"""SamplingTAR: Training-free Concept Localization for Typographic Attack Defense.

Interactive demo of the training-free defense from
"Towards Robustness against Typographic Attack with Training-free Concept Localization".

The method mines attention heads responsible for reading text in images using
randomly-initialised Sparse Autoencoders (SAEs), then ablates those heads at
inference to recover the true object prediction.

We use the HuggingFace transformers CLIP model (openai/clip-vit-large-patch14)
for a stable, well-defined attention interface.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import

import math
import random
from contextlib import contextmanager
from functools import partial

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.font_manager
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms

import gradio as gr
from transformers import CLIPModel, CLIPProcessor

# ---------------------------------------------------------------------------
# Model loading at module scope (ZeroGPU rule: load eagerly, .to("cuda"))
# ---------------------------------------------------------------------------

MODEL_ID = "openai/clip-vit-large-patch14"

print(f"Loading CLIP model from {MODEL_ID}…")
_clip_model = CLIPModel.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
_processor = CLIPProcessor.from_pretrained(MODEL_ID)

_vision = _clip_model.vision_model
_vision.eval()
_vision = _vision.to("cuda")

# The projection to the shared embedding space
_visual_projection = _clip_model.visual_projection
_visual_projection.eval()
_visual_projection = _visual_projection.to("cuda")

# Text encoder
_text_model = _clip_model.text_model
_text_projection = _clip_model.text_projection
_text_model.eval()
_text_projection.eval()
_text_model = _text_model.to("cuda")
_text_projection = _text_projection.to("cuda")

# Architecture constants
PATCH_SIZE = _vision.config.patch_size  # 14
IMAGE_SIZE = _vision.config.image_size  # 224
NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2  # 256
NUM_HEADS = _vision.config.num_attention_heads  # 16
HIDDEN_DIM = _vision.config.hidden_size  # 1024
HEAD_DIM = HIDDEN_DIM // NUM_HEADS  # 64
NUM_LAYERS = _vision.config.num_hidden_layers  # 24
# The paper uses the top ~20% of layers for circuit mining.
MINE_LAYERS = [x for x in range(NUM_LAYERS) if x >= NUM_LAYERS * 0.8]
print(f"CLIP loaded.  image_size={IMAGE_SIZE}  patch={PATCH_SIZE}  "
      f"patches={NUM_PATCHES}  heads={NUM_HEADS}  layers={NUM_LAYERS}  "
      f"mine_layers={MINE_LAYERS}")

# CLIP normalisation constants
CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD = [0.26862954, 0.26130258, 0.27577711]

# ---------------------------------------------------------------------------
# Minimal SAE (faithful to model_training/sae.py — only what circuit mining needs)
# ---------------------------------------------------------------------------

class TopKSAE(nn.Module):
    """Minimal TopK SAE — only the pieces circuit mining touches."""
    def __init__(self, n_heads, head_dim, hidden_dim, k, device="cuda"):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.hidden_dim = hidden_dim
        self.k = k
        self.per_head_recon = False
        self.encoders = nn.ModuleList()
        for _ in range(n_heads):
            enc = nn.ModuleList([
                nn.Linear(head_dim, hidden_dim // n_heads, bias=True),
            ])
            self.encoders.append(enc)
        self.to(device)


def build_random_sae_list(n_heads, head_dim, hidden_dim, k, layers, device):
    """Build per-layer randomly-initialised SAEs (as the repo does)."""
    sae_list = {}
    for layer in layers:
        sae = TopKSAE(n_heads, head_dim, hidden_dim, k, device=device)
        for i, encoder in enumerate(sae.encoders):
            encoder[0].weight.data = torch.randn_like(encoder[0].weight.data).to(device)
            encoder[0].bias.data = torch.randn_like(encoder[0].bias.data).to(device)
            encoder[0].weight.data = F.normalize(encoder[0].weight.data, dim=1, p=2)
        sae.per_head_recon = False
        sae_list[layer] = sae
    return sae_list


# ---------------------------------------------------------------------------
# Preprocessing
# ---------------------------------------------------------------------------

_preprocess = transforms.Compose([
    transforms.Resize(IMAGE_SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.CenterCrop(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
])


# ---------------------------------------------------------------------------
# Core attribution functions (ported from eval_utils.py, adapted for transformers)
# ---------------------------------------------------------------------------

def get_attention_scores(vision_model, input_tensor):
    """Extract Q, K, V and attention scores from all transformer layers.

    For transformers CLIPVisionModel, each layer's attention module is
    CLIPAttention with q_proj, k_proj, v_proj, out_proj.
    We hook on the encoder layer (not self_attn) because CLIPAttention.forward()
    takes only kwargs, so forward hooks on it receive empty input tuples.
    """
    qs, ks, vs, xs = {}, {}, {}, {}

    def get_qkv_hook(i, module, input, output):
        # module is CLIPEncoderLayer, input[0] is hidden_states
        hidden = input[0]
        xs[i] = hidden
        B, L, C = hidden.shape
        # Apply layer norm before attention (as the layer does internally)
        x = module.layer_norm1(hidden)
        attn = module.self_attn
        q = attn.q_proj(x)
        k = attn.k_proj(x)
        v = attn.v_proj(x)
        q = q.view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)
        k = k.view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)
        v = v.view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)
        qs[i], ks[i], vs[i] = q, k, v
        return output

    hooks = []
    for i, layer in enumerate(vision_model.encoder.layers):
        hook = layer.register_forward_hook(partial(get_qkv_hook, i))
        hooks.append(hook)
    try:
        with torch.no_grad():
            _ = vision_model(pixel_values=input_tensor)
    finally:
        for hook in hooks:
            hook.remove()

    attention_scores = {}
    for i in qs:
        d_k = qs[i].size(-1)
        scores = torch.matmul(qs[i], ks[i].transpose(-2, -1)) / np.sqrt(d_k)
        attention_scores[i] = scores
    return xs, qs, ks, vs, attention_scores


def get_attribution_map(pre_softmax_attention, vs, sae_w, sae_b, layer, head):
    """Compute the attribution map for one (layer, head) — from the repo."""
    cls_token_index = 0
    original_scores_head = pre_softmax_attention[layer][:, head]
    values_head = vs[layer][:, head]
    sae_neuron_weight = sae_w.T
    cls_scores = original_scores_head[:, cls_token_index, :]
    cls_attention_softmax = F.softmax(cls_scores, dim=-1)
    value_contributions = torch.matmul(values_head, sae_neuron_weight)

    avg_value_contribution = torch.sum(
        cls_attention_softmax.unsqueeze(-1) * value_contributions, dim=1, keepdim=True
    )
    analytical_gradient = cls_attention_softmax.unsqueeze(-1) * (
        value_contributions - avg_value_contribution
    )
    return cls_attention_softmax, analytical_gradient


def create_masks(token_shape, text_depth=2):
    """Binary masks for the four text-border regions (top/bottom/left/right)."""
    d = int(math.floor(math.sqrt(token_shape)))
    masks = []
    for i in range(4):
        mask = torch.zeros(d, d)
        if i == 0:      # Top
            mask[:text_depth, :] = 1
        elif i == 1:    # Bottom
            mask[-text_depth:, :] = 1
        elif i == 2:    # Left
            mask[:, :text_depth] = 1
        elif i == 3:    # Right
            mask[:, -text_depth:] = 1
        masks.append(mask.flatten())
    return torch.stack(masks)


def calculate_score(text_loc, score_maps, all_masks):
    """Score how strongly a head localises to the text border vs. the object."""
    masks = all_masks[text_loc]  # (token_shape,)
    score_maps = score_maps[:, 1:]  # remove cls token -> (B, T, C)
    score_maps = torch.clamp(score_maps, min=0.0, max=1.0)
    score_maps = score_maps.permute(0, 2, 1)  # (B, C, T)
    pos_scores = (score_maps * masks).sum(-1)  # (B, C, T) * (T,) -> (B, C)
    neg_scores = (score_maps * (1 - masks)).sum(-1)
    final_scores = pos_scores / (pos_scores + neg_scores + 1e-6)
    return final_scores  # (B, C)


# ---------------------------------------------------------------------------
# Head ablation context manager (adapted for transformers CLIPAttention)
# ---------------------------------------------------------------------------

@contextmanager
def fix_attn_head_list(vision_model, layer_spec, alpha=1.0):
    """Ablate the CLS->text attention in the specified heads (training-free defense).

    Registers a forward hook on each targeted CLIPEncoderLayer that recomputes
    the attention output with the specified heads' CLS-token attention to
    non-CLS patches neutralised, then passes through the MLP normally.
    """
    hooks = []

    def hook_fn(module, input, output, layer_idx, heads):
        # module is CLIPEncoderLayer, input[0] is hidden_states
        hidden = input[0]
        B, L, C = hidden.shape
        attn = module.self_attn
        # Apply layer norm before attention (as the layer does internally)
        x = module.layer_norm1(hidden)
        q = attn.q_proj(x).view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)
        k = attn.k_proj(x).view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)
        v = attn.v_proj(x).view(B, L, attn.num_heads, attn.head_dim).transpose(1, 2)

        att = q @ k.transpose(-2, -1) / math.sqrt(attn.head_dim)
        att = att.softmax(dim=-1)

        # Redistribute CLS-token attention: neutralise text-patch contribution.
        factors = att[:, :, :1, 1:].sum(dim=-1, keepdim=True)
        for head in heads:
            att[:, head, :1, 0] = alpha
            att[:, head, :1, 1:] = att[:, head, :1, 1:] * (1 - alpha) / (
                factors[:, head, :1, :] + 1e-6
            )

        # Recompute attention output
        v_out = att @ v  # (B, heads, L, head_dim)
        ctx = v_out.transpose(1, 2).reshape(B, L, C)
        attn_out = attn.out_proj(ctx)

        # Residual + MLP (as the original layer would do)
        hidden_out = hidden + attn_out
        hidden_out = hidden_out + module.mlp(module.layer_norm2(hidden_out))
        return hidden_out

    for layer in layer_spec:
        heads = layer_spec[layer]
        hook = vision_model.encoder.layers[layer].register_forward_hook(
            partial(hook_fn, layer_idx=layer, heads=heads)
        )
        hooks.append(hook)
    try:
        yield
    finally:
        for hook in hooks:
            hook.remove()


# ---------------------------------------------------------------------------
# Text border transform (for circuit mining)
# ---------------------------------------------------------------------------

def draw_text_border(img, text, edge_idx=0, border_ratio=0.2, target_size=(224, 224)):
    """Draw a text strip along one edge of the image."""
    target_w, target_h = target_size
    border_px_h = int(target_h * border_ratio)
    border_px_w = int(target_w * border_ratio)

    canvas = Image.new("RGB", (target_w, target_h), (255, 255, 255))
    color = (random.randint(0, 150), random.randint(0, 150), random.randint(0, 150))

    if edge_idx == 0:  # Top
        img_resized = img.resize((target_w, target_h - border_px_h), Image.Resampling.LANCZOS)
        canvas.paste(img_resized, (0, border_px_h))
    elif edge_idx == 1:  # Bottom
        img_resized = img.resize((target_w, target_h - border_px_h), Image.Resampling.LANCZOS)
        canvas.paste(img_resized, (0, 0))
    elif edge_idx == 2:  # Left
        img_resized = img.resize((target_w - border_px_w, target_h), Image.Resampling.LANCZOS)
        canvas.paste(img_resized, (border_px_w, 0))
    elif edge_idx == 3:  # Right
        img_resized = img.resize((target_w - border_px_w, target_h), Image.Resampling.LANCZOS)
        canvas.paste(img_resized, (0, 0))

    draw = ImageDraw.Draw(canvas)
    font_paths = matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext="ttf")
    font = ImageFont.truetype(random.choice(font_paths), max(8, int(border_px_h * 0.7))) if font_paths else ImageFont.load_default()

    text_repeated = (text + "  ") * 5
    if edge_idx in (0, 1):
        draw.text((10, border_px_h // 4 if edge_idx == 0 else target_h - border_px_h + border_px_h // 4),
                  text_repeated, fill=color, font=font)
    else:
        strip = Image.new("RGBA", (border_px_w, target_h), (0, 0, 0, 0))
        strip_draw = ImageDraw.Draw(strip)
        strip_draw.text((border_px_w // 4, 10), text_repeated, fill=color, font=font)
        strip = strip.rotate(90, expand=True, resample=Image.BICUBIC)
        if edge_idx == 2:
            canvas.paste(strip, (0, 0), strip)
        else:
            canvas.paste(strip, (target_w - border_px_w, 0), strip)

    return canvas, edge_idx


# ---------------------------------------------------------------------------
# Circuit mining
# ---------------------------------------------------------------------------

def mine_head_scores(pil_img, class_labels, n_samples=8):
    """Mine head scores for a single image by creating text-bordered variants."""
    head_latent_dim = 64  # hidden_dim // n_heads for the SAE
    sae_hidden_dim = head_latent_dim * NUM_HEADS

    sae_list = build_random_sae_list(
        n_heads=NUM_HEADS,
        head_dim=HEAD_DIM,
        hidden_dim=sae_hidden_dim,
        k=128,
        layers=MINE_LAYERS,
        device="cuda",
    )

    all_masks = create_masks(NUM_PATCHES, text_depth=int(np.ceil(IMAGE_SIZE * 0.2 / PATCH_SIZE))).to("cuda")

    head_scores = {}
    for layer in MINE_LAYERS:
        for head in range(NUM_HEADS):
            head_scores[(layer, head)] = torch.zeros(head_latent_dim).to("cuda")

    for s in range(n_samples):
        edge = s % 4
        label = random.choice(class_labels) if class_labels else "object"
        bordered_img, text_loc = draw_text_border(
            pil_img, label, edge_idx=edge, border_ratio=0.2,
            target_size=(IMAGE_SIZE, IMAGE_SIZE)
        )
        input_tensor = _preprocess(bordered_img).unsqueeze(0).to("cuda")

        with torch.inference_mode():
            _, _, _, vs, attention_scores = get_attention_scores(_vision, input_tensor)

        for layer in MINE_LAYERS:
            sae = sae_list[layer]
            for head in range(NUM_HEADS):
                sae_w = sae.encoders[head][0].weight.data
                sae_b = sae.encoders[head][0].bias.data
                cls_attn, analytical_grad = get_attribution_map(
                    attention_scores, vs, sae_w, sae_b, layer=layer, head=head
                )
                final_scores = calculate_score(text_loc, analytical_grad, all_masks).sum(0)
                head_scores[(layer, head)] += final_scores

    for layer in MINE_LAYERS:
        for head in range(NUM_HEADS):
            head_scores[(layer, head)] = head_scores[(layer, head)].mean().item() / n_samples

    return head_scores


def build_layer_spec(head_scores, layers, n_heads, sigma):
    """Select heads with score > layer_mean + sigma * layer_std."""
    layer_spec = {}
    for layer in layers:
        layer_spec[layer] = []
        layer_scores = np.array([head_scores[(layer, head)] for head in range(n_heads)])
        layer_mean, layer_std = layer_scores.mean(), layer_scores.std()
        for head in range(n_heads):
            if head_scores[(layer, head)] > layer_mean + layer_std * sigma:
                layer_spec[layer].append(head)
    return layer_spec


# ---------------------------------------------------------------------------
# CLIP zero-shot classification + attention heatmap
# ---------------------------------------------------------------------------

def encode_text_prompts(labels):
    """Encode 'a photo of a <label>.' prompts into normalised text embeddings."""
    prompts = ["a photo of a " + l + "." for l in labels]
    inputs = _processor.tokenizer(prompts, padding=True, return_tensors="pt").to("cuda")
    with torch.no_grad():
        text_outputs = _text_model(**inputs)
        text_features = _text_projection(text_outputs.pooler_output)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
    return text_features


def classify(image_tensor, class_embeddings, layer_spec=None):
    """Zero-shot classify an image; optionally with head ablation (defense)."""
    from contextlib import nullcontext
    ctx = fix_attn_head_list(_vision, layer_spec, alpha=1.0) if layer_spec else nullcontext()
    with torch.inference_mode(), ctx:
        vision_outputs = _vision(pixel_values=image_tensor)
        image_features = _visual_projection(vision_outputs.pooler_output)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        sims = image_features @ class_embeddings.t()
        probs = sims.softmax(dim=-1)
    return probs.squeeze(0).cpu()


def get_cls_attention_heatmap(image_tensor, layer, head):
    """Extract the CLS-token attention map for a single (layer, head)."""
    with torch.inference_mode():
        _, qs, ks, _, _ = get_attention_scores(_vision, image_tensor)
    d_k = qs[layer].size(-1)
    scores = torch.matmul(qs[layer], ks[layer].transpose(-2, -1)) / np.sqrt(d_k)
    attn = F.softmax(scores, dim=-1)
    cls_attn = attn[0, head, 0, 1:]
    grid_size = int(math.sqrt(cls_attn.shape[0]))
    heatmap = cls_attn.cpu().numpy().reshape(grid_size, grid_size)
    return heatmap


def overlay_heatmap(pil_img, heatmap, alpha=0.5):
    """Overlay a normalised heatmap on a PIL image for visualisation."""
    w, h = pil_img.size
    hm = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    ax.imshow(pil_img)
    ax.imshow(hm, cmap="jet", alpha=alpha, extent=(0, w, h, 0))
    ax.axis("off")
    fig.tight_layout(pad=0)
    fig.canvas.draw()
    buf = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    a = buf[:, :, 0:1]
    rgb = buf[:, :, 1:4]
    rgba = np.concatenate([rgb, a], axis=2)
    from PIL import Image as PILImage
    img = PILImage.fromarray(rgba, "RGBA")
    plt.close(fig)
    return img.convert("RGB")


# ---------------------------------------------------------------------------
# Main inference function
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
def defend(image, labels_text, sigma=1.0, n_mining_samples=8):
    """Run the SamplingTAR defense on an uploaded image.

    Args:
        image: Input image (PIL Image).
        labels_text: Comma-separated candidate labels for zero-shot classification.
        sigma: Z-threshold for head selection (higher = fewer heads ablated).
        n_mining_samples: Number of text-bordered variants to generate for circuit mining.

    Returns:
        Tuple of (result_markdown, undefended_attn_img, defended_attn_img, ablated_heads_text).
    """
    if image is None:
        return "Please upload an image.", None, None, ""
    if not labels_text or not labels_text.strip():
        return "Please provide candidate labels (comma-separated).", None, None, ""

    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
    if len(labels) < 2:
        return "Please provide at least 2 candidate labels.", None, None, ""

    pil_img = image.convert("RGB")
    input_tensor = _preprocess(pil_img).unsqueeze(0).to("cuda")

    # Encode text labels
    class_emb = encode_text_prompts(labels)

    # --- Undefended classification ---
    probs_undef = classify(input_tensor, class_emb, layer_spec=None)
    pred_undef = labels[probs_undef.argmax().item()]

    # --- Mine text-reading heads ---
    head_scores = mine_head_scores(pil_img, labels, n_samples=int(n_mining_samples))
    layer_spec = build_layer_spec(head_scores, MINE_LAYERS, NUM_HEADS, sigma)

    # --- Defended classification ---
    probs_def = classify(input_tensor, class_emb, layer_spec=layer_spec)
    pred_def = labels[probs_def.argmax().item()]

    # --- Attention heatmaps ---
    best_lh = max(head_scores, key=head_scores.get)
    best_layer, best_head = best_lh

    undef_heatmap = get_cls_attention_heatmap(input_tensor, best_layer, best_head)
    with fix_attn_head_list(_vision, layer_spec, alpha=1.0):
        def_heatmap = get_cls_attention_heatmap(input_tensor, best_layer, best_head)

    undef_attn_img = overlay_heatmap(pil_img, undef_heatmap)
    def_attn_img = overlay_heatmap(pil_img, def_heatmap)

    # --- Format results ---
    top_k = min(5, len(labels))
    undef_top = torch.topk(probs_undef, top_k)
    def_top = torch.topk(probs_def, top_k)

    result_md = "## Results\n\n### Without Defense (Undefended CLIP)\n| Label | Probability |\n|---|---|\n"
    for i in range(top_k):
        idx = undef_top.indices[i].item()
        result_md += f"| {labels[idx]} | {undef_top.values[i].item():.4f} |\n"
    result_md += f"\n**Predicted: {pred_undef}**\n"

    result_md += "\n### With SamplingTAR Defense\n| Label | Probability |\n|---|---|\n"
    for i in range(top_k):
        idx = def_top.indices[i].item()
        result_md += f"| {labels[idx]} | {def_top.values[i].item():.4f} |\n"
    result_md += f"\n**Predicted: {pred_def}**\n"

    if pred_undef != pred_def:
        result_md += f"\n✅ **Defense changed the prediction** from '{pred_undef}' to '{pred_def}'"
    else:
        result_md += f"\n⚠️ Prediction unchanged ('{pred_def}'), but probability distribution may have shifted."

    ablated = []
    for layer in layer_spec:
        for head in layer_spec[layer]:
            ablated.append(f"L{layer}H{head}")
    ablated_text = f"Sigma={sigma:.1f}, Ablated {len(ablated)} heads: {', '.join(ablated[:20])}"
    if len(ablated) > 20:
        ablated_text += f" ... (+{len(ablated)-20} more)"

    return result_md, undef_attn_img, def_attn_img, ablated_text


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLE_LABELS = "bonnet, green mamba, langur, Doberman, gyromitra, vacuum, window screen, grass snake"

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        "# SamplingTAR: Training-free Defense against Typographic Attacks\n\n"
        "This demo implements the training-free concept localization method from "
        "[Towards Robustness against Typographic Attack with Training-free Concept Localization]"
        "(https://huggingface.co/papers/2607.02494). "
        "Upload an image containing a typographic attack (text overlaid on an image to fool CLIP), "
        "provide candidate labels, and the defense will ablate the attention heads responsible "
        "for reading text to recover the true object prediction."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(type="pil", label="Input Image", height=300)
                labels_input = gr.Textbox(
                    label="Candidate Labels (comma-separated)",
                    value=EXAMPLE_LABELS,
                    placeholder="e.g. cat, dog, bird",
                )
                run_btn = gr.Button("Run Defense", variant="primary")

                with gr.Accordion("Advanced Settings", open=False):
                    sigma_slider = gr.Slider(
                        label="Z-threshold (sigma)", minimum=0.0, maximum=3.0, value=1.0, step=0.5,
                        info="Higher = fewer heads ablated (more conservative). Lower = more heads ablated.",
                    )
                    n_samples_slider = gr.Slider(
                        label="Mining samples", minimum=4, maximum=16, value=8, step=2,
                        info="Number of text-bordered variants for circuit mining. More = better head scores.",
                    )

            with gr.Column(scale=1):
                result_md = gr.Markdown(label="Results")
                ablated_info = gr.Textbox(label="Ablated Heads", interactive=False)

        with gr.Row():
            undef_attn = gr.Image(label="Undefended Attention (top text-reading head)", height=250)
            def_attn = gr.Image(label="Defended Attention (same head, ablated)", height=250)

    run_btn.click(
        fn=defend,
        inputs=[image_input, labels_input, sigma_slider, n_samples_slider],
        outputs=[result_md, undef_attn, def_attn, ablated_info],
        api_name="defend",
    )

    gr.Examples(
        examples=[
            ["examples/0000.png", EXAMPLE_LABELS, 1.0, 8],
            ["examples/0001.png", EXAMPLE_LABELS, 1.0, 8],
            ["examples/0002.png", EXAMPLE_LABELS, 1.0, 8],
            ["examples/0003.png", EXAMPLE_LABELS, 1.0, 8],
        ],
        inputs=[image_input, labels_input, sigma_slider, n_samples_slider],
        outputs=[result_md, undef_attn, def_attn, ablated_info],
        fn=defend,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True)