#!/usr/bin/env python3
"""Downstream causal evaluation for Matrix SAEs.

Measures how well SAE-reconstructed GDN recurrent states preserve model behavior.
Split each sequence into a prefix (builds up state) and suffix (measured).
Replace the recurrent state for one layer with the SAE reconstruction, then
compare suffix perplexity against the unmodified baseline.

This isolates the SAE's information loss: if reconstruction is perfect, suffix
perplexity matches baseline exactly.

Architecture note: Qwen3.5's GDN layers only use cached recurrent states when
seq_len == 1 (token-by-token decoding). For multi-token suffix processing, we
patch chunk_gated_delta_rule to inject the initial state directly.
"""

import json
import time
from contextlib import contextmanager
from typing import TypedDict
from pathlib import Path

import numpy as np
import torch
from tqdm import tqdm  # type: ignore[import-untyped]

from sae import BilinearMatrixSAE, FlatSAE, MatrixSAE, build_sae_from_config

TrainableSAE = FlatSAE | MatrixSAE | BilinearMatrixSAE


# Runtime container holding a loaded SAE plus tag/type/train-mse metadata.
# Distinct from core.types.SAEConfig (training/checkpoint config schema).
class LoadedSAEEntry(TypedDict):
    tag: str
    sae: TrainableSAE
    sae_type: str
    train_mse: float | None


def load_sae_from_checkpoint(
    ckpt_path: str,
    config_path: str,
    device: str = "cuda",
):
    """Load a trained SAE from checkpoint, returning (sae, cfg, train_mse)."""
    cfg = json.loads(Path(config_path).read_text())
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    sae = build_sae_from_config(cfg, state_dict=ckpt["model_state_dict"])

    sae.load_state_dict(ckpt["model_state_dict"])
    sae = sae.to(device).eval()

    train_mse = ckpt.get("val_mse")
    return sae, cfg, train_mse


def reconstruct_state_head(sae: TrainableSAE, state_head: torch.Tensor, sae_type: str) -> torch.Tensor:
    """Reconstruct a single head's state (d_k, d_v) through the SAE.

    Args:
        sae: trained SAE model
        state_head: (d_k, d_v) tensor, the recurrent state for one head
        sae_type: "flat", "rank1", "bilinear", or "bilinear_tied"

    Returns:
        Reconstructed (d_k, d_v) tensor
    """
    d_k, d_v = state_head.shape
    if sae_type == "flat":
        x = state_head.reshape(1, d_k * d_v)
        out = sae(x)
        return out.reconstruction.reshape(d_k, d_v)
    else:
        x = state_head.unsqueeze(0)
        out = sae(x)
        return out.reconstruction.squeeze(0)


@contextmanager
def _patch_gdn_initial_states(model, gdn_layer_indices: list[int], states: dict[int, torch.Tensor]):
    """Temporarily patch GDN layers' chunk_gated_delta_rule to inject initial states.

    The Qwen3.5 GDN layer calls chunk_gated_delta_rule with initial_state=None
    when seq_len > 1. This patches all specified GDN layers so the suffix pass
    correctly continues from the prefix's accumulated recurrent states.

    Args:
        model: HuggingFace CausalLM with model.model.layers
        gdn_layer_indices: list of layer indices to patch
        states: {layer_idx: (batch, n_heads, d_k, d_v)} tensors to inject
    """
    model_layers = model.model.layers if hasattr(model, "model") else model.layers
    originals = {}

    for idx in gdn_layer_indices:
        if idx not in states:
            continue
        gdn = model_layers[idx].linear_attn
        originals[idx] = gdn.chunk_gated_delta_rule
        state = states[idx]

        def _make_patched(fn, s):
            def patched(*args, **kwargs):
                kwargs["initial_state"] = s
                return fn(*args, **kwargs)
            return patched

        gdn.chunk_gated_delta_rule = _make_patched(originals[idx], state)

    try:
        yield
    finally:
        for idx, fn in originals.items():
            model_layers[idx].linear_attn.chunk_gated_delta_rule = fn


@torch.no_grad()
def compute_split_perplexity(
    model,
    input_ids: torch.Tensor,
    split_pos: int,
    gdn_layer_indices: list[int],
    target_layer_idx: int | None = None,
    sae: TrainableSAE | None = None,
    sae_type: str | None = None,
    head_idx: int = 0,
) -> dict[str, float | int]:
    """Compute perplexity on the suffix of a sequence, optionally with SAE state reconstruction.

    Protocol:
      1. Run prefix (tokens[:split_pos]) with use_cache=True to build all caches
      2. Extract recurrent states for ALL GDN layers from the cache
      3. If sae is provided, reconstruct target_layer's head_idx state through the SAE
      4. Patch ALL GDN layers' chunk_gated_delta_rule to inject their prefix states
      5. Run suffix with the patched model + attention KV cache
      6. Compute cross-entropy loss on suffix tokens

    All GDN layers are patched so the suffix pass correctly continues from each
    layer's accumulated state. Only the target layer gets SAE reconstruction.

    Args:
        model: HuggingFace CausalLM model
        input_ids: (batch, seq_len) token ids
        split_pos: position to split prefix/suffix
        gdn_layer_indices: all GDN layer indices in the model
        target_layer_idx: GDN layer to apply SAE on (None = no SAE, pure baseline)
        sae: trained SAE model (None = no modification)
        sae_type: SAE type string for dispatch
        head_idx: which head to reconstruct (default 0)

    Returns:
        dict with keys: loss, perplexity, n_tokens
    """
    batch_size, seq_len = input_ids.shape
    assert split_pos > 0 and split_pos < seq_len

    prefix = input_ids[:, :split_pos]
    suffix = input_ids[:, split_pos:]

    prefix_out = model(input_ids=prefix, use_cache=True)
    cache = prefix_out.past_key_values

    gdn_states = {}
    for idx in gdn_layer_indices:
        layer_cache = cache.layers[idx]
        if hasattr(layer_cache, "recurrent_states") and layer_cache.recurrent_states is not None:
            gdn_states[idx] = layer_cache.recurrent_states.clone()

    if sae is not None and target_layer_idx is not None and sae_type is not None:
        if target_layer_idx in gdn_states:
            state = gdn_states[target_layer_idx]
            for b in range(batch_size):
                original_head = state[b, head_idx].float()
                reconstructed = reconstruct_state_head(sae, original_head, sae_type)
                state[b, head_idx] = reconstructed.to(state.dtype)

    with _patch_gdn_initial_states(model, gdn_layer_indices, gdn_states):
        suffix_out = model(
            input_ids=suffix,
            past_key_values=cache,
            use_cache=False,
            labels=suffix,
        )

    loss = suffix_out.loss
    n_tokens = suffix.shape[1] - 1

    return {
        "loss": loss.item(),
        "perplexity": float(torch.exp(loss).item()),
        "n_tokens": n_tokens,
    }


def _get_gdn_layer_indices(model) -> list[int]:
    """Get indices of GDN (linear_attention) layers from model config."""
    config = getattr(model.config, "text_config", model.config)
    return [i for i, t in enumerate(config.layer_types) if t == "linear_attention"]


@torch.no_grad()
def evaluate_downstream(
    model,
    tokenizer,
    corpus_batches: list[torch.Tensor],
    layer_idx: int,
    sae_configs: list[LoadedSAEEntry],
    head_idx: int = 0,
    split_fraction: float = 0.5,
    device: str = "cuda",
) -> dict[str, object]:
    """Run the full downstream evaluation for one layer across multiple SAEs.

    Args:
        model: HuggingFace CausalLM
        tokenizer: tokenizer (unused, kept for API consistency)
        corpus_batches: list of (batch_size, seq_len) token tensors
        layer_idx: GDN layer index to evaluate
        sae_configs: list of dicts with keys: tag, sae, sae_type, train_mse
        head_idx: which head SAEs are trained on
        split_fraction: fraction of sequence used as prefix (default 0.5)
        device: compute device

    Returns:
        dict with baseline and per-SAE results
    """
    seq_len = corpus_batches[0].shape[1]
    split_pos = int(seq_len * split_fraction)
    n_sequences = sum(b.shape[0] for b in corpus_batches)
    gdn_layers = _get_gdn_layer_indices(model)

    print(f"Downstream evaluation: layer {layer_idx}, head {head_idx}")
    print(f"  {n_sequences} sequences x {seq_len} tokens, split at position {split_pos}")
    print(f"  {len(gdn_layers)} GDN layers patched for state forwarding")
    print(f"  Evaluating {len(sae_configs)} SAE checkpoints + baseline")

    results: dict[str, object] = {
        "layer": layer_idx,
        "head": head_idx,
        "seq_len": seq_len,
        "split_pos": split_pos,
        "n_sequences": n_sequences,
        "gdn_layers": gdn_layers,
    }

    print("\nComputing baseline perplexity (original state, no SAE)...")
    baseline_losses = []
    baseline_tokens = 0
    t0 = time.time()

    for batch in tqdm(corpus_batches, desc="Baseline"):
        batch = batch.to(device)
        for i in range(batch.shape[0]):
            seq = batch[i:i+1]
            r = compute_split_perplexity(
                model, seq, split_pos,
                gdn_layer_indices=gdn_layers,
            )
            baseline_losses.append(r["loss"] * r["n_tokens"])
            baseline_tokens += r["n_tokens"]
        torch.cuda.empty_cache()

    baseline_avg_loss = sum(baseline_losses) / max(baseline_tokens, 1)
    baseline_ppl = float(np.exp(baseline_avg_loss))
    baseline_time = time.time() - t0

    results["baseline"] = {
        "loss": baseline_avg_loss,
        "perplexity": baseline_ppl,
        "n_tokens": baseline_tokens,
        "time_s": round(baseline_time, 1),
    }
    print(f"  Baseline perplexity: {baseline_ppl:.2f} (loss={baseline_avg_loss:.4f}, {baseline_time:.1f}s)")

    results["sae_results"] = {}

    for sae_cfg in sae_configs:
        tag = sae_cfg["tag"]
        sae = sae_cfg["sae"]
        sae_type = sae_cfg["sae_type"]
        train_mse = sae_cfg.get("train_mse")

        print(f"\nEvaluating: {tag}")
        sae_losses = []
        sae_tokens = 0
        t0 = time.time()

        for batch in tqdm(corpus_batches, desc=f"  {tag}"):
            batch = batch.to(device)
            for i in range(batch.shape[0]):
                seq = batch[i:i+1]
                r = compute_split_perplexity(
                    model, seq, split_pos,
                    gdn_layer_indices=gdn_layers,
                    target_layer_idx=layer_idx, sae=sae,
                    sae_type=sae_type, head_idx=head_idx,
                )
                sae_losses.append(r["loss"] * r["n_tokens"])
                sae_tokens += r["n_tokens"]
            torch.cuda.empty_cache()

        sae_avg_loss = sum(sae_losses) / max(sae_tokens, 1)
        sae_ppl = float(np.exp(sae_avg_loss))
        delta_pct = (sae_ppl - baseline_ppl) / baseline_ppl * 100
        sae_time = time.time() - t0

        results["sae_results"][tag] = {
            "loss": sae_avg_loss,
            "perplexity": sae_ppl,
            "delta_pct": delta_pct,
            "train_mse": train_mse,
            "n_tokens": sae_tokens,
            "time_s": round(sae_time, 1),
        }
        print(f"  Perplexity: {sae_ppl:.2f} (delta={delta_pct:+.2f}%, loss={sae_avg_loss:.4f}, {sae_time:.1f}s)")

    return results


@torch.no_grad()
def compute_split_perplexity_allheads(
    model,
    input_ids: torch.Tensor,
    split_pos: int,
    gdn_layer_indices: list[int],
    target_layer_idx: int | None = None,
    sae: TrainableSAE | None = None,
    sae_type: str | None = None,
    n_heads: int = 32,
) -> dict[str, float | int]:
    """Like compute_split_perplexity but reconstructs ALL heads through one SAE.

    Uses the same head-0-trained SAE for every head. This amplifies the
    structured-vs-flat MSE gap by 32x compared to single-head replacement.

    Args:
        model: HuggingFace CausalLM model
        input_ids: (batch, seq_len) token ids
        split_pos: position to split prefix/suffix
        gdn_layer_indices: all GDN layer indices in the model
        target_layer_idx: GDN layer to apply SAE on (None = no SAE, pure baseline)
        sae: trained SAE model (None = no modification)
        sae_type: SAE type string for dispatch
        n_heads: number of heads to reconstruct

    Returns:
        dict with keys: loss, perplexity, n_tokens
    """
    batch_size, seq_len = input_ids.shape
    assert split_pos > 0 and split_pos < seq_len

    prefix = input_ids[:, :split_pos]
    suffix = input_ids[:, split_pos:]

    prefix_out = model(input_ids=prefix, use_cache=True)
    cache = prefix_out.past_key_values

    gdn_states = {}
    for idx in gdn_layer_indices:
        layer_cache = cache.layers[idx]
        if hasattr(layer_cache, "recurrent_states") and layer_cache.recurrent_states is not None:
            gdn_states[idx] = layer_cache.recurrent_states.clone()

    if sae is not None and target_layer_idx is not None and sae_type is not None:
        if target_layer_idx in gdn_states:
            state = gdn_states[target_layer_idx]
            actual_n_heads = min(n_heads, state.shape[1])
            for b in range(batch_size):
                for h in range(actual_n_heads):
                    original_head = state[b, h].float()
                    reconstructed = reconstruct_state_head(sae, original_head, sae_type)
                    state[b, h] = reconstructed.to(state.dtype)

    with _patch_gdn_initial_states(model, gdn_layer_indices, gdn_states):
        suffix_out = model(
            input_ids=suffix,
            past_key_values=cache,
            use_cache=False,
            labels=suffix,
        )

    loss = suffix_out.loss
    n_tokens = suffix.shape[1] - 1

    return {
        "loss": loss.item(),
        "perplexity": float(torch.exp(loss).item()),
        "n_tokens": n_tokens,
    }


@torch.no_grad()
def evaluate_downstream_allheads(
    model,
    tokenizer,
    corpus_batches: list[torch.Tensor],
    layer_idx: int,
    sae_configs: list[LoadedSAEEntry],
    n_heads: int = 32,
    split_fraction: float = 0.5,
    device: str = "cuda",
) -> dict[str, object]:
    """Run downstream evaluation replacing ALL heads through each SAE.

    Same protocol as evaluate_downstream but reconstructs every head's state
    through the (head-0-trained) SAE. This amplifies reconstruction error
    into a larger perplexity delta.

    Args:
        model: HuggingFace CausalLM
        tokenizer: tokenizer (unused, kept for API consistency)
        corpus_batches: list of (batch_size, seq_len) token tensors
        layer_idx: GDN layer index to evaluate
        sae_configs: list of dicts with keys: tag, sae, sae_type, train_mse
        n_heads: number of heads to reconstruct (default 32)
        split_fraction: fraction of sequence used as prefix (default 0.5)
        device: compute device

    Returns:
        dict with baseline and per-SAE results
    """
    seq_len = corpus_batches[0].shape[1]
    split_pos = int(seq_len * split_fraction)
    n_sequences = sum(b.shape[0] for b in corpus_batches)
    gdn_layers = _get_gdn_layer_indices(model)

    print(f"Downstream evaluation (ALL HEADS): layer {layer_idx}, {n_heads} heads")
    print(f"  {n_sequences} sequences x {seq_len} tokens, split at position {split_pos}")
    print(f"  {len(gdn_layers)} GDN layers patched for state forwarding")
    print(f"  Evaluating {len(sae_configs)} SAE checkpoints + baseline")

    results: dict[str, object] = {
        "layer": layer_idx,
        "n_heads_replaced": n_heads,
        "all_heads": True,
        "shared_head0_sae_across_heads": True,
        "approximation_note": (
            "Each checkpoint was trained on head 0 and then applied to every head in the "
            "target layer. Treat this as a stress test, not as per-head matched reconstruction."
        ),
        "seq_len": seq_len,
        "split_pos": split_pos,
        "n_sequences": n_sequences,
        "gdn_layers": gdn_layers,
    }

    print("\nComputing baseline perplexity (original state, no SAE)...")
    baseline_losses = []
    baseline_tokens = 0
    t0 = time.time()

    for batch in tqdm(corpus_batches, desc="Baseline"):
        batch = batch.to(device)
        for i in range(batch.shape[0]):
            seq = batch[i:i+1]
            r = compute_split_perplexity_allheads(
                model, seq, split_pos,
                gdn_layer_indices=gdn_layers,
            )
            baseline_losses.append(r["loss"] * r["n_tokens"])
            baseline_tokens += r["n_tokens"]
        torch.cuda.empty_cache()

    baseline_avg_loss = sum(baseline_losses) / max(baseline_tokens, 1)
    baseline_ppl = float(np.exp(baseline_avg_loss))
    baseline_time = time.time() - t0

    results["baseline"] = {
        "loss": baseline_avg_loss,
        "perplexity": baseline_ppl,
        "n_tokens": baseline_tokens,
        "time_s": round(baseline_time, 1),
    }
    print(f"  Baseline perplexity: {baseline_ppl:.2f} (loss={baseline_avg_loss:.4f}, {baseline_time:.1f}s)")

    results["sae_results"] = {}

    for sae_cfg in sae_configs:
        tag = sae_cfg["tag"]
        sae = sae_cfg["sae"]
        sae_type = sae_cfg["sae_type"]
        train_mse = sae_cfg.get("train_mse")

        print(f"\nEvaluating (all {n_heads} heads): {tag}")
        sae_losses = []
        sae_tokens = 0
        t0 = time.time()

        for batch in tqdm(corpus_batches, desc=f"  {tag}"):
            batch = batch.to(device)
            for i in range(batch.shape[0]):
                seq = batch[i:i+1]
                r = compute_split_perplexity_allheads(
                    model, seq, split_pos,
                    gdn_layer_indices=gdn_layers,
                    target_layer_idx=layer_idx, sae=sae,
                    sae_type=sae_type, n_heads=n_heads,
                )
                sae_losses.append(r["loss"] * r["n_tokens"])
                sae_tokens += r["n_tokens"]
            torch.cuda.empty_cache()

        sae_avg_loss = sum(sae_losses) / max(sae_tokens, 1)
        sae_ppl = float(np.exp(sae_avg_loss))
        delta_pct = (sae_ppl - baseline_ppl) / baseline_ppl * 100
        sae_time = time.time() - t0

        results["sae_results"][tag] = {
            "loss": sae_avg_loss,
            "perplexity": sae_ppl,
            "delta_pct": delta_pct,
            "train_mse": train_mse,
            "n_tokens": sae_tokens,
            "time_s": round(sae_time, 1),
        }
        print(f"  Perplexity: {sae_ppl:.2f} (delta={delta_pct:+.2f}%, loss={sae_avg_loss:.4f}, {sae_time:.1f}s)")

    return results


@torch.no_grad()
def compute_split_perplexity_perhead_matched(
    model,
    input_ids: torch.Tensor,
    split_pos: int,
    gdn_layer_indices: list[int],
    target_layer_idx: int | None = None,
    head_saes: dict[int, tuple] | None = None,
    n_heads: int = 16,
) -> dict:
    """Like compute_split_perplexity_allheads but uses a MATCHED SAE per head.

    Instead of one head-0 SAE applied to all heads, each head gets its own
    SAE trained on that head's data. This gives a clean per-head reconstruction
    without cross-head spectral mismatch.

    Args:
        model: HuggingFace CausalLM model
        input_ids: (batch, seq_len) token ids
        split_pos: position to split prefix/suffix
        gdn_layer_indices: all GDN layer indices in the model
        target_layer_idx: GDN layer to apply SAEs on (None = no SAE, pure baseline)
        head_saes: {head_idx: (sae, sae_type)} dict mapping each head to its SAE
        n_heads: number of heads to reconstruct

    Returns:
        dict with keys: loss, perplexity, n_tokens
    """
    batch_size, seq_len = input_ids.shape
    assert split_pos > 0 and split_pos < seq_len

    prefix = input_ids[:, :split_pos]
    suffix = input_ids[:, split_pos:]

    prefix_out = model(input_ids=prefix, use_cache=True)
    cache = prefix_out.past_key_values

    gdn_states = {}
    for idx in gdn_layer_indices:
        layer_cache = cache.layers[idx]
        if hasattr(layer_cache, "recurrent_states") and layer_cache.recurrent_states is not None:
            gdn_states[idx] = layer_cache.recurrent_states.clone()

    if head_saes is not None and target_layer_idx is not None:
        if target_layer_idx in gdn_states:
            state = gdn_states[target_layer_idx]
            actual_n_heads = min(n_heads, state.shape[1])
            for b in range(batch_size):
                for h in range(actual_n_heads):
                    if h not in head_saes:
                        continue
                    sae, sae_type = head_saes[h]
                    original_head = state[b, h].float()
                    reconstructed = reconstruct_state_head(sae, original_head, sae_type)
                    state[b, h] = reconstructed.to(state.dtype)

    with _patch_gdn_initial_states(model, gdn_layer_indices, gdn_states):
        suffix_out = model(
            input_ids=suffix,
            past_key_values=cache,
            use_cache=False,
            labels=suffix,
        )

    loss = suffix_out.loss
    n_tokens = suffix.shape[1] - 1

    return {
        "loss": loss.item(),
        "perplexity": float(torch.exp(loss).item()),
        "n_tokens": n_tokens,
    }


@torch.no_grad()
def evaluate_downstream_perhead_matched(
    model,
    tokenizer,
    corpus_batches: list[torch.Tensor],
    layer_idx: int,
    sae_type_configs: dict[str, dict[int, tuple]],
    n_heads: int = 16,
    split_fraction: float = 0.5,
    device: str = "cuda",
    baseline_result: dict | None = None,
) -> dict:
    """Run downstream evaluation with per-head matched SAEs.

    Each SAE type gets 16 SAEs (one per head), each trained on that head's data.
    All 16 heads are reconstructed simultaneously through their matched SAEs.

    Args:
        model: HuggingFace CausalLM
        tokenizer: tokenizer (unused, kept for API consistency)
        corpus_batches: list of (batch_size, seq_len) token tensors
        layer_idx: GDN layer index to evaluate
        sae_type_configs: {sae_type_tag: {head_idx: (sae, sae_type_str)}}
            e.g. {"flat": {0: (sae0, "flat"), 1: (sae1, "flat"), ...}}
        n_heads: number of heads (default 16)
        split_fraction: fraction of sequence used as prefix (default 0.5)
        device: compute device

    Returns:
        dict with baseline and per-SAE-type results
    """
    seq_len = corpus_batches[0].shape[1]
    split_pos = int(seq_len * split_fraction)
    n_sequences = sum(b.shape[0] for b in corpus_batches)
    gdn_layers = _get_gdn_layer_indices(model)

    print(f"Downstream evaluation (PER-HEAD MATCHED): layer {layer_idx}, {n_heads} heads")
    print(f"  {n_sequences} sequences x {seq_len} tokens, split at position {split_pos}")
    print(f"  {len(gdn_layers)} GDN layers patched for state forwarding")
    print(f"  Evaluating {len(sae_type_configs)} SAE types + baseline")

    results: dict[str, object] = {
        "layer": layer_idx,
        "n_heads_replaced": n_heads,
        "all_heads": True,
        "per_head_matched": True,
        "seq_len": seq_len,
        "split_pos": split_pos,
        "n_sequences": n_sequences,
        "gdn_layers": gdn_layers,
    }

    if baseline_result is None:
        print("\nComputing baseline perplexity (original state, no SAE)...")
        baseline_losses = []
        baseline_tokens = 0
        t0 = time.time()

        for batch in tqdm(corpus_batches, desc="Baseline"):
            batch = batch.to(device)
            for i in range(batch.shape[0]):
                seq = batch[i:i+1]
                r = compute_split_perplexity_perhead_matched(
                    model, seq, split_pos,
                    gdn_layer_indices=gdn_layers,
                )
                baseline_losses.append(r["loss"] * r["n_tokens"])
                baseline_tokens += r["n_tokens"]
            torch.cuda.empty_cache()

        baseline_avg_loss = sum(baseline_losses) / max(baseline_tokens, 1)
        baseline_ppl = float(np.exp(baseline_avg_loss))
        baseline_time = time.time() - t0

        results["baseline"] = {
            "loss": baseline_avg_loss,
            "perplexity": baseline_ppl,
            "n_tokens": baseline_tokens,
            "time_s": round(baseline_time, 1),
        }
        print(
            f"  Baseline perplexity: {baseline_ppl:.2f} "
            f"(loss={baseline_avg_loss:.4f}, {baseline_time:.1f}s)"
        )
    else:
        baseline_avg_loss = float(baseline_result["loss"])
        baseline_ppl = float(baseline_result["perplexity"])
        results["baseline"] = dict(baseline_result)
        print(
            f"\nReusing baseline perplexity: {baseline_ppl:.2f} "
            f"(loss={baseline_avg_loss:.4f})"
        )

    results["sae_results"] = {}

    for tag, head_saes in sae_type_configs.items():
        n_loaded = len(head_saes)
        print(f"\nEvaluating (per-head matched, {n_loaded}/{n_heads} heads): {tag}")
        sae_losses = []
        sae_tokens = 0
        t0 = time.time()

        for batch in tqdm(corpus_batches, desc=f"  {tag}"):
            batch = batch.to(device)
            for i in range(batch.shape[0]):
                seq = batch[i:i+1]
                r = compute_split_perplexity_perhead_matched(
                    model, seq, split_pos,
                    gdn_layer_indices=gdn_layers,
                    target_layer_idx=layer_idx,
                    head_saes=head_saes,
                    n_heads=n_heads,
                )
                sae_losses.append(r["loss"] * r["n_tokens"])
                sae_tokens += r["n_tokens"]
            torch.cuda.empty_cache()

        sae_avg_loss = sum(sae_losses) / max(sae_tokens, 1)
        sae_ppl = float(np.exp(sae_avg_loss))
        delta_pct = (sae_ppl - baseline_ppl) / baseline_ppl * 100
        sae_time = time.time() - t0

        results["sae_results"][tag] = {
            "loss": sae_avg_loss,
            "perplexity": sae_ppl,
            "delta_pct": delta_pct,
            "train_mse": None,
            "n_heads_loaded": n_loaded,
            "n_tokens": sae_tokens,
            "time_s": round(sae_time, 1),
        }
        print(f"  Perplexity: {sae_ppl:.2f} (delta={delta_pct:+.2f}%, loss={sae_avg_loss:.4f}, {sae_time:.1f}s)")

    return results


def format_results_table(results: dict) -> str:
    """Format evaluation results as a readable table."""
    lines = []
    layer = results["layer"]
    n_seq = results["n_sequences"]
    seq_len = results["seq_len"]
    split_pos = results["split_pos"]

    if results.get("all_heads"):
        n_heads = results.get("n_heads_replaced", "?")
        if results.get("per_head_matched"):
            lines.append(f"Downstream Evaluation (PER-HEAD MATCHED, {n_heads} HEADS): layer {layer}, {n_seq} sequences x {seq_len} tokens (split at {split_pos})")
        else:
            lines.append(f"Downstream Evaluation (ALL {n_heads} HEADS): layer {layer}, {n_seq} sequences x {seq_len} tokens (split at {split_pos})")
        if results.get("approximation_note"):
            lines.append(f"Note: {results['approximation_note']}")
    else:
        lines.append(f"Downstream Evaluation: layer {layer}, {n_seq} sequences x {seq_len} tokens (split at {split_pos})")
    lines.append(f"Baseline perplexity: {results['baseline']['perplexity']:.2f}")
    lines.append("")
    lines.append(f"{'SAE':<45} | {'Perplexity':>10} | {'Delta (%)':>10} | {'MSE (train)':>12}")
    lines.append("-" * 85)

    for tag, r in sorted(results["sae_results"].items()):
        mse_str = f"{r['train_mse']:.2e}" if r["train_mse"] is not None else "N/A"
        lines.append(
            f"{tag:<45} | {r['perplexity']:>10.2f} | {r['delta_pct']:>+9.2f}% | {mse_str:>12}"
        )

    return "\n".join(lines)