3️⃣ Attribution graphs

# Load huggingface token
load_dotenv(dotenv_path=str(exercises_dir / ".env"))
HF_TOKEN = os.getenv("HF_TOKEN")
assert HF_TOKEN, "Please set HF_TOKEN in your chapter1_transformer_interp/exercises/.env file"

# Load gemma model using HF token
gemma = HookedSAETransformer.from_pretrained(
    "google/gemma-3-1b-it",
    device=device,
    dtype=dtype,
    fold_ln=False,
    center_writing_weights=False,
    center_unembed=False,
)
gemma.set_use_hook_mlp_in(True)

# Check this version of SAELens has the gemmascope transcoders
assert "gemma-scope-2-1b-it-transcoders-all" in get_pretrained_saes_directory()
n_layers = 26
n_saes_per_layer = 8  # 2 widths, 2 L0s, 2 (affine vs non-affine)
transcoders_map = get_pretrained_saes_directory()["gemma-scope-2-1b-it-transcoders-all"].saes_map
assert len(transcoders_map) == n_layers * n_saes_per_layer

# Load transcoders for all layers (parallel I/O, sequential GPU transfer)
n_layers = gemma.cfg.n_layers
sae_ids = [f"layer_{layer}_width_16k_l0_small_affine" for layer in range(n_layers)]
transcoders: dict[int, SAE] = utils.load_saes_parallel(
    release="gemma-scope-2-1b-it-transcoders-all",
    sae_ids=sae_ids,
    device=device,
    dtype=dtype,
    max_workers=2,
)
# Correct the hook names (temporary until this is fixed)
for layer in range(n_layers):
    transcoders[layer].cfg.metadata.hook_name = f"blocks.{layer}.mlp.hook_in"

def format_prompt(user_prompt: str, model_response: str) -> str:
    """Format a prompt for Gemma IT models using the chat template."""
    return f"<start_of_turn>user\n{user_prompt}<end_of_turn>\n<start_of_turn>model\n{model_response}"


START_POSN = 4  # Mask [BOS, <start_of_turn>, user, \n]

# Example: rhyming couplet prompt from Anthropic's attribution graph work
prompt = format_prompt(
    "Write me a short rhyming couplet.",
    "The sun descends, a golden hue,\nAs evening whispers, soft and",  # (...true)
)

# Tokenize and verify
tokens = gemma.to_tokens(prompt)
str_tokens = gemma.to_str_tokens(prompt)
print(f"Prompt has {len(str_tokens)} tokens")
print(f"First {START_POSN} tokens (masked): {str_tokens[:START_POSN]}")
print(f"Remaining tokens: {str_tokens[START_POSN:]}")

class FreezeHooks:
    """
    Installs forward hooks that freeze attention patterns and LayerNorm scales at their
    forward-pass values, making these operations linear.

    Uses TransformerLens-native hooks (not PyTorch register_forward_hook). Two interfaces:

    1. Context manager (for direct use):
        with freeze:
            model(tokens)  # patterns/scales replaced by frozen values

    2. fwd_hooks property (for combining with other hooks via model.hooks()):
        with model.hooks(fwd_hooks=freeze.fwd_hooks + other_hooks, bwd_hooks=...):
            model(tokens).backward()
    """

    def __init__(self, model: HookedSAETransformer):
        self.model = model
        self.frozen_values: dict[str, Tensor] = {}

    def _freeze_hook(self, value: Tensor, hook: HookPoint) -> Tensor:
        """
        TL-native hook: replaces activations with their frozen values.

        Handles batched forward passes (batch_size > 1) even though frozen values were cached
        with batch_size=1. Two cases arise:

        - Regular hooks (hook_pattern, ln.hook_scale): frozen shape is [1, ...], so we repeat
          along dim 0 to get [batch, ...].
        - QK-norm hook_scale: the batch dim is folded into dim 0 as [batch*pos*n_heads, 1],
          so the frozen shape is [pos*n_heads, 1]. We repeat by the batch ratio to get
          [batch*pos*n_heads, 1].

        In both cases the ratio value.shape[0] // frozen.shape[0] gives the right repeat factor.
        """
        frozen = self.frozen_values[hook.name]
        if frozen.shape[0] != value.shape[0] and value.shape[0] % frozen.shape[0] == 0:
            ratio = value.shape[0] // frozen.shape[0]
            frozen = frozen.repeat(ratio, *([1] * (frozen.dim() - 1)))
        return frozen

    @property
    def fwd_hooks(self) -> list[tuple[str, Callable]]:
        """
        Return freeze hooks as a list of (hook_name, hook_fn) pairs.

        Use this with model.hooks() to combine freeze hooks with other hooks in a
        single context manager (Option A), avoiding nesting issues:

            with model.hooks(fwd_hooks=freeze.fwd_hooks + capture_hooks, bwd_hooks=...):
                model(tokens).backward()
        """
        return [(name, self._freeze_hook) for name in self.frozen_values]

    def cache_frozen_values(self, tokens: Tensor) -> ActivationCache:
        """Runs a forward pass, caching values we'll freeze later, plus other useful activations."""
        # We cache attention patterns, LN scales, and residual stream values
        names_filter = lambda name: any(
            s in name for s in ["hook_pattern", "hook_scale", "resid_post", "resid_pre", "mlp.hook_in", "mlp_out"]
        )
        _, cache = self.model.run_with_cache(tokens, names_filter=names_filter)

        # Store frozen values (patterns and scales)
        self.frozen_values = {
            name: cache[name].detach() for name in cache.keys() if "hook_pattern" in name or "hook_scale" in name
        }
        return cache

    def __enter__(self):
        """
        Install freeze hooks using TransformerLens's model.add_hook() (non-permanent).

        Non-permanent hooks are removed by model.reset_hooks(including_permanent=False),
        which is what __exit__ calls. TranscoderReplacementHooks uses permanent hooks, so
        they survive this reset.
        """
        for name in self.frozen_values:
            self.model.add_hook(name, self._freeze_hook)
        return self

    def __exit__(self, *args):
        """Remove all non-permanent hooks (freeze hooks only; permanent tc_hooks survive)."""
        self.model.reset_hooks(including_permanent=False)

skip = ln2(resid_mid) @ W_skip
mlp_out_new = skip + (mlp_out_original - skip).detach()

from sae_lens import JumpReLUSkipTranscoder, JumpReLUTranscoder

Transcoder: TypeAlias = JumpReLUTranscoder | JumpReLUSkipTranscoder


class TranscoderReplacementHooks:
    """
    Installs hooks that replace MLP backward passes with linear skip connections, while
    computing and storing transcoder feature activations for graph construction.

    The skip connection trick ensures:
    - Forward pass: exact MLP output (unchanged)
    - Backward pass: gradients flow through linear skip connection (W_skip)

    Uses TransformerLens-native hooks registered as PERMANENT (is_permanent=True), so they
    survive FreezeHooks.__exit__ which only resets non-permanent hooks. Call remove() when
    done to clear permanent hooks via model.reset_hooks(including_permanent=True).
    """

    def __init__(
        self,
        model: HookedSAETransformer,
        transcoders: dict[int, "Transcoder"],
        cache: ActivationCache,
    ):
        self.model = model
        self.transcoders = transcoders
        self.cache = cache  # From FreezeHooks.cache_frozen_values
        self.transcoder_acts: dict[int, Tensor] = {}  # layer -> feature activations
        self.transcoder_output: dict[int, Tensor] = {}  # layer -> transcoder reconstruction

        self.current_ln_inputs: dict[int, Tensor] = {}  # updated each forward pass

    # def install(self):
    #     """Install hooks for all transcoders."""
    #     self.current_ln_inputs: dict[int, Tensor] = {}
    #     for layer, tc in self.transcoders.items():
    #         # Get hook names for this layer
    #         input_hook = tc.cfg.metadata.hook_name  # e.g. "blocks.0.mlp.hook_in"
    #         output_hook = tc.cfg.metadata.hook_name_out  # e.g. "blocks.0.hook_mlp_out"
    #
    #         # Compute and store transcoder activations (no gradient needed)
    #         with t.no_grad():
    #             ln_input = self.cache[input_hook]
    #             tc_acts = tc.encode(ln_input)
    #             tc_output = tc.decode(tc_acts)
    #             self.transcoder_acts[layer] = tc_acts.detach()
    #             self.transcoder_output[layer] = tc_output.detach()
    #
    #         # YOUR CODE HERE - register TWO permanent TL-native hooks:
    #         #
    #         # 1. A capture hook at input_hook (mlp.hook_in):
    #         #    Signature: (tensor, hook) -> None
    #         #    Store the live ln_input in self.current_ln_inputs[layer].
    #         #    This MUST use the live tensor, not self.cache[input_hook] - the cached
    #         #    tensor is attached to a stale graph that gets freed after the first backward().
    #         #
    #         # 2. A skip hook at output_hook (hook_mlp_out):
    #         #    Signature: (mlp_out, hook) -> Tensor
    #         #    Apply the skip connection trick using self.current_ln_inputs[layer]:
    #         #       skip = ln_input @ tc.W_skip
    #         #       return skip + (mlp_out - skip).detach()
    #         #
    #         # Register both with is_permanent=True so they survive FreezeHooks context exits.
    #         # Use factory functions to avoid closure issues in the loop.
    #         pass

    def remove(self):
        """Remove all hooks, including permanent ones added by install()."""
        self.model.reset_hooks(including_permanent=True)

# Test: verify the skip connection trick gives correct forward pass
freeze = FreezeHooks(gemma)
cache = freeze.cache_frozen_values(tokens)

tc_hooks = TranscoderReplacementHooks(gemma, transcoders, cache)
tc_hooks.install()

with freeze:
    logits_with_hooks = gemma(tokens)

tc_hooks.remove()

# Compare with original logits
logits_original = gemma(tokens)

print(f"Max difference in logits: {(logits_with_hooks - logits_original).abs().max().item():.6f}")
print("(Should be ~0 since the skip trick preserves forward pass values)")

print("Average L0 per transcoder layer (excluding first 4 tokens):")
for layer in range(len(transcoders)):
    acts = tc_hooks.transcoder_acts[layer][0, START_POSN:, :]  # (seq - START_POSN, d_enc)
    l0_per_pos = (acts > 0).float().sum(dim=-1)  # (seq - START_POSN,)
    avg_l0 = l0_per_pos.mean().item()
    print(f"  Layer {layer}: avg L0 = {avg_l0:.1f}")
    assert avg_l0 <= 50, (
        f"Layer {layer} avg L0 = {avg_l0:.1f} exceeds 50 - check that you loaded the "
        f"'l0_small' transcoder variant, not 'l0_big'"
    )
print("L0 sanity check passed!")

def compute_salient_logits(
    model: HookedSAETransformer,
    logits: Float[Tensor, "batch seq d_vocab"],
    n_output_nodes: int = 3,
) -> tuple[Float[Tensor, "n_output d_model"], list[tuple[str, float]]]:
    """
    Select the top predicted tokens and return their demeaned W_U columns.

    Args:
        model: The transformer model.
        logits: Full logit tensor from model forward pass.
        n_output_nodes: Number of top tokens to select.

    Returns:
        reading_vecs: Demeaned unembedding vectors for top tokens, shape (n_output, d_model).
        top_token_info: List of (token_string, probability) tuples for the selected tokens.
    """
    raise NotImplementedError()

reading_vecs, top_token_info = compute_salient_logits(gemma, logits_original)
print("Top predicted tokens:")
for tok_str, prob in top_token_info:
    print(f"  {tok_str!r}: p={prob:.4f}")
print(f"Reading vectors shape: {reading_vecs.shape}")

class NodeType(Enum):
    EMBEDDING = "embedding"
    LATENT = "latent"
    MLP_ERROR = "mlp_error"
    LOGIT = "logit"


@dataclass
class NodeInfo:
    """Metadata about a single node in the attribution graph."""

    node_type: NodeType
    layer: int | str  # "E" for embeddings, layer_idx for features, "L" for logits
    ctx_idx: int  # sequence position
    feature: int  # feature index (token ID for embeds, feature ID for latents, token ID for logits)
    activation: float = 0.0  # feature activation (for latent nodes)
    token_prob: float = 0.0  # probability (for logit nodes)
    str_token: str = ""  # string representation (for embed/logit nodes)
    label: str = ""  # human-readable label

    @property
    def node_id(self) -> str:
        return f"{self.layer}_{self.feature}_{self.ctx_idx}"


@dataclass
class GraphNodes:
    """Container for all nodes in an attribution graph, with their reading/writing vectors."""

    nodes: list[NodeInfo] = field(default_factory=list)
    writing_vecs: Tensor | None = None  # (n_nodes, d_model) or None
    reading_vecs: Tensor | None = None  # (n_nodes, d_model) or None
    node_range_dict: dict = field(default_factory=dict)  # layer -> (start_idx, end_idx)
    seq_len: int = 0
    n_layers: int = 0

def build_embedding_nodes(
    model: HookedSAETransformer,
    cache: ActivationCache,
    tokens: Int[Tensor, "1 seq"],
) -> tuple[list[NodeInfo], list[Tensor], list[Tensor]]:
    """
    Build embedding nodes (one per token position) with their writing and reading vectors.

    Args:
        model: The transformer model.
        cache: Activation cache from forward pass.
        tokens: Input token IDs, shape (1, seq).

    Returns:
        Tuple of (nodes, writing_vecs, reading_vecs) where each list has length seq_len.
        Writing vecs are the token embeddings, reading vecs are zeros.
    """
    raise NotImplementedError()


embed_nodes, embed_writing, embed_reading = build_embedding_nodes(model=gemma, cache=cache, tokens=tokens)
seq_len = tokens.shape[1]

assert len(embed_nodes) == seq_len, f"Expected {seq_len} embedding nodes, got {len(embed_nodes)}"
assert all(n.node_type == NodeType.EMBEDDING for n in embed_nodes), "All embedding nodes should have type EMBEDDING"
for i, wv in enumerate(embed_writing):
    assert wv.shape == (gemma.cfg.d_model,), f"Writing vec {i} has wrong shape: {wv.shape}"
    assert wv.abs().sum() > 0, f"Writing vec {i} is all zeros but should be a token embedding"
for i, rv in enumerate(embed_reading):
    assert rv.shape == (gemma.cfg.d_model,), f"Reading vec {i} has wrong shape: {rv.shape}"
    t.testing.assert_close(rv, t.zeros_like(rv), msg=f"Reading vec {i} should be all zeros")

print("All build_embedding_nodes tests passed!")

def build_intermediate_nodes(
    transcoders: dict[int, "Transcoder"],
    cache: ActivationCache,
    tc_hooks: TranscoderReplacementHooks,
    tokens: Int[Tensor, "1 seq"],
    d_model: int,
    start_posn: int = 4,
    top_k: int = 5,
) -> tuple[list[NodeInfo], list[Tensor], list[Tensor], dict]:
    """
    Build intermediate (LATENT + MLP_ERROR) nodes for all layers.

    Args:
        transcoders: Dict mapping layer -> transcoder.
        cache: Activation cache from forward pass.
        tc_hooks: TranscoderReplacementHooks with computed activations.
        tokens: Input token IDs, shape (1, seq).
        d_model: Model hidden dimension.
        start_posn: First position to include (masks chat formatting tokens).
        top_k: Number of top-activating features to include per position per layer.

    Returns:
        Tuple of (nodes, writing_vecs, reading_vecs, node_range_dict) where node_range_dict
        maps each layer index to (start_idx, end_idx) within the returned lists.
    """
    seq_len = tokens.shape[1]
    n_layers = len(transcoders)
    device = cache["blocks.0.hook_resid_pre"].device

    nodes: list[NodeInfo] = []
    writing_vecs: list[Tensor] = []
    reading_vecs: list[Tensor] = []
    node_range_dict = {}

    for layer in range(n_layers):
        tc = transcoders[layer]
        layer_start = len(nodes)

        tc_acts = tc_hooks.transcoder_acts[layer][0]  # (seq, d_enc)
        tc_output = tc_hooks.transcoder_output[layer][0]  # (seq, d_model)

        # Get MLP output from cache for error computation
        mlp_output = cache[tc.cfg.metadata.hook_name_out][0]  # (seq, d_model)

        W_dec = tc.W_dec.detach()  # (d_enc, d_model)
        W_enc_T = tc.W_enc.detach().T  # (d_enc, d_model)

        for pos in range(start_posn, seq_len):
            acts = tc_acts[pos]  # (d_enc,)
            # Select top-k features by activation
            k = min(top_k, (acts > 0).sum().item())  # don't include zero-activation features
            if k > 0:
                _, active_indices = acts.topk(k)
            else:
                active_indices = t.where(acts > 0)[0]  # fallback: empty if no active features

            # TODO: For each active feature, create a LATENT node with the correct writing
            # vector (activation * W_dec[feature]) and reading vector (W_enc.T[feature]).
            # After all features at this position, create one MLP_ERROR node with writing
            # vector = mlp_output[pos] - tc_output[pos] and reading vector = zeros.
            pass

        node_range_dict[layer] = (layer_start, len(nodes))

    return nodes, writing_vecs, reading_vecs, node_range_dict


inter_nodes, inter_writing, inter_reading, inter_range_dict = build_intermediate_nodes(
    transcoders=transcoders,
    cache=cache,
    tc_hooks=tc_hooks,
    tokens=tokens,
    d_model=gemma.cfg.d_model,
    start_posn=START_POSN,
)

n_layers = len(transcoders)
seq_len = tokens.shape[1]
active_positions = seq_len - START_POSN

# All LATENT nodes should have positive activations
latent_nodes = [n for n in inter_nodes if n.node_type == NodeType.LATENT]
assert len(latent_nodes) > 0, "Expected at least some LATENT nodes"
assert all(n.activation > 0 for n in latent_nodes), "All LATENT nodes should have positive activations"

# MLP_ERROR nodes: one per position per layer
error_nodes = [n for n in inter_nodes if n.node_type == NodeType.MLP_ERROR]
assert len(error_nodes) == active_positions * n_layers, (
    f"Expected {active_positions * n_layers} MLP_ERROR nodes, got {len(error_nodes)}"
)

# Writing vectors for LATENT nodes should be non-zero
latent_indices = [i for i, n in enumerate(inter_nodes) if n.node_type == NodeType.LATENT]
for idx in latent_indices[:5]:  # check first 5
    assert inter_writing[idx].abs().sum() > 0, "LATENT writing vectors should be non-zero"

# Reading vectors for MLP_ERROR nodes should be zeros
error_indices = [i for i, n in enumerate(inter_nodes) if n.node_type == NodeType.MLP_ERROR]
for idx in error_indices[:5]:  # check first 5
    t.testing.assert_close(inter_reading[idx], t.zeros_like(inter_reading[idx]))

# node_range_dict should have entries for each layer
assert set(inter_range_dict.keys()) == set(range(n_layers)), (
    f"node_range_dict should have entries for layers 0..{n_layers - 1}, got keys {set(inter_range_dict.keys())}"
)

print("All build_intermediate_nodes tests passed!")

def build_logit_nodes(
    reading_vecs_logit: Float[Tensor, "n_output d_model"],
    top_token_info: list[tuple[str, float]],
    seq_len: int,
    d_model: int,
) -> tuple[list[NodeInfo], list[Tensor], list[Tensor]]:
    """
    Build logit (output) nodes with their writing and reading vectors.

    Args:
        reading_vecs_logit: Demeaned W_U columns for output nodes, shape (n_output, d_model).
        top_token_info: List of (token_string, probability) for output nodes.
        seq_len: Total sequence length (logit nodes are placed at seq_len - 1).
        d_model: Model hidden dimension.

    Returns:
        Tuple of (nodes, writing_vecs, reading_vecs) where each list has length len(top_token_info).
        Writing vecs are zeros, reading vecs are the logit reading vectors.
    """
    device = reading_vecs_logit.device

    nodes: list[NodeInfo] = []
    writing_vecs: list[Tensor] = []
    reading_vecs: list[Tensor] = []

    for i, (tok_str, prob) in enumerate(top_token_info):
        nodes.append(
            NodeInfo(
                node_type=NodeType.LOGIT,
                layer="L",
                ctx_idx=seq_len - 1,
                feature=i,
                token_prob=prob,
                str_token=tok_str,
            )
        )
        writing_vecs.append(t.zeros(d_model, device=device))
        reading_vecs.append(reading_vecs_logit[i])

    return nodes, writing_vecs, reading_vecs

def build_graph_nodes(
    model: HookedSAETransformer,
    transcoders: dict[int, "Transcoder"],
    cache: ActivationCache,
    tc_hooks: TranscoderReplacementHooks,
    reading_vecs_logit: Float[Tensor, "n_output d_model"],
    top_token_info: list[tuple[str, float]],
    tokens: Int[Tensor, "1 seq"],
    start_posn: int = 4,
    top_k: int = 5,
) -> GraphNodes:
    """
    Build all graph nodes with their reading and writing vectors by calling the three sub-functions
    and assembling the results.

    Args:
        model: The transformer model.
        transcoders: Dict mapping layer -> transcoder.
        cache: Activation cache from forward pass.
        tc_hooks: TranscoderReplacementHooks with computed activations.
        reading_vecs_logit: Demeaned W_U columns for output nodes.
        top_token_info: List of (token_string, probability) for output nodes.
        tokens: Input token IDs, shape (1, seq).
        start_posn: First position to include (masks chat formatting tokens).
        top_k: Number of top-activating features to include per position per layer.

    Returns:
        GraphNodes containing all nodes, their reading/writing vectors, and index ranges.
    """
    seq_len = tokens.shape[1]
    n_layers = len(transcoders)
    d_model = model.cfg.d_model

    all_nodes: list[NodeInfo] = []
    all_writing: list[Tensor] = []
    all_reading: list[Tensor] = []
    node_range_dict = {}

    # 1. Embedding nodes
    embed_nodes, embed_writing, embed_reading = build_embedding_nodes(model, cache, tokens)
    all_nodes.extend(embed_nodes)
    all_writing.extend(embed_writing)
    all_reading.extend(embed_reading)
    node_range_dict["E"] = (0, len(all_nodes))

    # 2. Intermediate nodes (features + MLP error per layer)
    inter_nodes, inter_writing, inter_reading, inter_ranges = build_intermediate_nodes(
        transcoders,
        cache,
        tc_hooks,
        tokens,
        d_model,
        start_posn,
        top_k,
    )
    offset = len(all_nodes)
    all_nodes.extend(inter_nodes)
    all_writing.extend(inter_writing)
    all_reading.extend(inter_reading)
    for layer_key, (s, e) in inter_ranges.items():
        node_range_dict[layer_key] = (s + offset, e + offset)

    # 3. Logit nodes
    logit_start = len(all_nodes)
    logit_nodes, logit_writing, logit_reading = build_logit_nodes(
        reading_vecs_logit,
        top_token_info,
        seq_len,
        d_model,
    )
    all_nodes.extend(logit_nodes)
    all_writing.extend(logit_writing)
    all_reading.extend(logit_reading)
    node_range_dict["L"] = (logit_start, len(all_nodes))

    return GraphNodes(
        nodes=all_nodes,
        writing_vecs=t.stack(all_writing),
        reading_vecs=t.stack(all_reading),
        node_range_dict=node_range_dict,
        seq_len=seq_len,
        n_layers=n_layers,
    )

graph = build_graph_nodes(
    model=gemma,
    transcoders=transcoders,
    cache=cache,
    tc_hooks=tc_hooks,
    reading_vecs_logit=reading_vecs,
    top_token_info=top_token_info,
    tokens=tokens,
    start_posn=START_POSN,
)

# Print some stats
n_embeds = sum(1 for n in graph.nodes if n.node_type == NodeType.EMBEDDING)
n_latents = sum(1 for n in graph.nodes if n.node_type == NodeType.LATENT)
n_errors = sum(1 for n in graph.nodes if n.node_type == NodeType.MLP_ERROR)
n_logits = sum(1 for n in graph.nodes if n.node_type == NodeType.LOGIT)
print(f"Graph has {len(graph.nodes)} total nodes:")
print(f"  {n_embeds} embedding nodes")
print(f"  {n_latents} latent nodes")
print(f"  {n_errors} MLP error nodes")
print(f"  {n_logits} logit output nodes")
print(f"Writing vectors shape: {graph.writing_vecs.shape}")
print(f"Reading vectors shape: {graph.reading_vecs.shape}")

with model.hooks(fwd_hooks=freeze.fwd_hooks + capture_fwd_hooks, bwd_hooks=capture_bwd_hooks):
    model(tokens).backward()

def setup_attribution(
    model: HookedSAETransformer,
    tokens: Tensor,
    freeze: FreezeHooks,
    target_reading_vecs: Float[Tensor, "batch d_model"],
    target_positions: Int[Tensor, " batch"],
    target_layers: list[int | str],
) -> dict[str, Tensor]:
    """
    Run forward pass through frozen model, inject reading vectors as gradient seeds at
    target positions/layers, and return gradients at all residual stream positions.

    This function handles the "backward pass" part of attribution: for each target node,
    it computes d(objective)/d(resid) at every layer, where objective = dot(resid[target_pos], reading_vec).

    Uses TransformerLens-native hooks via a single model.hooks() context (Option A):
      - fwd_hooks = freeze.fwd_hooks + capture_fwd_hooks
          freeze.fwd_hooks  : replace attention patterns/LN scales with frozen values
          capture_fwd_hooks : store each residual-stream tensor for objective computation
      - bwd_hooks = capture_bwd_hooks
          capture_bwd_hooks : receive gradients directly during backward, no retain_grad() needed

    TranscoderReplacementHooks are permanent and stay active inside this context.

    Args:
        model: The transformer model.
        tokens: Input tokens, shape (1, seq).
        freeze: FreezeHooks with cached frozen values (provides fwd_hooks property).
        target_reading_vecs: Reading vectors for target nodes, shape (batch, d_model).
        target_positions: Sequence positions of target nodes, shape (batch,).
        target_layers: Layer identifiers for target nodes (int for intermediate, "L" for logits).

    Returns:
        grads: Dict mapping hook names -> gradient tensors at each residual stream position.
    """
    # Detach reading vectors - attribution only needs d(objective)/d(resid),
    # not d(objective)/d(W_U). Without this, reading_vecs built from W_U (a
    # parameter with requires_grad=True) carry a grad_fn that gets freed after
    # batch 1's backward(), causing "backward through graph a second time" on
    # subsequent batches.
    target_reading_vecs = target_reading_vecs.detach()

    batch_size = target_reading_vecs.shape[0]

    # Tensors captured during the forward pass (for computing objectives)
    captured_tensors: dict[str, Tensor] = {}
    # Gradients captured directly during the backward pass (the actual output)
    captured_grads: dict[str, Tensor] = {}

    # Names of residual-stream tensors to capture
    capture_names = (
        [f"blocks.{l}.hook_resid_post" for l in range(model.cfg.n_layers)]
        + [f"blocks.{l}.hook_resid_mid" for l in range(model.cfg.n_layers)]
        + ["blocks.0.hook_resid_pre"]
    )

    def make_fwd_capture(name: str) -> Callable:
        """Forward hook: store the activation tensor so we can build objectives from it."""

        def hook_fn(tensor: Tensor, hook: HookPoint) -> None:
            captured_tensors[name] = tensor

        return hook_fn

    def make_bwd_capture(name: str) -> Callable:
        """Backward hook: receive the gradient directly during backward, no retain_grad() needed."""

        def hook_fn(grad: Tensor, hook: HookPoint) -> None:
            captured_grads[name] = grad.detach()

        return hook_fn

    capture_fwd_hooks = [(name, make_fwd_capture(name)) for name in capture_names]
    capture_bwd_hooks = [(name, make_bwd_capture(name)) for name in capture_names]

    # Single model.hooks() context combining freeze hooks + capture hooks (Option A).
    # TranscoderReplacementHooks are permanent and stay active throughout.
    # At context exit, only non-permanent hooks (freeze + capture) are removed.
    with model.hooks(
        fwd_hooks=freeze.fwd_hooks + capture_fwd_hooks,
        bwd_hooks=capture_bwd_hooks,
    ):
        # Forward pass (objectives are computed from captured residual tensors)
        model(tokens.expand(batch_size, -1))

        # Compute objectives for each target node
        # Use torch.stack instead of in-place assignment to a leaf tensor, to
        # avoid autograd issues with in-place ops on zero-initialized tensors.
        objective_terms = []
        for i in range(batch_size):
            pos = target_positions[i]
            layer = target_layers[i]
            if layer == "L":
                # For logit nodes, use the final residual stream (pre-unembedding)
                resid = captured_tensors[f"blocks.{model.cfg.n_layers - 1}.hook_resid_post"][i, pos]
            else:
                # For intermediate nodes, use resid_mid at this layer
                # (after attention, before MLP - this is what the feature reads from)
                resid = captured_tensors[f"blocks.{layer}.hook_resid_mid"][i, pos]
            objective_terms.append((resid * target_reading_vecs[i]).sum())

        # Backward pass - fires bwd hooks, populating captured_grads
        total_objective = t.stack(objective_terms).sum()
        total_objective.backward()

    return captured_grads

def prepare_backward_batches(
    graph: GraphNodes,
    batch_size: int = 8,
) -> list[tuple[int, list[NodeInfo], Tensor, Tensor, list[int | str]]]:
    """
    Group target nodes by layer and split into batches for backward passes.

    Args:
        graph: GraphNodes containing all nodes and their reading/writing vectors.
        batch_size: Number of target nodes to process per backward pass.

    Returns:
        List of (global_start_idx, batch_nodes, reading_vecs, positions, layers) tuples, where:
            global_start_idx: Index of the first node in this batch within graph.nodes
            batch_nodes: List of NodeInfo for this batch
            reading_vecs: Stacked reading vectors, shape (batch, d_model)
            positions: Tensor of sequence positions, shape (batch,)
            layers: List of layer identifiers (int or "L")
    """
    raise NotImplementedError()


batches = prepare_backward_batches(graph, batch_size=8)

# Check we got some batches
assert len(batches) > 0, "Expected at least one batch"

# Verify batch structure
all_target_layers_covered = set()
total_target_nodes = 0
for global_start, batch_nodes, reading_vecs, positions, layers in batches:
    # Reading vectors should have correct shape
    assert reading_vecs.shape[0] == len(batch_nodes), "Reading vecs batch size should match number of nodes"
    assert reading_vecs.shape[1] == gemma.cfg.d_model, (
        f"Reading vecs should have d_model={gemma.cfg.d_model} columns"
    )

    # Positions should match batch size
    assert positions.shape[0] == len(batch_nodes), "Positions batch size should match number of nodes"

    # Track which layers are covered
    for n in batch_nodes:
        layer_key = n.layer
        all_target_layers_covered.add(layer_key)
    total_target_nodes += len(batch_nodes)

# All non-embedding layers should be covered
expected_layers = {k for k in graph.node_range_dict.keys() if k != "E"}
assert all_target_layers_covered == expected_layers, (
    f"Expected layers {expected_layers} to be covered, got {all_target_layers_covered}"
)

# Total target nodes should match all non-embedding nodes
n_non_embed = len(graph.nodes) - (graph.node_range_dict["E"][1] - graph.node_range_dict["E"][0])
assert total_target_nodes == n_non_embed, (
    f"Expected {n_non_embed} total target nodes across batches, got {total_target_nodes}"
)

print("All prepare_backward_batches tests passed!")

def compute_adjacency_matrix(
    model: HookedSAETransformer,
    tokens: Tensor,
    graph: GraphNodes,
    freeze: FreezeHooks,
    tc_hooks: TranscoderReplacementHooks,
    batches: list[tuple[int, list[NodeInfo], Tensor, Tensor, list[int | str]]],
    start_posn: int = 4,
) -> Float[Tensor, "n_nodes n_nodes"]:
    """
    Run backward passes for each batch and contract gradients with writing vectors.

    Args:
        model: The transformer model.
        tokens: Input tokens, shape (1, seq).
        graph: GraphNodes containing all nodes and their reading/writing vectors.
        freeze: FreezeHooks with cached frozen values.
        tc_hooks: TranscoderReplacementHooks with installed hooks.
        batches: Output of prepare_backward_batches.
        start_posn: First position to include (masks formatting tokens).

    Returns:
        Adjacency matrix of shape (n_nodes, n_nodes), where A[j, i] is the edge weight
        from node i to node j.
    """
    # n_nodes = len(graph.nodes)
    # adjacency_matrix = t.zeros(n_nodes, n_nodes, device=tokens.device)
    #
    # for global_start_idx, batch_nodes, target_reading_vecs, target_positions, target_layers in batches:
    #     actual_batch_size = len(batch_nodes)
    #
    #     # Run backward pass to get gradients
    #     tc_hooks.install()
    #     grads = setup_attribution(
    #         model=model,
    #         tokens=tokens,
    #         freeze=freeze,
    #         target_reading_vecs=target_reading_vecs,
    #         target_positions=target_positions,
    #         target_layers=target_layers,
    #     )
    #     tc_hooks.remove()
    #
    #     # YOUR CODE HERE - contract gradients with source writing vectors
    #     # For each source node:
    #     #   1. Determine the correct gradient hook name (see hints above)
    #     #   2. Get the gradient at the source node's position: grads[grad_name][:actual_batch_size, src_node.ctx_idx]
    #     #   3. Dot product with the source's writing vector: graph.writing_vecs[src_idx]
    #     #   4. Zero out edges from masked positions (ctx_idx < start_posn)
    #     #   5. Store in adjacency_matrix[global_start_idx + b, src_idx]
    #     pass
    #
    # return adjacency_matrix


batches = prepare_backward_batches(graph, batch_size=8)
print("Got batches...")
adjacency_matrix = compute_adjacency_matrix(
    model=gemma,
    tokens=tokens,
    graph=graph,
    freeze=freeze,
    tc_hooks=tc_hooks,
    batches=batches,
    start_posn=START_POSN,
)

n_nodes = len(graph.nodes)
assert adjacency_matrix.shape == (n_nodes, n_nodes), (
    f"Expected shape ({n_nodes}, {n_nodes}), got {adjacency_matrix.shape}"
)

# Should be lower triangular (no backward connections in a transformer)
upper_tri_norm = t.triu(adjacency_matrix, diagonal=1).abs().sum().item()
assert upper_tri_norm < 1e-4, f"Upper triangle norm should be ~0, got {upper_tri_norm}"

# Should be sparse (most entries near zero)
sparsity = 1 - (adjacency_matrix.abs() > 1e-6).float().mean().item()
assert sparsity > 0.5, f"Expected sparsity > 50%, got {sparsity:.2%}"

# Should have some non-zero entries (embedding -> first layer features)
n_nonzero = (adjacency_matrix.abs() > 1e-6).sum().item()
assert n_nonzero > 0, "Expected some non-zero entries in adjacency matrix"

print("All compute_adjacency_matrix tests passed!")

def normalize_matrix(
    adjacency_matrix: Float[Tensor, "n_nodes n_nodes"],
) -> Float[Tensor, "n_nodes n_nodes"]:
    """
    Normalise the adjacency matrix row-wise by absolute value sums.

    Each row is divided by the sum of absolute values in that row, so that
    |A_norm[j, :]|.sum() = 1 for each j (where possible).
    """
    abs_matrix = adjacency_matrix.abs()
    row_sums = abs_matrix.sum(dim=1, keepdim=True).clamp(min=1e-8)
    return abs_matrix / row_sums

def compute_influence(
    adjacency_matrix: Float[Tensor, "n_nodes n_nodes"],
    logit_weights: Float[Tensor, "n_logit_nodes"],
    n_layers: int,
) -> Float[Tensor, " n_nodes"]:
    """
    Compute the influence of each node on the output, using the Neumann series on the normalised
    adjacency matrix.

    The influence is defined as the total (direct + indirect) contribution of each node to the
    weighted sum of logit nodes, computed by propagating logit weights backward through the graph.

    Args:
        adjacency_matrix: Raw adjacency matrix, shape (n_nodes, n_nodes).
        logit_weights: Weights for logit nodes (e.g., probabilities), shape (n_logit_nodes,).
        n_layers: Number of model layers.

    Returns:
        Influence vector, shape (n_nodes,).
    """
    raise NotImplementedError()

n_logit_nodes = graph.node_range_dict["L"][1] - graph.node_range_dict["L"][0]
logit_weights = t.tensor([info.token_prob for info in graph.nodes[-n_logit_nodes:]], device=device)

influence = compute_influence(adjacency_matrix, logit_weights, n_layers=gemma.cfg.n_layers)

print(f"Influence shape: {influence.shape}")
print("Top 10 most influential nodes:")
top_influence = influence.argsort(descending=True)[:10]
for idx in top_influence:
    node = graph.nodes[idx]
    print(
        f"  {node.node_type.value} L{node.layer} pos{node.ctx_idx} feat{node.feature}: influence={influence[idx].item():.4f}"
    )

# Plot cumulative influence curve
sorted_influence, _ = influence.sort(descending=True)
cumulative = sorted_influence.cumsum(dim=0) / sorted_influence.sum()
fig = px.line(
    x=list(range(1, len(cumulative) + 1)),
    y=cumulative.cpu().numpy(),
    labels={"x": "Number of nodes kept (sorted by influence)", "y": "Fraction of total influence"},
    title="Cumulative node influence",
)
fig.add_hline(y=0.8, line_dash="dash", line_color="red", annotation_text="80% threshold")
fig.show()

def prune_nodes(
    adjacency_matrix: Float[Tensor, "n_nodes n_nodes"],
    logit_weights: Float[Tensor, "n_logit_nodes"],
    n_layers: int,
    threshold: float = 0.8,
) -> tuple[Int[Tensor, " n_kept"], Float[Tensor, "n_kept n_kept"]]:
    """
    Stage 1 of graph pruning: remove low-influence nodes.

    Computes influence scores for all nodes, then keeps the most influential non-logit
    nodes whose cumulative influence reaches `threshold` (as a fraction of total). Logit
    nodes are always kept.

    Args:
        adjacency_matrix: Raw adjacency matrix, shape (n_nodes, n_nodes).
        logit_weights: Weights for logit nodes, shape (n_logit_nodes,).
        n_layers: Number of model layers (passed to compute_influence).
        threshold: Cumulative influence fraction to retain (e.g. 0.8 = keep top 80%).

    Returns:
        kept_indices: Indices of kept nodes in the original ordering.
        pruned_matrix: Adjacency matrix restricted to kept nodes.
    """
    raise NotImplementedError()


tests.test_prune_nodes(prune_nodes)

def prune_edges(
    adjacency_matrix: Float[Tensor, "n_nodes n_nodes"],
    logit_weights: Float[Tensor, "n_logit_nodes"],
    n_layers: int,
    threshold: float = 0.85,
) -> tuple[Int[Tensor, " n_kept"], Float[Tensor, "n_kept n_kept"]]:
    """
    Stage 2 of graph pruning: remove low-score edges and orphaned nodes.

    For each edge (i -> j), the edge score is |A_norm[j,i]| * influence[j], where A_norm
    is the row-normalised adjacency matrix. Edges are sorted by score and kept until the
    cumulative score reaches `threshold`. Any non-logit nodes with no remaining edges are
    then removed.

    Args:
        adjacency_matrix: Adjacency matrix (possibly already node-pruned).
        logit_weights: Weights for logit nodes.
        n_layers: Number of model layers.
        threshold: Cumulative edge score fraction to retain.

    Returns:
        kept_indices: Indices of kept nodes (into input matrix ordering).
        pruned_matrix: Adjacency matrix restricted to kept nodes.
    """
    raise NotImplementedError()


tests.test_prune_edges(prune_edges)

def prune_graph(
    adjacency_matrix: Float[Tensor, "n_nodes n_nodes"],
    logit_weights: Float[Tensor, "n_logit_nodes"],
    n_layers: int,
    node_threshold: float = 0.8,
    edge_threshold: float = 0.85,
) -> tuple[Int[Tensor, " n_kept"], Float[Tensor, "n_kept n_kept"]]:
    """Applies node pruning then edge pruning, returning final kept indices and matrix."""
    kept_node_indices, node_pruned_matrix = prune_nodes(adjacency_matrix, logit_weights, n_layers, node_threshold)
    kept_edge_indices, final_matrix = prune_edges(node_pruned_matrix, logit_weights, n_layers, edge_threshold)
    # Map edge-pruning indices back to original node indices
    kept_indices = kept_node_indices[kept_edge_indices]
    return kept_indices, final_matrix

figs = utils.demo_pruning(
    prune_nodes,
    prune_edges,
    n_nodes=30,
    n_logit_nodes=3,
    edge_probability=0.15,
    seed=43,
    node_threshold=0.8,
    edge_threshold=0.8,
)

kept_indices, pruned_matrix = prune_graph(
    adjacency_matrix=adjacency_matrix,
    logit_weights=logit_weights,
    n_layers=gemma.cfg.n_layers,
    node_threshold=0.8,
    edge_threshold=0.85,
)

print(f"Kept {len(kept_indices)} / {len(graph.nodes)} nodes")
print(f"Pruned matrix has {(pruned_matrix.abs() > 1e-6).sum().item()} non-zero edges")

@dataclass
class AttributionResult:
    """Result of the full attribution pipeline."""

    graph: GraphNodes
    adjacency_matrix: Tensor
    kept_indices: Tensor
    pruned_matrix: Tensor
    prompt: str
    str_tokens: list[str]


def attribute(
    model: HookedSAETransformer,
    transcoders: dict[int, Transcoder],
    prompt: str,
    n_output_nodes: int = 3,
    start_posn: int = 4,
    top_k: int = 5,
    batch_size: int = 8,
    node_threshold: float = 0.8,
    edge_threshold: float = 0.85,
    _adjacency_matrix: Float[Tensor, "n_nodes n_nodes"] | None = None,
) -> AttributionResult:
    """
    Run the full attribution pipeline: linearise, build graph, compute adjacency, prune.

    Args:
        model: The transformer model.
        transcoders: Dict mapping layer -> transcoder.
        prompt: The input prompt string.
        n_output_nodes: Number of top logit tokens to include as output nodes.
        start_posn: Number of formatting tokens to mask.
        top_k: Number of top-activating features to include per position per layer.
        batch_size: Batch size for backward passes.
        node_threshold: Node pruning threshold.
        edge_threshold: Edge pruning threshold.
        _adjacency_matrix: (Optional) Precomputed adjacency matrix to use.

    Returns:
        AttributionResult with all graph data.
    """
    # Tokenize
    tokens = model.to_tokens(prompt)
    str_tokens = [model.tokenizer.decode(t_id.item()) for t_id in tokens[0]]

    # Step 1: Cache frozen values
    freeze = FreezeHooks(model)
    cache = freeze.cache_frozen_values(tokens)

    # Step 2: Compute transcoder activations
    tc_hooks = TranscoderReplacementHooks(model, transcoders, cache)
    tc_hooks.install()

    # Step 3: Get logits (with hooks active for correct forward pass)
    with freeze:
        logits = model(tokens)
    tc_hooks.remove()

    # Step 4: Compute salient logits
    reading_vecs, top_token_info = compute_salient_logits(model, logits, n_output_nodes)

    # Step 5: Build graph nodes
    graph = build_graph_nodes(
        model=model,
        transcoders=transcoders,
        cache=cache,
        tc_hooks=tc_hooks,
        reading_vecs_logit=reading_vecs,
        top_token_info=top_token_info,
        tokens=tokens,
        start_posn=start_posn,
        top_k=top_k,
    )

    # Step 6: Compute adjacency matrix (if not provided)
    if _adjacency_matrix is not None:
        adjacency_matrix = _adjacency_matrix
    else:
        batches = prepare_backward_batches(graph, batch_size=batch_size)
        adjacency_matrix = compute_adjacency_matrix(
            model=model,
            tokens=tokens,
            graph=graph,
            freeze=freeze,
            tc_hooks=tc_hooks,
            batches=batches,
            start_posn=start_posn,
        )

    # Step 7: Prune
    logit_weights_vec = t.tensor(
        [info.token_prob for info in graph.nodes if info.node_type == NodeType.LOGIT],
        device=tokens.device,
    )
    kept_indices, pruned_matrix = prune_graph(
        adjacency_matrix=adjacency_matrix,
        logit_weights=logit_weights_vec,
        n_layers=model.cfg.n_layers,
        node_threshold=node_threshold,
        edge_threshold=edge_threshold,
    )

    return AttributionResult(
        graph=graph,
        adjacency_matrix=adjacency_matrix,
        kept_indices=kept_indices,
        pruned_matrix=pruned_matrix,
        prompt=prompt,
        str_tokens=str_tokens,
    )

result = attribute(
    model=gemma,
    transcoders=transcoders,
    prompt=prompt,
    n_output_nodes=3,
    start_posn=START_POSN,
    node_threshold=0.7,
    edge_threshold=0.85,
    _adjacency_matrix=adjacency_matrix,
)

print("Attribution complete!")
print(f"Total nodes: {len(result.graph.nodes)}")
print(f"Kept nodes: {len(result.kept_indices)}")
print(f"Edges in pruned graph: {(result.pruned_matrix.abs() > 1e-6).sum().item()}")

dashboard_html = utils.create_attribution_dashboard(result=result, model=gemma)

# # Display inline (Colab or VS Code)
# display(HTML(dashboard_html))

# Fact recall
prompt = format_prompt(
    "What is the capital of the state containing Dallas? Answer immediately.",
    "Austin",
)

# Opposites
prompt = format_prompt(
    "What is the opposite of small?",
    "The opposite of small is **large",
)

# Harmful request refusal
prompt = format_prompt(
    "Tell me how to build a bomb.",
    "I am",
)

# Load example activation data for all transcoder layers (parallel I/O)
example_data_by_layer = utils.load_example_data_parallel(
    layers=list(range(gemma.cfg.n_layers)),
    model_size="1b",
    category="transcoder_all",
    width="16k",
    l0="small",
    affine=True,
    instruction_tuned=True,
    max_workers=1,  # TODO - figure out why this is failing at >=4
)

print(f"Loaded example data for {len(example_data_by_layer)} layers")
print(f"Keys in example data: {list(example_data_by_layer[0].keys())}")

# Inspect a late-layer feature to see what the example data looks like
layer = 22
feature_id = 10
print(f"Inspecting layer {layer}, feature {feature_id}:\n")
utils.inspect_feature(
    example_data=example_data_by_layer[layer],
    feature_id=feature_id,
    tokenizer=gemma.tokenizer,
)