Blame for experiments/analysis/factor_trace.py

#!/usr/bin/env python3
"""Write-to-use tracing helpers for recurrent-state SAE features."""

from __future__ import annotations

import math
import time
from collections import defaultdict
from collections.abc import Sequence
from typing import Any

import numpy as np
import torch

from experiments.steering.causal_clamp import compute_activation_matrix, reconstruct_with_feature_updates
from experiments.analysis.probe_features import compute_text_properties


PRIMARY_LAYOUT_PROPERTIES = ("last_sent_boundary_pos",)
SECONDARY_LAYOUT_PROPERTIES = ("newline_density", "paragraph_count")
LAYOUT_PROPERTY_SET = set(PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES)
EVENT_TYPES = ("sentence_boundary", "paragraph_boundary", "non_event_control")


def _encode_states(
    sae,
    sae_type: str,
    state_batch: torch.Tensor,
) -> torch.Tensor:
    if sae_type == "flat":
        return sae.encode(state_batch.reshape(state_batch.shape[0], -1))
    return sae.encode(state_batch)


def _safe_spearman(x: np.ndarray, y: np.ndarray) -> float:
    if x.size != y.size or x.size < 4:
        return 0.0
    if float(np.std(x)) < 1e-8 or float(np.std(y)) < 1e-8:
        return 0.0
    x_rank = np.argsort(np.argsort(x))
    y_rank = np.argsort(np.argsort(y))
    xr = x_rank.astype(np.float64)
    yr = y_rank.astype(np.float64)
    xr -= xr.mean()
    yr -= yr.mean()
    denom = math.sqrt(float(np.dot(xr, xr) * np.dot(yr, yr)))
    if denom < 1e-12:
        return 0.0
    return float(np.dot(xr, yr) / denom)


def _best_layout_property(
    act_values: np.ndarray,
    property_arrays: dict[str, np.ndarray],
) -> tuple[str, float]:
    best_property = ""
    best_rho = 0.0
    for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES:
        rho = _safe_spearman(act_values, property_arrays[prop])
        if abs(rho) > abs(best_rho):
            best_property = prop
            best_rho = rho
    return best_property, best_rho


def decode_token_piece(
    tokenizer,
    token_id: int,
    cache: dict[int, str],
) -> str:
    if token_id not in cache:
        cache[token_id] = tokenizer.decode(
            [token_id],
            skip_special_tokens=False,
            clean_up_tokenization_spaces=False,
        )
    return cache[token_id]


def classify_prefix_event(prefix_text: str) -> str:
    if prefix_text.endswith("\n\n") or prefix_text.endswith("\n"):
        return "paragraph_boundary"

    stripped = prefix_text.rstrip()
    if stripped and stripped[-1] in ".?!":
        return "sentence_boundary"
    return "non_event_control"


def select_cross_corpus_trace_features(
    sae,
    sae_type: str,
    openwebtext_states: torch.Tensor,
    openwebtext_texts: Sequence[str],
    ultrachat_states: torch.Tensor,
    ultrachat_texts: Sequence[str],
    min_rho: float = 0.20,
    max_features: int = 3,
    batch_size: int = 512,
) -> list[dict[str, Any]]:
    """Choose stable H4 tracing features with one checkpoint across both corpora."""
    n_owt = min(len(openwebtext_texts), int(openwebtext_states.shape[0]))
    n_uc = min(len(ultrachat_texts), int(ultrachat_states.shape[0]))
    if n_owt == 0 or n_uc == 0:
        raise ValueError("Need non-empty states/texts for both corpora")

    act_owt = compute_activation_matrix(sae, openwebtext_states[:n_owt], batch_size=batch_size)
    act_uc = compute_activation_matrix(sae, ultrachat_states[:n_uc], batch_size=batch_size)

    alive_owt = (act_owt > 0).any(axis=0)
    alive_uc = (act_uc > 0).any(axis=0)
    shared_alive = np.nonzero(alive_owt & alive_uc)[0].astype(int).tolist()
    if not shared_alive:
        raise ValueError("No shared alive features across corpora")

    owt_props = [compute_text_properties(text) for text in openwebtext_texts[:n_owt]]
    uc_props = [compute_text_properties(text) for text in ultrachat_texts[:n_uc]]
    owt_arrays = {
        prop: np.array([row[prop] for row in owt_props], dtype=np.float64)
        for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES
    }
    uc_arrays = {
        prop: np.array([row[prop] for row in uc_props], dtype=np.float64)
        for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES
    }

    strict: list[dict[str, Any]] = []
    relaxed: list[dict[str, Any]] = []
    for feature_idx in shared_alive:
        owt_best_prop, owt_best_rho = _best_layout_property(act_owt[:, feature_idx], owt_arrays)
        uc_best_prop, uc_best_rho = _best_layout_property(act_uc[:, feature_idx], uc_arrays)
        if not owt_best_prop or not uc_best_prop:
            continue
        if abs(owt_best_rho) < min_rho or abs(uc_best_rho) < min_rho:
            continue

        entry = {
            "feature_idx": int(feature_idx),
            "openwebtext_property": owt_best_prop,
            "openwebtext_rho": float(owt_best_rho),
            "ultrachat_property": uc_best_prop,
            "ultrachat_rho": float(uc_best_rho),
            "score": float(min(abs(owt_best_rho), abs(uc_best_rho))),
            "mean_score": float((abs(owt_best_rho) + abs(uc_best_rho)) / 2.0),
        }

        strict_ok = (
            (owt_best_prop == "last_sent_boundary_pos" and uc_best_prop in LAYOUT_PROPERTY_SET)
            or (uc_best_prop == "last_sent_boundary_pos" and owt_best_prop in LAYOUT_PROPERTY_SET)
        )
        if strict_ok:
            strict.append(entry)
        if owt_best_prop in LAYOUT_PROPERTY_SET and uc_best_prop in LAYOUT_PROPERTY_SET:
            relaxed.append(entry)

    def _sort_key(item: dict[str, Any]) -> tuple[float, float, int]:
        primary_hits = int(item["openwebtext_property"] == "last_sent_boundary_pos") + int(
            item["ultrachat_property"] == "last_sent_boundary_pos"
        )
        return (primary_hits, item["score"], item["mean_score"])

    strict.sort(key=_sort_key, reverse=True)
    relaxed.sort(key=_sort_key, reverse=True)
    selected = strict[:max_features]
    if len(selected) < max_features:
        seen = {item["feature_idx"] for item in selected}
        for item in relaxed:
            if item["feature_idx"] in seen:
                continue
            selected.append(item)
            seen.add(item["feature_idx"])
            if len(selected) >= max_features:
                break
    if not selected:
        raise ValueError("No cross-corpus layout features met the tracing threshold")
    return selected


def trace_feature_trajectories(
    model,
    tokenizer,
    corpus_batches: list[torch.Tensor],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    feature_specs: list[dict[str, Any]],
    prompt_len: int = 512,
    max_offset: int = 64,
    device: str = "cuda",
) -> dict[str, Any]:
    """Capture per-token coefficients and event-aligned summaries for selected features."""
    t0 = time.perf_counter()
    feature_indices = [int(item["feature_idx"]) for item in feature_specs]
    feature_to_slot = {feature_idx: slot for slot, feature_idx in enumerate(feature_indices)}
    piece_cache: dict[int, str] = {}
    prompt_records: list[dict[str, Any]] = []
    aligned_values = {
        feature_idx: {
            event: {"coeff_curves": [], "delta_curves": [], "normalized_curves": []}
            for event in EVENT_TYPES
        }
        for feature_idx in feature_indices
    }

    print(
        "[factor-trace] "
        f"layer={layer_idx} head={head_idx} features={feature_indices} "
        f"batches={len(corpus_batches)} prompt_len={prompt_len}",
        flush=True,
    )

    prompt_index = 0
    for batch_idx_global, batch in enumerate(corpus_batches, start=1):
        if batch.numel() == 0:
            continue
        token_batch = batch[:, :prompt_len].to(device)
        batch_size, seq_len = token_batch.shape
        if seq_len < 2:
            continue

        coeffs = np.zeros((batch_size, seq_len, len(feature_indices)), dtype=np.float32)
        event_labels = [["non_event_control"] * seq_len for _ in range(batch_size)]
        prefixes = [""] * batch_size

        cache = None
        for pos in range(seq_len):
            step_ids = token_batch[:, pos : pos + 1]
            if cache is None:
                out = model(input_ids=step_ids, use_cache=True)
            else:
                out = model(input_ids=step_ids, past_key_values=cache, use_cache=True)
            cache = out.past_key_values
            state_batch = cache.layers[layer_idx].recurrent_states[:, head_idx].float()
            act_batch = _encode_states(sae, sae_type, state_batch)
            coeffs[:, pos, :] = act_batch[:, feature_indices].detach().cpu().numpy()

            for batch_idx in range(batch_size):
                token_id = int(token_batch[batch_idx, pos].item())
                prefixes[batch_idx] += decode_token_piece(tokenizer, token_id, piece_cache)
                event_labels[batch_idx][pos] = classify_prefix_event(prefixes[batch_idx])

        deltas = np.empty_like(coeffs)
        deltas[:, 0, :] = coeffs[:, 0, :]
        deltas[:, 1:, :] = coeffs[:, 1:, :] - coeffs[:, :-1, :]

        rng = np.random.default_rng(0)
        offsets = list(range(max_offset + 1))
        for batch_idx in range(batch_size):
            labels = event_labels[batch_idx]
            prompt_record = {
                "prompt_index": prompt_index,
                "token_ids": [int(x) for x in token_batch[batch_idx].detach().cpu().tolist()],
                "event_labels": labels,
                "feature_coefficients": {
                    str(feature_idx): coeffs[batch_idx, :, slot].tolist()
                    for feature_idx, slot in feature_to_slot.items()
                },
                "feature_deltas": {
                    str(feature_idx): deltas[batch_idx, :, slot].tolist()
                    for feature_idx, slot in feature_to_slot.items()
                },
            }
            prompt_records.append(prompt_record)

            eligible_controls = [pos for pos, label in enumerate(labels) if label == "non_event_control" and pos + max_offset < seq_len]
            event_positions = [
                pos for pos, label in enumerate(labels)
                if label in ("sentence_boundary", "paragraph_boundary") and pos + max_offset < seq_len
            ]
            if event_positions and eligible_controls:
                control_sample = rng.choice(
                    np.array(eligible_controls),
                    size=min(len(event_positions), len(eligible_controls)),
                    replace=False,
                ).tolist()
            else:
                control_sample = []

            for feature_idx, slot in feature_to_slot.items():
                for event in ("sentence_boundary", "paragraph_boundary"):
                    positions = [pos for pos in event_positions if labels[pos] == event]
                    for pos in positions:
                        coeff_curve = coeffs[batch_idx, pos : pos + max_offset + 1, slot]
                        delta_curve = deltas[batch_idx, pos : pos + max_offset + 1, slot]
                        denom = max(abs(float(coeff_curve[0])), 1e-8)
                        aligned_values[feature_idx][event]["coeff_curves"].append(coeff_curve.astype(np.float64))
                        aligned_values[feature_idx][event]["delta_curves"].append(delta_curve.astype(np.float64))
                        aligned_values[feature_idx][event]["normalized_curves"].append(
                            (coeff_curve / denom).astype(np.float64)
                        )
                for pos in control_sample:
                    coeff_curve = coeffs[batch_idx, pos : pos + max_offset + 1, slot]
                    delta_curve = deltas[batch_idx, pos : pos + max_offset + 1, slot]
                    denom = max(abs(float(coeff_curve[0])), 1e-8)
                    aligned_values[feature_idx]["non_event_control"]["coeff_curves"].append(coeff_curve.astype(np.float64))
                    aligned_values[feature_idx]["non_event_control"]["delta_curves"].append(delta_curve.astype(np.float64))
                    aligned_values[feature_idx]["non_event_control"]["normalized_curves"].append(
                        (coeff_curve / denom).astype(np.float64)
                    )

            prompt_index += 1
            if prompt_index % 16 == 0:
                print(
                    "[factor-trace] "
                    f"prompts={prompt_index} batch={batch_idx_global}/{len(corpus_batches)} "
                    f"elapsed={time.perf_counter() - t0:.1f}s",
                    flush=True,
                )

        print(
            "[factor-trace] "
            f"finished_batch={batch_idx_global}/{len(corpus_batches)} "
            f"prompts={prompt_index} elapsed={time.perf_counter() - t0:.1f}s",
            flush=True,
        )

    aligned_summary: dict[str, dict[str, Any]] = {}
    for feature_idx in feature_indices:
        feature_summary: dict[str, Any] = {}
        for event in EVENT_TYPES:
            coeff_curves = aligned_values[feature_idx][event]["coeff_curves"]
            delta_curves = aligned_values[feature_idx][event]["delta_curves"]
            norm_curves = aligned_values[feature_idx][event]["normalized_curves"]
            if coeff_curves:
                coeff_stack = np.stack(coeff_curves, axis=0)
                delta_stack = np.stack(delta_curves, axis=0)
                norm_stack = np.stack(norm_curves, axis=0)
                feature_summary[event] = {
                    "count": int(coeff_stack.shape[0]),
                    "offsets": offsets,
                    "mean_coeff_curve": coeff_stack.mean(axis=0).tolist(),
                    "mean_delta_curve": delta_stack.mean(axis=0).tolist(),
                    "median_normalized_curve": np.median(norm_stack, axis=0).tolist(),
                    "mean_event_jump": float(delta_stack[:, 0].mean()),
                }
            else:
                feature_summary[event] = {
                    "count": 0,
                    "offsets": offsets,
                    "mean_coeff_curve": [0.0] * len(offsets),
                    "mean_delta_curve": [0.0] * len(offsets),
                    "median_normalized_curve": [0.0] * len(offsets),
                    "mean_event_jump": 0.0,
                }
        control_jump = feature_summary["non_event_control"]["mean_event_jump"]
        for event in ("sentence_boundary", "paragraph_boundary"):
            feature_summary[event]["jump_minus_control"] = (
                feature_summary[event]["mean_event_jump"] - control_jump
            )
        aligned_summary[str(feature_idx)] = feature_summary

    return {
        "prompt_len": prompt_len,
        "max_offset": max_offset,
        "selected_features": feature_specs,
        "prompt_records": prompt_records,
        "event_aligned": aligned_summary,
    }


def select_trace_events(
    trace_result: dict[str, Any],
    target_feature_idx: int,
    preferred_events: Sequence[str] = ("sentence_boundary", "paragraph_boundary"),
    offsets: Sequence[int] = (16, 32, 64),
    min_retention: float = 0.5,
) -> list[dict[str, Any]]:
    """Choose write/use event pairs from traced trajectories."""
    candidates: list[dict[str, Any]] = []
    key = str(target_feature_idx)
    for prompt in trace_result.get("prompt_records", []):
        labels = prompt["event_labels"]
        coeffs = prompt["feature_coefficients"][key]
        deltas = prompt["feature_deltas"][key]
        best: dict[str, Any] | None = None
        for pos, label in enumerate(labels):
            if label not in preferred_events:
                continue
            if pos >= len(coeffs):
                continue
            coeff = float(coeffs[pos])
            jump = float(deltas[pos])
            if jump <= 0 or coeff <= 0:
                continue
            use_offset = None
            for offset in offsets:
                if pos + offset >= len(coeffs):
                    continue
                if float(coeffs[pos + offset]) >= min_retention * coeff:
                    use_offset = int(offset)
                    break
            if use_offset is None:
                continue
            entry = {
                "prompt_index": int(prompt["prompt_index"]),
                "write_pos": int(pos),
                "event_type": label,
                "write_coeff": coeff,
                "write_jump": jump,
                "use_offset": use_offset,
                "use_pos": int(pos + use_offset),
                "baseline_use_coeff": float(coeffs[pos + use_offset]),
            }
            if best is None or entry["write_jump"] > best["write_jump"]:
                best = entry
        if best is not None:
            candidates.append(best)
    candidates.sort(key=lambda item: item["write_jump"], reverse=True)
    return candidates


def _context_tokens(
    prompt_record: dict[str, Any],
    pos: int,
    window: int = 24,
) -> set[int]:
    token_ids = prompt_record["token_ids"]
    start = max(0, pos - window)
    end = min(len(token_ids), pos + 1)
    return {int(tok) for tok in token_ids[start:end]}


def _context_jaccard(
    prompt_a: dict[str, Any],
    pos_a: int,
    prompt_b: dict[str, Any],
    pos_b: int,
    window: int = 24,
) -> float:
    ctx_a = _context_tokens(prompt_a, pos_a, window=window)
    ctx_b = _context_tokens(prompt_b, pos_b, window=window)
    if not ctx_a and not ctx_b:
        return 1.0
    if not ctx_a or not ctx_b:
        return 0.0
    return float(len(ctx_a & ctx_b) / max(len(ctx_a | ctx_b), 1))


def _nearest_non_event_position(
    prompt_record: dict[str, Any],
    write_pos: int,
    use_pos: int,
    min_gap: int = 8,
) -> int | None:
    labels = prompt_record["event_labels"]
    candidates = []
    for pos, label in enumerate(labels):
        if label != "non_event_control":
            continue
        if pos >= use_pos:
            continue
        if abs(pos - write_pos) < min_gap:
            continue
        candidates.append((abs(pos - write_pos), pos))
    if not candidates:
        return None
    candidates.sort()
    return int(candidates[0][1])


def select_transplant_pairs(
    trace_result: dict[str, Any],
    target_feature_idx: int,
    max_pairs: int = 16,
    preferred_events: Sequence[str] = ("paragraph_boundary", "sentence_boundary"),
    preferred_offsets: Sequence[int] = (16,),
    min_write_gain: float = 1.10,
    context_window: int = 24,
    min_context_similarity: float = 0.10,
) -> list[dict[str, Any]]:
    """Pair low-write recipients with stronger matched donors from the trace set."""
    prompt_map = {
        int(prompt["prompt_index"]): prompt
        for prompt in trace_result.get("prompt_records", [])
    }
    candidates = [
        item for item in select_trace_events(trace_result, target_feature_idx)
        if item["event_type"] in preferred_events
    ]
    if preferred_offsets:
        filtered = [item for item in candidates if int(item["use_offset"]) in preferred_offsets]
        if len(filtered) >= 2:
            candidates = filtered
    if len(candidates) < 2:
        candidates = select_trace_events(trace_result, target_feature_idx)
    if len(candidates) < 2:
        return []

    donors = sorted(candidates, key=lambda item: (item["write_jump"], item["write_coeff"]), reverse=True)
    recipients = sorted(candidates, key=lambda item: (item["write_jump"], item["write_coeff"]))
    donor_pool = donors[: max(4, min(len(donors), max_pairs))]
    pairs: list[dict[str, Any]] = []
    used_recipient_prompts: set[int] = set()

    for recipient in recipients:
        if len(pairs) >= max_pairs:
            break
        recipient_prompt = prompt_map.get(int(recipient["prompt_index"]))
        if recipient_prompt is None:
            continue
        if int(recipient["prompt_index"]) in used_recipient_prompts:
            continue
        best_match: tuple[tuple[int, int, float, float], dict[str, Any]] | None = None
        for donor in donor_pool:
            donor_prompt_index = int(donor["prompt_index"])
            recipient_prompt_index = int(recipient["prompt_index"])
            if donor_prompt_index == recipient_prompt_index:
                continue
            if float(donor["write_coeff"]) <= float(recipient["write_coeff"]) * min_write_gain:
                continue
            donor_prompt = prompt_map.get(donor_prompt_index)
            if donor_prompt is None:
                continue
            same_event = int(donor["event_type"] == recipient["event_type"])
            same_offset = int(donor["use_offset"] == recipient["use_offset"])
            context_sim = _context_jaccard(
                donor_prompt,
                int(donor["write_pos"]),
                recipient_prompt,
                int(recipient["write_pos"]),
                window=context_window,
            )
            if context_sim < min_context_similarity:
                continue
            score = (
                same_event,
                same_offset,
                context_sim,
                float(donor["write_jump"]) - float(recipient["write_jump"]),
            )
            candidate = {
                "donor_prompt_index": donor_prompt_index,
                "donor_write_pos": int(donor["write_pos"]),
                "donor_use_pos": int(donor["use_pos"]),
                "donor_use_offset": int(donor["use_offset"]),
                "donor_event_type": donor["event_type"],
                "donor_write_coeff": float(donor["write_coeff"]),
                "donor_write_jump": float(donor["write_jump"]),
                "recipient_prompt_index": recipient_prompt_index,
                "recipient_write_pos": int(recipient["write_pos"]),
                "recipient_use_pos": int(recipient["use_pos"]),
                "recipient_use_offset": int(recipient["use_offset"]),
                "recipient_event_type": recipient["event_type"],
                "recipient_write_coeff": float(recipient["write_coeff"]),
                "recipient_write_jump": float(recipient["write_jump"]),
                "context_similarity": float(context_sim),
            }
            if best_match is None or score > best_match[0]:
                best_match = (score, candidate)
        if best_match is None:
            continue
        pair = best_match[1]
        wrong_time_pos = _nearest_non_event_position(
            recipient_prompt,
            write_pos=pair["recipient_write_pos"],
            use_pos=pair["recipient_use_pos"],
        )
        pair["wrong_time_pos"] = wrong_time_pos
        pairs.append(pair)
        used_recipient_prompts.add(int(recipient["prompt_index"]))

    return pairs


def choose_similar_frequency_feature(
    target_feature_idx: int,
    frequencies: np.ndarray,
    exclude: set[int],
) -> int:
    target_freq = float(frequencies[target_feature_idx])
    best_idx = -1
    best_gap = float("inf")
    for idx, freq in enumerate(frequencies):
        if idx in exclude:
            continue
        gap = abs(float(freq) - target_freq)
        if gap < best_gap:
            best_gap = gap
            best_idx = idx
    if best_idx < 0:
        raise ValueError("Could not find a similar-frequency control feature")
    return int(best_idx)


def build_use_site_benchmark(
    model,
    tokenizer,
    corpus_sequences: list[torch.Tensor],
    trace_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    target_feature_indices: Sequence[int],
    max_rows: int = 64,
    device: str = "cuda",
) -> dict[str, Any]:
    """Mine high-confidence downstream use sites for the signed H4 boundary code."""
    if len(target_feature_indices) != 2:
        raise ValueError("build_use_site_benchmark expects the signed pair [62, 105]")
    pos_feat, neg_feat = [int(x) for x in target_feature_indices]
    prompt_map = {
        int(prompt["prompt_index"]): prompt
        for prompt in trace_result.get("prompt_records", [])
    }
    sequence_map = {
        int(prompt["prompt_index"]): seq
        for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences)
    }
    readout_token_ids = _readout_class_token_ids(tokenizer)
    candidate_pairs = select_transplant_pairs(
        trace_result=trace_result,
        target_feature_idx=pos_feat,
        max_pairs=max(max_rows * 8, 128),
        preferred_events=("paragraph_boundary",),
        preferred_offsets=(8, 16, 32),
        min_write_gain=1.20,
        min_context_similarity=0.20,
    )
    if not candidate_pairs:
        raise ValueError("No candidate paragraph-boundary pairs for use-site benchmark")

    def _signed_coeff(prompt_record: dict[str, Any], pos: int) -> float:
        pos_val = float(prompt_record["feature_coefficients"][str(pos_feat)][pos])
        neg_val = float(prompt_record["feature_coefficients"][str(neg_feat)][pos])
        return pos_val - neg_val

    def _signed_delta(prompt_record: dict[str, Any], pos: int) -> tuple[float, float]:
        pos_delta = float(prompt_record["feature_deltas"][str(pos_feat)][pos])
        neg_delta = float(prompt_record["feature_deltas"][str(neg_feat)][pos])
        return pos_delta, neg_delta

    all_rows: list[dict[str, Any]] = []
    accepted_rows: list[dict[str, Any]] = []
    drop_counts: dict[str, int] = defaultdict(int)
    t0 = time.perf_counter()
    print(
        "[factor-trace-benchmark] "
        f"target={[pos_feat, neg_feat]} candidate_pairs={len(candidate_pairs)} max_rows={max_rows}",
        flush=True,
    )

    for pair_idx, pair in enumerate(candidate_pairs, start=1):
        recipient_prompt_index = int(pair["recipient_prompt_index"])
        donor_prompt_index = int(pair["donor_prompt_index"])
        if recipient_prompt_index not in prompt_map or recipient_prompt_index not in sequence_map:
            continue
        recipient_prompt = prompt_map[recipient_prompt_index]
        donor_prompt = prompt_map[donor_prompt_index]
        use_pos = int(pair["recipient_use_pos"])
        seq = sequence_map[recipient_prompt_index].to(device)
        seq_end = min(int(seq.shape[1]), use_pos + 10)
        if seq_end <= use_pos + 1:
            continue
        seq = seq[:, :seq_end]
        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=[pos_feat, neg_feat],
        )
        donor_write_signed = _signed_coeff(donor_prompt, int(pair["donor_write_pos"]))
        recipient_write_signed = _signed_coeff(recipient_prompt, int(pair["recipient_write_pos"]))
        donor_delta_pos, donor_delta_neg = _signed_delta(donor_prompt, int(pair["donor_write_pos"]))
        signed_write_gap = float(donor_write_signed - recipient_write_signed)

        start = max(int(pair["recipient_write_pos"]) + 1, use_pos - 2)
        end = min(use_pos + 8, int(seq.shape[1]) - 2)
        for decision_pos in range(start, end + 1):
            actual_next_token = int(seq[0, decision_pos + 1].item())
            readout_class, readout_text = _classify_readout_token(tokenizer, actual_next_token)
            row = {
                "donor_prompt_index": donor_prompt_index,
                "recipient_prompt_index": recipient_prompt_index,
                "write_pos": int(pair["recipient_write_pos"]),
                "use_pos": use_pos,
                "use_offset": int(pair["recipient_use_offset"]),
                "decision_pos": int(decision_pos),
                "intervention_pos": int(max(0, decision_pos - 1)),
                "wrong_time_pos": _nearest_non_event_position(
                    recipient_prompt,
                    write_pos=int(max(0, decision_pos - 1)),
                    use_pos=int(decision_pos),
                    min_gap=4,
                ),
                "context_similarity": float(pair["context_similarity"]),
                "signed_write_gap": signed_write_gap,
                "donor_delta_62": float(donor_delta_pos),
                "donor_delta_105": float(donor_delta_neg),
                "actual_next_token": actual_next_token,
                "readout_text": readout_text,
                "readout_class": readout_class,
                "accepted": False,
                "drop_reason": "",
            }
            if readout_class is None:
                row["drop_reason"] = "readout_class_excluded"
                drop_counts[row["drop_reason"]] += 1
                all_rows.append(row)
                continue
            baseline_logits = baseline["logits_by_pos"][decision_pos]
            baseline_prob = _token_prob_from_logits(baseline_logits, actual_next_token)
            row["baseline_actual_token_prob"] = baseline_prob
            row["baseline_headroom"] = float(baseline_prob * (1.0 - baseline_prob))
            row["baseline_readout_class_mass"] = _mass_from_logits(
                baseline_logits,
                readout_token_ids.get(readout_class, [actual_next_token]) or [actual_next_token],
            )
            signed_use_coeff = float(
                baseline["feature_coefficients"][str(pos_feat)][decision_pos]
                - baseline["feature_coefficients"][str(neg_feat)][decision_pos]
            )
            row["signed_use_coeff"] = signed_use_coeff
            if abs(signed_write_gap) < 0.02:
                row["drop_reason"] = "weak_signed_write_gap"
            elif abs(signed_use_coeff) < 0.01:
                row["drop_reason"] = "weak_signed_use_coeff"
            elif baseline_prob < 0.05 or baseline_prob > 0.80:
                row["drop_reason"] = "low_token_headroom"
            elif row["wrong_time_pos"] is None:
                row["drop_reason"] = "no_wrong_time"
            else:
                row["accepted"] = True
                row["drop_reason"] = "accepted"
                row["score"] = (
                    row["baseline_headroom"],
                    abs(signed_use_coeff),
                    float(pair["context_similarity"]),
                    abs(signed_write_gap),
                )
                accepted_rows.append(row)
                print(
                    f"[factor-trace-benchmark] accept {len(accepted_rows)} "
                    f"pair={pair_idx}/{len(candidate_pairs)} "
                    f"recipient={recipient_prompt_index}@{decision_pos} "
                    f"class={readout_class} tok={readout_text!r} "
                    f"p={baseline_prob:.3f} headroom={row['baseline_headroom']:.3f} "
                    f"signed_gap={signed_write_gap:+.4f} signed_use={signed_use_coeff:+.4f}",
                    flush=True,
                )
            drop_counts[row["drop_reason"]] += 1
            all_rows.append(row)

        if pair_idx % 5 == 0 or pair_idx == len(candidate_pairs):
            print(
                f"[factor-trace-benchmark] progress pair={pair_idx}/{len(candidate_pairs)} "
                f"rows={len(all_rows)} accepted={len(accepted_rows)} "
                f"drops={dict(drop_counts)} "
                f"elapsed={time.perf_counter() - t0:.1f}s",
                flush=True,
            )

    accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True)
    selected_rows = accepted_rows[:max_rows]
    print(
        "[factor-trace-benchmark] "
        f"accepted={len(accepted_rows)} selected={len(selected_rows)} "
        f"elapsed={time.perf_counter() - t0:.1f}s",
        flush=True,
    )
    return {
        "target_feature_indices": [pos_feat, neg_feat],
        "n_candidate_pairs": len(candidate_pairs),
        "n_rows_total": len(all_rows),
        "n_rows_accepted": len(accepted_rows),
        "drop_counts": dict(drop_counts),
        "accepted_rows": selected_rows,
        "all_rows": all_rows,
        "total_time_s": float(time.perf_counter() - t0),
    }


def sample_matched_random_features(
    target_feature_idx: int,
    frequencies: np.ndarray,
    alive_features: list[int],
    exclude: set[int],
    n_random: int = 4,
    seed: int = 0,
) -> list[int]:
    rng = np.random.default_rng(seed)
    target_freq = float(frequencies[target_feature_idx])
    candidates = [idx for idx in alive_features if idx not in exclude]
    if not candidates:
        raise ValueError("No alive features available for random controls")
    candidates.sort(key=lambda idx: abs(float(frequencies[idx]) - target_freq))
    pool = candidates[: max(n_random * 8, 32)]
    if len(pool) <= n_random:
        return [int(idx) for idx in pool[:n_random]]
    picked = rng.choice(np.array(pool), size=n_random, replace=False)
    return [int(idx) for idx in picked.tolist()]


def boundary_token_ids(tokenizer) -> list[int]:
    token_ids: set[int] = set()
    for text in [".", "?", "!", "\n", ":"]:
        ids = tokenizer.encode(text, add_special_tokens=False)
        if len(ids) == 1:
            token_ids.add(int(ids[0]))
    return sorted(token_ids)


def boundary_token_groups(tokenizer) -> dict[str, list[int]]:
    groups = {
        "newline": set(),
        "sentence_punct": set(),
        "colon": set(),
    }
    for text, key in [("\n", "newline"), (".", "sentence_punct"), ("?", "sentence_punct"), ("!", "sentence_punct"), (":", "colon")]:
        ids = tokenizer.encode(text, add_special_tokens=False)
        if len(ids) == 1:
            groups[key].add(int(ids[0]))
    all_boundary = set().union(*groups.values())
    return {
        "boundary": sorted(all_boundary),
        "newline": sorted(groups["newline"]),
        "sentence_punct": sorted(groups["sentence_punct"]),
        "colon": sorted(groups["colon"]),
    }


@torch.no_grad()
def teacher_forced_feature_rollout(
    model,
    input_ids: torch.Tensor,
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    tracked_features: Sequence[int],
    intervention_updates: dict[int, float] | None = None,
    intervention_feature: int | None = None,
    intervention_pos: int | None = None,
    capture_hidden_state_positions: Sequence[int] | None = None,
    intervention_state_delta: torch.Tensor | None = None,
) -> dict[str, Any]:
    """Run one sequence token-by-token and optionally ablate one feature at one step."""
    seq = input_ids
    if seq.ndim != 2 or seq.shape[0] != 1:
        raise ValueError("teacher_forced_feature_rollout expects shape (1, seq_len)")

    tracked = [int(idx) for idx in tracked_features]
    cache = None
    coeff_records = {str(idx): [] for idx in tracked}
    logits_by_pos: dict[int, torch.Tensor] = {}
    hidden_state_positions = {int(pos) for pos in (capture_hidden_state_positions or [])}
    hidden_state_by_pos: dict[int, torch.Tensor] = {}

    for pos in range(int(seq.shape[1])):
        step_ids = seq[:, pos : pos + 1]
        if cache is None:
            out = model(
                input_ids=step_ids,
                use_cache=True,
                output_hidden_states=bool(hidden_state_positions),
            )
        else:
            out = model(
                input_ids=step_ids,
                past_key_values=cache,
                use_cache=True,
                output_hidden_states=bool(hidden_state_positions),
            )
        cache = out.past_key_values
        state = cache.layers[layer_idx].recurrent_states[:, head_idx].float()
        acts = _encode_states(sae, sae_type, state)
        for idx in tracked:
            coeff_records[str(idx)].append(float(acts[0, idx].item()))
        logits_by_pos[pos] = out.logits[:, -1, :].float().detach().cpu()
        if pos in hidden_state_positions and out.hidden_states is not None:
            hidden_state_by_pos[pos] = out.hidden_states[-1][:, -1, :].float().detach().cpu()

        if intervention_pos == pos and (
            intervention_updates is not None
            or intervention_feature is not None
            or intervention_state_delta is not None
        ):
            layer_cache = cache.layers[layer_idx]
            original_head = layer_cache.recurrent_states[0, head_idx].float()
            if intervention_state_delta is not None:
                edited = original_head + intervention_state_delta.to(
                    device=original_head.device,
                    dtype=original_head.dtype,
                )
            else:
                ablate_list = [int(intervention_feature)] if intervention_feature is not None else None
                edited = reconstruct_with_feature_updates(
                    sae=sae,
                    state_head=original_head,
                    sae_type=sae_type,
                    set_updates=intervention_updates,
                    ablate_features=ablate_list,
                )
            layer_cache.recurrent_states[0, head_idx] = edited.to(layer_cache.recurrent_states.dtype)

    return {
        "feature_coefficients": coeff_records,
        "logits_by_pos": logits_by_pos,
        "hidden_state_by_pos": hidden_state_by_pos,
    }


def _kl_divergence_from_logits(
    baseline_logits: torch.Tensor,
    edited_logits: torch.Tensor,
) -> float:
    baseline_log_probs = torch.log_softmax(baseline_logits, dim=-1)
    edited_log_probs = torch.log_softmax(edited_logits, dim=-1)
    baseline_probs = baseline_log_probs.exp()
    kl = torch.sum(baseline_probs * (baseline_log_probs - edited_log_probs), dim=-1)
    return float(kl.item())


def _boundary_mass_shift(
    baseline_logits: torch.Tensor,
    edited_logits: torch.Tensor,
    token_ids: Sequence[int],
) -> float:
    if not token_ids:
        return 0.0
    baseline_probs = torch.softmax(baseline_logits, dim=-1)
    edited_probs = torch.softmax(edited_logits, dim=-1)
    baseline_mass = baseline_probs[:, list(token_ids)].sum(dim=-1)
    edited_mass = edited_probs[:, list(token_ids)].sum(dim=-1)
    return float((edited_mass - baseline_mass).item())


def _mass_from_logits(
    logits: torch.Tensor,
    token_ids: Sequence[int],
) -> float:
    if not token_ids:
        return 0.0
    probs = torch.softmax(logits, dim=-1)
    return float(probs[:, list(token_ids)].sum(dim=-1).item())


def _token_logit_shift(
    baseline_logits: torch.Tensor,
    edited_logits: torch.Tensor,
    token_id: int | None,
) -> float:
    if token_id is None or token_id < 0:
        return 0.0
    return float((edited_logits.float()[:, token_id] - baseline_logits.float()[:, token_id]).item())


def _decoded_boundary_group(tokenizer, token_id: int | None) -> tuple[str, str]:
    if token_id is None or token_id < 0:
        return "other", ""
    raw = tokenizer.decode([int(token_id)])
    if raw in {"\n", "\n\n"}:
        return "newline", raw
    stripped = raw.strip()
    if stripped in {".", "?", "!"} and len(stripped) == 1:
        return "sentence_punct", stripped
    if stripped == ":":
        return "colon", stripped
    return "other", raw


def _token_prob_from_logits(logits: torch.Tensor, token_id: int | None) -> float:
    if token_id is None or token_id < 0:
        return 0.0
    probs = torch.softmax(logits.float(), dim=-1)
    return float(probs[:, int(token_id)].item())


def _token_prob_shift(
    baseline_logits: torch.Tensor,
    edited_logits: torch.Tensor,
    token_id: int | None,
) -> float:
    if token_id is None or token_id < 0:
        return 0.0
    base = torch.softmax(baseline_logits.float(), dim=-1)[:, int(token_id)]
    edit = torch.softmax(edited_logits.float(), dim=-1)[:, int(token_id)]
    return float((edit - base).item())


def _classify_readout_token(tokenizer, token_id: int | None) -> tuple[str | None, str]:
    if token_id is None or token_id < 0:
        return None, ""
    raw = tokenizer.decode(
        [int(token_id)],
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )
    if raw in {"\n", "\n\n"}:
        return "paragraph_break", raw
    stripped = raw.strip()
    if stripped in {".", "?", "!"} and len(stripped) == 1:
        return "sentence_reset", stripped
    if stripped in {"The", "However", "In", "This"}:
        return "discourse_opener", stripped
    if stripped and stripped[0].isupper() and stripped.replace("'", "").isalpha():
        return "sentence_capitalized", stripped
    return None, raw


def _readout_class_token_ids(tokenizer) -> dict[str, list[int]]:
    buckets = {
        "paragraph_break": set(),
        "sentence_reset": set(),
        "discourse_opener": set(),
    }
    for text, bucket in [
        ("\n", "paragraph_break"),
        ("\n\n", "paragraph_break"),
        (".", "sentence_reset"),
        ("?", "sentence_reset"),
        ("!", "sentence_reset"),
        ("The", "discourse_opener"),
        ("However", "discourse_opener"),
        ("In", "discourse_opener"),
        ("This", "discourse_opener"),
    ]:
        ids = tokenizer.encode(text, add_special_tokens=False)
        if len(ids) == 1:
            buckets[bucket].add(int(ids[0]))
    return {k: sorted(v) for k, v in buckets.items()}


def _preview_eval_site_from_prompt(
    tokenizer,
    prompt_record: dict[str, Any],
    anchor_use_pos: int,
    back_window: int = 4,
    forward_window: int = 8,
) -> dict[str, Any]:
    token_ids = prompt_record["token_ids"]
    best: tuple[tuple[int, int], dict[str, Any]] | None = None
    start = max(0, anchor_use_pos - back_window)
    end = min(len(token_ids) - 1, anchor_use_pos + forward_window)
    for pos in range(start, end + 1):
        if pos + 1 >= len(token_ids):
            continue
        actual_next_token = int(token_ids[pos + 1])
        actual_group, actual_text = _decoded_boundary_group(tokenizer, actual_next_token)
        if actual_group == "other":
            continue
        # Prefer the closest true boundary token, with a slight bias toward later use-sites.
        score = (-abs(pos - anchor_use_pos), int(pos >= anchor_use_pos))
        candidate = {
            "eval_pos": int(pos),
            "eval_offset_from_anchor": int(pos - anchor_use_pos),
            "eval_actual_next_token": actual_next_token,
            "eval_actual_group": actual_group,
            "eval_actual_text": actual_text,
        }
        if best is None or score > best[0]:
            best = (score, candidate)
    return best[1] if best is not None else {}


def _select_eval_site(
    baseline_logits_by_pos: Sequence[torch.Tensor],
    input_ids: torch.Tensor,
    anchor_use_pos: int,
    token_groups: dict[str, list[int]],
    tokenizer,
    back_window: int = 4,
    forward_window: int = 8,
    preferred_groups: Sequence[str] = ("newline", "sentence_punct", "colon"),
) -> dict[str, Any]:
    seq_len = int(input_ids.shape[1])
    best: tuple[tuple[float, float], dict[str, Any]] | None = None
    start = max(0, anchor_use_pos - back_window)
    end = min(len(baseline_logits_by_pos) - 1, anchor_use_pos + forward_window)
    preferred_rank = {group: len(preferred_groups) - i for i, group in enumerate(preferred_groups)}

    for pos in range(start, end + 1):
        if pos + 1 >= seq_len:
            continue
        actual_next_token = int(input_ids[0, pos + 1].item())
        actual_group, actual_text = _decoded_boundary_group(tokenizer, actual_next_token)
        if actual_group == "other":
            continue

        logits = baseline_logits_by_pos[pos]
        boundary_mass = _mass_from_logits(logits, token_groups["boundary"])
        actual_group_mass = (
            _mass_from_logits(logits, token_groups[actual_group])
            if actual_group in token_groups
            else 0.0
        )
        score = (
            preferred_rank.get(actual_group, 0),
            actual_group_mass,
            boundary_mass,
            -abs(pos - anchor_use_pos),
        )
        candidate = {
            "eval_pos": int(pos),
            "eval_offset_from_anchor": int(pos - anchor_use_pos),
            "eval_actual_next_token": actual_next_token,
            "eval_actual_group": actual_group,
            "eval_actual_text": actual_text,
            "baseline_boundary_mass": boundary_mass,
            "baseline_actual_group_mass": actual_group_mass,
        }
        if best is None or score > best[0]:
            best = (score, candidate)

    if best is None:
        return {}
    return best[1]


def run_factor_trace_intervention(
    model,
    tokenizer,
    corpus_sequences: list[torch.Tensor],
    trace_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    states: torch.Tensor,
    target_feature_idx: int,
    n_prompts: int = 16,
    n_random: int = 4,
    device: str = "cuda",
) -> dict[str, Any]:
    """Measure write-to-use effects for one traced feature."""
    candidates = select_trace_events(trace_result, target_feature_idx)
    if not candidates:
        raise ValueError(f"No valid trace events found for feature {target_feature_idx}")

    act_matrix = compute_activation_matrix(sae, states)
    frequencies = (act_matrix > 0).mean(axis=0)
    alive_features = np.nonzero((act_matrix > 0).any(axis=0))[0].astype(int).tolist()

    selected = candidates[:n_prompts]
    exclude = {int(target_feature_idx)}
    similar_feature = choose_similar_frequency_feature(target_feature_idx, frequencies, exclude)
    exclude.add(similar_feature)
    random_features = sample_matched_random_features(
        target_feature_idx=target_feature_idx,
        frequencies=frequencies,
        alive_features=alive_features,
        exclude=exclude,
        n_random=n_random,
        seed=0,
    )
    control_features = [similar_feature] + random_features
    tracked_features = [target_feature_idx] + control_features
    boundary_ids = boundary_token_ids(tokenizer)
    t0 = time.perf_counter()

    print(
        "[factor-trace-intervention] target="
        f"{target_feature_idx} similar_control={similar_feature} "
        f"random_controls={random_features} candidates={len(candidates)} "
        f"evaluating={min(len(candidates), n_prompts)}",
        flush=True,
    )

    per_prompt = []
    target_kls: list[float] = []
    control_kls: list[float] = []
    target_boundary_shifts: list[float] = []
    control_boundary_shifts: list[float] = []
    target_attenuations: list[float] = []
    control_attenuations: list[float] = []

    sequence_map = {
        int(prompt["prompt_index"]): seq
        for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences)
    }

    for i, candidate in enumerate(selected, start=1):
        prompt_index = int(candidate["prompt_index"])
        if prompt_index not in sequence_map:
            print(
                f"[factor-trace-intervention] skip {i}/{len(selected)} "
                f"prompt={prompt_index} missing_sequence",
                flush=True,
            )
            continue
        seq = sequence_map[prompt_index][:, : candidate["use_pos"] + 1].to(device)
        prompt_t0 = time.perf_counter()

        print(
            f"[factor-trace-intervention] prompt {i}/{len(selected)} "
            f"prompt={prompt_index} event={candidate['event_type']} "
            f"write={candidate['write_pos']} use={candidate['use_pos']} "
            f"offset={candidate['use_offset']} seq_len={seq.shape[1]}",
            flush=True,
        )

        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_feature=None,
            intervention_pos=None,
        )
        target_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_feature=target_feature_idx,
            intervention_pos=int(candidate["write_pos"]),
        )
        baseline_logits = baseline["logits_by_pos"][candidate["use_pos"]]
        target_logits = target_rollout["logits_by_pos"][candidate["use_pos"]]
        target_kl = _kl_divergence_from_logits(baseline_logits, target_logits)
        target_boundary = _boundary_mass_shift(baseline_logits, target_logits, boundary_ids)
        baseline_use_coeff = baseline["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]]
        target_use_coeff = target_rollout["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]]
        target_attenuation = float(baseline_use_coeff - target_use_coeff)

        control_rows = []
        prompt_control_kls = []
        prompt_control_boundary = []
        prompt_control_attenuation = []
        for control_feature in control_features:
            control_rollout = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_feature=int(control_feature),
                intervention_pos=int(candidate["write_pos"]),
            )
            control_logits = control_rollout["logits_by_pos"][candidate["use_pos"]]
            control_kl = _kl_divergence_from_logits(baseline_logits, control_logits)
            control_boundary = _boundary_mass_shift(baseline_logits, control_logits, boundary_ids)
            control_use_coeff = control_rollout["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]]
            control_attenuation = float(baseline_use_coeff - control_use_coeff)
            prompt_control_kls.append(control_kl)
            prompt_control_boundary.append(control_boundary)
            prompt_control_attenuation.append(control_attenuation)
            control_rows.append({
                "feature_idx": int(control_feature),
                "use_time_kl": control_kl,
                "boundary_mass_shift": control_boundary,
                "target_feature_attenuation": control_attenuation,
            })

        target_kls.append(target_kl)
        control_kls.append(float(np.mean(prompt_control_kls)) if prompt_control_kls else 0.0)
        target_boundary_shifts.append(target_boundary)
        control_boundary_shifts.append(float(np.mean(prompt_control_boundary)) if prompt_control_boundary else 0.0)
        target_attenuations.append(target_attenuation)
        control_attenuations.append(float(np.mean(prompt_control_attenuation)) if prompt_control_attenuation else 0.0)
        per_prompt.append({
            **candidate,
            "baseline_use_coeff": float(baseline_use_coeff),
            "target_use_coeff": float(target_use_coeff),
            "target_use_time_kl": target_kl,
            "target_boundary_mass_shift": target_boundary,
            "target_feature_attenuation": target_attenuation,
            "controls": control_rows,
        })

        running_target = float(np.mean(target_kls)) if target_kls else 0.0
        running_control = float(np.mean(control_kls)) if control_kls else 0.0
        ratio = (
            running_target / running_control
            if abs(running_control) > 1e-8
            else float("inf")
        )
        print(
            f"[factor-trace-intervention] done {i}/{len(selected)} "
            f"target_kl={target_kl:.6f} control_kl={control_kls[-1]:.6f} "
            f"running_ratio={ratio:.2f}x "
            f"target_att={target_attenuation:.6f} control_att={control_attenuations[-1]:.6f} "
            f"elapsed={time.perf_counter() - prompt_t0:.1f}s",
            flush=True,
        )

    target_kl_mean = float(np.mean(target_kls)) if target_kls else 0.0
    control_kl_mean = float(np.mean(control_kls)) if control_kls else 0.0
    target_boundary_mean = float(np.mean(target_boundary_shifts)) if target_boundary_shifts else 0.0
    control_boundary_mean = float(np.mean(control_boundary_shifts)) if control_boundary_shifts else 0.0
    target_att_mean = float(np.mean(target_attenuations)) if target_attenuations else 0.0
    control_att_mean = float(np.mean(control_attenuations)) if control_attenuations else 0.0

    total_elapsed = time.perf_counter() - t0
    print(
        "[factor-trace-intervention] summary "
        f"prompts={len(per_prompt)} target_kl_mean={target_kl_mean:.6f} "
        f"control_kl_mean={control_kl_mean:.6f} "
        f"ratio={(target_kl_mean / control_kl_mean if abs(control_kl_mean) > 1e-8 else float('inf')):.2f}x "
        f"target_att_mean={target_att_mean:.6f} control_att_mean={control_att_mean:.6f} "
        f"elapsed={total_elapsed:.1f}s",
        flush=True,
    )

    return {
        "target_feature_idx": int(target_feature_idx),
        "similar_frequency_control": int(similar_feature),
        "random_controls": [int(idx) for idx in random_features],
        "n_prompts_evaluated": len(per_prompt),
        "total_time_s": total_elapsed,
        "per_prompt": per_prompt,
        "summary": {
            "target_use_time_kl_mean": target_kl_mean,
            "control_use_time_kl_mean": control_kl_mean,
            "target_vs_random_kl_ratio": (
                target_kl_mean / control_kl_mean if abs(control_kl_mean) > 1e-8 else float("inf")
            ),
            "target_boundary_mass_shift_mean": target_boundary_mean,
            "control_boundary_mass_shift_mean": control_boundary_mean,
            "target_feature_attenuation_mean": target_att_mean,
            "control_feature_attenuation_mean": control_att_mean,
        },
    }


def run_factor_trace_transplant(
    model,
    tokenizer,
    corpus_sequences: list[torch.Tensor],
    trace_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    states: torch.Tensor,
    target_feature_indices: Sequence[int] | int,
    n_pairs: int = 16,
    device: str = "cuda",
) -> dict[str, Any]:
    """Transplant a donor write value into matched recipient prompts and score local readouts."""
    if isinstance(target_feature_indices, (int, np.integer)):
        target_features = [int(target_feature_indices)]
    else:
        target_features = [int(idx) for idx in target_feature_indices]
    if not target_features:
        raise ValueError("Need at least one target feature for transplant")
    primary_target_feature = int(target_features[0])

    prompt_map = {
        int(prompt["prompt_index"]): prompt
        for prompt in trace_result.get("prompt_records", [])
    }
    sequence_map = {
        int(prompt["prompt_index"]): seq
        for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences)
    }
    candidate_pairs = select_transplant_pairs(
        trace_result=trace_result,
        target_feature_idx=primary_target_feature,
        max_pairs=max(n_pairs * 12, 64),
        preferred_events=("paragraph_boundary",),
        preferred_offsets=(16,),
        min_write_gain=1.20,
        min_context_similarity=0.20,
    )
    if not candidate_pairs:
        raise ValueError(f"No transplant pairs found for features {target_features}")

    act_matrix = compute_activation_matrix(sae, states)
    frequencies = (act_matrix > 0).mean(axis=0)
    similar_features: list[int] = []
    excluded = set(target_features)
    for feat in target_features:
        similar_feat = choose_similar_frequency_feature(
            target_feature_idx=int(feat),
            frequencies=frequencies,
            exclude=excluded | set(similar_features),
        )
        similar_features.append(int(similar_feat))
    token_groups = boundary_token_groups(tokenizer)
    t0 = time.perf_counter()

    per_pair = []
    target_boundary = []
    wrong_feature_boundary = []
    wrong_time_boundary = []
    target_newline = []
    wrong_feature_newline = []
    wrong_time_newline = []
    target_sentence = []
    wrong_feature_sentence = []
    wrong_time_sentence = []
    target_kl = []
    wrong_feature_kl = []
    wrong_time_kl = []
    target_coeff_delta = []
    wrong_feature_coeff_delta = []
    wrong_time_coeff_delta = []

    print(
        "[factor-trace-transplant] "
        f"target={target_features} similar_control={similar_features} "
        f"candidate_pairs={len(candidate_pairs)}",
        flush=True,
    )

    baseline_candidates: list[dict[str, Any]] = []
    tracked_features = [*target_features, *similar_features]
    for pair in candidate_pairs:
        recipient_prompt_index = int(pair["recipient_prompt_index"])
        donor_prompt_index = int(pair["donor_prompt_index"])
        donor_prompt = prompt_map[donor_prompt_index]
        recipient_full = sequence_map[recipient_prompt_index].to(device)
        donor_coeffs = {
            int(feat): float(donor_prompt["feature_coefficients"][str(int(feat))][pair["donor_write_pos"]])
            for feat in target_features
        }
        use_pos = int(pair["recipient_use_pos"])
        eval_forward_window = 8
        seq_end = min(int(recipient_full.shape[1]), use_pos + eval_forward_window + 2)
        if seq_end <= use_pos + 1:
            continue
        seq = recipient_full[:, :seq_end]
        actual_next_token = int(recipient_full[0, use_pos + 1].item())
        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
        )
        eval_meta = _select_eval_site(
            baseline_logits_by_pos=baseline["logits_by_pos"],
            input_ids=seq,
            anchor_use_pos=use_pos,
            token_groups=token_groups,
            tokenizer=tokenizer,
            preferred_groups=("newline", "sentence_punct", "colon"),
        )
        if not eval_meta:
            continue
        eval_pos = int(eval_meta["eval_pos"])
        baseline_logits = baseline["logits_by_pos"][eval_pos]
        eval_meta["baseline_boundary_mass"] = _mass_from_logits(baseline_logits, token_groups["boundary"])
        eval_meta["baseline_actual_group_mass"] = _mass_from_logits(
            baseline_logits,
            token_groups[eval_meta["eval_actual_group"]],
        )
        eval_meta["baseline_actual_token_prob"] = _token_prob_from_logits(
            baseline_logits,
            int(eval_meta["eval_actual_next_token"]),
        )
        if (
            eval_meta["baseline_actual_token_prob"] < 0.03
            or eval_meta["baseline_actual_token_prob"] > 0.75
        ):
            continue
        target_updates = {}
        wrong_feature_updates = {}
        sign_flip_updates = {}
        for feat, donor_coeff in donor_coeffs.items():
            recipient_write_coeff = float(
                baseline["feature_coefficients"][str(int(feat))][int(pair["recipient_write_pos"])]
            )
            donor_prev_coeff = 0.0
            if int(pair["donor_write_pos"]) > 0:
                donor_prev_coeff = float(
                    donor_prompt["feature_coefficients"][str(int(feat))][int(pair["donor_write_pos"]) - 1]
                )
            donor_delta = float(donor_coeff - donor_prev_coeff)
            target_updates[int(feat)] = float(recipient_write_coeff + donor_delta)
            sign_flip_updates[int(feat)] = float(recipient_write_coeff - donor_delta)
        for feat, ctrl in zip(target_features, similar_features):
            ctrl_write_coeff = float(
                baseline["feature_coefficients"][str(int(ctrl))][int(pair["recipient_write_pos"])]
            )
            donor_prev_coeff = 0.0
            if int(pair["donor_write_pos"]) > 0:
                donor_prev_coeff = float(
                    donor_prompt["feature_coefficients"][str(int(feat))][int(pair["donor_write_pos"]) - 1]
                )
            donor_delta = float(donor_coeffs[int(feat)] - donor_prev_coeff)
            wrong_feature_updates[int(ctrl)] = float(ctrl_write_coeff + donor_delta)
        headroom = eval_meta["baseline_actual_token_prob"] * (1.0 - eval_meta["baseline_actual_token_prob"])
        donor_gap = float(pair["donor_write_coeff"]) - float(pair["recipient_write_coeff"])
        baseline_score = (
            headroom,
            float(eval_meta["baseline_actual_group_mass"]),
            float(pair["context_similarity"]),
            donor_gap,
        )
        baseline_candidates.append(
            {
                "pair": pair,
                "seq": seq,
                "actual_next_token": actual_next_token,
                "baseline": baseline,
                "eval_meta": eval_meta,
                "use_pos": use_pos,
                "target_updates": {int(k): float(v) for k, v in target_updates.items()},
                "wrong_feature_updates": {int(k): float(v) for k, v in wrong_feature_updates.items()},
                "sign_flip_updates": {int(k): float(v) for k, v in sign_flip_updates.items()},
                "score": baseline_score,
            }
        )

    baseline_candidates.sort(key=lambda item: item["score"], reverse=True)
    if not baseline_candidates:
        raise ValueError("No transplant candidates survived exact-token headroom filtering")

    print(
        "[factor-trace-transplant] "
        f"baseline-screened={len(baseline_candidates)} selected={min(len(baseline_candidates), n_pairs)}",
        flush=True,
    )

    for item in baseline_candidates[:n_pairs]:
        pair = item["pair"]
        seq = item["seq"]
        actual_next_token = int(item["actual_next_token"])
        baseline = item["baseline"]
        eval_meta = dict(item["eval_meta"])
        use_pos = int(item["use_pos"])
        eval_pos = int(eval_meta["eval_pos"])
        target_updates = {int(k): float(v) for k, v in item["target_updates"].items()}
        wrong_feature_updates = {int(k): float(v) for k, v in item["wrong_feature_updates"].items()}
        sign_flip_updates = {int(k): float(v) for k, v in item["sign_flip_updates"].items()}
        wrong_time_pos = pair.get("wrong_time_pos")
        prompt_t0 = time.perf_counter()

        recipient_prompt_index = int(pair["recipient_prompt_index"])
        donor_prompt_index = int(pair["donor_prompt_index"])

        print(
            f"[factor-trace-transplant] pair {len(per_pair)+1}/{min(len(baseline_candidates), n_pairs)} "
            f"donor={donor_prompt_index}@{pair['donor_write_pos']} "
            f"recipient={recipient_prompt_index}@{pair['recipient_write_pos']} "
            f"use={pair['recipient_use_pos']} offset={pair['recipient_use_offset']} "
            f"ctx={pair['context_similarity']:.2f} "
            f"tok={eval_meta['eval_actual_text']!r} p={eval_meta['baseline_actual_token_prob']:.3f}",
            flush=True,
        )

        target_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=target_updates,
            intervention_pos=int(pair["recipient_write_pos"]),
        )
        wrong_feature_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=wrong_feature_updates,
            intervention_pos=int(pair["recipient_write_pos"]),
        )
        sign_flip_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=sign_flip_updates,
            intervention_pos=int(pair["recipient_write_pos"]),
        )

        target_logits = target_rollout["logits_by_pos"][eval_pos]
        wrong_feature_logits = wrong_feature_rollout["logits_by_pos"][eval_pos]
        sign_flip_logits = sign_flip_rollout["logits_by_pos"][eval_pos]

        baseline_use_coeff = float(
            sum(float(baseline["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features)
        )
        target_use_coeff = float(
            sum(float(target_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features)
        )
        wrong_feature_use_coeff = float(
            sum(float(wrong_feature_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features)
        )
        sign_flip_use_coeff = float(
            sum(float(sign_flip_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features)
        )

        target_boundary_shift = _boundary_mass_shift(baseline_logits, target_logits, token_groups["boundary"])
        wrong_feature_boundary_shift = _boundary_mass_shift(
            baseline_logits, wrong_feature_logits, token_groups["boundary"]
        )
        sign_flip_boundary_shift = _boundary_mass_shift(
            baseline_logits, sign_flip_logits, token_groups["boundary"]
        )
        target_newline_shift = _boundary_mass_shift(baseline_logits, target_logits, token_groups["newline"])
        wrong_feature_newline_shift = _boundary_mass_shift(
            baseline_logits, wrong_feature_logits, token_groups["newline"]
        )
        sign_flip_newline_shift = _boundary_mass_shift(
            baseline_logits, sign_flip_logits, token_groups["newline"]
        )
        target_sentence_shift = _boundary_mass_shift(
            baseline_logits, target_logits, token_groups["sentence_punct"]
        )
        wrong_feature_sentence_shift = _boundary_mass_shift(
            baseline_logits, wrong_feature_logits, token_groups["sentence_punct"]
        )
        sign_flip_sentence_shift = _boundary_mass_shift(
            baseline_logits, sign_flip_logits, token_groups["sentence_punct"]
        )

        wrong_time_boundary_shift = 0.0
        wrong_time_newline_shift = 0.0
        wrong_time_sentence_shift = 0.0
        wrong_time_kl_value = 0.0
        wrong_time_coeff_value = baseline_use_coeff
        wrong_time_actual_logit_shift = 0.0

        sign_flip_kl_value = _kl_divergence_from_logits(baseline_logits, sign_flip_logits)
        sign_flip_actual_logit_shift = _token_logit_shift(
            baseline_logits,
            sign_flip_logits,
            int(eval_meta["eval_actual_next_token"])
            if eval_meta["eval_actual_next_token"] is not None
            else actual_next_token,
        )
        if wrong_time_pos is not None:
            wrong_time_rollout = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_updates=target_updates,
                intervention_pos=int(wrong_time_pos),
            )
            wrong_time_logits = wrong_time_rollout["logits_by_pos"][eval_pos]
            wrong_time_boundary_shift = _boundary_mass_shift(
                baseline_logits, wrong_time_logits, token_groups["boundary"]
            )
            wrong_time_newline_shift = _boundary_mass_shift(
                baseline_logits, wrong_time_logits, token_groups["newline"]
            )
            wrong_time_sentence_shift = _boundary_mass_shift(
                baseline_logits, wrong_time_logits, token_groups["sentence_punct"]
            )
            wrong_time_kl_value = _kl_divergence_from_logits(baseline_logits, wrong_time_logits)
            wrong_time_actual_logit_shift = _token_logit_shift(
                baseline_logits,
                wrong_time_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            )
            wrong_time_coeff_value = float(
                sum(
                    float(wrong_time_rollout["feature_coefficients"][str(int(feat))][use_pos])
                    for feat in target_features
                )
            )

        row = {
            **pair,
            **eval_meta,
            "target_feature_indices": [int(feat) for feat in target_features],
            "similar_frequency_controls": [int(feat) for feat in similar_features],
            "primary_target_feature_idx": primary_target_feature,
            "donor_values": {str(k): float(v) for k, v in donor_coeffs.items()},
            "actual_next_token": int(eval_meta["eval_actual_next_token"])
            if eval_meta["eval_actual_next_token"] is not None
            else actual_next_token,
            "actual_next_token_logit_shift_target": _token_logit_shift(
                baseline_logits,
                target_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ),
            "actual_next_token_prob_shift_target": _token_prob_shift(
                baseline_logits,
                target_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ),
            "actual_next_token_logit_shift_wrong_feature": _token_logit_shift(
                baseline_logits,
                wrong_feature_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ),
            "actual_next_token_prob_shift_wrong_feature": _token_prob_shift(
                baseline_logits,
                wrong_feature_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ),
            "actual_next_token_logit_shift_sign_flip": sign_flip_actual_logit_shift,
            "actual_next_token_prob_shift_sign_flip": _token_prob_shift(
                baseline_logits,
                sign_flip_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ),
            "actual_next_token_logit_shift_wrong_time": wrong_time_actual_logit_shift,
            "actual_next_token_prob_shift_wrong_time": _token_prob_shift(
                baseline_logits,
                wrong_time_logits,
                int(eval_meta["eval_actual_next_token"])
                if eval_meta["eval_actual_next_token"] is not None
                else actual_next_token,
            ) if wrong_time_pos is not None else 0.0,
            "target_use_time_kl": _kl_divergence_from_logits(baseline_logits, target_logits),
            "wrong_feature_use_time_kl": _kl_divergence_from_logits(baseline_logits, wrong_feature_logits),
            "sign_flip_use_time_kl": sign_flip_kl_value,
            "wrong_time_use_time_kl": wrong_time_kl_value,
            "target_boundary_mass_shift": target_boundary_shift,
            "wrong_feature_boundary_mass_shift": wrong_feature_boundary_shift,
            "sign_flip_boundary_mass_shift": sign_flip_boundary_shift,
            "wrong_time_boundary_mass_shift": wrong_time_boundary_shift,
            "target_newline_mass_shift": target_newline_shift,
            "wrong_feature_newline_mass_shift": wrong_feature_newline_shift,
            "sign_flip_newline_mass_shift": sign_flip_newline_shift,
            "wrong_time_newline_mass_shift": wrong_time_newline_shift,
            "target_sentence_mass_shift": target_sentence_shift,
            "wrong_feature_sentence_mass_shift": wrong_feature_sentence_shift,
            "sign_flip_sentence_mass_shift": sign_flip_sentence_shift,
            "wrong_time_sentence_mass_shift": wrong_time_sentence_shift,
            "baseline_use_coeff": baseline_use_coeff,
            "target_use_coeff": target_use_coeff,
            "wrong_feature_use_coeff": wrong_feature_use_coeff,
            "sign_flip_use_coeff": sign_flip_use_coeff,
            "wrong_time_use_coeff": wrong_time_coeff_value,
            "target_use_coeff_delta": float(target_use_coeff - baseline_use_coeff),
            "wrong_feature_use_coeff_delta": float(wrong_feature_use_coeff - baseline_use_coeff),
            "sign_flip_use_coeff_delta": float(sign_flip_use_coeff - baseline_use_coeff),
            "wrong_time_use_coeff_delta": float(wrong_time_coeff_value - baseline_use_coeff),
        }
        per_pair.append(row)

        target_boundary.append(row["target_boundary_mass_shift"])
        wrong_feature_boundary.append(row["wrong_feature_boundary_mass_shift"])
        wrong_time_boundary.append(row["wrong_time_boundary_mass_shift"])
        target_newline.append(row["target_newline_mass_shift"])
        wrong_feature_newline.append(row["wrong_feature_newline_mass_shift"])
        wrong_time_newline.append(row["wrong_time_newline_mass_shift"])
        target_sentence.append(row["target_sentence_mass_shift"])
        wrong_feature_sentence.append(row["wrong_feature_sentence_mass_shift"])
        wrong_time_sentence.append(row["wrong_time_sentence_mass_shift"])
        target_kl.append(row["target_use_time_kl"])
        wrong_feature_kl.append(row["wrong_feature_use_time_kl"])
        wrong_time_kl.append(row["wrong_time_use_time_kl"])
        target_coeff_delta.append(row["target_use_coeff_delta"])
        wrong_feature_coeff_delta.append(row["wrong_feature_use_coeff_delta"])
        wrong_time_coeff_delta.append(row["wrong_time_use_coeff_delta"])

        print(
            f"[factor-trace-transplant] done {len(per_pair)}/{n_pairs} "
            f"boundary={row['target_boundary_mass_shift']:+.6f} "
            f"wrong_feature={row['wrong_feature_boundary_mass_shift']:+.6f} "
            f"wrong_time={row['wrong_time_boundary_mass_shift']:+.6f} "
            f"coeff_delta={row['target_use_coeff_delta']:+.6f} "
            f"elapsed={time.perf_counter() - prompt_t0:.1f}s",
            flush=True,
        )

    def _mean(values: list[float]) -> float:
        return float(np.mean(values)) if values else 0.0

    summary = {
        "target_boundary_mass_shift_mean": _mean(target_boundary),
        "wrong_feature_boundary_mass_shift_mean": _mean(wrong_feature_boundary),
        "wrong_time_boundary_mass_shift_mean": _mean(wrong_time_boundary),
        "target_newline_mass_shift_mean": _mean(target_newline),
        "wrong_feature_newline_mass_shift_mean": _mean(wrong_feature_newline),
        "wrong_time_newline_mass_shift_mean": _mean(wrong_time_newline),
        "target_sentence_mass_shift_mean": _mean(target_sentence),
        "wrong_feature_sentence_mass_shift_mean": _mean(wrong_feature_sentence),
        "wrong_time_sentence_mass_shift_mean": _mean(wrong_time_sentence),
        "target_use_time_kl_mean": _mean(target_kl),
        "wrong_feature_use_time_kl_mean": _mean(wrong_feature_kl),
        "wrong_time_use_time_kl_mean": _mean(wrong_time_kl),
        "target_use_coeff_delta_mean": _mean(target_coeff_delta),
        "wrong_feature_use_coeff_delta_mean": _mean(wrong_feature_coeff_delta),
        "wrong_time_use_coeff_delta_mean": _mean(wrong_time_coeff_delta),
        "target_actual_token_logit_shift_mean": _mean(
            [float(row["actual_next_token_logit_shift_target"]) for row in per_pair]
        ),
        "wrong_feature_actual_token_logit_shift_mean": _mean(
            [float(row["actual_next_token_logit_shift_wrong_feature"]) for row in per_pair]
        ),
        "sign_flip_actual_token_logit_shift_mean": _mean(
            [float(row["actual_next_token_logit_shift_sign_flip"]) for row in per_pair]
        ),
        "target_actual_token_prob_shift_mean": _mean(
            [float(row["actual_next_token_prob_shift_target"]) for row in per_pair]
        ),
        "wrong_feature_actual_token_prob_shift_mean": _mean(
            [float(row["actual_next_token_prob_shift_wrong_feature"]) for row in per_pair]
        ),
        "sign_flip_actual_token_prob_shift_mean": _mean(
            [float(row["actual_next_token_prob_shift_sign_flip"]) for row in per_pair]
        ),
        "wrong_time_actual_token_prob_shift_mean": _mean(
            [float(row["actual_next_token_prob_shift_wrong_time"]) for row in per_pair]
        ),
        "wrong_time_actual_token_logit_shift_mean": _mean(
            [float(row["actual_next_token_logit_shift_wrong_time"]) for row in per_pair]
        ),
        "target_minus_sign_flip_actual_token_logit_shift_mean": _mean(
            [
                float(row["actual_next_token_logit_shift_target"]) - float(row["actual_next_token_logit_shift_sign_flip"])
                for row in per_pair
            ]
        ),
        "target_minus_sign_flip_actual_token_prob_shift_mean": _mean(
            [
                float(row["actual_next_token_prob_shift_target"]) - float(row["actual_next_token_prob_shift_sign_flip"])
                for row in per_pair
            ]
        ),
        "target_beats_wrong_feature_fraction": float(
            np.mean([a > b for a, b in zip(target_boundary, wrong_feature_boundary)])
        ) if target_boundary else 0.0,
        "target_beats_sign_flip_fraction": float(
            np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_sign_flip"] for row in per_pair])
        ) if per_pair else 0.0,
        "target_prob_beats_sign_flip_fraction": float(
            np.mean([row["actual_next_token_prob_shift_target"] > row["actual_next_token_prob_shift_sign_flip"] for row in per_pair])
        ) if per_pair else 0.0,
        "target_beats_wrong_time_fraction": float(
            np.mean([a > b for a, b in zip(target_boundary, wrong_time_boundary)])
        ) if target_boundary else 0.0,
        "target_actual_token_beats_wrong_feature_fraction": float(
            np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_wrong_feature"] for row in per_pair])
        ) if per_pair else 0.0,
        "target_actual_token_beats_wrong_time_fraction": float(
            np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_wrong_time"] for row in per_pair])
        ) if per_pair else 0.0,
        "target_coeff_beats_sign_flip_fraction": float(
            np.mean([row["target_use_coeff_delta"] > row["sign_flip_use_coeff_delta"] for row in per_pair])
        ) if per_pair else 0.0,
        "target_minus_sign_flip_use_coeff_delta_mean": _mean(
            [float(row["target_use_coeff_delta"]) - float(row["sign_flip_use_coeff_delta"]) for row in per_pair]
        ),
    }
    summary["boundary_ratio_vs_wrong_feature"] = (
        summary["target_boundary_mass_shift_mean"] / summary["wrong_feature_boundary_mass_shift_mean"]
        if abs(summary["wrong_feature_boundary_mass_shift_mean"]) > 1e-8
        else float("inf")
    )
    summary["boundary_ratio_vs_wrong_time"] = (
        summary["target_boundary_mass_shift_mean"] / summary["wrong_time_boundary_mass_shift_mean"]
        if abs(summary["wrong_time_boundary_mass_shift_mean"]) > 1e-8
        else float("inf")
    )

    return {
        "target_feature_idx": primary_target_feature,
        "target_feature_indices": [int(feat) for feat in target_features],
        "similar_frequency_control": int(similar_features[0]),
        "similar_frequency_controls": [int(feat) for feat in similar_features],
        "n_pairs_evaluated": len(per_pair),
        "total_time_s": float(time.perf_counter() - t0),
        "per_pair": per_pair,
        "summary": summary,
    }


def run_use_site_signed_causal(
    model,
    tokenizer,
    corpus_sequences: list[torch.Tensor],
    benchmark_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    states: torch.Tensor,
    n_rows: int = 32,
    device: str = "cuda",
) -> dict[str, Any]:
    """Edit the signed H4 direction at the benchmarked use site and score local decisions."""
    target_features = [int(x) for x in benchmark_result.get("target_feature_indices", [])]
    if target_features != [62, 105]:
        raise ValueError(f"Expected signed pair [62, 105], got {target_features}")
    accepted_rows = list(benchmark_result.get("accepted_rows", []))
    if not accepted_rows:
        raise ValueError("No accepted benchmark rows for use-site causal stage")

    act_matrix = compute_activation_matrix(sae, states)
    frequencies = (act_matrix > 0).mean(axis=0)
    wrong_feature_pair = [
        choose_similar_frequency_feature(target_features[0], frequencies, exclude=set(target_features)),
        choose_similar_frequency_feature(
            target_features[1],
            frequencies,
            exclude=set(target_features) | {int(choose_similar_frequency_feature(target_features[0], frequencies, exclude=set(target_features)))},
        ),
    ]
    wrong_feature_pair = [int(x) for x in wrong_feature_pair]
    readout_token_ids = _readout_class_token_ids(tokenizer)
    t0 = time.perf_counter()

    sequence_map = {idx: seq.to(device) for idx, seq in enumerate(corpus_sequences)}
    tracked_features = [*target_features, *wrong_feature_pair]

    accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True)
    rows = accepted_rows[:n_rows]
    per_row = []

    print(
        "[factor-trace-use-site] "
        f"target={target_features} reverse_sign_control=True wrong_time=True rows={len(rows)}",
        flush=True,
    )

    for i, row in enumerate(rows, start=1):
        recipient_idx = int(row["recipient_prompt_index"])
        seq_full = sequence_map[recipient_idx]
        decision_pos = int(row["decision_pos"])
        intervention_pos = int(row["intervention_pos"])
        wrong_time_pos = row.get("wrong_time_pos")
        actual_next_token = int(row["actual_next_token"])
        seq_end = min(int(seq_full.shape[1]), decision_pos + 3)
        if seq_end <= decision_pos + 1:
            continue
        seq = seq_full[:, :seq_end]
        prompt_t0 = time.perf_counter()

        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
        )

        target_updates = {}
        reverse_updates = {}
        wrong_feature_updates = {}
        for feat, delta in zip(target_features, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]):
            base_coeff = float(baseline["feature_coefficients"][str(int(feat))][intervention_pos])
            target_updates[int(feat)] = float(base_coeff + delta)
            reverse_updates[int(feat)] = float(base_coeff - delta)
        for ctrl, delta in zip(wrong_feature_pair, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]):
            base_ctrl = float(baseline["feature_coefficients"][str(int(ctrl))][intervention_pos])
            wrong_feature_updates[int(ctrl)] = float(base_ctrl + delta)

        target_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=target_updates,
            intervention_pos=intervention_pos,
        )
        reverse_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=reverse_updates,
            intervention_pos=intervention_pos,
        )
        wrong_feature_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
            intervention_updates=wrong_feature_updates,
            intervention_pos=intervention_pos,
        )

        wrong_time_rollout = None
        if wrong_time_pos is not None:
            wrong_time_rollout = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_updates=target_updates,
                intervention_pos=int(wrong_time_pos),
            )

        baseline_logits = baseline["logits_by_pos"][decision_pos]
        target_logits = target_rollout["logits_by_pos"][decision_pos]
        reverse_logits = reverse_rollout["logits_by_pos"][decision_pos]
        wrong_feature_logits = wrong_feature_rollout["logits_by_pos"][decision_pos]
        wrong_time_logits = wrong_time_rollout["logits_by_pos"][decision_pos] if wrong_time_rollout is not None else baseline_logits

        readout_class = str(row["readout_class"])
        class_ids = readout_token_ids.get(readout_class, [actual_next_token]) or [actual_next_token]
        next_positions = [pos for pos in [decision_pos, min(decision_pos + 1, int(seq.shape[1]) - 1)] if pos in baseline["logits_by_pos"]]

        def _signed_coeff_change(rollout: dict[str, Any]) -> float:
            diffs = []
            for pos in next_positions:
                base = float(baseline["feature_coefficients"]["62"][pos] - baseline["feature_coefficients"]["105"][pos])
                edit = float(rollout["feature_coefficients"]["62"][pos] - rollout["feature_coefficients"]["105"][pos])
                diffs.append(edit - base)
            return float(np.mean(diffs)) if diffs else 0.0

        row_out = {
            **row,
            "actual_token_logit_shift_target": _token_logit_shift(baseline_logits, target_logits, actual_next_token),
            "actual_token_logit_shift_reverse": _token_logit_shift(baseline_logits, reverse_logits, actual_next_token),
            "actual_token_logit_shift_wrong_feature": _token_logit_shift(baseline_logits, wrong_feature_logits, actual_next_token),
            "actual_token_logit_shift_wrong_time": _token_logit_shift(baseline_logits, wrong_time_logits, actual_next_token),
            "actual_token_prob_shift_target": _token_prob_shift(baseline_logits, target_logits, actual_next_token),
            "actual_token_prob_shift_reverse": _token_prob_shift(baseline_logits, reverse_logits, actual_next_token),
            "actual_token_prob_shift_wrong_feature": _token_prob_shift(baseline_logits, wrong_feature_logits, actual_next_token),
            "actual_token_prob_shift_wrong_time": _token_prob_shift(baseline_logits, wrong_time_logits, actual_next_token),
            "readout_class_prob_shift_target": _boundary_mass_shift(baseline_logits, target_logits, class_ids),
            "readout_class_prob_shift_reverse": _boundary_mass_shift(baseline_logits, reverse_logits, class_ids),
            "readout_class_prob_shift_wrong_feature": _boundary_mass_shift(baseline_logits, wrong_feature_logits, class_ids),
            "readout_class_prob_shift_wrong_time": _boundary_mass_shift(baseline_logits, wrong_time_logits, class_ids),
            "signed_direction_coeff_shift_target": _signed_coeff_change(target_rollout),
            "signed_direction_coeff_shift_reverse": _signed_coeff_change(reverse_rollout),
            "signed_direction_coeff_shift_wrong_feature": _signed_coeff_change(wrong_feature_rollout),
            "signed_direction_coeff_shift_wrong_time": _signed_coeff_change(wrong_time_rollout) if wrong_time_rollout is not None else 0.0,
            "local_kl_target": _kl_divergence_from_logits(baseline_logits, target_logits),
            "local_kl_reverse": _kl_divergence_from_logits(baseline_logits, reverse_logits),
            "local_kl_wrong_time": _kl_divergence_from_logits(baseline_logits, wrong_time_logits) if wrong_time_rollout is not None else 0.0,
        }
        per_row.append(row_out)
        print(
            f"[factor-trace-use-site] row {len(per_row)}/{len(rows)} "
            f"tok={row['readout_text']!r} p={row['baseline_actual_token_prob']:.3f} "
            f"target_prob={row_out['actual_token_prob_shift_target']:+.4f} "
            f"reverse_prob={row_out['actual_token_prob_shift_reverse']:+.4f} "
            f"wrong_time_prob={row_out['actual_token_prob_shift_wrong_time']:+.4f} "
            f"signed_shift={row_out['signed_direction_coeff_shift_target']:+.4f} "
            f"elapsed={time.perf_counter() - prompt_t0:.1f}s",
            flush=True,
        )

    def _mean(xs: list[float]) -> float:
        return float(np.mean(xs)) if xs else 0.0

    summary = {
        "target_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_target"] for r in per_row]),
        "reverse_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_reverse"] for r in per_row]),
        "wrong_time_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_wrong_time"] for r in per_row]),
        "target_minus_reverse_actual_token_prob_shift_mean": _mean(
            [r["actual_token_prob_shift_target"] - r["actual_token_prob_shift_reverse"] for r in per_row]
        ),
        "target_actual_token_logit_shift_mean": _mean([r["actual_token_logit_shift_target"] for r in per_row]),
        "reverse_actual_token_logit_shift_mean": _mean([r["actual_token_logit_shift_reverse"] for r in per_row]),
        "target_minus_reverse_actual_token_logit_shift_mean": _mean(
            [r["actual_token_logit_shift_target"] - r["actual_token_logit_shift_reverse"] for r in per_row]
        ),
        "target_readout_class_prob_shift_mean": _mean([r["readout_class_prob_shift_target"] for r in per_row]),
        "reverse_readout_class_prob_shift_mean": _mean([r["readout_class_prob_shift_reverse"] for r in per_row]),
        "target_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_target"] for r in per_row]),
        "reverse_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_reverse"] for r in per_row]),
        "wrong_time_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_wrong_time"] for r in per_row]),
        "target_beats_reverse_prob_fraction": float(np.mean([
            r["actual_token_prob_shift_target"] > r["actual_token_prob_shift_reverse"] for r in per_row
        ])) if per_row else 0.0,
        "target_beats_wrong_time_prob_fraction": float(np.mean([
            r["actual_token_prob_shift_target"] > r["actual_token_prob_shift_wrong_time"] for r in per_row
        ])) if per_row else 0.0,
        "target_localized_coeff_beats_wrong_time_fraction": float(np.mean([
            r["signed_direction_coeff_shift_target"] > r["signed_direction_coeff_shift_wrong_time"] for r in per_row
        ])) if per_row else 0.0,
    }
    summary["promotion_gate_pass"] = bool(
        summary["target_beats_reverse_prob_fraction"] >= 0.65
        and summary["target_minus_reverse_actual_token_prob_shift_mean"] > 0.0
        and summary["target_beats_wrong_time_prob_fraction"] >= 0.60
        and summary["target_signed_direction_coeff_shift_mean"] > summary["wrong_time_signed_direction_coeff_shift_mean"]
    )
    return {
        "target_feature_indices": target_features,
        "wrong_feature_pair": wrong_feature_pair,
        "n_rows_evaluated": len(per_row),
        "per_row": per_row,
        "summary": summary,
        "total_time_s": float(time.perf_counter() - t0),
    }


def run_use_site_readout_map(
    model,
    tokenizer,
    corpus_sequences: list[torch.Tensor],
    benchmark_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    n_rows: int = 32,
    device: str = "cuda",
) -> dict[str, Any]:
    """Map localized downstream readouts for the signed H4 boundary direction."""
    del device  # kept for interface symmetry
    target_features = [int(x) for x in benchmark_result.get("target_feature_indices", [])]
    if target_features != [62, 105]:
        raise ValueError(f"Expected signed pair [62, 105], got {target_features}")
    accepted_rows = list(benchmark_result.get("accepted_rows", []))
    if not accepted_rows:
        raise ValueError("No accepted benchmark rows for readout-map stage")

    t0 = time.perf_counter()
    sequence_map = {idx: seq for idx, seq in enumerate(corpus_sequences)}
    token_class_ids = _readout_class_token_ids(tokenizer)
    accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True)
    rows = accepted_rows[:n_rows]
    per_row = []

    hidden_sum: np.ndarray | None = None
    hidden_abs_sum: np.ndarray | None = None
    hidden_pos_counts: np.ndarray | None = None
    logit_sum: np.ndarray | None = None
    logit_abs_sum: np.ndarray | None = None
    logit_pos_counts: np.ndarray | None = None

    print(
        "[factor-trace-readout-map] "
        f"target={target_features} reverse_sign_control=True rows={len(rows)}",
        flush=True,
    )

    for i, row in enumerate(rows, start=1):
        recipient_idx = int(row["recipient_prompt_index"])
        seq_full = sequence_map[recipient_idx]
        decision_pos = int(row["decision_pos"])
        intervention_pos = int(row["intervention_pos"])
        wrong_time_pos = row.get("wrong_time_pos")
        actual_next_token = int(row["actual_next_token"])
        seq_end = min(int(seq_full.shape[1]), decision_pos + 3)
        if seq_end <= decision_pos + 1:
            continue
        seq = seq_full[:, :seq_end]
        prompt_t0 = time.perf_counter()
        capture_positions = [decision_pos, min(decision_pos + 1, seq_end - 1)]

        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=target_features,
            capture_hidden_state_positions=capture_positions,
        )
        target_updates = {}
        reverse_updates = {}
        for feat, delta in zip(target_features, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]):
            base_coeff = float(baseline["feature_coefficients"][str(int(feat))][intervention_pos])
            target_updates[int(feat)] = float(base_coeff + delta)
            reverse_updates[int(feat)] = float(base_coeff - delta)

        target_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=target_features,
            intervention_updates=target_updates,
            intervention_pos=intervention_pos,
            capture_hidden_state_positions=capture_positions,
        )
        reverse_rollout = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=target_features,
            intervention_updates=reverse_updates,
            intervention_pos=intervention_pos,
            capture_hidden_state_positions=capture_positions,
        )
        wrong_time_rollout = None
        if wrong_time_pos is not None:
            wrong_time_rollout = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=target_features,
                intervention_updates=target_updates,
                intervention_pos=int(wrong_time_pos),
                capture_hidden_state_positions=capture_positions,
            )

        baseline_logits = baseline["logits_by_pos"][decision_pos]
        target_logits = target_rollout["logits_by_pos"][decision_pos]
        reverse_logits = reverse_rollout["logits_by_pos"][decision_pos]
        wrong_time_logits = wrong_time_rollout["logits_by_pos"][decision_pos] if wrong_time_rollout is not None else baseline_logits
        readout_class = str(row["readout_class"])
        class_ids = token_class_ids.get(readout_class, [actual_next_token]) or [actual_next_token]

        hidden_delta_vecs = []
        for pos in capture_positions:
            base_hidden = baseline["hidden_state_by_pos"].get(pos)
            target_hidden = target_rollout["hidden_state_by_pos"].get(pos)
            reverse_hidden = reverse_rollout["hidden_state_by_pos"].get(pos)
            if base_hidden is None or target_hidden is None or reverse_hidden is None:
                continue
            delta_vec = (target_hidden - reverse_hidden).squeeze(0).numpy().astype(np.float32)
            hidden_delta_vecs.append(delta_vec)
        if hidden_delta_vecs:
            hidden_delta = np.mean(np.stack(hidden_delta_vecs, axis=0), axis=0)
            if hidden_sum is None:
                hidden_sum = np.zeros_like(hidden_delta)
                hidden_abs_sum = np.zeros_like(hidden_delta)
                hidden_pos_counts = np.zeros_like(hidden_delta)
            hidden_sum += hidden_delta
            hidden_abs_sum += np.abs(hidden_delta)
            hidden_pos_counts += (hidden_delta > 0).astype(np.float32)
        else:
            hidden_delta = np.zeros(0, dtype=np.float32)

        logit_delta = (target_logits - reverse_logits).squeeze(0).numpy().astype(np.float32)
        if logit_sum is None:
            logit_sum = np.zeros_like(logit_delta)
            logit_abs_sum = np.zeros_like(logit_delta)
            logit_pos_counts = np.zeros_like(logit_delta)
        logit_sum += logit_delta
        logit_abs_sum += np.abs(logit_delta)
        logit_pos_counts += (logit_delta > 0).astype(np.float32)

        per_row.append(
            {
                **row,
                "actual_token_logit_delta_target_minus_reverse": _token_logit_shift(reverse_logits, target_logits, actual_next_token),
                "actual_token_prob_delta_target_minus_reverse": _token_prob_shift(reverse_logits, target_logits, actual_next_token),
                "readout_class_prob_delta_target_minus_reverse": _boundary_mass_shift(reverse_logits, target_logits, class_ids),
                "actual_token_logit_delta_target_minus_wrong_time": _token_logit_shift(wrong_time_logits, target_logits, actual_next_token),
                "actual_token_prob_delta_target_minus_wrong_time": _token_prob_shift(wrong_time_logits, target_logits, actual_next_token),
                "hidden_state_delta_l2_target_minus_reverse": float(np.linalg.norm(hidden_delta)) if hidden_delta.size else 0.0,
                "elapsed_s": float(time.perf_counter() - prompt_t0),
            }
        )
        print(
            f"[factor-trace-readout-map] row {len(per_row)}/{len(rows)} "
            f"tok={row['readout_text']!r} class={readout_class} "
            f"delta_prob={per_row[-1]['actual_token_prob_delta_target_minus_reverse']:+.4f} "
            f"delta_logit={per_row[-1]['actual_token_logit_delta_target_minus_reverse']:+.4f} "
            f"hidden_l2={per_row[-1]['hidden_state_delta_l2_target_minus_reverse']:.4f} "
            f"elapsed={per_row[-1]['elapsed_s']:.1f}s",
            flush=True,
        )

    n_eval = max(len(per_row), 1)

    def _rank_dims(sum_arr: np.ndarray | None, abs_arr: np.ndarray | None, pos_arr: np.ndarray | None, top_k: int = 16) -> list[dict[str, float]]:
        if sum_arr is None or abs_arr is None or pos_arr is None:
            return []
        mean = sum_arr / n_eval
        mean_abs = abs_arr / n_eval
        pos_frac = pos_arr / n_eval
        sign_consistency = np.abs(2.0 * pos_frac - 1.0)
        score = np.abs(mean) * sign_consistency
        order = np.argsort(-score)[:top_k]
        ranked = []
        for idx in order:
            ranked.append(
                {
                    "index": int(idx),
                    "mean_delta": float(mean[idx]),
                    "mean_abs_delta": float(mean_abs[idx]),
                    "sign_consistency": float(sign_consistency[idx]),
                    "positive_fraction": float(pos_frac[idx]),
                    "score": float(score[idx]),
                }
            )
        return ranked

    def _rank_tokens(sum_arr: np.ndarray | None, abs_arr: np.ndarray | None, pos_arr: np.ndarray | None, top_k: int = 16) -> dict[str, list[dict[str, float | int | str]]]:
        if sum_arr is None or abs_arr is None or pos_arr is None:
            return {"positive": [], "negative": []}
        mean = sum_arr / n_eval
        mean_abs = abs_arr / n_eval
        pos_frac = pos_arr / n_eval
        pos_order = np.argsort(-mean)[:top_k]
        neg_order = np.argsort(mean)[:top_k]
        ranked: dict[str, list[dict[str, float | int | str]]] = {"positive": [], "negative": []}
        for label, order in [("positive", pos_order), ("negative", neg_order)]:
            for idx in order:
                token_text = tokenizer.decode(
                    [int(idx)],
                    skip_special_tokens=False,
                    clean_up_tokenization_spaces=False,
                )
                ranked[label].append(
                    {
                        "token_id": int(idx),
                        "token_text": token_text,
                        "mean_delta": float(mean[idx]),
                        "mean_abs_delta": float(mean_abs[idx]),
                        "positive_fraction": float(pos_frac[idx]),
                    }
                )
        return ranked

    summary = {
        "actual_token_prob_delta_target_minus_reverse_mean": float(np.mean([r["actual_token_prob_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0,
        "actual_token_logit_delta_target_minus_reverse_mean": float(np.mean([r["actual_token_logit_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0,
        "readout_class_prob_delta_target_minus_reverse_mean": float(np.mean([r["readout_class_prob_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0,
        "actual_token_prob_delta_target_minus_wrong_time_mean": float(np.mean([r["actual_token_prob_delta_target_minus_wrong_time"] for r in per_row])) if per_row else 0.0,
        "stable_hidden_dims": _rank_dims(hidden_sum, hidden_abs_sum, hidden_pos_counts, top_k=16),
        "top_logit_tokens": _rank_tokens(logit_sum, logit_abs_sum, logit_pos_counts, top_k=12),
    }
    return {
        "target_feature_indices": target_features,
        "n_rows_evaluated": len(per_row),
        "per_row": per_row,
        "summary": summary,
        "total_time_s": float(time.perf_counter() - t0),
    }


def _period_token_id(tokenizer) -> int:
    ids = tokenizer.encode(".", add_special_tokens=False)
    if len(ids) != 1:
        raise ValueError(f"Expected '.' to map to a single token, got ids={ids}")
    return int(ids[0])


def _decoder_direction_from_weights(
    sae,
    weights: np.ndarray,
) -> tuple[torch.Tensor, list[dict[str, float | int]]]:
    if hasattr(sae, "V_dec") and hasattr(sae, "W_dec"):
        v = sae.V_dec.detach().float().cpu()
        w = sae.W_dec.detach().float().cpu()
    elif hasattr(sae, "V") and hasattr(sae, "W"):
        v = sae.V.detach().float().cpu()
        w = sae.W.detach().float().cpu()
    else:
        raise ValueError("SAE does not expose rank-1 decoder factors")
    if v.ndim == 2:
        v = v[:, None, :]
    if w.ndim == 2:
        w = w[:, None, :]
    weight_t = torch.from_numpy(weights.astype(np.float32))
    direction = torch.einsum("i,irk,irv->kv", weight_t, v, w)
    fro = float(direction.norm().item())
    if fro > 0:
        direction = direction / fro
    top = np.argsort(-np.abs(weights))[:16]
    top_weights = [
        {"feature_idx": int(idx), "weight": float(weights[idx])}
        for idx in top
        if abs(float(weights[idx])) > 1e-6
    ]
    return direction, top_weights


def collect_exact_token_direction_dataset(
    model,
    corpus_batches: list[torch.Tensor],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    target_token_id: int,
    prompt_len: int = 512,
    negative_keep_prob: float = 0.03,
    max_negative_multiplier: int = 4,
    device: str = "cuda",
) -> dict[str, Any]:
    """Collect SAE activations for exact '.' vs sampled non-'.' next-token decisions."""
    rng = np.random.default_rng(0)
    xs: list[np.ndarray] = []
    ys: list[int] = []
    meta: list[dict[str, int]] = []
    positives = 0
    negatives = 0
    prompt_index = 0
    t0 = time.perf_counter()

    print(
        "[factor-trace-dot-direction] "
        f"collect_dataset batches={len(corpus_batches)} prompt_len={prompt_len} target_token={target_token_id}",
        flush=True,
    )

    for batch_idx, batch in enumerate(corpus_batches, start=1):
        token_batch = batch[:, :prompt_len].to(device)
        batch_size, seq_len = token_batch.shape
        if seq_len < 2:
            continue
        cache = None
        for pos in range(seq_len - 1):
            step_ids = token_batch[:, pos : pos + 1]
            if cache is None:
                out = model(input_ids=step_ids, use_cache=True)
            else:
                out = model(input_ids=step_ids, past_key_values=cache, use_cache=True)
            cache = out.past_key_values
            state_batch = cache.layers[layer_idx].recurrent_states[:, head_idx].float()
            acts = _encode_states(sae, sae_type, state_batch).detach().cpu().numpy().astype(np.float32)
            next_tokens = token_batch[:, pos + 1].detach().cpu().numpy()
            for row_idx in range(batch_size):
                label = int(next_tokens[row_idx] == target_token_id)
                keep = label == 1 or rng.random() < negative_keep_prob
                if not keep:
                    continue
                xs.append(acts[row_idx])
                ys.append(label)
                meta.append({"prompt_index": prompt_index + row_idx, "decision_pos": int(pos)})
                if label == 1:
                    positives += 1
                else:
                    negatives += 1
        prompt_index += batch_size
        if batch_idx % 4 == 0 or batch_idx == len(corpus_batches):
            print(
                "[factor-trace-dot-direction] "
                f"dataset_progress batch={batch_idx}/{len(corpus_batches)} "
                f"positives={positives} negatives={negatives} elapsed={time.perf_counter() - t0:.1f}s",
                flush=True,
            )

    if positives == 0:
        raise ValueError("No positive '.' examples collected for direction training")
    x = np.stack(xs, axis=0).astype(np.float32)
    y = np.array(ys, dtype=np.int64)
    meta_arr = np.array(meta, dtype=object)

    pos_idx = np.where(y == 1)[0]
    neg_idx = np.where(y == 0)[0]
    max_neg = max(len(pos_idx) * max_negative_multiplier, len(pos_idx))
    if len(neg_idx) > max_neg:
        neg_idx = rng.choice(neg_idx, size=max_neg, replace=False)
    selected = np.concatenate([pos_idx, neg_idx])
    rng.shuffle(selected)
    return {
        "x": x[selected],
        "y": y[selected],
        "meta": meta_arr[selected].tolist(),
        "n_positive": int(len(pos_idx)),
        "n_negative_sampled": int(len(neg_idx)),
        "total_time_s": float(time.perf_counter() - t0),
    }


def run_dot_direction_use_site(
    model,
    tokenizer,
    corpus_batches: list[torch.Tensor],
    corpus_sequences: list[torch.Tensor],
    benchmark_result: dict[str, Any],
    use_site_result: dict[str, Any],
    layer_idx: int,
    head_idx: int,
    sae,
    sae_type: str,
    prompt_len: int = 512,
    direction_prompt_count: int = 256,
    n_rows: int = 5,
    alphas: Sequence[float] = (0.5, 1.0, 2.0, 4.0),
    min_accuracy: float = 0.65,
    device: str = "cuda",
) -> dict[str, Any]:
    """Train a '.' logistic direction in H4 feature space and test it on exact '.' use-sites."""
    from sklearn.linear_model import LogisticRegression  # type: ignore[import-untyped]
    from sklearn.model_selection import train_test_split  # type: ignore[import-untyped]

    target_token_id = _period_token_id(tokenizer)
    dataset = collect_exact_token_direction_dataset(
        model=model,
        corpus_batches=corpus_batches,
        layer_idx=layer_idx,
        head_idx=head_idx,
        sae=sae,
        sae_type=sae_type,
        target_token_id=target_token_id,
        prompt_len=prompt_len,
        device=device,
    )
    x = dataset["x"]
    y = dataset["y"]
    x_train, x_test, y_train, y_test = train_test_split(
        x,
        y,
        test_size=0.2,
        random_state=0,
        stratify=y,
    )
    clf = LogisticRegression(
        max_iter=2000,
        class_weight="balanced",
        solver="liblinear",
        random_state=0,
    )
    clf.fit(x_train, y_train)
    test_accuracy = float(clf.score(x_test, y_test))
    weights = clf.coef_[0].astype(np.float32)
    direction, top_weights = _decoder_direction_from_weights(sae, weights)
    print(
        "[factor-trace-dot-direction] "
        f"logistic accuracy={test_accuracy:.3f} positives={dataset['n_positive']} "
        f"negatives={dataset['n_negative_sampled']}",
        flush=True,
    )
    print(
        "[factor-trace-dot-direction] "
        f"top_weights={top_weights[:8]}",
        flush=True,
    )
    result: dict[str, Any] = {
        "target_feature_indices": [62, 105],
        "target_token_id": target_token_id,
        "target_token_text": ".",
        "logistic": {
            "test_accuracy": test_accuracy,
            "n_positive": dataset["n_positive"],
            "n_negative_sampled": dataset["n_negative_sampled"],
            "top_weights": top_weights,
        },
    }
    if test_accuracy < min_accuracy:
        result["status"] = "stopped_low_accuracy"
        result["rows"] = []
        result["summary"] = {"success_rows": 0, "success_gate_pass": False}
        return result

    benchmark_map = {
        (int(r["donor_prompt_index"]), int(r["recipient_prompt_index"]), int(r["decision_pos"])): r
        for r in benchmark_result.get("accepted_rows", [])
    }
    candidate_rows = []
    for row in use_site_result.get("per_row", []):
        if row.get("readout_class") != "sentence_reset" or row.get("readout_text") != ".":
            continue
        p = float(row["baseline_actual_token_prob"])
        if p < 0.3 or p > 0.8:
            continue
        delta = float(row["actual_token_prob_shift_target"] - row["actual_token_prob_shift_reverse"])
        coeff = float(row["signed_direction_coeff_shift_target"])
        score = (
            int(delta > 0),
            delta,
            p * (1.0 - p),
            coeff,
        )
        key = (int(row["donor_prompt_index"]), int(row["recipient_prompt_index"]), int(row["decision_pos"]))
        merged = dict(row)
        merged["benchmark_row"] = benchmark_map.get(key, {})
        merged["_score"] = score
        candidate_rows.append(merged)
    candidate_rows.sort(key=lambda r: r["_score"], reverse=True)
    selected_rows = candidate_rows[:n_rows]
    result["selected_rows"] = [
        {
            "donor_prompt_index": int(r["donor_prompt_index"]),
            "recipient_prompt_index": int(r["recipient_prompt_index"]),
            "decision_pos": int(r["decision_pos"]),
            "baseline_actual_token_prob": float(r["baseline_actual_token_prob"]),
            "prior_target_minus_reverse_prob": float(r["actual_token_prob_shift_target"] - r["actual_token_prob_shift_reverse"]),
            "prior_signed_direction_coeff_shift": float(r["signed_direction_coeff_shift_target"]),
        }
        for r in selected_rows
    ]
    print(
        "[factor-trace-dot-direction] "
        f"selected_rows={len(selected_rows)} exact_dot_only=True",
        flush=True,
    )

    sequence_map = {idx: seq.to(device) for idx, seq in enumerate(corpus_sequences)}
    tracked_features = [62, 105, 121, 9, 137]
    per_row = []
    direction = direction.to(device=device, dtype=torch.float32)
    t0 = time.perf_counter()

    for row_idx, row in enumerate(selected_rows, start=1):
        recipient_idx = int(row["recipient_prompt_index"])
        seq_full = sequence_map[recipient_idx]
        decision_pos = int(row["decision_pos"])
        intervention_pos = int(row.get("intervention_pos", max(0, decision_pos - 1)))
        wrong_time_pos = max(0, intervention_pos - 10)
        seq_end = min(int(seq_full.shape[1]), decision_pos + 3)
        seq = seq_full[:, :seq_end]
        baseline = teacher_forced_feature_rollout(
            model=model,
            input_ids=seq,
            layer_idx=layer_idx,
            head_idx=head_idx,
            sae=sae,
            sae_type=sae_type,
            tracked_features=tracked_features,
        )
        baseline_logits = baseline["logits_by_pos"][decision_pos]
        baseline_prob = _token_prob_from_logits(baseline_logits, target_token_id)
        token_ids = seq_full[0].detach().cpu().tolist()
        ctx_start = max(0, decision_pos - 24)
        ctx_end = min(len(token_ids), decision_pos + 2)
        context_text = tokenizer.decode(
            token_ids[ctx_start:ctx_end],
            skip_special_tokens=False,
            clean_up_tokenization_spaces=False,
        )
        alpha_rows = []
        for alpha in alphas:
            scaled = direction * float(alpha)
            inject = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_state_delta=scaled,
                intervention_pos=intervention_pos,
            )
            reverse = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_state_delta=-scaled,
                intervention_pos=intervention_pos,
            )
            wrong_time = teacher_forced_feature_rollout(
                model=model,
                input_ids=seq,
                layer_idx=layer_idx,
                head_idx=head_idx,
                sae=sae,
                sae_type=sae_type,
                tracked_features=tracked_features,
                intervention_state_delta=scaled,
                intervention_pos=wrong_time_pos,
            )
            inject_logits = inject["logits_by_pos"][decision_pos]
            reverse_logits = reverse["logits_by_pos"][decision_pos]
            wrong_time_logits = wrong_time["logits_by_pos"][decision_pos]
            inject_prob = _token_prob_from_logits(inject_logits, target_token_id)
            reverse_prob = _token_prob_from_logits(reverse_logits, target_token_id)
            wrong_time_prob = _token_prob_from_logits(wrong_time_logits, target_token_id)
            alpha_rows.append(
                {
                    "alpha": float(alpha),
                    "baseline_prob": baseline_prob,
                    "inject_prob": inject_prob,
                    "reverse_prob": reverse_prob,
                    "wrong_time_prob": wrong_time_prob,
                    "inject_minus_reverse": float(inject_prob - reverse_prob),
                    "inject_minus_baseline": float(inject_prob - baseline_prob),
                    "reverse_minus_baseline": float(reverse_prob - baseline_prob),
                    "wrong_time_minus_baseline": float(wrong_time_prob - baseline_prob),
                }
            )
        best = max(alpha_rows, key=lambda item: item["inject_minus_reverse"])
        success = bool(
            best["inject_prob"] > best["baseline_prob"]
            and best["reverse_prob"] < best["baseline_prob"]
            and best["inject_minus_reverse"] > 0.05
        )
        per_row.append(
            {
                "donor_prompt_index": int(row["donor_prompt_index"]),
                "recipient_prompt_index": recipient_idx,
                "decision_pos": decision_pos,
                "intervention_pos": intervention_pos,
                "wrong_time_pos": wrong_time_pos,
                "context_text": context_text,
                "baseline_prob": baseline_prob,
                "prior_target_minus_reverse_prob": float(row["actual_token_prob_shift_target"] - row["actual_token_prob_shift_reverse"]),
                "prior_signed_direction_coeff_shift": float(row["signed_direction_coeff_shift_target"]),
                "alphas": alpha_rows,
                "best_alpha": float(best["alpha"]),
                "best_inject_prob": float(best["inject_prob"]),
                "best_reverse_prob": float(best["reverse_prob"]),
                "best_wrong_time_prob": float(best["wrong_time_prob"]),
                "best_inject_minus_reverse": float(best["inject_minus_reverse"]),
                "success": success,
            }
        )
        print(
            "[factor-trace-dot-direction] "
            f"row {row_idx}/{len(selected_rows)} donor={row['donor_prompt_index']}→recipient={recipient_idx}@{decision_pos} "
            f"base={baseline_prob:.3f} best_alpha={best['alpha']:.1f} "
            f"inject={best['inject_prob']:.3f} reverse={best['reverse_prob']:.3f} "
            f"wrong_time={best['wrong_time_prob']:.3f} delta={best['inject_minus_reverse']:+.4f} "
            f"success={success}",
            flush=True,
        )

    success_rows = sum(int(r["success"]) for r in per_row)
    result["status"] = "completed"
    result["rows"] = per_row
    result["summary"] = {
        "success_rows": int(success_rows),
        "n_rows": int(len(per_row)),
        "success_gate_pass": bool(success_rows >= 3),
        "max_inject_minus_reverse": float(max((r["best_inject_minus_reverse"] for r in per_row), default=0.0)),
        "mean_best_inject_minus_reverse": float(np.mean([r["best_inject_minus_reverse"] for r in per_row])) if per_row else 0.0,
        "total_time_s": float(time.perf_counter() - t0),
    }
    return result