"""
knowledge_suppression_trace.py — internal SAE-activation comparison.

For each bathroom-only image (object perceptually absent, scene present):

  1. The base model free-greedy-decodes K tokens. This sequence is the
     **matched input** shared across all methods. Using one fixed
     sequence isolates the effect of *weights* from the confound of each
     method generating a different continuation.

  2. Each method (base / ours = ΔW@all from the LoRA adapter / Nullu) is
     teacher-forced on the matched K-token sequence under a fresh
     forward pass.

  3. At every selected layer, the residual stream is captured **only at
     the K generated-text positions** (sliced from the tail, so any
     image-token expansion in the middle of the sequence is irrelevant
     to the slice). The residuals are SAE-encoded and the pre-selected
     "confident toilet" features (per layer) are gathered.

  4. Per-layer scalar = aggregator over (K positions × top_k features).
     A single PNG per image plots one curve per method; a population
     summary aggregates across all images.

Why text positions only: we hypothesise the toilet *knowledge* lives in
the LLM's text-side computation. So we probe at the residuals carrying
the assistant's continuation, not at the image-patch positions.

No τ_c threshold — the claim is comparative: at every layer, the
"ours" curve should sit below Nullu/EFUF. Output-suppression methods
(Nullu/EFUF) leave a mid-network hump in the trajectory; a method that
removes the knowledge from the weights should keep the curve flat
throughout.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import traceback
from contextlib import contextmanager
from typing import Dict, List

# Make the repo importable (mirrors delta_w_feature_trace.py).
_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.delta_w_feature_trace import (
    CATEGORY_CHOICES,
    PROMPT_TEMPLATE,
    build_hook_name,
    build_image_index,
    filter_samples,
    sae_lookup,
    teacher_forced_capture,
)
from mechanistic_interp.lora_delta import applied_lora_pairs, load_lora_pairs
from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from sae.SAE_Tools import load_sae_model


# ── Method registry (plot styles + labels) ───────────────────────────────────

COLOR_BASE = "#1f77b4"
COLOR_OURS = "#2ca02c"
COLOR_NULLU = "#ff7f0e"
COLOR_EFUF = "#9467bd"

METHOD_ORDER = ("base", "ours", "nullu", "efuf")
METHOD_STYLES = {
    "base":  dict(color=COLOR_BASE,  linestyle="-",   marker="o", label="base"),
    "ours":  dict(color=COLOR_OURS,  linestyle="--",  marker="s", label="ΔW@all (ours)"),
    "nullu": dict(color=COLOR_NULLU, linestyle="-.",  marker="^", label="Nullu"),
    "efuf":  dict(color=COLOR_EFUF,  linestyle=":",   marker="D", label="EFUF"),
}

AGG_CHOICES = ("max", "mean")


# ── Nullu: full per-layer splice (all 9 LlamaDecoderLayer weights) ───────────
# Mirrors Nullu/scripts/eval_relation.py:_splice_edited_layers. Nullu edits all
# decoder-block weights (self_attn.{q,k,v,o}_proj, mlp.{gate,up,down}_proj,
# input_layernorm, post_attention_layernorm) for the configured layer range;
# the previous applied_nullu_down_proj swapped only mlp.down_proj, which is
# incorrect — that gives Nullu zero credit for the q/k/v/gate/up edits.
#
# Nullu's checkpoint uses liuhaotian-style keys (model.layers.{L}.*); HF LLaVA
# stores the same modules at model.language_model.layers.{L}.*. We map by
# attribute access on the existing GPU parameter tensors so we never hold a
# second 7B model on GPU.

_NULLU_LAYER_PARAM_NAMES = (
    "input_layernorm.weight",
    "post_attention_layernorm.weight",
    "self_attn.q_proj.weight",
    "self_attn.k_proj.weight",
    "self_attn.v_proj.weight",
    "self_attn.o_proj.weight",
    "mlp.gate_proj.weight",
    "mlp.up_proj.weight",
    "mlp.down_proj.weight",
)


def _get_layer_param(model, layer_idx: int, param_name: str) -> torch.nn.Parameter:
    """Resolve `model.model.language_model.layers[layer_idx].<param_name>`."""
    mod = model.model.language_model.layers[layer_idx]
    for attr in param_name.split("."):
        mod = getattr(mod, attr)
    return mod  # final attr is a Parameter (the .weight tensor)


def _load_nullu_layer_weights(
    nullu_model_path: str, layer_indices: List[int]
) -> Dict[int, Dict[str, "torch.Tensor"]]:
    """Read Nullu's edited per-layer weights from a HF-format directory. Mirrors
    ``Nullu/scripts/eval_relation.py:_splice_edited_layers``: walks the
    safetensors shard map and pulls every key matching ``model.layers.{L}.*``
    for L in ``layer_indices``. Returns ``{L: {param_name: cpu_tensor}}``.
    """
    import os
    from safetensors import safe_open

    index_path = os.path.join(nullu_model_path, "model.safetensors.index.json")
    single_path = os.path.join(nullu_model_path, "model.safetensors")
    prefixes = tuple(f"model.layers.{L}." for L in layer_indices)

    out: Dict[int, Dict[str, "torch.Tensor"]] = {L: {} for L in layer_indices}

    def _stash_key(key: str, tensor: "torch.Tensor"):
        # key like "model.layers.16.mlp.down_proj.weight" → L=16, param="mlp.down_proj.weight"
        if not key.startswith("model.layers."):
            return
        rest = key[len("model.layers."):]
        layer_str, _, param_name = rest.partition(".")
        try:
            L = int(layer_str)
        except ValueError:
            return
        if L not in out:
            return
        if param_name in _NULLU_LAYER_PARAM_NAMES:
            out[L][param_name] = tensor.detach().cpu()

    if os.path.exists(index_path):
        with open(index_path) as f:
            weight_map = json.load(f)["weight_map"]
        # Group target keys by shard.
        shards: Dict[str, List[str]] = {}
        for key in weight_map.keys():
            if key.startswith(prefixes):
                shards.setdefault(weight_map[key], []).append(key)
        for shard, keys in shards.items():
            with safe_open(os.path.join(nullu_model_path, shard),
                           framework="pt", device="cpu") as f:
                for key in keys:
                    _stash_key(key, f.get_tensor(key))
    elif os.path.exists(single_path):
        with safe_open(single_path, framework="pt", device="cpu") as f:
            for key in f.keys():
                if key.startswith(prefixes):
                    _stash_key(key, f.get_tensor(key))
    else:
        raise FileNotFoundError(
            f"No safetensors in {nullu_model_path}. "
            f"Expected model.safetensors.index.json or model.safetensors."
        )

    missing = {L: [p for p in _NULLU_LAYER_PARAM_NAMES if p not in out[L]]
               for L in layer_indices}
    missing = {L: ps for L, ps in missing.items() if ps}
    if missing:
        raise RuntimeError(
            f"Nullu ckpt missing per-layer params:\n  " +
            "\n  ".join(f"L{L}: {ps}" for L, ps in missing.items())
        )
    return out


@contextmanager
def applied_nullu_layers(
    model,
    edited_cpu: Dict[int, Dict[str, "torch.Tensor"]],
    original_cpu: Dict[int, Dict[str, "torch.Tensor"]],
):
    """Swap Nullu's full edited decoder layers (8-32 by default) in-place on
    the existing GPU model; restore originals on exit. Never holds a second
    7B model on GPU."""
    try:
        for L, params in edited_cpu.items():
            for pname, src in params.items():
                tgt = _get_layer_param(model, L, pname)
                tgt.data.copy_(src.to(device=tgt.device, dtype=tgt.dtype))
        yield
    finally:
        for L, params in original_cpu.items():
            for pname, src in params.items():
                tgt = _get_layer_param(model, L, pname)
                tgt.data.copy_(src.to(device=tgt.device, dtype=tgt.dtype))


# ── EFUF: in-place MM-projector swap (CPU↔GPU, parity with Nullu pattern) ────
# liuhaotian-format LLaVA keys (in the EFUF ckpt) map to HF-format keys
# (in our HookedSAELlavaConditionalGeneration) as:
#   model.mm_projector.0.{weight,bias}  ->  model.multi_modal_projector.linear_1.{weight,bias}
#   model.mm_projector.2.{weight,bias}  ->  model.multi_modal_projector.linear_2.{weight,bias}
EFUF_KEY_MAP = {
    "model.mm_projector.0.weight": "linear_1.weight",
    "model.mm_projector.0.bias":   "linear_1.bias",
    "model.mm_projector.2.weight": "linear_2.weight",
    "model.mm_projector.2.bias":   "linear_2.bias",
}


def _load_efuf_projector_weights(ckpt_path: str) -> Dict[str, "torch.Tensor"]:
    """Load EFUF's edited mm_projector weights, mapped to HF names. Returns
    ``{'linear_1.weight': T, 'linear_1.bias': T, 'linear_2.weight': T, 'linear_2.bias': T}``
    on CPU."""
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
    out = {}
    for liuhao_key, hf_name in EFUF_KEY_MAP.items():
        if liuhao_key not in sd:
            raise KeyError(
                f"EFUF ckpt {ckpt_path!r} missing {liuhao_key!r}. "
                f"Available: {[k for k in sd.keys() if 'projector' in k]}"
            )
        out[hf_name] = sd[liuhao_key].detach().cpu().clone()
    return out


@contextmanager
def applied_efuf_projector(
    model, edited_cpu: Dict[str, torch.Tensor], original_cpu: Dict[str, torch.Tensor],
):
    """Swap the HF multi_modal_projector's 4 weights with EFUF's edited
    versions, then restore the originals on exit. CPU→GPU one-shot copies
    into existing GPU tensors — never holds two model copies on GPU.
    """
    mmp = model.model.multi_modal_projector
    submap = {
        "linear_1.weight": mmp.linear_1.weight,
        "linear_1.bias":   mmp.linear_1.bias,
        "linear_2.weight": mmp.linear_2.weight,
        "linear_2.bias":   mmp.linear_2.bias,
    }
    try:
        for k, src in edited_cpu.items():
            tgt = submap[k]
            tgt.data.copy_(src.to(device=tgt.device, dtype=tgt.dtype))
        yield
    finally:
        for k, src in original_cpu.items():
            tgt = submap[k]
            tgt.data.copy_(src.to(device=tgt.device, dtype=tgt.dtype))


# ── Capture + aggregation primitives ─────────────────────────────────────────

@torch.no_grad()
def capture_text_pos_acts(
    *,
    model,
    sae,
    sae_batch,
    device,
    dtype_attn,
    full_ids: torch.Tensor,
    pixel_values,
    new_len: int,
    layers: List[int],
    hook_type: str,
    selected: Dict[int, List[int]],
) -> Dict[int, torch.Tensor]:
    """Teacher-force ``full_ids`` (1, T_input) once, hook every selected
    layer, then take the last ``new_len`` residual positions (the assistant's
    K generated tokens — unambiguously after any image-token expansion).
    SAE-encode and gather ``selected[L]``. Returns ``{L: (new_len, top_k_L)}``.
    """
    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] = {}
    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:]  # last new_len = generated text positions
        feats = selected.get(L, [])
        if feats:
            acts[L] = sae_lookup(slice_, feats, sae, sae_batch, device)
        else:
            acts[L] = torch.zeros(new_len, 0)
    return acts


def per_layer_scalar(
    act_map: Dict[int, torch.Tensor], layers: List[int], agg: str
) -> np.ndarray:
    """Collapse ``{L: (K, top_k)}`` → ``(n_layers,)`` per-layer scalar."""
    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.max()) if agg == "max" else float(tf.mean())
    return out


# ── Per-image work ───────────────────────────────────────────────────────────

@torch.no_grad()
def trace_one_image(
    *,
    sample,
    image_index,
    prompt,
    processor,
    model,
    sae,
    lora_pairs,
    lora_scale,
    nullu_payload,
    efuf_payload,
    selected,
    layers,
    hook_type,
    gen_tokens,
    sae_batch,
    device,
    dtype_attn,
    fixed_assistant_prefix: str = "",
):
    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}"

    text = PROMPT_TEMPLATE.format(question=prompt)
    inputs = processor(images=image, text=text, return_tensors="pt").to(device)
    prompt_ids = inputs["input_ids"]
    pixel_values = inputs["pixel_values"]
    prompt_len = int(prompt_ids.shape[1])

    if fixed_assistant_prefix:
        # 1a) Deterministic prefix mode — every method sees the SAME assistant
        # text. No free-gen needed: tokenize the full string (prompt + the
        # fixed assistant prefix) and the assistant-side tokens are the last
        # `new_len` positions we capture under each method.
        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, "fixed_assistant_prefix tokenized to 0 new tokens"
        matched_text = processor.tokenizer.decode(matched_ids[0, prompt_len:])
    else:
        # 1b) Free-gen mode — base produces the matched K-token sequence.
        gen = model.generate(
            **inputs,
            do_sample=False, num_beams=1, use_cache=True,
            max_new_tokens=gen_tokens,
        )
        matched_ids = gen[:, : prompt_len + gen_tokens]
        new_len = int(matched_ids.shape[1] - prompt_len)
        if new_len <= 0:
            return None, "base produced no new tokens"
        matched_text = processor.tokenizer.decode(matched_ids[0, prompt_len:])

    out = {
        "image_id": image_id,
        "category": sample.get("category"),
        "matched_text": matched_text,
        "new_len": new_len,
        "feature_ids_per_layer": {L: selected[L] for L in layers if L in selected},
        "acts": {},  # {method_key: {L: (new_len, top_k_L)}}
    }

    def _capture():
        return capture_text_pos_acts(
            model=model, sae=sae, sae_batch=sae_batch,
            device=device, dtype_attn=dtype_attn,
            full_ids=matched_ids, pixel_values=pixel_values, new_len=new_len,
            layers=layers, hook_type=hook_type, selected=selected,
        )

    # 2) base — no edit applied.
    out["acts"]["base"] = _capture()

    # 3) ours — ΔW@all from the LoRA adapter, applied in-place to LM layers.
    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,
        ):
            out["acts"]["ours"] = _capture()

    # 4) Nullu — full per-layer splice (8-32) swapped in-place.
    if nullu_payload is not None:
        with applied_nullu_layers(
            model,
            edited_cpu=nullu_payload["edited_cpu"],
            original_cpu=nullu_payload["original_cpu"],
        ):
            out["acts"]["nullu"] = _capture()

    # 5) EFUF — edited MM-projector weights swapped in-place (4 tensors).
    if efuf_payload is not None:
        with applied_efuf_projector(
            model,
            edited_cpu=efuf_payload["edited_cpu"],
            original_cpu=efuf_payload["original_cpu"],
        ):
            out["acts"]["efuf"] = _capture()

    return out, None


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

def _draw_curves(ax, scalars: Dict[str, np.ndarray], layers, xs):
    for key in METHOD_ORDER:
        ys = scalars.get(key)
        if ys is None:
            continue
        style = METHOD_STYLES[key]
        ax.plot(
            xs, ys,
            color=style["color"], linestyle=style["linestyle"], marker=style["marker"],
            linewidth=2, markersize=4, label=style["label"],
        )
    ax.set_xticks(xs)
    ax.set_xticklabels([str(L) for L in layers], fontsize=8)
    ax.set_xlabel("capture layer")
    ax.legend(loc="best")
    ax.grid(True, linestyle=":", alpha=0.4)


def plot_per_sample(out, layers, graph_dir, agg):
    image_id = out["image_id"]
    sample_dir = os.path.join(graph_dir, str(image_id))
    os.makedirs(sample_dir, exist_ok=True)

    xs = np.arange(len(layers))
    scalars = {
        key: per_layer_scalar(out["acts"][key], layers, agg)
        for key in out["acts"]
    }

    fig, ax = plt.subplots(figsize=(11, 5))
    _draw_curves(ax, scalars, layers, xs)
    ax.set_ylabel(f"per-layer scalar  [{agg} over (K positions × top-k features)]")
    title = (
        f"{image_id} — internal toilet-feature activation (matched input)\n"
        f"matched continuation: {out['matched_text'].strip()[:80]!r}"
    )
    ax.set_title(title, fontsize=10, loc="left")
    out_path = os.path.join(sample_dir, "internal_activation_trace.png")
    fig.savefig(out_path, dpi=110, bbox_inches="tight")
    plt.close(fig)
    return out_path


def plot_summary(per_image_scalars, layers, graph_dir, agg, n_images):
    fig, ax = plt.subplots(figsize=(11, 5))
    xs = np.arange(len(layers))
    for key in METHOD_ORDER:
        lst = per_image_scalars.get(key) or []
        if not lst:
            continue
        arr = np.stack(lst, axis=0)  # (N, n_layers)
        med = np.nanmedian(arr, axis=0)
        lo = np.nanpercentile(arr, 25, axis=0)
        hi = np.nanpercentile(arr, 75, axis=0)
        style = METHOD_STYLES[key]
        ax.plot(
            xs, med,
            color=style["color"], linestyle=style["linestyle"], marker=style["marker"],
            linewidth=2, markersize=4, label=f"{style['label']} (N={len(lst)})",
        )
        ax.fill_between(xs, lo, hi, color=style["color"], alpha=0.15)
    ax.set_xticks(xs)
    ax.set_xticklabels([str(L) for L in layers], fontsize=8)
    ax.set_xlabel("capture layer")
    ax.set_ylabel(f"per-layer scalar  [{agg} over (K positions × top-k features)]")
    ax.set_title(
        f"Population summary across {n_images} bathroom-only images "
        f"(median ± IQR)",
        fontsize=11, loc="left",
    )
    ax.legend(loc="best")
    ax.grid(True, linestyle=":", alpha=0.4)
    out_path = os.path.join(graph_dir, "_summary.png")
    os.makedirs(graph_dir, exist_ok=True)
    fig.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close(fig)
    return out_path


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

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--features_json", required=True,
                   help="Per-layer top-k confident toilet features "
                        "(same format as select_features.py output).")
    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("--nullu_model_path", default=None,
                   help="Path to Nullu's edited model dir. If unset, Nullu is skipped.")
    p.add_argument("--nullu_lowest_layer", type=int, default=16)
    p.add_argument("--nullu_highest_layer", type=int, default=32)
    p.add_argument("--efuf_ckpt", default=None,
                   help="Path to an EFUF epoch_XXX.pth (liuhaotian-format state-dict "
                        "containing model.mm_projector.{0,2}.{weight,bias}). "
                        "If unset, the EFUF curve is omitted.")

    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="0 = all category-matching samples.")
    p.add_argument("--gen_tokens", type=int, default=32,
                   help="K — length of the matched base continuation. "
                        "Ignored when --fixed_assistant_prefix is non-empty.")
    p.add_argument("--fixed_assistant_prefix", default="",
                   help="If non-empty, skip base free-generation and use this "
                        "string as the assistant-side text under every method. "
                        "K = number of tokens it produces. Example: "
                        "\"In this image there is a toilet\"")
    p.add_argument("--hook_type", default="post", choices=["pre", "mid", "post"])
    p.add_argument("--sae_batch", type=int, default=2048)
    p.add_argument("--category", choices=CATEGORY_CHOICES, default="bathroom_only",
                   help="Default 'bathroom_only': D_{A,¬c} from the spec.")
    p.add_argument("--agg", choices=AGG_CHOICES, default="max",
                   help="Per-layer aggregator. 'max' = peak across (positions × features); "
                        "'mean' = average across the same set.")

    p.add_argument("--out_dir", required=True)
    p.add_argument("--graph_dir", required=True)
    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)
    os.makedirs(args.graph_dir, exist_ok=True)
    torch.set_grad_enabled(False)

    # Confident toilet 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())
    if not layers:
        raise ValueError("no per-layer features in features_json")
    print(f"Confident toilet features for {len(layers)} layers "
          f"(top_k={len(selected[layers[0]])})")

    # Filter samples — bathroom_only by default. We use gen_mode='scratch'
    # so the base=T/lora=F predicate (which only makes sense in the
    # prefix-mode workflow) is bypassed; every category-matching sample
    # is included.
    keep = filter_samples(args.samples_json, args.prompt, args.category, gen_mode="scratch")
    print(f"Filter category={args.category}: {len(keep)} samples")
    if args.n_samples > 0:
        keep = keep[: args.n_samples]
        print(f"Capped to first {len(keep)} (--n_samples={args.n_samples})")

    # Model + LoRA + Nullu + 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 = sum(1 for mp in pairs if "language_model" in mp)
    print(f"LoRA pairs: {len(pairs)} (language_model: {n_lm}) | scale={lora_scale}")

    nullu_payload = None
    if args.nullu_model_path:
        n_total = model.config.text_config.num_hidden_layers
        if not (0 <= args.nullu_lowest_layer < args.nullu_highest_layer <= n_total):
            raise ValueError(
                f"need 0 <= nullu_lowest_layer < nullu_highest_layer <= {n_total}; "
                f"got {args.nullu_lowest_layer}-{args.nullu_highest_layer}"
            )
        nullu_idxs = list(range(args.nullu_lowest_layer, args.nullu_highest_layer))
        print(f"Loading Nullu full layer splice for layers "
              f"{nullu_idxs[0]}-{nullu_idxs[-1]} (inclusive) — "
              f"{len(_NULLU_LAYER_PARAM_NAMES)} params × {len(nullu_idxs)} layers")
        edited_cpu = _load_nullu_layer_weights(args.nullu_model_path, nullu_idxs)
        original_cpu = {
            L: {
                pname: _get_layer_param(model, L, pname).detach().cpu().clone()
                for pname in _NULLU_LAYER_PARAM_NAMES
            }
            for L in nullu_idxs
        }
        # Shape sanity-check against the actual model.
        for L in nullu_idxs:
            for pname in _NULLU_LAYER_PARAM_NAMES:
                want = original_cpu[L][pname].shape
                got = edited_cpu[L][pname].shape
                if want != got:
                    raise ValueError(
                        f"Nullu L{L} {pname}: edited shape {tuple(got)} != "
                        f"model shape {tuple(want)}"
                    )
        nullu_payload = {
            "edited_cpu": edited_cpu,
            "original_cpu": original_cpu,
            "layer_indices": nullu_idxs,
        }

    efuf_payload = None
    if args.efuf_ckpt:
        print(f"Loading EFUF mm_projector weights from {args.efuf_ckpt}")
        edited_proj = _load_efuf_projector_weights(args.efuf_ckpt)
        mmp = model.model.multi_modal_projector
        original_proj = {
            "linear_1.weight": mmp.linear_1.weight.detach().cpu().clone(),
            "linear_1.bias":   mmp.linear_1.bias.detach().cpu().clone(),
            "linear_2.weight": mmp.linear_2.weight.detach().cpu().clone(),
            "linear_2.bias":   mmp.linear_2.bias.detach().cpu().clone(),
        }
        # Sanity-check shapes.
        for k, t in edited_proj.items():
            if t.shape != original_proj[k].shape:
                raise ValueError(
                    f"EFUF {k} shape {tuple(t.shape)} does not match HF model "
                    f"shape {tuple(original_proj[k].shape)}"
                )
        efuf_payload = {"edited_cpu": edited_proj, "original_cpu": original_proj}

    print("Loading SAE …")
    sae = load_sae_model(args.sae_ckpt, model_type="llava", hook_type="text", device=args.device)

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

    # Per-image trace + accumulate population summary.
    per_image_scalars: Dict[str, list] = {k: [] for k in METHOD_ORDER}
    summary = []
    n_ok = n_skip = 0
    for s in tqdm(keep, desc="samples"):
        out, err = trace_one_image(
            sample=s, image_index=image_index, prompt=args.prompt,
            processor=processor, model=model, sae=sae,
            lora_pairs=pairs, lora_scale=lora_scale,
            nullu_payload=nullu_payload,
            efuf_payload=efuf_payload,
            selected=selected, layers=layers,
            hook_type=args.hook_type, gen_tokens=args.gen_tokens,
            sae_batch=args.sae_batch, device=args.device, dtype_attn=torch.long,
            fixed_assistant_prefix=args.fixed_assistant_prefix,
        )
        if out is None:
            n_skip += 1
            print(f"  skip {s['image_id']}: {err}")
            continue

        out_path = os.path.join(args.out_dir, f"{out['image_id']}.pt")
        torch.save(out, out_path)
        summary.append({"image_id": out["image_id"], "path": out_path})
        n_ok += 1

        try:
            png = plot_per_sample(out, layers, args.graph_dir, args.agg)
            print(f"  saved {png}")
        except Exception as e:
            print(f"  plot_per_sample failed for {out['image_id']}: {e}")
            traceback.print_exc()

        for key in METHOD_ORDER:
            if key in out["acts"]:
                per_image_scalars[key].append(per_layer_scalar(out["acts"][key], layers, args.agg))

    if n_ok > 0:
        try:
            png = plot_summary(per_image_scalars, layers, args.graph_dir, args.agg, n_ok)
            print(f"  saved {png}")
        except Exception as e:
            print(f"  plot_summary failed: {e}")
            traceback.print_exc()

    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()