"""
Per-layer ΔW effect — matched fixed-prefix probe.

Input pattern matches ``probe_output_trace.py``:

  * ``--prompt``                  → the user-side question (e.g. "Describe this image.")
  * ``--fixed_assistant_prefix``  → the assistant-side text teacher-forced after the prompt

For every filter-passing image, teacher-force ``prompt + " " +
fixed_assistant_prefix`` once under base weights, once per ``L_int`` under
ΔW@L_int (single-layer), and once under ΔW@all-layers. **No autoregressive
decoding** — the residuals come from a single forward pass per condition.

At every selected capture layer ``L_cap``:
  residual @ K assistant-prefix positions  →  FrozenSAEEncoder  →  (K, d_sae)
  pool over K (max|mean, ``--pool``)       →  (d_sae,)
  * gather the layer's selected features and aggregate (mean|max,
    ``--agg``) over the top-k feature dimension → ``per_layer_effect_L*.png``
  * apply the per-layer linear probe head (``probes.{L}.weight``,
    ``probes.{L}.bias``) → ``per_layer_effect_probe_L*.png``

The 32 ``L_int`` panels are **split across multiple PNGs** (``_GROUP_SIZE``
panels per PNG, default 4, 2×2 layout) so each panel is large enough to
read. Files land in ``{graph_dir}/{image_id}/`` as
``per_layer_effect_L00-03.png`` … ``per_layer_effect_L28-31.png`` and the
matching ``per_layer_effect_probe_L*.png`` set.

A companion summary plot ``per_layer_effect_all.png`` is also emitted: a
1×2 layout (features left, probes right) comparing base vs ΔW@all-layers
simultaneously.

For the same prompt + assistant-prefix + image, the probe numbers reported
here are identical to ``probe_output_trace.py`` (same SAE wrapper, same
max-pool, same float32 GPU torch.dot procedure).
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import traceback
from typing import Dict, List, Tuple

# Make `hallucination.*` importable for sae.SAE_Tools etc.
_PARENT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if _PARENT not in sys.path:
    sys.path.insert(0, _PARENT)
_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm import tqdm
from transformers import LlavaProcessor

from mechanistic_interp.constants import (
    AGG_CHOICES,
    CATEGORY_CHOICES,
    COLOR_BASE,
    COLOR_LORA,
    POOL_CHOICES,
    PROBE_OUTPUT_CHOICES,
    PROMPT_TEMPLATE,
)
from mechanistic_interp.delta_w_feature_trace import (
    build_hook_name,
    build_image_index,
    filter_samples,
    teacher_forced_capture,
)
from mechanistic_interp.lora_delta import applied_lora_pairs, load_lora_pairs
from experiment.training.finetune_adv import FrozenSAEEncoder
from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration


# ── SAE encode + per-feature gather + probe head (single residual) ───────────

@torch.no_grad()
def _sae_dense_at(resid: torch.Tensor, sae, pool: str, device) -> torch.Tensor:
    """SAE-encode a ``(K, D)`` residual slice via ``FrozenSAEEncoder``
    (JumpReLU dense), pool over K, return ``(d_sae,)`` device float32.

    Mirrors ``probe_output_trace.capture_probe_p`` exactly — same encoder,
    same pool, same float32 device tensor — so the downstream GPU
    ``torch.dot(dense, w) + b`` reports identical numbers to
    ``probe_output_trace`` for the same residual.
    """
    if resid.numel() == 0:
        d_sae = getattr(sae, "d_sae", None) or getattr(getattr(sae, "cfg", None), "d_sae", 0)
        return torch.zeros(int(d_sae), dtype=torch.float32, device=device)
    latents = sae(resid.to(device))                       # (K, d_sae)
    if pool == "max":
        pooled = latents.max(dim=0).values                # (d_sae,)
    else:
        pooled = latents.mean(dim=0)                      # (d_sae,)
    return pooled.float()


@torch.no_grad()
def _capture_acts_and_probes(
    *, model, sae, device, dtype_attn,
    full_ids, pixel_values, new_len: int,
    layers: List[int], hook_type: str, pool: str,
    selected: Dict[int, List[int]],
    probes: Dict[int, Dict[str, torch.Tensor]],
) -> Tuple[Dict[int, torch.Tensor], Dict[int, float]]:
    """Teacher-force once, hook every selected capture layer at
    ``hook_resid_<hook_type>``, slice the last ``new_len`` positions (the
    assistant-prefix tokens), pool, and return:
      * ``acts[L]``     — (top_k_L,) gathered feature values, CPU float32
      * ``probe_out[L]`` — float, probe logit = w[L] @ pooled + b[L]
                          (sigmoid applied later at plot time so the
                          ``--probe_output`` flag still works).

    Same slice + pool + GPU float32 ``torch.dot`` as
    ``probe_output_trace.capture_probe_p``.
    """
    hook_names = {build_hook_name(L, hook_type) for L in layers}
    cache = teacher_forced_capture(model, full_ids, pixel_values, hook_names, dtype_attn)
    acts: Dict[int, torch.Tensor] = {}
    probe_out: Dict[int, float] = {}
    for L in layers:
        hp = build_hook_name(L, hook_type)
        cache_t = cache.get(hp)
        if cache_t is None:
            continue
        slice_ = cache_t[0, -new_len:]                       # (K, D)
        dense = _sae_dense_at(slice_, sae, pool, device)     # (d_sae,) device f32
        feats = selected.get(L, [])
        if feats:
            fid = torch.tensor(feats, dtype=torch.long, device=dense.device)
            acts[L] = dense.index_select(0, fid).detach().cpu().clone()
        else:
            acts[L] = torch.zeros(0, dtype=torch.float32)
        p = probes.get(L)
        if p is not None:
            w = p["w"].to(device=dense.device, dtype=torch.float32)
            b = p["b"].to(device=dense.device, dtype=torch.float32)
            logit = torch.dot(dense, w) + b
            probe_out[L] = float(logit)
        else:
            probe_out[L] = float("nan")
    return acts, probe_out


# ── Probe head ───────────────────────────────────────────────────────────────

_PROBE_KEY_RE = re.compile(r"probes\.(\d+)\.weight$")


def load_probe_heads(probes_path: str) -> Dict[int, Dict[str, torch.Tensor]]:
    """Return ``{L: {'w': (d_sae,), 'b': scalar}}`` from a probes state-dict.

    Mirrors the parse used by ``select_features.select_by_probe`` — bias is
    stored as a 0-d tensor on CPU; weight is float32 (1, d_sae) squeezed to
    (d_sae,).
    """
    sd = torch.load(probes_path, map_location="cpu", weights_only=False)
    if hasattr(sd, "state_dict"):
        sd = sd.state_dict()
    out: Dict[int, Dict[str, torch.Tensor]] = {}
    for k, v in sd.items():
        m = _PROBE_KEY_RE.match(k)
        if m is None:
            continue
        L = int(m.group(1))
        out.setdefault(L, {})["w"] = v.squeeze(0).float().cpu()
    for k, v in sd.items():
        if k.endswith(".bias") and k.startswith("probes."):
            L = int(k.split(".")[1])
            if L in out:
                out[L]["b"] = v.float().cpu().reshape(())
    # Default missing biases to 0 just in case the dict was partial.
    for L, d in out.items():
        d.setdefault("b", torch.zeros((), dtype=torch.float32))
    return out


# ── Aggregation ──────────────────────────────────────────────────────────────

def _agg_layer_map(
    act_map: Dict[int, torch.Tensor], layers: List[int], agg: str,
) -> np.ndarray:
    """Reduce ``{L: (top_k_L,)}`` to a (n_layers,) array of scalars.

    ``mean`` averages the top-k features; ``max`` takes their peak.
    Missing/empty entries become NaN.
    """
    out = np.full(len(layers), np.nan, dtype=np.float32)
    for i, L in enumerate(layers):
        t = act_map.get(L)
        if t is None or t.numel() == 0:
            continue
        tf = t.float()
        out[i] = float(tf.mean()) if agg == "mean" else float(tf.max())
    return out


# ── Per-sample driver ────────────────────────────────────────────────────────

@torch.no_grad()
def trace_one_sample(
    *,
    sample,
    image_index,
    prompt,
    fixed_assistant_prefix: str,
    processor,
    model,
    sae,
    lora_pairs,
    lora_scale,
    selected: Dict[int, List[int]],
    layers: List[int],
    hook_type: str,
    pool: str,
    device,
    dtype_attn,
    probes: Dict[int, Dict[str, torch.Tensor]],
):
    image_id = sample["image_id"]
    image = image_index.get(str(image_id))
    if image is None:
        return None, f"image not found in HF split for {image_id}"

    # Tokenize prompt-only and prompt+assistant-prefix; the K assistant-prefix
    # positions are the last (matched_ids.shape[1] - prompt_len) tokens.
    # Same recipe as probe_output_trace.trace_one_image.
    text = PROMPT_TEMPLATE.format(question=prompt)
    inputs = processor(images=image, text=text, return_tensors="pt").to(device)
    prompt_len = int(inputs["input_ids"].shape[1])
    pixel_values = inputs["pixel_values"]

    full_text = text + " " + fixed_assistant_prefix
    full_inputs = processor(images=image, text=full_text,
                            return_tensors="pt").to(device)
    matched_ids = full_inputs["input_ids"]
    new_len = int(matched_ids.shape[1] - prompt_len)
    if new_len <= 0:
        return None, "prefix tokenized to 0 new tokens"
    matched_text = processor.tokenizer.decode(matched_ids[0, prompt_len:])

    def _capture():
        return _capture_acts_and_probes(
            model=model, sae=sae, device=device, dtype_attn=dtype_attn,
            full_ids=matched_ids, pixel_values=pixel_values, new_len=new_len,
            layers=layers, hook_type=hook_type, pool=pool,
            selected=selected, probes=probes,
        )

    # 1) Base.
    base_acts, base_probe = _capture()

    # 2) ΔW@L_int sweep — single-layer intervention.
    lora_acts: Dict[int, Dict[int, torch.Tensor]] = {}
    lora_probe: Dict[int, Dict[int, float]] = {}
    for L_int in tqdm(layers, desc=f"  L_int sweep ({image_id})", leave=False):
        if not selected.get(L_int):
            continue
        with applied_lora_pairs(
            model, lora_pairs, lora_scale,
            components="all", layers=[L_int], language_only=True,
            lowmem=False, save_device="cpu",
        ):
            a, p = _capture()
        lora_acts[L_int] = a
        lora_probe[L_int] = p

    # 3) ΔW@all-layers — every selected layer simultaneously.
    lora_all_acts: Dict[int, torch.Tensor] = {}
    lora_all_probe: Dict[int, float] = {}
    if any("language_model" in mp for mp in lora_pairs):
        with applied_lora_pairs(
            model, lora_pairs, lora_scale,
            components="all", layers=layers, language_only=True,
            lowmem=True,
        ):
            lora_all_acts, lora_all_probe = _capture()

    return {
        "image_id": image_id,
        "category": sample.get("category"),
        "prompt": prompt,
        "fixed_assistant_prefix": fixed_assistant_prefix,
        "matched_text": matched_text,
        "new_len": new_len,
        "base_acts": base_acts,
        "lora_acts": lora_acts,
        "lora_all_acts": lora_all_acts,
        "base_probe": base_probe,
        "lora_probe": lora_probe,
        "lora_all_probe": lora_all_probe,
        "feature_ids_per_layer": {L: selected[L] for L in layers if L in selected},
    }, None


# ── Plotting ─────────────────────────────────────────────────────────────────

# Panels per PNG. 32 L_int values → 8 PNGs at GROUP_SIZE=4. Layout is 2×2
# per PNG so each panel is large enough to read individual layer values.
_GROUP_SIZE = 4
_GROUP_ROWS = 2
_GROUP_COLS = 2


def _layer_groups(layers: List[int], group_size: int = _GROUP_SIZE) -> List[List[int]]:
    """Yield consecutive chunks of ``layers`` of size ``group_size``."""
    return [layers[i: i + group_size] for i in range(0, len(layers), group_size)]


def _group_tag(group: List[int]) -> str:
    """``[0,1,2,3]`` → ``'L00-03'`` (zero-padded so filenames sort)."""
    if not group:
        return "Lempty"
    return f"L{group[0]:02d}-{group[-1]:02d}"


def _render_per_layer_grid(
    *,
    out_path: str,
    layers: List[int],
    panel_layers: List[int],
    base_curve: np.ndarray,
    lora_curves: Dict[int, np.ndarray],
    image_id: str,
    suptitle_extra: str,
    y_label: str,
):
    """One PNG covering ``panel_layers`` (one panel per L_int) with the
    full ``layers`` set on the x-axis. Y-axis auto-scales per panel so
    early-layer detail isn't swallowed by late-layer spikes.

    ``base_curve`` has shape ``(len(layers),)``; ``lora_curves[L_int]``
    same. NaNs render as gaps. Aggregation choice (mean/max or sigmoid
    probe) is decided by the caller.
    """
    if not layers or not panel_layers:
        return
    rows = _GROUP_ROWS
    cols = _GROUP_COLS
    fig, axes = plt.subplots(
        rows, cols,
        figsize=(cols * 4.5, rows * 4.0),
        squeeze=False, sharex=True, sharey=False,
    )
    xs = np.arange(len(layers))
    xtick_step = max(1, len(layers) // 8)
    xtick_idx = np.arange(0, len(layers), xtick_step)
    base_line = lora_line = None
    for idx, L_int in enumerate(panel_layers):
        r = idx // cols
        c = idx % cols
        ax = axes[r][c]
        lora_curve = lora_curves.get(
            L_int, np.full(len(layers), np.nan, dtype=np.float32),
        )
        bl, = ax.plot(
            xs, base_curve,
            color=COLOR_BASE, linestyle="-", marker="o",
            markersize=4, linewidth=1.6, label="base",
        )
        ll, = ax.plot(
            xs, lora_curve,
            color=COLOR_LORA, linestyle="--", marker="s",
            markersize=4, linewidth=1.6, label="ΔW@L_int",
        )
        base_line = bl
        lora_line = ll
        if L_int in layers:
            ax.axvline(
                layers.index(L_int), color="black", linestyle=":",
                linewidth=0.9, alpha=0.5,
            )
        ax.set_title(f"L_int = {L_int}", fontsize=11, pad=3)
        ax.grid(True, linestyle=":", alpha=0.45)
        ax.set_xticks(xtick_idx)
        ax.set_xticklabels([str(layers[i]) for i in xtick_idx], fontsize=9)
        ax.tick_params(axis="y", labelsize=9)
        if r == rows - 1:
            ax.set_xlabel("capture layer L", fontsize=10)
        if c == 0:
            ax.set_ylabel(y_label, fontsize=10)
    for idx in range(len(panel_layers), rows * cols):
        r = idx // cols
        c = idx % cols
        axes[r][c].axis("off")
    if base_line is not None and lora_line is not None:
        fig.legend(
            [base_line, lora_line], ["base", "ΔW@L_int"],
            loc="upper center", ncol=2, bbox_to_anchor=(0.5, 0.975),
            fontsize=11, frameon=False,
        )
    fig.suptitle(
        f"{image_id} — per-layer ΔW effect, fixed-prefix probe   "
        f"(L_int ∈ {{{', '.join(str(L) for L in panel_layers)}}})   "
        f"{suptitle_extra}",
        fontsize=11, y=0.995,
    )
    fig.subplots_adjust(top=0.90, hspace=0.30, wspace=0.22)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    fig.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close(fig)


def _render_per_layer_groups(
    *,
    out_dir: str,
    filename_stem: str,
    layers: List[int],
    base_curve: np.ndarray,
    lora_curves: Dict[int, np.ndarray],
    image_id: str,
    suptitle_extra: str,
    y_label: str,
) -> List[str]:
    """Iterate ``_layer_groups`` and emit one PNG per group.

    Returns the list of paths written so the caller can log them.
    """
    written = []
    for group in _layer_groups(layers):
        tag = _group_tag(group)
        out_path = os.path.join(out_dir, f"{filename_stem}_{tag}.png")
        _render_per_layer_grid(
            out_path=out_path,
            layers=layers,
            panel_layers=group,
            base_curve=base_curve,
            lora_curves=lora_curves,
            image_id=image_id,
            suptitle_extra=suptitle_extra,
            y_label=y_label,
        )
        if os.path.exists(out_path):
            written.append(out_path)
    return written


def plot_per_layer_effect(
    *,
    out_dir: str,
    layers: List[int],
    base_acts: Dict[int, torch.Tensor],
    lora_acts: Dict[int, Dict[int, torch.Tensor]],
    image_id: str,
    matched_text: str,
    agg: str,
) -> List[str]:
    """Top-k feature activation, aggregated per layer with ``--agg``.

    Emits one PNG per ``L_int`` group under ``out_dir`` and returns the
    list of paths actually written.
    """
    base_curve = _agg_layer_map(base_acts, layers, agg)
    lora_curves = {
        L_int: _agg_layer_map(lora_acts.get(L_int, {}), layers, agg)
        for L_int in layers
    }
    return _render_per_layer_groups(
        out_dir=out_dir,
        filename_stem="per_layer_effect",
        layers=layers,
        base_curve=base_curve,
        lora_curves=lora_curves,
        image_id=image_id,
        suptitle_extra=f"(agg={agg}; prefix: {matched_text.strip()[:90]!r})",
        y_label=f"{agg}(top-k feats)",
    )


def _probe_curve(
    probe_map: Dict[int, float], layers: List[int], probe_output: str,
) -> np.ndarray:
    """Collapse ``{L: logit}`` to ``(n_layers,)`` of float, applying
    sigmoid if ``probe_output == 'prob'``. NaN cells stay NaN."""
    out = np.full(len(layers), np.nan, dtype=np.float32)
    for i, L in enumerate(layers):
        v = probe_map.get(L) if probe_map else None
        if v is None:
            continue
        z = float(v)
        if not np.isfinite(z):
            continue
        if probe_output == "prob":
            # Stable sigmoid: protect against overflow on large |z|.
            out[i] = float(1.0 / (1.0 + np.exp(-z))) if z >= 0 else float(
                np.exp(z) / (1.0 + np.exp(z))
            )
        else:
            out[i] = z
    return out


def plot_per_layer_effect_probe(
    *,
    out_dir: str,
    layers: List[int],
    base_probe: Dict[int, float],
    lora_probe: Dict[int, Dict[int, float]],
    image_id: str,
    matched_text: str,
    probe_output: str,
) -> List[str]:
    """Per-layer linear-probe output (probability or logit) across L_cap,
    one panel per L_int. Same group splitting as ``plot_per_layer_effect``.
    """
    base_curve = _probe_curve(base_probe, layers, probe_output)
    lora_curves = {
        L_int: _probe_curve(lora_probe.get(L_int, {}), layers, probe_output)
        for L_int in layers
    }
    y_label = "p(toilet)" if probe_output == "prob" else "probe logit"
    return _render_per_layer_groups(
        out_dir=out_dir,
        filename_stem="per_layer_effect_probe",
        layers=layers,
        base_curve=base_curve,
        lora_curves=lora_curves,
        image_id=image_id,
        suptitle_extra=f"(probe={probe_output}; prefix: {matched_text.strip()[:90]!r})",
        y_label=y_label,
    )


def plot_all_layer_effect(
    *,
    out_path: str,
    layers: List[int],
    base_acts: Dict[int, torch.Tensor],
    lora_all_acts: Dict[int, torch.Tensor],
    base_probe: Dict[int, float],
    lora_all_probe: Dict[int, float],
    image_id: str,
    matched_text: str,
    agg: str,
    probe_output: str,
):
    """Summary comparison plot: 1×2 layout, features (left) + probes (right),
    each showing ``base`` vs ``ΔW@all-layers`` across capture layers.

    Companion to the per-L_int split PNGs — lets readers see the full
    intervention's effect at a glance, then return to the per-L_int views
    to locate which layer is responsible.
    """
    if not layers:
        return
    feat_base = _agg_layer_map(base_acts, layers, agg)
    feat_lora = _agg_layer_map(lora_all_acts, layers, agg)
    prob_base = _probe_curve(base_probe, layers, probe_output)
    prob_lora = _probe_curve(lora_all_probe, layers, probe_output)

    fig, axes = plt.subplots(
        1, 2, figsize=(13.0, 4.6),
        squeeze=False, sharex=True, sharey=False,
    )
    xs = np.arange(len(layers))
    xtick_step = max(1, len(layers) // 8)
    xtick_idx = np.arange(0, len(layers), xtick_step)

    # ── Features panel ──────────────────────────────────────────────────
    ax_f = axes[0][0]
    ax_f.plot(xs, feat_base, color=COLOR_BASE, linestyle="-",  marker="o",
              markersize=4, linewidth=1.8, label="base")
    ax_f.plot(xs, feat_lora, color=COLOR_LORA, linestyle="--", marker="s",
              markersize=4, linewidth=1.8, label="ΔW@all-layers")
    ax_f.set_title("Features", fontsize=12, pad=4)
    ax_f.set_xlabel("capture layer L", fontsize=11)
    ax_f.set_ylabel(f"{agg}(top-k feats)", fontsize=11)
    ax_f.grid(True, linestyle=":", alpha=0.45)
    ax_f.set_xticks(xtick_idx)
    ax_f.set_xticklabels([str(layers[i]) for i in xtick_idx], fontsize=9)
    ax_f.tick_params(axis="y", labelsize=9)

    # ── Probes panel ────────────────────────────────────────────────────
    ax_p = axes[0][1]
    ax_p.plot(xs, prob_base, color=COLOR_BASE, linestyle="-",  marker="o",
              markersize=4, linewidth=1.8, label="base")
    ax_p.plot(xs, prob_lora, color=COLOR_LORA, linestyle="--", marker="s",
              markersize=4, linewidth=1.8, label="ΔW@all-layers")
    probe_y_label = "p(toilet)" if probe_output == "prob" else "probe logit"
    ax_p.set_title("Probes", fontsize=12, pad=4)
    ax_p.set_xlabel("capture layer L", fontsize=11)
    ax_p.set_ylabel(probe_y_label, fontsize=11)
    ax_p.grid(True, linestyle=":", alpha=0.45)
    ax_p.set_xticks(xtick_idx)
    ax_p.set_xticklabels([str(layers[i]) for i in xtick_idx], fontsize=9)
    ax_p.tick_params(axis="y", labelsize=9)

    handles, labels_ = ax_f.get_legend_handles_labels()
    fig.legend(
        handles, labels_,
        loc="upper center", ncol=2, bbox_to_anchor=(0.5, 0.975),
        fontsize=11, frameon=False,
    )
    fig.suptitle(
        f"{image_id} — ΔW@all-layers vs base, fixed-prefix probe   "
        f"(agg={agg}, probe={probe_output}; prefix: {matched_text.strip()[:90]!r})",
        fontsize=11, y=0.995,
    )
    fig.subplots_adjust(top=0.86, wspace=0.22)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    fig.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close(fig)


def plot_sample(result: dict, graph_dir: str, agg: str, probe_output: str):
    image_id = result["image_id"]
    layers = sorted(result["feature_ids_per_layer"].keys())
    sample_dir = os.path.join(graph_dir, str(image_id))
    os.makedirs(sample_dir, exist_ok=True)
    matched_text = result.get("matched_text", "")

    try:
        written = plot_per_layer_effect(
            out_dir=sample_dir,
            layers=layers,
            base_acts=result["base_acts"],
            lora_acts=result["lora_acts"],
            image_id=str(image_id),
            matched_text=matched_text,
            agg=agg,
        )
        for p in written:
            print(f"  saved {os.path.abspath(p)}")
    except Exception as e:
        print(f"  per_layer_effect plot failed for {image_id}: {e}")
        traceback.print_exc()

    base_probe = result.get("base_probe") or {}
    lora_probe = result.get("lora_probe") or {}
    if base_probe and lora_probe:
        try:
            written = plot_per_layer_effect_probe(
                out_dir=sample_dir,
                layers=layers,
                base_probe=base_probe,
                lora_probe=lora_probe,
                image_id=str(image_id),
                matched_text=matched_text,
                probe_output=probe_output,
            )
            for p in written:
                print(f"  saved {os.path.abspath(p)}")
        except Exception as e:
            print(f"  per_layer_effect_probe plot failed for {image_id}: {e}")
            traceback.print_exc()

    # Companion summary: ΔW@all-layers vs base, features + probes side by side.
    lora_all_acts = result.get("lora_all_acts") or {}
    lora_all_probe = result.get("lora_all_probe") or {}
    if lora_all_acts and lora_all_probe and base_probe:
        all_path = os.path.join(sample_dir, "per_layer_effect_all.png")
        try:
            plot_all_layer_effect(
                out_path=all_path,
                layers=layers,
                base_acts=result["base_acts"],
                lora_all_acts=lora_all_acts,
                base_probe=base_probe,
                lora_all_probe=lora_all_probe,
                image_id=str(image_id),
                matched_text=matched_text,
                agg=agg,
                probe_output=probe_output,
            )
            if os.path.exists(all_path):
                print(f"  saved {os.path.abspath(all_path)}")
        except Exception as e:
            print(f"  per_layer_effect_all plot failed for {image_id}: {e}")
            traceback.print_exc()


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--features_json", required=True,
                   help="Output of select_features.py (per-layer top-k).")
    p.add_argument("--samples_json", default="mechanistic_interp/toilet-bathroom/lora_adapter/samples.json")
    p.add_argument("--prompt", default="Describe this image.")
    p.add_argument("--hf_dataset", default="pbcong/bathroom-toilet")
    p.add_argument("--hf_split", default="validation")
    p.add_argument("--id_col", default="image_id")

    p.add_argument("--adapter_path",
                   default="mechanistic_interp/toilet-bathroom/lora_adapter/adapter_model.safetensors")
    p.add_argument("--adapter_cfg",
                   default="mechanistic_interp/toilet-bathroom/lora_adapter/adapter_config.json")

    p.add_argument("--sae_ckpt", required=True)

    p.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    p.add_argument("--device", default="cuda:0")
    p.add_argument("--dtype", default="bfloat16",
                   choices=["float32", "float16", "bfloat16"])

    p.add_argument("--n_samples", type=int, default=0,
                   help="Max filter-passing samples to process. 0 (default) = all.")
    p.add_argument("--hook_type", default="post", choices=["pre", "mid", "post"])
    p.add_argument("--pool", choices=POOL_CHOICES, default="max",
                   help="Pool over the K assistant-prefix positions before "
                        "applying the probe / gathering features. Must match "
                        "probe-training pool ('max' matches train_probe_gen).")
    p.add_argument("--fixed_assistant_prefix", required=True,
                   help="Assistant-side text teacher-forced after the question; "
                        "every condition (base / ΔW@L_int / ΔW@all) sees the same "
                        "tokens. Example: 'In this bathroom there is a shower, a "
                        "sink and'")

    p.add_argument("--out_dir", required=True)
    p.add_argument("--graph_dir", default="graph",
                   help="Plots written to {graph_dir}/{image_id}/: "
                        "per_layer_effect_L*.png (features) + "
                        "per_layer_effect_probe_L*.png (probes) + "
                        "per_layer_effect_all.png (base vs ΔW@all summary). "
                        "Pass --no_plots to skip.")
    p.add_argument("--no_plots", action="store_true",
                   help="Skip inline plotting.")
    p.add_argument("--agg", choices=AGG_CHOICES, default="max",
                   help="Reducer over the top-k feature dim per (L_int, L_cap). "
                        "'max' (default) for peak intensity; 'mean' for average.")
    p.add_argument("--object_name", default="toilet",
                   help="Primary object name for category filtering (e.g., 'toilet', 'oven', 'tv').")
    p.add_argument("--object2_name", default=None,
                   help="Second object name for category filtering (e.g., 'bathroom' when "
                        "object_name='toilet'). If provided, enables category choices like "
                        "'{object1}_only', '{object2}_only', '{object1}_{object2}'.")
    p.add_argument("--category", choices=CATEGORY_CHOICES, default="any",
                   help="Filter samples by per-sample 'category' field. "
                        "'any' (default) or '{object1}_only' | '{object2}_only' | '{object1}_{object2}' "
                        "when --object2_name is provided.")
    p.add_argument("--probes_path",
                   default="mechanistic_interp/probes/probes_gen_bathroom_toilet.pt",
                   help="Per-layer linear probe state-dict "
                        "(probes.{L}.weight: (1, d_sae), probes.{L}.bias: (1,)). "
                        "Used for the per_layer_effect_probe.png plot.")
    p.add_argument("--probe_output", choices=PROBE_OUTPUT_CHOICES, default="prob",
                   help="Probe output to plot. 'prob' (default) applies sigmoid "
                        "and shows p(toilet) in [0,1]; 'logit' plots the raw "
                        "linear-probe score.")
    args = p.parse_args()

    dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
    dtype = dtype_map[args.dtype]

    os.makedirs(args.out_dir, exist_ok=True)
    torch.set_grad_enabled(False)

    # ── Selected features ───────────────────────────────────────────────────
    with open(args.features_json) as f:
        feat_json = json.load(f)
    selected: Dict[int, List[int]] = {}
    for k, v in feat_json.items():
        if not k.startswith("layer_"):
            continue
        L = int(k.split("_")[1])
        selected[L] = list(map(int, v["features"]))
    layers = sorted(selected.keys())
    print(f"Selected features for {len(layers)} layers (top_k = {len(selected[layers[0]])})")

    # ── Filtered samples ────────────────────────────────────────────────────
    # gen_mode='scratch' bypasses the base=T/lora=F predicate that only makes
    # sense in the prefix-decode flow — here we teacher-force a fixed prefix.
    keep = filter_samples(args.samples_json, args.prompt, args.category,
                          gen_mode="scratch", allowed_categories=CATEGORY_CHOICES)
    print(f"Filter category={args.category}: {len(keep)} samples "
          f"(prompt={args.prompt!r}, prefix={args.fixed_assistant_prefix!r})")
    if args.n_samples > 0:
        keep = keep[: args.n_samples]
        print(f"Processing first {len(keep)} (capped by --n_samples={args.n_samples})")
    else:
        print(f"Processing all {len(keep)} samples (--n_samples=0)")

    # ── Model + LoRA + SAE ──────────────────────────────────────────────────
    print("Loading model …")
    processor = LlavaProcessor.from_pretrained(args.model_name)
    model = HookedSAELlavaConditionalGeneration.from_pretrained(
        args.model_name, attn_implementation="eager",
    ).to(args.device, dtype=dtype).eval()

    cfg = json.loads(open(args.adapter_cfg).read())
    lora_scale = cfg["lora_alpha"] / cfg["r"]
    pairs = load_lora_pairs(args.adapter_path)
    n_lm_pairs = sum(1 for mp in pairs if "language_model" in mp)
    print(f"LoRA (A,B) pairs: {len(pairs)} (language_model: {n_lm_pairs})  | scale={lora_scale}")

    print(f"Loading FrozenSAEEncoder from {args.sae_ckpt} (JumpReLU dense; "
          "matches probe-training SAE)")
    sae = FrozenSAEEncoder.from_checkpoint(args.sae_ckpt, torch.device(args.device))

    print(f"Loading probes ({args.probes_path}) …")
    probes = load_probe_heads(args.probes_path)
    missing = [L for L in layers if L not in probes]
    if missing:
        print(f"  WARN: no probe head for layers {missing} — those panels show NaN.")
    # Move once to device, matching probe_output_trace.py:404-405 — the
    # per-call float32 dot product runs on GPU without extra .to() churn.
    for L, d in probes.items():
        d["w"] = d["w"].to(args.device)
        d["b"] = d["b"].to(args.device)
    print(f"  loaded {len(probes)} probe heads (moved to {args.device})")

    # ── HF image index ──────────────────────────────────────────────────────
    needed_ids = {str(s["image_id"]) for s in keep}
    image_index = build_image_index(args.hf_dataset, args.hf_split, args.id_col)
    have = needed_ids & set(image_index.keys())
    print(f"HF images indexed: {len(image_index)} | needed: {len(needed_ids)} | resolved: {len(have)}")

    # ── Per-sample loop ─────────────────────────────────────────────────────
    summary = []
    n_ok = n_skip = 0
    for s in tqdm(keep, desc="samples"):
        result, err = trace_one_sample(
            sample=s,
            image_index=image_index,
            prompt=args.prompt,
            fixed_assistant_prefix=args.fixed_assistant_prefix,
            processor=processor,
            model=model,
            sae=sae,
            lora_pairs=pairs,
            lora_scale=lora_scale,
            selected=selected,
            layers=layers,
            hook_type=args.hook_type,
            pool=args.pool,
            device=args.device,
            dtype_attn=torch.long,
            probes=probes,
        )
        if result is None:
            n_skip += 1
            print(f"  skip {s['image_id']}: {err}")
            continue
        result["agg"] = args.agg
        result["probe_output"] = args.probe_output
        out_path = os.path.join(args.out_dir, f"{result['image_id']}.pt")
        torch.save(result, out_path)
        summary.append({"image_id": result["image_id"], "path": out_path})
        n_ok += 1

        if not args.no_plots:
            plot_sample(
                result, args.graph_dir,
                agg=args.agg, probe_output=args.probe_output,
            )

    with open(os.path.join(args.out_dir, "summary.json"), "w") as f:
        json.dump({
            "config": vars(args),
            "n_ok": n_ok,
            "n_skipped": n_skip,
            "samples": summary,
        }, f, indent=2)

    print(f"Done. ok={n_ok} skipped={n_skip}. Output → {args.out_dir}")


if __name__ == "__main__":
    main()