run_residual_patching_batch.py

#!/usr/bin/env python3
"""Run residual-stream patching for one or more snippet pairs and save plots.

The sweep starts at the first prompt-token mismatch. For a causal decoder,
positions before that mismatch have identical prefixes, so patching them is
normally wasted work.
"""

from __future__ import annotations

import argparse
import ast
import csv
import difflib
import gc
import json
import math
import re
import shutil
import subprocess
import sys
import tempfile
import textwrap
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from scripts.methodology_core import (
    PAIR_FAMILIES,
    PROMPT_ROLES,
    StimulusPair,
    classify_behavior,
    clean_final_output_text,
    iter_pair_dirs,
    load_pair_dirs_from_args,
    load_stimulus_pair,
)

pd = None
plt = None
torch = None
tqdm = None

SOURCE_RECOVERY_CMAP = "source_recovery"
DESTINATION_RECOVERY_CMAP = "destination_recovery"
SOURCE_LOGIT_MARGIN_CMAP = "source_logit_margin"
DESTINATION_LOGIT_MARGIN_CMAP = "destination_logit_margin"
SIGNED_SOURCE_POSITIVE_CMAP = "signed_source_positive"
SIGNED_DESTINATION_POSITIVE_CMAP = "signed_destination_positive"
CLEAN_RECOVERY_CMAP = SOURCE_RECOVERY_CMAP
CORRUPTED_RECOVERY_CMAP = DESTINATION_RECOVERY_CMAP
SEMANTIC_ROLE_CHOICES = (
    "comment",
    "docstring",
    "function_name",
    "variable_name",
    "changed_cue_region",
    "changed_implementation_region",
    "input_example_object",
    "output_readout",
    "code_body",
    "unknown",
)



DEFAULT_UNIT_TEST_PROMPT_TEMPLATE = (
    "You are given a Python snippet. Generate one concise unit test (pytest style) with exactly "
    "3 test cases for the main function defined in it. Each test case must be a single assertion "
    "of the form `assert <expression> == <literal expected value>`. Do not use bare boolean "
    "assertions such as `assert f(x)` or `assert not f(x)`; write `assert f(x) == True` or "
    "`assert f(x) == False` instead. Return exactly one fenced Python code block. End the code "
    "block with closing triple backticks.\n\n"
    "```python\n{code}\n```\n\n"
    "Unit test:\n"
)

UNIT_TEST_ASSERT_LIMIT = 3

MODEL_NAME_BY_SLUG = {
    "codellama_7b_python_hf": "codellama/CodeLlama-7b-Python-hf",
    "qwen2_5_7b_instruct": "Qwen/Qwen2.5-7B-Instruct",
    "llama_3_1_8b_instruct": "meta-llama/Llama-3.1-8B-Instruct",
    "mistral_7b_instruct_v0_1": "mistralai/Mistral-7B-Instruct-v0.1",
}
MODEL_SLUG_BY_NAME = {name: slug for slug, name in MODEL_NAME_BY_SLUG.items()}

RQ_STAGES = ("rq1_1", "rq1_2", "rq2_1", "rq2_2")
RQ_STAGE_DIRS = {
    "rq1_1": "rq1_1_final_output_generation",
    "rq1_2": "rq1_2_final_output_residual_patching",
    "rq2_1": "rq2_1_unit_test_generation",
    "rq2_2": "rq2_2_unit_test_residual_patching",
}


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def load_runtime_dependencies() -> None:
    global pd, plt, torch, tqdm

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as _plt
    import pandas as _pd
    import torch as _torch
    from tqdm.auto import tqdm as _tqdm

    _plt.rcParams.update({"pdf.fonttype": 42, "ps.fonttype": 42})

    pd = _pd
    plt = _plt
    torch = _torch
    tqdm = _tqdm


def load_hooked_transformer_class():
    try:
        from transformer_lens import HookedTransformer
    except Exception as e:  # pragma: no cover - import guard for user environment
        raise ImportError(
            "transformer_lens is required. Install it with: pip install transformer_lens"
        ) from e
    return HookedTransformer


def cleanup_runtime_memory() -> None:
    gc.collect()
    if torch is None:
        return
    if hasattr(torch, "cuda") and torch.cuda.is_available():
        torch.cuda.empty_cache()
        if hasattr(torch.cuda, "ipc_collect"):
            torch.cuda.ipc_collect()
    if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
        try:
            torch.mps.empty_cache()
        except RuntimeError:
            pass


def pair_dir_display_name(pair_dir: Path, stimulus_root: Path) -> str:
    try:
        return pair_dir.relative_to(stimulus_root).as_posix()
    except ValueError:
        return pair_dir.name


def stimulus_pair_to_runner_pair(pair: StimulusPair, pair_dir: Path) -> Dict[str, Any]:
    return {
        "pair_dir": pair_dir,
        "pair_id": pair.pair_id,
        "pair_family": pair.pair_family,
        "prompts": pair.prompts,
        "patch_directions": pair.patch_directions,
        "meta": {
            "pair_id": pair.pair_id,
            "pair_family": pair.pair_family,
            "prompts": {
                role: {
                    "file": prompt.file,
                    "prompt_role": prompt.role,
                    "execution_output": prompt.execution_output,
                    "cue_output": prompt.cue_output,
                }
                for role, prompt in pair.prompts.items()
            },
        },
        "prompt_codes": {role: prompt.code for role, prompt in pair.prompts.items()},
    }


def resolve_pair_dir(stimulus_root: Path, pair_name: str) -> Path:
    matches = load_pair_dirs_from_args(
        stimulus_root=stimulus_root,
        names=[pair_name],
        use_all=False,
    )
    if not matches:
        raise FileNotFoundError(f"Stimulus pair folder not found: {pair_name}")
    return matches[0]


def load_pair(snippet_dir: Path, pair_name: str) -> Dict:
    pair_dir = resolve_pair_dir(snippet_dir, pair_name)
    return stimulus_pair_to_runner_pair(load_stimulus_pair(pair_dir), pair_dir)


def list_pairs(snippet_dir: Path, pair_family: Optional[str] = None) -> List[str]:
    names = []
    for pair_dir in iter_pair_dirs(snippet_dir):
        pair = load_stimulus_pair(pair_dir)
        if pair_family and pair.pair_family != pair_family:
            continue
        names.append(pair_dir_display_name(pair_dir, snippet_dir))
    return sorted(names)


def build_prompt(code: str, template: str) -> str:
    return template.format(code=code.rstrip())


def fenced_code_match(text: str) -> Optional[re.Match[str]]:
    match = re.search(r"```[^\n`]*\n(.*?)\n```", text, flags=re.DOTALL)
    if match:
        return match
    return re.search(r"```(.*?)```", text, flags=re.DOTALL)


def split_fenced_prompt(prompt: str) -> Dict[str, str]:
    match = fenced_code_match(prompt)
    if not match:
        return {
            "prompt_text": prompt.strip(),
            "snippet_text": "",
            "prefix_text": prompt,
            "suffix_text": "",
        }
    prefix = prompt[: match.start()]
    snippet = match.group(1).strip("\n")
    suffix = prompt[match.end() :]
    return {
        "prompt_text": prefix.strip(),
        "snippet_text": snippet,
        "prefix_text": prefix,
        "suffix_text": suffix,
    }


def prompt_output_label(prompt_parts: Dict[str, str], fallback: str = "Generated Answer") -> str:
    label = str(prompt_parts.get("suffix_text") or "").strip()
    return label or fallback


def fenced_code_body_span(prompt: str) -> Optional[Tuple[int, int]]:
    match = fenced_code_match(prompt)
    if not match:
        return None
    return int(match.start(1)), int(match.end(1))


def fenced_code_and_output_cue_span(prompt: str) -> Optional[Tuple[int, int]]:
    match = fenced_code_match(prompt)
    if not match:
        return None
    return int(match.start(1)), len(prompt)


def display_model_token(token: str) -> str:
    return str(token).replace("\u0120", " ").replace("\u2581", " ").replace("Ċ", "\n")


def normalize_token_match_text(text: str) -> str:
    return display_model_token(text).replace("\r\n", "\n").replace("\r", "\n").strip()


def token_span_in_text(source: str, start_index: int, raw_token: str) -> Tuple[int, int]:
    token = display_model_token(raw_token)
    if not token:
        return start_index, start_index
    if source.startswith(token, start_index):
        return start_index, start_index + len(token)
    if token.startswith(" ") and source.startswith(token[1:], start_index):
        return start_index, start_index + len(token[1:])
    if token and all(ch.isspace() for ch in token):
        end = start_index
        limit = min(len(source), start_index + max(1, len(token)))
        while end < limit and source[end].isspace():
            end += 1
        return start_index, end

    normalized_token = normalize_token_match_text(token)
    if not normalized_token:
        return start_index, start_index
    search_limit = min(len(source), start_index + 96)
    best_start = source.find(token, start_index, search_limit)
    if best_start >= 0:
        return best_start, best_start + len(token)
    stripped = token.strip()
    if stripped:
        best_start = source.find(stripped, start_index, search_limit)
        if best_start >= 0:
            return best_start, best_start + len(stripped)
    return start_index, start_index


def token_positions_overlapping_span(
    source: str,
    token_strs: Sequence[str],
    char_span: Optional[Tuple[int, int]],
) -> List[int]:
    if char_span is None:
        return []
    span_start, span_end = int(char_span[0]), int(char_span[1])
    positions: List[int] = []
    cursor = 0
    for pos, raw_token in enumerate(token_strs):
        token_start, token_end = token_span_in_text(source, cursor, raw_token)
        if token_end > span_start and token_start < span_end:
            positions.append(int(pos))
        cursor = max(cursor, token_end)
    return positions


def filter_patch_units_to_position_sets(
    patch_units: Sequence["PatchUnit"],
    clean_positions: Sequence[int],
    corrupted_positions: Sequence[int],
) -> List["PatchUnit"]:
    clean_set = {int(pos) for pos in clean_positions}
    corrupted_set = {int(pos) for pos in corrupted_positions}
    filtered: List[PatchUnit] = []
    for unit in patch_units:
        if not unit.clean_positions or not unit.corrupted_positions:
            continue
        if all(int(pos) in clean_set for pos in unit.clean_positions) and all(
            int(pos) in corrupted_set for pos in unit.corrupted_positions
        ):
            filtered.append(unit)
    return filtered


def extract_completion(full_text: str, prompt: str) -> str:
    if full_text.startswith(prompt):
        return full_text[len(prompt) :].strip()
    return full_text.strip()


def truncate_after_first_fenced_block(text: str) -> str:
    opener = text.find("```")
    if opener < 0:
        return text

    closer = text.find("```", opener + 3)
    if closer < 0:
        return text
    return text[: closer + 3]


def contains_complete_fenced_block(text: str) -> bool:
    opener = text.find("```")
    return opener >= 0 and text.find("```", opener + 3) >= 0


def extract_first_fenced_block(text: str) -> str:
    match = re.search(r"```[^\n`]*\n(.*?)```", text, flags=re.DOTALL)
    if not match:
        match = re.search(r"```(.*?)```", text, flags=re.DOTALL)
    if not match:
        return text.strip()
    return match.group(1).strip()


def extract_first_fenced_block_or_empty(text: str) -> str:
    match = re.search(r"```[^\n`]*\n(.*?)```", text, flags=re.DOTALL)
    if not match:
        match = re.search(r"```(.*?)```", text, flags=re.DOTALL)
    if not match:
        return ""
    return match.group(1).strip()


def extract_first_fenced_block_or_text(text: str) -> str:
    extracted = extract_first_fenced_block_or_empty(text)
    return extracted if extracted else text.strip()


def split_first_fenced_block_prefix_body(text: str) -> Optional[Tuple[str, str, bool]]:
    opener = text.find("```")
    if opener < 0:
        return None
    header_end = text.find("\n", opener + 3)
    if header_end < 0:
        return None
    body_start = header_end + 1
    closer = text.find("```", body_start)
    if closer < 0:
        return text[:body_start], text[body_start:], False
    return text[:body_start], text[body_start:closer], True


def parsed_assertion_spans(source: str) -> Optional[List[Tuple[int, int]]]:
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return None

    spans: List[Tuple[int, int]] = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Assert):
            continue
        span = node_char_span(source, node)
        if span is not None:
            spans.append((int(span[0]), int(span[1])))
    return sorted(spans)


def parsed_supported_equality_assertion_spans(source: str) -> Optional[List[Tuple[int, int]]]:
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return None

    spans: List[Tuple[int, int]] = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Assert):
            continue
        test = node.test
        if not isinstance(test, ast.Compare):
            continue
        if len(test.ops) != 1 or not isinstance(test.ops[0], ast.Eq):
            continue
        if len(test.comparators) != 1:
            continue
        span = node_char_span(source, node)
        if span is not None:
            spans.append((int(span[0]), int(span[1])))
    return sorted(spans)


def cap_unit_test_source_to_assert_limit(
    test_source: str,
    max_asserts: int = UNIT_TEST_ASSERT_LIMIT,
) -> str:
    source = str(test_source or "").strip()
    if max_asserts <= 0:
        return ""

    supported_spans = parsed_supported_equality_assertion_spans(source)
    if supported_spans is not None and len(supported_spans) >= max_asserts:
        return source[: supported_spans[max_asserts - 1][1]].rstrip()

    spans = parsed_assertion_spans(source)
    if spans is None or len(spans) <= max_asserts:
        return source
    if supported_spans:
        return source
    return source[: spans[max_asserts - 1][1]].rstrip()


def unit_test_assert_limit_reached(
    completion: str,
    max_asserts: int = UNIT_TEST_ASSERT_LIMIT,
) -> bool:
    parts = split_first_fenced_block_prefix_body(completion)
    candidate = parts[1] if parts is not None else completion
    spans = parsed_supported_equality_assertion_spans(str(candidate or "").strip())
    return spans is not None and len(spans) >= max_asserts


def cap_unit_test_completion_to_assert_limit(
    completion: str,
    max_asserts: int = UNIT_TEST_ASSERT_LIMIT,
) -> str:
    completion = str(completion or "")
    parts = split_first_fenced_block_prefix_body(completion)
    if parts is None:
        return cap_unit_test_source_to_assert_limit(completion, max_asserts=max_asserts).strip()

    prefix, body, has_closer = parts
    capped_body = cap_unit_test_source_to_assert_limit(body, max_asserts=max_asserts)
    capped_spans = parsed_supported_equality_assertion_spans(capped_body)
    reached_limit = capped_spans is not None and len(capped_spans) >= max_asserts
    if reached_limit:
        return f"{prefix}{capped_body}\n```".strip()
    if has_closer:
        return truncate_after_first_fenced_block(completion).strip()
    return completion.strip()


def extract_first_json_object(text: str) -> str:
    stripped = extract_first_fenced_block_or_text(text).strip()
    if stripped.startswith("{") and stripped.endswith("}"):
        return stripped
    start = stripped.find("{")
    if start < 0:
        return stripped
    depth = 0
    in_string = False
    escape = False
    for index, ch in enumerate(stripped[start:], start=start):
        if in_string:
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == '"':
                in_string = False
            continue
        if ch == '"':
            in_string = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return stripped[start : index + 1]
    return stripped[start:]


def compact_token_id_label(token_ids: Optional[Sequence[int]]) -> str:
    ids = [int(token_id) for token_id in (token_ids or [])]
    if not ids:
        return ""
    if len(ids) == 1:
        return str(ids[0])
    shown = ",".join(str(token_id) for token_id in ids[:3])
    return shown if len(ids) <= 3 else shown + ",..."


def format_whitespace_for_axis(text: str) -> str:
    if text == " ":
        return "<space>"
    if text and set(text) == {" "}:
        return f"<{len(text)} spaces>"
    if text == "\n":
        return "<newline>"
    if text and set(text) == {"\n"}:
        return f"<{len(text)} newlines>"
    if text == "\t":
        return "<tab>"
    if text and set(text) == {"\t"}:
        return f"<{len(text)} tabs>"
    escaped = (
        text.replace(" ", "·")
        .replace("\n", "\\n")
        .replace("\r", "\\r")
        .replace("\t", "\\t")
    )
    return f"<ws:{escaped}>"


def format_text_for_axis(
    text: str,
    max_chars: int = 18,
    token_ids: Optional[Sequence[int]] = None,
) -> str:
    raw_text = str(text).replace("\u0120", " ").replace("\u2581", " ")
    if raw_text == "":
        label = "<empty>"
    elif raw_text.isspace():
        label = format_whitespace_for_axis(raw_text)
    else:
        label = raw_text.replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t")
        if raw_text.strip() != raw_text:
            label = repr(label)
    if len(label) > max_chars:
        label = label[: max_chars - 3] + "..."
    return label


def format_token_for_axis(
    token: str,
    max_chars: int = 18,
    token_ids: Optional[Sequence[int]] = None,
) -> str:
    return format_text_for_axis(token, max_chars=max_chars, token_ids=token_ids)


def format_span_for_axis(
    text: str,
    max_chars: int = 18,
    token_ids: Optional[Sequence[int]] = None,
) -> str:
    return format_text_for_axis(text, max_chars=max_chars, token_ids=token_ids)


def recovery_cmap_for_flow_direction(flow_direction: str) -> str:
    if str(flow_direction).startswith(("aligned_cue_to", "aligned_implementation_to")):
        return SOURCE_RECOVERY_CMAP
    return DESTINATION_RECOVERY_CMAP


def get_variant_code(pair: Dict, variant: str) -> str:
    prompt_codes = pair.get("prompt_codes")
    if isinstance(prompt_codes, dict) and variant in prompt_codes:
        return str(prompt_codes[variant])
    raise ValueError(f"Unknown prompt_role: {variant!r}.")


def execute_snippet_and_get_final_printed_output(code: str) -> Optional[str]:
    import io
    from contextlib import redirect_stdout

    stdout_buffer = io.StringIO()
    exec_scope = {"__builtins__": __builtins__}
    try:
        with redirect_stdout(stdout_buffer):
            exec(code, exec_scope, exec_scope)
    except Exception:
        return None

    lines = stdout_buffer.getvalue().splitlines()
    return lines[-1] if lines else None


@dataclass
class AnswerDecisionSpec:
    answer_resolution_schema: str
    clean_positive_answer: str
    clean_negative_answer: Optional[str]
    corrupted_positive_answer: Optional[str]
    corrupted_negative_answer: Optional[str]
    forward_positive_answer: str
    forward_negative_answer: Optional[str]
    reverse_positive_answer: Optional[str]
    reverse_negative_answer: Optional[str]
    clean_expected_output: Optional[str]
    corrupted_expected_output: Optional[str]


def optional_str(value: Any) -> Optional[str]:
    return None if value is None else str(value)


def variant_metadata_value(meta: Dict[str, Any], mapping_key: str, variant: str) -> Optional[str]:
    mapping = meta.get(mapping_key)
    if isinstance(mapping, dict) and variant in mapping:
        return optional_str(mapping[variant])
    return None


def flat_variant_metadata_value(meta: Dict[str, Any], variant: str, suffix: str) -> Optional[str]:
    return optional_str(meta.get(f"{variant}_{suffix}"))


def resolve_variant_answer(meta: Dict[str, Any], variant: str) -> Optional[str]:
    return (
        variant_metadata_value(meta, "answers", variant)
        or variant_metadata_value(meta, "variant_answers", variant)
        or flat_variant_metadata_value(meta, variant, "answer")
    )


def resolve_variant_wrong_answer(meta: Dict[str, Any], variant: str) -> Optional[str]:
    return (
        variant_metadata_value(meta, "wrong_answers", variant)
        or variant_metadata_value(meta, "variant_wrong_answers", variant)
        or flat_variant_metadata_value(meta, variant, "wrong_answer")
    )


def has_variant_answer_metadata(meta: Dict[str, Any]) -> bool:
    return any(
        key in meta
        for key in (
            "answers",
            "variant_answers",
            "wrong_answers",
            "variant_wrong_answers",
            "control_answer",
            "control_wrong_answer",
            "treatment_answer",
            "treatment_wrong_answer",
        )
    )


def explicit_direction_decision(
    meta: Dict[str, Any],
    source_variant: str,
    destination_variant: str,
    role_key: str,
) -> Tuple[Optional[str], Optional[str]]:
    decisions = meta.get("answer_decisions")
    if not isinstance(decisions, dict):
        return None, None

    for key in (
        f"{source_variant}_into_{destination_variant}",
        role_key,
    ):
        decision = decisions.get(key)
        if not isinstance(decision, dict):
            continue
        positive = optional_str(
            decision.get("positive_answer", decision.get("positive", decision.get("answer")))
        )
        negative = optional_str(
            decision.get(
                "negative_answer",
                decision.get("negative", decision.get("wrong_answer")),
            )
        )
        return positive, negative
    return None, None


def legacy_answer_decision_spec(
    *,
    target: str,
    wrong: Optional[str],
    clean_expected_output: Optional[str],
    corrupted_expected_output: Optional[str],
    answer_resolution_schema: str,
) -> AnswerDecisionSpec:
    return AnswerDecisionSpec(
        answer_resolution_schema=answer_resolution_schema,
        clean_positive_answer=target,
        clean_negative_answer=wrong,
        corrupted_positive_answer=wrong,
        corrupted_negative_answer=target if wrong is not None else None,
        forward_positive_answer=target,
        forward_negative_answer=wrong,
        reverse_positive_answer=wrong,
        reverse_negative_answer=target if wrong is not None else None,
        clean_expected_output=clean_expected_output,
        corrupted_expected_output=corrupted_expected_output,
    )


def resolve_answer_decision_spec(
    pair: Dict,
    clean_variant: str,
    corrupted_variant: str,
    clean_code: str,
    corrupted_code: str,
    use_gold_answer: bool = True,
    manual_target_answer: Optional[str] = None,
    manual_wrong_answer: Optional[str] = None,
    force_target_to_clean_output: bool = True,
    auto_wrong_from_corrupted_output: bool = True,
) -> AnswerDecisionSpec:
    meta = pair["meta"]
    clean_output = execute_snippet_and_get_final_printed_output(clean_code)
    corrupted_output = execute_snippet_and_get_final_printed_output(corrupted_code)

    prompts = pair.get("prompts")
    if isinstance(prompts, dict) and clean_variant in prompts and corrupted_variant in prompts:
        source_prompt = prompts[clean_variant]
        destination_prompt = prompts[corrupted_variant]
        pair_family = str(pair.get("pair_family") or meta.get("pair_family") or "")
        if pair_family == "cue_varied":
            source_value = str(source_prompt.cue_output)
            destination_value = str(destination_prompt.cue_output)
        elif pair_family == "implementation_varied":
            source_value = str(source_prompt.execution_output)
            destination_value = str(destination_prompt.execution_output)
        else:
            raise ValueError(f"Unsupported pair_family: {pair_family!r}")
        if clean_output is not None and clean_output != source_prompt.execution_output:
            raise ValueError("runtime_metadata_mismatch")
        if corrupted_output is not None and corrupted_output != destination_prompt.execution_output:
            raise ValueError("runtime_metadata_mismatch")
        return AnswerDecisionSpec(
            answer_resolution_schema="methodology_schema",
            clean_positive_answer=source_value,
            clean_negative_answer=destination_value,
            corrupted_positive_answer=destination_value,
            corrupted_negative_answer=source_value,
            forward_positive_answer=source_value,
            forward_negative_answer=destination_value,
            reverse_positive_answer=destination_value,
            reverse_negative_answer=source_value,
            clean_expected_output=str(source_prompt.execution_output),
            corrupted_expected_output=str(destination_prompt.execution_output),
        )

    if not use_gold_answer or manual_target_answer is not None or manual_wrong_answer is not None:
        target = meta.get("answer") if use_gold_answer and manual_target_answer is None else manual_target_answer
        if target is None:
            raise ValueError("No target answer resolved. Set --manual-target-answer or use metadata.")
        target = str(target)

        wrong = meta.get("wrong_answer") if manual_wrong_answer is None else manual_wrong_answer
        wrong = optional_str(wrong)

        if (
            force_target_to_clean_output
            and use_gold_answer
            and manual_target_answer is None
            and clean_output is not None
            and clean_output != target
        ):
            target = clean_output

        if (
            auto_wrong_from_corrupted_output
            and manual_wrong_answer is None
            and corrupted_output is not None
            and corrupted_output != target
        ):
            wrong = corrupted_output

        return legacy_answer_decision_spec(
            target=target,
            wrong=wrong,
            clean_expected_output=clean_output,
            corrupted_expected_output=corrupted_output,
            answer_resolution_schema="manual_or_legacy",
        )

    forward_positive, forward_negative = explicit_direction_decision(
        meta,
        source_variant=clean_variant,
        destination_variant=corrupted_variant,
        role_key="clean_into_corrupted",
    )
    reverse_positive, reverse_negative = explicit_direction_decision(
        meta,
        source_variant=corrupted_variant,
        destination_variant=clean_variant,
        role_key="corrupted_into_clean",
    )
    if forward_positive is not None or reverse_positive is not None:
        if forward_positive is None:
            raise ValueError("answer_decisions is missing the forward positive answer.")
        return AnswerDecisionSpec(
            answer_resolution_schema="answer_decisions",
            clean_positive_answer=forward_positive,
            clean_negative_answer=forward_negative,
            corrupted_positive_answer=reverse_positive,
            corrupted_negative_answer=reverse_negative,
            forward_positive_answer=forward_positive,
            forward_negative_answer=forward_negative,
            reverse_positive_answer=reverse_positive,
            reverse_negative_answer=reverse_negative,
            clean_expected_output=clean_output,
            corrupted_expected_output=corrupted_output,
        )

    if has_variant_answer_metadata(meta):
        clean_positive = resolve_variant_answer(meta, clean_variant)
        corrupted_positive = resolve_variant_answer(meta, corrupted_variant)

        if force_target_to_clean_output and clean_output is not None:
            clean_positive = clean_output
        if force_target_to_clean_output and corrupted_output is not None:
            corrupted_positive = corrupted_output

        if clean_positive is None:
            clean_positive = optional_str(meta.get("answer"))
        if corrupted_positive is None:
            corrupted_positive = corrupted_output

        if clean_positive is None:
            raise ValueError(f"No positive answer resolved for {clean_variant!r}.")
        if corrupted_positive is None:
            raise ValueError(f"No positive answer resolved for {corrupted_variant!r}.")

        clean_negative = resolve_variant_wrong_answer(meta, clean_variant)
        corrupted_negative = resolve_variant_wrong_answer(meta, corrupted_variant)
        if clean_negative is None and corrupted_positive != clean_positive:
            clean_negative = corrupted_positive
        if corrupted_negative is None and clean_positive != corrupted_positive:
            corrupted_negative = clean_positive

        return AnswerDecisionSpec(
            answer_resolution_schema="variant_answers",
            clean_positive_answer=str(clean_positive),
            clean_negative_answer=clean_negative,
            corrupted_positive_answer=str(corrupted_positive),
            corrupted_negative_answer=corrupted_negative,
            forward_positive_answer=str(clean_positive),
            forward_negative_answer=clean_negative,
            reverse_positive_answer=str(corrupted_positive),
            reverse_negative_answer=corrupted_negative,
            clean_expected_output=clean_output,
            corrupted_expected_output=corrupted_output,
        )

    target = meta.get("answer")
    if target is None:
        raise ValueError("No target answer resolved. Set --manual-target-answer or use metadata.")
    target = str(target)

    wrong = optional_str(meta.get("wrong_answer"))
    if (
        force_target_to_clean_output
        and clean_output is not None
        and clean_output != target
    ):
        target = clean_output

    if (
        auto_wrong_from_corrupted_output
        and corrupted_output is not None
        and corrupted_output != target
    ):
        wrong = corrupted_output

    return legacy_answer_decision_spec(
        target=target,
        wrong=wrong,
        clean_expected_output=clean_output,
        corrupted_expected_output=corrupted_output,
        answer_resolution_schema="legacy_answer_wrong_answer",
    )


def resolve_target_and_wrong_answers(
    pair: Dict,
    clean_code: str,
    corrupted_code: str,
    use_gold_answer: bool = True,
    manual_target_answer: Optional[str] = None,
    manual_wrong_answer: Optional[str] = None,
    force_target_to_clean_output: bool = True,
    auto_wrong_from_corrupted_output: bool = True,
) -> Tuple[str, Optional[str], Optional[str], Optional[str]]:
    spec = resolve_answer_decision_spec(
        pair=pair,
        clean_variant="control",
        corrupted_variant="treatment",
        clean_code=clean_code,
        corrupted_code=corrupted_code,
        use_gold_answer=use_gold_answer,
        manual_target_answer=manual_target_answer,
        manual_wrong_answer=manual_wrong_answer,
        force_target_to_clean_output=force_target_to_clean_output,
        auto_wrong_from_corrupted_output=auto_wrong_from_corrupted_output,
    )
    return (
        spec.forward_positive_answer,
        spec.forward_negative_answer,
        spec.clean_expected_output,
        spec.corrupted_expected_output,
    )


def slugify(value: str) -> str:
    value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
    return value.strip("._") or "pair"


def role_kind(role: Any) -> str:
    text = str(role or "")
    if text.startswith("aligned_"):
        return "aligned"
    if text.startswith("conflicting_"):
        return "conflicting"
    return "unknown"


def parse_stage_csv(value: Optional[str], *, default: Sequence[str] = ()) -> Set[str]:
    if value is None or str(value).strip() == "":
        return set(default)
    stages = {
        part.strip().lower().replace(".", "_")
        for part in str(value).split(",")
        if part.strip()
    }
    invalid = sorted(stages.difference(RQ_STAGES))
    if invalid:
        raise ValueError(
            f"Invalid RQ stage(s): {', '.join(invalid)}. "
            f"Expected comma-separated values from: {', '.join(RQ_STAGES)}"
        )
    return stages


def resolve_model_name_and_slug(
    model_arg: Optional[str],
    model_name: str,
    model_slug: Optional[str],
) -> Tuple[str, str]:
    selected = str(model_arg or model_name or "").strip()
    if model_arg and selected in MODEL_NAME_BY_SLUG:
        resolved_name = MODEL_NAME_BY_SLUG[selected]
        resolved_slug = selected
    else:
        resolved_name = MODEL_NAME_BY_SLUG.get(selected, selected)
        resolved_slug = model_slug or MODEL_SLUG_BY_NAME.get(resolved_name) or slugify(resolved_name).lower()
    if model_slug:
        resolved_slug = str(model_slug)
    return resolved_name, resolved_slug


def rq_pair_output_dir(
    output_root: Path,
    model_slug: str,
    pair_family: str,
    pair_id: str,
) -> Path:
    return output_root / str(model_slug) / str(pair_family) / str(pair_id)


def rq_stage_dir(pair_dir: Path, stage: str) -> Path:
    return pair_dir / RQ_STAGE_DIRS[stage]


def load_json_dict(path: Path) -> Optional[Dict[str, Any]]:
    if not path.exists():
        return None
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return None
    return payload if isinstance(payload, dict) else None


def write_json_dict(path: Path, payload: Dict[str, Any]) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8")
    return path


def stage_is_complete(stage_dir: Path, stage: str) -> bool:
    if stage == "rq1_1":
        return all((stage_dir / name).exists() for name in ("generations.json", "behavior.json", "prompts.json"))
    if stage == "rq1_2":
        summary = load_json_dict(stage_dir / "summary.json")
        if not isinstance(summary, dict):
            return False
        if summary.get("status") == "skipped_tokenization_mismatch":
            return (stage_dir / "patch_units.csv").exists()
        return all(
            (stage_dir / name).exists()
            for name in ("summary.json", "residual_scores.csv", "top_sites.csv", "patch_units.csv", "output_contexts.json")
        )
    if stage == "rq2_1":
        return all((stage_dir / name).exists() for name in ("generations.json", "assertions.csv", "behavior.json"))
    if stage == "rq2_2":
        summary = load_json_dict(stage_dir / "summary.json")
        if not isinstance(summary, dict):
            return False
        status = str(summary.get("status") or "")
        if status.startswith("skipped_"):
            return True
        return all((stage_dir / name).exists() for name in ("summary.json", "residual_scores.csv", "top_sites.csv", "presentations.json"))
    raise ValueError(f"Unknown stage: {stage!r}")


def clear_stage_dir(stage_dir: Path) -> None:
    if stage_dir.exists():
        shutil.rmtree(stage_dir)
    stage_dir.mkdir(parents=True, exist_ok=True)


def ensure_stage_dir(stage_dir: Path) -> None:
    stage_dir.mkdir(parents=True, exist_ok=True)


def sample_evenly(items: Sequence[int], max_items: Optional[int]) -> List[int]:
    if max_items is None or max_items <= 0 or len(items) <= max_items:
        return [int(x) for x in items]

    out: List[int] = []
    seen = set()
    for step in range(max_items):
        if max_items == 1:
            item_index = 0
        else:
            item_index = round(step * (len(items) - 1) / (max_items - 1))
        value = int(items[int(item_index)])
        if value not in seen:
            out.append(value)
            seen.add(value)

    # Rounded spacing can duplicate indices for short ranges. Fill deterministically.
    if len(out) < max_items:
        for value in items:
            value = int(value)
            if value not in seen:
                out.append(value)
                seen.add(value)
            if len(out) >= max_items:
                break
    return out


@dataclass
class PatchUnit:
    unit_id: int
    unit_kind: str
    clean_positions: List[int]
    corrupted_positions: List[int]
    clean_tokens: List[str]
    corrupted_tokens: List[str]
    clean_token_ids: List[int]
    corrupted_token_ids: List[int]
    clean_text: str
    corrupted_text: str
    is_asymmetric: bool


@dataclass
class AdditionalGeneration:
    task: str
    prompt_role: str
    patch_role: str
    max_new_tokens: int
    prompt: str
    full_text: str
    completion: str
    extracted_output: str

    @property
    def variant(self) -> str:
        return self.prompt_role

    @property
    def variant_role(self) -> str:
        return self.patch_role


@dataclass
class AdditionalGenerations:
    pair_name: str
    source_role: str
    destination_role: str
    source_reference_output: str
    destination_reference_output: Optional[str]
    source_execution_output: Optional[str]
    destination_execution_output: Optional[str]
    source_model_output: str
    destination_model_output: str
    source_answer_generation: AdditionalGeneration
    destination_answer_generation: AdditionalGeneration
    unit_tests: Dict[str, AdditionalGeneration]

    @property
    def target_answer(self) -> str:
        return self.source_reference_output

    @property
    def clean_expected_output(self) -> Optional[str]:
        return self.source_execution_output

    @property
    def corrupted_expected_output(self) -> Optional[str]:
        return self.destination_execution_output

    @property
    def clean_model_answer(self) -> str:
        return self.source_model_output

    @property
    def corrupted_model_answer(self) -> str:
        return self.destination_model_output

    @property
    def clean_answer_generation(self) -> AdditionalGeneration:
        return self.source_answer_generation

    @property
    def corrupted_answer_generation(self) -> AdditionalGeneration:
        return self.destination_answer_generation


def token_str_at(tokens: Sequence[str], pos: Optional[int]) -> str:
    if pos is None or pos < 0 or pos >= len(tokens):
        return ""
    return str(tokens[pos])


def token_id_at(token_ids: Sequence[int], pos: Optional[int]) -> Optional[int]:
    if pos is None or pos < 0 or pos >= len(token_ids):
        return None
    return int(token_ids[pos])


def unit_sort_pos(unit: PatchUnit) -> int:
    positions = unit.clean_positions + unit.corrupted_positions
    return min(positions) if positions else unit.unit_id


def unit_destination_positions(unit: PatchUnit, direction: str) -> List[int]:
    if direction == "forward":
        return [int(pos) for pos in unit.corrupted_positions]
    if direction == "reverse":
        return [int(pos) for pos in unit.clean_positions]
    raise ValueError(f"Unknown patch direction: {direction!r}")


def unit_source_positions(unit: PatchUnit, direction: str) -> List[int]:
    if direction == "forward":
        return [int(pos) for pos in unit.clean_positions]
    if direction == "reverse":
        return [int(pos) for pos in unit.corrupted_positions]
    raise ValueError(f"Unknown patch direction: {direction!r}")


def representative_pos(unit: PatchUnit, direction: str = "forward") -> Optional[int]:
    positions = unit_destination_positions(unit, direction)
    if not positions:
        return None
    return int(positions[0])


def format_patch_unit_axis_label(unit: PatchUnit, direction: str, max_chars: int = 32) -> str:
    if direction == "reverse":
        source_text = unit.corrupted_text
        source_token_ids = unit.corrupted_token_ids
        destination_text = unit.clean_text
        destination_token_ids = unit.clean_token_ids
    else:
        source_text = unit.clean_text
        source_token_ids = unit.clean_token_ids
        destination_text = unit.corrupted_text
        destination_token_ids = unit.corrupted_token_ids

    source_label = format_span_for_axis(
        source_text,
        max_chars=14,
        token_ids=source_token_ids,
    )
    if source_text == destination_text and source_token_ids == destination_token_ids:
        return source_label

    destination_label = format_span_for_axis(
        destination_text,
        max_chars=14,
        token_ids=destination_token_ids,
    )
    label = f"{source_label} -> {destination_label}"
    if len(label) > max_chars:
        label = label[: max_chars - 3] + "..."
    return label


def build_alignment_patch_units(
    clean_ids: Sequence[int],
    corrupted_ids: Sequence[int],
    clean_strs: Sequence[str],
    corrupted_strs: Sequence[str],
) -> Tuple[List[PatchUnit], List[Dict[str, Any]], Optional[int]]:
    matcher = difflib.SequenceMatcher(
        a=[int(token_id) for token_id in clean_ids],
        b=[int(token_id) for token_id in corrupted_ids],
        autojunk=False,
    )
    units: List[PatchUnit] = []
    alignment_rows: List[Dict[str, Any]] = []
    first_changed_pos: Optional[int] = None

    def add_alignment_row(
        clean_pos: Optional[int],
        corrupted_pos: Optional[int],
        token_ids_equal: bool,
        patch_unit_id: Optional[int],
        change_unit_id: Optional[int],
    ) -> None:
        pos = corrupted_pos if corrupted_pos is not None else clean_pos
        alignment_rows.append(
            {
                "alignment_row": int(len(alignment_rows)),
                "clean_pos": None if clean_pos is None else int(clean_pos),
                "corrupted_pos": None if corrupted_pos is None else int(corrupted_pos),
                "pos": None if pos is None else int(pos),
                "is_first_changed": False,
                "clean_token_id": token_id_at(clean_ids, clean_pos),
                "corrupted_token_id": token_id_at(corrupted_ids, corrupted_pos),
                "clean_token_str": token_str_at(clean_strs, clean_pos),
                "corrupted_token_str": token_str_at(corrupted_strs, corrupted_pos),
                "token_ids_equal": bool(token_ids_equal),
                "patch_unit_id": None if patch_unit_id is None else int(patch_unit_id),
                "change_unit_id": None if change_unit_id is None else int(change_unit_id),
            }
        )

    for tag, clean_start, clean_end, corrupted_start, corrupted_end in matcher.get_opcodes():
        clean_positions = list(range(clean_start, clean_end))
        corrupted_positions = list(range(corrupted_start, corrupted_end))
        if tag == "equal":
            for clean_pos, corrupted_pos in zip(clean_positions, corrupted_positions):
                unit = PatchUnit(
                    unit_id=len(units),
                    unit_kind="token",
                    clean_positions=[int(clean_pos)],
                    corrupted_positions=[int(corrupted_pos)],
                    clean_tokens=[token_str_at(clean_strs, clean_pos)],
                    corrupted_tokens=[token_str_at(corrupted_strs, corrupted_pos)],
                    clean_token_ids=[int(clean_ids[clean_pos])],
                    corrupted_token_ids=[int(corrupted_ids[corrupted_pos])],
                    clean_text=token_str_at(clean_strs, clean_pos),
                    corrupted_text=token_str_at(corrupted_strs, corrupted_pos),
                    is_asymmetric=False,
                )
                units.append(unit)
                add_alignment_row(clean_pos, corrupted_pos, True, unit.unit_id, None)
            continue

        if first_changed_pos is None:
            first_changed_pos = int(
                corrupted_start if corrupted_start < len(corrupted_ids) else clean_start
            )

        unit = PatchUnit(
            unit_id=len(units),
            unit_kind="changed_region",
            clean_positions=[int(pos) for pos in clean_positions],
            corrupted_positions=[int(pos) for pos in corrupted_positions],
            clean_tokens=[token_str_at(clean_strs, pos) for pos in clean_positions],
            corrupted_tokens=[token_str_at(corrupted_strs, pos) for pos in corrupted_positions],
            clean_token_ids=[int(clean_ids[pos]) for pos in clean_positions],
            corrupted_token_ids=[int(corrupted_ids[pos]) for pos in corrupted_positions],
            clean_text="".join(token_str_at(clean_strs, pos) for pos in clean_positions),
            corrupted_text="".join(token_str_at(corrupted_strs, pos) for pos in corrupted_positions),
            is_asymmetric=len(clean_positions) != len(corrupted_positions),
        )
        units.append(unit)
        row_count = max(len(clean_positions), len(corrupted_positions))
        for offset in range(row_count):
            clean_pos = clean_positions[offset] if offset < len(clean_positions) else None
            corrupted_pos = (
                corrupted_positions[offset] if offset < len(corrupted_positions) else None
            )
            add_alignment_row(clean_pos, corrupted_pos, False, unit.unit_id, unit.unit_id)

    if first_changed_pos is not None:
        for row in alignment_rows:
            if not row["token_ids_equal"]:
                row["is_first_changed"] = True
                break

    return units, alignment_rows, first_changed_pos


def patch_units_from_first_change(
    units: Sequence[PatchUnit],
    first_changed_pos: Optional[int],
    context_window: int = 0,
) -> List[PatchUnit]:
    if first_changed_pos is None:
        return []

    start = max(0, int(first_changed_pos) - max(int(context_window), 0))
    return [unit for unit in units if unit_sort_pos(unit) >= start]


def sample_patch_units(units: Sequence[PatchUnit], max_units: Optional[int]) -> List[PatchUnit]:
    units = list(units)
    if max_units is None or max_units <= 0 or len(units) <= max_units:
        return units

    changed_units = [unit for unit in units if unit.unit_kind == "changed_region"]
    token_units = [unit for unit in units if unit.unit_kind != "changed_region"]
    remaining = int(max_units) - len(changed_units)
    if remaining <= 0:
        return sorted(changed_units, key=lambda unit: unit.unit_id)

    sampled_token_indexes = set(sample_evenly(list(range(len(token_units))), remaining))
    selected = changed_units + [
        unit for index, unit in enumerate(token_units) if index in sampled_token_indexes
    ]
    return sorted(selected, key=lambda unit: unit.unit_id)


def patch_unit_to_row(unit: PatchUnit, is_swept: bool = False) -> Dict[str, Any]:
    return {
        "unit_id": int(unit.unit_id),
        "unit_kind": unit.unit_kind,
        "pos": representative_pos(unit, "forward"),
        "start_pos": representative_pos(unit, "forward"),
        "end_pos": (
            max(unit.corrupted_positions)
            if unit.corrupted_positions
            else representative_pos(unit, "forward")
        ),
        "span_len": int(len(unit.corrupted_positions)),
        "positions": json.dumps([int(pos) for pos in unit.corrupted_positions]),
        "clean_positions": json.dumps([int(pos) for pos in unit.clean_positions]),
        "corrupted_positions": json.dumps([int(pos) for pos in unit.corrupted_positions]),
        "clean_tokens": json.dumps(list(unit.clean_tokens)),
        "corrupted_tokens": json.dumps(list(unit.corrupted_tokens)),
        "clean_token_ids": json.dumps([int(token_id) for token_id in unit.clean_token_ids]),
        "corrupted_token_ids": json.dumps([int(token_id) for token_id in unit.corrupted_token_ids]),
        "clean_text": unit.clean_text,
        "corrupted_text": unit.corrupted_text,
        "is_asymmetric": bool(unit.is_asymmetric),
        "is_swept": bool(is_swept),
    }


def alignment_row_to_csv_row(row: Dict[str, Any], swept_unit_ids: Sequence[int]) -> Dict[str, Any]:
    swept_unit_id_set = set(int(unit_id) for unit_id in swept_unit_ids)
    swept = row.get("patch_unit_id") in swept_unit_id_set
    return {
        **row,
        "pos": "" if row.get("pos") is None else int(row["pos"]),
        "clean_pos": "" if row.get("clean_pos") is None else int(row["clean_pos"]),
        "corrupted_pos": ""
        if row.get("corrupted_pos") is None
        else int(row["corrupted_pos"]),
        "is_swept": bool(swept),
        "clean_token_id": ""
        if row.get("clean_token_id") is None
        else int(row["clean_token_id"]),
        "corrupted_token_id": ""
        if row.get("corrupted_token_id") is None
        else int(row["corrupted_token_id"]),
        "change_unit_id": ""
        if row.get("change_unit_id") is None
        else int(row["change_unit_id"]),
        "patch_unit_id": ""
        if row.get("patch_unit_id") is None
        else int(row["patch_unit_id"]),
    }


RESIDUAL_SCORE_ARTIFACT_NAMES = (
    "residual_scores.csv",
    "top_residual_sites.csv",
    "reverse_residual_scores.csv",
    "top_reverse_residual_sites.csv",
    "residual_scores_heatmap.pdf",
    "residual_recovery_heatmap.pdf",
    "reverse_residual_scores_heatmap.pdf",
    "reverse_residual_recovery_heatmap.pdf",
)

LEGACY_SPAN_ARTIFACT_NAMES = (
    "residual_span_scores.csv",
    "top_residual_spans.csv",
    "reverse_residual_span_scores.csv",
    "top_reverse_residual_spans.csv",
    "residual_span_recovery_heatmap.pdf",
    "reverse_residual_span_recovery_heatmap.pdf",
)

SKIP_ARTIFACT_NAMES = ("skipped_tokenization_mismatch.txt",)



def remove_artifacts(pair_output_dir: Path, names: Sequence[str]) -> None:
    for name in names:
        (pair_output_dir / name).unlink(missing_ok=True)


def changed_region_units(units: Sequence[PatchUnit]) -> List[PatchUnit]:
    return [unit for unit in units if unit.unit_kind == "changed_region"]


def changed_region_summary(units: Sequence[PatchUnit]) -> List[Dict[str, Any]]:
    rows = []
    for unit in changed_region_units(units):
        rows.append(
            {
                "patch_unit_id": int(unit.unit_id),
                "patch_unit_type": unit.unit_kind,
                "patch_unit_text": unit.corrupted_text,
                "span_id": int(unit.unit_id),
                "unit_id": int(unit.unit_id),
                "start_pos": representative_pos(unit, "forward"),
                "end_pos": max(unit.corrupted_positions) if unit.corrupted_positions else None,
                "span_len": int(len(unit.corrupted_positions)),
                "positions": [int(pos) for pos in unit.corrupted_positions],
                "source_positions": [int(pos) for pos in unit.clean_positions],
                "destination_positions": [int(pos) for pos in unit.corrupted_positions],
                "clean_positions": [int(pos) for pos in unit.clean_positions],
                "corrupted_positions": [int(pos) for pos in unit.corrupted_positions],
                "source_text": unit.clean_text,
                "destination_text": unit.corrupted_text,
                "clean_text": unit.clean_text,
                "corrupted_text": unit.corrupted_text,
                "is_asymmetric": bool(unit.is_asymmetric),
            }
        )
    return rows


def patch_units_to_rows(
    units: Sequence[PatchUnit],
    swept_unit_ids: Sequence[int],
) -> List[Dict[str, Any]]:
    swept = set(int(unit_id) for unit_id in swept_unit_ids)
    return [patch_unit_to_row(unit, is_swept=unit.unit_id in swept) for unit in units]


def methodology_patch_unit_to_row(unit: PatchUnit, is_swept: bool = False) -> Dict[str, Any]:
    return {
        "patch_unit_id": int(unit.unit_id),
        "patch_unit_type": unit.unit_kind,
        "patch_unit_text": unit.corrupted_text,
        "source_positions": json.dumps([int(pos) for pos in unit.clean_positions]),
        "destination_positions": json.dumps([int(pos) for pos in unit.corrupted_positions]),
        "source_tokens": json.dumps(list(unit.clean_tokens)),
        "destination_tokens": json.dumps(list(unit.corrupted_tokens)),
        "source_token_ids": json.dumps([int(token_id) for token_id in unit.clean_token_ids]),
        "destination_token_ids": json.dumps([int(token_id) for token_id in unit.corrupted_token_ids]),
        "source_text": unit.clean_text,
        "destination_text": unit.corrupted_text,
        "is_asymmetric": bool(unit.is_asymmetric),
        "is_swept": bool(is_swept),
    }


def methodology_patch_units_to_rows(
    units: Sequence[PatchUnit],
    swept_unit_ids: Sequence[int],
) -> List[Dict[str, Any]]:
    swept = set(int(unit_id) for unit_id in swept_unit_ids)
    return [
        methodology_patch_unit_to_row(unit, is_swept=unit.unit_id in swept)
        for unit in units
    ]


def methodology_alignment_row_to_csv_row(
    row: Dict[str, Any],
    swept_unit_ids: Sequence[int],
) -> Dict[str, Any]:
    swept_unit_id_set = set(int(unit_id) for unit_id in swept_unit_ids)
    patch_unit_id = row.get("patch_unit_id")
    swept = patch_unit_id in swept_unit_id_set
    source_pos = row.get("clean_pos")
    destination_pos = row.get("corrupted_pos")
    pos = row.get("pos")
    if pos is None:
        pos = destination_pos if destination_pos is not None else source_pos
    return {
        "alignment_row": ""
        if row.get("alignment_row") is None
        else int(row["alignment_row"]),
        "pos": "" if pos is None else int(pos),
        "source_pos": "" if source_pos is None else int(source_pos),
        "destination_pos": ""
        if destination_pos is None
        else int(destination_pos),
        "is_swept": bool(swept),
        "is_first_changed": bool(row.get("is_first_changed", False)),
        "source_token_id": ""
        if row.get("clean_token_id") is None
        else int(row["clean_token_id"]),
        "destination_token_id": ""
        if row.get("corrupted_token_id") is None
        else int(row["corrupted_token_id"]),
        "source_token_str": row.get("clean_token_str", ""),
        "destination_token_str": row.get("corrupted_token_str", ""),
        "token_ids_equal": bool(row.get("token_ids_equal", False)),
        "patch_unit_id": ""
        if patch_unit_id is None
        else int(patch_unit_id),
        "change_unit_id": ""
        if row.get("change_unit_id") is None
        else int(row["change_unit_id"]),
    }


METHODOLOGY_RESIDUAL_SCORE_COLUMNS = [
    "score_metric",
    "pair_id",
    "pair_family",
    "patch_direction",
    "source_role",
    "destination_role",
    "source_file",
    "destination_file",
    "source_execution_output",
    "source_cue_output",
    "destination_execution_output",
    "destination_cue_output",
    "source_output",
    "destination_output",
    "source_output_label",
    "destination_output_label",
    "source_is_accurate",
    "destination_is_accurate",
    "source_token",
    "destination_token",
    "source_token_id",
    "destination_token_id",
    "divergent_token_index",
    "common_target_prefix",
    "destination_margin",
    "source_margin",
    "layer",
    "pos",
    "patch_unit_id",
    "patch_unit_type",
    "patch_unit_text",
    "start_pos",
    "end_pos",
    "span_len",
    "positions",
    "source_positions",
    "destination_positions",
    "source_tokens",
    "destination_tokens",
    "source_token_ids",
    "destination_token_ids",
    "source_text",
    "destination_text",
    "is_asymmetric",
    "patched_margin",
    "margin_delta",
    "signed_recovery",
    "distance_recovery",
]


def methodology_residual_score_df(df):
    public_df = pd.DataFrame(index=df.index)
    for column in METHODOLOGY_RESIDUAL_SCORE_COLUMNS:
        if column in df.columns:
            public_df[column] = df[column]
        elif column == "patch_unit_id" and "unit_id" in df.columns:
            public_df[column] = df["unit_id"]
        elif column == "patch_unit_type" and "unit_kind" in df.columns:
            public_df[column] = df["unit_kind"]
        elif column == "patch_unit_text":
            if "corrupted_text" in df.columns:
                public_df[column] = df["corrupted_text"]
            elif "clean_text" in df.columns:
                public_df[column] = df["clean_text"]
        elif column == "source_positions" and "clean_positions" in df.columns:
            public_df[column] = df["clean_positions"]
        elif column == "destination_positions" and "corrupted_positions" in df.columns:
            public_df[column] = df["corrupted_positions"]
        elif column == "source_tokens" and "clean_tokens" in df.columns:
            public_df[column] = df["clean_tokens"]
        elif column == "destination_tokens" and "corrupted_tokens" in df.columns:
            public_df[column] = df["corrupted_tokens"]
        elif column == "source_token_ids" and "clean_token_ids" in df.columns:
            public_df[column] = df["clean_token_ids"]
        elif column == "destination_token_ids" and "corrupted_token_ids" in df.columns:
            public_df[column] = df["corrupted_token_ids"]
        elif column == "source_text" and "clean_text" in df.columns:
            public_df[column] = df["clean_text"]
        elif column == "destination_text" and "corrupted_text" in df.columns:
            public_df[column] = df["corrupted_text"]
        elif column == "patched_margin" and "score" in df.columns:
            public_df[column] = df["score"]
        elif column == "distance_recovery" and "distance_based_recovery" in df.columns:
            public_df[column] = df["distance_based_recovery"]
        elif column == "destination_margin" and "baseline_margin" in df.columns:
            public_df[column] = df["baseline_margin"]
    return public_df


def write_methodology_residual_scores(df, path: Path) -> None:
    methodology_residual_score_df(df).to_csv(path, index=False)


def best_changed_region_row(top_sites_df) -> Dict[str, Any]:
    if top_sites_df.empty or "unit_kind" not in top_sites_df.columns:
        return {}
    changed_df = top_sites_df[top_sites_df["unit_kind"] == "changed_region"]
    if changed_df.empty:
        return {}
    return changed_df.iloc[0].to_dict()


def json_positions_from_row(row: Dict[str, Any], key: str = "positions") -> List[int]:
    if not row:
        return []
    value = row.get(key, "[]")
    if isinstance(value, list):
        return [int(pos) for pos in value]
    if value in ("", None):
        return []
    return [int(pos) for pos in json.loads(str(value))]


def parse_bool_cell(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() in {"1", "true", "yes", "y"}


def maybe_float_cell(value: Any) -> Optional[float]:
    if value in ("", None):
        return None
    try:
        output = float(value)
    except (TypeError, ValueError):
        return None
    return output if math.isfinite(output) else None




def line_start_offsets(text: str) -> List[int]:
    offsets = [0]
    for match in re.finditer("\n", text):
        offsets.append(match.end())
    return offsets


def node_char_span(source: str, node: ast.AST) -> Optional[Tuple[int, int]]:
    lineno = getattr(node, "lineno", None)
    col_offset = getattr(node, "col_offset", None)
    end_lineno = getattr(node, "end_lineno", None)
    end_col_offset = getattr(node, "end_col_offset", None)
    if None in (lineno, col_offset, end_lineno, end_col_offset):
        return None

    offsets = line_start_offsets(source)
    if lineno < 1 or end_lineno < 1 or lineno > len(offsets) or end_lineno > len(offsets):
        return None
    return (
        int(offsets[int(lineno) - 1] + int(col_offset)),
        int(offsets[int(end_lineno) - 1] + int(end_col_offset)),
    )


def source_segment(source: str, node: ast.AST) -> str:
    segment = ast.get_source_segment(source, node)
    if segment is not None:
        return segment
    try:
        return ast.unparse(node)
    except Exception:
        return ""


def collect_supported_assert_equalities(test_source: str) -> List[Dict[str, Any]]:
    tree = ast.parse(test_source)
    setup_by_assert_span: Dict[Tuple[int, int], str] = {}

    def collect_setup_sources(statements: Sequence[ast.stmt]) -> None:
        for index, statement in enumerate(statements):
            if isinstance(statement, ast.Assert):
                statement_span = node_char_span(test_source, statement)
                if statement_span is not None:
                    setup_segments = []
                    for setup_statement in statements[:index]:
                        if isinstance(setup_statement, ast.Assert):
                            continue
                        segment = source_segment(test_source, setup_statement)
                        if segment:
                            setup_segments.append(textwrap.dedent(segment).strip())
                    setup_by_assert_span[statement_span] = "\n".join(
                        segment for segment in setup_segments if segment
                    )
                continue
            nested_body = getattr(statement, "body", None)
            if isinstance(nested_body, list):
                collect_setup_sources(nested_body)

    collect_setup_sources(tree.body)
    assertions: List[Dict[str, Any]] = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Assert):
            continue
        if not isinstance(node.test, ast.Compare):
            continue
        compare = node.test
        if len(compare.ops) != 1 or not isinstance(compare.ops[0], ast.Eq):
            continue
        if len(compare.comparators) != 1:
            continue

        expr_node = compare.left
        expected_node = compare.comparators[0]
        assertion_span = node_char_span(test_source, node)
        expr_span = node_char_span(test_source, expr_node)
        expected_span = node_char_span(test_source, expected_node)
        if assertion_span is None or expr_span is None or expected_span is None:
            continue

        assertions.append(
            {
                "assertion_index": int(len(assertions)),
                "lineno": int(getattr(node, "lineno", 0)),
                "col_offset": int(getattr(node, "col_offset", 0)),
                "assertion_start": int(assertion_span[0]),
                "assertion_end": int(assertion_span[1]),
                "expr_start": int(expr_span[0]),
                "expr_end": int(expr_span[1]),
                "expected_start": int(expected_span[0]),
                "expected_end": int(expected_span[1]),
                "assertion_source": test_source[assertion_span[0] : assertion_span[1]],
                "expr_source": source_segment(test_source, expr_node).strip(),
                "expected_source": source_segment(test_source, expected_node).strip(),
                "setup_source": setup_by_assert_span.get(assertion_span, ""),
            }
        )
    assertions = sorted(assertions, key=lambda row: (row["assertion_start"], row["assertion_end"]))
    for assertion_index, row in enumerate(assertions):
        row["assertion_index"] = int(assertion_index)
    return assertions


def execute_snippet_namespace(code: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
    import io
    from contextlib import redirect_stdout

    namespace: Dict[str, Any] = {"__builtins__": __builtins__}
    try:
        with redirect_stdout(io.StringIO()):
            exec(code, namespace, namespace)
    except Exception as e:
        return None, str(e)
    return namespace, None


def cue_oracle_namespace_for_pair(pair_id: Optional[str]) -> Tuple[Optional[Dict[str, Any]], str]:
    normalized_pair_id = Path(str(pair_id or "")).name
    if not normalized_pair_id:
        return None, "cue_oracle_pair_id_missing"
    try:
        from scripts.analyze_rq2_1_cue_assertion_labels import CUE_ORACLES
    except Exception as e:  # pragma: no cover - import guard for non-analysis installs
        return None, f"cue_oracle_import_error:{type(e).__name__}:{e}"

    factory = CUE_ORACLES.get(normalized_pair_id)
    if factory is None:
        return None, f"cue_oracle_missing:{normalized_pair_id}"
    try:
        return factory(), "loaded"
    except Exception as e:  # pragma: no cover - preserve oracle construction failures in artifacts
        return None, f"cue_oracle_build_error:{type(e).__name__}:{e}"


def evaluate_expression_in_namespace(
    expression_source: str,
    namespace: Dict[str, Any],
    setup_source: str = "",
) -> Tuple[Any, Optional[str]]:
    import io
    import sys
    from contextlib import redirect_stdout

    try:
        eval_namespace = dict(namespace)
        eval_namespace.setdefault("io", io)
        eval_namespace.setdefault("StringIO", io.StringIO)
        eval_namespace.setdefault("sys", sys)
        compiled = compile(expression_source, "<unit_test_assertion_expr>", "eval")
        with redirect_stdout(io.StringIO()):
            if setup_source.strip():
                exec(
                    compile(
                        textwrap.dedent(setup_source),
                        "<unit_test_assertion_setup>",
                        "exec",
                    ),
                    eval_namespace,
                    eval_namespace,
                )
            return eval(compiled, eval_namespace, eval_namespace), None
    except Exception as e:
        return None, str(e)


def literal_eval_source(source: str) -> Tuple[Any, Optional[str]]:
    try:
        return ast.literal_eval(source), None
    except Exception as e:
        return None, str(e)


def first_divergent_token_ids(
    positive_ids: Sequence[int],
    negative_ids: Sequence[int],
) -> Optional[Tuple[int, int, int, List[int]]]:
    positive = [int(token_id) for token_id in positive_ids]
    negative = [int(token_id) for token_id in negative_ids]
    limit = min(len(positive), len(negative))
    for index in range(limit):
        if positive[index] != negative[index]:
            return int(index), int(positive[index]), int(negative[index]), positive[:index]
    if len(positive) != len(negative):
        if limit >= len(positive) or limit >= len(negative):
            return None
        return int(limit), int(positive[limit]), int(negative[limit]), positive[:limit]
    return None


def extract_unit_test_decisions(
    clean_code: str,
    corrupted_code: str,
    clean_test_code: str,
    corrupted_test_code: str,
    cue_oracle_namespace: Optional[Dict[str, Any]] = None,
    require_cue_oracle_targets: bool = False,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]:
    skips: List[Dict[str, Any]] = []
    clean_test_code = cap_unit_test_source_to_assert_limit(clean_test_code)
    corrupted_test_code = cap_unit_test_source_to_assert_limit(corrupted_test_code)
    try:
        clean_assertions = collect_supported_assert_equalities(clean_test_code)
    except SyntaxError as e:
        clean_assertions = []
        skips.append(
            {
                "variant_role": "clean",
                "assertion_index": "",
                "stage": "parse",
                "reason": "syntax_error",
                "error": str(e),
            }
        )
    try:
        corrupted_assertions = collect_supported_assert_equalities(corrupted_test_code)
    except SyntaxError as e:
        corrupted_assertions = []
        skips.append(
            {
                "variant_role": "corrupted",
                "assertion_index": "",
                "stage": "parse",
                "reason": "syntax_error",
                "error": str(e),
            }
        )

    corrupted_namespace, corrupted_exec_error = execute_snippet_namespace(corrupted_code)
    if corrupted_exec_error is not None:
        skips.append(
            {
                "variant_role": "corrupted",
                "assertion_index": "",
                "stage": "exec_snippet",
                "reason": "snippet_execution_error",
                "error": corrupted_exec_error,
            }
        )

    clean_namespace, clean_exec_error = execute_snippet_namespace(clean_code)
    if clean_exec_error is not None:
        skips.append(
            {
                "variant_role": "clean",
                "assertion_index": "",
                "stage": "exec_snippet",
                "reason": "snippet_execution_error",
                "error": clean_exec_error,
            }
        )

    decisions: List[Dict[str, Any]] = []
    if corrupted_namespace is None:
        return decisions, skips, {
            "clean_assertions": int(len(clean_assertions)),
            "corrupted_assertions": int(len(corrupted_assertions)),
        }
    if require_cue_oracle_targets and cue_oracle_namespace is None:
        skips.append(
            {
                "variant_role": "cue_oracle",
                "assertion_index": "",
                "stage": "target_selection",
                "reason": "cue_oracle_missing",
            }
        )
        return decisions, skips, {
            "clean_assertions": int(len(clean_assertions)),
            "corrupted_assertions": int(len(corrupted_assertions)),
        }

    for assertion_index, corrupted_assertion in enumerate(corrupted_assertions):
        decision_id = f"assertion_{assertion_index}"
        corrupted_runtime_value, corrupted_runtime_error = evaluate_expression_in_namespace(
            corrupted_assertion["expr_source"],
            corrupted_namespace,
            setup_source=corrupted_assertion.get("setup_source", ""),
        )
        if corrupted_runtime_error is not None:
            skips.append(
                {
                    "variant_role": "corrupted",
                    "assertion_index": int(assertion_index),
                    "decision_id": decision_id,
                    "stage": "eval_assertion_lhs",
                    "reason": "lhs_runtime_eval_error",
                    "error": corrupted_runtime_error,
                    "expr_source": corrupted_assertion["expr_source"],
                }
            )
            continue

        same_expression_matches = [
            assertion
            for assertion in clean_assertions
            if assertion["expr_source"] == corrupted_assertion["expr_source"]
        ]
        if len(same_expression_matches) == 1:
            clean_assertion = same_expression_matches[0]
            clean_match_type = "same_expression"
        elif assertion_index < len(clean_assertions):
            clean_assertion = clean_assertions[assertion_index]
            clean_match_type = "same_index"
        else:
            clean_assertion = None
            clean_match_type = "none"
        if clean_assertion is None:
            skips.append(
                {
                    "variant_role": "clean",
                    "assertion_index": int(assertion_index),
                    "decision_id": decision_id,
                    "stage": "match_clean_assertion",
                    "reason": "no_clean_assertion_at_same_index",
                }
            )
            continue

        clean_expected_source = clean_assertion["expected_source"]
        if clean_namespace is None:
            skips.append(
                {
                    "variant_role": "clean",
                    "assertion_index": int(assertion_index),
                    "decision_id": decision_id,
                    "stage": "eval_assertion_lhs",
                    "reason": "clean_namespace_unavailable",
                    "expr_source": clean_assertion["expr_source"],
                }
            )
            continue

        clean_runtime_value, clean_runtime_error = evaluate_expression_in_namespace(
            clean_assertion["expr_source"],
            clean_namespace,
            setup_source=clean_assertion.get("setup_source", ""),
        )
        if clean_runtime_error is not None:
            skips.append(
                {
                    "variant_role": "clean",
                    "assertion_index": int(assertion_index),
                    "decision_id": decision_id,
                    "stage": "eval_assertion_lhs",
                    "reason": "lhs_runtime_eval_error",
                    "error": clean_runtime_error,
                    "expr_source": clean_assertion["expr_source"],
                }
            )
            continue

        clean_runtime_source = repr(clean_runtime_value)
        corrupted_runtime_source = repr(corrupted_runtime_value)
        clean_cue_oracle_source = ""
        corrupted_cue_oracle_source = ""
        forward_positive_source = ""
        forward_positive_source_kind = ""
        forward_negative_source = ""
        forward_negative_source_kind = ""
        reverse_positive_source = ""
        reverse_positive_source_kind = ""
        reverse_negative_source = ""
        reverse_negative_source_kind = ""
        target_mode = ""
        reverse_unavailable_reason = ""
        reverse_available = False

        if require_cue_oracle_targets:
            corrupted_cue_value, corrupted_cue_error = evaluate_expression_in_namespace(
                corrupted_assertion["expr_source"],
                cue_oracle_namespace or {},
                setup_source=corrupted_assertion.get("setup_source", ""),
            )
            if corrupted_cue_error is not None:
                skips.append(
                    {
                        "variant_role": "cue_oracle",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "eval_assertion_lhs",
                        "reason": "cue_oracle_lhs_eval_error",
                        "error": corrupted_cue_error,
                        "expr_source": corrupted_assertion["expr_source"],
                    }
                )
                continue
            clean_cue_value, clean_cue_error = evaluate_expression_in_namespace(
                clean_assertion["expr_source"],
                cue_oracle_namespace or {},
                setup_source=clean_assertion.get("setup_source", ""),
            )
            if clean_cue_error is not None:
                skips.append(
                    {
                        "variant_role": "cue_oracle",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "eval_assertion_lhs",
                        "reason": "cue_oracle_lhs_eval_error",
                        "error": clean_cue_error,
                        "expr_source": clean_assertion["expr_source"],
                    }
                )
                continue

            clean_cue_oracle_source = repr(clean_cue_value)
            corrupted_cue_oracle_source = repr(corrupted_cue_value)
            if (
                clean_runtime_value == clean_cue_value
                and corrupted_runtime_value == corrupted_cue_value
            ):
                skips.append(
                    {
                        "variant_role": "cue_oracle",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "target_selection",
                        "reason": "execution_and_cue_oracle_identical",
                        "expr_source": corrupted_assertion["expr_source"],
                        "execution_source": corrupted_runtime_source,
                        "cue_oracle_source": corrupted_cue_oracle_source,
                    }
                )
                continue

            target_mode = "cue_oracle_pair"
            forward_positive_source = corrupted_runtime_source
            forward_positive_source_kind = "destination_execution_runtime"
            forward_negative_source = corrupted_cue_oracle_source
            forward_negative_source_kind = "destination_cue_oracle"
            reverse_available = clean_runtime_value != clean_cue_value
            reverse_positive_source = clean_cue_oracle_source
            reverse_positive_source_kind = "source_cue_oracle"
            reverse_negative_source = clean_runtime_source
            reverse_negative_source_kind = "source_execution_runtime"
            reverse_unavailable_reason = (
                "" if reverse_available else "source_execution_and_cue_oracle_identical"
            )
        elif clean_runtime_value != corrupted_runtime_value:
            target_mode = "runtime_pair"
            forward_positive_source = clean_runtime_source
            forward_positive_source_kind = "clean_runtime"
            forward_negative_source = corrupted_runtime_source
            forward_negative_source_kind = "corrupted_runtime"
            reverse_positive_source = corrupted_runtime_source
            reverse_positive_source_kind = "corrupted_runtime"
            reverse_negative_source = clean_runtime_source
            reverse_negative_source_kind = "clean_runtime"
            reverse_available = True
        else:
            expected_value, expected_error = literal_eval_source(
                corrupted_assertion["expected_source"]
            )
            if expected_error is not None:
                skips.append(
                    {
                        "variant_role": "corrupted",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "literal_eval_expected",
                        "reason": "generated_expected_unparseable",
                        "error": expected_error,
                        "expected_source": corrupted_assertion["expected_source"],
                    }
                )
                continue
            clean_expected_value, clean_expected_error = literal_eval_source(clean_expected_source)
            if clean_expected_error is not None:
                skips.append(
                    {
                        "variant_role": "clean",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "literal_eval_expected",
                        "reason": "clean_generated_expected_unparseable",
                        "error": clean_expected_error,
                        "expected_source": clean_expected_source,
                    }
                )
                continue
            if clean_expected_value == expected_value:
                skips.append(
                    {
                        "variant_role": "both",
                        "assertion_index": int(assertion_index),
                        "decision_id": decision_id,
                        "stage": "target_selection",
                        "reason": "clean_and_corrupted_generated_expected_identical",
                        "expr_source": corrupted_assertion["expr_source"],
                        "expected_source": corrupted_assertion["expected_source"],
                    }
                )
                continue
            else:
                target_mode = "generated_expected"
                forward_positive_source = clean_runtime_source
                forward_positive_source_kind = "common_runtime"
                forward_negative_source = corrupted_assertion["expected_source"]
                forward_negative_source_kind = "corrupted_generated_expected"
                reverse_available = True
                reverse_positive_source = corrupted_assertion["expected_source"]
                reverse_positive_source_kind = "corrupted_generated_expected"
                reverse_negative_source = clean_expected_source
                reverse_negative_source_kind = "clean_generated_expected"

        decisions.append(
            {
                "decision_id": decision_id,
                "assertion_index": int(assertion_index),
                "corrupted_assertion": corrupted_assertion,
                "clean_assertion": clean_assertion,
                "corrupted_expr_source": corrupted_assertion["expr_source"],
                "corrupted_expected_source": corrupted_assertion["expected_source"],
                "runtime_correct_source": corrupted_runtime_source,
                "runtime_correct_type": type(corrupted_runtime_value).__name__,
                "clean_expected_source": clean_expected_source,
                "clean_runtime_source": clean_runtime_source,
                "corrupted_runtime_source": corrupted_runtime_source,
                "clean_cue_oracle_source": clean_cue_oracle_source,
                "corrupted_cue_oracle_source": corrupted_cue_oracle_source,
                "target_mode": target_mode,
                "forward_positive_source": forward_positive_source,
                "forward_positive_source_kind": forward_positive_source_kind,
                "forward_negative_source": forward_negative_source,
                "forward_negative_source_kind": forward_negative_source_kind,
                "clean_match_type": clean_match_type,
                "reverse_available": bool(reverse_available),
                "reverse_positive_source": reverse_positive_source,
                "reverse_positive_source_kind": reverse_positive_source_kind,
                "reverse_negative_source": reverse_negative_source,
                "reverse_negative_source_kind": reverse_negative_source_kind,
                "reverse_unavailable_reason": reverse_unavailable_reason,
            }
        )

    return decisions, skips, {
        "clean_assertions": int(len(clean_assertions)),
        "corrupted_assertions": int(len(corrupted_assertions)),
    }


def normalize_recovery(score: float, baseline: float, clean_ref: float, eps: float = 1e-9) -> float:
    return (float(score) - float(baseline)) / (float(clean_ref) - float(baseline) + eps)


def signed_recovery_matrix(
    scores,
    baseline: float,
    source_ref: float,
    eps: float = 1e-9,
):
    return (scores - float(baseline)) / (float(source_ref) - float(baseline) + eps)


def distance_based_recovery(score: float, baseline: float, source_ref: float, eps: float = 1e-9) -> float:
    return 1.0 - abs(float(score) - float(source_ref)) / (
        abs(float(baseline) - float(source_ref)) + eps
    )


def distance_based_recovery_matrix(
    scores,
    baseline: float,
    source_ref: float,
    eps: float = 1e-9,
):
    return 1.0 - torch.abs(scores - float(source_ref)) / (
        abs(float(baseline) - float(source_ref)) + eps
    )


def recovery_matrix_for_metric(
    recovery_metric: str,
    scores,
    baseline: float,
    source_ref: float,
):
    if recovery_metric == "distance":
        return distance_based_recovery_matrix(scores, baseline, source_ref)
    if recovery_metric == "signed":
        return signed_recovery_matrix(scores, baseline, source_ref)
    raise ValueError(f"Unknown recovery metric: {recovery_metric!r}")


def recovery_metric_column(recovery_metric: str) -> str:
    if recovery_metric == "distance":
        return "distance_recovery"
    if recovery_metric == "signed":
        return "signed_recovery"
    raise ValueError(f"Unknown recovery metric: {recovery_metric!r}")


def recovery_heatmap_kwargs(
    recovery_metric: str,
    recovery_scores,
    cmap: Optional[str],
) -> Dict[str, Any]:
    if recovery_metric == "distance":
        return {
            "colorbar_label": "distance recovery",
            "vmin": 0.0,
            "vmax": 1.0,
            "cmap": cmap,
        }

    finite_scores = recovery_scores[torch.isfinite(recovery_scores)]
    max_abs = 1.0
    if finite_scores.numel():
        max_abs = max(max_abs, float(torch.max(torch.abs(finite_scores)).item()))
    return {
        "colorbar_label": "signed recovery",
        "vmin": -max_abs,
        "vmax": max_abs,
        "cmap": SIGNED_DESTINATION_POSITIVE_CMAP
        if cmap == CORRUPTED_RECOVERY_CMAP
        else SIGNED_SOURCE_POSITIVE_CMAP,
    }


def patch_gain_toward_clean(score: float, baseline: float, clean_ref: float) -> float:
    return abs(float(baseline) - float(clean_ref)) - abs(float(score) - float(clean_ref))


class ResidualPatchRunner:
    def __init__(
        self,
        model,
        question_template: str,
        unit_test_prompt_template: str = DEFAULT_UNIT_TEST_PROMPT_TEMPLATE,
        prepend_bos: bool = True,
        hook_name: str = "hook_resid_post",
        patch_batch_size: int = 8,
    ) -> None:
        if torch is None:
            load_runtime_dependencies()
        self.model = model
        self.question_template = question_template
        self.unit_test_prompt_template = unit_test_prompt_template
        self.prepend_bos = prepend_bos
        self.hook_name = hook_name
        self.patch_batch_size = max(1, int(patch_batch_size))
        self.n_layers = int(model.cfg.n_layers)
        self.n_heads = int(model.cfg.n_heads)

    def to_tokens(self, text: str):
        return self.model.to_tokens(text, prepend_bos=self.prepend_bos)

    def token_strs(self, tokens) -> List[str]:
        return self.model.to_str_tokens(tokens)

    def answer_token_ids(self, answer: str):
        answer_toks = self.to_tokens(answer)
        start_idx = 1 if self.prepend_bos else 0
        return answer_toks[0, start_idx:]

    def metric_from_logits_with_target_answer(self, logits, target: str) -> float:
        answer_ids = self.answer_token_ids(target)
        first_tok = answer_ids[0].item()
        return float(logits[0, -1, first_tok].item())

    def decode_tokens(self, tokens) -> str:
        token_ids = tokens.detach().cpu().tolist()
        return self.model.to_string(token_ids)

    def generate_completion(
        self,
        prompt: str,
        max_new_tokens: int,
        stop_condition: Optional[Callable[[str], bool]] = None,
        postprocess_completion: Optional[Callable[[str], str]] = None,
    ) -> Tuple[str, str, str]:
        tokens = self.to_tokens(prompt)
        prompt_len = int(tokens.shape[1])
        with torch.no_grad():
            for _ in range(int(max_new_tokens)):
                logits = self.model(tokens)
                next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
                tokens = torch.cat([tokens, next_token], dim=-1)
                completion = self.decode_tokens(tokens[0, prompt_len:])
                if contains_complete_fenced_block(completion):
                    break
                if stop_condition is not None and stop_condition(completion):
                    break

        completion = self.decode_tokens(tokens[0, prompt_len:])
        if postprocess_completion is not None:
            completion = postprocess_completion(completion)
        else:
            completion = truncate_after_first_fenced_block(completion).strip()
        full_text = prompt + completion
        extracted_output = extract_first_fenced_block(completion)
        return full_text, completion, extracted_output


    def additional_generation_specs(
        self,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        unit_test_max_new_tokens: int,
    ) -> List[Dict]:
        variants = [
            ("source", clean_variant, clean_code),
            ("destination", corrupted_variant, corrupted_code),
        ]
        specs = []
        for patch_role, prompt_role, code in variants:
            specs.append(
                {
                    "task": "unit_test",
                    "patch_role": patch_role,
                    "prompt_role": prompt_role,
                    "max_new_tokens": int(unit_test_max_new_tokens),
                    "prompt": build_prompt(code, self.unit_test_prompt_template),
                }
            )
        return specs

    def generate_extracted_answer(self, code: str, max_new_tokens: int) -> str:
        prompt = build_prompt(code, self.question_template)
        _full_text, completion, _extracted_output = self.generate_completion(
            prompt=prompt,
            max_new_tokens=max_new_tokens,
        )
        return extract_first_fenced_block_or_text(completion)

    def build_answer_generations(
        self,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        answer_max_new_tokens: int,
    ) -> Tuple[AdditionalGeneration, AdditionalGeneration, str, str]:
        clean_answer_prompt = build_prompt(clean_code, self.question_template)
        clean_answer_full_text, clean_answer_completion, _clean_answer_extracted = (
            self.generate_completion(
                prompt=clean_answer_prompt,
                max_new_tokens=answer_max_new_tokens,
            )
        )
        clean_model_answer = clean_final_output_text(
            extract_first_fenced_block_or_text(clean_answer_completion)
        )
        clean_answer_generation = AdditionalGeneration(
            task="answer",
            patch_role="source",
            prompt_role=clean_variant,
            max_new_tokens=int(answer_max_new_tokens),
            prompt=clean_answer_prompt,
            full_text=clean_answer_full_text,
            completion=clean_answer_completion,
            extracted_output=clean_model_answer,
        )

        corrupted_answer_prompt = build_prompt(corrupted_code, self.question_template)
        corrupted_answer_full_text, corrupted_answer_completion, _corrupted_answer_extracted = (
            self.generate_completion(
                prompt=corrupted_answer_prompt,
                max_new_tokens=answer_max_new_tokens,
            )
        )
        corrupted_model_answer = clean_final_output_text(
            extract_first_fenced_block_or_text(corrupted_answer_completion)
        )
        corrupted_answer_generation = AdditionalGeneration(
            task="answer",
            patch_role="destination",
            prompt_role=corrupted_variant,
            max_new_tokens=int(answer_max_new_tokens),
            prompt=corrupted_answer_prompt,
            full_text=corrupted_answer_full_text,
            completion=corrupted_answer_completion,
            extracted_output=corrupted_model_answer,
        )
        return (
            clean_answer_generation,
            corrupted_answer_generation,
            clean_model_answer,
            corrupted_model_answer,
        )

    def build_unit_test_generations(
        self,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        unit_test_max_new_tokens: int,
    ) -> Dict[str, AdditionalGeneration]:
        specs = self.additional_generation_specs(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            unit_test_max_new_tokens=unit_test_max_new_tokens,
        )
        unit_tests: Dict[str, AdditionalGeneration] = {}
        for spec in specs:
            full_text, completion, extracted_output = self.generate_completion(
                prompt=spec["prompt"],
                max_new_tokens=spec["max_new_tokens"],
                stop_condition=unit_test_assert_limit_reached,
                postprocess_completion=cap_unit_test_completion_to_assert_limit,
            )
            extracted_output = cap_unit_test_source_to_assert_limit(extracted_output)
            completion = cap_unit_test_completion_to_assert_limit(completion)
            full_text = str(spec["prompt"]) + completion
            unit_tests[str(spec["patch_role"])] = AdditionalGeneration(
                task=str(spec["task"]),
                patch_role=str(spec["patch_role"]),
                prompt_role=str(spec["prompt_role"]),
                max_new_tokens=int(spec["max_new_tokens"]),
                prompt=str(spec["prompt"]),
                full_text=full_text,
                completion=completion,
                extracted_output=extracted_output,
            )
        return unit_tests

    def build_additional_generations(
        self,
        pair_name: str,
        target_answer: str,
        destination_reference_output: Optional[str],
        clean_expected_output: Optional[str],
        corrupted_expected_output: Optional[str],
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
    ) -> AdditionalGenerations:
        (
            clean_answer_generation,
            corrupted_answer_generation,
            clean_model_answer,
            corrupted_model_answer,
        ) = self.build_answer_generations(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            answer_max_new_tokens=answer_max_new_tokens,
        )
        unit_tests = self.build_unit_test_generations(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            unit_test_max_new_tokens=unit_test_max_new_tokens,
        )

        return AdditionalGenerations(
            pair_name=pair_name,
            source_role=clean_variant,
            destination_role=corrupted_variant,
            source_reference_output=target_answer,
            destination_reference_output=destination_reference_output,
            source_execution_output=clean_expected_output,
            destination_execution_output=corrupted_expected_output,
            source_model_output=clean_model_answer,
            destination_model_output=corrupted_model_answer,
            source_answer_generation=clean_answer_generation,
            destination_answer_generation=corrupted_answer_generation,
            unit_tests=unit_tests,
        )

    def write_additional_generation_dump(
        self,
        output_path: Path,
        generations: AdditionalGenerations,
    ) -> Path:
        output_path.parent.mkdir(parents=True, exist_ok=True)

        lines = [
            "# Additional task generations",
            "",
            f"Pair: {generations.pair_name}",
            f"Patch direction: {generations.source_role} -> {generations.destination_role}",
            f"Source prompt role: {generations.source_role}",
            f"Destination prompt role: {generations.destination_role}",
            f"Source reference output: {generations.source_reference_output!r}",
            f"Destination reference output: {generations.destination_reference_output!r}",
            f"Source execution output: {generations.source_execution_output!r}",
            f"Destination execution output: {generations.destination_execution_output!r}",
            f"Source model output: {generations.source_model_output!r}",
            f"Destination model output: {generations.destination_model_output!r}",
            "",
        ]
        for generation in generations.unit_tests.values():
            lines.extend(
                [
                    "=" * 88,
                    f"Task: {generation.task}",
                    f"Patch role: {generation.patch_role}",
                    f"Prompt role: {generation.prompt_role}",
                    f"Max new tokens: {generation.max_new_tokens}",
                    "-" * 88,
                    "Prompt:",
                    generation.prompt.rstrip(),
                    "-" * 88,
                    "Extracted output:",
                    generation.extracted_output.rstrip() or "<empty>",
                    "-" * 88,
                    "",
                ]
            )

        output_path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
        return output_path

    def save_additional_generation_dump(
        self,
        output_path: Path,
        pair_name: str,
        target_answer: str,
        destination_reference_output: Optional[str],
        clean_expected_output: Optional[str],
        corrupted_expected_output: Optional[str],
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
    ) -> Path:
        generations = self.build_additional_generations(
            pair_name=pair_name,
            target_answer=target_answer,
            destination_reference_output=destination_reference_output,
            clean_expected_output=clean_expected_output,
            corrupted_expected_output=corrupted_expected_output,
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            answer_max_new_tokens=answer_max_new_tokens,
            unit_test_max_new_tokens=unit_test_max_new_tokens,
        )
        return self.write_additional_generation_dump(output_path, generations)

    def resolve_block_hook_name(self, layer: int, hook_name: Optional[str] = None) -> str:
        hook_name = hook_name or self.hook_name
        if hook_name.startswith("blocks."):
            if hook_name in self.model.hook_dict:
                return hook_name
            raise KeyError(f"Hook not found: {hook_name}")

        base_names = [hook_name]
        if hook_name.startswith("hook_"):
            base_names.append(hook_name.removeprefix("hook_"))
        else:
            base_names.append(f"hook_{hook_name}")

        for base in base_names:
            full_name = f"blocks.{layer}.{base}"
            if full_name in self.model.hook_dict:
                return full_name

        available = [name for name in self.model.hook_dict.keys() if name.startswith(f"blocks.{layer}.")]
        raise KeyError(
            f"Could not resolve hook {hook_name!r} for layer {layer}. "
            f"Example available hooks: {available[:10]}"
        )

    @staticmethod
    def first_changed_token_pos(clean_tokens, corrupted_tokens, aligned_seq_len: int) -> Optional[int]:
        clean_ids = clean_tokens[0, :aligned_seq_len].detach().cpu().tolist()
        corrupted_ids = corrupted_tokens[0, :aligned_seq_len].detach().cpu().tolist()
        for pos, (clean_id, corrupted_id) in enumerate(zip(clean_ids, corrupted_ids)):
            if int(clean_id) != int(corrupted_id):
                return int(pos)
        return None

    def residual_positions_from_first_change(
        self,
        clean_tokens,
        corrupted_tokens,
        aligned_seq_len: int,
        context_window: int = 0,
        max_positions: Optional[int] = None,
    ) -> Tuple[List[int], Optional[int]]:
        first_changed = self.first_changed_token_pos(clean_tokens, corrupted_tokens, aligned_seq_len)
        if first_changed is None:
            return [], None

        start = max(0, first_changed - max(int(context_window), 0))
        positions = list(range(start, aligned_seq_len))
        return sample_evenly(positions, max_positions), first_changed

    @staticmethod
    def changed_token_spans(
        clean_tokens,
        corrupted_tokens,
        positions: Sequence[int],
    ) -> List[Dict[str, Any]]:
        spans: List[Dict[str, Any]] = []
        current: List[int] = []
        clean_ids = clean_tokens[0].detach().cpu().tolist()
        corrupted_ids = corrupted_tokens[0].detach().cpu().tolist()

        for pos in sorted(int(p) for p in positions):
            if pos >= len(clean_ids) or pos >= len(corrupted_ids):
                continue
            is_changed = int(clean_ids[pos]) != int(corrupted_ids[pos])
            if is_changed and (not current or pos == current[-1] + 1):
                current.append(pos)
                continue
            if current:
                spans.append(
                    {
                        "span_id": len(spans),
                        "positions": list(current),
                        "start_pos": int(current[0]),
                        "end_pos": int(current[-1]),
                        "span_len": int(len(current)),
                    }
                )
                current = []
            if is_changed:
                current.append(pos)

        if current:
            spans.append(
                {
                    "span_id": len(spans),
                    "positions": list(current),
                    "start_pos": int(current[0]),
                    "end_pos": int(current[-1]),
                    "span_len": int(len(current)),
                }
            )
        return spans

    @staticmethod
    def enrich_span_metadata(
        spans: Sequence[Dict[str, Any]],
        clean_strs: Sequence[str],
        corrupted_strs: Sequence[str],
    ) -> List[Dict[str, Any]]:
        enriched: List[Dict[str, Any]] = []
        for span in spans:
            positions = [int(pos) for pos in span["positions"]]
            clean_tokens = [str(clean_strs[pos]) for pos in positions if pos < len(clean_strs)]
            corrupted_tokens = [
                str(corrupted_strs[pos]) for pos in positions if pos < len(corrupted_strs)
            ]
            enriched.append(
                {
                    **span,
                    "positions": positions,
                    "clean_tokens": clean_tokens,
                    "corrupted_tokens": corrupted_tokens,
                    "clean_text": "".join(clean_tokens),
                    "corrupted_text": "".join(corrupted_tokens),
                }
            )
        return enriched

    def patch_resid_and_score(
        self,
        layer: int,
        pos: int,
        source_cache,
        destination_tokens,
        answer: str,
    ) -> float:
        name = self.resolve_block_hook_name(layer)

        def patch_hook(act, hook):
            act[:, pos, :] = source_cache[name][:, pos, :]
            return act

        with torch.no_grad():
            patched_logits = self.model.run_with_hooks(
                destination_tokens,
                fwd_hooks=[(name, patch_hook)],
            )
        return self.metric_from_logits_with_target_answer(patched_logits, answer)

    def patch_resid_positions_and_score(
        self,
        layer: int,
        positions: Sequence[int],
        source_cache,
        destination_tokens,
        answer: str,
    ) -> float:
        pos_list = [int(pos) for pos in positions]
        if not pos_list:
            return float("nan")

        name = self.resolve_block_hook_name(layer)

        def patch_hook(act, hook):
            act[:, pos_list, :] = source_cache[name][:, pos_list, :]
            return act

        with torch.no_grad():
            patched_logits = self.model.run_with_hooks(
                destination_tokens,
                fwd_hooks=[(name, patch_hook)],
            )
        return self.metric_from_logits_with_target_answer(patched_logits, answer)

    def patch_resid_unit_and_score(
        self,
        layer: int,
        source_positions: Sequence[int],
        destination_positions: Sequence[int],
        source_cache,
        destination_tokens,
        answer: str,
    ) -> float:
        source_pos_list = [int(pos) for pos in source_positions]
        destination_pos_list = [int(pos) for pos in destination_positions]
        if not source_pos_list or not destination_pos_list:
            return float("nan")
        if len(source_pos_list) != len(destination_pos_list):
            raise ValueError(
                "Patch unit source and destination lengths differ: "
                f"{source_pos_list} -> {destination_pos_list}"
            )

        name = self.resolve_block_hook_name(layer)

        def patch_hook(act, hook):
            act[:, destination_pos_list, :] = source_cache[name][:, source_pos_list, :]
            return act

        with torch.no_grad():
            patched_logits = self.model.run_with_hooks(
                destination_tokens,
                fwd_hooks=[(name, patch_hook)],
            )
        return self.metric_from_logits_with_target_answer(patched_logits, answer)

    def residual_patch_unit_sweep(
        self,
        source_cache,
        destination_tokens,
        answer: str,
        patch_units: Sequence[PatchUnit],
        direction: str,
        desc: str = "Residual layers",
    ) -> torch.Tensor:
        max_unit_id = max((int(unit.unit_id) for unit in patch_units), default=-1)
        scores = torch.full((self.n_layers, max_unit_id + 1), float("nan"), dtype=torch.float32)
        for layer in tqdm(range(self.n_layers), desc=desc, leave=False):
            for unit in patch_units:
                scores[layer, int(unit.unit_id)] = self.patch_resid_unit_and_score(
                    layer=layer,
                    source_positions=unit_source_positions(unit, direction),
                    destination_positions=unit_destination_positions(unit, direction),
                    source_cache=source_cache,
                    destination_tokens=destination_tokens,
                    answer=answer,
                )
        return scores

    def continuation_token_ids_after_context(
        self,
        context: str,
        continuation: str,
    ) -> Tuple[List[int], Optional[str]]:
        context_ids = self.to_tokens(context)[0].detach().cpu().tolist()
        full_ids = self.to_tokens(context + continuation)[0].detach().cpu().tolist()
        if full_ids[: len(context_ids)] != context_ids:
            return [], "continuation_tokenization_crossed_context_boundary"
        return [int(token_id) for token_id in full_ids[len(context_ids) :]], None

    def target_boundary_candidates(
        self,
        context: str,
        positive_text: str,
        negative_text: str,
    ) -> List[Tuple[str, str, str, str]]:
        candidates = [(context, positive_text, negative_text, "")]
        trailing_match = re.search(r"\s+$", context)
        if trailing_match:
            moved_suffix = trailing_match.group(0)
            candidates.append(
                (
                    context[: trailing_match.start()],
                    moved_suffix + positive_text,
                    moved_suffix + negative_text,
                    moved_suffix,
                )
            )
        return candidates

    @staticmethod
    def move_context_suffix_into_target(context: str, moved_suffix: str) -> str:
        if moved_suffix and context.endswith(moved_suffix):
            return context[: -len(moved_suffix)]
        return context

    def decision_target_info(
        self,
        context: str,
        positive_text: str,
        negative_text: str,
    ) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
        if positive_text == negative_text:
            return None, "positive_and_negative_text_identical"

        last_error = "unknown_target_tokenization_error"
        for terminator in ("", "\n"):
            tokenized_positive_text = positive_text + terminator
            tokenized_negative_text = negative_text + terminator
            for candidate_context, candidate_positive, candidate_negative, moved_suffix in (
                self.target_boundary_candidates(
                    context,
                    tokenized_positive_text,
                    tokenized_negative_text,
                )
            ):
                positive_ids, positive_error = self.continuation_token_ids_after_context(
                    candidate_context,
                    candidate_positive,
                )
                if positive_error is not None:
                    last_error = f"positive_{positive_error}"
                    continue
                negative_ids, negative_error = self.continuation_token_ids_after_context(
                    candidate_context,
                    candidate_negative,
                )
                if negative_error is not None:
                    last_error = f"negative_{negative_error}"
                    continue
                divergent = first_divergent_token_ids(positive_ids, negative_ids)
                if divergent is None:
                    last_error = "no_first_divergent_token"
                    continue
                break
            else:
                continue
            break
        else:
            return None, last_error

        divergent_index, positive_token_id, negative_token_id, common_token_ids = divergent
        common_text = self.model.to_string(common_token_ids) if common_token_ids else ""
        return (
            {
                "divergent_token_index": int(divergent_index),
                "positive_token_id": int(positive_token_id),
                "negative_token_id": int(negative_token_id),
                "positive_token_str": self.model.to_string([int(positive_token_id)]),
                "negative_token_str": self.model.to_string([int(negative_token_id)]),
                "common_token_ids": [int(token_id) for token_id in common_token_ids],
                "common_text": common_text,
                "context_text": candidate_context,
                "context_boundary_moved_text": moved_suffix,
                "tokenized_positive_text": candidate_positive,
                "tokenized_negative_text": candidate_negative,
                "answer_terminator_text": terminator,
                "target_selection_mode": "first_divergent_token"
                if not terminator
                else "first_divergent_token_with_answer_terminator",
            },
            None,
        )

    def answer_decision_target_info(
        self,
        source_context: str,
        destination_context: str,
        positive_text: str,
        negative_text: str,
        direction_label: str,
    ) -> Dict[str, Any]:
        target_info, target_error = self.decision_target_info(
            destination_context,
            positive_text=positive_text,
            negative_text=negative_text,
        )
        if target_error is not None or target_info is None:
            raise ValueError(
                "Answer logit-margin target selection failed for "
                f"{direction_label}: {target_error or 'unknown_target_tokenization_error'}"
            )

        common_text = str(target_info["common_text"])
        moved_suffix = str(target_info.get("context_boundary_moved_text") or "")
        source_context_for_scoring = self.move_context_suffix_into_target(
            source_context,
            moved_suffix,
        )
        destination_context_for_scoring = str(target_info.get("context_text") or destination_context)
        _source_common_ids, source_common_error = self.continuation_token_ids_after_context(
            source_context_for_scoring,
            common_text,
        )
        if source_common_error is not None:
            raise ValueError(
                "Answer logit-margin target selection failed for "
                f"{direction_label}: source_common_prefix_{source_common_error}"
            )
        _destination_common_ids, destination_common_error = (
            self.continuation_token_ids_after_context(
                destination_context_for_scoring,
                common_text,
            )
        )
        if destination_common_error is not None:
            raise ValueError(
                "Answer logit-margin target selection failed for "
                f"{direction_label}: destination_common_prefix_{destination_common_error}"
            )
        target_info["source_context_text"] = source_context_for_scoring
        target_info["destination_context_text"] = destination_context_for_scoring
        return target_info

    @staticmethod
    def metric_from_logits_token_margin(
        logits,
        positive_token_id: int,
        negative_token_id: int,
    ) -> float:
        return float(
            logits[0, -1, int(positive_token_id)].item()
            - logits[0, -1, int(negative_token_id)].item()
        )

    def patch_resid_unit_and_score_token_margin(
        self,
        layer: int,
        source_positions: Sequence[int],
        destination_positions: Sequence[int],
        source_cache,
        destination_tokens,
        positive_token_id: int,
        negative_token_id: int,
    ) -> float:
        source_pos_list = [int(pos) for pos in source_positions]
        destination_pos_list = [int(pos) for pos in destination_positions]
        if not source_pos_list or not destination_pos_list:
            return float("nan")
        if len(source_pos_list) != len(destination_pos_list):
            raise ValueError(
                "Patch unit source and destination lengths differ: "
                f"{source_pos_list} -> {destination_pos_list}"
            )

        name = self.resolve_block_hook_name(layer)

        def patch_hook(act, hook):
            act[:, destination_pos_list, :] = source_cache[name][:, source_pos_list, :]
            return act

        with torch.no_grad():
            patched_logits = self.model.run_with_hooks(
                destination_tokens,
                fwd_hooks=[(name, patch_hook)],
            )
        return self.metric_from_logits_token_margin(
            patched_logits,
            positive_token_id=positive_token_id,
            negative_token_id=negative_token_id,
        )

    def patch_resid_unit_batch_and_score_token_margin(
        self,
        layer: int,
        units: Sequence[PatchUnit],
        direction: str,
        source_cache,
        destination_tokens,
        positive_token_id: int,
        negative_token_id: int,
    ):
        if not units:
            return torch.empty(0, dtype=torch.float32)

        patch_specs = []
        for unit in units:
            source_pos_list = unit_source_positions(unit, direction)
            destination_pos_list = unit_destination_positions(unit, direction)
            if not source_pos_list or not destination_pos_list:
                patch_specs.append(None)
                continue
            if len(source_pos_list) != len(destination_pos_list):
                raise ValueError(
                    "Patch unit source and destination lengths differ: "
                    f"{source_pos_list} -> {destination_pos_list}"
                )
            patch_specs.append(
                (
                    [int(pos) for pos in source_pos_list],
                    [int(pos) for pos in destination_pos_list],
                )
            )

        name = self.resolve_block_hook_name(layer)
        batch_size = len(units)
        batched_destination_tokens = destination_tokens.repeat(batch_size, 1)

        def patch_hook(act, hook):
            for batch_index, spec in enumerate(patch_specs):
                if spec is None:
                    continue
                source_pos_list, destination_pos_list = spec
                act[batch_index, destination_pos_list, :] = source_cache[name][
                    0, source_pos_list, :
                ]
            return act

        with torch.no_grad():
            patched_logits = self.model.run_with_hooks(
                batched_destination_tokens,
                fwd_hooks=[(name, patch_hook)],
            )
        margins = (
            patched_logits[:, -1, int(positive_token_id)]
            - patched_logits[:, -1, int(negative_token_id)]
        )
        margins = margins.detach().float().cpu()
        for batch_index, spec in enumerate(patch_specs):
            if spec is None:
                margins[batch_index] = float("nan")
        return margins

    def residual_patch_unit_margin_sweep(
        self,
        source_cache,
        destination_tokens,
        positive_token_id: int,
        negative_token_id: int,
        patch_units: Sequence[PatchUnit],
        direction: str = "forward",
        desc: str = "Residual decision patch units",
    ) -> torch.Tensor:
        max_unit_id = max((int(unit.unit_id) for unit in patch_units), default=-1)
        scores = torch.full((self.n_layers, max_unit_id + 1), float("nan"), dtype=torch.float32)
        for layer in tqdm(range(self.n_layers), desc=desc, leave=False):
            for start in range(0, len(patch_units), self.patch_batch_size):
                unit_batch = list(patch_units[start : start + self.patch_batch_size])
                batch_scores = self.patch_resid_unit_batch_and_score_token_margin(
                    layer=layer,
                    units=unit_batch,
                    direction=direction,
                    source_cache=source_cache,
                    destination_tokens=destination_tokens,
                    positive_token_id=positive_token_id,
                    negative_token_id=negative_token_id,
                )
                for unit, value in zip(unit_batch, batch_scores.tolist()):
                    scores[layer, int(unit.unit_id)] = float(value)
        return scores


    def build_unit_test_flow_score_tables(
        self,
        scores: torch.Tensor,
        patch_units: Sequence[PatchUnit],
        baseline: float,
        source_ref: float,
        metadata: Dict[str, Any],
    ) -> Tuple:
        score_df, top_sites_df = self.build_residual_score_tables(
            scores=scores,
            patch_units=patch_units,
            direction="forward",
            baseline=baseline,
            source_ref=source_ref,
            distance_column="distance_to_source_ref",
        )
        for df in (score_df, top_sites_df):
            for key, value in reversed(list(metadata.items())):
                df.insert(0, key, value)
            if "clean_positions" in df.columns:
                df["source_positions"] = df["clean_positions"]
            if "corrupted_positions" in df.columns:
                df["destination_positions"] = df["corrupted_positions"]
            if "clean_text" in df.columns:
                df["source_text"] = df["clean_text"]
            if "corrupted_text" in df.columns:
                df["destination_text"] = df["corrupted_text"]
        return score_df, top_sites_df

    @staticmethod
    def unit_test_combined_axis_label(row: Dict[str, Any]) -> str:
        text = next(
            (
                value
                for value in (
                    row.get("destination_text"),
                    row.get("corrupted_text"),
                    row.get("clean_text"),
                    row.get("source_text"),
                    row.get("positive_text"),
                )
                if value not in ("", None)
            ),
            "",
        )
        return format_span_for_axis(str(text or ""), max_chars=12)

    @staticmethod
    def unit_test_combined_axis_color(row: Dict[str, Any]) -> str:
        source_family = str(
            row.get("derived_from_family") or row.get("family") or ""
        )
        if source_family == "prompt_to_output":
            return "#1f5f99"
        return "#a23b3b"

    @staticmethod
    def unit_test_combined_unit_sort_key(row: Dict[str, Any]) -> Tuple[int, int, int]:
        source_family = str(
            row.get("derived_from_family") or row.get("family") or ""
        )
        family_order = 0 if source_family == "prompt_to_output" else 1
        try:
            positions = json_positions_from_row(row, "positions")
        except Exception:
            positions = []
        if positions:
            pos = int(min(positions))
        else:
            pos = safe_int_cell(row.get("pos"), safe_int_cell(row.get("source_unit_id"), 0))
        return (
            family_order,
            pos,
            safe_int_cell(row.get("source_unit_id"), safe_int_cell(row.get("unit_id"), 0)),
        )

    def plot_unit_test_combined_context_heatmaps(
        self,
        combined_scores_df,
        flow_dir: Path,
        token_label_recovery_threshold: float,
    ) -> int:
        if combined_scores_df.empty:
            return 0

        plot_count = 0
        for (decision_id, flow_direction), group in combined_scores_df.groupby(
            ["decision_id", "flow_direction"]
        ):
            records = group.to_dict("records")
            if not records:
                continue

            unit_rows: Dict[int, Dict[str, Any]] = {}
            for row in records:
                unit_id = safe_int_cell(row.get("unit_id"))
                if unit_id not in unit_rows:
                    unit_rows[unit_id] = row
            ordered_unit_ids = [
                unit_id
                for unit_id, _row in sorted(
                    unit_rows.items(),
                    key=lambda item: self.unit_test_combined_unit_sort_key(item[1]),
                )
            ]
            if not ordered_unit_ids:
                continue

            layers = [
                safe_int_cell(row.get("layer"))
                for row in records
                if row.get("layer") not in ("", None)
            ]
            if not layers:
                continue
            matrix_rows = max(layers) + 1
            matrix_cols = len(ordered_unit_ids)
            score_matrix = torch.full((matrix_rows, matrix_cols), float("nan"))
            recovery_matrix = torch.full((matrix_rows, matrix_cols), float("nan"))
            unit_to_col = {unit_id: col for col, unit_id in enumerate(ordered_unit_ids)}
            forced_tick_positions = list(range(matrix_cols))

            for row in records:
                unit_id = safe_int_cell(row.get("unit_id"))
                if unit_id not in unit_to_col or row.get("layer") in ("", None):
                    continue
                layer = safe_int_cell(row.get("layer"))
                col = unit_to_col[unit_id]
                score = maybe_float_cell(row.get("score"))
                recovery = maybe_float_cell(row.get("distance_based_recovery"))
                if score is not None:
                    score_matrix[layer, col] = float(score)
                if recovery is not None:
                    recovery_matrix[layer, col] = float(recovery)

            label_by_col = {
                col: self.unit_test_combined_axis_label(unit_rows[unit_id])
                for unit_id, col in unit_to_col.items()
            }
            label_color_by_col = {
                col: self.unit_test_combined_axis_color(unit_rows[unit_id])
                for unit_id, col in unit_to_col.items()
            }
            highlight_cols = set()
            for row in records:
                unit_id = safe_int_cell(row.get("unit_id"))
                recovery = maybe_float_cell(row.get("distance_based_recovery"))
                if (
                    unit_id in unit_to_col
                    and recovery is not None
                    and recovery >= float(token_label_recovery_threshold)
                ):
                    highlight_cols.add(unit_to_col[unit_id])
            highlight_cols = sorted(highlight_cols)
            if highlight_cols:
                forced_tick_positions = sorted(set(forced_tick_positions + highlight_cols))

            heatmap_prefix = flow_dir / f"{decision_id}_{flow_direction}_combined_context"
            self.plot_heatmap(
                score_matrix,
                positions=list(range(matrix_cols)),
                colorbar_label="decision logit margin",
                output_path=heatmap_prefix.with_name(
                    heatmap_prefix.name + "_scores_heatmap.pdf"
                ),
                token_labels=label_by_col,
                token_label_colors=label_color_by_col,
                token_label_color_legend={
                    "code context": "#1f5f99",
                    "assertion context": "#a23b3b",
                },
                forced_tick_positions=forced_tick_positions,
                xlabel="Combined context patch unit",
                cmap=recovery_cmap_for_flow_direction(str(flow_direction)),
                prompt_structure_end_label=None,
            )
            self.plot_heatmap(
                recovery_matrix,
                positions=list(range(matrix_cols)),
                colorbar_label="decision distance-based recovery",
                output_path=heatmap_prefix.with_name(
                    heatmap_prefix.name + "_recovery_heatmap.pdf"
                ),
                token_labels=label_by_col,
                token_label_colors=label_color_by_col,
                token_label_color_legend={
                    "code context": "#1f5f99",
                    "assertion context": "#a23b3b",
                },
                forced_tick_positions=forced_tick_positions,
                xlabel="Combined context patch unit",
                vmin=0.0,
                vmax=1.0,
                cmap=recovery_cmap_for_flow_direction(str(flow_direction)),
                prompt_structure_end_label=None,
            )
            plot_count += 2
        return plot_count

    def text_patch_units(
        self,
        source_text: str,
        destination_text: str,
        context_window: int,
        max_units: Optional[int],
        from_first_change: bool = True,
        code_only: bool = False,
    ) -> Tuple[List[PatchUnit], List[PatchUnit], Optional[int], List[Dict[str, Any]]]:
        source_tokens = self.to_tokens(source_text)
        destination_tokens = self.to_tokens(destination_text)
        source_strs = self.token_strs(source_tokens)
        destination_strs = self.token_strs(destination_tokens)
        source_ids = source_tokens[0].detach().cpu().tolist()
        destination_ids = destination_tokens[0].detach().cpu().tolist()
        units, alignment_rows, first_changed_pos = build_alignment_patch_units(
            clean_ids=source_ids,
            corrupted_ids=destination_ids,
            clean_strs=source_strs,
            corrupted_strs=destination_strs,
        )
        patchable_units = (
            patch_units_from_first_change(
                units,
                first_changed_pos=first_changed_pos,
                context_window=context_window,
            )
            if from_first_change
            else list(units)
        )
        if code_only:
            source_code_positions = token_positions_overlapping_span(
                source_text,
                source_strs,
                fenced_code_body_span(source_text),
            )
            destination_code_positions = token_positions_overlapping_span(
                destination_text,
                destination_strs,
                fenced_code_body_span(destination_text),
            )
            patchable_units = filter_patch_units_to_position_sets(
                patchable_units,
                source_code_positions,
                destination_code_positions,
            )
        sampled_units = sample_patch_units(patchable_units, max_units)
        return sampled_units, patchable_units, first_changed_pos, alignment_rows

    def run_unit_test_flow_patching(
        self,
        pair_output_dir: Path,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        pair_id: Optional[str],
        pair_family: str,
        additional_generations: AdditionalGenerations,
        max_positions: Optional[int],
        pos_context: int,
        token_label_recovery_threshold: float,
        flow_dir_override: Optional[Path] = None,
        clean_schema: bool = False,
    ) -> Dict[str, Any]:
        flow_dir = flow_dir_override or (pair_output_dir / "unit_test_flow")
        flow_dir.mkdir(parents=True, exist_ok=True)
        for stale_heatmap in flow_dir.glob("*_heatmap.pdf"):
            stale_heatmap.unlink(missing_ok=True)

        clean_generation = additional_generations.unit_tests.get(
            "source",
            additional_generations.unit_tests.get("clean"),
        )
        corrupted_generation = additional_generations.unit_tests.get(
            "destination",
            additional_generations.unit_tests.get("corrupted"),
        )
        if clean_generation is None or corrupted_generation is None:
            summary = {
                "status": "skipped_missing_unit_test_generations",
                "parsed_assertions": 0,
                "analyzed_forward_decisions": 0,
                "analyzed_reverse_decisions": 0,
                "skipped_decisions": 1,
                "output_dir": str(flow_dir),
            }
            summary_name = "summary.json" if clean_schema else "unit_test_flow_summary.json"
            (flow_dir / summary_name).write_text(
                json.dumps(summary, indent=2) + "\n",
                encoding="utf-8",
            )
            return summary

        use_cue_oracle_targets = str(pair_family) == "cue_varied"
        cue_oracle_namespace = None
        cue_oracle_status = "not_required"
        if use_cue_oracle_targets:
            cue_oracle_namespace, cue_oracle_status = cue_oracle_namespace_for_pair(
                pair_id or additional_generations.pair_name
            )

        decisions, skip_rows, assertion_counts = extract_unit_test_decisions(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_test_code=clean_generation.extracted_output,
            corrupted_test_code=corrupted_generation.extracted_output,
            cue_oracle_namespace=cue_oracle_namespace,
            require_cue_oracle_targets=use_cue_oracle_targets,
        )
        decision_rows: List[Dict[str, Any]] = []
        prompt_score_frames = []
        prompt_top_frames = []
        response_score_frames = []
        response_top_frames = []
        context_token_rows: List[Dict[str, Any]] = []
        context_presentations: List[Dict[str, Any]] = []
        analyzed_forward_decision_ids = set()
        analyzed_reverse_decision_ids = set()
        forward_direction_name = f"{clean_variant}_to_{corrupted_variant}"
        reverse_direction_name = f"{corrupted_variant}_to_{clean_variant}"

        def unit_test_response_fields(
            *,
            candidate: str,
            execution_value: str,
            cue_value: str,
        ) -> Dict[str, Any]:
            is_accurate = str(candidate) == str(execution_value)
            if pair_family == "cue_varied":
                if str(cue_value) != "":
                    status = (
                        "classified_with_cue_oracle"
                        if str(cue_value) != str(execution_value)
                        else "classified"
                    )
                    return {
                        "is_accurate": bool(is_accurate),
                        "response_label": classify_behavior(
                            candidate,
                            execution_value,
                            cue_value,
                        ),
                        "response_label_status": status,
                    }
                return {
                    "is_accurate": bool(is_accurate),
                    "response_label": "",
                    "response_label_status": "skipped_no_independent_cue_expected_value",
                }
            return {
                "is_accurate": bool(is_accurate),
                "response_label": classify_behavior(candidate, execution_value, cue_value),
                "response_label_status": "classified",
            }

        def source_role_for_direction(flow_direction: str) -> str:
            return clean_variant if flow_direction == forward_direction_name else corrupted_variant

        def destination_role_for_direction(flow_direction: str) -> str:
            return corrupted_variant if flow_direction == forward_direction_name else clean_variant

        def decision_prefix(test_code: str, assertion: Optional[Dict[str, Any]]) -> str:
            if assertion is None:
                return ""
            return test_code[: int(assertion["expected_start"])]

        def global_top(df):
            if df.empty:
                return df
            return (
                df.sort_values(
                    ["distance_reduction_gain", "distance_to_source_ref"],
                    ascending=[False, True],
                )
                .head(50)
                .reset_index(drop=True)
            )

        def compact_best(top_df) -> Dict[str, Any]:
            if top_df.empty:
                return {}
            row = top_df.iloc[0].to_dict()
            out: Dict[str, Any] = {}
            for key in (
                "decision_id",
                "patch_direction",
                "family",
                "layer",
                "patch_unit_id",
                "patch_unit_type",
                "positions",
                "positive_text",
                "negative_text",
                "patched_margin",
                "distance_recovery",
                "distance_reduction_gain",
            ):
                value = row.get(key)
                if value is None and key == "patch_unit_id":
                    value = row.get("unit_id")
                if value is None and key == "patch_unit_type":
                    value = row.get("unit_kind")
                if value is None and key == "patched_margin":
                    value = row.get("score")
                if value is None and key == "distance_recovery":
                    value = row.get("distance_based_recovery")
                if value is None:
                    continue
                if hasattr(value, "item"):
                    value = value.item()
                out[key] = value
            return out

        def append_skip(
            decision: Dict[str, Any],
            flow_direction: str,
            family: str,
            stage: str,
            reason: str,
            error: str = "",
        ) -> None:
            skip_rows.append(
                {
                    "decision_id": decision.get("decision_id", ""),
                    "assertion_index": decision.get("assertion_index", ""),
                    "flow_direction": flow_direction,
                    "patch_direction": flow_direction,
                    "source_role": source_role_for_direction(flow_direction),
                    "destination_role": destination_role_for_direction(flow_direction),
                    "family": family,
                    "stage": stage,
                    "reason": reason,
                    "error": error,
                }
            )

        def run_family(
            decision: Dict[str, Any],
            flow_direction: str,
            family: str,
            source_generation: AdditionalGeneration,
            destination_generation: AdditionalGeneration,
            source_test_prefix: str,
            destination_test_prefix: str,
            positive_text: str,
            negative_text: str,
        ) -> bool:
            if family == "prompt_to_output":
                source_base_context = source_generation.prompt + destination_test_prefix
                destination_base_context = destination_generation.prompt + destination_test_prefix
            elif family == "response_carry":
                source_base_context = source_generation.prompt + source_test_prefix
                destination_base_context = (
                    destination_generation.prompt + destination_test_prefix
                )
            else:
                raise ValueError(f"Unknown unit-test flow family: {family!r}")

            target_info, target_error = self.decision_target_info(
                destination_base_context,
                positive_text=positive_text,
                negative_text=negative_text,
            )
            if target_error is not None or target_info is None:
                append_skip(
                    decision,
                    flow_direction,
                    family,
                    "target_tokenization",
                    target_error or "unknown_target_tokenization_error",
                )
                return False

            moved_suffix = str(target_info.get("context_boundary_moved_text") or "")
            source_base_context = self.move_context_suffix_into_target(
                source_base_context,
                moved_suffix,
            )
            destination_base_context = str(
                target_info.get("context_text") or destination_base_context
            )
            destination_test_prefix_for_display = self.move_context_suffix_into_target(
                destination_test_prefix,
                moved_suffix,
            )

            common_text = str(target_info["common_text"])
            _source_common_ids, source_common_error = self.continuation_token_ids_after_context(
                source_base_context,
                common_text,
            )
            if source_common_error is not None:
                append_skip(
                    decision,
                    flow_direction,
                    family,
                    "target_tokenization",
                    f"source_common_prefix_{source_common_error}",
                )
                return False
            _destination_common_ids, destination_common_error = (
                self.continuation_token_ids_after_context(
                    destination_base_context,
                    common_text,
                )
            )
            if destination_common_error is not None:
                append_skip(
                    decision,
                    flow_direction,
                    family,
                    "target_tokenization",
                    f"destination_common_prefix_{destination_common_error}",
                )
                return False

            if family == "prompt_to_output":
                sampled_units, patchable_units, _first_changed, _alignment_rows = self.text_patch_units(
                    source_text=source_generation.prompt,
                    destination_text=destination_generation.prompt,
                    context_window=pos_context,
                    max_units=max_positions,
                    from_first_change=True,
                    code_only=True,
                )
                asymmetric_units = [
                    unit
                    for unit in changed_region_units(patchable_units)
                    if unit.is_asymmetric
                ]
                if asymmetric_units:
                    append_skip(
                        decision,
                        flow_direction,
                        family,
                        "token_alignment",
                        "asymmetric_prompt_changed_region",
                        json.dumps(changed_region_summary(asymmetric_units)),
                    )
                    return False
            elif family == "response_carry":
                source_prompt_len = int(self.to_tokens(source_generation.prompt).shape[1])
                destination_prompt_len = int(
                    self.to_tokens(destination_generation.prompt).shape[1]
                )
                _sampled_units, all_units, _first_changed, _alignment_rows = self.text_patch_units(
                    source_text=source_base_context,
                    destination_text=destination_base_context,
                    context_window=0,
                    max_units=None,
                    from_first_change=False,
                )
                response_units = []
                asymmetric_units = []
                for unit in all_units:
                    if not unit.clean_positions or not unit.corrupted_positions:
                        continue
                    source_in_response = min(unit.clean_positions) >= source_prompt_len
                    destination_in_response = (
                        min(unit.corrupted_positions) >= destination_prompt_len
                    )
                    if not source_in_response or not destination_in_response:
                        continue
                    if unit.is_asymmetric:
                        asymmetric_units.append(unit)
                    else:
                        response_units.append(unit)
                if asymmetric_units:
                    append_skip(
                        decision,
                        flow_direction,
                        family,
                        "token_alignment",
                        "asymmetric_response_prefix_region",
                        json.dumps(changed_region_summary(asymmetric_units)),
                    )
                    return False
                sampled_units = sample_patch_units(response_units, max_positions)

            if not sampled_units:
                append_skip(
                    decision,
                    flow_direction,
                    family,
                    "token_alignment",
                    "no_patchable_units",
                )
                return False
            source_context = source_base_context + common_text
            destination_context = destination_base_context + common_text
            source_tokens = self.to_tokens(source_context)
            destination_tokens = self.to_tokens(destination_context)
            destination_prompt_token_count = int(
                self.to_tokens(destination_generation.prompt).shape[1]
            )
            destination_base_token_count = int(self.to_tokens(destination_base_context).shape[1])
            destination_token_ids = destination_tokens[0].detach().cpu().tolist()
            destination_token_strs = self.token_strs(destination_tokens)
            destination_prompt_parts = split_fenced_prompt(destination_generation.prompt)
            destination_snippet_span = fenced_code_body_span(destination_generation.prompt)
            destination_prompt_char_count = len(destination_generation.prompt)
            destination_base_char_count = len(destination_base_context)
            destination_snippet_positions = set(
                token_positions_overlapping_span(
                    destination_generation.prompt,
                    self.token_strs(self.to_tokens(destination_generation.prompt)),
                    fenced_code_body_span(destination_generation.prompt),
                )
            )
            context_presentations.append(
                {
                    "decision_id": decision["decision_id"],
                    "assertion_index": int(decision["assertion_index"]),
                    "flow_direction": flow_direction,
                    "patch_direction": flow_direction,
                    "family": family,
                    "source_role": source_generation.variant,
                    "destination_role": destination_generation.variant,
                    "destination_variant": destination_generation.variant,
                    "destination_generation_role": destination_generation.variant_role,
                    "prompt_role": destination_generation.variant,
                    "response_label": decision.get(
                        f"{destination_generation.variant}_response_label",
                        "",
                    ),
                    "response_label_status": decision.get(
                        f"{destination_generation.variant}_response_label_status",
                        "",
                    ),
                    "is_accurate": decision.get(
                        f"{destination_generation.variant}_is_accurate",
                        "",
                    ),
                    "text": destination_base_context + common_text,
                    "prompt_text": destination_prompt_parts["prompt_text"],
                    "snippet_text": destination_prompt_parts["snippet_text"],
                    "snippet_start": ""
                    if destination_snippet_span is None
                    else int(destination_snippet_span[0]),
                    "snippet_end": ""
                    if destination_snippet_span is None
                    else int(destination_snippet_span[1]),
                    "output_label": prompt_output_label(
                        destination_prompt_parts,
                        "Generated Answer",
                    ),
                    "generated_output": destination_generation.extracted_output,
                    "generated_context": destination_test_prefix_for_display + common_text,
                    "generated_prefix": destination_test_prefix_for_display,
                    "common_target_prefix": common_text,
                    "prompt_end": int(destination_prompt_char_count),
                    "generated_prefix_start": int(destination_prompt_char_count),
                    "generated_prefix_end": int(destination_base_char_count),
                    "common_target_prefix_start": int(destination_base_char_count),
                    "common_target_prefix_end": int(
                        destination_base_char_count + len(common_text)
                    ),
                    "positive_target": positive_text,
                    "negative_target": negative_text,
                    "target_selection_mode": target_info.get("target_selection_mode", ""),
                    "answer_terminator_text": target_info.get("answer_terminator_text", ""),
                    "positive_token_str": target_info["positive_token_str"],
                    "negative_token_str": target_info["negative_token_str"],
                    "next_token_scored": (
                        f"{target_info['positive_token_str']} vs "
                        f"{target_info['negative_token_str']}"
                    ),
                }
            )
            for token_pos, token_id in enumerate(destination_token_ids):
                if token_pos < destination_prompt_token_count:
                    segment = "snippet" if token_pos in destination_snippet_positions else "prompt"
                elif token_pos < destination_base_token_count:
                    segment = "generated_prefix"
                else:
                    segment = "common_target_prefix"
                context_token_rows.append(
                    {
                        "decision_id": decision["decision_id"],
                        "assertion_index": int(decision["assertion_index"]),
                        "flow_direction": flow_direction,
                        "patch_direction": flow_direction,
                        "family": family,
                        "source_role": source_generation.variant,
                        "destination_role": destination_generation.variant,
                        "destination_variant": destination_generation.variant,
                        "destination_generation_role": destination_generation.variant_role,
                        "prompt_role": destination_generation.variant,
                        "response_label": decision.get(
                            f"{destination_generation.variant}_response_label",
                            "",
                        ),
                        "response_label_status": decision.get(
                            f"{destination_generation.variant}_response_label_status",
                            "",
                        ),
                        "is_accurate": decision.get(
                            f"{destination_generation.variant}_is_accurate",
                            "",
                        ),
                        "token_pos": int(token_pos),
                        "token_id": int(token_id),
                        "token_str": destination_token_strs[token_pos],
                        "segment": segment,
                        "is_prompt": bool(segment == "prompt"),
                        "is_snippet": bool(segment == "snippet"),
                        "is_generated_prefix": bool(segment == "generated_prefix"),
                        "is_common_target_prefix": bool(segment == "common_target_prefix"),
                        "positive_text": positive_text,
                        "negative_text": negative_text,
                        "positive_token_id": int(target_info["positive_token_id"]),
                        "negative_token_id": int(target_info["negative_token_id"]),
                    }
                )
            with torch.no_grad():
                source_logits, source_cache = self.model.run_with_cache(source_tokens)
                destination_logits = self.model(destination_tokens)

            positive_token_id = int(target_info["positive_token_id"])
            negative_token_id = int(target_info["negative_token_id"])
            baseline = self.metric_from_logits_token_margin(
                destination_logits,
                positive_token_id=positive_token_id,
                negative_token_id=negative_token_id,
            )
            source_ref = self.metric_from_logits_token_margin(
                source_logits,
                positive_token_id=positive_token_id,
                negative_token_id=negative_token_id,
            )
            scores = self.residual_patch_unit_margin_sweep(
                source_cache=source_cache,
                destination_tokens=destination_tokens,
                positive_token_id=positive_token_id,
                negative_token_id=negative_token_id,
                patch_units=sampled_units,
                desc=f"{flow_direction} {family} decision units",
            )

            metadata = {
                "decision_id": decision["decision_id"],
                "assertion_index": int(decision["assertion_index"]),
                "patch_direction": flow_direction,
                "flow_direction": flow_direction,
                "family": family,
                "source_role": source_generation.variant,
                "destination_role": destination_generation.variant,
                "prompt_role": destination_generation.variant,
                "response_label": decision.get(
                    f"{destination_generation.variant}_response_label",
                    "",
                ),
                "response_label_status": decision.get(
                    f"{destination_generation.variant}_response_label_status",
                    "",
                ),
                "is_accurate": decision.get(
                    f"{destination_generation.variant}_is_accurate",
                    "",
                ),
                "source_output": positive_text,
                "destination_output": negative_text,
                "positive_text": positive_text,
                "negative_text": negative_text,
                "source_token_id": positive_token_id,
                "destination_token_id": negative_token_id,
                "source_token": target_info["positive_token_str"],
                "destination_token": target_info["negative_token_str"],
                "divergent_token_index": int(target_info["divergent_token_index"]),
                "common_target_prefix": common_text,
                "baseline_margin": float(baseline),
                "source_margin": float(source_ref),
                "source_ref_margin": float(source_ref),
            }
            score_df, top_sites_df = self.build_unit_test_flow_score_tables(
                scores=scores,
                patch_units=sampled_units,
                baseline=baseline,
                source_ref=source_ref,
                metadata=metadata,
            )
            if family == "prompt_to_output":
                prompt_score_frames.append(score_df)
                prompt_top_frames.append(top_sites_df)
            else:
                response_score_frames.append(score_df)
                response_top_frames.append(top_sites_df)
            if flow_direction == forward_direction_name:
                analyzed_forward_decision_ids.add(str(decision["decision_id"]))
            else:
                analyzed_reverse_decision_ids.add(str(decision["decision_id"]))
            return True

        def append_derived_combined_contexts() -> None:
            presentation_keys = {
                (
                    row.get("decision_id"),
                    row.get("flow_direction"),
                    row.get("family"),
                )
                for row in context_presentations
            }
            for row in list(context_presentations):
                if row.get("family") != "response_carry":
                    continue
                key = (row.get("decision_id"), row.get("flow_direction"), "combined_context")
                if key in presentation_keys:
                    continue
                combined_row = dict(row)
                combined_row["family"] = "combined_context"
                context_presentations.append(combined_row)
                presentation_keys.add(key)

            token_keys = {
                (
                    row.get("decision_id"),
                    row.get("flow_direction"),
                    row.get("family"),
                    row.get("token_pos"),
                )
                for row in context_token_rows
            }
            for row in list(context_token_rows):
                if row.get("family") != "response_carry":
                    continue
                key = (
                    row.get("decision_id"),
                    row.get("flow_direction"),
                    "combined_context",
                    row.get("token_pos"),
                )
                if key in token_keys:
                    continue
                combined_row = dict(row)
                combined_row["family"] = "combined_context"
                context_token_rows.append(combined_row)
                token_keys.add(key)

        def derive_combined_scores(
            prompt_scores_df,
            response_scores_df,
        ):
            if response_scores_df.empty:
                return pd.DataFrame()

            reference_by_decision: Dict[Tuple[str, str], Tuple[float, float]] = {}
            for row in response_scores_df.to_dict("records"):
                key = (str(row.get("decision_id")), str(row.get("flow_direction")))
                if key in reference_by_decision:
                    continue
                try:
                    baseline = float(row["baseline_margin"])
                    source_ref = float(row["source_ref_margin"])
                except (KeyError, TypeError, ValueError):
                    continue
                if math.isfinite(baseline) and math.isfinite(source_ref):
                    reference_by_decision[key] = (baseline, source_ref)

            if not reference_by_decision:
                return pd.DataFrame()

            combined_rows: List[Dict[str, Any]] = []

            def append_rows(source_df, source_family: str, unit_offset: int) -> None:
                if source_df.empty:
                    return
                for row in source_df.to_dict("records"):
                    key = (str(row.get("decision_id")), str(row.get("flow_direction")))
                    reference = reference_by_decision.get(key)
                    if reference is None:
                        continue
                    baseline, source_ref = reference
                    try:
                        score = float(row["score"])
                    except (KeyError, TypeError, ValueError):
                        continue
                    if not math.isfinite(score):
                        continue

                    combined_row = dict(row)
                    original_unit_id = combined_row.get("unit_id")
                    try:
                        combined_row["unit_id"] = int(original_unit_id) + int(unit_offset)
                        combined_row["source_unit_id"] = int(original_unit_id)
                    except (TypeError, ValueError):
                        combined_row["source_unit_id"] = original_unit_id
                    combined_row["family"] = "combined_context"
                    combined_row["derived_from_family"] = source_family
                    combined_row["baseline_margin"] = float(baseline)
                    combined_row["source_ref_margin"] = float(source_ref)
                    combined_row["normalized_recovery"] = normalize_recovery(
                        score,
                        baseline,
                        source_ref,
                    )
                    combined_row["distance_based_recovery"] = distance_based_recovery(
                        score,
                        baseline,
                        source_ref,
                    )
                    combined_row["distance_reduction_gain"] = patch_gain_toward_clean(
                        score,
                        baseline,
                        source_ref,
                    )
                    combined_row["distance_to_source_ref"] = abs(score - source_ref)
                    combined_rows.append(combined_row)

            append_rows(prompt_scores_df, "prompt_to_output", 0)
            append_rows(response_scores_df, "response_carry", 1_000_000)
            return pd.DataFrame(combined_rows)

        for decision in decisions:
            corrupted_assertion = decision["corrupted_assertion"]
            clean_assertion = decision["clean_assertion"]
            corrupted_prefix = decision_prefix(
                corrupted_generation.extracted_output,
                corrupted_assertion,
            )
            clean_prefix = decision_prefix(
                clean_generation.extracted_output,
                clean_assertion,
            )
            source_response = unit_test_response_fields(
                candidate=decision["clean_expected_source"],
                execution_value=decision["clean_runtime_source"],
                cue_value=decision["clean_runtime_source"],
            )
            destination_response = unit_test_response_fields(
                candidate=decision["corrupted_expected_source"],
                execution_value=decision["corrupted_runtime_source"],
                cue_value=decision.get("corrupted_cue_oracle_source")
                if use_cue_oracle_targets
                else decision["clean_runtime_source"],
            )
            decision[clean_variant + "_is_accurate"] = source_response["is_accurate"]
            decision[clean_variant + "_response_label"] = source_response["response_label"]
            decision[clean_variant + "_response_label_status"] = source_response[
                "response_label_status"
            ]
            decision[corrupted_variant + "_is_accurate"] = destination_response["is_accurate"]
            decision[corrupted_variant + "_response_label"] = destination_response[
                "response_label"
            ]
            decision[corrupted_variant + "_response_label_status"] = destination_response[
                "response_label_status"
            ]
            decision_rows.append(
                {
                    "decision_id": decision["decision_id"],
                    "assertion_index": int(decision["assertion_index"]),
                    "pair_family": pair_family,
                    "prompt_role": corrupted_variant,
                    "source_role": clean_variant,
                    "destination_role": corrupted_variant,
                    "corrupted_expr_source": decision["corrupted_expr_source"],
                    "runtime_correct_source": decision["runtime_correct_source"],
                    "corrupted_expected_source": decision["corrupted_expected_source"],
                    "clean_expected_source": decision["clean_expected_source"],
                    "clean_runtime_source": decision["clean_runtime_source"],
                    "corrupted_runtime_source": decision["corrupted_runtime_source"],
                    "clean_cue_oracle_source": decision.get("clean_cue_oracle_source", ""),
                    "corrupted_cue_oracle_source": decision.get(
                        "corrupted_cue_oracle_source",
                        "",
                    ),
                    "execution_expected_value": decision["corrupted_runtime_source"],
                    "cue_expected_value": decision.get("corrupted_cue_oracle_source", "")
                    if use_cue_oracle_targets
                    else decision["clean_runtime_source"],
                    "source_cue_expected_value": decision["clean_runtime_source"],
                    "destination_cue_expected_value": decision.get(
                        "corrupted_cue_oracle_source",
                        "",
                    )
                    if use_cue_oracle_targets
                    else decision["clean_runtime_source"],
                    "is_accurate": destination_response["is_accurate"],
                    "response_label": destination_response["response_label"],
                    "response_label_status": destination_response["response_label_status"],
                    "source_is_accurate": source_response["is_accurate"],
                    "source_response_label": source_response["response_label"],
                    "source_response_label_status": source_response[
                        "response_label_status"
                    ],
                    "destination_is_accurate": destination_response["is_accurate"],
                    "destination_response_label": destination_response["response_label"],
                    "destination_response_label_status": destination_response[
                        "response_label_status"
                    ],
                    "target_mode": decision["target_mode"],
                    "forward_positive_source": decision["forward_positive_source"],
                    "forward_positive_source_kind": decision["forward_positive_source_kind"],
                    "forward_negative_source": decision["forward_negative_source"],
                    "forward_negative_source_kind": decision["forward_negative_source_kind"],
                    "clean_match_type": decision["clean_match_type"],
                    "reverse_available": bool(decision["reverse_available"]),
                    "reverse_positive_source": decision["reverse_positive_source"],
                    "reverse_positive_source_kind": decision["reverse_positive_source_kind"],
                    "reverse_negative_source": decision["reverse_negative_source"],
                    "reverse_negative_source_kind": decision["reverse_negative_source_kind"],
                    "reverse_unavailable_reason": decision["reverse_unavailable_reason"],
                    "corrupted_expected_start": int(corrupted_assertion["expected_start"]),
                    "corrupted_expected_end": int(corrupted_assertion["expected_end"]),
                    "clean_expected_start": ""
                    if clean_assertion is None
                    else int(clean_assertion["expected_start"]),
                    "clean_expected_end": ""
                    if clean_assertion is None
                    else int(clean_assertion["expected_end"]),
                }
            )

            run_family(
                decision=decision,
                flow_direction=forward_direction_name,
                family="prompt_to_output",
                source_generation=clean_generation,
                destination_generation=corrupted_generation,
                source_test_prefix=clean_prefix,
                destination_test_prefix=corrupted_prefix,
                positive_text=decision["forward_positive_source"],
                negative_text=decision["forward_negative_source"],
            )
            if clean_assertion is not None:
                run_family(
                    decision=decision,
                    flow_direction=forward_direction_name,
                    family="response_carry",
                    source_generation=clean_generation,
                    destination_generation=corrupted_generation,
                    source_test_prefix=clean_prefix,
                    destination_test_prefix=corrupted_prefix,
                    positive_text=decision["forward_positive_source"],
                    negative_text=decision["forward_negative_source"],
                )

            if clean_assertion is not None and decision["reverse_available"]:
                run_family(
                    decision=decision,
                    flow_direction=reverse_direction_name,
                    family="prompt_to_output",
                    source_generation=corrupted_generation,
                    destination_generation=clean_generation,
                    source_test_prefix=corrupted_prefix,
                    destination_test_prefix=clean_prefix,
                    positive_text=decision["reverse_positive_source"],
                    negative_text=decision["reverse_negative_source"],
                )
                run_family(
                    decision=decision,
                    flow_direction=reverse_direction_name,
                    family="response_carry",
                    source_generation=corrupted_generation,
                    destination_generation=clean_generation,
                    source_test_prefix=corrupted_prefix,
                    destination_test_prefix=clean_prefix,
                    positive_text=decision["reverse_positive_source"],
                    negative_text=decision["reverse_negative_source"],
                )

        prompt_scores_df = (
            pd.concat(prompt_score_frames, ignore_index=True)
            if prompt_score_frames
            else pd.DataFrame()
        )
        response_scores_df = (
            pd.concat(response_score_frames, ignore_index=True)
            if response_score_frames
            else pd.DataFrame()
        )
        append_derived_combined_contexts()
        combined_scores_df = derive_combined_scores(prompt_scores_df, response_scores_df)
        prompt_top_df = global_top(prompt_scores_df)
        response_top_df = global_top(response_scores_df)
        combined_top_df = global_top(combined_scores_df)
        combined_context_plot_files = self.plot_unit_test_combined_context_heatmaps(
            combined_scores_df=combined_scores_df,
            flow_dir=flow_dir,
            token_label_recovery_threshold=token_label_recovery_threshold,
        )

        if clean_schema:
            all_score_frames = [
                frame for frame in (prompt_scores_df, response_scores_df, combined_scores_df) if not frame.empty
            ]
            all_top_frames = [
                frame for frame in (prompt_top_df, response_top_df, combined_top_df) if not frame.empty
            ]
            residual_scores_df = (
                pd.concat(all_score_frames, ignore_index=True) if all_score_frames else pd.DataFrame()
            )
            top_sites_df = (
                pd.concat(all_top_frames, ignore_index=True) if all_top_frames else pd.DataFrame()
            )
            cleaned_presentations = [
                {
                    key: value
                    for key, value in row.items()
                    if key not in {"flow_direction", "destination_variant", "prompt_role"}
                }
                for row in context_presentations
            ]
            cleaned_tokens = [
                {
                    key: value
                    for key, value in row.items()
                    if key not in {"flow_direction", "destination_variant", "prompt_role"}
                }
                for row in context_token_rows
            ]
            cleaned_decisions = [
                {key: value for key, value in row.items() if key != "prompt_role"}
                for row in decision_rows
            ]
            cleaned_skips = [
                {key: value for key, value in row.items() if key not in {"flow_direction", "prompt_role"}}
                for row in skip_rows
            ]
            write_dict_rows_csv(flow_dir / "decisions.csv", cleaned_decisions)
            write_dict_rows_csv(flow_dir / "decision_skips.csv", cleaned_skips)
            write_dict_rows_csv(flow_dir / "context_tokens.csv", cleaned_tokens)
            (flow_dir / "presentations.json").write_text(
                json.dumps(cleaned_presentations, indent=2, default=str) + "\n",
                encoding="utf-8",
            )
            write_clean_dataframe_csv(residual_scores_df, flow_dir / "residual_scores.csv")
            write_clean_dataframe_csv(top_sites_df, flow_dir / "top_sites.csv")
        else:
            write_dict_rows_csv(flow_dir / "unit_test_decisions.csv", decision_rows)
            write_dict_rows_csv(flow_dir / "unit_test_decision_skips.csv", skip_rows)
            write_dict_rows_csv(flow_dir / "unit_test_context_tokens.csv", context_token_rows)
            (flow_dir / "unit_test_presentations.json").write_text(
                json.dumps(context_presentations, indent=2, default=str) + "\n",
                encoding="utf-8",
            )
            prompt_scores_df.to_csv(flow_dir / "prompt_to_output_scores.csv", index=False)
            prompt_top_df.to_csv(flow_dir / "top_prompt_to_output_sites.csv", index=False)
            response_scores_df.to_csv(flow_dir / "response_carry_scores.csv", index=False)
            response_top_df.to_csv(flow_dir / "top_response_carry_sites.csv", index=False)
            combined_scores_df.to_csv(flow_dir / "combined_context_scores.csv", index=False)
            combined_top_df.to_csv(flow_dir / "top_combined_context_sites.csv", index=False)

        summary = {
            "status": "completed",
            "pair_family": pair_family,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "pair_id": pair_id or "",
            "cue_oracle_targets": bool(use_cue_oracle_targets),
            "cue_oracle_status": cue_oracle_status,
            "forward_patch_direction": forward_direction_name,
            "reverse_patch_direction": reverse_direction_name,
            "parsed_clean_assertions": int(assertion_counts["clean_assertions"]),
            "parsed_corrupted_assertions": int(assertion_counts["corrupted_assertions"]),
            "parsed_assertions": int(
                assertion_counts["clean_assertions"]
                + assertion_counts["corrupted_assertions"]
            ),
            "decision_rows": int(len(decision_rows)),
            "analyzed_forward_decisions": int(len(analyzed_forward_decision_ids)),
            "analyzed_reverse_decisions": int(len(analyzed_reverse_decision_ids)),
            "skipped_decisions": int(len(skip_rows)),
            "prompt_to_output_score_rows": int(len(prompt_scores_df)),
            "response_carry_score_rows": int(len(response_scores_df)),
            "combined_context_score_rows": int(len(combined_scores_df)),
            "combined_context_score_source": "derived_from_prompt_to_output_and_response_carry",
            "combined_context_plot_files": int(combined_context_plot_files),
            "patch_batch_size": int(self.patch_batch_size),
            "best_prompt_to_output_site": compact_best(prompt_top_df),
            "best_response_carry_site": compact_best(response_top_df),
            "best_combined_context_site": compact_best(combined_top_df),
            "output_dir": str(flow_dir),
        }
        summary_name = "summary.json" if clean_schema else "unit_test_flow_summary.json"
        (flow_dir / summary_name).write_text(
            json.dumps(summary, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        return summary

    def residual_span_patch_sweep(
        self,
        source_cache,
        destination_tokens,
        answer: str,
        spans: Sequence[Dict[str, Any]],
        desc: str = "Residual changed spans",
    ) -> torch.Tensor:
        scores = torch.full((self.n_layers, len(spans)), float("nan"), dtype=torch.float32)
        for layer in tqdm(range(self.n_layers), desc=desc, leave=False):
            for span_index, span in enumerate(spans):
                scores[layer, int(span_index)] = self.patch_resid_positions_and_score(
                    layer=layer,
                    positions=span["positions"],
                    source_cache=source_cache,
                    destination_tokens=destination_tokens,
                    answer=answer,
                )
        return scores

    def plot_heatmap(
        self,
        matrix: torch.Tensor,
        positions: Sequence[int],
        colorbar_label: str,
        output_path: Path,
        token_labels: Optional[Dict[int, str]] = None,
        token_label_colors: Optional[Dict[int, str]] = None,
        token_label_bold_positions: Optional[Sequence[int]] = None,
        token_label_color_legend: Optional[Dict[str, str]] = None,
        forced_tick_positions: Optional[Sequence[int]] = None,
        xlabel: str = "Token position",
        vmin: Optional[float] = None,
        vmax: Optional[float] = None,
        cmap: Optional[str] = None,
        show_prompt_structure_ticks: bool = True,
        prompt_structure_end_label: Optional[str] = None,
    ) -> None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path = output_path.with_suffix(".pdf")
        positions = [int(p) for p in positions]
        tick_cols: List[float] = []
        tick_labels: List[str] = []
        tick_colors: List[Optional[str]] = []
        tick_bold_flags: List[bool] = []
        if positions:
            pos_to_col = {pos: col for col, pos in enumerate(positions)}
            bold_position_set = {
                int(pos) for pos in (token_label_bold_positions or [])
            }

            def tick_label(col: int) -> str:
                pos = positions[col]
                if token_labels and pos in token_labels:
                    return str(token_labels[pos])
                return str(pos)

            all_tick_cols = list(range(len(positions)))
            all_tick_labels = [tick_label(col) for col in all_tick_cols]
            if self.all_heatmap_xtick_labels_fit(len(positions), all_tick_labels):
                tick_cols = all_tick_cols
            elif forced_tick_positions:
                tick_cols = [
                    pos_to_col[int(pos)]
                    for pos in forced_tick_positions
                    if int(pos) in pos_to_col
                ]
                tick_cols = sorted(set(tick_cols))
            tick_labels = [tick_label(col) for col in tick_cols]
            tick_colors = [
                str(token_label_colors[positions[int(col)]])
                if token_label_colors and positions[int(col)] in token_label_colors
                else None
                for col in tick_cols
            ]
            tick_bold_flags = [
                int(positions[int(col)]) in bold_position_set
                for col in tick_cols
            ]
            if show_prompt_structure_ticks:
                tick_cols = [-2.0, -1.0, *tick_cols]
                tick_labels = ["[Context:]", "[...]", *tick_labels]
                tick_colors = [None, None, *tick_colors]
                tick_bold_flags = [False, False, *tick_bold_flags]
                if prompt_structure_end_label:
                    tick_cols.append(float(len(positions)))
                    tick_labels.append(str(prompt_structure_end_label))
                    tick_colors.append(None)
                    tick_bold_flags.append(False)

        fig, ax = plt.subplots(
            figsize=self.heatmap_figsize(
                n_positions=len(positions),
                tick_labels=tick_labels,
            )
        )

        if len(positions) == 0:
            ax.text(
                0.5,
                0.5,
                "No patchable prompt positions after the first token difference.",
                ha="center",
                va="center",
                transform=ax.transAxes,
            )
            ax.set_axis_off()
        else:
            data = matrix[:, positions].detach().cpu().numpy()
            im = ax.imshow(
                data,
                aspect="auto",
                interpolation="nearest",
                origin="lower",
                vmin=vmin,
                vmax=vmax,
                cmap=self.resolve_heatmap_cmap(cmap),
            )
            cbar = fig.colorbar(im, ax=ax, fraction=0.04, pad=0.02)
            cbar.set_label(colorbar_label, fontsize=9, labelpad=4)
            if "signed recovery" in str(colorbar_label):
                signed_ticks = [-1.0, -0.5, 0.0, 0.5, 1.0]
                cbar.set_ticks(signed_ticks)
                cbar.set_ticklabels(["-1", "-0.5", "0", "0.5", "1"])
            cbar.ax.tick_params(labelsize=9, length=3, pad=2)
            if xlabel:
                ax.set_xlabel(xlabel, fontsize=14, labelpad=24)
            ax.set_ylabel("Layer", fontsize=14, labelpad=5)
            ax.tick_params(axis="both", labelsize=11, length=3, pad=2)
            ax.tick_params(axis="x", pad=0)
            if show_prompt_structure_ticks:
                right_margin = 0.5 if prompt_structure_end_label else -0.5
                ax.set_xlim(-2.5, len(positions) + right_margin)

            ax.set_xticks(tick_cols)
            ax.set_xticklabels(
                tick_labels,
                rotation=90,
                ha="center",
                va="top",
            )
            for tick_label in ax.get_xticklabels():
                tick_label.set_y(self.heatmap_tick_label_y(tick_labels))
            for tick_label, color, is_bold in zip(
                ax.get_xticklabels(),
                tick_colors,
                tick_bold_flags,
            ):
                if color:
                    tick_label.set_color(color)
                if color or is_bold:
                    tick_label.set_fontweight("bold")
            if xlabel:
                ax.xaxis.set_label_coords(0.5, self.heatmap_xlabel_y(tick_labels))
            if token_label_color_legend:
                from matplotlib.lines import Line2D

                handles = [
                    Line2D(
                        [0],
                        [0],
                        marker="s",
                        linestyle="",
                        markersize=5.2,
                        markerfacecolor=str(color),
                        markeredgecolor=str(color),
                        label=str(label),
                    )
                    for label, color in token_label_color_legend.items()
                ]
                ax.legend(
                    handles=handles,
                    loc="upper left",
                    bbox_to_anchor=(0.0, 1.02),
                    frameon=False,
                    fontsize=8.8,
                    ncol=len(handles),
                    handletextpad=0.3,
                    columnspacing=0.8,
                )

        fig.tight_layout(pad=0.25)
        if len(positions) != 0:
            fig.subplots_adjust(bottom=self.heatmap_bottom_margin(tick_labels))
        fig.savefig(output_path, bbox_inches="tight", pad_inches=0.02)
        plt.close(fig)

    @staticmethod
    def all_heatmap_xtick_labels_fit(
        n_positions: int,
        tick_labels: Optional[Sequence[str]] = None,
    ) -> bool:
        # With vertical labels the horizontal footprint is mostly the font height.
        longest_label = (
            max(len(str(label)) for label in tick_labels)
            if tick_labels
            else 0
        )
        if longest_label >= 24:
            return int(n_positions) <= 36
        if longest_label >= 16:
            return int(n_positions) <= 42
        return int(n_positions) <= 48

    @staticmethod
    def heatmap_figsize(n_positions: int, tick_labels: Sequence[str]) -> Tuple[float, float]:
        width = min(15.5, max(8.0, 2.6 + (0.18 * max(int(n_positions), 1))))
        if not tick_labels:
            return width, 3.2

        longest_label = max(len(str(label)) for label in tick_labels)
        height = min(9.2, max(5.0, 3.8 + (0.11 * min(longest_label, 48))))
        return width, height

    @staticmethod
    def heatmap_bottom_margin(tick_labels: Sequence[str]) -> float:
        if not tick_labels:
            return 0.24
        longest_label = max(len(str(label)) for label in tick_labels)
        return min(0.82, max(0.48, 0.34 + (0.014 * min(longest_label, 48))))

    @staticmethod
    def heatmap_xlabel_y(tick_labels: Sequence[str]) -> float:
        if not tick_labels:
            return -0.16
        longest_label = max(len(str(label)) for label in tick_labels)
        return -min(0.62, max(0.34, 0.26 + (0.01 * min(longest_label, 48))))

    @staticmethod
    def heatmap_tick_label_y(tick_labels: Sequence[str]) -> float:
        if not tick_labels:
            return -0.025
        longest_label = max(len(str(label)) for label in tick_labels)
        return -min(0.052, max(0.028, 0.02 + (0.0008 * min(longest_label, 40))))

    @staticmethod
    def resolve_heatmap_cmap(cmap: Optional[str]) -> Any:
        def with_white_bad_and_under(resolved_cmap):
            resolved_cmap.set_bad("#ffffff")
            resolved_cmap.set_under("#ffffff")
            return resolved_cmap

        if cmap == SOURCE_LOGIT_MARGIN_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    SOURCE_LOGIT_MARGIN_CMAP,
                    [
                        (0.0, "#ffffff"),
                        (0.14, "#ffffff"),
                        (1.0, "#2f74b5"),
                    ],
                )
            )
        if cmap == DESTINATION_LOGIT_MARGIN_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    DESTINATION_LOGIT_MARGIN_CMAP,
                    [
                        (0.0, "#ffffff"),
                        (0.14, "#ffffff"),
                        (1.0, "#c43d3d"),
                    ],
                )
            )
        if cmap == CLEAN_RECOVERY_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    CLEAN_RECOVERY_CMAP,
                    ["#ffffff", "#2f74b5"],
                )
            )
        if cmap == CORRUPTED_RECOVERY_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    CORRUPTED_RECOVERY_CMAP,
                    ["#ffffff", "#c43d3d"],
                )
            )
        if cmap == SIGNED_SOURCE_POSITIVE_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    SIGNED_SOURCE_POSITIVE_CMAP,
                    ["#c43d3d", "#ffffff", "#2f74b5"],
                )
            )
        if cmap == SIGNED_DESTINATION_POSITIVE_CMAP:
            import matplotlib.colors as mcolors

            return with_white_bad_and_under(
                mcolors.LinearSegmentedColormap.from_list(
                    SIGNED_DESTINATION_POSITIVE_CMAP,
                    ["#2f74b5", "#ffffff", "#c43d3d"],
                )
            )
        return cmap

    @staticmethod
    def positions_over_recovery_threshold(
        recovery_scores: torch.Tensor,
        positions: Sequence[int],
        threshold: float,
    ) -> List[int]:
        high_positions = []
        for pos in positions:
            column = recovery_scores[:, int(pos)]
            finite_column = column[torch.isfinite(column)]
            if finite_column.numel() and float(finite_column.max().item()) > float(threshold):
                high_positions.append(int(pos))
        return high_positions

    @staticmethod
    def positions_over_abs_recovery_threshold(
        recovery_scores: torch.Tensor,
        positions: Sequence[int],
        threshold: float,
    ) -> List[int]:
        high_positions = []
        for pos in positions:
            column = recovery_scores[:, int(pos)]
            finite_column = column[torch.isfinite(column)]
            if finite_column.numel() and float(torch.max(torch.abs(finite_column)).item()) > float(threshold):
                high_positions.append(int(pos))
        return high_positions

    def build_residual_score_tables(
        self,
        scores: torch.Tensor,
        patch_units: Sequence[PatchUnit],
        direction: str,
        baseline: float,
        source_ref: float,
        distance_column: str,
        ranking_metric: str = "signed",
    ) -> Tuple:
        score_rows = []
        for layer in range(scores.shape[0]):
            for unit in patch_units:
                value = float(scores[layer, int(unit.unit_id)].item())
                if not math.isfinite(value):
                    continue
                destination_positions = unit_destination_positions(unit, direction)
                start_pos = min(destination_positions) if destination_positions else None
                end_pos = max(destination_positions) if destination_positions else None
                score_rows.append(
                    {
                        "layer": int(layer),
                        "pos": "" if start_pos is None else int(start_pos),
                        "patch_unit_id": int(unit.unit_id),
                        "patch_unit_type": unit.unit_kind,
                        "patch_unit_text": unit.corrupted_text
                        if direction == "forward"
                        else unit.clean_text,
                        "unit_id": int(unit.unit_id),
                        "unit_kind": unit.unit_kind,
                        "start_pos": "" if start_pos is None else int(start_pos),
                        "end_pos": "" if end_pos is None else int(end_pos),
                        "span_len": int(len(destination_positions)),
                        "positions": json.dumps([int(pos) for pos in destination_positions]),
                        "clean_positions": json.dumps([int(pos) for pos in unit.clean_positions]),
                        "corrupted_positions": json.dumps(
                            [int(pos) for pos in unit.corrupted_positions]
                        ),
                        "clean_tokens": json.dumps(list(unit.clean_tokens)),
                        "corrupted_tokens": json.dumps(list(unit.corrupted_tokens)),
                        "clean_token_ids": json.dumps([int(token_id) for token_id in unit.clean_token_ids]),
                        "corrupted_token_ids": json.dumps(
                            [int(token_id) for token_id in unit.corrupted_token_ids]
                        ),
                        "clean_text": unit.clean_text,
                        "corrupted_text": unit.corrupted_text,
                        "is_asymmetric": bool(unit.is_asymmetric),
                        "score": value,
                        "patched_margin": value,
                        "margin_delta": value - baseline,
                        "signed_recovery": normalize_recovery(value, baseline, source_ref),
                        "distance_recovery": distance_based_recovery(value, baseline, source_ref),
                        "normalized_recovery": normalize_recovery(value, baseline, source_ref),
                        "distance_based_recovery": distance_based_recovery(value, baseline, source_ref),
                        "distance_reduction_gain": patch_gain_toward_clean(value, baseline, source_ref),
                        distance_column: abs(value - source_ref),
                    }
                )

        score_df = pd.DataFrame(score_rows)
        if score_df.empty:
            top_sites_df = pd.DataFrame(
                columns=[
                    "layer",
                    "pos",
                    "patch_unit_id",
                    "patch_unit_type",
                    "patch_unit_text",
                    "unit_id",
                    "unit_kind",
                    "start_pos",
                    "end_pos",
                    "span_len",
                    "positions",
                    "clean_positions",
                    "corrupted_positions",
                    "clean_tokens",
                    "corrupted_tokens",
                    "clean_token_ids",
                    "corrupted_token_ids",
                    "clean_text",
                    "corrupted_text",
                    "is_asymmetric",
                    "score",
                    "patched_margin",
                    "margin_delta",
                    "signed_recovery",
                    "distance_recovery",
                    "normalized_recovery",
                    "distance_based_recovery",
                    "distance_reduction_gain",
                    distance_column,
                ]
            )
        else:
            if ranking_metric == "distance":
                sort_columns = ["distance_recovery", distance_column]
                sort_ascending = [False, True]
            else:
                sort_columns = ["signed_recovery", "margin_delta"]
                sort_ascending = [False, False]
            top_sites_df = (
                score_df.sort_values(sort_columns, ascending=sort_ascending)
                .head(50)
                .reset_index(drop=True)
            )
        return score_df, top_sites_df

    def build_residual_span_score_tables(
        self,
        scores: torch.Tensor,
        spans: Sequence[Dict[str, Any]],
        baseline: float,
        source_ref: float,
        distance_column: str,
        ranking_metric: str = "signed",
    ) -> Tuple:
        score_rows = []
        for layer in range(scores.shape[0]):
            for span_index, span in enumerate(spans):
                value = float(scores[layer, int(span_index)].item())
                if not math.isfinite(value):
                    continue
                score_rows.append(
                    {
                        "layer": int(layer),
                        "span_id": int(span["span_id"]),
                        "start_pos": int(span["start_pos"]),
                        "end_pos": int(span["end_pos"]),
                        "span_len": int(span["span_len"]),
                        "positions": json.dumps([int(pos) for pos in span["positions"]]),
                        "clean_tokens": json.dumps(list(span.get("clean_tokens", []))),
                        "corrupted_tokens": json.dumps(list(span.get("corrupted_tokens", []))),
                        "clean_text": str(span.get("clean_text", "")),
                        "corrupted_text": str(span.get("corrupted_text", "")),
                        "score": value,
                        "signed_recovery": normalize_recovery(value, baseline, source_ref),
                        "distance_recovery": distance_based_recovery(value, baseline, source_ref),
                        "normalized_recovery": normalize_recovery(value, baseline, source_ref),
                        "distance_based_recovery": distance_based_recovery(value, baseline, source_ref),
                        "distance_reduction_gain": patch_gain_toward_clean(value, baseline, source_ref),
                        distance_column: abs(value - source_ref),
                    }
                )

        columns = [
            "layer",
            "span_id",
            "start_pos",
            "end_pos",
            "span_len",
            "positions",
            "clean_tokens",
            "corrupted_tokens",
            "clean_text",
            "corrupted_text",
            "score",
            "signed_recovery",
            "distance_recovery",
            "normalized_recovery",
            "distance_based_recovery",
            "distance_reduction_gain",
            distance_column,
        ]
        score_df = pd.DataFrame(score_rows, columns=columns)
        if score_df.empty:
            top_spans_df = pd.DataFrame(columns=columns)
        else:
            if ranking_metric == "distance":
                sort_columns = ["distance_recovery", distance_column]
                sort_ascending = [False, True]
            else:
                sort_columns = ["signed_recovery"]
                sort_ascending = [False]
            top_spans_df = (
                score_df.sort_values(sort_columns, ascending=sort_ascending)
                .head(50)
                .reset_index(drop=True)
            )
        return score_df, top_spans_df

    def save_tokenization_mismatch_skip_artifacts(
        self,
        pair_name: str,
        pair: Dict,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        answer_spec: AnswerDecisionSpec,
        clean_tokens,
        corrupted_tokens,
        aligned_seq_len: int,
        all_patch_units: Sequence[PatchUnit],
        patchable_units: Sequence[PatchUnit],
        alignment_rows: Sequence[Dict[str, Any]],
        first_changed_pos: Optional[int],
        output_dir: Path,
        pair_output_dir_override: Optional[Path] = None,
    ) -> Dict:
        pair_output_dir = pair_output_dir_override or (output_dir / slugify(pair_name))
        pair_output_dir.mkdir(parents=True, exist_ok=True)
        remove_artifacts(
            pair_output_dir,
            [
                *RESIDUAL_SCORE_ARTIFACT_NAMES,
                *LEGACY_SPAN_ARTIFACT_NAMES,
            ],
        )

        pd.DataFrame(
            [methodology_alignment_row_to_csv_row(row, []) for row in alignment_rows]
        ).to_csv(pair_output_dir / "token_alignment.csv", index=False)
        pd.DataFrame(methodology_patch_units_to_rows(patchable_units, [])).to_csv(
            pair_output_dir / "patch_units.csv",
            index=False,
        )

        asymmetric_units = [
            unit for unit in changed_region_units(patchable_units) if unit.is_asymmetric
        ]
        asymmetric_rows = changed_region_summary(asymmetric_units)
        report_lines = [
            f"Pair: {pair_name}",
            f"Status: skipped_tokenization_mismatch",
            f"Reason: at least one changed region has different source/destination token counts.",
            f"Source role: {clean_variant}",
            f"Destination role: {corrupted_variant}",
            f"Source sequence length: {int(clean_tokens.shape[1])}",
            f"Destination sequence length: {int(corrupted_tokens.shape[1])}",
            f"Aligned min sequence length: {int(aligned_seq_len)}",
            f"First changed position: {first_changed_pos}",
            "",
            "Asymmetric changed regions:",
        ]
        for unit in asymmetric_units:
            report_lines.extend(
                [
                    f"- unit_id={unit.unit_id}",
                    f"  source_positions={unit.clean_positions}",
                    f"  destination_positions={unit.corrupted_positions}",
                    f"  source_text={unit.clean_text!r}",
                    f"  destination_text={unit.corrupted_text!r}",
                ]
            )
        report_lines.extend(
            [
                "",
                f"Source snippet ({clean_variant}):",
                clean_code.rstrip(),
                "",
                f"Destination snippet ({corrupted_variant}):",
                corrupted_code.rstrip(),
                "",
            ]
        )
        (pair_output_dir / "skipped_tokenization_mismatch.txt").write_text(
            "\n".join(report_lines),
            encoding="utf-8",
        )

        summary = {
            "status": "skipped_tokenization_mismatch",
            "pair": pair_name,
            "phenomenon": pair["meta"].get("phenomenon", ""),
            "pair_id": pair.get("pair_id") or pair_name,
            "pair_family": pair.get("pair_family") or pair["meta"].get("pair_family", ""),
            "patch_direction": f"{clean_variant}_to_{corrupted_variant}",
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "source_output": answer_spec.forward_positive_answer,
            "destination_output": answer_spec.forward_negative_answer,
            "answer_resolution_schema": answer_spec.answer_resolution_schema,
            "clean_positive_answer": answer_spec.clean_positive_answer,
            "clean_negative_answer": answer_spec.clean_negative_answer,
            "corrupted_positive_answer": answer_spec.corrupted_positive_answer,
            "corrupted_negative_answer": answer_spec.corrupted_negative_answer,
            "forward_direction": f"{clean_variant}_into_{corrupted_variant}",
            "forward_scored_answer": answer_spec.forward_positive_answer,
            "forward_positive_answer": answer_spec.forward_positive_answer,
            "forward_negative_answer": answer_spec.forward_negative_answer,
            "reverse_positive_answer": answer_spec.reverse_positive_answer,
            "reverse_negative_answer": answer_spec.reverse_negative_answer,
            "clean_expected_output": answer_spec.clean_expected_output,
            "corrupted_expected_output": answer_spec.corrupted_expected_output,
            "source_seq_len": int(clean_tokens.shape[1]),
            "destination_seq_len": int(corrupted_tokens.shape[1]),
            "clean_seq_len": int(clean_tokens.shape[1]),
            "corrupted_seq_len": int(corrupted_tokens.shape[1]),
            "aligned_seq_len": int(aligned_seq_len),
            "first_changed_pos": None if first_changed_pos is None else int(first_changed_pos),
            "patching_unit_mode": "mixed_token_region",
            "patching_prompt_region": "code_and_output_cue",
            "n_patch_units": int(len(patchable_units)),
            "n_changed_spans": int(len(changed_region_units(patchable_units))),
            "changed_spans": changed_region_summary(patchable_units),
            "n_asymmetric_changed_regions": int(len(asymmetric_units)),
            "asymmetric_changed_regions": asymmetric_rows,
            "skip_report": str(pair_output_dir / "skipped_tokenization_mismatch.txt"),
            "output_dir": str(pair_output_dir),
        }
        (pair_output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
        return summary

    def save_pair_artifacts(
        self,
        pair_name: str,
        pair: Dict,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        clean_prompt: str,
        corrupted_prompt: str,
        target_answer: str,
        wrong_answer: Optional[str],
        clean_expected_output: Optional[str],
        corrupted_expected_output: Optional[str],
        clean_tokens,
        corrupted_tokens,
        clean_logits,
        corrupted_logits,
        clean_cache,
        corrupted_cache,
        aligned_seq_len: int,
        positions: Sequence[int],
        first_changed_pos: Optional[int],
        output_dir: Path,
        save_additional_tasks: bool,
        unit_test_flow_patching: bool,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
        token_label_recovery_threshold: float,
        recovery_metric: str,
        max_positions: Optional[int],
        pos_context: int,
        pair_output_dir_override: Optional[Path] = None,
    ) -> Dict:
        pair_output_dir = output_dir / slugify(pair_name)
        pair_output_dir.mkdir(parents=True, exist_ok=True)

        baseline = self.metric_from_logits_with_target_answer(corrupted_logits, target_answer)
        clean_ref = self.metric_from_logits_with_target_answer(clean_logits, target_answer)

        scores = self.residual_patch_sweep(
            source_cache=clean_cache,
            destination_tokens=corrupted_tokens,
            answer=target_answer,
            aligned_seq_len=aligned_seq_len,
            positions=positions,
            desc="Clean -> violation residual layers",
        )
        distance_recovery_scores = distance_based_recovery_matrix(scores, baseline, clean_ref)
        selected_recovery_scores = recovery_matrix_for_metric(
            recovery_metric,
            scores,
            baseline,
            clean_ref,
        )

        clean_strs = self.token_strs(clean_tokens)
        corrupted_strs = self.token_strs(corrupted_tokens)
        high_recovery_positions = self.positions_over_recovery_threshold(
            recovery_scores=selected_recovery_scores,
            positions=positions,
            threshold=token_label_recovery_threshold,
        )
        high_recovery_token_labels = {
            int(pos): format_token_for_axis(corrupted_strs[int(pos)])
            for pos in high_recovery_positions
            if int(pos) < len(corrupted_strs)
        }
        token_label_by_pos = {
            int(pos): format_token_for_axis(corrupted_strs[int(pos)])
            for pos in positions
            if int(pos) < len(corrupted_strs)
        }

        self.plot_heatmap(
            scores,
            positions=positions,
            colorbar_label="first-answer-token logit",
            output_path=pair_output_dir / "residual_scores_heatmap.pdf",
            token_labels=token_label_by_pos,
            forced_tick_positions=high_recovery_positions,
            cmap=SOURCE_LOGIT_MARGIN_CMAP,
        )
        self.plot_heatmap(
            selected_recovery_scores,
            positions=positions,
            colorbar_label=recovery_heatmap_kwargs(
                recovery_metric,
                selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["colorbar_label"],
            output_path=pair_output_dir / "residual_recovery_heatmap.pdf",
            token_labels=token_label_by_pos,
            forced_tick_positions=high_recovery_positions,
            vmin=recovery_heatmap_kwargs(
                recovery_metric,
                selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["vmin"],
            vmax=recovery_heatmap_kwargs(
                recovery_metric,
                selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["vmax"],
            cmap=recovery_heatmap_kwargs(
                recovery_metric,
                selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["cmap"],
        )

        swept_positions = set(int(p) for p in positions)
        token_rows = []
        for pos in range(aligned_seq_len):
            token_rows.append(
                {
                    "pos": int(pos),
                    "is_swept": bool(pos in swept_positions),
                    "is_first_changed": bool(first_changed_pos is not None and pos == first_changed_pos),
                    "clean_token_id": int(clean_tokens[0, pos].item()),
                    "corrupted_token_id": int(corrupted_tokens[0, pos].item()),
                    "clean_token_str": clean_strs[pos],
                    "corrupted_token_str": corrupted_strs[pos],
                    "token_ids_equal": bool(clean_tokens[0, pos].item() == corrupted_tokens[0, pos].item()),
                }
            )
        pd.DataFrame(token_rows).to_csv(pair_output_dir / "token_alignment.csv", index=False)

        score_df, top_sites_df = self.build_residual_score_tables(
            scores=scores,
            positions=positions,
            baseline=baseline,
            source_ref=clean_ref,
            distance_column="distance_to_clean",
            ranking_metric=recovery_metric,
        )
        score_df.to_csv(pair_output_dir / "residual_scores.csv", index=False)
        top_sites_df.to_csv(pair_output_dir / "top_residual_sites.csv", index=False)

        changed_spans = self.enrich_span_metadata(
            self.changed_token_spans(clean_tokens, corrupted_tokens, positions),
            clean_strs,
            corrupted_strs,
        )
        span_ids = [int(span["span_id"]) for span in changed_spans]
        span_scores = self.residual_span_patch_sweep(
            source_cache=clean_cache,
            destination_tokens=corrupted_tokens,
            answer=target_answer,
            spans=changed_spans,
            desc="Clean -> violation residual changed spans",
        )
        span_distance_recovery_scores = distance_based_recovery_matrix(
            span_scores,
            baseline,
            clean_ref,
        )
        span_selected_recovery_scores = recovery_matrix_for_metric(
            recovery_metric,
            span_scores,
            baseline,
            clean_ref,
        )
        high_recovery_span_ids = self.positions_over_recovery_threshold(
            recovery_scores=span_selected_recovery_scores,
            positions=span_ids,
            threshold=token_label_recovery_threshold,
        )
        span_label_by_id = {
            int(span["span_id"]): format_span_for_axis(str(span.get("corrupted_text", "")))
            for span in changed_spans
        }

        self.plot_heatmap(
            span_selected_recovery_scores,
            positions=span_ids,
            colorbar_label=f"changed-span {recovery_metric_column(recovery_metric).replace('_', ' ')}",
            output_path=pair_output_dir / "residual_span_recovery_heatmap.pdf",
            token_labels=span_label_by_id,
            forced_tick_positions=high_recovery_span_ids,
            xlabel="Changed token span",
            vmin=recovery_heatmap_kwargs(
                recovery_metric,
                span_selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["vmin"],
            vmax=recovery_heatmap_kwargs(
                recovery_metric,
                span_selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["vmax"],
            cmap=recovery_heatmap_kwargs(
                recovery_metric,
                span_selected_recovery_scores,
                CLEAN_RECOVERY_CMAP,
            )["cmap"],
        )

        span_score_df, top_spans_df = self.build_residual_span_score_tables(
            scores=span_scores,
            spans=changed_spans,
            baseline=baseline,
            source_ref=clean_ref,
            distance_column="distance_to_clean",
            ranking_metric=recovery_metric,
        )
        span_score_df.to_csv(pair_output_dir / "residual_span_scores.csv", index=False)
        top_spans_df.to_csv(pair_output_dir / "top_residual_spans.csv", index=False)

        best_span = top_spans_df.iloc[0].to_dict() if not top_spans_df.empty else {}
        span_summary = {
            "n_changed_spans": int(len(changed_spans)),
            "changed_spans": [
                {
                    "span_id": int(span["span_id"]),
                    "start_pos": int(span["start_pos"]),
                    "end_pos": int(span["end_pos"]),
                    "span_len": int(span["span_len"]),
                    "positions": [int(pos) for pos in span["positions"]],
                    "clean_text": str(span.get("clean_text", "")),
                    "corrupted_text": str(span.get("corrupted_text", "")),
                }
                for span in changed_spans
            ],
            "best_span_id": None if not best_span else int(best_span["span_id"]),
            "best_span_positions": []
            if not best_span
            else json.loads(str(best_span["positions"])),
            "best_span_start_pos": None if not best_span else int(best_span["start_pos"]),
            "best_span_end_pos": None if not best_span else int(best_span["end_pos"]),
            "best_span_len": None if not best_span else int(best_span["span_len"]),
            "best_span_score": None if not best_span else float(best_span["score"]),
            "best_span_normalized_recovery": None
            if not best_span
            else float(best_span["normalized_recovery"]),
            "best_span_distance_based_recovery": None
            if not best_span
            else float(best_span["distance_based_recovery"]),
            "best_span_distance_reduction_gain": None
            if not best_span
            else float(best_span["distance_reduction_gain"]),
        }

        reverse_summary = {
            "wrong_answer": wrong_answer,
            "reverse_direction": f"{corrupted_variant}_into_{clean_variant}",
            "reverse_scored_answer": wrong_answer,
            "reverse_baseline_wrong_answer_score": None,
            "reverse_source_ref_wrong_answer_score": None,
            "reverse_best_layer": None,
            "reverse_best_pos": None,
            "reverse_best_score": None,
            "reverse_best_normalized_recovery": None,
            "reverse_best_distance_based_recovery": None,
            "reverse_best_distance_reduction_gain": None,
            "reverse_token_labeled_positions": [],
            "reverse_token_labeled_tokens": {},
            "reverse_best_span_id": None,
            "reverse_best_span_positions": [],
            "reverse_best_span_start_pos": None,
            "reverse_best_span_end_pos": None,
            "reverse_best_span_len": None,
            "reverse_best_span_score": None,
            "reverse_best_span_normalized_recovery": None,
            "reverse_best_span_distance_based_recovery": None,
            "reverse_best_span_distance_reduction_gain": None,
        }
        if wrong_answer is None:
            reverse_summary["reverse_patch_error"] = "No wrong answer was available for reverse-direction scoring."
        else:
            reverse_baseline = self.metric_from_logits_with_target_answer(clean_logits, wrong_answer)
            reverse_source_ref = self.metric_from_logits_with_target_answer(corrupted_logits, wrong_answer)
            reverse_scores = self.residual_patch_sweep(
                source_cache=corrupted_cache,
                destination_tokens=clean_tokens,
                answer=wrong_answer,
                aligned_seq_len=aligned_seq_len,
                positions=positions,
                desc="Violation -> clean residual layers",
            )
            reverse_distance_recovery_scores = distance_based_recovery_matrix(
                reverse_scores,
                reverse_baseline,
                reverse_source_ref,
            )
            reverse_selected_recovery_scores = recovery_matrix_for_metric(
                recovery_metric,
                reverse_scores,
                reverse_baseline,
                reverse_source_ref,
            )
            reverse_high_recovery_positions = self.positions_over_recovery_threshold(
                recovery_scores=reverse_selected_recovery_scores,
                positions=positions,
                threshold=token_label_recovery_threshold,
            )
            reverse_high_recovery_token_labels = {
                int(pos): format_token_for_axis(clean_strs[int(pos)])
                for pos in reverse_high_recovery_positions
                if int(pos) < len(clean_strs)
            }
            reverse_token_label_by_pos = {
                int(pos): format_token_for_axis(clean_strs[int(pos)])
                for pos in positions
                if int(pos) < len(clean_strs)
            }

            self.plot_heatmap(
                reverse_scores,
                positions=positions,
                colorbar_label="wrong-answer first-token logit",
                output_path=pair_output_dir / "reverse_residual_scores_heatmap.pdf",
                token_labels=reverse_token_label_by_pos,
                forced_tick_positions=reverse_high_recovery_positions,
                cmap=DESTINATION_LOGIT_MARGIN_CMAP,
            )
            self.plot_heatmap(
                reverse_selected_recovery_scores,
                positions=positions,
                colorbar_label=recovery_metric_column(recovery_metric).replace("_", " "),
                output_path=pair_output_dir / "reverse_residual_recovery_heatmap.pdf",
                token_labels=reverse_token_label_by_pos,
                forced_tick_positions=reverse_high_recovery_positions,
                vmin=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["vmin"],
                vmax=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["vmax"],
                cmap=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["cmap"],
            )

            reverse_score_df, reverse_top_sites_df = self.build_residual_score_tables(
                scores=reverse_scores,
                positions=positions,
                baseline=reverse_baseline,
                source_ref=reverse_source_ref,
                distance_column="distance_to_violation_ref",
                ranking_metric=recovery_metric,
            )
            reverse_score_df.to_csv(pair_output_dir / "reverse_residual_scores.csv", index=False)
            reverse_top_sites_df.to_csv(pair_output_dir / "top_reverse_residual_sites.csv", index=False)

            reverse_span_scores = self.residual_span_patch_sweep(
                source_cache=corrupted_cache,
                destination_tokens=clean_tokens,
                answer=wrong_answer,
                spans=changed_spans,
                desc="Violation -> clean residual changed spans",
            )
            reverse_span_distance_recovery_scores = distance_based_recovery_matrix(
                reverse_span_scores,
                reverse_baseline,
                reverse_source_ref,
            )
            reverse_span_selected_recovery_scores = recovery_matrix_for_metric(
                recovery_metric,
                reverse_span_scores,
                reverse_baseline,
                reverse_source_ref,
            )
            reverse_high_recovery_span_ids = self.positions_over_recovery_threshold(
                recovery_scores=reverse_span_selected_recovery_scores,
                positions=span_ids,
                threshold=token_label_recovery_threshold,
            )
            reverse_span_label_by_id = {
                int(span["span_id"]): format_span_for_axis(str(span.get("clean_text", "")))
                for span in changed_spans
            }
            self.plot_heatmap(
                reverse_span_selected_recovery_scores,
                positions=span_ids,
                colorbar_label=f"changed-span {recovery_metric_column(recovery_metric).replace('_', ' ')}",
                output_path=pair_output_dir / "reverse_residual_span_recovery_heatmap.pdf",
                token_labels=reverse_span_label_by_id,
                forced_tick_positions=reverse_high_recovery_span_ids,
                xlabel="Changed token span",
                vmin=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_span_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["vmin"],
                vmax=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_span_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["vmax"],
                cmap=recovery_heatmap_kwargs(
                    recovery_metric,
                    reverse_span_selected_recovery_scores,
                    CORRUPTED_RECOVERY_CMAP,
                )["cmap"],
            )

            reverse_span_score_df, reverse_top_spans_df = self.build_residual_span_score_tables(
                scores=reverse_span_scores,
                spans=changed_spans,
                baseline=reverse_baseline,
                source_ref=reverse_source_ref,
                distance_column="distance_to_violation_ref",
                ranking_metric=recovery_metric,
            )
            reverse_span_score_df.to_csv(
                pair_output_dir / "reverse_residual_span_scores.csv",
                index=False,
            )
            reverse_top_spans_df.to_csv(
                pair_output_dir / "top_reverse_residual_spans.csv",
                index=False,
            )

            reverse_best = reverse_top_sites_df.iloc[0].to_dict() if not reverse_top_sites_df.empty else {}
            reverse_best_span = (
                reverse_top_spans_df.iloc[0].to_dict() if not reverse_top_spans_df.empty else {}
            )
            reverse_summary.update(
                {
                    "reverse_baseline_wrong_answer_score": float(reverse_baseline),
                    "reverse_source_ref_wrong_answer_score": float(reverse_source_ref),
                    "reverse_best_layer": None if not reverse_best else int(reverse_best["layer"]),
                    "reverse_best_pos": None if not reverse_best else int(reverse_best["pos"]),
                    "reverse_best_score": None if not reverse_best else float(reverse_best["score"]),
                    "reverse_best_normalized_recovery": None
                    if not reverse_best
                    else float(reverse_best["normalized_recovery"]),
                    "reverse_best_distance_based_recovery": None
                    if not reverse_best
                    else float(reverse_best["distance_based_recovery"]),
                    "reverse_best_distance_reduction_gain": None
                    if not reverse_best
                    else float(reverse_best["distance_reduction_gain"]),
                    "reverse_token_labeled_positions": [int(pos) for pos in reverse_high_recovery_positions],
                    "reverse_token_labeled_tokens": {
                        str(pos): label for pos, label in reverse_high_recovery_token_labels.items()
                    },
                    "reverse_best_span_id": None
                    if not reverse_best_span
                    else int(reverse_best_span["span_id"]),
                    "reverse_best_span_positions": []
                    if not reverse_best_span
                    else json.loads(str(reverse_best_span["positions"])),
                    "reverse_best_span_start_pos": None
                    if not reverse_best_span
                    else int(reverse_best_span["start_pos"]),
                    "reverse_best_span_end_pos": None
                    if not reverse_best_span
                    else int(reverse_best_span["end_pos"]),
                    "reverse_best_span_len": None
                    if not reverse_best_span
                    else int(reverse_best_span["span_len"]),
                    "reverse_best_span_score": None
                    if not reverse_best_span
                    else float(reverse_best_span["score"]),
                    "reverse_best_span_normalized_recovery": None
                    if not reverse_best_span
                    else float(reverse_best_span["normalized_recovery"]),
                    "reverse_best_span_distance_based_recovery": None
                    if not reverse_best_span
                    else float(reverse_best_span["distance_based_recovery"]),
                    "reverse_best_span_distance_reduction_gain": None
                    if not reverse_best_span
                    else float(reverse_best_span["distance_reduction_gain"]),
                }
            )

        best = top_sites_df.iloc[0].to_dict() if not top_sites_df.empty else {}
        summary = {
            "pair": pair_name,
            "phenomenon": pair["meta"].get("phenomenon", ""),
            "answer": target_answer,
            "forward_direction": f"{clean_variant}_into_{corrupted_variant}",
            "forward_scored_answer": target_answer,
            "clean_expected_output": clean_expected_output,
            "corrupted_expected_output": corrupted_expected_output,
            "clean_seq_len": int(clean_tokens.shape[1]),
            "corrupted_seq_len": int(corrupted_tokens.shape[1]),
            "aligned_seq_len": int(aligned_seq_len),
            "first_changed_pos": None if first_changed_pos is None else int(first_changed_pos),
            "swept_positions": [int(p) for p in positions],
            "n_swept_positions": int(len(positions)),
            "baseline_proxy_score": float(baseline),
            "clean_ref_proxy_score": float(clean_ref),
            "best_layer": None if not best else int(best["layer"]),
            "best_pos": None if not best else int(best["pos"]),
            "best_score": None if not best else float(best["score"]),
            "best_normalized_recovery": None if not best else float(best["normalized_recovery"]),
            "best_distance_based_recovery": None if not best else float(best["distance_based_recovery"]),
            "best_distance_reduction_gain": None if not best else float(best["distance_reduction_gain"]),
            "token_label_recovery_threshold": float(token_label_recovery_threshold),
            "token_label_recovery_metric": recovery_metric_column(recovery_metric),
            "token_labeled_positions": [int(pos) for pos in high_recovery_positions],
            "token_labeled_tokens": {str(pos): label for pos, label in high_recovery_token_labels.items()},
            "output_dir": str(pair_output_dir),
        }
        summary.update(span_summary)
        summary.update(reverse_summary)

        additional_generations: Optional[AdditionalGenerations] = None
        if save_additional_tasks or unit_test_flow_patching:
            try:
                additional_generations = self.build_additional_generations(
                    pair_name=pair_name,
                    target_answer=target_answer,
                    destination_reference_output=wrong_answer,
                    clean_expected_output=clean_expected_output,
                    corrupted_expected_output=corrupted_expected_output,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    answer_max_new_tokens=answer_max_new_tokens,
                    unit_test_max_new_tokens=unit_test_max_new_tokens,
                )
            except Exception as e:
                summary["additional_generations_error"] = str(e)

        clean_prompt_parts = split_fenced_prompt(clean_prompt)
        corrupted_prompt_parts = split_fenced_prompt(corrupted_prompt)
        forward_direction_name = f"{clean_variant}_to_{corrupted_variant}"
        reverse_direction_name = f"{corrupted_variant}_to_{clean_variant}"
        output_contexts = {
            forward_direction_name: {
                "patch_direction": forward_direction_name,
                "destination_role": corrupted_variant,
                "prompt_text": corrupted_prompt_parts["prompt_text"],
                "snippet_text": corrupted_prompt_parts["snippet_text"] or corrupted_code.rstrip(),
                "output_label": prompt_output_label(corrupted_prompt_parts, "Generated Answer"),
                "generated_output": ""
                if additional_generations is None
                else additional_generations.corrupted_model_answer,
                "generated_prefix": "",
                "common_target_prefix": forward_common_text,
                "source_output": target_answer,
                "destination_output": wrong_answer,
                "source_token": forward_target_info["positive_token_str"],
                "destination_token": forward_target_info["negative_token_str"],
                "next_token_scored": (
                    f"{forward_target_info['positive_token_str']} vs "
                    f"{forward_target_info['negative_token_str']}"
                ),
            },
            reverse_direction_name: {
                "patch_direction": reverse_direction_name,
                "destination_role": clean_variant,
                "prompt_text": clean_prompt_parts["prompt_text"],
                "snippet_text": clean_prompt_parts["snippet_text"] or clean_code.rstrip(),
                "output_label": prompt_output_label(clean_prompt_parts, "Generated Answer"),
                "generated_output": ""
                if additional_generations is None
                else additional_generations.clean_model_answer,
                "generated_prefix": "",
                "common_target_prefix": reverse_common_text,
                "source_output": wrong_answer,
                "destination_output": target_answer,
                "source_token": reverse_target_info["positive_token_str"],
                "destination_token": reverse_target_info["negative_token_str"],
                "next_token_scored": (
                    f"{reverse_target_info['positive_token_str']} vs "
                    f"{reverse_target_info['negative_token_str']}"
                ),
            },
        }
        (pair_output_dir / "output_contexts.json").write_text(
            json.dumps(output_contexts, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        summary["output_contexts"] = str(pair_output_dir / "output_contexts.json")

        if additional_generations is not None:
            try:
                generation_json_path = write_additional_generations_json(
                    pair_output_dir / "unit_test_generations.json",
                    additional_generations,
                )
                summary["unit_test_generations"] = str(generation_json_path)
                attach_task1_outputs(
                    summary,
                    pair,
                    clean_variant,
                    corrupted_variant,
                    additional_generations,
                )
            except Exception as e:
                summary["unit_test_generations_error"] = str(e)

        if save_additional_tasks and additional_generations is not None:
            try:
                task_dump_path = self.write_additional_generation_dump(
                    output_path=pair_output_dir / "additional_task_generations.txt",
                    generations=additional_generations,
                )
                summary["additional_task_dump"] = str(task_dump_path)
            except Exception as e:
                summary["additional_task_dump_error"] = str(e)

        if unit_test_flow_patching and additional_generations is not None:
            try:
                flow_summary = self.run_unit_test_flow_patching(
                    pair_output_dir=pair_output_dir,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    pair_id=str(pair.get("pair_id") or pair_name),
                    pair_family=str(pair.get("pair_family") or pair["meta"].get("pair_family", "")),
                    additional_generations=additional_generations,
                    max_positions=max_positions,
                    pos_context=pos_context,
                    token_label_recovery_threshold=token_label_recovery_threshold,
                )
                summary["unit_test_flow"] = flow_summary
                summary["unit_test_flow_parsed_assertions"] = int(
                    flow_summary.get("parsed_assertions", 0)
                )
                summary["unit_test_flow_analyzed_forward_decisions"] = int(
                    flow_summary.get("analyzed_forward_decisions", 0)
                )
                summary["unit_test_flow_analyzed_reverse_decisions"] = int(
                    flow_summary.get("analyzed_reverse_decisions", 0)
                )
                summary["unit_test_flow_skipped_decisions"] = int(
                    flow_summary.get("skipped_decisions", 0)
                )
                summary["unit_test_flow_best_prompt_to_output_site"] = json.dumps(
                    flow_summary.get("best_prompt_to_output_site", {})
                )
                summary["unit_test_flow_best_response_carry_site"] = json.dumps(
                    flow_summary.get("best_response_carry_site", {})
                )
                summary["unit_test_flow_best_combined_context_site"] = json.dumps(
                    flow_summary.get("best_combined_context_site", {})
                )
            except Exception as e:
                summary["unit_test_flow_error"] = str(e)

        (pair_output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
        return summary

    def save_patch_unit_pair_artifacts(
        self,
        pair_name: str,
        pair: Dict,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        clean_prompt: str,
        corrupted_prompt: str,
        answer_spec: AnswerDecisionSpec,
        clean_tokens,
        corrupted_tokens,
        clean_logits,
        corrupted_logits,
        clean_cache,
        corrupted_cache,
        aligned_seq_len: int,
        patch_units: Sequence[PatchUnit],
        patchable_units: Sequence[PatchUnit],
        alignment_rows: Sequence[Dict[str, Any]],
        first_changed_pos: Optional[int],
        output_dir: Path,
        save_additional_tasks: bool,
        unit_test_flow_patching: bool,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
        token_label_recovery_threshold: float,
        recovery_metric: str,
        max_positions: Optional[int],
        pos_context: int,
        pair_output_dir_override: Optional[Path] = None,
    ) -> Dict:
        target_answer = answer_spec.forward_positive_answer
        wrong_answer = answer_spec.forward_negative_answer
        reverse_positive_answer = answer_spec.reverse_positive_answer
        reverse_negative_answer = answer_spec.reverse_negative_answer
        clean_expected_output = answer_spec.clean_expected_output
        corrupted_expected_output = answer_spec.corrupted_expected_output

        pair_output_dir = pair_output_dir_override or (output_dir / slugify(pair_name))
        pair_output_dir.mkdir(parents=True, exist_ok=True)
        remove_artifacts(
            pair_output_dir,
            [*LEGACY_SPAN_ARTIFACT_NAMES, *SKIP_ARTIFACT_NAMES],
        )

        swept_unit_ids = [int(unit.unit_id) for unit in patch_units]
        pd.DataFrame(
            [
                methodology_alignment_row_to_csv_row(row, swept_unit_ids)
                for row in alignment_rows
            ]
        ).to_csv(pair_output_dir / "token_alignment.csv", index=False)
        pd.DataFrame(
            methodology_patch_units_to_rows(patchable_units, swept_unit_ids)
        ).to_csv(
            pair_output_dir / "patch_units.csv",
            index=False,
        )

        if wrong_answer is None:
            raise ValueError(
                "Answer patching now uses a contrastive logit margin and requires "
                "a forward negative answer. Provide --manual-wrong-answer, metadata "
                "wrong_answer, or metadata wrong_answers for the clean/source variant."
            )
        if reverse_positive_answer is None or reverse_negative_answer is None:
            raise ValueError(
                "Reverse answer patching uses a contrastive logit margin and requires "
                "positive and negative answers for the corrupted/source variant. Provide "
                "metadata answers/wrong_answers for both variants or explicit answer_decisions."
            )

        forward_target_info = self.answer_decision_target_info(
            source_context=clean_prompt,
            destination_context=corrupted_prompt,
            positive_text=target_answer,
            negative_text=wrong_answer,
            direction_label="clean_into_corrupted",
        )
        forward_common_text = str(forward_target_info["common_text"])
        forward_clean_context = str(forward_target_info.get("source_context_text") or clean_prompt)
        forward_corrupted_context = str(
            forward_target_info.get("destination_context_text") or corrupted_prompt
        )
        forward_clean_tokens = self.to_tokens(forward_clean_context + forward_common_text)
        forward_corrupted_tokens = self.to_tokens(forward_corrupted_context + forward_common_text)
        with torch.no_grad():
            forward_clean_logits, forward_clean_cache = self.model.run_with_cache(
                forward_clean_tokens
            )
            forward_corrupted_logits, forward_corrupted_cache = self.model.run_with_cache(
                forward_corrupted_tokens
            )

        forward_positive_token_id = int(forward_target_info["positive_token_id"])
        forward_negative_token_id = int(forward_target_info["negative_token_id"])
        baseline = self.metric_from_logits_token_margin(
            forward_corrupted_logits,
            positive_token_id=forward_positive_token_id,
            negative_token_id=forward_negative_token_id,
        )
        clean_ref = self.metric_from_logits_token_margin(
            forward_clean_logits,
            positive_token_id=forward_positive_token_id,
            negative_token_id=forward_negative_token_id,
        )

        scores = self.residual_patch_unit_margin_sweep(
            source_cache=forward_clean_cache,
            destination_tokens=forward_corrupted_tokens,
            positive_token_id=forward_positive_token_id,
            negative_token_id=forward_negative_token_id,
            patch_units=patch_units,
            direction="forward",
            desc="Clean -> violation residual patch units",
        )
        distance_recovery_scores = distance_based_recovery_matrix(scores, baseline, clean_ref)
        selected_recovery_scores = recovery_matrix_for_metric(
            recovery_metric,
            scores,
            baseline,
            clean_ref,
        )

        high_recovery_unit_ids = self.positions_over_recovery_threshold(
            recovery_scores=selected_recovery_scores,
            positions=swept_unit_ids,
            threshold=token_label_recovery_threshold,
        )
        bold_recovery_unit_ids = self.positions_over_abs_recovery_threshold(
            recovery_scores=selected_recovery_scores,
            positions=swept_unit_ids,
            threshold=0.2,
        )
        emphasized_recovery_unit_ids = sorted(
            set(high_recovery_unit_ids) | set(bold_recovery_unit_ids)
        )
        unit_by_id = {int(unit.unit_id): unit for unit in patch_units}
        unit_label_by_id = {
            int(unit_id): format_patch_unit_axis_label(unit_by_id[int(unit_id)], "forward")
            for unit_id in swept_unit_ids
            if int(unit_id) in unit_by_id
        }

        self.plot_heatmap(
            scores,
            positions=swept_unit_ids,
            colorbar_label="answer logit margin",
            output_path=pair_output_dir / "residual_scores_heatmap.pdf",
            token_labels=unit_label_by_id,
            token_label_bold_positions=bold_recovery_unit_ids,
            forced_tick_positions=emphasized_recovery_unit_ids,
            xlabel="",
            cmap=SOURCE_LOGIT_MARGIN_CMAP,
        )
        recovery_plot = recovery_heatmap_kwargs(
            recovery_metric,
            selected_recovery_scores,
            CLEAN_RECOVERY_CMAP,
        )
        self.plot_heatmap(
            selected_recovery_scores,
            positions=swept_unit_ids,
            colorbar_label=recovery_plot["colorbar_label"],
            output_path=pair_output_dir / "residual_recovery_heatmap.pdf",
            token_labels=unit_label_by_id,
            token_label_bold_positions=bold_recovery_unit_ids,
            forced_tick_positions=emphasized_recovery_unit_ids,
            xlabel="",
            vmin=recovery_plot["vmin"],
            vmax=recovery_plot["vmax"],
            cmap=recovery_plot["cmap"],
        )

        score_df, top_sites_df = self.build_residual_score_tables(
            scores=scores,
            patch_units=patch_units,
            direction="forward",
            baseline=baseline,
            source_ref=clean_ref,
            distance_column="distance_to_clean",
            ranking_metric=recovery_metric,
        )
        pair_id = str(pair.get("pair_id") or pair_name)
        pair_family = str(pair.get("pair_family") or pair["meta"].get("pair_family", ""))
        forward_direction_name = f"{clean_variant}_to_{corrupted_variant}"
        reverse_direction_name = f"{corrupted_variant}_to_{clean_variant}"
        source_prompt = pair.get("prompts", {}).get(clean_variant)
        destination_prompt = pair.get("prompts", {}).get(corrupted_variant)
        forward_score_metadata = {
            "score_metric": "logit_margin",
            "pair_id": pair_id,
            "pair_family": pair_family,
            "patch_direction": forward_direction_name,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "source_file": "" if source_prompt is None else source_prompt.file,
            "destination_file": "" if destination_prompt is None else destination_prompt.file,
            "source_execution_output": clean_expected_output,
            "source_cue_output": "" if source_prompt is None else source_prompt.cue_output,
            "destination_execution_output": corrupted_expected_output,
            "destination_cue_output": "" if destination_prompt is None else destination_prompt.cue_output,
            "source_output": target_answer,
            "destination_output": wrong_answer,
            "source_output_label": ""
            if source_prompt is None
            else classify_behavior(target_answer, source_prompt.execution_output, source_prompt.cue_output),
            "destination_output_label": ""
            if destination_prompt is None
            else classify_behavior(wrong_answer, destination_prompt.execution_output, destination_prompt.cue_output),
            "source_is_accurate": ""
            if source_prompt is None
            else str(target_answer == source_prompt.execution_output),
            "destination_is_accurate": ""
            if destination_prompt is None
            else str(wrong_answer == destination_prompt.execution_output),
            "source_token": forward_target_info["positive_token_str"],
            "destination_token": forward_target_info["negative_token_str"],
            "source_token_id": forward_positive_token_id,
            "destination_token_id": forward_negative_token_id,
            "divergent_token_index": int(forward_target_info["divergent_token_index"]),
            "common_target_prefix": forward_common_text,
            "target_selection_mode": forward_target_info.get("target_selection_mode", ""),
            "answer_terminator_text": forward_target_info.get("answer_terminator_text", ""),
            "baseline_margin": float(baseline),
            "source_margin": float(clean_ref),
        }
        for df in (score_df, top_sites_df):
            for key, value in reversed(list(forward_score_metadata.items())):
                df.insert(0, key, value)
        write_methodology_residual_scores(
            score_df,
            pair_output_dir / "residual_scores.csv",
        )
        write_methodology_residual_scores(
            top_sites_df,
            pair_output_dir / "top_residual_sites.csv",
        )

        best = top_sites_df.iloc[0].to_dict() if not top_sites_df.empty else {}
        best_span = best_changed_region_row(top_sites_df)
        changed_spans = changed_region_summary(patchable_units)
        span_summary = {
            "n_changed_spans": int(len(changed_spans)),
            "changed_spans": changed_spans,
            "best_span_id": None if not best_span else int(best_span["unit_id"]),
            "best_span_positions": json_positions_from_row(best_span),
            "best_span_start_pos": None if not best_span else int(best_span["start_pos"]),
            "best_span_end_pos": None if not best_span else int(best_span["end_pos"]),
            "best_span_len": None if not best_span else int(best_span["span_len"]),
            "best_span_score": None if not best_span else float(best_span["score"]),
            "best_span_normalized_recovery": None
            if not best_span
            else float(best_span["normalized_recovery"]),
            "best_span_distance_based_recovery": None
            if not best_span
            else float(best_span["distance_based_recovery"]),
            "best_span_distance_reduction_gain": None
            if not best_span
            else float(best_span["distance_reduction_gain"]),
        }

        reverse_summary = {
            "reverse_patch_direction": f"{corrupted_variant}_to_{clean_variant}",
            "reverse_direction": f"{corrupted_variant}_to_{clean_variant}",
            "reverse_score_metric": "logit_margin",
            "reverse_source_role": corrupted_variant,
            "reverse_destination_role": clean_variant,
            "reverse_source_output": reverse_positive_answer,
            "reverse_destination_output": reverse_negative_answer,
            "reverse_baseline_margin": None,
            "reverse_source_margin": None,
            "reverse_best_layer": None,
            "reverse_best_pos": None,
            "reverse_best_patch_unit_id": None,
            "reverse_best_patch_unit_type": None,
            "reverse_best_patch_unit_text": None,
            "reverse_best_unit_id": None,
            "reverse_best_unit_kind": None,
            "reverse_best_positions": [],
            "reverse_best_score": None,
            "reverse_best_normalized_recovery": None,
            "reverse_best_distance_based_recovery": None,
            "reverse_best_distance_reduction_gain": None,
            "reverse_token_labeled_positions": [],
            "reverse_token_labeled_unit_ids": [],
            "reverse_token_labeled_tokens": {},
            "reverse_best_span_id": None,
            "reverse_best_span_positions": [],
            "reverse_best_span_start_pos": None,
            "reverse_best_span_end_pos": None,
            "reverse_best_span_len": None,
            "reverse_best_span_score": None,
            "reverse_best_span_normalized_recovery": None,
            "reverse_best_span_distance_based_recovery": None,
            "reverse_best_span_distance_reduction_gain": None,
        }
        reverse_target_info = self.answer_decision_target_info(
            source_context=corrupted_prompt,
            destination_context=clean_prompt,
            positive_text=reverse_positive_answer,
            negative_text=reverse_negative_answer,
            direction_label="corrupted_into_clean",
        )
        reverse_common_text = str(reverse_target_info["common_text"])
        reverse_corrupted_context = str(
            reverse_target_info.get("source_context_text") or corrupted_prompt
        )
        reverse_clean_context = str(reverse_target_info.get("destination_context_text") or clean_prompt)
        reverse_corrupted_tokens = self.to_tokens(reverse_corrupted_context + reverse_common_text)
        reverse_clean_tokens = self.to_tokens(reverse_clean_context + reverse_common_text)
        with torch.no_grad():
            reverse_corrupted_logits, reverse_corrupted_cache = self.model.run_with_cache(
                reverse_corrupted_tokens
            )
            reverse_clean_logits = self.model(reverse_clean_tokens)

        reverse_positive_token_id = int(reverse_target_info["positive_token_id"])
        reverse_negative_token_id = int(reverse_target_info["negative_token_id"])
        reverse_baseline = self.metric_from_logits_token_margin(
            reverse_clean_logits,
            positive_token_id=reverse_positive_token_id,
            negative_token_id=reverse_negative_token_id,
        )
        reverse_source_ref = self.metric_from_logits_token_margin(
            reverse_corrupted_logits,
            positive_token_id=reverse_positive_token_id,
            negative_token_id=reverse_negative_token_id,
        )
        reverse_scores = self.residual_patch_unit_margin_sweep(
            source_cache=reverse_corrupted_cache,
            destination_tokens=reverse_clean_tokens,
            positive_token_id=reverse_positive_token_id,
            negative_token_id=reverse_negative_token_id,
            patch_units=patch_units,
            direction="reverse",
            desc="Violation -> clean residual patch units",
        )
        reverse_distance_recovery_scores = distance_based_recovery_matrix(
            reverse_scores,
            reverse_baseline,
            reverse_source_ref,
        )
        reverse_selected_recovery_scores = recovery_matrix_for_metric(
            recovery_metric,
            reverse_scores,
            reverse_baseline,
            reverse_source_ref,
        )
        reverse_high_recovery_unit_ids = self.positions_over_recovery_threshold(
            recovery_scores=reverse_selected_recovery_scores,
            positions=swept_unit_ids,
            threshold=token_label_recovery_threshold,
        )
        reverse_bold_recovery_unit_ids = self.positions_over_abs_recovery_threshold(
            recovery_scores=reverse_selected_recovery_scores,
            positions=swept_unit_ids,
            threshold=0.2,
        )
        reverse_emphasized_recovery_unit_ids = sorted(
            set(reverse_high_recovery_unit_ids) | set(reverse_bold_recovery_unit_ids)
        )
        reverse_unit_label_by_id = {
            int(unit_id): format_patch_unit_axis_label(unit_by_id[int(unit_id)], "reverse")
            for unit_id in swept_unit_ids
            if int(unit_id) in unit_by_id
        }

        self.plot_heatmap(
            reverse_scores,
            positions=swept_unit_ids,
            colorbar_label="answer logit margin",
            output_path=pair_output_dir / "reverse_residual_scores_heatmap.pdf",
            token_labels=reverse_unit_label_by_id,
            token_label_bold_positions=reverse_bold_recovery_unit_ids,
            forced_tick_positions=reverse_emphasized_recovery_unit_ids,
            xlabel="",
            cmap=DESTINATION_LOGIT_MARGIN_CMAP,
        )
        reverse_recovery_plot = recovery_heatmap_kwargs(
            recovery_metric,
            reverse_selected_recovery_scores,
            CORRUPTED_RECOVERY_CMAP,
        )
        self.plot_heatmap(
            reverse_selected_recovery_scores,
            positions=swept_unit_ids,
            colorbar_label=recovery_metric_column(recovery_metric).replace("_", " "),
            output_path=pair_output_dir / "reverse_residual_recovery_heatmap.pdf",
            token_labels=reverse_unit_label_by_id,
            token_label_bold_positions=reverse_bold_recovery_unit_ids,
            forced_tick_positions=reverse_emphasized_recovery_unit_ids,
            xlabel="",
            vmin=reverse_recovery_plot["vmin"],
            vmax=reverse_recovery_plot["vmax"],
            cmap=reverse_recovery_plot["cmap"],
        )

        reverse_score_df, reverse_top_sites_df = self.build_residual_score_tables(
            scores=reverse_scores,
            patch_units=patch_units,
            direction="reverse",
            baseline=reverse_baseline,
            source_ref=reverse_source_ref,
            distance_column="distance_to_violation_ref",
            ranking_metric=recovery_metric,
        )
        reverse_source_prompt = pair.get("prompts", {}).get(corrupted_variant)
        reverse_destination_prompt = pair.get("prompts", {}).get(clean_variant)
        reverse_score_metadata = {
            "score_metric": "logit_margin",
            "pair_id": pair_id,
            "pair_family": pair_family,
            "patch_direction": reverse_direction_name,
            "source_role": corrupted_variant,
            "destination_role": clean_variant,
            "source_file": "" if reverse_source_prompt is None else reverse_source_prompt.file,
            "destination_file": ""
            if reverse_destination_prompt is None
            else reverse_destination_prompt.file,
            "source_execution_output": corrupted_expected_output,
            "source_cue_output": ""
            if reverse_source_prompt is None
            else reverse_source_prompt.cue_output,
            "destination_execution_output": clean_expected_output,
            "destination_cue_output": ""
            if reverse_destination_prompt is None
            else reverse_destination_prompt.cue_output,
            "source_output": reverse_positive_answer,
            "destination_output": reverse_negative_answer,
            "source_output_label": ""
            if reverse_source_prompt is None
            else classify_behavior(
                reverse_positive_answer,
                reverse_source_prompt.execution_output,
                reverse_source_prompt.cue_output,
            ),
            "destination_output_label": ""
            if reverse_destination_prompt is None
            else classify_behavior(
                reverse_negative_answer,
                reverse_destination_prompt.execution_output,
                reverse_destination_prompt.cue_output,
            ),
            "source_is_accurate": ""
            if reverse_source_prompt is None
            else str(reverse_positive_answer == reverse_source_prompt.execution_output),
            "destination_is_accurate": ""
            if reverse_destination_prompt is None
            else str(reverse_negative_answer == reverse_destination_prompt.execution_output),
            "source_token": reverse_target_info["positive_token_str"],
            "destination_token": reverse_target_info["negative_token_str"],
            "source_token_id": reverse_positive_token_id,
            "destination_token_id": reverse_negative_token_id,
            "divergent_token_index": int(reverse_target_info["divergent_token_index"]),
            "common_target_prefix": reverse_common_text,
            "target_selection_mode": reverse_target_info.get("target_selection_mode", ""),
            "answer_terminator_text": reverse_target_info.get("answer_terminator_text", ""),
            "baseline_margin": float(reverse_baseline),
            "source_margin": float(reverse_source_ref),
        }
        for df in (reverse_score_df, reverse_top_sites_df):
            for key, value in reversed(list(reverse_score_metadata.items())):
                df.insert(0, key, value)
        write_methodology_residual_scores(
            reverse_score_df,
            pair_output_dir / "reverse_residual_scores.csv",
        )
        write_methodology_residual_scores(
            reverse_top_sites_df,
            pair_output_dir / "top_reverse_residual_sites.csv",
        )

        reverse_best = (
            reverse_top_sites_df.iloc[0].to_dict()
            if not reverse_top_sites_df.empty
            else {}
        )
        reverse_best_span = best_changed_region_row(reverse_top_sites_df)
        reverse_summary.update(
            {
                "reverse_positive_token_id": reverse_positive_token_id,
                "reverse_negative_token_id": reverse_negative_token_id,
                "reverse_positive_token_str": reverse_target_info["positive_token_str"],
                "reverse_negative_token_str": reverse_target_info["negative_token_str"],
                "reverse_common_target_prefix": reverse_common_text,
                "reverse_divergent_token_index": int(reverse_target_info["divergent_token_index"]),
                "reverse_target_selection_mode": reverse_target_info.get("target_selection_mode", ""),
                "reverse_answer_terminator_text": reverse_target_info.get("answer_terminator_text", ""),
                "reverse_baseline_margin": float(reverse_baseline),
                "reverse_source_margin": float(reverse_source_ref),
                "reverse_source_ref_margin": float(reverse_source_ref),
                "reverse_best_layer": None
                if not reverse_best
                else int(reverse_best["layer"]),
                "reverse_best_pos": None
                if not reverse_best
                else int(reverse_best["pos"]),
                "reverse_best_patch_unit_id": None
                if not reverse_best
                else int(reverse_best["unit_id"]),
                "reverse_best_patch_unit_type": None
                if not reverse_best
                else str(reverse_best["unit_kind"]),
                "reverse_best_patch_unit_text": None
                if not reverse_best
                else str(reverse_best["patch_unit_text"]),
                "reverse_best_unit_id": None
                if not reverse_best
                else int(reverse_best["unit_id"]),
                "reverse_best_unit_kind": None
                if not reverse_best
                else str(reverse_best["unit_kind"]),
                "reverse_best_positions": json_positions_from_row(reverse_best),
                "reverse_best_score": None
                if not reverse_best
                else float(reverse_best["score"]),
                "reverse_best_normalized_recovery": None
                if not reverse_best
                else float(reverse_best["normalized_recovery"]),
                "reverse_best_distance_based_recovery": None
                if not reverse_best
                else float(reverse_best["distance_based_recovery"]),
                "reverse_best_distance_reduction_gain": None
                if not reverse_best
                else float(reverse_best["distance_reduction_gain"]),
                "reverse_token_labeled_positions": [
                    int(unit_id) for unit_id in reverse_emphasized_recovery_unit_ids
                ],
                "reverse_token_labeled_unit_ids": [
                    int(unit_id) for unit_id in reverse_emphasized_recovery_unit_ids
                ],
                "reverse_token_label_bold_abs_threshold": 0.2,
                "reverse_token_label_bold_unit_ids": [
                    int(unit_id) for unit_id in reverse_bold_recovery_unit_ids
                ],
                "reverse_token_labeled_tokens": {
                    str(unit_id): reverse_unit_label_by_id[int(unit_id)]
                    for unit_id in reverse_emphasized_recovery_unit_ids
                    if int(unit_id) in reverse_unit_label_by_id
                },
                "reverse_best_span_id": None
                if not reverse_best_span
                else int(reverse_best_span["unit_id"]),
                "reverse_best_span_positions": json_positions_from_row(reverse_best_span),
                "reverse_best_span_start_pos": None
                if not reverse_best_span
                else int(reverse_best_span["start_pos"]),
                "reverse_best_span_end_pos": None
                if not reverse_best_span
                else int(reverse_best_span["end_pos"]),
                "reverse_best_span_len": None
                if not reverse_best_span
                else int(reverse_best_span["span_len"]),
                "reverse_best_span_score": None
                if not reverse_best_span
                else float(reverse_best_span["score"]),
                "reverse_best_span_normalized_recovery": None
                if not reverse_best_span
                else float(reverse_best_span["normalized_recovery"]),
                "reverse_best_span_distance_based_recovery": None
                if not reverse_best_span
                else float(reverse_best_span["distance_based_recovery"]),
                "reverse_best_span_distance_reduction_gain": None
                if not reverse_best_span
                else float(reverse_best_span["distance_reduction_gain"]),
            }
        )

        summary = {
            "status": "completed",
            "pair_id": pair_id,
            "pair_family": pair_family,
            "pair": pair_id,
            "phenomenon": pair["meta"].get("phenomenon", ""),
            "patch_direction": forward_direction_name,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "source_output": target_answer,
            "destination_output": wrong_answer,
            "source_execution_output": clean_expected_output,
            "destination_execution_output": corrupted_expected_output,
            "source_cue_output": "" if source_prompt is None else source_prompt.cue_output,
            "destination_cue_output": ""
            if destination_prompt is None
            else destination_prompt.cue_output,
            "source_output_label": ""
            if source_prompt is None
            else classify_behavior(target_answer, source_prompt.execution_output, source_prompt.cue_output),
            "destination_output_label": ""
            if destination_prompt is None
            else classify_behavior(
                wrong_answer,
                destination_prompt.execution_output,
                destination_prompt.cue_output,
            ),
            "source_is_accurate": ""
            if source_prompt is None
            else bool(target_answer == source_prompt.execution_output),
            "destination_is_accurate": ""
            if destination_prompt is None
            else bool(wrong_answer == destination_prompt.execution_output),
            "answer_resolution_schema": answer_spec.answer_resolution_schema,
            "clean_positive_answer": answer_spec.clean_positive_answer,
            "clean_negative_answer": answer_spec.clean_negative_answer,
            "corrupted_positive_answer": answer_spec.corrupted_positive_answer,
            "corrupted_negative_answer": answer_spec.corrupted_negative_answer,
            "forward_direction": f"{clean_variant}_into_{corrupted_variant}",
            "answer_score_metric": "logit_margin",
            "forward_positive_answer": target_answer,
            "forward_negative_answer": wrong_answer,
            "forward_positive_token_id": forward_positive_token_id,
            "forward_negative_token_id": forward_negative_token_id,
            "forward_positive_token_str": forward_target_info["positive_token_str"],
            "forward_negative_token_str": forward_target_info["negative_token_str"],
            "forward_common_target_prefix": forward_common_text,
            "forward_divergent_token_index": int(forward_target_info["divergent_token_index"]),
            "forward_target_selection_mode": forward_target_info.get("target_selection_mode", ""),
            "forward_answer_terminator_text": forward_target_info.get("answer_terminator_text", ""),
            "clean_expected_output": clean_expected_output,
            "corrupted_expected_output": corrupted_expected_output,
            "clean_seq_len": int(clean_tokens.shape[1]),
            "corrupted_seq_len": int(corrupted_tokens.shape[1]),
            "aligned_seq_len": int(aligned_seq_len),
            "first_changed_pos": None if first_changed_pos is None else int(first_changed_pos),
            "patching_unit_mode": "mixed_token_region",
            "patching_prompt_region": "code_and_output_cue",
            "patch_batch_size": int(self.patch_batch_size),
            "swept_unit_ids": swept_unit_ids,
            "swept_positions": [
                int(pos)
                for pos in (
                    representative_pos(unit, "forward") for unit in patch_units
                )
                if pos is not None
            ],
            "n_swept_patch_units": int(len(patch_units)),
            "n_swept_positions": int(len(patch_units)),
            "n_patch_units": int(len(patchable_units)),
            "baseline_margin": float(baseline),
            "source_margin": float(clean_ref),
            "baseline_proxy_score": float(baseline),
            "source_ref_margin": float(clean_ref),
            "clean_ref_proxy_score": float(clean_ref),
            "best_layer": None if not best else int(best["layer"]),
            "best_pos": None if not best else int(best["pos"]),
            "best_patch_unit_id": None if not best else int(best["unit_id"]),
            "best_patch_unit_type": None if not best else str(best["unit_kind"]),
            "best_patch_unit_text": None if not best else str(best["patch_unit_text"]),
            "best_unit_id": None if not best else int(best["unit_id"]),
            "best_unit_kind": None if not best else str(best["unit_kind"]),
            "best_positions": json_positions_from_row(best),
            "best_score": None if not best else float(best["score"]),
            "best_patched_margin": None if not best else float(best["patched_margin"]),
            "best_signed_recovery": None
            if not best
            else float(best["signed_recovery"]),
            "best_distance_recovery": None
            if not best
            else float(best["distance_recovery"]),
            "best_normalized_recovery": None
            if not best
            else float(best["normalized_recovery"]),
            "best_distance_based_recovery": None
            if not best
            else float(best["distance_based_recovery"]),
            "best_distance_reduction_gain": None
            if not best
            else float(best["distance_reduction_gain"]),
            "token_label_recovery_threshold": float(token_label_recovery_threshold),
            "token_label_recovery_metric": recovery_metric_column(recovery_metric),
            "token_label_bold_abs_threshold": 0.2,
            "token_labeled_positions": [int(unit_id) for unit_id in emphasized_recovery_unit_ids],
            "token_labeled_unit_ids": [int(unit_id) for unit_id in emphasized_recovery_unit_ids],
            "token_label_bold_unit_ids": [int(unit_id) for unit_id in bold_recovery_unit_ids],
            "token_labeled_tokens": {
                str(unit_id): unit_label_by_id[int(unit_id)]
                for unit_id in emphasized_recovery_unit_ids
                if int(unit_id) in unit_label_by_id
            },
            "output_dir": str(pair_output_dir),
        }
        summary.update(span_summary)
        summary.update(reverse_summary)

        additional_generations: Optional[AdditionalGenerations] = None
        if save_additional_tasks or unit_test_flow_patching:
            try:
                additional_generations = self.build_additional_generations(
                    pair_name=pair_name,
                    target_answer=target_answer,
                    destination_reference_output=wrong_answer,
                    clean_expected_output=clean_expected_output,
                    corrupted_expected_output=corrupted_expected_output,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    answer_max_new_tokens=answer_max_new_tokens,
                    unit_test_max_new_tokens=unit_test_max_new_tokens,
                )
            except Exception as e:
                summary["additional_generations_error"] = str(e)

        clean_prompt_parts = split_fenced_prompt(clean_prompt)
        corrupted_prompt_parts = split_fenced_prompt(corrupted_prompt)
        output_contexts = {
            forward_direction_name: {
                "patch_direction": forward_direction_name,
                "destination_role": corrupted_variant,
                "prompt_text": corrupted_prompt_parts["prompt_text"],
                "snippet_text": corrupted_prompt_parts["snippet_text"] or corrupted_code.rstrip(),
                "output_label": prompt_output_label(corrupted_prompt_parts, "Generated Answer"),
                "generated_output": ""
                if additional_generations is None
                else additional_generations.corrupted_model_answer,
                "generated_prefix": "",
                "common_target_prefix": forward_common_text,
                "source_output": target_answer,
                "destination_output": wrong_answer,
                "source_token": forward_target_info["positive_token_str"],
                "destination_token": forward_target_info["negative_token_str"],
                "next_token_scored": (
                    f"{forward_target_info['positive_token_str']} vs "
                    f"{forward_target_info['negative_token_str']}"
                ),
            },
            reverse_direction_name: {
                "patch_direction": reverse_direction_name,
                "destination_role": clean_variant,
                "prompt_text": clean_prompt_parts["prompt_text"],
                "snippet_text": clean_prompt_parts["snippet_text"] or clean_code.rstrip(),
                "output_label": prompt_output_label(clean_prompt_parts, "Generated Answer"),
                "generated_output": ""
                if additional_generations is None
                else additional_generations.clean_model_answer,
                "generated_prefix": "",
                "common_target_prefix": reverse_common_text,
                "source_output": reverse_positive_answer,
                "destination_output": reverse_negative_answer,
                "source_token": reverse_target_info["positive_token_str"],
                "destination_token": reverse_target_info["negative_token_str"],
                "next_token_scored": (
                    f"{reverse_target_info['positive_token_str']} vs "
                    f"{reverse_target_info['negative_token_str']}"
                ),
            },
        }
        (pair_output_dir / "output_contexts.json").write_text(
            json.dumps(output_contexts, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        summary["output_contexts"] = str(pair_output_dir / "output_contexts.json")

        if additional_generations is not None:
            try:
                generation_json_path = write_additional_generations_json(
                    pair_output_dir / "unit_test_generations.json",
                    additional_generations,
                )
                summary["unit_test_generations"] = str(generation_json_path)
                attach_task1_outputs(
                    summary,
                    pair,
                    clean_variant,
                    corrupted_variant,
                    additional_generations,
                )
            except Exception as e:
                summary["unit_test_generations_error"] = str(e)

        if save_additional_tasks and additional_generations is not None:
            try:
                task_dump_path = self.write_additional_generation_dump(
                    output_path=pair_output_dir / "additional_task_generations.txt",
                    generations=additional_generations,
                )
                summary["additional_task_dump"] = str(task_dump_path)
            except Exception as e:
                summary["additional_task_dump_error"] = str(e)

        if unit_test_flow_patching and additional_generations is not None:
            try:
                flow_summary = self.run_unit_test_flow_patching(
                    pair_output_dir=pair_output_dir,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    pair_id=str(pair.get("pair_id") or pair_name),
                    pair_family=pair_family,
                    additional_generations=additional_generations,
                    max_positions=max_positions,
                    pos_context=pos_context,
                    token_label_recovery_threshold=token_label_recovery_threshold,
                )
                summary["unit_test_flow"] = flow_summary
                summary["unit_test_flow_parsed_assertions"] = int(
                    flow_summary.get("parsed_assertions", 0)
                )
                summary["unit_test_flow_analyzed_forward_decisions"] = int(
                    flow_summary.get("analyzed_forward_decisions", 0)
                )
                summary["unit_test_flow_analyzed_reverse_decisions"] = int(
                    flow_summary.get("analyzed_reverse_decisions", 0)
                )
                summary["unit_test_flow_skipped_decisions"] = int(
                    flow_summary.get("skipped_decisions", 0)
                )
                summary["unit_test_flow_best_prompt_to_output_site"] = json.dumps(
                    flow_summary.get("best_prompt_to_output_site", {})
                )
                summary["unit_test_flow_best_response_carry_site"] = json.dumps(
                    flow_summary.get("best_response_carry_site", {})
                )
                summary["unit_test_flow_best_combined_context_site"] = json.dumps(
                    flow_summary.get("best_combined_context_site", {})
                )
            except Exception as e:
                summary["unit_test_flow_error"] = str(e)

        (pair_output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
        return summary

    def run_pair(
        self,
        snippet_dir: Path,
        output_dir: Path,
        pair_name: str,
        clean_variant: str,
        corrupted_variant: str,
        max_positions: Optional[int],
        pos_context: int,
        use_gold_answer: bool,
        manual_target_answer: Optional[str],
        manual_wrong_answer: Optional[str],
        force_target_to_clean_output: bool,
        auto_wrong_from_corrupted_output: bool,
        save_additional_tasks: bool,
        unit_test_flow_patching: bool,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
        token_label_recovery_threshold: float,
        recovery_metric: str,
        pair_output_dir_override: Optional[Path] = None,
    ) -> Dict:
        pair = load_pair(snippet_dir, pair_name)
        clean_code = get_variant_code(pair, clean_variant)
        corrupted_code = get_variant_code(pair, corrupted_variant)
        clean_prompt = build_prompt(clean_code, self.question_template)
        corrupted_prompt = build_prompt(corrupted_code, self.question_template)
        answer_spec = resolve_answer_decision_spec(
            pair=pair,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            use_gold_answer=use_gold_answer,
            manual_target_answer=manual_target_answer,
            manual_wrong_answer=manual_wrong_answer,
            force_target_to_clean_output=force_target_to_clean_output,
            auto_wrong_from_corrupted_output=auto_wrong_from_corrupted_output,
        )
        if answer_spec.forward_negative_answer is None:
            raise ValueError(
                "Answer patching uses a contrastive logit margin and requires a wrong answer. "
                "Provide --manual-wrong-answer, metadata wrong_answer, or metadata wrong_answers."
            )
        if answer_spec.reverse_positive_answer is None or answer_spec.reverse_negative_answer is None:
            raise ValueError(
                "Reverse answer patching uses a contrastive logit margin and requires positive "
                "and negative answers for the corrupted/source variant. Provide metadata "
                "answers/wrong_answers for both variants or explicit answer_decisions."
            )

        clean_tokens = self.to_tokens(clean_prompt)
        corrupted_tokens = self.to_tokens(corrupted_prompt)
        aligned_seq_len = int(min(clean_tokens.shape[1], corrupted_tokens.shape[1]))
        if aligned_seq_len == 0:
            raise ValueError("Tokenization produced empty prompts.")

        clean_ids = clean_tokens[0].detach().cpu().tolist()
        corrupted_ids = corrupted_tokens[0].detach().cpu().tolist()
        clean_strs = self.token_strs(clean_tokens)
        corrupted_strs = self.token_strs(corrupted_tokens)
        all_patch_units, alignment_rows, first_changed_pos = build_alignment_patch_units(
            clean_ids=clean_ids,
            corrupted_ids=corrupted_ids,
            clean_strs=clean_strs,
            corrupted_strs=corrupted_strs,
        )
        patchable_units = patch_units_from_first_change(
            all_patch_units,
            first_changed_pos=first_changed_pos,
            context_window=pos_context,
        )
        clean_patching_region_positions = token_positions_overlapping_span(
            clean_prompt,
            clean_strs,
            fenced_code_and_output_cue_span(clean_prompt),
        )
        corrupted_patching_region_positions = token_positions_overlapping_span(
            corrupted_prompt,
            corrupted_strs,
            fenced_code_and_output_cue_span(corrupted_prompt),
        )
        prompt_patchable_units = filter_patch_units_to_position_sets(
            patchable_units,
            clean_patching_region_positions,
            corrupted_patching_region_positions,
        )
        if not prompt_patchable_units:
            raise ValueError("No patchable code/output-cue prompt token units were found.")
        sampled_patch_units = sample_patch_units(prompt_patchable_units, max_positions)
        if any(unit.is_asymmetric for unit in changed_region_units(prompt_patchable_units)):
            summary = self.save_tokenization_mismatch_skip_artifacts(
                pair_name=pair_name,
                pair=pair,
                clean_code=clean_code,
                corrupted_code=corrupted_code,
                clean_variant=clean_variant,
                corrupted_variant=corrupted_variant,
                answer_spec=answer_spec,
                clean_tokens=clean_tokens,
                corrupted_tokens=corrupted_tokens,
                aligned_seq_len=aligned_seq_len,
                all_patch_units=all_patch_units,
                patchable_units=prompt_patchable_units,
                alignment_rows=alignment_rows,
                first_changed_pos=first_changed_pos,
                output_dir=output_dir,
                pair_output_dir_override=pair_output_dir_override,
            )
            pair_output_dir = pair_output_dir_override or (output_dir / slugify(pair_name))
            additional_generations: Optional[AdditionalGenerations] = None
            if save_additional_tasks or unit_test_flow_patching:
                try:
                    additional_generations = self.build_additional_generations(
                        pair_name=pair_name,
                        target_answer=answer_spec.forward_positive_answer,
                        destination_reference_output=answer_spec.forward_negative_answer,
                        clean_expected_output=answer_spec.clean_expected_output,
                        corrupted_expected_output=answer_spec.corrupted_expected_output,
                        clean_code=clean_code,
                        corrupted_code=corrupted_code,
                        clean_variant=clean_variant,
                        corrupted_variant=corrupted_variant,
                        answer_max_new_tokens=answer_max_new_tokens,
                        unit_test_max_new_tokens=unit_test_max_new_tokens,
                    )
                except Exception as e:
                    summary["additional_generations_error"] = str(e)
            if save_additional_tasks and additional_generations is not None:
                try:
                    task_dump_path = self.write_additional_generation_dump(
                        output_path=pair_output_dir / "additional_task_generations.txt",
                        generations=additional_generations,
                    )
                    summary["additional_task_dump"] = str(task_dump_path)
                except Exception as e:
                    summary["additional_task_dump_error"] = str(e)
            if unit_test_flow_patching and additional_generations is not None:
                try:
                    flow_summary = self.run_unit_test_flow_patching(
                        pair_output_dir=pair_output_dir,
                        clean_code=clean_code,
                        corrupted_code=corrupted_code,
                        clean_variant=clean_variant,
                        corrupted_variant=corrupted_variant,
                        pair_id=str(pair.get("pair_id") or pair_name),
                        pair_family=str(pair.get("pair_family") or pair["meta"].get("pair_family", "")),
                        additional_generations=additional_generations,
                        max_positions=max_positions,
                        pos_context=pos_context,
                        token_label_recovery_threshold=token_label_recovery_threshold,
                    )
                    summary["unit_test_flow"] = flow_summary
                    summary["unit_test_flow_parsed_assertions"] = int(
                        flow_summary.get("parsed_assertions", 0)
                    )
                    summary["unit_test_flow_analyzed_forward_decisions"] = int(
                        flow_summary.get("analyzed_forward_decisions", 0)
                    )
                    summary["unit_test_flow_analyzed_reverse_decisions"] = int(
                        flow_summary.get("analyzed_reverse_decisions", 0)
                    )
                    summary["unit_test_flow_skipped_decisions"] = int(
                        flow_summary.get("skipped_decisions", 0)
                    )
                    summary["unit_test_flow_best_prompt_to_output_site"] = json.dumps(
                        flow_summary.get("best_prompt_to_output_site", {})
                    )
                    summary["unit_test_flow_best_response_carry_site"] = json.dumps(
                        flow_summary.get("best_response_carry_site", {})
                    )
                    summary["unit_test_flow_best_combined_context_site"] = json.dumps(
                        flow_summary.get("best_combined_context_site", {})
                    )
                except Exception as e:
                    summary["unit_test_flow_error"] = str(e)
            (pair_output_dir / "summary.json").write_text(
                json.dumps(summary, indent=2) + "\n",
                encoding="utf-8",
            )
            return summary

        with torch.no_grad():
            clean_logits, clean_cache = self.model.run_with_cache(clean_tokens)
            corrupted_logits, corrupted_cache = self.model.run_with_cache(corrupted_tokens)

        return self.save_patch_unit_pair_artifacts(
            pair_name=pair_name,
            pair=pair,
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            clean_prompt=clean_prompt,
            corrupted_prompt=corrupted_prompt,
            answer_spec=answer_spec,
            clean_tokens=clean_tokens,
            corrupted_tokens=corrupted_tokens,
            clean_logits=clean_logits,
            corrupted_logits=corrupted_logits,
            clean_cache=clean_cache,
            corrupted_cache=corrupted_cache,
            aligned_seq_len=aligned_seq_len,
            patch_units=sampled_patch_units,
            patchable_units=prompt_patchable_units,
            alignment_rows=alignment_rows,
            first_changed_pos=first_changed_pos,
            output_dir=output_dir,
            save_additional_tasks=save_additional_tasks,
            unit_test_flow_patching=unit_test_flow_patching,
            answer_max_new_tokens=answer_max_new_tokens,
            unit_test_max_new_tokens=unit_test_max_new_tokens,
            token_label_recovery_threshold=token_label_recovery_threshold,
            recovery_metric=recovery_metric,
            max_positions=max_positions,
            pos_context=pos_context,
            pair_output_dir_override=pair_output_dir_override,
        )

    def write_rq1_1_outputs(
        self,
        stage_dir: Path,
        pair: Dict[str, Any],
        pair_name: str,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        answer_max_new_tokens: int,
    ) -> Dict[str, Any]:
        ensure_stage_dir(stage_dir)
        (
            source_generation,
            destination_generation,
            source_output,
            destination_output,
        ) = self.build_answer_generations(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            answer_max_new_tokens=answer_max_new_tokens,
        )
        pair_id = str(pair.get("pair_id") or pair_name)
        pair_family = str(pair.get("pair_family") or pair["meta"].get("pair_family", ""))
        prompt_specs = pair.get("prompts", {})
        sides = [
            ("source", clean_variant, source_generation, source_output),
            ("destination", corrupted_variant, destination_generation, destination_output),
        ]
        behavior_rows: List[Dict[str, Any]] = []
        generations: Dict[str, Any] = {}
        prompts: Dict[str, Any] = {}
        for side, role, generation, model_output in sides:
            prompt_spec = prompt_specs.get(role)
            execution_output = "" if prompt_spec is None else str(prompt_spec.execution_output)
            cue_output = "" if prompt_spec is None else str(prompt_spec.cue_output)
            label = (
                classify_behavior(model_output, execution_output, cue_output)
                if execution_output and cue_output
                else ""
            )
            is_correct = label in {"both_consistent", "execution_consistent"} if label else None
            generations[side] = {
                "task": "final_output",
                "source_role": clean_variant,
                "destination_role": corrupted_variant,
                "side": side,
                "role": role,
                "max_new_tokens": int(answer_max_new_tokens),
                "prompt": generation.prompt,
                "full_text": generation.full_text,
                "completion": generation.completion,
                "model_output_raw": extract_first_fenced_block_or_text(generation.completion),
                "model_output_clean": model_output,
                "model_output_cleaning_applied": bool(
                    extract_first_fenced_block_or_text(generation.completion) != model_output
                ),
            }
            prompts[side] = {
                "role": role,
                "prompt": generation.prompt,
                "code": clean_code if side == "source" else corrupted_code,
                "execution_output": execution_output,
                "cue_output": cue_output,
            }
            behavior_rows.append(
                {
                    "pair_id": pair_id,
                    "pair_family": pair_family,
                    "side": side,
                    "role": role,
                    "role_kind": role_kind(role),
                    "model_output": model_output,
                    "model_output_raw": generations[side]["model_output_raw"],
                    "model_output_cleaning_applied": generations[side][
                        "model_output_cleaning_applied"
                    ],
                    "execution_output": execution_output,
                    "cue_output": cue_output,
                    "response_label": label,
                    "is_correct": "" if is_correct is None else bool(is_correct),
                    "is_incorrect_misleading": bool(label == "cue_consistent") if label else "",
                }
            )

        payload = {
            "stage": "rq1_1",
            "pair_name": pair_name,
            "pair_id": pair_id,
            "pair_family": pair_family,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "generations": generations,
        }
        write_json_dict(stage_dir / "generations.json", payload)
        write_json_dict(
            stage_dir / "behavior.json",
            {
                "stage": "rq1_1",
                "pair_id": pair_id,
                "pair_family": pair_family,
                "rows": behavior_rows,
            },
        )
        write_json_dict(
            stage_dir / "prompts.json",
            {
                "stage": "rq1_1",
                "pair_id": pair_id,
                "pair_family": pair_family,
                "prompts": prompts,
            },
        )
        return payload

    def write_rq2_1_outputs(
        self,
        stage_dir: Path,
        pair: Dict[str, Any],
        pair_name: str,
        clean_code: str,
        corrupted_code: str,
        clean_variant: str,
        corrupted_variant: str,
        unit_test_max_new_tokens: int,
    ) -> AdditionalGenerations:
        ensure_stage_dir(stage_dir)
        unit_tests = self.build_unit_test_generations(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            unit_test_max_new_tokens=unit_test_max_new_tokens,
        )
        pair_id = str(pair.get("pair_id") or pair_name)
        pair_family = str(pair.get("pair_family") or pair["meta"].get("pair_family", ""))
        generation_payload = {
            "stage": "rq2_1",
            "pair_name": pair_name,
            "pair_id": pair_id,
            "pair_family": pair_family,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "generations": {
                key: generation_to_dict(value)
                for key, value in sorted(unit_tests.items())
            },
        }
        write_json_dict(stage_dir / "generations.json", generation_payload)

        source_generation = unit_tests.get("source")
        destination_generation = unit_tests.get("destination")
        assertion_rows: List[Dict[str, Any]] = []
        behavior_rows: List[Dict[str, Any]] = []
        skip_rows: List[Dict[str, Any]] = []
        if source_generation is not None and destination_generation is not None:
            decisions, skip_rows, assertion_counts = extract_unit_test_decisions(
                clean_code=clean_code,
                corrupted_code=corrupted_code,
                clean_test_code=source_generation.extracted_output,
                corrupted_test_code=destination_generation.extracted_output,
            )
            for decision in decisions:
                destination_is_accurate = (
                    str(decision["corrupted_expected_source"])
                    == str(decision["corrupted_runtime_source"])
                )
                response_label = ""
                response_label_status = "skipped_no_independent_cue_expected_value"
                if pair_family != "cue_varied":
                    response_label = classify_behavior(
                        decision["corrupted_expected_source"],
                        decision["corrupted_runtime_source"],
                        decision["clean_runtime_source"],
                    )
                    response_label_status = "classified"
                row = {
                    "decision_id": decision["decision_id"],
                    "assertion_index": int(decision["assertion_index"]),
                    "pair_id": pair_id,
                    "pair_family": pair_family,
                    "source_role": clean_variant,
                    "destination_role": corrupted_variant,
                    "assertion_expression": decision["corrupted_expr_source"],
                    "generated_expected_value": decision["corrupted_expected_source"],
                    "execution_expected_value": decision["corrupted_runtime_source"],
                    "cue_expected_value": decision["clean_runtime_source"],
                    "is_accurate": bool(destination_is_accurate),
                    "response_label": response_label,
                    "response_label_status": response_label_status,
                    "target_mode": decision["target_mode"],
                    "clean_match_type": decision["clean_match_type"],
                    "reverse_available": bool(decision["reverse_available"]),
                }
                assertion_rows.append(row)
                behavior_rows.append(row)
        else:
            assertion_counts = {"clean_assertions": 0, "corrupted_assertions": 0}
            skip_rows.append({"reason": "missing_source_or_destination_unit_test_generation"})

        write_dict_rows_csv(stage_dir / "assertions.csv", assertion_rows)
        write_dict_rows_csv(stage_dir / "decision_skips.csv", skip_rows)
        write_json_dict(
            stage_dir / "behavior.json",
            {
                "stage": "rq2_1",
                "pair_id": pair_id,
                "pair_family": pair_family,
                "source_role": clean_variant,
                "destination_role": corrupted_variant,
                "parsed_source_assertions": int(assertion_counts.get("clean_assertions", 0)),
                "parsed_destination_assertions": int(
                    assertion_counts.get("corrupted_assertions", 0)
                ),
                "rows": behavior_rows,
            },
        )
        empty_answer_generation = AdditionalGeneration(
            task="final_output",
            patch_role="",
            prompt_role="",
            max_new_tokens=0,
            prompt="",
            full_text="",
            completion="",
            extracted_output="",
        )
        return AdditionalGenerations(
            pair_name=pair_name,
            source_role=clean_variant,
            destination_role=corrupted_variant,
            source_reference_output="",
            destination_reference_output="",
            source_execution_output=None,
            destination_execution_output=None,
            source_model_output="",
            destination_model_output="",
            source_answer_generation=empty_answer_generation,
            destination_answer_generation=empty_answer_generation,
            unit_tests=unit_tests,
        )

    def load_rq2_1_generations(self, stage_dir: Path, pair_name: str) -> AdditionalGenerations:
        payload = load_json_dict(stage_dir / "generations.json")
        if payload is None:
            raise FileNotFoundError(f"Missing RQ2.1 generations: {stage_dir / 'generations.json'}")
        unit_tests_payload = payload.get("generations")
        if not isinstance(unit_tests_payload, dict):
            raise ValueError(f"Invalid RQ2.1 generations payload: {stage_dir / 'generations.json'}")
        unit_tests = {
            str(key): cap_unit_test_generation(
                generation_from_dict(value if isinstance(value, dict) else {})
            )
            for key, value in unit_tests_payload.items()
        }
        empty_answer_generation = AdditionalGeneration(
            task="final_output",
            patch_role="",
            prompt_role="",
            max_new_tokens=0,
            prompt="",
            full_text="",
            completion="",
            extracted_output="",
        )
        return AdditionalGenerations(
            pair_name=pair_name,
            source_role=str(payload.get("source_role") or ""),
            destination_role=str(payload.get("destination_role") or ""),
            source_reference_output="",
            destination_reference_output="",
            source_execution_output=None,
            destination_execution_output=None,
            source_model_output="",
            destination_model_output="",
            source_answer_generation=empty_answer_generation,
            destination_answer_generation=empty_answer_generation,
            unit_tests=unit_tests,
        )

    def normalize_rq1_2_outputs(self, stage_dir: Path) -> None:
        frames = []
        top_frames = []
        for name in ("residual_scores.csv", "reverse_residual_scores.csv"):
            path = stage_dir / name
            if path.exists() and path.stat().st_size > 0:
                frames.append(pd.read_csv(path))
        for name in ("top_residual_sites.csv", "top_reverse_residual_sites.csv"):
            path = stage_dir / name
            if path.exists() and path.stat().st_size > 0:
                top_frames.append(pd.read_csv(path))
        if frames:
            write_clean_dataframe_csv(pd.concat(frames, ignore_index=True), stage_dir / "residual_scores.csv")
        if top_frames:
            write_clean_dataframe_csv(pd.concat(top_frames, ignore_index=True), stage_dir / "top_sites.csv")
        for stale_name in (
            "reverse_residual_scores.csv",
            "top_residual_sites.csv",
            "top_reverse_residual_sites.csv",
        ):
            (stage_dir / stale_name).unlink(missing_ok=True)
        summary = load_json_dict(stage_dir / "summary.json")
        if isinstance(summary, dict):
            destination_margin = summary.get("baseline_margin", summary.get("baseline_proxy_score"))
            source_margin = summary.get("source_margin", summary.get("source_ref_margin"))
            for key in (
                "baseline_margin",
                "baseline_proxy_score",
                "source_ref_margin",
                "clean_ref_proxy_score",
                "best_unit_id",
                "best_unit_kind",
                "best_positions",
                "best_normalized_recovery",
                "best_distance_based_recovery",
                "reverse_baseline_margin",
                "reverse_source_ref_margin",
                "reverse_best_unit_id",
                "reverse_best_unit_kind",
                "reverse_best_positions",
                "reverse_best_normalized_recovery",
                "reverse_best_distance_based_recovery",
            ):
                summary.pop(key, None)
            if destination_margin not in ("", None):
                summary["destination_margin"] = destination_margin
            if source_margin not in ("", None):
                summary["source_margin"] = source_margin
            write_json_dict(stage_dir / "summary.json", summary)

    def run_pair_rq_stages(
        self,
        *,
        snippet_dir: Path,
        output_root: Path,
        model_slug: str,
        model_name: str,
        pair_name: str,
        clean_variant: str,
        corrupted_variant: str,
        stages: Set[str],
        overwrite_stages: Set[str],
        max_positions: Optional[int],
        pos_context: int,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
        token_label_recovery_threshold: float,
        recovery_metric: str,
    ) -> Dict[str, Any]:
        pair = load_pair(snippet_dir, pair_name)
        clean_code = get_variant_code(pair, clean_variant)
        corrupted_code = get_variant_code(pair, corrupted_variant)
        pair_id = str(pair.get("pair_id") or pair_name)
        pair_family = str(pair.get("pair_family") or pair["meta"].get("pair_family", ""))
        pair_dir = rq_pair_output_dir(output_root, model_slug, pair_family, pair_id)
        pair_dir.mkdir(parents=True, exist_ok=True)
        metadata = {
            "schema": "rq_stage_v1",
            "model_slug": model_slug,
            "model_name": model_name,
            "pair_name": pair_name,
            "pair_id": pair_id,
            "pair_family": pair_family,
            "source_role": clean_variant,
            "destination_role": corrupted_variant,
            "stages_requested": sorted(stages),
            "overwrite_stages": sorted(overwrite_stages),
        }
        write_json_dict(pair_dir / "metadata.json", metadata)
        manifest = load_json_dict(pair_dir / "manifest.json") or {
            "schema": "rq_stage_v1",
            "model_slug": model_slug,
            "model_name": model_name,
            "pair_id": pair_id,
            "pair_family": pair_family,
            "stages": {},
        }
        manifest["path"] = str(pair_dir / "manifest.json")

        def mark(stage: str, status: str, detail: str = "") -> None:
            manifest.setdefault("stages", {})[stage] = {
                "status": status,
                "detail": detail,
                "path": str(rq_stage_dir(pair_dir, stage)),
            }
            write_json_dict(pair_dir / "manifest.json", manifest)

        def prepare_stage(stage: str) -> Tuple[Path, bool]:
            stage_dir = rq_stage_dir(pair_dir, stage)
            if stage in overwrite_stages:
                clear_stage_dir(stage_dir)
                return stage_dir, True
            ensure_stage_dir(stage_dir)
            if stage_is_complete(stage_dir, stage):
                mark(stage, "reused")
                return stage_dir, False
            return stage_dir, True

        rq2_1_generations: Optional[AdditionalGenerations] = None

        if "rq1_1" in stages:
            stage_dir, should_compute = prepare_stage("rq1_1")
            if should_compute:
                self.write_rq1_1_outputs(
                    stage_dir=stage_dir,
                    pair=pair,
                    pair_name=pair_name,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    answer_max_new_tokens=answer_max_new_tokens,
                )
                mark("rq1_1", "computed")

        if "rq1_2" in stages:
            stage_dir, should_compute = prepare_stage("rq1_2")
            if should_compute:
                self.run_pair(
                    snippet_dir=snippet_dir,
                    output_dir=pair_dir,
                    pair_name=pair_name,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    max_positions=max_positions,
                    pos_context=pos_context,
                    use_gold_answer=True,
                    manual_target_answer=None,
                    manual_wrong_answer=None,
                    force_target_to_clean_output=True,
                    auto_wrong_from_corrupted_output=True,
                    save_additional_tasks=False,
                    unit_test_flow_patching=False,
                    answer_max_new_tokens=answer_max_new_tokens,
                    unit_test_max_new_tokens=unit_test_max_new_tokens,
                    token_label_recovery_threshold=token_label_recovery_threshold,
                    recovery_metric=recovery_metric,
                    pair_output_dir_override=stage_dir,
                )
                self.normalize_rq1_2_outputs(stage_dir)
                mark("rq1_2", "computed")

        if "rq2_1" in stages:
            stage_dir, should_compute = prepare_stage("rq2_1")
            if should_compute:
                rq2_1_generations = self.write_rq2_1_outputs(
                    stage_dir=stage_dir,
                    pair=pair,
                    pair_name=pair_name,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    unit_test_max_new_tokens=unit_test_max_new_tokens,
                )
                mark("rq2_1", "computed")
            else:
                rq2_1_generations = self.load_rq2_1_generations(stage_dir, pair_name)

        if "rq2_2" in stages:
            rq2_1_dir = rq_stage_dir(pair_dir, "rq2_1")
            if rq2_1_generations is None:
                if not stage_is_complete(rq2_1_dir, "rq2_1"):
                    raise FileNotFoundError(
                        "rq2_2 requires rq2_1_unit_test_generation artifacts. "
                        "Run with --stages rq2_1,rq2_2 or provide an existing RQ2.1 stage."
                    )
                rq2_1_generations = self.load_rq2_1_generations(rq2_1_dir, pair_name)
            stage_dir, should_compute = prepare_stage("rq2_2")
            if should_compute:
                self.run_unit_test_flow_patching(
                    pair_output_dir=pair_dir,
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    pair_id=pair_id,
                    pair_family=pair_family,
                    additional_generations=rq2_1_generations,
                    max_positions=max_positions,
                    pos_context=pos_context,
                    token_label_recovery_threshold=token_label_recovery_threshold,
                    flow_dir_override=stage_dir,
                    clean_schema=True,
                )
                mark("rq2_2", "computed")

        return manifest

    def run_pair_task1_only(
        self,
        snippet_dir: Path,
        output_dir: Path,
        pair_name: str,
        clean_variant: str,
        corrupted_variant: str,
        save_additional_tasks: bool,
        answer_max_new_tokens: int,
    ) -> Dict:
        pair = load_pair(snippet_dir, pair_name)
        pair_dir = pair_output_dir(output_dir, pair_name)
        summary = load_existing_pair_summary(output_dir, pair_name)
        if summary is None:
            raise FileNotFoundError(
                f"Task-only final-output reruns require an existing summary.json for {pair_name}."
            )

        clean_code = get_variant_code(pair, clean_variant)
        corrupted_code = get_variant_code(pair, corrupted_variant)
        answer_spec = resolve_answer_decision_spec(
            pair=pair,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            use_gold_answer=True,
            manual_target_answer=None,
            manual_wrong_answer=None,
            force_target_to_clean_output=True,
            auto_wrong_from_corrupted_output=True,
        )
        (
            clean_answer_generation,
            corrupted_answer_generation,
            clean_model_answer,
            corrupted_model_answer,
        ) = self.build_answer_generations(
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            answer_max_new_tokens=answer_max_new_tokens,
        )
        existing_generations = load_additional_generations_json(pair_dir)
        generations = AdditionalGenerations(
            pair_name=pair_name,
            source_role=clean_variant,
            destination_role=corrupted_variant,
            source_reference_output=answer_spec.forward_positive_answer,
            destination_reference_output=answer_spec.forward_negative_answer,
            source_execution_output=answer_spec.clean_expected_output,
            destination_execution_output=answer_spec.corrupted_expected_output,
            source_model_output=clean_model_answer,
            destination_model_output=corrupted_model_answer,
            source_answer_generation=clean_answer_generation,
            destination_answer_generation=corrupted_answer_generation,
            unit_tests={} if existing_generations is None else dict(existing_generations.unit_tests),
        )
        generation_json_path = write_additional_generations_json(
            pair_dir / "unit_test_generations.json",
            generations,
        )
        summary["unit_test_generations"] = str(generation_json_path)
        attach_task1_outputs(summary, pair, clean_variant, corrupted_variant, generations)
        summary["task1_only_rerun"] = True
        if save_additional_tasks:
            task_dump_path = self.write_additional_generation_dump(
                output_path=pair_dir / "additional_task_generations.txt",
                generations=generations,
            )
            summary["additional_task_dump"] = str(task_dump_path)
        (pair_dir / "summary.json").write_text(
            json.dumps(summary, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        return summary

    def run_pair_task2_only(
        self,
        snippet_dir: Path,
        output_dir: Path,
        pair_name: str,
        clean_variant: str,
        corrupted_variant: str,
        max_positions: Optional[int],
        pos_context: int,
        token_label_recovery_threshold: float,
        save_additional_tasks: bool,
        answer_max_new_tokens: int,
        unit_test_max_new_tokens: int,
        regenerate_unit_tests: bool = False,
    ) -> Dict:
        pair = load_pair(snippet_dir, pair_name)
        pair_dir = pair_output_dir(output_dir, pair_name)
        summary = load_existing_pair_summary(output_dir, pair_name)
        if summary is None:
            raise FileNotFoundError(
                f"Task-only unit-test reruns require an existing summary.json for {pair_name}."
            )

        additional_generations = load_additional_generations_json(pair_dir)
        if additional_generations is None:
            clean_code = get_variant_code(pair, clean_variant)
            corrupted_code = get_variant_code(pair, corrupted_variant)
            answer_spec = resolve_answer_decision_spec(
                pair=pair,
                clean_variant=clean_variant,
                corrupted_variant=corrupted_variant,
                clean_code=clean_code,
                corrupted_code=corrupted_code,
                use_gold_answer=True,
                manual_target_answer=None,
                manual_wrong_answer=None,
                force_target_to_clean_output=True,
                auto_wrong_from_corrupted_output=True,
            )
            additional_generations = self.build_additional_generations(
                pair_name=pair_name,
                target_answer=answer_spec.forward_positive_answer,
                destination_reference_output=answer_spec.forward_negative_answer,
                clean_expected_output=answer_spec.clean_expected_output,
                corrupted_expected_output=answer_spec.corrupted_expected_output,
                clean_code=clean_code,
                corrupted_code=corrupted_code,
                clean_variant=clean_variant,
                corrupted_variant=corrupted_variant,
                answer_max_new_tokens=answer_max_new_tokens,
                unit_test_max_new_tokens=unit_test_max_new_tokens,
            )
            generation_json_path = write_additional_generations_json(
                pair_dir / "unit_test_generations.json",
                additional_generations,
            )
            summary["unit_test_generations"] = str(generation_json_path)
            summary["task2_only_generated_unit_tests"] = True
            summary["task2_only_regenerated_unit_tests"] = bool(regenerate_unit_tests)
            if save_additional_tasks:
                task_dump_path = self.write_additional_generation_dump(
                    output_path=pair_dir / "additional_task_generations.txt",
                    generations=additional_generations,
                )
                summary["additional_task_dump"] = str(task_dump_path)
        else:
            clean_code = get_variant_code(pair, clean_variant)
            corrupted_code = get_variant_code(pair, corrupted_variant)
            required_unit_test_keys = {"source", "destination"}
            missing_unit_tests = required_unit_test_keys.difference(
                str(key) for key in additional_generations.unit_tests.keys()
            )
            if regenerate_unit_tests or missing_unit_tests:
                new_unit_tests = self.build_unit_test_generations(
                    clean_code=clean_code,
                    corrupted_code=corrupted_code,
                    clean_variant=clean_variant,
                    corrupted_variant=corrupted_variant,
                    unit_test_max_new_tokens=unit_test_max_new_tokens,
                )
                if regenerate_unit_tests:
                    additional_generations.unit_tests = new_unit_tests
                else:
                    additional_generations.unit_tests.update(new_unit_tests)
                generation_json_path = write_additional_generations_json(
                    pair_dir / "unit_test_generations.json",
                    additional_generations,
                )
                summary["unit_test_generations"] = str(generation_json_path)
                summary["task2_only_generated_unit_tests"] = True
                summary["task2_only_regenerated_unit_tests"] = bool(regenerate_unit_tests)
                if save_additional_tasks:
                    task_dump_path = self.write_additional_generation_dump(
                        output_path=pair_dir / "additional_task_generations.txt",
                        generations=additional_generations,
                    )
                    summary["additional_task_dump"] = str(task_dump_path)
            else:
                summary["task2_only_generated_unit_tests"] = False
                summary["task2_only_regenerated_unit_tests"] = False

        flow_summary = self.run_unit_test_flow_patching(
            pair_output_dir=pair_dir,
            clean_code=clean_code,
            corrupted_code=corrupted_code,
            clean_variant=clean_variant,
            corrupted_variant=corrupted_variant,
            pair_id=str(pair.get("pair_id") or pair_name),
            pair_family=str(pair.get("pair_family") or pair["meta"].get("pair_family", "")),
            additional_generations=additional_generations,
            max_positions=max_positions,
            pos_context=pos_context,
            token_label_recovery_threshold=token_label_recovery_threshold,
        )
        attach_unit_test_flow_summary(summary, flow_summary)
        summary["task2_only_rerun"] = True
        summary["task2_only_rerun_source"] = str(pair_dir / "unit_test_generations.json")
        (pair_dir / "summary.json").write_text(
            json.dumps(summary, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        return summary


def load_pair_names(args, snippet_dir: Path) -> List[str]:
    names: List[str] = []
    if args.all:
        names.extend(list_pairs(snippet_dir, pair_family=getattr(args, "pair_family", None)))
    explicit_names: List[str] = []
    if args.pairs:
        explicit_names.extend(args.pairs)
    if args.pair_file:
        pair_file = Path(args.pair_file)
        explicit_names.extend(
            line.strip()
            for line in pair_file.read_text().splitlines()
            if line.strip() and not line.strip().startswith("#")
        )
    if explicit_names:
        explicit_pair_dirs = load_pair_dirs_from_args(
            stimulus_root=snippet_dir,
            names=explicit_names,
            use_all=False,
            pair_family=getattr(args, "pair_family", None),
        )
        names.extend(pair_dir_display_name(path, snippet_dir) for path in explicit_pair_dirs)

    deduped = []
    seen = set()
    for name in names:
        if name not in seen:
            deduped.append(name)
            seen.add(name)
    if not deduped:
        raise ValueError("Specify at least one pair with --pairs, --pair-file, or --all.")
    return deduped


def pair_output_dir(output_dir: Path, pair_name: str) -> Path:
    return output_dir / slugify(pair_name)


def load_existing_pair_summary(output_dir: Path, pair_name: str) -> Optional[Dict]:
    summary_path = pair_output_dir(output_dir, pair_name) / "summary.json"
    if not summary_path.exists():
        return None

    try:
        summary = json.loads(summary_path.read_text())
    except json.JSONDecodeError:
        return None

    if summary.get("pair") not in {pair_name, summary.get("pair_id")}:
        return None
    return summary


def load_json_file(path: Path) -> Optional[Dict[str, Any]]:
    if not path.exists():
        return None
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return None
    return payload if isinstance(payload, dict) else None


def attach_unit_test_flow_summary(summary: Dict[str, Any], flow_summary: Dict[str, Any]) -> None:
    summary["unit_test_flow"] = flow_summary
    summary["unit_test_flow_parsed_assertions"] = int(flow_summary.get("parsed_assertions", 0))
    summary["unit_test_flow_analyzed_forward_decisions"] = int(
        flow_summary.get("analyzed_forward_decisions", 0)
    )
    summary["unit_test_flow_analyzed_reverse_decisions"] = int(
        flow_summary.get("analyzed_reverse_decisions", 0)
    )
    summary["unit_test_flow_skipped_decisions"] = int(flow_summary.get("skipped_decisions", 0))
    summary["unit_test_flow_best_prompt_to_output_site"] = json.dumps(
        flow_summary.get("best_prompt_to_output_site", {})
    )
    summary["unit_test_flow_best_response_carry_site"] = json.dumps(
        flow_summary.get("best_response_carry_site", {})
    )
    summary["unit_test_flow_best_combined_context_site"] = json.dumps(
        flow_summary.get("best_combined_context_site", {})
    )
    summary.pop("unit_test_flow_error", None)


def attach_task1_outputs(
    summary: Dict[str, Any],
    pair: Dict[str, Any],
    clean_variant: str,
    corrupted_variant: str,
    generations: AdditionalGenerations,
) -> None:
    prompts = pair.get("prompts") if isinstance(pair.get("prompts"), dict) else {}
    source_prompt = prompts.get(clean_variant)
    destination_prompt = prompts.get(corrupted_variant)
    source_output = str(generations.source_model_output)
    destination_output = str(generations.destination_model_output)
    summary["source_role"] = clean_variant
    summary["destination_role"] = corrupted_variant
    summary["source_output"] = source_output
    summary["destination_output"] = destination_output
    if source_prompt is not None:
        summary["source_execution_output"] = source_prompt.execution_output
        summary["source_cue_output"] = source_prompt.cue_output
        summary["source_output_label"] = classify_behavior(
            source_output,
            source_prompt.execution_output,
            source_prompt.cue_output,
        )
        summary["source_is_accurate"] = bool(source_output == source_prompt.execution_output)
    if destination_prompt is not None:
        summary["destination_execution_output"] = destination_prompt.execution_output
        summary["destination_cue_output"] = destination_prompt.cue_output
        summary["destination_output_label"] = classify_behavior(
            destination_output,
            destination_prompt.execution_output,
            destination_prompt.cue_output,
        )
        summary["destination_is_accurate"] = bool(
            destination_output == destination_prompt.execution_output
        )


def task1_is_complete(summary: Optional[Dict[str, Any]]) -> bool:
    if not isinstance(summary, dict):
        return False
    for side in ("source", "destination"):
        if str(summary.get(f"{side}_output") or "") == "":
            return False
        if str(summary.get(f"{side}_output_label") or "") == "":
            return False
    return True


def unit_test_flow_is_complete(pair_dir: Path) -> bool:
    flow_summary = load_json_file(pair_dir / "unit_test_flow" / "unit_test_flow_summary.json")
    if not isinstance(flow_summary, dict) or flow_summary.get("status") != "completed":
        return False
    decisions_path = pair_dir / "unit_test_flow" / "unit_test_decisions.csv"
    skips_path = pair_dir / "unit_test_flow" / "unit_test_decision_skips.csv"
    if decisions_path.exists() and decisions_path.stat().st_size > 0:
        return True
    return (
        int(flow_summary.get("decision_rows") or 0) == 0
        and int(flow_summary.get("skipped_decisions") or 0) > 0
        and skips_path.exists()
        and skips_path.stat().st_size > 0
    )


def write_dict_rows_csv(path: Path, rows: Sequence[Dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames: List[str] = []
    seen = set()
    for row in rows:
        for key in row.keys():
            key_name = str(key)
            if key_name not in seen:
                fieldnames.append(key_name)
                seen.add(key_name)

    if not fieldnames:
        path.write_text("", encoding="utf-8")
        return

    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows({str(key): value for key, value in row.items()} for row in rows)


DROP_PUBLIC_ALIAS_COLUMNS = {
    "flow_direction",
    "prompt_role",
    "destination_variant",
    "source_variant",
    "unit_id",
    "unit_kind",
    "score",
    "normalized_recovery",
    "distance_based_recovery",
    "source_ref_margin",
    "baseline_margin",
    "clean_positions",
    "corrupted_positions",
    "clean_tokens",
    "corrupted_tokens",
    "clean_token_ids",
    "corrupted_token_ids",
    "clean_text",
    "corrupted_text",
}


def clean_public_dataframe(df):
    if df is None:
        return df
    out = df.copy()
    if "patch_unit_id" not in out.columns and "unit_id" in out.columns:
        out["patch_unit_id"] = out["unit_id"]
    if "patch_unit_type" not in out.columns and "unit_kind" in out.columns:
        out["patch_unit_type"] = out["unit_kind"]
    if "patched_margin" not in out.columns and "score" in out.columns:
        out["patched_margin"] = out["score"]
    if "distance_recovery" not in out.columns and "distance_based_recovery" in out.columns:
        out["distance_recovery"] = out["distance_based_recovery"]
    if "source_margin" not in out.columns and "source_ref_margin" in out.columns:
        out["source_margin"] = out["source_ref_margin"]
    if "destination_margin" not in out.columns and "baseline_margin" in out.columns:
        out["destination_margin"] = out["baseline_margin"]
    if "destination_positions" not in out.columns and "corrupted_positions" in out.columns:
        out["destination_positions"] = out["corrupted_positions"]
    if "source_positions" not in out.columns and "clean_positions" in out.columns:
        out["source_positions"] = out["clean_positions"]
    if "destination_tokens" not in out.columns and "corrupted_tokens" in out.columns:
        out["destination_tokens"] = out["corrupted_tokens"]
    if "source_tokens" not in out.columns and "clean_tokens" in out.columns:
        out["source_tokens"] = out["clean_tokens"]
    if "destination_token_ids" not in out.columns and "corrupted_token_ids" in out.columns:
        out["destination_token_ids"] = out["corrupted_token_ids"]
    if "source_token_ids" not in out.columns and "clean_token_ids" in out.columns:
        out["source_token_ids"] = out["clean_token_ids"]
    if "destination_text" not in out.columns and "corrupted_text" in out.columns:
        out["destination_text"] = out["corrupted_text"]
    if "source_text" not in out.columns and "clean_text" in out.columns:
        out["source_text"] = out["clean_text"]
    return out.drop(columns=[col for col in DROP_PUBLIC_ALIAS_COLUMNS if col in out.columns])


def write_clean_dataframe_csv(df, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    clean_public_dataframe(df).to_csv(path, index=False)


def read_dict_rows_csv(path: Path) -> List[Dict[str, Any]]:
    if not path.exists() or path.stat().st_size == 0:
        return []
    with path.open(newline="", encoding="utf-8") as f:
        return [dict(row) for row in csv.DictReader(f)]


def bool_cell(value: Any) -> Optional[bool]:
    if isinstance(value, bool):
        return value
    if value is None:
        return None
    text = str(value).strip().lower()
    if text in {"true", "1", "yes"}:
        return True
    if text in {"false", "0", "no"}:
        return False
    return None


def generation_to_dict(generation: AdditionalGeneration) -> Dict[str, Any]:
    return {
        "task": generation.task,
        "prompt_role": generation.prompt_role,
        "patch_role": generation.patch_role,
        "max_new_tokens": int(generation.max_new_tokens),
        "prompt": generation.prompt,
        "full_text": generation.full_text,
        "completion": generation.completion,
        "extracted_output": generation.extracted_output,
    }


def additional_generations_to_dict(generations: AdditionalGenerations) -> Dict[str, Any]:
    return {
        "pair_name": generations.pair_name,
        "source_role": generations.source_role,
        "destination_role": generations.destination_role,
        "source_reference_output": generations.source_reference_output,
        "destination_reference_output": generations.destination_reference_output,
        "source_execution_output": generations.source_execution_output,
        "destination_execution_output": generations.destination_execution_output,
        "source_model_output": generations.source_model_output,
        "destination_model_output": generations.destination_model_output,
        "source_answer_generation": generation_to_dict(generations.source_answer_generation),
        "destination_answer_generation": generation_to_dict(generations.destination_answer_generation),
        "unit_tests": {
            str(key): generation_to_dict(generation)
            for key, generation in generations.unit_tests.items()
        },
    }


def generation_from_dict(payload: Dict[str, Any]) -> AdditionalGeneration:
    return AdditionalGeneration(
        task=str(payload.get("task") or ""),
        prompt_role=str(payload.get("prompt_role") or ""),
        patch_role=str(payload.get("patch_role") or ""),
        max_new_tokens=safe_int_cell(payload.get("max_new_tokens"), 0),
        prompt=str(payload.get("prompt") or ""),
        full_text=str(payload.get("full_text") or ""),
        completion=str(payload.get("completion") or ""),
        extracted_output=str(payload.get("extracted_output") or ""),
    )


def cap_unit_test_generation(generation: AdditionalGeneration) -> AdditionalGeneration:
    if generation.task != "unit_test":
        return generation
    generation.completion = cap_unit_test_completion_to_assert_limit(generation.completion)
    extracted = generation.extracted_output or extract_first_fenced_block(generation.completion)
    generation.extracted_output = cap_unit_test_source_to_assert_limit(extracted)
    if generation.prompt and generation.completion:
        generation.full_text = generation.prompt + generation.completion
    return generation


def additional_generations_from_dict(payload: Dict[str, Any]) -> AdditionalGenerations:
    unit_tests_payload = payload.get("unit_tests")
    if not isinstance(unit_tests_payload, dict):
        unit_tests_payload = {}
    return AdditionalGenerations(
        pair_name=str(payload.get("pair_name") or ""),
        source_role=str(payload.get("source_role") or ""),
        destination_role=str(payload.get("destination_role") or ""),
        source_reference_output=str(payload.get("source_reference_output") or ""),
        destination_reference_output=optional_str(payload.get("destination_reference_output")),
        source_execution_output=optional_str(payload.get("source_execution_output")),
        destination_execution_output=optional_str(payload.get("destination_execution_output")),
        source_model_output=str(payload.get("source_model_output") or ""),
        destination_model_output=str(payload.get("destination_model_output") or ""),
        source_answer_generation=generation_from_dict(
            payload.get("source_answer_generation")
            if isinstance(payload.get("source_answer_generation"), dict)
            else {}
        ),
        destination_answer_generation=generation_from_dict(
            payload.get("destination_answer_generation")
            if isinstance(payload.get("destination_answer_generation"), dict)
            else {}
        ),
        unit_tests={
            str(key): cap_unit_test_generation(
                generation_from_dict(value if isinstance(value, dict) else {})
            )
            for key, value in unit_tests_payload.items()
        },
    )


def load_additional_generations_json(pair_dir: Path) -> Optional[AdditionalGenerations]:
    payload = load_generation_json(pair_dir)
    if payload is None:
        return None
    return additional_generations_from_dict(payload)


def write_additional_generations_json(
    output_path: Path,
    generations: AdditionalGenerations,
) -> Path:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(additional_generations_to_dict(generations), indent=2, default=str) + "\n",
        encoding="utf-8",
    )
    return output_path


def tail_text(text: str, max_chars: int = 4000) -> str:
    text = str(text or "")
    return text if len(text) <= max_chars else text[-max_chars:]


def parse_pytest_count(output: str, label: str) -> int:
    matches = re.findall(rf"(\d+)\s+{re.escape(label)}(?:\s|,|\.|$)", output)
    return sum(int(match) for match in matches)


def parse_pytest_counts(stdout: str, stderr: str) -> Dict[str, int]:
    combined = f"{stdout}\n{stderr}"
    return {
        "passed": parse_pytest_count(combined, "passed"),
        "failed": parse_pytest_count(combined, "failed"),
        "skipped": parse_pytest_count(combined, "skipped"),
        "errors": parse_pytest_count(combined, "error") + parse_pytest_count(combined, "errors"),
    }


def execute_generated_pytest_suite(
    *,
    snippet_code: str,
    test_code: str,
    timeout_seconds: int,
) -> Dict[str, Any]:
    timeout_seconds = max(1, int(timeout_seconds))
    if not str(test_code or "").strip():
        return {
            "pytest_exit_code": "",
            "passed": 0,
            "failed": 0,
            "skipped": 0,
            "errors": 0,
            "outcome": "missing_test_code",
            "error": "missing_test_code",
            "stdout_tail": "",
            "stderr_tail": "",
        }

    with tempfile.TemporaryDirectory(prefix="generated_pytest_") as tmp:
        test_path = Path(tmp) / "test_generated.py"
        test_path.write_text(
            str(snippet_code).rstrip() + "\n\n" + str(test_code).rstrip() + "\n",
            encoding="utf-8",
        )
        try:
            completed = subprocess.run(
                [sys.executable, "-m", "pytest", "-q", str(test_path)],
                check=False,
                capture_output=True,
                text=True,
                timeout=timeout_seconds,
            )
        except subprocess.TimeoutExpired as e:
            stdout = e.stdout if isinstance(e.stdout, str) else ""
            stderr = e.stderr if isinstance(e.stderr, str) else ""
            return {
                "pytest_exit_code": "",
                "passed": 0,
                "failed": 0,
                "skipped": 0,
                "errors": 0,
                "outcome": "timeout",
                "error": f"timeout_after_{timeout_seconds}s",
                "stdout_tail": tail_text(stdout),
                "stderr_tail": tail_text(stderr),
            }

    stdout = completed.stdout or ""
    stderr = completed.stderr or ""
    counts = parse_pytest_counts(stdout, stderr)
    if completed.returncode == 5:
        outcome = "no_tests_collected"
    elif completed.returncode == 0:
        outcome = "passed"
    elif counts["failed"] > 0:
        outcome = "failed"
    elif counts["errors"] > 0 or completed.returncode in {2, 3, 4}:
        outcome = "error"
    else:
        outcome = "failed"
    return {
        "pytest_exit_code": int(completed.returncode),
        "passed": int(counts["passed"]),
        "failed": int(counts["failed"]),
        "skipped": int(counts["skipped"]),
        "errors": int(counts["errors"]),
        "outcome": outcome,
        "error": "pytest_error" if completed.returncode != 0 and outcome == "error" else "",
        "stdout_tail": tail_text(stdout),
        "stderr_tail": tail_text(stderr),
    }


def increment_count(container: Dict[str, Any], path: Sequence[str], amount: int = 1) -> None:
    current = container
    for key in path[:-1]:
        current = current.setdefault(str(key), {})
    leaf = str(path[-1])
    current[leaf] = int(current.get(leaf, 0)) + int(amount)


def add_rate_fields(group: Dict[str, Any]) -> None:
    total = int(group.get("total_pairs", 0) or 0)
    if total <= 0:
        return
    rates = group.setdefault("rates", {})
    for label, count in (group.get("status_counts") or {}).items():
        rates[f"status_{label}"] = float(count) / float(total)
    for side, labels in (group.get("answer_label_counts") or {}).items():
        side_total = sum(int(value) for value in labels.values())
        if side_total <= 0:
            continue
        for label, count in labels.items():
            rates[f"{side}_{label}"] = float(count) / float(side_total)


def flatten_stats_rows(stats: Dict[str, Any]) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []

    def visit(prefix: List[str], value: Any) -> None:
        if isinstance(value, dict):
            for key in sorted(value.keys()):
                visit(prefix + [str(key)], value[key])
            return
        rows.append(
            {
                "section": prefix[0] if prefix else "",
                "group": ".".join(prefix[1:-1]),
                "metric": prefix[-1] if prefix else "",
                "value": value,
            }
        )

    visit([], stats)
    return rows


def batch_group_template() -> Dict[str, Any]:
    return {
        "total_pairs": 0,
        "status_counts": {},
        "answer_label_counts": {
            "source": {},
            "destination": {},
            "combined": {},
        },
        "accuracy_counts": {
            "source": {},
            "destination": {},
        },
        "unit_test_flow_status_counts": {},
        "parsed_unit_test_decision_label_counts": {},
    }


def aggregate_summary_row(group: Dict[str, Any], summary: Dict[str, Any]) -> None:
    group["total_pairs"] = int(group.get("total_pairs", 0)) + 1
    status = str(summary.get("status") or "completed")
    increment_count(group, ["status_counts", status])
    for side, key in (("source", "source_output_label"), ("destination", "destination_output_label")):
        label = str(summary.get(key) or "")
        if label:
            increment_count(group, ["answer_label_counts", side, label])
            increment_count(group, ["answer_label_counts", "combined", label])
    for side, key in (("source", "source_is_accurate"), ("destination", "destination_is_accurate")):
        value = bool_cell(summary.get(key))
        if value is not None:
            increment_count(group, ["accuracy_counts", side, "accurate" if value else "inaccurate"])
    unit_status = ""
    unit_summary = summary.get("unit_test_flow")
    if isinstance(unit_summary, dict):
        unit_status = str(unit_summary.get("status") or "")
    unit_status = unit_status or ("error" if summary.get("unit_test_flow_error") else "")
    if unit_status:
        increment_count(group, ["unit_test_flow_status_counts", unit_status])


def collect_parsed_unit_test_decision_stats(pair_dir: Path) -> Dict[str, int]:
    rows = read_dict_rows_csv(pair_dir / "unit_test_flow" / "unit_test_decisions.csv")
    counts: Dict[str, int] = {}
    for row in rows:
        label = str(row.get("response_label") or "")
        if label:
            counts[label] = int(counts.get(label, 0)) + 1
    return counts


def load_generation_json(pair_dir: Path) -> Optional[Dict[str, Any]]:
    path = pair_dir / "unit_test_generations.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return None


def task1_generation_output(generation_payload: Optional[Dict[str, Any]], side: str) -> str:
    if not isinstance(generation_payload, dict):
        return ""
    generation = generation_payload.get(f"{side}_answer_generation")
    if isinstance(generation, dict):
        extracted = str(generation.get("extracted_output") or "")
        if extracted:
            return extracted
    return str(generation_payload.get(f"{side}_model_output") or "")


def summary_lookup_keys(summary: Dict[str, Any]) -> Set[str]:
    keys: Set[str] = set()
    pair_text = str(summary.get("pair") or "")
    pair_id = str(summary.get("pair_id") or "")
    pair_family = str(summary.get("pair_family") or "")
    for value in (pair_text, pair_id):
        if value:
            keys.add(value)
    if pair_family and pair_id:
        keys.add(f"{pair_family}/{pair_id}")
        keys.add(slugify(f"{pair_family}/{pair_id}"))
    return keys


def build_summary_lookup(summaries: Sequence[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
    lookup: Dict[str, Dict[str, Any]] = {}
    for summary in summaries:
        for key in summary_lookup_keys(summary):
            lookup[key] = dict(summary)
    return lookup


def load_pair_summary_from_disk(output_dir: Path, pair_name: str) -> Optional[Dict[str, Any]]:
    summary_path = pair_output_dir(output_dir, pair_name) / "summary.json"
    if not summary_path.exists():
        return None
    try:
        payload = json.loads(summary_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return None
    return payload if isinstance(payload, dict) else None


def overlay_cleaned_task1_generation_labels(
    summary: Dict[str, Any],
    generation_payload: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
    cleaned_summary = dict(summary)
    for side in ("source", "destination"):
        raw_output = task1_generation_output(generation_payload, side)
        if not raw_output:
            continue
        model_output = clean_final_output_text(raw_output)
        execution_output = str(cleaned_summary.get(f"{side}_execution_output") or "")
        cue_output = str(cleaned_summary.get(f"{side}_cue_output") or execution_output)
        if not model_output or not execution_output or not cue_output:
            continue
        label = classify_behavior(model_output, execution_output, cue_output)
        cleaned_summary[f"{side}_model_output_raw"] = raw_output
        cleaned_summary[f"{side}_model_output_clean"] = model_output
        cleaned_summary[f"{side}_output"] = model_output
        cleaned_summary[f"{side}_output_label"] = label
        cleaned_summary[f"{side}_is_accurate"] = label in {
            "both_consistent",
            "execution_consistent",
        }
    return cleaned_summary


def generated_unit_test_entries(generation_payload: Dict[str, Any]) -> List[Dict[str, Any]]:
    unit_tests = generation_payload.get("unit_tests")
    if not isinstance(unit_tests, dict):
        return []
    entries = []
    for key, payload in unit_tests.items():
        if not isinstance(payload, dict):
            continue
        entries.append(
            {
                "generated_test_key": str(key),
                "generated_patch_role": str(payload.get("patch_role") or key),
                "generated_prompt_role": str(payload.get("prompt_role") or ""),
                "test_code": str(payload.get("extracted_output") or ""),
            }
        )
    return entries


def run_batch_descriptive_stats(
    *,
    output_dir: Path,
    snippet_dir: Path,
    pair_names: Sequence[str],
    summaries: Sequence[Dict[str, Any]],
    unit_test_execution_timeout: int,
) -> Dict[str, Any]:
    stats: Dict[str, Any] = {
        "overall": batch_group_template(),
        "by_pair_family": {},
        "unit_test_execution": {
            "outcome_counts": {},
            "by_pair_family": {},
            "by_generated_patch_role": {},
            "by_code_role": {},
            "cross_matrix": {},
        },
    }
    execution_rows: List[Dict[str, Any]] = []
    summary_by_pair = build_summary_lookup(summaries)

    for pair_name in pair_names:
        summary = summary_by_pair.get(str(pair_name)) or summary_by_pair.get(slugify(str(pair_name)))
        if summary is None:
            summary = load_pair_summary_from_disk(output_dir, str(pair_name))
        if summary is None:
            continue
        pair_dir = pair_output_dir(output_dir, str(pair_name))
        generation_payload = load_generation_json(pair_dir)
        summary_for_stats = overlay_cleaned_task1_generation_labels(summary, generation_payload)
        pair_family = str(summary.get("pair_family") or "unknown")
        family_group = stats["by_pair_family"].setdefault(pair_family, batch_group_template())
        aggregate_summary_row(stats["overall"], summary_for_stats)
        aggregate_summary_row(family_group, summary_for_stats)

        for label, count in collect_parsed_unit_test_decision_stats(pair_dir).items():
            increment_count(stats["overall"], ["parsed_unit_test_decision_label_counts", label], count)
            increment_count(family_group, ["parsed_unit_test_decision_label_counts", label], count)

        if generation_payload is None:
            continue
        try:
            pair = load_pair(snippet_dir, str(pair_name))
        except Exception as e:
            execution_rows.append(
                {
                    "pair": pair_name,
                    "pair_family": pair_family,
                    "generated_test_key": "",
                    "generated_patch_role": "",
                    "generated_prompt_role": "",
                    "code_role": "",
                    "pytest_exit_code": "",
                    "passed": 0,
                    "failed": 0,
                    "skipped": 0,
                    "errors": 0,
                    "outcome": "pair_load_error",
                    "error": str(e),
                    "stdout_tail": "",
                    "stderr_tail": "",
                }
            )
            continue
        code_roles = [
            str(summary.get("source_role") or generation_payload.get("source_role") or ""),
            str(summary.get("destination_role") or generation_payload.get("destination_role") or ""),
        ]
        code_roles = [role for role in code_roles if role]
        for test_entry in generated_unit_test_entries(generation_payload):
            for code_role in code_roles:
                try:
                    result = execute_generated_pytest_suite(
                        snippet_code=get_variant_code(pair, code_role),
                        test_code=test_entry["test_code"],
                        timeout_seconds=unit_test_execution_timeout,
                    )
                except Exception as e:
                    result = {
                        "pytest_exit_code": "",
                        "passed": 0,
                        "failed": 0,
                        "skipped": 0,
                        "errors": 0,
                        "outcome": "execution_error",
                        "error": str(e),
                        "stdout_tail": "",
                        "stderr_tail": "",
                    }
                row = {
                    "pair": pair_name,
                    "pair_family": pair_family,
                    "generated_test_key": test_entry["generated_test_key"],
                    "generated_patch_role": test_entry["generated_patch_role"],
                    "generated_prompt_role": test_entry["generated_prompt_role"],
                    "code_role": code_role,
                    **result,
                }
                execution_rows.append(row)
                outcome = str(result.get("outcome") or "unknown")
                generated_patch_role = str(test_entry["generated_patch_role"] or "unknown")
                generated_prompt_role = str(test_entry["generated_prompt_role"] or "unknown")
                increment_count(stats["unit_test_execution"], ["outcome_counts", outcome])
                increment_count(stats["unit_test_execution"], ["by_pair_family", pair_family, outcome])
                increment_count(stats["unit_test_execution"], ["by_generated_patch_role", generated_patch_role, outcome])
                increment_count(stats["unit_test_execution"], ["by_code_role", code_role, outcome])
                matrix_key = f"{generated_patch_role}:{generated_prompt_role}->run_on:{code_role}"
                increment_count(stats["unit_test_execution"], ["cross_matrix", matrix_key, outcome])

    add_rate_fields(stats["overall"])
    for group in stats["by_pair_family"].values():
        add_rate_fields(group)

    write_dict_rows_csv(output_dir / "unit_test_execution_summary.csv", execution_rows)
    write_dict_rows_csv(output_dir / "descriptive_stats.csv", flatten_stats_rows(stats))
    (output_dir / "descriptive_stats.json").write_text(
        json.dumps(stats, indent=2, default=str) + "\n",
        encoding="utf-8",
    )
    return stats


def safe_int_cell(value: Any, default: int = 0) -> int:
    try:
        return int(float(value))
    except (TypeError, ValueError):
        return int(default)




def parse_args() -> argparse.Namespace:
    base_dir = repo_root()
    config_path = base_dir / "config" / "experiment_config.json"
    config = json.loads(config_path.read_text())

    parser = argparse.ArgumentParser(
        description="Run residual patching from the first changed token onward for methodology stimulus pairs."
    )
    parser.add_argument("--pairs", nargs="+", help="Stimulus pair folder names or pair_id values to run.")
    parser.add_argument("--pair-file", help="Text file containing one stimulus pair name per line.")
    parser.add_argument("--all", action="store_true", help="Run all stimulus pair folders.")
    parser.add_argument(
        "--stimulus-root",
        default=str(base_dir / "snippets_working"),
        help="Directory containing cue_varied and implementation_varied folders.",
    )
    parser.add_argument("--pair-family", choices=PAIR_FAMILIES)
    parser.add_argument(
        "--patch-direction",
        help=(
            "Patch direction name. Defaults to the aligned-to-conflicting direction "
            "for each pair family while still computing the reverse direction."
        ),
    )
    parser.add_argument(
        "--output-root",
        default=str(base_dir / "eval" / "model_sweep_clean_schema"),
        help="Clean RQ-stage artifact root.",
    )
    parser.add_argument(
        "--stages",
        default="rq1_1,rq1_2,rq2_1,rq2_2",
        help="Comma-separated RQ stages to run: rq1_1,rq1_2,rq2_1,rq2_2.",
    )
    parser.add_argument(
        "--overwrite-stages",
        default="",
        help="Comma-separated RQ stages to recompute even if their artifacts are complete.",
    )
    parser.add_argument(
        "--summary-label",
        help=(
            "Optional label for split runs. When set, writes batch_summary_<label>.csv "
            "and split-specific error/skip CSVs."
        ),
    )
    parser.add_argument(
        "--model",
        help="Model slug or HuggingFace model name. Slugs are resolved to configured model names.",
    )
    parser.add_argument("--model-name", default=config.get("default_model_name", "gpt2"))
    parser.add_argument("--model-slug", help="Artifact folder slug for the selected model.")
    parser.add_argument("--device", default=config.get("default_device", "cuda"))
    parser.add_argument("--hook-name", default="hook_resid_post")
    parser.add_argument(
        "--patch-batch-size",
        type=int,
        default=int(config.get("patch_batch_size", 8)),
        help="Number of independent patch-unit interventions to evaluate per forward pass.",
    )
    parser.add_argument("--max-positions", type=int, default=None, help="Evenly subsample swept patch units.")
    parser.add_argument("--pos-context", type=int, default=0, help="Include this many positions before first mismatch.")
    parser.add_argument("--no-prepend-bos", action="store_true")
    parser.add_argument("--recovery-metric", choices=["signed", "distance"], default="signed")
    parser.add_argument(
        "--no-additional-task-dump",
        action="store_true",
        help="Skip saving unit-test prompt/output dumps.",
    )
    parser.add_argument(
        "--descriptive-stats",
        dest="descriptive_stats",
        action="store_true",
        default=True,
        help="Write batch-level descriptive_stats.{json,csv} and unit_test_execution_summary.csv.",
    )
    parser.add_argument(
        "--no-descriptive-stats",
        dest="descriptive_stats",
        action="store_false",
        help="Skip batch-level descriptive statistics and generated pytest execution summaries.",
    )
    parser.add_argument(
        "--unit-test-execution-timeout",
        type=int,
        default=30,
        help="Seconds to allow each generated pytest suite/code-under-test execution.",
    )
    parser.add_argument(
        "--answer-max-new-tokens",
        type=int,
        default=32,
        help="Generation length for the answer-only prompt in the additional task dump.",
    )
    parser.add_argument(
        "--unit-test-max-new-tokens",
        type=int,
        default=int(config.get("unit_test_max_new_tokens", 768)),
        help="Fallback generation cap for unit-test generation; compliant fenced outputs stop earlier at closing backticks.",
    )
    parser.add_argument(
        "--token-label-recovery-threshold",
        type=float,
        default=0.3,
        help=(
            "When a heatmap has too many x-axis labels to show all of them, only show "
            "ticks where the selected recovery metric exceeds this value."
        ),
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    base_dir = repo_root()
    snippet_dir = Path(args.stimulus_root).expanduser()
    if not snippet_dir.is_absolute():
        snippet_dir = base_dir / snippet_dir
    snippet_dir = snippet_dir.resolve()
    output_root_arg = args.output_root
    output_root = Path(output_root_arg).expanduser()
    if not output_root.is_absolute():
        output_root = base_dir / output_root
    output_root = output_root.resolve()
    stages = parse_stage_csv(args.stages, default=RQ_STAGES)
    overwrite_stages = parse_stage_csv(args.overwrite_stages)
    model_name, model_slug = resolve_model_name_and_slug(
        model_arg=args.model,
        model_name=args.model_name,
        model_slug=args.model_slug,
    )

    pair_names = load_pair_names(args, snippet_dir)
    errors = []
    manifest_by_pair = {}
    processed_pair_names = []
    skipped_pair_names = []
    pending_pair_names = []
    selected_direction_by_pair: Dict[str, Any] = {}

    for pair_name in pair_names:
        try:
            pair_for_direction = load_pair(snippet_dir, pair_name)
            directions = pair_for_direction["patch_directions"]
            if args.patch_direction:
                selected_direction = next(
                    (
                        direction
                        for direction in directions
                        if direction.name == args.patch_direction
                    ),
                    None,
                )
                if selected_direction is None:
                    raise ValueError(
                        f"Patch direction {args.patch_direction!r} is not valid for "
                        f"{pair_for_direction['pair_id']}"
                    )
            else:
                selected_direction = directions[0]
            selected_direction_by_pair[pair_name] = selected_direction
            pair_id = str(pair_for_direction.get("pair_id") or pair_name)
            pair_family = str(
                pair_for_direction.get("pair_family")
                or pair_for_direction["meta"].get("pair_family", "")
            )
            clean_pair_dir = rq_pair_output_dir(output_root, model_slug, pair_family, pair_id)
            needs_work = any(
                stage in overwrite_stages
                or not stage_is_complete(rq_stage_dir(clean_pair_dir, stage), stage)
                for stage in stages
            )
            if needs_work:
                pending_pair_names.append(pair_name)
            else:
                manifest = load_json_dict(clean_pair_dir / "manifest.json") or {}
                manifest_by_pair[pair_name] = manifest
                skipped_pair_names.append(pair_name)
        except Exception as e:
            errors.append({"pair": pair_name, "error": str(e)})

    if pending_pair_names:
        load_runtime_dependencies()
        print(f"Loading model {model_name!r} on {args.device!r}...")
        HookedTransformer = load_hooked_transformer_class()
        model = HookedTransformer.from_pretrained(model_name, device=args.device)
        model.eval()

        config = json.loads((base_dir / "config" / "experiment_config.json").read_text())
        runner = ResidualPatchRunner(
            model=model,
            question_template=config["question_template"],
            unit_test_prompt_template=config.get("unit_test_prompt_template", DEFAULT_UNIT_TEST_PROMPT_TEMPLATE),
            prepend_bos=not args.no_prepend_bos,
            hook_name=args.hook_name,
            patch_batch_size=args.patch_batch_size,
        )

        for pair_name in tqdm(pending_pair_names, desc="Snippet pairs"):
            try:
                selected_direction = selected_direction_by_pair[pair_name]
                manifest = runner.run_pair_rq_stages(
                    snippet_dir=snippet_dir,
                    output_root=output_root,
                    model_slug=model_slug,
                    model_name=model_name,
                    pair_name=pair_name,
                    clean_variant=selected_direction.source_role,
                    corrupted_variant=selected_direction.destination_role,
                    stages=stages,
                    overwrite_stages=overwrite_stages,
                    max_positions=args.max_positions,
                    pos_context=args.pos_context,
                    answer_max_new_tokens=args.answer_max_new_tokens,
                    unit_test_max_new_tokens=args.unit_test_max_new_tokens,
                    token_label_recovery_threshold=args.token_label_recovery_threshold,
                    recovery_metric=args.recovery_metric,
                )
                manifest_by_pair[pair_name] = manifest
                processed_pair_names.append(pair_name)
            except Exception as e:
                errors.append({"pair": pair_name, "error": str(e)})
            finally:
                cleanup_runtime_memory()
        cleanup_runtime_memory()

    output_root.mkdir(parents=True, exist_ok=True)
    summary_suffix = f"_{args.summary_label}" if args.summary_label else ""
    batch_rows = []
    for pair_name in pair_names:
        manifest = manifest_by_pair.get(pair_name) or {}
        batch_rows.append(
            {
                "pair_name": pair_name,
                "model_slug": model_slug,
                "model_name": model_name,
                "stages_requested": ",".join(sorted(stages)),
                "overwrite_stages": ",".join(sorted(overwrite_stages)),
                "status": "processed" if pair_name in processed_pair_names else "reused",
                "manifest_path": manifest.get("path", ""),
                "stage_statuses": json.dumps(
                    {
                        stage: (manifest.get("stages", {}).get(stage, {}) or {}).get("status", "")
                        for stage in RQ_STAGES
                    },
                    sort_keys=True,
                ),
            }
        )
    write_dict_rows_csv(output_root / f"batch_summary{summary_suffix}.csv", batch_rows)
    if errors:
        write_dict_rows_csv(output_root / f"batch_errors{summary_suffix}.csv", errors)

    print(f"Processed pairs: {len(processed_pair_names)}")
    print(f"Skipped existing pairs: {len(skipped_pair_names)}")
    print(f"Failed pairs: {len(errors)}")
    print(f"Saved outputs to: {output_root}")
    return 1 if errors and not processed_pair_names else 0


if __name__ == "__main__":
    raise SystemExit(main())