"""
compare_activations.py — Compare SAE activation statistics between original
and modified LLaVA models on probe-selected features.

For each layer in --layers:
  1. Loads the linear probe(s) from a folder, selects top-k features by probe weight.
  2. Identifies sub features: all features NOT in the top-k set.
  3. Runs forward passes through BOTH models on the same inputs.
  4. Computes:
     - Per-feature activation statistics for top-k probe features
       (mean, std, mean |diff| between original and modified).
     - Aggregate |activation diff| across sub-threshold features.
     - Per-position JSD between original and modified activation distributions
       over sub-threshold features.

Data: HF "pbcong/toilet_bathroom" filtered by --filter_mode.
DDP: Supported via torchrun.

Usage
-----
  # Bathroom-only images, single GPU
  python training/compare_activations.py \\
      --original_model llava-hf/llava-1.5-7b-hf \\
      --modified_model /path/to/modified_model \\
      --sae_ckpt       training/multilayer_sae_ckpt/last.ckpt \\
      --probe_dir      training/multilayer_sae_ckpt \\
      --image_folder   /path/to/cc3m_images/train \\
      --filter_mode    bathroom \\
      --layers 0 1 2 3 4 5 6 \\
      --output_dir     outputs/compare_activations \\
      --device_id 0

  # Multi-GPU via torchrun
  torchrun --nproc_per_node=8 -m training.compare_activations \\
      --original_model llava-hf/llava-1.5-7b-hf \\
      --modified_model /path/to/modified_model \\
      --sae_ckpt       training/multilayer_sae_ckpt/last.ckpt \\
      --probe_dir      training/multilayer_sae_ckpt \\
      --image_folder   /path/to/cc3m_images/train \\
      --filter_mode    bathroom \\
      --layers 0 1 2 3 4 5 6 \\
      --output_dir     outputs/compare_activations
"""

import sys
import os
import json
import argparse
import datetime
from pathlib import Path
from typing import Dict, List, Tuple

import torch
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from PIL import Image
from tqdm import tqdm



sys.path.insert(0, str(Path(__file__).parent.parent))

from hallucination.extra_materials.mechanistic_interp.probe.probing import LinearProbe
from sae.SAE_Tools import load_sae_model
from sae.autoencoder.Utils import standardize as sae_standardize
from model.llava.hooked_llava import (
    HookedSAELlavaConditionalGeneration,
    load_nullu_model,
)
from model.llava.hooked_lora_llava import HookedLoRALlava
from sae.Training_Utils import str_to_torch_dtype
from transformers import LlavaProcessor


# ─────────────────────────────────────────────────────────────────────────────
# Distributed helpers
# ─────────────────────────────────────────────────────────────────────────────

def setup_distributed(timeout_hours: int = 8):
    if "RANK" in os.environ:
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])
        dist.init_process_group(
            backend="nccl",
            timeout=datetime.timedelta(hours=timeout_hours),
        )
        return rank, world_size, local_rank
    return 0, 1, 0


def cleanup_distributed():
    if dist.is_initialized():
        dist.destroy_process_group()


# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────

IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tiff"}


# ─────────────────────────────────────────────────────────────────────────────
# Dataset
# ─────────────────────────────────────────────────────────────────────────────

class SingleImageDataset(Dataset):
    """Single image for a quick one-off comparison pass."""

    def __init__(self, image_path: str):
        path = Path(image_path)
        if not path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")
        self.samples = [{"path": str(path), "imgid": path.stem, "caption": ""}]

    def __len__(self):
        return 1

    def __getitem__(self, idx):
        s = self.samples[idx]
        return {"image": Image.open(s["path"]).convert("RGB"),
                "imgid": s["imgid"], "caption": s["caption"]}


class FilteredBathroomToiletDataset(Dataset):
    """
    HF pbcong/bathroom-toilet filtered by bathroom/toilet labels.

    filter_mode:
      bathroom — bathroom == 1  (regardless of toilet value)
      toilet   — toilet   == 1  (regardless of bathroom value)
      both     — toilet   == 1 OR bathroom == 1
      all      — no filtering (every row with a matching local image)

    caption_mode:
      generated — caption=""; collate_fn uses prompt-only text.
      caption   — caption taken from HF "caption" field (CC3M caption).
    """

    def __init__(
        self,
        image_folder: str,
        filter_mode: str,
        hf_dataset: str = "pbcong/bathroom-toilet",
        caption_mode: str = "generated",
    ):
        from datasets import load_dataset as _load_dataset

        assert caption_mode in ("generated", "caption"), caption_mode
        ds = _load_dataset(hf_dataset, split="train+validation")

        stem_to_path: Dict[str, str] = {}
        for fname in os.listdir(image_folder):
            if Path(fname).suffix.lower() in IMG_EXTS:
                stem_to_path[fname.rsplit(".", 1)[0]] = os.path.join(
                    image_folder, fname
                )

        self.samples: List[dict] = []
        for row in ds:
            bathroom = row.get("bathroom", 0)
            toilet   = row.get("toilet",   0)
            if filter_mode == "bathroom" and bathroom != 1:
                continue
            elif filter_mode == "toilet" and toilet != 1:
                continue
            elif filter_mode == "both" and not (toilet == 1 or bathroom == 1):
                continue
            # "all": include every row

            img_id = str(row["image_id"])
            if img_id not in stem_to_path:
                continue

            caption = (row.get("caption", "") or "") if caption_mode == "caption" else ""
            self.samples.append(
                {"path": stem_to_path[img_id], "imgid": img_id, "caption": caption}
            )

        if not self.samples:
            raise RuntimeError(
                f"No images found for filter_mode={filter_mode!r} "
                f"in {image_folder!r}."
            )

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        s = self.samples[idx]
        return {
            "image":   Image.open(s["path"]).convert("RGB"),
            "imgid":   s["imgid"],
            "caption": s["caption"],
        }


# ─────────────────────────────────────────────────────────────────────────────
# Unified dataloader
# ─────────────────────────────────────────────────────────────────────────────

def create_dataloader(args, processor, rank: int = 0, world_size: int = 1):
    """
    Unified dataloader supporting three data modes.

    --data_mode:
        single  — one image at --image_path (caption_mode ignored).
        toilet  — pbcong/bathroom-toilet filtered by --filter_mode.
                  Requires --image_folder.
        cc3m    — full CC3M via --hf_dataset + --local_val_path.
        coco    — COCO via --hf_dataset + --local_val_path.

    --caption_mode:
        generated — original model generates the full sequence; both models
                    then run a teacher-forced forward pass on that sequence so
                    activations at every token position (prompt + answer) can
                    be compared.
        caption   — both models run a forward pass on prompt + stored caption
                    (CC3M txt / COCO sentences / pbcong/bathroom-toilet caption
                    field), capturing activations at every token position.
    """
    data_mode    = getattr(args, "data_mode",    "toilet")
    caption_mode = getattr(args, "caption_mode", "generated")

    if data_mode == "single":
        dataset = SingleImageDataset(args.image_path)

    elif data_mode == "toilet":
        dataset = FilteredBathroomToiletDataset(
            image_folder = args.image_folder,
            filter_mode  = getattr(args, "filter_mode", "bathroom"),
            hf_dataset   = getattr(args, "hf_dataset",  "pbcong/bathroom-toilet"),
            caption_mode = caption_mode,
        )

    elif data_mode in ("cc3m", "coco"):
        from sae.SAE_Trainer import DataConfig
        from sae.Load_Data import cc3m_dataset, coco_dataset
        data_cfg = DataConfig(
            hf_dataset       = args.hf_dataset,
            local_val_path   = args.local_val_path,
            local_train_path = args.local_val_path,
            processor        = args.original_model,
            batch_size       = args.batch_size,
            num_workers      = args.num_workers,
        )
        split   = getattr(args, "split", "train")
        dataset = (cc3m_dataset if data_mode == "cc3m" else coco_dataset)(data_cfg, split)

    else:
        raise ValueError(f"Unknown data_mode: {data_mode!r}")

    # single mode: dataset has exactly 1 sample. DistributedSampler with
    # drop_last=True would truncate floor(1/world_size)*world_size = 0 items,
    # leaving every rank with an empty loader. Instead let all ranks share the
    # same 1-sample loader; the reduce sums identical contributions so means
    # remain correct.
    sampler = (
        DistributedSampler(dataset, num_replicas=world_size, rank=rank,
                           shuffle=False, drop_last=True)
        if world_size > 1 and data_mode != "single" else None
    )

    prompt = "USER: <image>\nDescribe this image. \nASSISTANT:"

    def collate_fn(batch):
        batch = [b for b in batch if b is not None]
        if not batch:
            return None
        images = [b["image"] for b in batch]
        if caption_mode == "caption":
            texts = [
                (f"USER: <image>\nDescribe this image. \nASSISTANT: {b['caption']}"
                 if b.get("caption") else prompt)
                for b in batch
            ]
        else:
            # generated mode: prompt-only for the generation step
            texts = [prompt] * len(batch)
        processed = processor(
            images=images, text=texts, return_tensors="pt", padding=True
        )
        return {
            "input_ids":      processed["input_ids"],
            "attention_mask": processed["attention_mask"],
            "pixel_values":   processed["pixel_values"],
            "_images":        images,   # raw PIL images for re-processing after generation
        }

    dataloader = DataLoader(
        dataset,
        batch_size  = args.batch_size,
        shuffle     = False,
        sampler     = sampler,
        num_workers = args.num_workers,
        collate_fn  = collate_fn,
    )
    return dataset, dataloader


# ─────────────────────────────────────────────────────────────────────────────
# SAE activation extraction (memory-efficient targeted extraction)
# ─────────────────────────────────────────────────────────────────────────────

@torch.no_grad()
def get_targeted_sae_activations(sae, flat_acts, feature_ids, device, sae_batch=4096):
    """
    Extract pre-topk ReLU'd SAE activations for *specific* features only.

    Instead of computing the full (N, d_sae) tensor, this selects the relevant
    rows of the encoder weight matrix and computes only (N, len(feature_ids)).

    Returns: (N, n_features) tensor on CPU.
    """
    feature_ids_t = torch.tensor(feature_ids, dtype=torch.long)
    W_target = sae.encoder.weight[feature_ids_t].to(device)  # (n_features, d_in)
    bias = sae.pre_encoder_bias.to(device)  # (d_in,)

    all_acts = []
    for i in range(0, flat_acts.shape[0], sae_batch):
        batch = flat_acts[i : i + sae_batch].to(device)
        if sae.cfg.standardize:
            batch, _ = sae_standardize(batch)
        hidden = (batch - bias) @ W_target.T  # (chunk, n_features)
        all_acts.append(torch.relu(hidden).cpu())
    return torch.cat(all_acts, dim=0)


# ─────────────────────────────────────────────────────────────────────────────
# JSD
# ─────────────────────────────────────────────────────────────────────────────

def jensen_shannon_divergence(p, q, eps=1e-10):
    """
    Per-row Jensen-Shannon divergence.

    p, q: (N, D) non-negative tensors. Each row is normalised to a
    probability distribution before computing JSD.

    Returns: (N,) tensor of JSD values in [0, ln2].
    """
    p = p.float() + eps
    q = q.float() + eps
    p = p / p.sum(dim=-1, keepdim=True)
    q = q / q.sum(dim=-1, keepdim=True)
    m = (p + q) / 2
    kl_pm = (p * (p / m).log()).sum(dim=-1)
    kl_qm = (q * (q / m).log()).sum(dim=-1)
    return (kl_pm + kl_qm) / 2



# ─────────────────────────────────────────────────────────────────────────────
# KL divergence 
# ─────────────────────────────────────────────────────────────────────────────

def kl_divergence(p, q, eps=1e-10):
    """
    Per-row Kullback-Leibler divergence.

    p, q: (N, D) non-negative tensors. Each row is normalised to a
    probability distribution before computing KL divergence.

    Returns: (N,) tensor of KL divergence values.
    """
    p = p.float() + eps
    q = q.float() + eps
    p = p / p.sum(dim=-1, keepdim=True)
    q = q / q.sum(dim=-1, keepdim=True)
    kl = (p * (p / q).log()).sum(dim=-1)
    return kl


# ─────────────────────────────────────────────────────────────────────────────
# Print helpers
# ─────────────────────────────────────────────────────────────────────────────

_ARROW_EPS = 1e-6


def _dir_arrow(new_val: float, ref_val: float) -> str:
    """Direction of change relative to ref_val.

    ↓  suppressed (new < ref)
    ↑  enhanced   (new > ref)
    =  unchanged
    —  both values are effectively zero
    """
    if ref_val < _ARROW_EPS and new_val < _ARROW_EPS:
        return "—"
    if new_val < ref_val - _ARROW_EPS:
        return "↓"
    if new_val > ref_val + _ARROW_EPS:
        return "↑"
    return "="


def _build_answer_mask(input_ids, attention_mask, L_prompt, pad_side):
    """Per-position bool mask that is True only on answer tokens — i.e.,
    strictly after the 'USER: <image> … ASSISTANT:' prompt prefix — and
    False on image + prompt tokens and on padding.

    Same prompt template is assumed for every sample (constant L_prompt).
    Supports both right- and left-padded sequences.
    """
    B, T = input_ids.shape
    attn = attention_mask.detach().cpu()
    mask = torch.zeros(B, T, dtype=torch.bool)
    for i in range(B):
        R_i = int(attn[i].sum())
        L_ans = max(0, R_i - L_prompt)
        if L_ans == 0:
            continue
        if pad_side == "left":
            mask[i, T - L_ans : T] = True
        else:  # right pad (default)
            mask[i, L_prompt : L_prompt + L_ans] = True
    return mask


# ─────────────────────────────────────────────────────────────────────────────
# Probe loading
# ─────────────────────────────────────────────────────────────────────────────

def load_probe_features(
    probe_dir: str,
    layer: int,
    top_k: int,
    input_dim: int = 65536,
) -> Tuple[List[int], List[float], List[int]]:
    """
    Load a linear probe and partition features into top-k and the rest.

    Returns:
        topk_ids       — feature indices with the k highest probe weights
        topk_weights   — corresponding probe weight values
        sub_ids        — all remaining feature indices not in top-k
    """
    ckp_path = (
        Path(probe_dir)
        / f"probe_model.language_model.layers.{layer}.hook_resid_post.pt"
    )
    if not ckp_path.exists():
        raise FileNotFoundError(f"Probe checkpoint not found: {ckp_path}")

    ckpt = torch.load(ckp_path, map_location="cpu")
    probe = LinearProbe(input_dim=input_dim, num_outputs=1)
    probe.load_state_dict(ckpt)
    probe.eval()

    weights = probe.weights.squeeze()  # (input_dim,)

    top = torch.topk(weights, k=top_k)
    topk_ids = top.indices.cpu().tolist()
    topk_weights = weights[top.indices].cpu().tolist()

    topk_set = set(topk_ids)
    sub_ids = [i for i in range(input_dim) if i not in topk_set]

    return topk_ids, topk_weights, sub_ids


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

def main():
    ap = argparse.ArgumentParser(
        description="Compare SAE activations between original and modified LLaVA.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # ── Models ───────────────────────────────────────────────────────────────
    ap.add_argument("--original_model", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--modified_model", required=True,
                    help="HF model name or local path for the modified (LoRA) model.")
    ap.add_argument("--nullu_model_path", default=None,
                    help="Path to Nullu's edited checkpoint directory "
                         "(e.g. Nullu/output/edited_model/LLaVA-7B-top4-0-32-test). "
                         "If provided, a third Nullu-edited model is added to the comparison.")
    ap.add_argument("--nullu_lowest_layer", type=int, default=16,
                    help="Inclusive lower bound of Nullu edited layer range.")
    ap.add_argument("--nullu_highest_layer", type=int, default=32,
                    help="Exclusive upper bound of Nullu edited layer range.")
    ap.add_argument("--sae_ckpt", required=True)
    ap.add_argument("--device_id", type=int, default=0)
    ap.add_argument("--dtype", default="float16")

    # ── Probe ────────────────────────────────────────────────────────────────
    ap.add_argument("--probe_dir", required=True,
                    help="Dir with probe_model.language_model.layers.*.pt files.")
    ap.add_argument("--probe_input_dim", type=int, default=65536)
    ap.add_argument("--layers", type=int, nargs="+", default=list(range(7)))
    ap.add_argument("--top_probe_k", type=int, default=1000,
                    help="Number of top probe features to track (by probe weight). "
                         "All are saved to JSON; only --top_print_k are printed.")
    ap.add_argument("--top_print_k", type=int, default=20,
                    help="Number of top features to print per-row in the log. "
                         "The rest are still computed and saved to JSON.")

    # ── Data mode ────────────────────────────────────────────────────────────
    ap.add_argument(
        "--data_mode", default="toilet", choices=["single", "toilet", "cc3m", "coco"],
        help=(
            "single: one image at --image_path. "
            "toilet: pbcong/bathroom-toilet filtered by --filter_mode (needs --image_folder). "
            "cc3m:   CC3M via --hf_dataset + --local_val_path. "
            "coco:   COCO via --hf_dataset + --local_val_path."
        ),
    )

    # ── Data — single image ──────────────────────────────────────────────────
    ap.add_argument("--image_path", default=None,
                    help="[single] Path to a single image file.")

    # ── Data — toilet ────────────────────────────────────────────────────────
    ap.add_argument("--image_folder", default=None,
                    help="[toilet] Local folder with CC3M (or other) images.")
    ap.add_argument("--filter_mode", default="bathroom",
                    choices=["bathroom", "toilet", "both", "all"],
                    help="[toilet] bathroom: bath==1; toilet: toilet==1; "
                         "both: toilet==1 OR bath==1; all: no filtering.")
    ap.add_argument("--hf_dataset", default="pbcong/bathroom-toilet",
                    help="[toilet] HF dataset name, or CC3M/COCO path for cc3m/coco mode.")

    # ── Data — CC3M / COCO ───────────────────────────────────────────────────
    ap.add_argument("--local_val_path", default=None,
                    help="[cc3m/coco] Local image root for the HF dataset split.")
    ap.add_argument("--split", default="train",
                    help="[cc3m/coco] HF dataset split.")

    # ── Caption mode ─────────────────────────────────────────────────────────
    ap.add_argument(
        "--caption_mode", default="generated", choices=["generated", "caption"],
        help=(
            "generated: original model generates full sequence; both models then "
            "run a teacher-forced forward pass on it (activations at every token). "
            "caption:   both models run a forward pass on prompt + stored caption "
            "(activations at every token)."
        ),
    )
    ap.add_argument("--max_new_tokens", type=int, default=256,
                    help="Max tokens to generate per image (only with --caption_mode generated).")

    # ── Distribution measure ──────────────────────────────────────────────────
    ap.add_argument(
        "--measure_distribution",
        default="jensen_shannon_divergence",
        choices=["kl_divergence", "jensen_shannon_divergence"],
        help="Divergence measure for sub-threshold feature distributions.",
    )

    # ── Diff norm ─────────────────────────────────────────────────────────────
    ap.add_argument(
        "--diff_norm",
        default="l1",
        choices=["l1", "l2"],
        help=(
            "Norm used for the Diff column and sub-threshold aggregate diff. "
            "l1: per-feature mean|Δ| (top-k) / per-position 1/N Σ‖Δ‖₁ (sub-threshold). "
            "l2: per-feature RMSE (top-k) / per-position 1/N Σ‖Δ‖₂ (sub-threshold)."
        ),
    )

    # ── Processing ───────────────────────────────────────────────────────────
    ap.add_argument("--batch_size", type=int, default=4)
    ap.add_argument("--sae_batch", type=int, default=4096)
    ap.add_argument("--num_workers", type=int, default=4)
    ap.add_argument("--max_batches", type=int, default=None)

    # ── Output ───────────────────────────────────────────────────────────────
    ap.add_argument("--output_dir", default="outputs/compare_activations")
    ap.add_argument("--dist_timeout_hours", type=int, default=8,
                    help="NCCL/store timeout in hours for distributed runs.")

    args = ap.parse_args()

    if args.data_mode == "single" and not args.image_path:
        ap.error("--data_mode single requires --image_path.")
    if args.data_mode == "toilet" and not args.image_folder:
        ap.error("--data_mode toilet requires --image_folder.")
    if args.data_mode in ("cc3m", "coco") and not args.hf_dataset:
        ap.error(f"--data_mode {args.data_mode} requires --hf_dataset (+ --local_val_path).")

    # ── Distributed setup ────────────────────────────────────────────────────
    rank, world_size, local_rank = setup_distributed(args.dist_timeout_hours)
    is_distributed = world_size > 1

    device = torch.device(
        f"cuda:{local_rank}"
        if is_distributed
        else f"cuda:{args.device_id}" if torch.cuda.is_available() else "cpu"
    )
    dtype = str_to_torch_dtype(args.dtype)

    if rank == 0:
        print(f"Device: {device}  |  world_size: {world_size}")
        print(f"Original model: {args.original_model}")
        print(f"Modified model: {args.modified_model}")
        if args.nullu_model_path:
            print(
                f"Nullu model:    {args.nullu_model_path}  "
                f"(layers {args.nullu_lowest_layer}-{args.nullu_highest_layer})"
            )

    # ── Load SAE ─────────────────────────────────────────────────────────────
    sae = load_sae_model(
        args.sae_ckpt, model_type="llava", hook_type="text", device=device
    )
    sae.eval()
    sae_dtype = next(sae.parameters()).dtype

    # ── Load both LLaVA models ───────────────────────────────────────────────
    if rank == 0:
        print("Loading original model …")
    model_orig = HookedSAELlavaConditionalGeneration.from_pretrained(
        args.original_model
    )
    model_orig.to(device, dtype=dtype).eval()

    if rank == 0:
        print("Loading modified model …")
        
    model_mod = HookedLoRALlava.from_pretrained(args.original_model)
    model_mod.load_lora_adapter(adapter_path=args.modified_model, merge=True)
    model_mod.to(device, dtype=dtype).eval()

    model_nullu = None
    if args.nullu_model_path:
        if rank == 0:
            print(
                f"Loading nullu model … (layers "
                f"{args.nullu_lowest_layer}-{args.nullu_highest_layer})"
            )
        model_nullu = load_nullu_model(
            lowest_layer=args.nullu_lowest_layer,
            highest_layer=args.nullu_highest_layer,
            edited_model_path=args.nullu_model_path,
            base_model_name=args.original_model,
            torch_dtype=dtype,
            device=device,
        )
        model_nullu.eval()

    processor = LlavaProcessor.from_pretrained(args.original_model)

    # Padding side + prompt template — used to build an answer-only token mask
    # so all stats are accumulated strictly over generated/caption answer
    # tokens (image + prompt tokens and padding are excluded).
    pad_side = getattr(processor.tokenizer, "padding_side", "right")
    prompt_template = "USER: <image>\nDescribe this image. \nASSISTANT:"
    L_prompt: int = None  # lazily computed on the first batch

    # ── Dataset & DataLoader ─────────────────────────────────────────────────
    dataset, dataloader = create_dataloader(args, processor, rank, world_size)

    if rank == 0:
        print(
            f"Dataset: {len(dataset)} images, "
            f"{len(dataloader)} batches/rank, "
            f"data_mode={args.data_mode}"
        )

    # ── Pre-load all probes & init accumulators ──────────────────────────────
    all_results: Dict[str, dict] = {}
    layer_meta   = {}   # layer -> {topk_ids, topk_wts, sub_ids, all_ids, n_topk, n_sub, hook_point}
    accumulators = {}   # layer -> dict of running tensors

    for layer in args.layers:
        hook_point = f"model.language_model.layers.{layer}.hook_resid_post"
        topk_ids, topk_wts, sub_ids = load_probe_features(
            args.probe_dir,
            layer,
            args.top_probe_k,
            args.probe_input_dim,
        )
        all_ids = topk_ids + sub_ids
        n_topk  = len(topk_ids)
        n_sub   = len(sub_ids)

        layer_meta[layer] = dict(
            hook_point=hook_point,
            topk_ids=topk_ids, topk_wts=topk_wts,
            sub_ids=sub_ids, all_ids=all_ids,
            n_topk=n_topk, n_sub=n_sub,
        )
        accumulators[layer] = dict(
            topk_orig_sum  = torch.zeros(n_topk, dtype=torch.float64),
            topk_orig_sq   = torch.zeros(n_topk, dtype=torch.float64),
            topk_mod_sum   = torch.zeros(n_topk, dtype=torch.float64),
            topk_mod_sq    = torch.zeros(n_topk, dtype=torch.float64),
            sub_abs_diff_sum = torch.zeros(1, dtype=torch.float64),
            sub_div_sum    = torch.zeros(1, dtype=torch.float64),
            sub_div_sq_sum = torch.zeros(1, dtype=torch.float64),
            total_positions  = torch.zeros(1, dtype=torch.long),
        )
        if model_nullu is not None:
            accumulators[layer].update(
                topk_nullu_sum     = torch.zeros(n_topk, dtype=torch.float64),
                topk_nullu_sq      = torch.zeros(n_topk, dtype=torch.float64),
                sub_nullu_abs_diff_sum = torch.zeros(1, dtype=torch.float64),
                sub_nullu_div_sum      = torch.zeros(1, dtype=torch.float64),
                sub_nullu_div_sq_sum   = torch.zeros(1, dtype=torch.float64),
            )

        if rank == 0:
            print(f"\n{'='*60}")
            print(f"[Layer {layer}]  hook = {hook_point}")
            print(f"  top-{args.top_probe_k} features (printing top-{args.top_print_k}): {topk_ids[:args.top_print_k]}")
            print(f"  sub features (all except top-{args.top_probe_k}): {n_sub}")

    # Hook pattern that captures all target layer residual streams in one pass
    target_hook_points = {m["hook_point"] for m in layer_meta.values()}

    def make_hook_fn(cache: dict):
        def hook_fn(act, hook):
            if hook.name in target_hook_points:
                cache[hook.name] = act.detach().cpu()
        return hook_fn

    # ── Single-pass batch loop ────────────────────────────────────────────────
    n_batches = (
        len(dataloader)
        if args.max_batches is None
        else min(args.max_batches, len(dataloader))
    )

    pbar = tqdm(dataloader, total=n_batches, disable=(rank != 0),
                desc="Batches")
    for batch_idx, batch in enumerate(pbar):
        if args.max_batches is not None and batch_idx >= args.max_batches:
            break

        if batch is None:
            continue

        raw_images = batch.pop("_images")   # List[PIL.Image]
        model_inputs = {
            k: batch[k].to(device)
            for k in ("input_ids", "attention_mask")
        }
        model_inputs["pixel_values"] = batch["pixel_values"].to(
            device, dtype=dtype
        )

        # ── Generated mode: generate ONCE per batch ──────────────────────────
        # Use the original model to generate the full sequence, then
        # teacher-force BOTH models with that sequence. Stats are then
        # accumulated only on the answer tokens (post-"ASSISTANT:");
        # image + prompt tokens are masked out below.
        if args.caption_mode == "generated":
            with torch.no_grad():
                gen_ids = model_orig.generate(
                    **model_inputs,
                    do_sample=False,
                    num_beams=1,
                    use_cache=True,
                    max_new_tokens=args.max_new_tokens,
                )
            full_texts = processor.batch_decode(gen_ids, skip_special_tokens=True)
            captions = [
                txt.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in txt
                else txt.strip()
                for txt in full_texts
            ]
            forced_texts = [
                f"USER: <image>\nDescribe this image. \nASSISTANT: {cap}"
                for cap in captions
            ]
            re_proc = processor(
                images=raw_images, text=forced_texts,
                return_tensors="pt", padding=True,
            )
            model_inputs = {
                "input_ids":      re_proc["input_ids"].to(device),
                "attention_mask": re_proc["attention_mask"].to(device),
                "pixel_values":   re_proc["pixel_values"].to(device, dtype=dtype),
            }

        # ── One forward pass per model, all target layers cached at once ─────
        cache_orig: dict = {}
        cache_mod:  dict = {}
        cache_nullu: dict = {}
        with torch.no_grad():
            model_orig.run_with_hooks(
                model_inputs,
                fwd_hooks=[(lambda n: n in target_hook_points, make_hook_fn(cache_orig))],
            )
            model_mod.run_with_hooks(
                model_inputs,
                fwd_hooks=[(lambda n: n in target_hook_points, make_hook_fn(cache_mod))],
            )
            if model_nullu is not None:
                model_nullu.run_with_hooks(
                    model_inputs,
                    fwd_hooks=[(lambda n: n in target_hook_points, make_hook_fn(cache_nullu))],
                )

        # Answer-only mask: True at positions for the answer tokens
        # (post-"ASSISTANT:") and False on image + prompt tokens and padding.
        # Stats below are accumulated only over these positions.
        if L_prompt is None:
            _pp = processor(
                images=[raw_images[0]], text=[prompt_template],
                return_tensors="pt", padding=False,
            )
            L_prompt = int(_pp["input_ids"].shape[1])

        answer_mask_2d = _build_answer_mask(
            model_inputs["input_ids"],
            model_inputs["attention_mask"],
            L_prompt,
            pad_side,
        )
        attn_mask = answer_mask_2d.view(-1)  # (B*T,) on CPU
        n_real = attn_mask.sum().long()

        # ── Accumulate stats per layer from cached activations ────────────────
        for layer in args.layers:
            meta = layer_meta[layer]
            acc  = accumulators[layer]
            hp   = meta["hook_point"]

            if hp not in cache_orig or hp not in cache_mod:
                continue   # hook didn't fire (shouldn't happen)

            acts_orig = cache_orig[hp]   # (B, T, D) on CPU
            acts_mod  = cache_mod[hp]
            B, T, D = acts_orig.shape

            # Mask out padding positions
            flat_orig = acts_orig.reshape(B * T, D)[attn_mask].to(sae_dtype)
            flat_mod  = acts_mod.reshape(B * T, D)[attn_mask].to(sae_dtype)

            all_ids = meta["all_ids"]
            n_topk  = meta["n_topk"]
            n_sub   = meta["n_sub"]

            if not all_ids:
                acc["total_positions"] += n_real
                continue

            sae_orig = get_targeted_sae_activations(
                sae, flat_orig, all_ids, device, args.sae_batch
            )
            sae_mod = get_targeted_sae_activations(
                sae, flat_mod, all_ids, device, args.sae_batch
            )
            sae_nullu = None
            if model_nullu is not None and hp in cache_nullu:
                flat_nullu = cache_nullu[hp].reshape(B * T, D)[attn_mask].to(sae_dtype)
                sae_nullu = get_targeted_sae_activations(
                    sae, flat_nullu, all_ids, device, args.sae_batch
                )

            # Top-k feature stats
            orig_topk = sae_orig[:, :n_topk].double()
            mod_topk  = sae_mod[:, :n_topk].double()

            acc["topk_orig_sum"] += orig_topk.sum(dim=0)
            acc["topk_orig_sq"]  += (orig_topk**2).sum(dim=0)
            acc["topk_mod_sum"]  += mod_topk.sum(dim=0)
            acc["topk_mod_sq"]   += (mod_topk**2).sum(dim=0)

            if sae_nullu is not None:
                nullu_topk = sae_nullu[:, :n_topk].double()
                acc["topk_nullu_sum"] += nullu_topk.sum(dim=0)
                acc["topk_nullu_sq"]  += (nullu_topk**2).sum(dim=0)

            # Sub-threshold stats
            if n_sub > 0:
                orig_sub = sae_orig[:, n_topk:]
                mod_sub  = sae_mod[:, n_topk:]

                diff_sub = orig_sub - mod_sub
                if args.diff_norm == "l1":
                    # L1: sum of per-position L1 norms → 1/N Σ_n Σ_f |a_nf − a'_nf|
                    acc["sub_abs_diff_sum"] += (
                        diff_sub.abs().sum(dim=-1).sum().double()
                    )
                else:
                    # L2: sum of per-position L2 norms → 1/N Σ_n ‖a_n − a'_n‖₂
                    acc["sub_abs_diff_sum"] += (
                        (diff_sub ** 2).sum(dim=-1).sqrt().sum().double()
                    )
                div_fn = (
                    jensen_shannon_divergence
                    if args.measure_distribution == "jensen_shannon_divergence"
                    else kl_divergence
                )
                div_vals = div_fn(orig_sub, mod_sub)
                acc["sub_div_sum"]    += div_vals.sum().double()
                acc["sub_div_sq_sum"] += (div_vals**2).sum().double()

                if sae_nullu is not None:
                    nullu_sub = sae_nullu[:, n_topk:]
                    diff_nullu = orig_sub - nullu_sub
                    if args.diff_norm == "l1":
                        acc["sub_nullu_abs_diff_sum"] += (
                            diff_nullu.abs().sum(dim=-1).sum().double()
                        )
                    else:
                        acc["sub_nullu_abs_diff_sum"] += (
                            (diff_nullu ** 2).sum(dim=-1).sqrt().sum().double()
                        )
                    div_vals_nullu = div_fn(orig_sub, nullu_sub)
                    acc["sub_nullu_div_sum"]    += div_vals_nullu.sum().double()
                    acc["sub_nullu_div_sq_sum"] += (div_vals_nullu**2).sum().double()

            acc["total_positions"] += n_real

    # ── DDP reduce + aggregate per layer ─────────────────────────────────────
    for layer in args.layers:
        meta = layer_meta[layer]
        acc  = accumulators[layer]

        if is_distributed:
            for key in acc:
                t = acc[key].to(device)
                dist.reduce(t, dst=0)
                acc[key] = t.cpu()

        if rank == 0:
            N = acc["total_positions"].item()
            if N == 0:
                print(f"  WARNING: layer {layer} — no positions processed.")
                continue

            hook_point = meta["hook_point"]
            topk_ids   = meta["topk_ids"]
            topk_wts   = meta["topk_wts"]
            n_topk     = meta["n_topk"]
            n_sub      = meta["n_sub"]

            layer_result = {
                "hook_point": hook_point,
                "total_positions": N,
                "top_features": {},
                "sub_threshold": {},
            }

            print(f"\n{'='*60}")
            print(f"[Layer {layer}]  hook = {hook_point}")
            has_nullu = model_nullu is not None
            if has_nullu:
                print(
                    f"\n  {'Feature':>10}  {'ProbeW':>8}  "
                    f"{'OrigMean':>10}  {'ModMean':>11}  {'Δ Mod':>10}  "
                    f"{'NulluMean':>12}  {'Δ Nullu':>10}"
                )
                print(
                    f"  {'-'*10}  {'-'*8}  {'-'*10}  {'-'*11}  {'-'*10}  "
                    f"{'-'*12}  {'-'*10}"
                )
            else:
                print(f"\n  {'Feature':>10}  {'ProbeW':>8}  "
                      f"{'OrigMean':>10}  {'ModMean':>11}  {'Δ Mean':>10}")
                print(f"  {'-'*10}  {'-'*8}  {'-'*10}  {'-'*11}  {'-'*10}")

            for j, (fid, pw) in enumerate(zip(topk_ids, topk_wts)):
                o_mean = acc["topk_orig_sum"][j].item() / N
                m_mean = acc["topk_mod_sum"][j].item() / N
                o_var  = max(0, acc["topk_orig_sq"][j].item() / N - o_mean**2)
                m_var  = max(0, acc["topk_mod_sq"][j].item() / N - m_mean**2)
                delta  = o_mean - m_mean

                m_arrow = _dir_arrow(m_mean, o_mean)

                feat_entry = {
                    "rank": j + 1,
                    "probe_weight": pw,
                    "original_mean": o_mean,
                    "original_std": o_var**0.5,
                    "modified_mean": m_mean,
                    "modified_std": m_var**0.5,
                    "delta_mean": delta,
                    "mod_direction": m_arrow,
                }

                if has_nullu:
                    n_mean  = acc["topk_nullu_sum"][j].item() / N
                    n_var   = max(0, acc["topk_nullu_sq"][j].item() / N - n_mean**2)
                    delta_n = o_mean - n_mean
                    n_arrow = _dir_arrow(n_mean, o_mean)
                    feat_entry.update({
                        "nullu_mean": n_mean,
                        "nullu_std": n_var**0.5,
                        "delta_nullu": delta_n,
                        "nullu_direction": n_arrow,
                    })

                layer_result["top_features"][str(fid)] = feat_entry

                if j < args.top_print_k:
                    if has_nullu:
                        print(
                            f"  {fid:>10}  {pw:>+8.4f}  "
                            f"{o_mean:>10.4f}  {m_mean:>9.4f} {m_arrow}  {delta:>+10.4f}  "
                            f"{n_mean:>10.4f} {n_arrow}  {delta_n:>+10.4f}"
                        )
                    else:
                        print(
                            f"  {fid:>10}  {pw:>+8.4f}  "
                            f"{o_mean:>10.4f}  {m_mean:>9.4f} {m_arrow}  {delta:>+10.4f}"
                        )
                elif j == args.top_print_k:
                    print(f"  ... ({n_topk - args.top_print_k} more features saved to JSON)")

            if n_sub > 0:
                mean_div      = acc["sub_div_sum"].item() / N
                div_var       = max(0, acc["sub_div_sq_sum"].item() / N - mean_div**2)
                denom = N * n_sub
                mean_abs_diff = acc["sub_abs_diff_sum"].item() / denom
                measure_name  = args.measure_distribution

                diff_label = "mean_l1_per_feature" if args.diff_norm == "l1" else "mean_l2_per_feature"
                layer_result["sub_threshold"] = {
                    "n_features": n_sub,
                    diff_label: mean_abs_diff,
                    "diff_norm": args.diff_norm,
                    "divergence_measure": measure_name,
                    "mean_divergence": mean_div,
                    "std_divergence": div_var**0.5,
                }
                diff_print_label = (
                    "1/(N·F) Σ|Δ| (L1)" if args.diff_norm == "l1"
                    else "1/(N·√F) Σ‖Δ‖₂ (L2)"
                )
                print(
                    f"\n  Sub features ({n_sub} features, "
                    f"all except top-{args.top_probe_k}):"
                )
                print(f"    [Mod vs Orig] {diff_print_label}: {mean_abs_diff:.6f}")
                print(f"    [Mod vs Orig] mean {measure_name}: {mean_div:.6f} "
                      f"± {div_var**0.5:.6f}")

                if has_nullu:
                    mean_div_n = acc["sub_nullu_div_sum"].item() / N
                    div_var_n  = max(0, acc["sub_nullu_div_sq_sum"].item() / N - mean_div_n**2)
                    mean_abs_diff_n = acc["sub_nullu_abs_diff_sum"].item() / denom
                    layer_result["sub_threshold"].update({
                        f"nullu_{diff_label}": mean_abs_diff_n,
                        "nullu_mean_divergence": mean_div_n,
                        "nullu_std_divergence": div_var_n**0.5,
                    })
                    print(f"    [Nullu vs Orig] {diff_print_label}: {mean_abs_diff_n:.6f}")
                    print(f"    [Nullu vs Orig] mean {measure_name}: {mean_div_n:.6f} "
                          f"± {div_var_n**0.5:.6f}")

            all_results[f"layer_{layer}"] = layer_result

    # ── Save results ─────────────────────────────────────────────────────────
    if rank == 0:
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        result = {
            "config": {
                "original_model":  args.original_model,
                "modified_model":  args.modified_model,
                "nullu_model_path": args.nullu_model_path,
                "nullu_lowest_layer":  args.nullu_lowest_layer if args.nullu_model_path else None,
                "nullu_highest_layer": args.nullu_highest_layer if args.nullu_model_path else None,
                "sae_ckpt":        args.sae_ckpt,
                "probe_dir":       args.probe_dir,
                "layers":          args.layers,
                "top_probe_k":     args.top_probe_k,
                "data_mode":       args.data_mode,
                "filter_mode":     getattr(args, "filter_mode", None),
                "hf_dataset":      args.hf_dataset,
                "caption_mode":    args.caption_mode,
                "measure_distribution": args.measure_distribution,
                "diff_norm":       args.diff_norm,
                "n_images":        len(dataset),
            },
            "layers": all_results,
        }

        out_path = output_dir / "comparison.json"
        with open(out_path, "w") as f:
            json.dump(result, f, indent=2)
        print(f"\nResults saved to {out_path}")

    cleanup_distributed()


if __name__ == "__main__":
    main()