"""
visualize_multilayer_sae_features.py — Single-pass SAE feature visualizer for LLaVA
                               scoped to toilet / bathroom images only.

Loads only images whose IDs appear in the HuggingFace dataset
"pbcong/bathroom-toilet" and satisfy the requested mode:
  --object_mode toilet   → rows where toilet   == 1
  --object_mode bathroom → rows where bathroom == 1

Caption modes:
  --caption_mode generated  (default) — model is prompted with
      "USER: <image>\\nDescribe this image. \\nASSISTANT:"
      and the residual stream is captured before generation (prefill only).
  --caption_mode dataset — the caption stored in the HF row (field "caption"
      or "txt") is appended, matching the cc3m_dataset convention.

Usage
-----
  python training/visualize_multilayer_sae_features.py.py \\
      --sae_ckpt  training/multilayer_sae_ckpt/last.ckpt \\
      --image_folder  /path/to/cc3m_images/train \\
      --object_mode   toilet \\
      --caption_mode  generated \\
      --feature_ids   0 1 42 \\
      --hook_point    model.language_model.layers.19.hook_resid_post \\
      --output_dir    visualize/multilayer_sae_features.py \\
      --device_id 0
"""

import sys
import os
import re
import io
import json
import html as html_lib
import base64
import argparse
import heapq
import pickle
import shutil
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

import torch as t
import torch.distributed as dist
import torch.nn.functional as F
import numpy as np
from PIL import Image, ImageDraw
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from transformers import LlavaProcessor
from transformer_lens.hook_points import HookPoint

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

from sae.SAE_Tools import load_sae_model, get_sae_activations, get_loader
from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from sae.Training_Utils import str_to_torch_dtype

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────
IMG_EXTS = {".jpg", ".jpeg", ".png"}
IMAGE_TOKEN_ID = 32000          # LLaVA-1.5 <image> token
N_IMAGE_PATCHES = 576           # CLIP ViT-L/14 @ 336px → 24 × 24
PATCH_GRID = 24


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

def setup_distributed():
    """Initialize distributed process group. Returns (rank, world_size, local_rank)."""
    if "RANK" in os.environ and "WORLD_SIZE" 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")
        return rank, world_size, local_rank
    return 0, 1, 0


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


# ─────────────────────────────────────────────────────────────────────────────
# Data structures for top-k tracking
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class ImageRecord:
    """A top-activating image patch."""
    value: float
    sample_idx: int
    patch_idx: int          # 0–575
    hook_name: str


@dataclass
class TextRecord:
    """A top-activating text token with surrounding context."""
    value: float
    sample_idx: int
    hook_name: str
    context_ids: List[int]      # token IDs in context window
    context_acts: List[float]   # feature activations for context tokens
    target_relative: int        # index of activated token within context


class TopKHeap:
    """Min-heap that retains only the k largest entries."""

    def __init__(self, k: int):
        self.k = k
        self._heap: List[Tuple[float, int, object]] = []
        self._counter = 0

    def push(self, value: float, record):
        self._counter += 1
        entry = (value, self._counter, record)
        if len(self._heap) < self.k:
            heapq.heappush(self._heap, entry)
        elif value > self._heap[0][0]:
            heapq.heapreplace(self._heap, entry)

    def sorted_records(self) -> list:
        return [r for _, _, r in sorted(self._heap, reverse=True)]

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


# ─────────────────────────────────────────────────────────────────────────────
# Data loading
# ─────────────────────────────────────────────────────────────────────────────

class ImageFolderDataset(Dataset):
    """Simple dataset from a folder of images (no captions)."""

    def __init__(self, data_dir: str):
        self.paths = sorted(
            p for p in Path(data_dir).rglob("*") if p.suffix.lower() in IMG_EXTS
        )
        if not self.paths:
            raise FileNotFoundError(f"No images found under: {data_dir}")

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

    def __getitem__(self, idx):
        p = self.paths[idx]
        return {"image": Image.open(p).convert("RGB"), "imgid": p.stem, "caption": ""}


class ToiletWithNegativesDataset(Dataset):
    """
    Loads pbcong/bathroom-toilet positives (toilet_mode==1) from a local image
    folder, plus num_negatives random CC3M images from the same folder that are
    NOT in the HF dataset.

    caption_mode="caption":
        positives  → "caption" field from pbcong/bathroom-toilet (the CC3M caption
                     stored in the HF row for that image).
        negatives  → CC3M caption looked up via the CC3M HF dataset using __key__;
                     falls back to "" if cc3m_hf is not provided or key absent.
    caption_mode="generated":
        All captions are ""; single_pass generates them on-the-fly.
    """

    def __init__(
        self,
        image_folder: str,
        object_mode: str = "toilet",
        caption_mode: str = "generated",
        num_negatives: int = 10000,
        cc3m_hf: Optional[str] = None,
        cc3m_split: str = "train",
    ):
        import random
        from datasets import load_dataset as _lds

        assert object_mode  in ("toilet", "bathroom", "both"), object_mode
        assert caption_mode in ("caption", "generated"), caption_mode

        # ── Load pbcong/bathroom-toilet ──────────────────────────────────────
        bt_ds = _lds("pbcong/bathroom-toilet", split="train+validation")

        bt_labels:   Dict[str, int] = {}
        bt_captions: Dict[str, str] = {}
        for row in bt_ds:
            imgid = str(row["image_id"])
            if object_mode == "both":
                bt_labels[imgid] = 1 if (row.get("toilet", 0) == 1 or row.get("bathroom", 0) == 1) else 0
            else:
                bt_labels[imgid] = row.get(object_mode, 0)
            bt_captions[imgid] = row.get("caption", "") or ""
        bt_ids = set(bt_labels.keys())

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

        # ── Positives ────────────────────────────────────────────────────────
        positive_ids = [
            imgid for imgid in bt_ids
            if bt_labels.get(imgid) == 1 and imgid in stem_to_path
        ]

        # ── Negatives ────────────────────────────────────────────────────────
        non_bt_stems = [s for s in stem_to_path if s not in bt_ids]
        n_neg = min(num_negatives, len(non_bt_stems))
        negative_stems = random.sample(non_bt_stems, n_neg)

        # ── Caption lookup for negatives (CC3M HF) ───────────────────────────
        neg_captions: Dict[str, str] = {}
        if caption_mode == "caption" and cc3m_hf and "cc3m" in cc3m_hf.lower():
            neg_stems_set = set(negative_stems)
            cc3m_ds = _lds(cc3m_hf, split=cc3m_split)
            for row in cc3m_ds:
                key = str(row.get("__key__", ""))
                if key in neg_stems_set:
                    neg_captions[key] = row.get("txt", "") or ""
                    if len(neg_captions) == len(neg_stems_set):
                        break  # found all we need

        # ── Assemble sample list ──────────────────────────────────────────────
        self.samples: List[dict] = []
        for imgid in positive_ids:
            caption = bt_captions.get(imgid, "") if caption_mode == "caption" else ""
            self.samples.append({"path": stem_to_path[imgid], "imgid": imgid,
                                  "caption": caption, "label": 1})
        for stem in negative_stems:
            caption = neg_captions.get(stem, "") if caption_mode == "caption" else ""
            self.samples.append({"path": stem_to_path[stem], "imgid": stem,
                                  "caption": caption, "label": 0})

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

    def __len__(self) -> int:
        return len(self.samples)

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


class _IndexedWrapper(Dataset):
    """Wraps a dataset to include the global index in each sample."""

    def __init__(self, dataset: Dataset):
        self.dataset = dataset

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

    def __getitem__(self, idx):
        return {**self.dataset[idx], "_idx": idx}


def create_dataloader(
    args, processor: LlavaProcessor,
    rank: int = 0, world_size: int = 1,
) -> Tuple[Dataset, DataLoader]:
    """
    Unified dataloader for all three visualization scripts.

    --data_mode:
        "toilet"  — pbcong/bathroom-toilet positives (filtered by --object_mode) +
                    --num_negatives random CC3M images from --image_folder that are
                    NOT in the HF dataset.  Requires --image_folder.
                    For --caption_mode caption with negatives, also pass
                    --hf_dataset (CC3M) so their captions can be looked up.
        "cc3m"    — Full CC3M via --hf_dataset + --local_val_path.
        "coco"    — COCO (yerevann/coco-karpathy) via --hf_dataset + --local_val_path.
        "folder"  — Plain image folder via --data_dir (no captions).

    --caption_mode:
        "caption"   — teacher-forced with stored caption (CC3M txt / COCO sentences /
                      pbcong/bathroom-toilet caption field for positives; CC3M txt
                      looked up by __key__ for toilet negatives).
        "generated" — caption="" in every batch item; single_pass() generates the
                      caption first, then does a second teacher-forced activation pass.
    """
    data_mode    = getattr(args, "data_mode",    "toilet")
    caption_mode = getattr(args, "caption_mode", "generated")

    if data_mode == "toilet":
        dataset = ToiletWithNegativesDataset(
            image_folder  = args.image_folder,
            object_mode   = getattr(args, "object_mode",   "toilet"),
            caption_mode  = caption_mode,
            num_negatives = getattr(args, "num_negatives", 10000),
            cc3m_hf       = getattr(args, "hf_dataset",   None),
            cc3m_split    = getattr(args, "split",         "train"),
        )

    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.model_name,
            batch_size       = args.batch_size,
            num_workers      = getattr(args, "num_workers", 4),
        )
        split   = getattr(args, "split", "train")
        dataset = (cc3m_dataset if data_mode == "cc3m" else coco_dataset)(data_cfg, split)

    elif data_mode == "folder":
        dataset = ImageFolderDataset(args.data_dir)

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

    indexed = _IndexedWrapper(dataset)
    sampler = (
        DistributedSampler(indexed, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True)
        if world_size > 1 else None
    )

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

    def collate_fn(batch: List[dict]):
        batch = [b for b in batch if b is not None]
        if not batch:
            return None
        images      = [b["image"]  for b in batch]
        global_idxs = [b["_idx"]   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:
            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"],
            "global_idxs":    global_idxs,
        }

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


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def find_image_token_positions(input_ids: t.Tensor) -> t.Tensor:
    """Return position of <image> in each sample, or -1 if absent."""
    mask = input_ids == IMAGE_TOKEN_ID
    has = mask.any(dim=1)
    pos = mask.int().argmax(dim=1)
    pos[~has] = -1
    return pos


def expanded_to_original(exp_pos: int, ip: int) -> int:
    """Map expanded-sequence position to original input_ids position."""
    if ip < 0:
        return exp_pos
    if exp_pos < ip:
        return exp_pos
    if exp_pos < ip + N_IMAGE_PATCHES:
        return -1   # image patch — no original token
    return exp_pos - N_IMAGE_PATCHES + 1


def original_to_expanded(orig_pos: int, ip: int) -> int:
    """Map original input_ids position to expanded-sequence position."""
    if ip < 0:
        return orig_pos
    if orig_pos < ip:
        return orig_pos
    if orig_pos == ip:
        return -1   # <image> token expands to 576 patches
    return orig_pos + N_IMAGE_PATCHES - 1


def _short_hook(name: str) -> str:
    parts = name.split(".")
    for i, p in enumerate(parts):
        if p == "layers" and i + 1 < len(parts):
            layer_num = parts[i + 1]
            hook_type = (
                parts[-1].replace("hook_resid_", "")
                if parts[-1].startswith("hook_")
                else parts[-1]
            )
            return f"L{layer_num}.{hook_type}"
    return name[-20:]


# ─────────────────────────────────────────────────────────────────────────────
# Single pass
# ─────────────────────────────────────────────────────────────────────────────

@t.no_grad()
def single_pass(
    model, sae, processor: LlavaProcessor,
    dataloader: DataLoader, device: t.device,
    feature_ids: List[int], args,
) -> Tuple[
    Dict[int, TopKHeap],   # image heaps
    Dict[int, TopKHeap],   # text heaps
    Dict[int, int],        # activation counts per feature
]:
    """
    Single forward pass over the dataset.

    For each batch:
      1. Run LLaVA forward, capturing the single specified hook activation.
         In ``generated`` mode the model first generates a caption per image,
         then a second teacher-forced forward pass with the full
         ``prompt + generated_caption`` is used to collect activations.
      2. Run the SAE to get sparse feature activations.
      3. For each target feature, update top-k heaps and activation counts.
    """
    sae_dtype = next(sae.parameters()).dtype
    sae.eval()
    sae.to(device)

    image_heaps = {fi: TopKHeap(args.top_images) for fi in feature_ids}
    text_heaps = {fi: TopKHeap(args.top_texts) for fi in feature_ids}
    hook_counts: Dict[int, int] = {fi: 0 for fi in feature_ids}

    n_batches = (
        len(dataloader) if args.max_batches is None
        else min(args.max_batches, len(dataloader))
    )

    for batch_idx, batch in enumerate(tqdm(dataloader, total=n_batches, desc="Scanning")):
        if args.max_batches is not None and batch_idx >= args.max_batches:
            break

        B = batch["input_ids"].shape[0]
        model_inputs = {
            k: batch[k].to(device)
            for k in ("input_ids", "attention_mask", "pixel_values")
        }
        global_idxs = batch["global_idxs"]          # List[int]

        # ── Generated caption mode ───────────────────────────────────────────
        # Mirrors Train_Probe_SAE.py: generate → decode → rebuild prompt with
        # generated caption → re-process through processor → forward pass.
        # This avoids feeding raw gen_ids (with padding / EOS artefacts) back
        # into the model and guarantees a clean teacher-forced activation pass.
        if args.caption_mode == "generated":
            with t.no_grad():
                gen_ids = model.generate(
                    **model_inputs,
                    do_sample=False,
                    num_beams=1,
                    use_cache=True,
                    max_new_tokens=args.max_new_tokens,
                )   # (B, T_prompt + T_gen)

            # Decode each sample, extract only the ASSISTANT response
            full_texts = processor.batch_decode(gen_ids, skip_special_tokens=True)
            captions = []
            for txt in full_texts:
                if "ASSISTANT:" in txt:
                    captions.append(txt.split("ASSISTANT:")[-1].strip())
                else:
                    captions.append(txt.strip())

            # Rebuild clean inputs with the generated caption appended
            # We need the original PIL images from the dataset
            images = [
                dataloader.dataset.dataset[gi]["image"]   # _IndexedWrapper → base dataset
                for gi in global_idxs
            ]
            forced_texts = [
                f"USER: <image>\nDescribe this image. \nASSISTANT: {cap}"
                for cap in captions
            ]
            re_processed = processor(
                images=images, text=forced_texts,
                return_tensors="pt", padding=True,
            )
            model_inputs = {
                "input_ids":      re_processed["input_ids"].to(device),
                "attention_mask": re_processed["attention_mask"].to(device),
                "pixel_values":   re_processed["pixel_values"].to(device),
            }
            input_ids_cpu = re_processed["input_ids"]  # (B, T_forced)
            B = input_ids_cpu.shape[0]
        else:
            input_ids_cpu = batch["input_ids"]         # (B, T_orig)

        acts = model.get_activation_at(model_inputs, args.hook_point)     # (B, T, D)

        img_positions = find_image_token_positions(input_ids_cpu)    # (B,)

        
        _, T, D = acts.shape
        flat = acts.reshape(B * T, 1, D).to(device, dtype=sae_dtype)
        loader = get_loader(flat, batch_size=min(args.sae_batch, B * T))
        idxs, vals = get_sae_activations(sae, loader, device, no_tqdm=True)
        # (B*T, 1, k) → (B, T, k)
        idxs = idxs.squeeze(1).reshape(B, T, -1)
        vals = vals.squeeze(1).reshape(B, T, -1)
        for fi in feature_ids:
            mask = (idxs == fi)
            feat_acts = (vals * mask.float()).sum(dim=-1)   # (B, T)

            # Activation count for this feature
            above = feat_acts > args.threshold
            hook_counts[fi] += above.sum().item()

            # ── Per-sample top-k extraction ─────────────────────────
            for b in range(B):
                ip = img_positions[b].item()
                global_idx = global_idxs[b]
                acts_b = feat_acts[b]                       # (T,)

                # — Image patches —
                if ip >= 0:
                    img_acts = acts_b[ip : ip + N_IMAGE_PATCHES]
                    n_img = (img_acts > args.threshold).sum().item()
                    if n_img > 0:
                        k_img = min(args.top_images, n_img)
                        topk_v, topk_i = img_acts.topk(k_img)
                        for v, pidx in zip(topk_v, topk_i):
                            if v.item() < args.threshold:
                                break
                            image_heaps[fi].push(
                                v.item(),
                                ImageRecord(v.item(), global_idx,
                                            pidx.item(), args.hook_point),
                            )

                # — Text tokens (before and after image) —
                text_ranges: List[Tuple[int, int]] = []
                if ip >= 0:
                    if ip > 0:
                        text_ranges.append((0, ip))
                    if ip + N_IMAGE_PATCHES < T:
                        text_ranges.append((ip + N_IMAGE_PATCHES, T))
                else:
                    text_ranges.append((0, T))

                for rng_start, rng_end in text_ranges:
                    txt_acts = acts_b[rng_start:rng_end]
                    n_txt = (txt_acts > args.threshold).sum().item()
                    if n_txt == 0:
                        continue
                    k_txt = min(args.top_texts, n_txt)
                    topk_v, topk_i = txt_acts.topk(k_txt)

                    for v, rel_i in zip(topk_v, topk_i):
                        if v.item() < args.threshold:
                            break
                        exp_pos = rng_start + rel_i.item()
                        orig_pos = expanded_to_original(exp_pos, ip)
                        if orig_pos < 0:
                            continue
                        # Skip <image> placeholder tokens regardless of ip detection
                        if orig_pos < input_ids_cpu.shape[1] and input_ids_cpu[b, orig_pos].item() == IMAGE_TOKEN_ID:
                            continue

                        # Context window in original token space
                        # Clamp so <image> token (position ip) never leaks in
                        seq_len = input_ids_cpu.shape[1]
                        ctx_s = max(0, orig_pos - args.buffer)
                        ctx_e = min(seq_len, orig_pos + args.buffer + 1)
                        if ip >= 0:
                            if orig_pos < ip:
                                ctx_e = min(ctx_e, ip)
                            elif orig_pos > ip:
                                ctx_s = max(ctx_s, ip + 1)
                        ctx_ids = input_ids_cpu[b, ctx_s:ctx_e].tolist()

                        # Feature activations for context tokens
                        ctx_acts_list: List[float] = []
                        for op in range(ctx_s, ctx_e):
                            ep = original_to_expanded(op, ip)
                            if ep < 0 or ep >= T:
                                ctx_acts_list.append(0.0)
                            else:
                                ctx_acts_list.append(feat_acts[b, ep].item())

                        text_heaps[fi].push(
                            v.item(),
                            TextRecord(
                                v.item(), global_idx, args.hook_point,
                                ctx_ids, ctx_acts_list,
                                orig_pos - ctx_s,
                            ),
                        )

        del acts, flat, idxs, vals, feat_acts

    return image_heaps, text_heaps, hook_counts


# ─────────────────────────────────────────────────────────────────────────────
# Visualisation helpers
# ─────────────────────────────────────────────────────────────────────────────

def _img_to_b64(img: Image.Image, max_dim: int = 400) -> str:
    """Resize (thumbnail) and encode an image as base64 PNG."""
    img = img.copy()
    img.thumbnail((max_dim, max_dim))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def _patch_crop(
    img: Image.Image, patch_idx: int, context: int = 2,
) -> Image.Image:
    """Crop a region around the given patch from the original image."""
    W, H = img.size
    pw, ph = W / PATCH_GRID, H / PATCH_GRID
    row, col = patch_idx // PATCH_GRID, patch_idx % PATCH_GRID
    r0 = max(0, row - context)
    r1 = min(PATCH_GRID, row + context + 1)
    c0 = max(0, col - context)
    c1 = min(PATCH_GRID, col + context + 1)
    return img.crop((int(c0 * pw), int(r0 * ph), int(c1 * pw), int(r1 * ph)))


def _annotate_image(img: Image.Image, patch_idx: int) -> Image.Image:
    """Draw a red rectangle on the image around the given patch."""
    annotated = img.copy()
    W, H = annotated.size
    pw, ph = W / PATCH_GRID, H / PATCH_GRID
    row, col = patch_idx // PATCH_GRID, patch_idx % PATCH_GRID
    draw = ImageDraw.Draw(annotated)
    x0, y0 = int(col * pw), int(row * ph)
    x1, y1 = int((col + 1) * pw), int((row + 1) * ph)
    width = max(2, min(W, H) // 100)
    draw.rectangle([x0, y0, x1, y1], outline="red", width=width)
    return annotated


def build_image_html(records: List[ImageRecord], dataset: Dataset) -> str:
    """Build HTML showing top activating image patch crops."""
    if not records:
        return "<p>No image patch activations above threshold.</p>"

    cards: List[str] = []
    for rank, rec in enumerate(records):
        sample = dataset[rec.sample_idx]
        img = sample["image"]
        imgid = sample.get("imgid", rec.sample_idx)

        annotated = _annotate_image(img, rec.patch_idx)
        crop = _patch_crop(img, rec.patch_idx)

        full_b64 = _img_to_b64(annotated, max_dim=350)
        crop_b64 = _img_to_b64(crop, max_dim=200)

        row, col = rec.patch_idx // PATCH_GRID, rec.patch_idx % PATCH_GRID
        cards.append(
            f'<div class="image-card">'
            f'  <div style="display:flex;gap:8px;align-items:flex-start;">'
            f'    <img src="data:image/png;base64,{full_b64}" style="max-height:280px;"/>'
            f'    <img src="data:image/png;base64,{crop_b64}" '
            f'         style="max-height:180px;border:2px solid #e74c3c;"/>'
            f'  </div>'
            f'  <div class="meta">'
            f'    #{rank+1} &nbsp; act={rec.value:.3f} &nbsp; '
            f'    patch={rec.patch_idx} ({row},{col})'
            f'    &nbsp; {_short_hook(rec.hook_name)} &nbsp; imgid={imgid}'
            f'  </div>'
            f'</div>'
        )

    return "<h2>Top Activating Image Patches</h2>\n" + "\n".join(cards)


def build_text_html(
    records: List[TextRecord], processor: LlavaProcessor,
) -> str:
    """Build HTML showing top activating text tokens with context."""
    if not records:
        return "<p>No text activations above threshold.</p>"

    # Filter out records whose target token is the <image> placeholder
    records = [
        rec for rec in records
        if not (rec.target_relative < len(rec.context_ids)
                and rec.context_ids[rec.target_relative] == IMAGE_TOKEN_ID)
    ]
    if not records:
        return "<p>No text activations above threshold.</p>"

    entries: List[str] = []
    for rank, rec in enumerate(records):
        tokens: List[str] = []
        for tid in rec.context_ids:
            if tid == IMAGE_TOKEN_ID:
                tokens.append("")   # omit <image> from context display
            else:
                tokens.append(processor.tokenizer.decode(tid))

        max_act = max(rec.context_acts) if rec.context_acts else 1.0
        if max_act <= 0:
            max_act = 1.0

        spans: List[str] = []
        for i, (tok, act) in enumerate(zip(tokens, rec.context_acts)):
            intensity = min(1.0, act / max_act)
            g = int(200 * intensity)
            bg = f"rgba(0,{g},0,{intensity * 0.6:.2f})" if intensity > 0.05 else "transparent"
            bold = "font-weight:bold;" if i == rec.target_relative else ""
            underline = "border-bottom:2px solid #e74c3c;" if i == rec.target_relative else ""
            tok_safe = html_lib.escape(tok)
            spans.append(
                f'<span style="background:{bg};{bold}{underline}'
                f'padding:1px 3px;border-radius:2px;">{tok_safe}</span>'
            )

        entries.append(
            f'<div class="text-card">'
            f'  <div class="text-rank">#{rank+1}</div>'
            f'  <div style="flex:1;">'
            f'    <div class="text-context">{"".join(spans)}</div>'
            f'    <div class="meta">'
            f'      act={rec.value:.3f} &nbsp; {_short_hook(rec.hook_name)}'
            f'    </div>'
            f'  </div>'
            f'</div>'
        )

    return "<h2>Top Activating Text Tokens</h2>\n" + "\n".join(entries)


# ─────────────────────────────────────────────────────────────────────────────
# Per-feature HTML report
# ─────────────────────────────────────────────────────────────────────────────

def save_feature_report(
    feat_idx: int,
    image_records: List[ImageRecord],
    text_records: List[TextRecord],
    total_act: int,
    hook_point: str,
    dataset: Dataset,
    processor: LlavaProcessor,
    output_dir: Path,
) -> Optional[str]:
    """Generate and save a full HTML report for one feature."""
    feat_dir = output_dir / f"feature_{feat_idx}"
    feat_dir.mkdir(parents=True, exist_ok=True)

    if total_act == 0 and not image_records and not text_records:
        return None

    image_section = build_image_html(image_records, dataset)
    text_section = build_text_html(text_records, processor)

    html = f"""<!DOCTYPE html>
<html>
<head>
  <meta charset="UTF-8">
  <title>Feature {feat_idx}</title>
  <style>
    body {{ font-family: 'Segoe UI', sans-serif; max-width: 1200px; margin: auto;
           padding: 20px; background: #f8f9fa; }}
    h1   {{ color: #2c3e50; }}
    h2   {{ color: #34495e; margin-top: 24px; }}
    .info {{ background: #fff; border: 1px solid #dee2e6; border-radius: 6px;
             padding: 12px 20px; margin-bottom: 20px; }}
    .info td {{ padding: 4px 12px; }}
    .image-card {{ background: #fff; border: 1px solid #ddd; border-radius: 6px;
                   padding: 10px; margin: 8px 0; }}
    .text-card {{ background: #fff; border: 1px solid #ddd; border-radius: 6px;
                  padding: 8px 12px; margin: 6px 0; display: flex;
                  align-items: flex-start; gap: 10px; }}
    .text-rank {{ font-weight: bold; color: #888; min-width: 35px; padding-top: 2px; }}
    .text-context {{ font-family: monospace; font-size: 0.9rem; line-height: 1.6;
                     word-break: break-word; }}
    .meta {{ font-size: 0.8rem; color: #666; margin-top: 4px; }}
  </style>
</head>
<body>
  <h1>Feature {feat_idx}</h1>
  <div class="info">
    <table>
      <tr><td><strong>Hook</strong></td><td>{hook_point}</td></tr>
      <tr><td><strong>Total activations</strong></td><td>{total_act}</td></tr>
      <tr><td><strong>Image records</strong></td><td>{len(image_records)}</td></tr>
      <tr><td><strong>Text records</strong></td><td>{len(text_records)}</td></tr>
    </table>
  </div>

  {image_section}

  {text_section}
</body>
</html>"""

    report_path = feat_dir / "report.html"
    report_path.write_text(html, encoding="utf-8")
    return str(report_path)


# ─────────────────────────────────────────────────────────────────────────────
# Index page
# ─────────────────────────────────────────────────────────────────────────────

def save_index_page(
    feature_ids: List[int],
    report_paths: Dict[int, Optional[str]],
    hook_counts: Dict[int, int],
    hook_point: str,
    output_dir: Path,
):
    rows: List[str] = []
    for fi in feature_ids:
        path = report_paths.get(fi)
        total = hook_counts.get(fi, 0)
        if path:
            rel = os.path.relpath(path, output_dir)
            link = f'<a href="{rel}">feature_{fi}</a>'
        else:
            link = f'<span style="color:#999">feature_{fi} (no activations)</span>'
        rows.append(f"<tr><td>{fi}</td><td>{total}</td><td>{link}</td></tr>")

    html = f"""<!DOCTYPE html>
<html><head><meta charset="UTF-8"><title>SAE Feature Index</title>
<style>
  body {{ font-family: sans-serif; max-width: 900px; margin: auto; padding: 20px; }}
  table {{ border-collapse: collapse; width: 100%; }}
  th, td {{ border: 1px solid #ccc; padding: 6px 12px; text-align: left; }}
  th {{ background: #f0f0f0; }}
  tr:nth-child(even) {{ background: #fafafa; }}
</style></head><body>
  <h1>SAE Feature Index</h1>
  <p><strong>Hook:</strong> {hook_point}</p>
  <table>
    <thead><tr><th>Feature</th><th>Total Acts</th><th>Report</th></tr></thead>
    <tbody>{"".join(rows)}</tbody>
  </table>
</body></html>"""
    (output_dir / "index.html").write_text(html, encoding="utf-8")


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

def main():
    ap = argparse.ArgumentParser(
        description="Single-pass SAE feature visualizer for LLaVA.",
    )

    # Model / SAE
    ap.add_argument("--sae_ckpt", required=True)
    ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--device_id", type=int, default=0)
    ap.add_argument("--dtype", default="float16")

    # ── Data mode ────────────────────────────────────────────────────────────
    ap.add_argument(
        "--data_mode", default="toilet", choices=["toilet", "cc3m", "coco", "folder"],
        help=(
            "toilet: pbcong/bathroom-toilet positives + CC3M negatives. "
            "cc3m:   full CC3M via --hf_dataset + --local_val_path. "
            "coco:   COCO via --hf_dataset + --local_val_path. "
            "folder: plain image folder via --data_dir."
        ),
    )

    # ── Data — toilet ────────────────────────────────────────────────────────
    ap.add_argument("--image_folder",  default=None,
                    help="[toilet] Local CC3M image folder (positives + negative pool).")
    ap.add_argument("--object_mode",   default="toilet", choices=["toilet", "bathroom", "both"],
                    help="[toilet] Positive class: toilet==1, bathroom==1, or both (toilet==1 OR bathroom==1).")
    ap.add_argument("--num_negatives", type=int, default=10000,
                    help="[toilet] Number of random CC3M negatives to include.")

    # ── Data — CC3M / COCO / toilet-negative captions ────────────────────────
    ap.add_argument("--hf_dataset",    default=None,
                    help="HF dataset path (CC3M / COCO / also used for toilet-negative captions).")
    ap.add_argument("--local_val_path", default=None,
                    help="[cc3m/coco] Local image root for the HF dataset split.")
    ap.add_argument("--data_dir",       default=None,
                    help="[folder] Plain image folder.")
    ap.add_argument("--split",          default="train",
                    help="HF dataset split.")

    # ── Caption mode ─────────────────────────────────────────────────────────
    ap.add_argument(
        "--caption_mode", default="generated", choices=["generated", "caption"],
        help=(
            "generated: model generates caption first; teacher-forced pass collects acts. "
            "caption:   use stored caption (CC3M txt / COCO sentences / "
            "pbcong/bathroom-toilet caption field)."
        ),
    )
    ap.add_argument("--num_workers", type=int, default=4)

    # Features (required)
    ap.add_argument("--feature_ids", type=int, nargs="+", required=True,
                    help="SAE feature IDs to visualize.")

    # Processing
    ap.add_argument("--batch_size", type=int, default=4,
                    help="DataLoader batch size. Keep small (4-8) for LLaVA-7B.")
    ap.add_argument("--sae_batch", type=int, default=4096,
                    help="Sub-batch size for SAE processing.")
    ap.add_argument("--threshold", type=float, default=1e-3,
                    help="Minimum activation to count / record.")
    ap.add_argument("--max_batches", type=int, default=None,
                    help="Limit number of batches (useful for testing).")
    ap.add_argument("--max_new_tokens", type=int, default=128,
                    help="Max tokens to generate per image (only with --caption_mode generated).")
    ap.add_argument("--hook_point", type=str, required=True,
                    help="Exact hook name to visualize, e.g. "
                         "'model.language_model.layers.19.hook_resid_post'.")

    # Visualisation
    ap.add_argument("--output_dir", default="outputs/features")
    ap.add_argument("--top_images", type=int, default=20,
                    help="Number of top image patches to show per feature.")
    ap.add_argument("--top_texts", type=int, default=20,
                    help="Number of top text tokens to show per feature.")
    ap.add_argument("--buffer", type=int, default=10,
                    help="Token context window radius for text display.")

    args = ap.parse_args()

    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).")
    if args.data_mode == "folder" and not args.data_dir:
        ap.error("--data_mode folder requires --data_dir.")

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

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

    if rank == 0:
        print(f"Running with {world_size} GPU(s)")

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

    # ── Load model + processor ───────────────────────────────────────────────
    model = HookedSAELlavaConditionalGeneration.from_pretrained(args.model_name)
    model.to(device, dtype=dtype)
    model.eval()
    processor = LlavaProcessor.from_pretrained(args.model_name)

    # ── Load data ────────────────────────────────────────────────────────────
    dataset, dataloader = create_dataloader(args, processor, rank, world_size)
    if rank == 0:
        print(f"Dataset: {len(dataset)} samples, {len(dataloader)} batches/rank.")
        print(f"Features to visualize: {args.feature_ids}")

    # ── Single pass (each rank processes its shard) ──────────────────────────
    image_heaps, text_heaps, hook_counts = single_pass(
        model, sae, processor, dataloader, device, args.feature_ids, args,
    )

    # ── Gather results from all ranks ────────────────────────────────────────
    output_dir = Path(args.output_dir)
    if is_distributed:
        t.cuda.empty_cache()
        tmp_dir = output_dir / ".ddp_tmp"
        if rank == 0:
            tmp_dir.mkdir(parents=True, exist_ok=True)
        dist.barrier()

        with open(tmp_dir / f"rank_{rank}.pkl", "wb") as f:
            pickle.dump({
                "image_heaps": {fi: image_heaps[fi].sorted_records() for fi in args.feature_ids},
                "text_heaps": {fi: text_heaps[fi].sorted_records() for fi in args.feature_ids},
                "hook_counts": hook_counts,
            }, f)
        dist.barrier()

        if rank == 0:
            merged_img = {fi: TopKHeap(args.top_images) for fi in args.feature_ids}
            merged_txt = {fi: TopKHeap(args.top_texts) for fi in args.feature_ids}
            merged_cnt: Dict[int, int] = {fi: 0 for fi in args.feature_ids}

            for r in range(world_size):
                with open(tmp_dir / f"rank_{r}.pkl", "rb") as f:
                    data = pickle.load(f)
                for fi in args.feature_ids:
                    for rec in data["image_heaps"][fi]:
                        merged_img[fi].push(rec.value, rec)
                    for rec in data["text_heaps"][fi]:
                        merged_txt[fi].push(rec.value, rec)
                    merged_cnt[fi] += data["hook_counts"][fi]

            image_heaps = merged_img
            text_heaps = merged_txt
            hook_counts = merged_cnt
            shutil.rmtree(tmp_dir)

    # ── Only rank 0 generates reports ────────────────────────────────────────
    if rank == 0:
        output_dir.mkdir(parents=True, exist_ok=True)

        report_paths: Dict[int, Optional[str]] = {}
        summary: Dict[str, dict] = {}

        for fi in args.feature_ids:
            print(f"Feature {fi} ...", end=" ")
            img_recs = image_heaps[fi].sorted_records()
            txt_recs = text_heaps[fi].sorted_records()
            total = hook_counts[fi]

            path = save_feature_report(
                fi, img_recs, txt_recs, total, args.hook_point, dataset, processor, output_dir,
            )
            report_paths[fi] = path
            summary[str(fi)] = {"total_activations": total, "hook_point": args.hook_point}
            print("saved." if path else "skipped (no activations).")

        save_index_page(args.feature_ids, report_paths, hook_counts, args.hook_point, output_dir)
        with open(output_dir / "summary.json", "w") as f:
            json.dump(summary, f, indent=2)

        print(f"\nDone. Open {output_dir}/index.html to browse features.")

    cleanup_distributed()


if __name__ == "__main__":
    main()