"""
Train binary linear probe(s) on SAE feature activations of Llava at specified layers.

Two-phase pipeline:
  Phase 1 (one VLM pass + SAE): for each image generate caption, run forced
    forward pass, extract hook activations, SAE encode + max-pool, and keep only
    the pooled feature vectors in memory (one per hook point). No disk writes.
  Phase 2: gather all pooled vectors from all ranks → train binary linear probes
    per hook point/layer. All probes share the same positive/negative split.

Dataset: any HF dataset with binary labels and train/validation splits.
Negatives: pre-split image IDs from {"train": [...], "validation": [...]} JSON file.
Saved probes: {save_dir}/{hook_type}/probe_{layer}.pth
"""

import argparse
import json
import os
from tqdm import tqdm

import torch as t
import torch.distributed as dist
from datasets import load_dataset
from PIL import Image
from sklearn.metrics import (
    confusion_matrix, f1_score, precision_score, recall_score, roc_auc_score,
)
from transformers import AutoConfig, LlavaProcessor

from extra_materials.mechanistic_interp.interp_utils import *
from extra_materials.mechanistic_interp.probe.probing import train_binary_probe
from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from sae.SAE_Tools import get_loader, get_sae_activations, load_sae_model

import wandb
wandb.init(project="Probe_SAE")


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

def setup_distributed():
    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()


def gather_tensors(tensor: t.Tensor, world_size: int) -> t.Tensor:
    """Gather variable-length tensors from all ranks onto rank 0."""
    if world_size == 1:
        return tensor
    local_size = t.tensor([tensor.shape[0]], device=tensor.device)
    all_sizes  = [t.zeros(1, dtype=t.long, device=tensor.device) for _ in range(world_size)]
    dist.all_gather(all_sizes, local_size)
    max_size = max(s.item() for s in all_sizes)
    if tensor.shape[0] < max_size:
        pad = list(tensor.shape)
        pad[0] = max_size - tensor.shape[0]
        tensor = t.cat([tensor, t.zeros(pad, dtype=tensor.dtype, device=tensor.device)])
    gathered = [t.zeros_like(tensor) for _ in range(world_size)]
    dist.all_gather(gathered, tensor)
    return t.cat([g[: int(s.item())] for g, s in zip(gathered, all_sizes)], dim=0)


# ── GPU memory reservation ───────────────────────────────────────────────────

def reserve_gpu_memory(device, fraction: float = 0.90, verbose: bool = True) -> None:
    """Claim VRAM in PyTorch's allocator so competing CUDA processes can't take it.

    Allocating then deleting a large buffer populates PyTorch's per-process cache.
    Other processes' cudaMalloc calls see far less free VRAM and will fail instead
    of competing at peak-memory moments.  Never call torch.cuda.empty_cache() after
    this — it releases cached blocks back to the driver and undoes the reservation.
    Controlled via GPU_MEM_RESERVE_FRAC env var (default 0.90).
    """
    if not str(device).startswith("cuda"):
        return
    total = t.cuda.get_device_properties(device).total_memory
    already = t.cuda.memory_reserved(device)
    to_alloc = max(int(total * fraction) - already, 0)
    if to_alloc == 0:
        if verbose:
            print(f"  [mem-reserve] already at {already/2**30:.1f}/{total/2**30:.1f} GB — skip")
        return
    try:
        buf = t.empty(to_alloc // 2, dtype=t.int16, device=device)
        del buf
        reserved = t.cuda.memory_reserved(device)
        if verbose:
            print(f"  [mem-reserve] {reserved/2**30:.1f}/{total/2**30:.1f} GB reserved "
                  f"({100*reserved/total:.0f}%) — others see ≤{(total-reserved)/2**30:.1f} GB free")
    except t.cuda.OutOfMemoryError:
        t.cuda.empty_cache()
        print(f"  [mem-reserve] WARNING: reservation failed — GPU already heavily loaded "
              f"({t.cuda.memory_allocated(device)/2**30:.1f} GB in use)")


# ── Hook-point construction ──────────────────────────────────────────────────

_HOOK_SUFFIX = {"pre": "hook_resid_pre", "mid": "hook_resid_mid", "post": "hook_resid_post"}

def build_hook_points(layers: list, hook_type: str) -> list:
    suffix = _HOOK_SUFFIX[hook_type]
    return [f"model.language_model.layers.{layer}.{suffix}" for layer in layers]


# ── Phase 1: cache activations ───────────────────────────────────────────────

def phase1_cache(rank, world_size, model, processor, image_files,
                 hook_points, args, device, question, sae):
    """
    Batch VLM forward passes; for each image: generate caption, forced forward,
    SAE encode + max-pool activations. Returns {img_stem: {hp: pooled_vec}}.
    """
    shard = image_files[rank::world_size]
    processed = 0

    rank_cache = {}  # {img_stem: {hp: pooled_vec}}

    batch_iter = range(0, len(shard), args.batch_size)
    if rank == 0:
        batch_iter = tqdm(batch_iter, desc="Phase1", total=len(batch_iter))
    for i in batch_iter:
        batch_files = shard[i:i + args.batch_size]
        images, valid_files = [], []
        for img_file in batch_files:
            try:
                img = Image.open(os.path.join(args.image_folder, img_file)).convert("RGB")
                images.append(img)
                valid_files.append(img_file)
            except Exception:
                continue
        if not images:
            continue

        # Generate captions
        gen_prompts = [f"USER: <image>\n{question}\nASSISTANT:"] * len(images)
        gen_inputs = processor(images=images, text=gen_prompts, return_tensors="pt", padding=True).to(device)

        with t.no_grad():
            gen_outputs = model.generate(
                **gen_inputs, do_sample=False, num_beams=1,
                use_cache=True, max_new_tokens=args.max_new_tokens,
            )

        captions = processor.batch_decode(gen_outputs, skip_special_tokens=True)
        del gen_inputs, gen_outputs

        # Forced forward to get activations
        forced_texts = []
        for cap in captions:
            asst_text = cap.split("ASSISTANT:")[-1].strip()
            forced_texts.append(f"USER: <image>\n{question}\nASSISTANT: {asst_text}")

        act_buf: dict = {}
        def make_hook(name):
            def _fn(act: t.Tensor, hook):
                act_buf[name] = act.detach().cpu()
            return _fn

        fwd_inputs = processor(images=images, text=forced_texts, return_tensors="pt", padding=True).to(device)
        with t.no_grad():
            model.run_with_hooks(
                fwd_inputs,
                fwd_hooks=[(hp, make_hook(hp)) for hp in hook_points]
            )

        attn_mask = fwd_inputs['attention_mask'].cpu()
        seq_lens = attn_mask.sum(dim=1)
        # Keep input_ids for dry-run token-slice verification before deleting fwd_inputs
        input_ids_cpu = fwd_inputs['input_ids'].cpu() if args.dry_run else None
        del fwd_inputs

        # Collect activations per hook point for batch SAE encoding
        batch_acts_per_hp = {hp: [] for hp in hook_points}
        for idx in range(len(valid_files)):
            current_caption = captions[idx].split("ASSISTANT:")[-1].strip()
            # add_special_tokens=False: prevents a spurious BOS that would shift
            # the slice one token too early into the "ASSISTANT: " prefix.
            cap_len = len(processor.tokenizer(current_caption,
                                              add_special_tokens=False)["input_ids"])

            if args.dry_run and idx == 0 and i == 0:
                sl = input_ids_cpu[idx, seq_lens[idx] - cap_len : seq_lens[idx]]
                decoded = processor.tokenizer.decode(sl, skip_special_tokens=True)
                print(f"  [dry-run] token slice (len={cap_len}): '{decoded}'")
                print(f"  [dry-run] expected caption:            '{current_caption}'")

            for hp in hook_points:
                full = act_buf[hp][idx]
                # Slice to only the generated caption tokens at the end of the valid sequence.
                # Image patches + question tokens are excluded by taking the tail cap_len tokens.
                a = full[seq_lens[idx] - cap_len : seq_lens[idx]]  # (cap_len, D)
                batch_acts_per_hp[hp].append(a)

        act_buf.clear()

        # SAE encode + pool each hook point's batch
        for hp in hook_points:
            if not batch_acts_per_hp[hp]:
                continue
            pooled_list = sae_encode_and_pool(batch_acts_per_hp[hp], sae, args.sae_batch, device)

            # Store pooled vectors per image keyed by bare stem ('000949152')
            for img_file, pooled_vec in zip(valid_files, pooled_list):
                stem = os.path.splitext(os.path.basename(img_file))[0]
                if stem not in rank_cache:
                    rank_cache[stem] = {}
                rank_cache[stem][hp] = pooled_vec

        processed += len(valid_files)

    if rank == 0:
        print(f"Phase 1 complete — processed: {processed}")
    if world_size > 1:
        dist.barrier()

    return rank_cache


# ── Phase 2: SAE encode + probe training ─────────────────────────────────────

def sae_encode_and_pool(acts_list: list, sae, sae_batch: int, device) -> list:
    """
    Encode variable-length (T_i, D) activations through the SAE in a single pass,
    then max-pool over tokens per input using scatter_reduce.
    Returns a list of (d_sae,) CPU vectors, one per input.
    """
    lens = [a.shape[0] for a in acts_list]
    flat = t.cat(acts_list, dim=0).unsqueeze(1).to(device)  # (N_total, 1, D)
    loader = get_loader(flat, batch_size=min(sae_batch, flat.shape[0]))
    topk_idxs, topk_vals = get_sae_activations(sae, loader, device, no_tqdm=True)
    topk_idxs = topk_idxs.squeeze(1)          # (N_total, k)
    topk_vals = topk_vals.squeeze(1).float()  # (N_total, k)

    d_sae = sae.cfg.d_sae
    pooled_list = []
    offset = 0
    for length in lens:
        idxs = topk_idxs[offset:offset + length].reshape(-1)
        vals = topk_vals[offset:offset + length].reshape(-1)
        pooled = t.zeros(d_sae, dtype=t.float32, device=topk_vals.device)
        pooled.scatter_reduce_(0, idxs, vals, reduce="amax", include_self=True)
        pooled_list.append(pooled.cpu())
        offset += length
    return pooled_list


def prepare_phase2_data(merged_cache, train_ids_labels, val_ids_labels, hook_points, verbose=False):
    """
    Convert merged_cache: {img_id: {hp: pooled_vec}} →
    {"train": {hp: (Tensor[N, d_sae], Tensor[N])}, "val": ...}
    """
    if verbose:
        cache_sample  = list(merged_cache.keys())[:3]
        train_sample  = [img_id for img_id, _ in train_ids_labels[:3]]
        val_sample    = [img_id for img_id, _ in val_ids_labels[:3]]
        train_matches = sum(1 for img_id, _ in train_ids_labels if img_id in merged_cache)
        val_matches   = sum(1 for img_id, _ in val_ids_labels   if img_id in merged_cache)
        print(f"  [dry-run] cache   ({len(merged_cache)}): {cache_sample}")
        print(f"  [dry-run] train_labels ({len(train_ids_labels)}) hits={train_matches}: {train_sample}")
        print(f"  [dry-run] val_labels   ({len(val_ids_labels)})   hits={val_matches}: {val_sample}")

    train_feats = {hp: [] for hp in hook_points}
    train_lbls = {hp: [] for hp in hook_points}
    val_feats = {hp: [] for hp in hook_points}
    val_lbls = {hp: [] for hp in hook_points}

    for img_id, label in train_ids_labels:
        if img_id not in merged_cache:
            continue
        cached = merged_cache[img_id]
        for hp in hook_points:
            if hp in cached:
                train_feats[hp].append(cached[hp])  # already (d_sae,)
                train_lbls[hp].append(label)

    for img_id, label in val_ids_labels:
        if img_id not in merged_cache:
            continue
        cached = merged_cache[img_id]
        for hp in hook_points:
            if hp in cached:
                val_feats[hp].append(cached[hp])
                val_lbls[hp].append(label)

    return {
        "train": {hp: (t.stack(train_feats[hp]), t.tensor(train_lbls[hp]))
                  for hp in hook_points if train_feats[hp]},
        "val": {hp: (t.stack(val_feats[hp]), t.tensor(val_lbls[hp]))
                for hp in hook_points if val_feats[hp]},
    }


def phase2_train_probes(rank, hook_points, layers, phase2_data, args, device):
    if rank != 0:
        return

    if not args.dry_run:
        os.makedirs(args.save_dir, exist_ok=True)
    train_feats = phase2_data["train"]
    val_feats = phase2_data["val"]

    for hook_point, layer in zip(hook_points, layers):
        if hook_point not in train_feats or hook_point not in val_feats:
            print(f"Skipping layer {layer} ({hook_point}): no activations found.")
            continue

        print(f"\n{'='*64}")
        print(f"  Layer {layer} — {hook_point}")
        print(f"{'='*64}")

        train_data, train_label = train_feats[hook_point]
        val_data, val_label = val_feats[hook_point]

        probe = train_binary_probe(
            train_data, train_label,
            batch_size=args.probe_batch_size,
            num_epochs=args.probe_epochs,
            lr=args.probe_lr,
            device=device,
        )

        if args.dry_run:
            print(f"  [dry-run] skipping save")
        else:
            os.makedirs(os.path.join(args.save_dir, args.hook_type), exist_ok=True)
            save_path = os.path.join(args.save_dir, args.hook_type, f"probe_{layer}.pth")
            t.save(probe.state_dict(), save_path)
            print(f"  Saved → {save_path}")

        with t.no_grad():
            val_logits = probe(val_data.to(device))  # (N, 1)
            val_preds  = (val_logits > 0).float()    # (N, 1)
            val_acc    = (val_preds.squeeze(-1) == val_label.to(device)).float().mean().item()

        # Ensure 1D arrays for metrics by squeezing the (N, 1) tensors to (N,)
        y_true  = val_label.view(-1).cpu().numpy()
        y_score = val_logits.squeeze(-1).float().cpu().numpy()
        y_pred  = val_preds.squeeze(-1).cpu().numpy()

        try:
            auc = roc_auc_score(y_true, y_score)
        except ValueError:
            auc = float("nan")
        f1   = f1_score(y_true, y_pred, zero_division=0)
        prec = precision_score(y_true, y_pred, zero_division=0)
        rec  = recall_score(y_true, y_pred, zero_division=0)
        tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()

        print(f"\n── Validation Results ─────────────────────────────")
        print(f"  Accuracy  : {val_acc:.4f}")
        print(f"  AUC-ROC   : {auc:.4f}")
        print(f"  F1        : {f1:.4f}")
        print(f"  Precision : {prec:.4f}")
        print(f"  Recall    : {rec:.4f}")
        print(f"  TP={tp}  TN={tn}  FP={fp}  FN={fn}")
        print("───────────────────────────────────────────────────")


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

def main():
    ap = argparse.ArgumentParser(
        description="Train binary linear probe(s) on SAE feature activations "
                    "of Llava at specified hook points (two-phase: cache then train)."
    )

    # Model / SAE
    ap.add_argument("--sae_ckpt",   required=True,
                    help="Path to the SAE checkpoint (.ckpt).")
    ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--device_id",  type=int, default=0,
                    help="GPU id (single-GPU mode only; ignored with torchrun).")
    ap.add_argument("--dtype",      default="bfloat16",
                    choices=["float32", "float16", "bfloat16"])

    # Data
    ap.add_argument("--image_folder", required=True,
                    help="Folder of images to process.")
    ap.add_argument("--hf_dataset",   default="pbcong/bathroom-toilet",
                    help="HF dataset with binary relation labels and train/validation splits.")
    ap.add_argument("--id_col",       default="image_id",
                    help="Column in the HF dataset holding the image filename stem.")
    ap.add_argument("--probe_type",   default="toilet",
                    help="Binary label column used as positives (label==1).")
    ap.add_argument("--other_object", default=None,
                    help="Column name of a contrastive object in the HF dataset. "
                         "When set: positives = probe_type=1 (any other_object value); "
                         "HF negatives = other_object=1 & probe_type=0 (drawn from train/val splits). "
                         "Stacked on top of any --neg_jsonl negatives.")
    ap.add_argument("--neg_jsonl",    default=None,
                    help='JSON file: {"train": [...image_ids], "validation": [...image_ids]}. '
                         'Optional when --other_object is set.')

    # Hook / layers
    ap.add_argument("--layers",    type=int, nargs="+", default=None,
                    help="Layer indices; one probe trained per layer (e.g. --layers 10 15 20). "
                         "If omitted, trains on all layers (0 .. num_hidden_layers-1).")
    ap.add_argument("--hook_type", default="mid", choices=["pre", "mid", "post"],
                    help="Residual stream position: pre (before attn), mid (after attn), post (after mlp).")

    # SAE
    ap.add_argument("--sae_batch", type=int, default=2048,
                    help="Sub-batch size for SAE processing.")

    # Generation
    ap.add_argument("--max_new_tokens", type=int, default=128)
    ap.add_argument("--batch_size", type=int, default=32,
                    help="Batch size for Phase 1 image processing.")
    ap.add_argument("--question", type=str, default="Describe this image.",
                    help='Question to ask about the image (used in prompt: "USER: <image>\n{question}\nASSISTANT:")')

    # Probe training
    ap.add_argument("--probe_batch_size", type=int,   default=256)
    ap.add_argument("--probe_epochs",     type=int,   default=20)
    ap.add_argument("--probe_lr",         type=float, default=0.01)

    # Output
    ap.add_argument("--save_dir", required=True,
                    help="Directory to save trained probe checkpoints.")

    # Dry run
    ap.add_argument("--dry_run", action="store_true",
                    help="Smoke-test: process only --dry_run_images images and 3 probe epochs "
                         "to verify the full pipeline end-to-end before a real run.")
    ap.add_argument("--dry_run_images", type=int, default=8,
                    help="Number of images to process in dry-run mode (default 8).")

    args = ap.parse_args()

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

    dtype_map = {"float32": t.float32, "float16": t.float16, "bfloat16": t.bfloat16}
    dtype  = dtype_map[args.dtype]
    device = (
        f"cuda:{local_rank}" if is_distributed
        else (f"cuda:{args.device_id}" if t.cuda.is_available() else "cpu")
    )

    if rank == 0:
        print(f"GPUs: {world_size} | device: {device} | dtype: {dtype}")

    # ── Resolve --layers default to every layer in the model ────────────────
    if not args.layers:
        n_layers = AutoConfig.from_pretrained(args.model_name).text_config.num_hidden_layers
        args.layers = list(range(n_layers))
        if rank == 0:
            print(f"--layers not provided; defaulting to all {n_layers} layers.")

    # ── Hook points ──────────────────────────────────────────────────────────
    hook_points = build_hook_points(args.layers, args.hook_type)
    if rank == 0:
        print(f"Hook points: {hook_points}")

    if not args.other_object and not args.neg_jsonl:
        ap.error("At least one of --other_object or --neg_jsonl must be provided for negatives.")

    # ── Load HF dataset: positive and (optionally) contrastive negative IDs ──
    train_ds = load_dataset(args.hf_dataset, split="train")
    val_ds   = load_dataset(args.hf_dataset, split="validation")
    label_col = args.probe_type

    # Positives: probe_type=1 (includes both other_object=0 and other_object=1).
    train_pos_ids = [row[args.id_col] for row in train_ds if row[label_col] == 1]
    val_pos_ids   = [row[args.id_col] for row in val_ds   if row[label_col] == 1]

    # HF-derived negatives: other_object=1 & probe_type=0 (contrastive mode only).
    if args.other_object:
        train_hf_neg_ids = [row[args.id_col] for row in train_ds
                            if row[args.other_object] == 1 and row[label_col] == 0]
        val_hf_neg_ids   = [row[args.id_col] for row in val_ds
                            if row[args.other_object] == 1 and row[label_col] == 0]
    else:
        train_hf_neg_ids, val_hf_neg_ids = [], []

    # ── Load pre-split negative image IDs ────────────────────────────────────
    train_json_neg_ids, val_json_neg_ids = [], []
    if args.neg_jsonl:
        with open(args.neg_jsonl) as f:
            neg_split = json.load(f)
        train_json_neg_ids = neg_split.get("train", [])
        val_json_neg_ids   = neg_split.get("validation", [])

    train_neg_ids = train_hf_neg_ids + train_json_neg_ids
    val_neg_ids   = val_hf_neg_ids   + val_json_neg_ids

    # ── Filter to images present in the folder ───────────────────────────────
    stem_to_file = {}
    for root, dirs, files in os.walk(args.image_folder):
        for fn in files:
            if fn.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
                stem = os.path.splitext(fn)[0]
                # Store relative path from image_folder to the file
                rel_path = os.path.relpath(os.path.join(root, fn), args.image_folder)
                stem_to_file[stem] = rel_path

    def ids_in_folder(ids):
        # Normalise to bare stem ('000949152') regardless of input format.
        # No overlap between splits, so stems are unique keys.
        return [os.path.splitext(os.path.basename(img_id))[0]
                for img_id in ids
                if os.path.splitext(os.path.basename(img_id))[0] in stem_to_file]

    train_pos_ids = ids_in_folder(train_pos_ids)
    val_pos_ids   = ids_in_folder(val_pos_ids)
    train_neg_ids = ids_in_folder(train_neg_ids)
    val_neg_ids   = ids_in_folder(val_neg_ids)

    train_ids_labels = [(s, 1) for s in train_pos_ids] + [(s, 0) for s in train_neg_ids]
    val_ids_labels   = [(s, 1) for s in val_pos_ids]   + [(s, 0) for s in val_neg_ids]

    all_needed_ids   = set(train_pos_ids + val_pos_ids + train_neg_ids + val_neg_ids)
    all_needed_files = [stem_to_file[s] for s in all_needed_ids]

    if args.dry_run:
        # Sample from each of the four groups so both train and val splits are
        # guaranteed to have activations — a plain [:N] slice on the unordered
        # set can land entirely in one split, causing every layer to be skipped.
        n = max(2, args.dry_run_images // 4)
        dry_stems = set(
            train_pos_ids[:n] + train_neg_ids[:n] +
            val_pos_ids[:n]   + val_neg_ids[:n]
        )
        all_needed_files = [stem_to_file[s] for s in dry_stems][:args.dry_run_images]
        args.probe_epochs = 3
        if rank == 0:
            print(f"\n{'='*64}")
            print(f"  DRY RUN — up to {args.dry_run_images} images ({len(all_needed_files)} selected), 3 probe epochs")
            print(f"{'='*64}\n")

    if rank == 0:
        print(f"Train : {len(train_pos_ids)} pos + {len(train_neg_ids)} neg "
              f"({len(train_hf_neg_ids)} HF-contrastive + {len(train_json_neg_ids)} JSON)"
              f" = {len(train_ids_labels)}")
        print(f"Val   : {len(val_pos_ids)} pos + {len(val_neg_ids)} neg "
              f"({len(val_hf_neg_ids)} HF-contrastive + {len(val_json_neg_ids)} JSON)"
              f" = {len(val_ids_labels)}")
        print(f"Phase 1 images to process: {len(all_needed_files)}")

    # ── Phase 1: run VLM + SAE encode, collect pooled features in memory ────────
    model = HookedSAELlavaConditionalGeneration.from_pretrained(args.model_name)
    processor = LlavaProcessor.from_pretrained(args.model_name)
    model = model.to(device, dtype=dtype)
    model.language_model = t.compile(model.language_model, dynamic=True)

    # Load SAE on same device for Phase 1 encoding
    sae = load_sae_model(args.sae_ckpt, model_type="llava", hook_type="text", device=device)

    # Reserve GPU memory so competing processes on shared machines see this VRAM as taken.
    # Set GPU_MEM_RESERVE_FRAC=0 to disable (e.g. if you own the node).
    _mem_frac = float(os.environ.get("GPU_MEM_RESERVE_FRAC", "0.90"))
    if _mem_frac > 0:
        reserve_gpu_memory(device, fraction=_mem_frac, verbose=(rank == 0))

    rank_cache = phase1_cache(
        rank, world_size, model, processor, all_needed_files,
        hook_points, args, device, args.question, sae
    )

    # Free GPU memory
    del model, sae
    t.cuda.empty_cache()

    # ── Phase 2: gather pooled features, train probes ───────────────────────────
    if is_distributed:
        all_rank_caches = [None] * world_size
        dist.all_gather_object(all_rank_caches, rank_cache)
    else:
        all_rank_caches = [rank_cache]

    if rank == 0:
        merged_cache = {}
        for rc in all_rank_caches:
            merged_cache.update(rc)

        phase2_data = prepare_phase2_data(merged_cache, train_ids_labels, val_ids_labels, hook_points,
                                          verbose=args.dry_run)
        phase2_train_probes(rank, hook_points, args.layers, phase2_data, args, device)

    cleanup_distributed()


if __name__ == "__main__":
    main()