main.py

import os
from pathlib import Path

os.environ.setdefault("TRANSFORMERLENS_ALLOW_MPS", "1")

import math
from collections import Counter
import torch
from dataclasses import dataclass
from itertools import combinations
from typing import Dict, List, Optional, Set, Tuple

from transformer_lens import HookedTransformer

from benchmark_specs import BENCHMARK_SUITE, BenchmarkSpec

import frontier_lab
import milestone_interp

def preferred_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


DEVICE = preferred_device()

# Half precision on accelerator backends; CPU keeps float32 for speed and compatibility.
DTYPE = torch.float16 if DEVICE in ("cuda", "mps") else torch.float32

MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"


def _dump_run_source_snapshots(repo_dir: Path) -> None:
    """Echo key Python modules to stdout before any model work (repro ledger)."""
    files = (
        ("benchmark_specs.py", repo_dir / "benchmark_specs.py"),
        ("main.py", repo_dir / "main.py"),
        ("frontier_lab.py", repo_dir / "frontier_lab.py"),
        ("milestone_interp.py", repo_dir / "milestone_interp.py"),
    )
    for label, path in files:
        bar = "=" * 79
        print(f"\n{bar}\n# BEGIN FILE DUMP: {label} ({path})\n{bar}\n")
        src = path.read_text(encoding="utf-8")
        print(src, end="" if src.endswith("\n") else "\n")
        print(f"{bar}\n# END FILE DUMP: {label}\n{bar}\n")


_REPO_ROOT = Path(__file__).resolve().parent
_skip_snapshot = (
    os.environ.get("PG_SKIP_SOURCE_DUMP", "").strip().lower() in ("1", "true", "yes", "on")
)
if not _skip_snapshot:
    _dump_run_source_snapshots(_REPO_ROOT)

# ---------------------------------------------------------
# Load model
# ---------------------------------------------------------

model = HookedTransformer.from_pretrained_no_processing(
    MODEL_NAME,
    device=DEVICE,
    dtype=DTYPE,
)

# ---------------------------------------------------------
# Experiment config
# ---------------------------------------------------------

# What intervention localizes:
# divergence between *this* clean vs corrupt trajectory at the hook (not "Paris knowledge" carte blanche).

CLEAN_PROMPT = "The capital of France is"
CORRUPT_PROMPT = "The capital of France is Rome. The correct answer is"

TARGET_TOKEN = " Paris"

# Wrong-capital probe for margins (paired with factual target). Tune per falsification variant when spelling changes.
DISTRACTOR_TOKEN = " Rome"

# Surfaces with tiny |clean_tgt−corrupt_tgt| blow up normalized metrics and distort Jaccard — exclude them.
# ~1 logit separates "usable contrast" from borderline denominators where norm is still jittery (~1 gaps).
MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS = 1.5


# Attention head sweep uses single-head replacements — underestimates pairs / small ensembles (known limit).


# Scripted template mutants (same distractor scaffolding) — cheap automated stress beyond hand-picked multilingual rows.
AUTO_STRESS_PAIRS: List[Tuple[str, str, str]] = [
    (
        "stress_double_space",
        "The  capital of France is",
        "The  capital of France is Rome. The correct answer is",
    ),
    (
        "stress_lower",
        "the capital of france is",
        "the capital of france is rome. the correct answer is",
    ),
    (
        "stress_dash_distractor",
        "The capital of France is",
        "The capital of France is Rome — wrong choice. The correct answer is",
    ),
    (
        "stress_synonym_fact",
        "The nation France's administrative capital is",
        "The nation France's administrative capital is Rome. The correct answer is",
    ),
]


def robustness_surface_eligible(clean_tgt: float, corrupt_tgt: float) -> bool:
    """Narrow-band or inverted surfaces pollute normalized / Jaccard summaries."""
    if corrupt_tgt >= clean_tgt - 1e-9:
        return False
    gap = abs(clean_tgt - corrupt_tgt)
    if gap < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS:
        return False
    return True


def nt_nm_sign_conflict(nt: float, nm: float) -> bool:
    """Finite nt and nm with opposite sign: logit-trajectory disagrees with tgt−distractor margin trajectory."""
    if math.isnan(nt) or math.isnan(nm) or math.isinf(nt) or math.isinf(nm):
        return False
    if nt == 0.0 or nm == 0.0:
        return False
    return (nt > 0.0) != (nm > 0.0)


@dataclass
class PatchResult:
    site: str
    layer: int
    target_logit: float
    distractor_logit: float
    logit_margin: float
    normalized_effect: float
    #: Same geometry as normalized_effect but denom = Δ(clean_corr margin); often stabler thin tgt gaps.
    normalized_margin_effect: float
    head: int | None = None


def final_residual_layer_idx() -> int:
    """Last transformer block residual before unembedding (= trivial replay into LM head when patched alone)."""
    return model.cfg.n_layers - 1


def excluding_readout_residual_rows(results: List[PatchResult]) -> List[PatchResult]:
    """Drop hook_resid_post at final layer — replacement feeds donor residual directly into unembedding."""
    hi = final_residual_layer_idx()
    return [r for r in results if not (r.layer == hi and "resid_post" in r.site)]


def residual_readout_bound_row(rows: List[PatchResult]) -> PatchResult | None:
    """The hook_resid_post row at final block (analytical upper norm bound), if present."""
    hi = final_residual_layer_idx()
    for r in rows:
        if r.layer == hi and "resid_post" in r.site:
            return r
    return None


@dataclass(frozen=True)
class FalsificationPrompt:
    clean: str
    corrupt: str
    target_token: str
    distractor_token: str = DISTRACTOR_TOKEN
    label: str = ""


# English surfaces + multilingual parallels. Each row carries distractor spelling aligned with the corrupt wrong city token when it differs from " Rome".
PROMPT_VARIANTS: List[FalsificationPrompt] = [
    FalsificationPrompt(
        CLEAN_PROMPT, CORRUPT_PROMPT, TARGET_TOKEN, DISTRACTOR_TOKEN,
        label="default",
    ),
    FalsificationPrompt(
        "France's capital is",
        "France's capital is Rome. The correct answer is",
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="france_alt",
    ),
    FalsificationPrompt(
        "The city governing France is",
        "The city governing France is Rome. The correct answer is",
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="governing_city",
    ),
    FalsificationPrompt(
        "What is the capital of France?",
        "What is the capital of France? It is Rome. The correct answer is",
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="wh_question",
    ),
    FalsificationPrompt(
        "The French capital is",
        "The French capital is Rome. The correct answer is",
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="french_synonym_np",
    ),
    FalsificationPrompt(
        CLEAN_PROMPT,
        "The capital of France is London. The correct answer is",
        TARGET_TOKEN,
        " London",
        label="distractor_london",
    ),
    FalsificationPrompt(
        "La capitale de la France est",
        "La capitale de la France est Rome. La bonne réponse est",
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="french_sentence",
    ),
    FalsificationPrompt(
        "La capital de Francia es",
        "La capital de Francia es Roma. La respuesta correcta es",
        TARGET_TOKEN,
        " Roma",
        label="spanish_sentence",
    ),
    FalsificationPrompt(
        "Die Hauptstadt von Frankreich ist",
        "Die Hauptstadt von Frankreich ist Rom. Die richtige Antwort ist",
        TARGET_TOKEN,
        " Rom",
        label="german_sentence",
    ),
]


def _token_logit(logits: torch.Tensor, token: str) -> float:
    tid = model.to_single_token(token)
    return logits[0, -1, tid].item()


def normalized_denoise(
    patched_tgt: float, corrupt_tgt: float, clean_tgt: float,
) -> float:
    denom = clean_tgt - corrupt_tgt
    if abs(denom) < 1e-9 or denom < 0:
        # corrupt >= clean — scripted "recovery toward clean" is ill-posed / inverted
        return float("nan")
    return (patched_tgt - corrupt_tgt) / denom


def normalized_noise(
    patched_tgt: float, corrupt_tgt: float, clean_tgt: float,
) -> float:
    denom = clean_tgt - corrupt_tgt
    if abs(denom) < 1e-9 or denom < 0:
        return float("nan")
    return (clean_tgt - patched_tgt) / denom


def normalized_denoise_margin(
    patched_margin: float, corrupt_margin: float, clean_margin: float,
) -> float:
    """Fraction of (clean−corr) **target−distractor margin** closed by patching (contrast-normalized analogue)."""
    denom = clean_margin - corrupt_margin
    if abs(denom) < 1e-9 or denom <= 0:
        return float("nan")
    return (patched_margin - corrupt_margin) / denom


def normalized_noise_margin(
    patched_margin: float, corrupt_margin: float, clean_margin: float,
) -> float:
    denom = clean_margin - corrupt_margin
    if abs(denom) < 1e-9 or denom <= 0:
        return float("nan")
    return (clean_margin - patched_margin) / denom


def patch_metrics(
    patched_logits: torch.Tensor,
    target_token: str,
    distractor_token: str,
    corrupt_tgt_baseline: float,
    clean_tgt_baseline: float,
    corrupt_margin_baseline: float,
    clean_margin_baseline: float,
    *,
    direction: str,
) -> Tuple[float, float, float, float, float]:
    tgt = _token_logit(patched_logits, target_token)
    dist = _token_logit(patched_logits, distractor_token)
    margin = tgt - dist
    if direction == "denoise":
        ne = normalized_denoise(tgt, corrupt_tgt_baseline, clean_tgt_baseline)
        ne_margin = normalized_denoise_margin(
            margin, corrupt_margin_baseline, clean_margin_baseline,
        )
    elif direction == "noise":
        ne = normalized_noise(tgt, corrupt_tgt_baseline, clean_tgt_baseline)
        ne_margin = normalized_noise_margin(
            margin, corrupt_margin_baseline, clean_margin_baseline,
        )
    else:
        raise ValueError(direction)
    return tgt, dist, margin, ne, ne_margin


def patch_last_pos_inject_from_cache(
    run_act: torch.Tensor,
    hook,
    donor_cache,
) -> torch.Tensor:
    """Replace final position activation with donor cache slice (same hook name)."""
    donor = donor_cache[hook.name]
    run_act[:, -1, :] = donor[:, -1, :]
    return run_act


def zero_last_pos_2d(t: torch.Tensor, *, hook) -> torch.Tensor:
    t[:, -1, :] = 0
    return t


def patch_last_pos_attn_z_head_from_cache(
    run_act: torch.Tensor,
    hook,
    donor_cache,
    head: int,
) -> torch.Tensor:
    donor = donor_cache[hook.name]
    run_act[:, -1, head, :] = donor[:, -1, head, :]
    return run_act


def sweep_resid_post_denoise(
    corrupt_prompt: str,
    clean_activations,
    *,
    target_token: str,
    distractor_token: str,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
) -> List[PatchResult]:
    results: List[PatchResult] = []
    for layer in range(model.cfg.n_layers):
        hook_name = f"blocks.{layer}.hook_resid_post"
        patched_logits = model.run_with_hooks(
            corrupt_prompt,
            fwd_hooks=[
                (
                    hook_name,
                    lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(
                        act, hook, ca,
                    ),
                ),
            ],
        )
        tg, dd, marg, norm, norm_m = patch_metrics(
            patched_logits,
            target_token,
            distractor_token,
            corrupt_tgt,
            clean_tgt,
            corrupt_margin,
            clean_margin,
            direction="denoise",
        )
        results.append(
            PatchResult(
                site="hook_resid_post_den",
                layer=layer,
                target_logit=tg,
                distractor_logit=dd,
                logit_margin=marg,
                normalized_effect=norm,
                normalized_margin_effect=norm_m,
            ),
        )
    return results


def sweep_resid_post_noise(
    clean_prompt: str,
    corrupt_activations,
    *,
    target_token: str,
    distractor_token: str,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
) -> List[PatchResult]:
    results: List[PatchResult] = []
    for layer in range(model.cfg.n_layers):
        hook_name = f"blocks.{layer}.hook_resid_post"
        patched_logits = model.run_with_hooks(
            clean_prompt,
            fwd_hooks=[
                (
                    hook_name,
                    lambda act, *, hook, cc=corrupt_activations: patch_last_pos_inject_from_cache(
                        act, hook, cc,
                    ),
                ),
            ],
        )
        tg, dd, marg, norm, norm_m = patch_metrics(
            patched_logits,
            target_token,
            distractor_token,
            corrupt_tgt,
            clean_tgt,
            corrupt_margin,
            clean_margin,
            direction="noise",
        )
        results.append(
            PatchResult(
                site="hook_resid_post_noise",
                layer=layer,
                target_logit=tg,
                distractor_logit=dd,
                logit_margin=marg,
                normalized_effect=norm,
                normalized_margin_effect=norm_m,
            ),
        )
    return results


def sweep_hook_mlp_out_denoise(
    corrupt_prompt: str,
    clean_activations,
    *,
    target_token: str,
    distractor_token: str,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
) -> List[PatchResult]:
    results: List[PatchResult] = []
    for layer in range(model.cfg.n_layers):
        hook_name = f"blocks.{layer}.hook_mlp_out"
        patched_logits = model.run_with_hooks(
            corrupt_prompt,
            fwd_hooks=[
                (
                    hook_name,
                    lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(
                        act, hook, ca,
                    ),
                ),
            ],
        )
        tg, dd, marg, norm, norm_m = patch_metrics(
            patched_logits,
            target_token,
            distractor_token,
            corrupt_tgt,
            clean_tgt,
            corrupt_margin,
            clean_margin,
            direction="denoise",
        )
        results.append(
            PatchResult(
                site="hook_mlp_out_den",
                layer=layer,
                target_logit=tg,
                distractor_logit=dd,
                logit_margin=marg,
                normalized_effect=norm,
                normalized_margin_effect=norm_m,
            ),
        )
    return results


def sweep_attn_hook_z_denoise(
    corrupt_prompt: str,
    clean_activations,
    *,
    target_token: str,
    distractor_token: str,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
) -> List[PatchResult]:
    # Single-head protocol — misses multi-head substrates.
    results: List[PatchResult] = []
    n_heads = model.cfg.n_heads
    for layer in range(model.cfg.n_layers):
        hook_name = f"blocks.{layer}.attn.hook_z"
        for head in range(n_heads):
            patched_logits = model.run_with_hooks(
                corrupt_prompt,
                fwd_hooks=[
                    (
                        hook_name,
                        lambda act, *, hook, ca=clean_activations, hd=head: patch_last_pos_attn_z_head_from_cache(
                            act, hook, ca, hd,
                        ),
                    ),
                ],
            )
            tg, dd, marg, norm, norm_m = patch_metrics(
                patched_logits,
                target_token,
                distractor_token,
                corrupt_tgt,
                clean_tgt,
                corrupt_margin,
                clean_margin,
                direction="denoise",
            )
            results.append(
                PatchResult(
                    site="hook_z_den",
                    layer=layer,
                    head=head,
                    target_logit=tg,
                    distractor_logit=dd,
                    logit_margin=marg,
                    normalized_effect=norm,
                    normalized_margin_effect=norm_m,
                ),
            )
    return results


def top_k_scores(results: List[PatchResult], k: int) -> List[PatchResult]:
    def score(r: PatchResult) -> float:
        if not math.isnan(r.normalized_effect):
            return r.normalized_effect
        return r.logit_margin

    return sorted(results, key=score, reverse=True)[:k]


def degenerate_best(results: List[PatchResult]) -> PatchResult:
    """Prefer normalized_effect when finite; otherwise maximize target logit."""

    def key(r: PatchResult) -> Tuple[int, float, float]:
        fin = int(not math.isnan(r.normalized_effect))
        ne = r.normalized_effect if fin else float("-inf")
        return (fin, ne, r.target_logit)

    return max(results, key=key)


def degenerate_best_residual_competitive(rows: List[PatchResult]) -> PatchResult | None:
    """Best residual-site result excluding trivial final-block readout replay."""
    comp = excluding_readout_residual_rows(rows)
    return degenerate_best(comp) if comp else None


def top_k_layer_ids(
    results: List[PatchResult],
    k: int,
    *,
    exclude_layers: Optional[Set[int]] = None,
) -> List[int]:
    rows = (
        results
        if exclude_layers is None
        else [r for r in results if r.layer not in exclude_layers]
    )
    return [r.layer for r in top_k_scores(rows, k)]


READOUT_UPPER_BOUND_NOTE = (
    "Final block hook_resid_post is an analytical upper bound (donor residual → LM head), "
    "not a competitive causal site vs earlier layers."
)


def jaccard(a: Set[int], b: Set[int]) -> float:
    if not a and not b:
        return 1.0
    return len(a & b) / len(a | b)


def print_explanation_robustness(
    label_to_results: Dict[str, List[PatchResult]],
    k: int,
    surface_targets: Dict[str, Tuple[float, float]],
    min_abs_gap: float,
) -> None:
    """Quantify overlap of top-k causal layers across falsification surfaces."""
    eligible_labels: List[str] = []
    skipped: List[str] = []
    for lab, (c_raw, r_raw) in surface_targets.items():
        if lab not in label_to_results:
            continue
        if robustness_surface_eligible(c_raw, r_raw):
            eligible_labels.append(lab)
        else:
            reason = []
            if r_raw >= c_raw - 1e-9:
                reason.append("corrupt_tgt>=clean_tgt")
            elif abs(c_raw - r_raw) < min_abs_gap:
                reason.append(f"|gap|<{min_abs_gap}")
            skipped.append(f"{lab} ({', '.join(reason)})")

    if not eligible_labels:
        print(
            "\n=== EXPLANATION ROBUSTNESS (skipped — no eligible surfaces after filters) ==="
        )
        print("Skipped:", "; ".join(skipped) if skipped else "(none)")
        return

    filtered = {lab: label_to_results[lab] for lab in eligible_labels}
    labels = list(filtered.keys())
    sets: Dict[str, Set[int]] = {
        lab: set(
            top_k_layer_ids(
                filtered[lab],
                k,
                exclude_layers={final_residual_layer_idx()},
            ),
        )
        for lab in labels
    }
    pairs = list(combinations(labels, 2))
    mean_j = sum(jaccard(sets[a], sets[b]) for a, b in pairs) / max(len(pairs), 1)

    consensus: Set[int] = set.intersection(*sets.values()) if sets else set()
    n = len(labels)
    majority: Set[int] = {
        layer
        for layer in set().union(*sets.values())
        if sum(1 for s in sets.values() if layer in s) > n / 2
    }

    print(
        f"\n=== EXPLANATION ROBUSTNESS (top-{k} layers among blocks 0–{final_residual_layer_idx() - 1} only — "
        f"omit final residual as readout upper bound; eligible surfaces: "
        f"|clean_tgt−corrupt_tgt|≥{min_abs_gap}, corrupt<clean) ==="
    )
    print(f"Eligible (n={n}): {', '.join(labels)}")
    if skipped:
        print(f"Excluded from Jaccard: {', '.join(skipped)}")
    print(f"Mean pairwise Jaccard overlap of top-{k} layer sets: {mean_j:.3f}")
    if consensus:
        print(f"Layers in top-{k} on every eligible surface (consensus): {sorted(consensus)}")
    else:
        print(f"No layer in top-{k} across all {n} eligible surfaces (no strict consensus).")
    if majority:
        print(f"Majority (>50% eligible) top-{k} layers: {sorted(majority)}")


def format_patch_line(r: PatchResult) -> str:
    nt = (
        "nan"
        if math.isnan(r.normalized_effect)
        else f"{r.normalized_effect:.3f}"
    )
    nmar = (
        "nan"
        if math.isnan(r.normalized_margin_effect)
        else f"{r.normalized_margin_effect:.3f}"
    )
    return (
        f"L{r.layer:02d} tgt={r.target_logit:.3f} margin={r.logit_margin:.3f} "
        f"nt={nt} nm={nmar}"
    )


def format_head_line(r: PatchResult) -> str:
    nt = "nan" if math.isnan(r.normalized_effect) else f"{r.normalized_effect:.3f}"
    nmar = "nan" if math.isnan(r.normalized_margin_effect) else f"{r.normalized_margin_effect:.3f}"
    return (
        f"L{r.layer:02d} H{r.head:02d} tgt={r.target_logit:.3f} "
        f"margin={r.logit_margin:.3f} nt={nt} nm={nmar}"
    )


def format_layer_effect_pair(r: PatchResult, decimals: int = 2) -> str:
    """Compact L + nt/nm for tables (decimals=2) or summaries."""
    w = decimals
    nt = "nan" if math.isnan(r.normalized_effect) else f"{r.normalized_effect:.{w}f}"
    nmar = "nan" if math.isnan(r.normalized_margin_effect) else f"{r.normalized_margin_effect:.{w}f}"
    return f"L{r.layer} nt={nt} nm={nmar}"


def print_ranked_residual_mlp(
    title: str,
    rows: List[PatchResult],
    k: int,
    *,
    annotate_readout_bound: bool = True,
) -> None:
    print(f"\n=== {title} ===")
    competitive = excluding_readout_residual_rows(rows) if annotate_readout_bound else rows
    if annotate_readout_bound:
        print(f"(Top-{k} omits L{final_residual_layer_idx():02d} hook_resid_post: {READOUT_UPPER_BOUND_NOTE})")
    for r in top_k_scores(competitive, k):
        print(format_patch_line(r))
    if annotate_readout_bound:
        rb = residual_readout_bound_row(rows)
        if rb is not None:
            nt_rb = (
                "nan"
                if math.isnan(rb.normalized_effect)
                else f"{rb.normalized_effect:.3f}"
            )
            nmt_rb = (
                "nan"
                if math.isnan(rb.normalized_margin_effect)
                else f"{rb.normalized_margin_effect:.3f}"
            )
            print(
                f"L{rb.layer:02d} tgt={rb.target_logit:.3f} margin={rb.logit_margin:.3f} "
                f"nt={nt_rb} nm={nmt_rb}  [upper bound — not ranked vs earlier blocks]"
            )


def print_ranked_heads(title: str, rows: List[PatchResult], k: int) -> None:
    print(f"\n=== {title} (top {k} head sites) ===")
    print("(nt=logit frac; nm=margin frac — same semantics as residual tables.)")
    for r in top_k_scores(rows, k):
        print(format_head_line(r))


@dataclass
class BenchmarkQuickOutcome:
    spec: BenchmarkSpec
    clean_tgt: float
    corrupt_tgt: float
    margin_clean: float
    margin_corr: float
    top_den_layer: int
    top_den_norm: float
    top_den_norm_margin: float
    top_noise_layer: int
    top_noise_norm: float
    top_noise_norm_margin: float
    thin_gap_warn: bool
    inverted_gap: bool
    denoise_nt_nm_conflict: bool
    noise_nt_nm_conflict: bool


def try_benchmark_quick(spec: BenchmarkSpec) -> Optional[BenchmarkQuickOutcome]:
    try:
        model.to_single_token(spec.target_token)
        model.to_single_token(spec.distractor_token)
    except Exception as exc:
        print(f"[SKIP {spec.id}] token check failed ({exc})")
        return None

    cl, ca = model.run_with_cache(spec.clean_prompt)
    crl, cc = model.run_with_cache(spec.corrupt_prompt)
    c_tgt = _token_logit(cl, spec.target_token)
    r_tgt = _token_logit(crl, spec.target_token)
    c_dist = _token_logit(cl, spec.distractor_token)
    r_dist = _token_logit(crl, spec.distractor_token)
    m_c = c_tgt - c_dist
    m_r = r_tgt - r_dist
    inverted = bool(r_tgt >= c_tgt - 1e-9)
    thin = bool(abs(c_tgt - r_tgt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS)

    den = sweep_resid_post_denoise(
        spec.corrupt_prompt,
        ca,
        target_token=spec.target_token,
        distractor_token=spec.distractor_token,
        clean_tgt=c_tgt,
        corrupt_tgt=r_tgt,
        clean_margin=m_c,
        corrupt_margin=m_r,
    )
    nz = sweep_resid_post_noise(
        spec.clean_prompt,
        cc,
        target_token=spec.target_token,
        distractor_token=spec.distractor_token,
        clean_tgt=c_tgt,
        corrupt_tgt=r_tgt,
        clean_margin=m_c,
        corrupt_margin=m_r,
    )
    best_den = degenerate_best_residual_competitive(den)
    best_nz = degenerate_best_residual_competitive(nz)
    if best_den is None:
        best_den = degenerate_best(den)
    if best_nz is None:
        best_nz = degenerate_best(nz)
    dn = best_den.normalized_effect
    nn = best_nz.normalized_effect
    dnm = best_den.normalized_margin_effect
    nnm = best_nz.normalized_margin_effect
    return BenchmarkQuickOutcome(
        spec=spec,
        clean_tgt=c_tgt,
        corrupt_tgt=r_tgt,
        margin_clean=m_c,
        margin_corr=m_r,
        top_den_layer=int(best_den.layer),
        top_den_norm=float(dn if not math.isnan(dn) else float("nan")),
        top_den_norm_margin=float(dnm if not math.isnan(dnm) else float("nan")),
        top_noise_layer=int(best_nz.layer),
        top_noise_norm=float(nn if not math.isnan(nn) else float("nan")),
        top_noise_norm_margin=float(nnm if not math.isnan(nnm) else float("nan")),
        thin_gap_warn=thin,
        inverted_gap=inverted,
        denoise_nt_nm_conflict=nt_nm_sign_conflict(dn, dnm),
        noise_nt_nm_conflict=nt_nm_sign_conflict(nn, nnm),
    )


def run_print_benchmark_suite() -> None:
    print(
        "\n=== MULTI-DOMAIN BENCHMARK HARNESS (residual hooks only — cheap cross-task lens) ==="
    )
    print(
        "Each row scans hook_resid_post denoise/noise like the Primary France block; "
        "MLP/head sweeps stay on the canonical task to limit compute."
    )
    print(
        f"Best-layer column uses competitive blocks only (drops L{final_residual_layer_idx():02d} hook_resid_post "
        f"— readout replay); thin-gap flag uses |clean_tgt−corrupt_tgt|<{MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS}."
    )
    rows_use: List[BenchmarkQuickOutcome] = []
    for spec in BENCHMARK_SUITE:
        o = try_benchmark_quick(spec)
        if o:
            rows_use.append(o)

    hdr = (
        f"{'benchmark_id':<22} {'domain':<18} {'dn_L':>5} {'d_nt':>6} {'d_nm':>6} "
        f"{'nz_L':>5} {'n_nt':>6} {'n_nm':>6} {'gap':>6} {'flg':^9}"
    )
    print(hdr)
    print("-" * len(hdr))
    for o in rows_use:
        dn_s = (
            "  nan "
            if math.isnan(o.top_den_norm)
            else f"{o.top_den_norm:6.3f}"
        )
        dnm_s = (
            "  nan "
            if math.isnan(o.top_den_norm_margin)
            else f"{o.top_den_norm_margin:6.3f}"
        )
        nz_s = (
            "  nan "
            if math.isnan(o.top_noise_norm)
            else f"{o.top_noise_norm:6.3f}"
        )
        nnm_s = (
            "  nan "
            if math.isnan(o.top_noise_norm_margin)
            else f"{o.top_noise_norm_margin:6.3f}"
        )
        gap = o.clean_tgt - o.corrupt_tgt
        flags = ""
        if o.inverted_gap:
            flags += "I"
        if o.thin_gap_warn:
            flags += "T"
        if o.denoise_nt_nm_conflict:
            flags += "d"
        if o.noise_nt_nm_conflict:
            flags += "m"
        flags = flags or "-"
        print(
            f"{o.spec.id:<22} {o.spec.domain:<18} {o.top_den_layer:5d} {dn_s} {dnm_s} "
            f"{o.top_noise_layer:5d} {nz_s} {nnm_s} {gap:+6.2f} {flags:^9}"
        )
    print(
        "(flg: I=inverted tgt gap; T=|clean_tgt−corrupt_tgt| below robust threshold; "
        "d=denoise winner nt/nm opposite sign; m=noise winner nt/nm opposite sign. "
        "d_nt/n_nt=logit frac; d_nm/n_nm=margin frac (tgt−dist).)"
    )

    eligible = [
        r
        for r in rows_use
        if robustness_surface_eligible(r.clean_tgt, r.corrupt_tgt)
    ]
    if eligible:
        c_den = Counter(r.top_den_layer for r in eligible)
        c_nz = Counter(r.top_noise_layer for r in eligible)
        den_mode = sorted(c_den.items(), key=lambda x: (-x[1], x[0]))[:6]
        nz_mode = sorted(c_nz.items(), key=lambda x: (-x[1], x[0]))[:6]
        print(
            "\nAcross eligible benchmarks (invert/thin flagged rows excluded here), "
            "dominant-layer histograms:"
        )
        lr = final_residual_layer_idx()
        print(
            f"(Winner layers exclude readout-bound L{lr:02d}; cross-task concentration in penultimate/noise bests "
            f"often matters more than spread denoise bests — asymmetry ⇒ distinct sufficiency vs necessity geometry.)"
        )
        print(f"  denoise bests: {den_mode}")
        print(f"   noise bests: {nz_mode}")


# ---------------------------------------------------------
# Main experiment (default prompt pair)
# ---------------------------------------------------------

run_print_benchmark_suite()

clean_logits, clean_activations = model.run_with_cache(CLEAN_PROMPT)
corrupt_logits, corrupt_activations = model.run_with_cache(CORRUPT_PROMPT)

ct_clean = _token_logit(clean_logits, TARGET_TOKEN)
cd_clean = _token_logit(clean_logits, DISTRACTOR_TOKEN)
ct_corrupt = _token_logit(corrupt_logits, TARGET_TOKEN)
cd_corrupt = _token_logit(corrupt_logits, DISTRACTOR_TOKEN)

marg_clean = ct_clean - cd_clean
marg_corrupt = ct_corrupt - cd_corrupt
denom_rest = ct_clean - ct_corrupt

print("\n=== BASELINES (default pair) ===")
print(
    f"Clean run: tgt {TARGET_TOKEN!r} logit={ct_clean:.4f}, "
    f"distractor {DISTRACTOR_TOKEN!r}={cd_clean:.4f}, margin={marg_clean:.4f}"
)
print(
    f"Corrupt run: tgt={ct_corrupt:.4f}, distractor={cd_corrupt:.4f}, margin={marg_corrupt:.4f}"
)
print(
    f"Target-gap (clean_tgt - corrupt_tgt) for normalization: {denom_rest:.4f}"
)
print(
    f"Margin-gap (clean_margin - corrupt_margin) for contrast-normalized nm: "
    f"{marg_clean - marg_corrupt:.4f}"
)

results_resid_den = sweep_resid_post_denoise(
    CORRUPT_PROMPT,
    clean_activations,
    target_token=TARGET_TOKEN,
    distractor_token=DISTRACTOR_TOKEN,
    clean_tgt=ct_clean,
    corrupt_tgt=ct_corrupt,
    clean_margin=marg_clean,
    corrupt_margin=marg_corrupt,
)

print_ranked_residual_mlp(
    "DENoise — hook_resid_post (patch corrupt <- clean)",
    results_resid_den,
    10,
)
print(
    "\n(norm>1 is expected when the corrupt frame already pushes \"capital-like\" completions: "
    "clean activations grafted at readout stack with corrupt priming, so full Paris logit overshoot "
    "\"clean-alone\" baseline; treat 1.0 as illustrative, not a tight ceiling. "
    "nt=logit fraction; nm=margin fraction vs clean−corrupt margin.)"
)

results_resid_noise = sweep_resid_post_noise(
    CLEAN_PROMPT,
    corrupt_activations,
    target_token=TARGET_TOKEN,
    distractor_token=DISTRACTOR_TOKEN,
    clean_tgt=ct_clean,
    corrupt_tgt=ct_corrupt,
    clean_margin=marg_clean,
    corrupt_margin=marg_corrupt,
)

print_ranked_residual_mlp(
    "NOISE — hook_resid_post (patch clean <- corrupt; necessity-style)",
    results_resid_noise,
    10,
)

top_resid_den = degenerate_best_residual_competitive(results_resid_den)
top_resid_noise = degenerate_best_residual_competitive(results_resid_noise)
if top_resid_den is None:
    top_resid_den = degenerate_best(results_resid_den)
if top_resid_noise is None:
    top_resid_noise = degenerate_best(results_resid_noise)

print(
    "\n=== INTERPRETATION (defaults: where trajectories diverge) ==="
)
print(
    "Maps localize causal **differences** between scripted clean vs corrupt forwards — "
    'not "where Paris facts live"; wrong-answer / distractor coupling can dominate late sites.'
)
print(
    "Sufficiency (denoise) vs necessity (noise) routinely split — e.g. best injector ≠ bottleneck.\n"
    "Best competitive blocker layers (omit final residual = readout replay bound):\n"
    f"L{top_resid_den.layer} leads denoise (nt={top_resid_den.normalized_effect:.3f}, "
    f"nm={top_resid_den.normalized_margin_effect:.3f}); "
    f"L{top_resid_noise.layer} leads noise toward corrupt (nt={top_resid_noise.normalized_effect:.3f}, "
    f"nm={top_resid_noise.normalized_margin_effect:.3f})."
)
results_mlp = sweep_hook_mlp_out_denoise(
    CORRUPT_PROMPT,
    clean_activations,
    target_token=TARGET_TOKEN,
    distractor_token=DISTRACTOR_TOKEN,
    clean_tgt=ct_clean,
    corrupt_tgt=ct_corrupt,
    clean_margin=marg_clean,
    corrupt_margin=marg_corrupt,
)
print_ranked_residual_mlp(
    "DENoise — hook_mlp_out",
    results_mlp,
    10,
    annotate_readout_bound=False,
)

results_heads = sweep_attn_hook_z_denoise(
    CORRUPT_PROMPT,
    clean_activations,
    target_token=TARGET_TOKEN,
    distractor_token=DISTRACTOR_TOKEN,
    clean_tgt=ct_clean,
    corrupt_tgt=ct_corrupt,
    clean_margin=marg_clean,
    corrupt_margin=marg_corrupt,
)
print_ranked_heads(
    "DENoise — hook_z (single-head sweep limit)",
    results_heads,
    15,
)

top_head = max(
    results_heads,
    key=lambda r: r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf"),
)
top_mlp = max(
    results_mlp,
    key=lambda r: r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf"),
)

print("\n=== DEPTH SPLIT (hypothesis scaffold, not proof) ===")
print(
    f"Early/low-mid peaks in MLP (e.g., L{top_mlp.layer} nt={top_mlp.normalized_effect:.3f}, "
    f"nm={top_mlp.normalized_margin_effect:.3f}) vs "
    f"later residual dominance (denoise L{top_resid_den.layer}) is consonant with "
    "\"retrieve / consolidate\" upstream vs contextual override nearer readout."
)
print(
    f"Largest isolated single head L{top_head.layer} H{top_head.head} "
    f"nt={top_head.normalized_effect:.3f} nm={top_head.normalized_margin_effect:.3f}; "
    "multi-head combos are deliberately not searched here."
)

comp_resid_den = excluding_readout_residual_rows(results_resid_den)
ranked_resid = top_k_scores(comp_resid_den, min(len(comp_resid_den), model.cfg.n_layers))
if len(ranked_resid) >= 2:
    l1, l2 = ranked_resid[0].layer, ranked_resid[1].layer
    duo_logits = model.run_with_hooks(
        CORRUPT_PROMPT,
        fwd_hooks=[
            (
                f"blocks.{l1}.hook_resid_post",
                lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(act, hook, ca),
            ),
            (
                f"blocks.{l2}.hook_resid_post",
                lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(act, hook, ca),
            ),
        ],
    )
    dt, dd_, _, duo_norm, duo_nm = patch_metrics(
        duo_logits,
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        ct_corrupt,
        ct_clean,
        marg_corrupt,
        marg_clean,
        direction="denoise",
    )

    print("\n=== MULTI-SITE RESidual (defaults, joint top two competitive denoise layers) ===")
    print(
        f"Layers L{l1}, L{l2} jointly tgt={dt:.4f} margin={dt-dd_:.4f} "
        f"nt={duo_norm:.3f} nm={duo_nm:.3f}"
    )
    print("(Compare to singles above.)")

phl = int(top_head.layer)
phh = int(top_head.head)
pml = int(top_mlp.layer)
pr_late = int(top_resid_den.layer)

path_joint_logits = model.run_with_hooks(
    CORRUPT_PROMPT,
    fwd_hooks=[
        (
            f"blocks.{phl}.attn.hook_z",
            lambda act, *, hook, ca=clean_activations, h=phh: patch_last_pos_attn_z_head_from_cache(
                act, hook, ca, h,
            ),
        ),
        (
            f"blocks.{pml}.hook_mlp_out",
            lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(
                act, hook, ca,
            ),
        ),
    ],
)
pj_tgt, _pj_dd, pj_mr, pn_joint, pn_joint_m = patch_metrics(
    path_joint_logits,
    TARGET_TOKEN,
    DISTRACTOR_TOKEN,
    ct_corrupt,
    ct_clean,
    marg_corrupt,
    marg_clean,
    direction="denoise",
)

path_washout_logits = model.run_with_hooks(
    CORRUPT_PROMPT,
    fwd_hooks=[
        (
            f"blocks.{phl}.attn.hook_z",
            lambda act, *, hook, ca=clean_activations, h=phh: patch_last_pos_attn_z_head_from_cache(
                act, hook, ca, h,
            ),
        ),
        (
            f"blocks.{pml}.hook_mlp_out",
            lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(
                act, hook, ca,
            ),
        ),
        (
            f"blocks.{pr_late}.hook_resid_post",
            lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(
                act, hook, ca,
            ),
        ),
    ],
)
_, _, _, pn_wash, pn_wash_m = patch_metrics(
    path_washout_logits,
    TARGET_TOKEN,
    DISTRACTOR_TOKEN,
    ct_corrupt,
    ct_clean,
    marg_corrupt,
    marg_clean,
    direction="denoise",
)

single_late_den = next(r for r in results_resid_den if r.layer == pr_late)
print("\n=== PATH PATCH (donor head_z @ last pos + donor mlp_out — late residual NOT stacked) ===")
print(
    f"L{phl}H{phh} z + L{pml} mlp_out jointly -> tgt={pj_tgt:.4f} margin={pj_mr:.4f} "
    f"nt={pn_joint:.3f} nm={pn_joint_m:.3f}. "
    "Later blocks run on this perturbed stream (not Wang frozen-attn path patching)."
)
print(
    "\nContrast — stacking hook_resid_post at "
    f"L{pr_late} after z/mlp: full residual overwrite **erases** upstream grafts toward the donor run; "
    f"triple nt≈solo resid "
    f"(triple nt={pn_wash:.3f} nm={pn_wash_m:.3f} vs solo L{pr_late} "
    f"nt={single_late_den.normalized_effect:.3f} nm={single_late_den.normalized_margin_effect:.3f})."
)

frontier_lab.print_mediation_scaffold(
    path_joint_nt=pn_joint,
    path_joint_nm=pn_joint_m,
    head_layer=phl,
    head_idx=phh,
    mlp_layer=pml,
    resid_denoise_layer=int(top_resid_den.layer),
    resid_noise_layer=int(top_resid_noise.layer),
    ko_layer=int(top_resid_den.layer),
    wash_triple_nt=float(pn_wash),
    solo_resid_nt=float(single_late_den.normalized_effect),
)

KO_LAYER = top_resid_den.layer
ablate_hook_name = f"blocks.{KO_LAYER}.hook_resid_post"

ko_only_logits = model.run_with_hooks(
    CORRUPT_PROMPT,
    fwd_hooks=[(ablate_hook_name, zero_last_pos_2d)],
)
ko_tgt = _token_logit(ko_only_logits, TARGET_TOKEN)
ko_dist = _token_logit(ko_only_logits, DISTRACTOR_TOKEN)
print("\n=== LOCALIZED RESIDUAL KNOCKOUT (default corrupt prompt) ===")
print(f"Zero final-pos residual after block {KO_LAYER}.")
print(
    f"Corrupt tgt {ct_corrupt:.4f} -> {ko_tgt:.4f}; margins {marg_corrupt:.4f} -> {(ko_tgt-ko_dist):.4f}"
)
print(
    "If Paris logits fall but Rome falls faster here, you've removed a corridor that was loudly "
    "coupled to *wrong-city* continuation under this corrupt scaffold — downstream layers can "
    "still salvage Paris (seen in post-knockout sweeps), so treat as redundancy, not a single bottleneck."
)


def resid_falsification_sweep(
    v: FalsificationPrompt,
    *,
    ablate_hook_resid_layer: int | None = None,
) -> Tuple[float, float, float, float, List[PatchResult]]:
    """Returns clean_tgt, corrupt_tgt, denom, corrupt_margin, results."""
    cl, ca = model.run_with_cache(v.clean)
    crl, _cc = model.run_with_cache(v.corrupt)
    c_tgt = _token_logit(cl, v.target_token)
    r_tgt = _token_logit(crl, v.target_token)
    c_dst = _token_logit(cl, v.distractor_token)
    r_dst = _token_logit(crl, v.distractor_token)
    marg_c_clean = c_tgt - c_dst
    marg_corrupt_fwd = r_tgt - r_dst
    denom = c_tgt - r_tgt

    outs: List[PatchResult] = []
    base_hooks = []
    if ablate_hook_resid_layer is not None:
        ablate_name = f"blocks.{ablate_hook_resid_layer}.hook_resid_post"
        base_hooks.append((ablate_name, zero_last_pos_2d))

    for layer in range(model.cfg.n_layers):
        if ablate_hook_resid_layer is not None and layer == ablate_hook_resid_layer:
            continue
        hook_name = f"blocks.{layer}.hook_resid_post"
        fwd_hooks = list(base_hooks) + [
            (
                hook_name,
                lambda tens, *, hook, act_cache=ca: patch_last_pos_inject_from_cache(
                    tens, hook, act_cache,
                ),
            ),
        ]
        patched = model.run_with_hooks(v.corrupt, fwd_hooks=fwd_hooks)
        tg, dd, marg, norm, norm_m = patch_metrics(
            patched,
            v.target_token,
            v.distractor_token,
            r_tgt,
            c_tgt,
            marg_corrupt_fwd,
            marg_c_clean,
            direction="denoise",
        )
        outs.append(
            PatchResult(
                site="hook_resid_post_den_falsify",
                layer=layer,
                target_logit=tg,
                distractor_logit=dd,
                logit_margin=marg,
                normalized_effect=norm,
                normalized_margin_effect=norm_m,
            ),
        )

    corrupt_margin = marg_corrupt_fwd
    return c_tgt, r_tgt, denom, corrupt_margin, outs


results_by_label: Dict[str, List[PatchResult]] = {}
surface_targets: Dict[str, Tuple[float, float]] = {}

print("\n=== FALSIFICATION — denoise residuals across surfaces ===")
print(
    "Per-row norms get noisy when |clean_tgt−corrupt_tgt| is tiny; pooled Jaccard excludes those surfaces. "
    "Per-surface bests and top‑3 omit final hook_resid_post (readout replay bound). "
    "nt=logit frac vs tgt gap; nm=margin frac vs (clean−corr) tgt−dist margin."
)

for v in PROMPT_VARIANTS:
    c_tgt, r_tgt, denom, corr_marg, pr = resid_falsification_sweep(v)
    results_by_label[v.label] = pr
    surface_targets[v.label] = (c_tgt, r_tgt)
    best = degenerate_best_residual_competitive(pr) or degenerate_best(pr)
    pr_rank = excluding_readout_residual_rows(pr)
    pool = pr_rank if pr_rank else pr
    top3 = top_k_scores(pool, min(3, len(pool)))
    layers = ", ".join(format_layer_effect_pair(r, decimals=2) for r in top3)
    denoise_norm_degenerate = r_tgt >= c_tgt - 1e-9
    tiny_gap = abs(c_tgt - r_tgt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS
    print(
        f"\n[{v.label}] distractor={v.distractor_token!r}  "
        f"clean_tgt={c_tgt:.3f} corrupt_tgt={r_tgt:.3f} gap={denom:.3f} corrupt_margin={corr_marg:.3f}"
        + (
            " [corrupt tgt >= clean — denoise norm disabled]"
            if denoise_norm_degenerate
            else (" [|gap| small — noisy norm]" if tiny_gap else "")
        )
        + "\n"
        f"  best L{best.layer} nt={best.normalized_effect:.3f} nm={best.normalized_margin_effect:.3f} "
        f"tgt={best.target_logit:.3f} margin={best.logit_margin:.3f}  top3: {layers}"
    )

print_explanation_robustness(
    results_by_label,
    k=5,
    surface_targets=surface_targets,
    min_abs_gap=MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS,
)

_, _, _, _, pr_after_ko = resid_falsification_sweep(
    FalsificationPrompt(
        CLEAN_PROMPT,
        CORRUPT_PROMPT,
        TARGET_TOKEN,
        DISTRACTOR_TOKEN,
        label="defaults",
    ),
    ablate_hook_resid_layer=KO_LAYER,
)
best_after = degenerate_best_residual_competitive(pr_after_ko) or degenerate_best(pr_after_ko)

rank_pool = excluding_readout_residual_rows(pr_after_ko) or list(pr_after_ko)
sorted_after = sorted(
    rank_pool,
    key=lambda r: (
        r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf")
    ),
    reverse=True,
)


def _floored(row: PatchResult, floor_tgt: float) -> bool:
    return row.target_logit > floor_tgt + 0.05


above_floor_simple = [r for r in sorted_after if _floored(r, ko_tgt)]

print("\n=== KNOCKOUT + RE-SWEEP (defaults; knockout still applied) ===")
print(
    f"Knockout L{KO_LAYER}; best survivor L{best_after.layer} "
    f"nt={best_after.normalized_effect:.3f} nm={best_after.normalized_margin_effect:.3f} "
    f"tgt={best_after.target_logit:.3f}"
)
print(
    "Above knockout tgt floor (~+0.05): "
    + ", ".join(
        f"L{r.layer}(t={r.target_logit:.2f})" for r in above_floor_simple[:8]
    )
)
print("(Downstream patches reinject signal after earlier hook zeroing.)")


print("\n=== AUTOMATED TEMPLATE STRESS (mutants vs default scaffolding) ===")
for stress_lab, cp_stress, corp_stress in AUTO_STRESS_PAIRS:
    v_stress = FalsificationPrompt(cp_stress, corp_stress, TARGET_TOKEN, DISTRACTOR_TOKEN, label=stress_lab)
    s_ct, s_rt, _, _s_cm, spr = resid_falsification_sweep(v_stress)
    sb = degenerate_best_residual_competitive(spr) or degenerate_best(spr)
    tag = ""
    if s_rt >= s_ct - 1e-9:
        tag = " [inverted tgt gap]"
    elif abs(s_ct - s_rt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS:
        tag = " [|gap|<robust threshold]"
    print(
        f"[{stress_lab}] gap={s_ct - s_rt:.3f} best L{sb.layer} "
        f"nt={sb.normalized_effect:.3f} nm={sb.normalized_margin_effect:.3f} "
        f"tgt={sb.target_logit:.3f} margin={sb.logit_margin:.3f}{tag}"
    )


def print_method_note() -> None:
    print("\n=== WHAT QUESTION THIS STACK ANSWERS ===")
    print(
        "Operational question: causal **circuit discovery anchored to one scripted distracting surface**, "
        "then **stress-tested across paraphrases + light template mutation** — task-local mediation, "
        "not an invariant Paris module."
    )
    print(
        "Benchmark harness: declarative prompts live in benchmark_specs.py; extend BENCHMARK_SUITE for "
        "new domains while keeping residual sweeps bounded—switch on full MLP/head ladders per task if compute allows."
    )
    print(
        "Frontier kit (PG_FRONTIER=1): subspace Δ-patch + random orthonormal baseline, holdout pairs, "
        "pairwise |cos Δ| across surfaces, toy SAE latent transfer, fingerprint attack w/ patch metrics, "
        "resid ranking + reroute when blocking denoise site; tune PG_RANDOM_SUBSPACE_TRIALS, "
        "PG_SUBSPACE_HOLDOUT_PAIRS, PG_FRONTIER_PAIRWISE_COS, PG_SAE_STEPS / PG_ATTACK_TRIALS."
    )
    print(
        "Milestone (PG_MILESTONE=1 or PG_FROZEN_ATTN=1 / PG_SAELENS=1): frozen clean→corrupt attn.hook_pattern "
        "graft; SAELens top-k xfer at pretrained hook (PG_SAE_RELEASE+PG_SAE_ID or PG_SAE_DISK)."
    )


print_method_note()

_SCRAF = os.environ.get("PG_FRONTIER", "").strip().lower()
if _SCRAF in ("1", "true", "yes", "on", "all"):
    _pairs = [(v.clean, v.corrupt) for v in PROMPT_VARIANTS[:8]]
    frontier_lab.run_frontier_suite(
        model,
        clean_prompt=CLEAN_PROMPT,
        corrupt_prompt=CORRUPT_PROMPT,
        target_token=TARGET_TOKEN,
        distractor_token=DISTRACTOR_TOKEN,
        clean_act=clean_activations,
        corr_act=corrupt_activations,
        clean_tgt=ct_clean,
        corrupt_tgt=ct_corrupt,
        clean_margin=marg_clean,
        corrupt_margin=marg_corrupt,
        exclude_readout_layer=final_residual_layer_idx(),
        focal_resid_layer=int(top_resid_den.layer),
        prompt_pairs_for_contrast=_pairs,
        path_joint_nt=float(pn_joint),
        path_joint_nm=float(pn_joint_m),
        wash_triple_nt=float(pn_wash),
        solo_resid_nt=float(single_late_den.normalized_effect),
        reroute_exclude_layer=int(top_resid_den.layer),
    )

if milestone_interp.milestone_any_env():
    milestone_interp.run_milestone_bundle(
        model,
        corrupt_prompt=CORRUPT_PROMPT,
        target_token=TARGET_TOKEN,
        distractor_token=DISTRACTOR_TOKEN,
        clean_act=clean_activations,
        corr_act=corrupt_activations,
        clean_tgt=ct_clean,
        corrupt_tgt=ct_corrupt,
        clean_margin=marg_clean,
        corrupt_margin=marg_corrupt,
        default_attn_layer_for_frozen_pattern=int(phl),
    )