frontier_lab.py

"""Forward-momentum probes: subspace residual patch, toy sparse AE, explanation fingerprint attack.

Enable with ``PG_FRONTIER=1``. Mediation narrative is printed from ``main.py`` (always; no extra forwards).

Env knobs: ``PG_SAE_STEPS``, ``PG_SAE_L1``, ``PG_SAE_TOPK``, ``PG_ATTACK_TRIALS``, ``PG_ATTACK_SEED``,
``PG_RANDOM_SUBSPACE_TRIALS`` (random rank-matched orthonormal baselines; default 3),
``PG_SUBSPACE_HOLDOUT_PAIRS`` (exclude last K pairs from Δ-stack for SVD fit),
``PG_FRONTIER_PAIRWISE_COS`` (print focal-layer |cos| between per-surface Δs; default on if ≥2 pairs),
``PG_REROUTE_BLOCK_KO`` (set from main: extra layers excluded from fingerprint sweep).
"""

from __future__ import annotations

import math
import os
import random
import re
import statistics
from typing import List, Optional, Sequence, Set, Tuple

import torch
import torch.nn as nn

from transformer_lens import HookedTransformer


def contrast_basis_rows(
    delta_columns: Sequence[torch.Tensor],
    *,
    rank: int,
    eps: float = 1e-6,
) -> torch.Tensor:
    """Columns = clean−corrupt last-token vectors at one hook. Return Q [rank, d] row-orthonormal."""
    if not delta_columns:
        raise ValueError("delta_columns empty")
    X = torch.stack([v.reshape(-1).float() for v in delta_columns], dim=1)
    if X.shape[1] == 0:
        raise ValueError("need ≥1 delta")
    u, _, _ = torch.linalg.svd(X, full_matrices=False)
    rr = max(1, min(rank, u.shape[0], u.shape[1]))
    rows = u[:, :rr].T.clone()
    rows = rows / torch.clamp(torch.norm(rows, dim=1, keepdim=True), min=eps)
    return rows.detach().cpu()


def random_row_orthonormal_basis(
    *,
    dim: int,
    rank: int,
    generator: torch.Generator,
    eps: float = 1e-6,
) -> torch.Tensor:
    """``rank`` orthonormal rows in R^dim (matched to ``contrast_basis_rows`` geometry)."""
    rr = max(1, min(rank, dim))
    x = torch.randn(rr, dim, generator=generator, dtype=torch.float32)
    q, _ = torch.linalg.qr(x.T, mode="reduced")
    rows = q.T[:rr].clone()
    rows = rows / torch.clamp(torch.norm(rows, dim=1, keepdim=True), min=eps)
    return rows.detach().cpu()


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


def _nt_nm_denoise(
    model: HookedTransformer,
    logits: torch.Tensor,
    *,
    target_token: str,
    distractor_token: str,
    corrupt_tgt: float,
    clean_tgt: float,
    corrupt_margin: float,
    clean_margin: float,
) -> Tuple[float, float]:
    tgt = _token_logit(model, logits, target_token)
    dist = _token_logit(model, logits, distractor_token)
    margin = tgt - dist
    den = clean_tgt - corrupt_tgt
    nt = float("nan") if abs(den) < 1e-9 or den <= 0 else (tgt - corrupt_tgt) / den
    d_m = clean_margin - corrupt_margin
    nm = float("nan") if abs(d_m) < 1e-9 or d_m <= 0 else (margin - corrupt_margin) / d_m
    return nt, nm


def _fwd_full_residual_patch(clean_activation_cache):

    def _fn(act: torch.Tensor, *, hook, ca=clean_activation_cache) -> torch.Tensor:
        donor = ca[hook.name]
        act[:, -1, :] = donor[:, -1, :].to(act.dtype)
        return act

    return _fn


def _make_subspace_residual_hook(clean_activation_cache, q_rows_cpu: torch.Tensor):

    ca = clean_activation_cache


    def _fn(act: torch.Tensor, *, hook) -> torch.Tensor:
        clean = ca[hook.name][:, -1, :].detach()
        rn = act[:, -1, :].detach()
        dlt = (clean.float() - rn.float()).reshape(-1, 1)
        q_basis = q_rows_cpu.to(device=rn.device, dtype=torch.float32)
        coeffs = q_basis @ dlt
        recon = (q_basis.T @ coeffs).reshape_as(rn.float())
        act[:, -1, :] = rn + recon.to(rn.dtype)
        return act

    return _fn


def fingerprint_argmax_competitive_residual(
    model: HookedTransformer,
    corrupt_prompt: str,
    *,
    donor_cache_clean,
    corrupt_tgt: float,
    clean_tgt: float,
    corrupt_margin: float,
    clean_margin: float,
    target_token: str,
    distractor_token: str,
    exclude_layers: Set[int],
) -> int:
    ranked = fingerprint_ranked_competitive_residual(
        model,
        corrupt_prompt,
        donor_cache_clean=donor_cache_clean,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
        target_token=target_token,
        distractor_token=distractor_token,
        exclude_layers=exclude_layers,
    )
    return ranked[0][0] if ranked else -1


def fingerprint_ranked_competitive_residual(
    model: HookedTransformer,
    corrupt_prompt: str,
    *,
    donor_cache_clean,
    corrupt_tgt: float,
    clean_tgt: float,
    corrupt_margin: float,
    clean_margin: float,
    target_token: str,
    distractor_token: str,
    exclude_layers: Set[int],
) -> List[Tuple[int, float, float]]:
    """Return sorted (layer, nt, nm) descending by nt, excluding ``exclude_layers``."""
    rows: List[Tuple[int, float, float]] = []
    for layer in range(model.cfg.n_layers):
        if layer in exclude_layers:
            continue
        hk = f"blocks.{layer}.hook_resid_post"
        logits = model.run_with_hooks(
            corrupt_prompt,
            fwd_hooks=[(hk, _fwd_full_residual_patch(donor_cache_clean))],
        )
        nt, nm = _nt_nm_denoise(
            model,
            logits,
            target_token=target_token,
            distractor_token=distractor_token,
            corrupt_tgt=corrupt_tgt,
            clean_tgt=clean_tgt,
            corrupt_margin=corrupt_margin,
            clean_margin=clean_margin,
        )
        sc = nt if not math.isnan(nt) else float("-inf")
        rows.append((layer, float(sc), float(nm) if not math.isnan(nm) else float("nan")))
    rows.sort(key=lambda t: t[1], reverse=True)
    return rows


class TinySparseAE(nn.Module):
    def __init__(self, dim: int, latent_dim: int) -> None:
        super().__init__()
        self.encoder = nn.Linear(dim, latent_dim, bias=True)
        self.decoder = nn.Linear(latent_dim, dim, bias=True)

    def encode_latent_nonneg(self, x: torch.Tensor) -> torch.Tensor:
        return torch.relu(self.encoder(x))

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        z = self.encode_latent_nonneg(x)
        return self.decoder(z), z


def train_tiny_sae(
    model_ae: TinySparseAE,
    x: torch.Tensor,
    *,
    steps: int,
    lr: float,
    l1: float,
) -> None:
    opt = torch.optim.Adam(model_ae.parameters(), lr=lr)
    for _ in range(steps):
        opt.zero_grad()
        xh, z = model_ae(x.float())
        rec = nn.functional.mse_loss(xh, x.float())
        sparse = z.abs().mean() * l1
        (rec + sparse).backward()
        opt.step()


def _make_latent_transfer_hook(
    clean_ca,
    corr_ca,
    ae: TinySparseAE,
    k_latents: int,
):

    def _fn(act: torch.Tensor, *, hook) -> torch.Tensor:
        corr_ref = corr_ca[hook.name][:, -1, :].float()
        cl = clean_ca[hook.name][:, -1, :].float()
        rn = act[:, -1, :].detach().float()
        z_rn = ae.encode_latent_nonneg(rn)
        z_cl = ae.encode_latent_nonneg(cl)
        z_cr = ae.encode_latent_nonneg(corr_ref)
        focus = torch.abs(z_cl - z_cr)
        kk = max(1, min(k_latents, int(focus.numel())))
        _, ix = torch.topk(focus.reshape(-1), kk)
        z_patch = z_rn.clone()
        flat = z_patch.reshape(-1)
        flat[ix] = z_cl.reshape(-1)[ix]
        z_patch = flat.reshape_as(z_patch)
        recon = ae.decoder(z_patch)
        act[:, -1, :] = recon.to(act.dtype).to(act.device)
        return act

    return _fn


def collect_residual_deltas(
    model: HookedTransformer,
    layers: Sequence[int],
    prompt_pairs: Sequence[Tuple[str, str]],
) -> List[torch.Tensor]:
    out: List[torch.Tensor] = []
    for clp, corp in prompt_pairs:
        _, ca = model.run_with_cache(clp)
        _, cb = model.run_with_cache(corp)
        for layer in layers:
            hk = f"blocks.{layer}.hook_resid_post"
            dv = ca[hk][:, -1, :].detach().float() - cb[hk][:, -1, :].detach().float()
            out.append(dv.reshape(-1).cpu())
    return out


def per_pair_unit_deltas_lastpos(
    model: HookedTransformer,
    layer: int,
    prompt_pairs: Sequence[Tuple[str, str]],
) -> List[torch.Tensor]:
    """Last-position clean−corrupt Δ per surface, L2-normalized (CPU)."""
    hk = f"blocks.{layer}.hook_resid_post"
    out: List[torch.Tensor] = []
    for clp, corp in prompt_pairs:
        _, ca = model.run_with_cache(clp)
        _, cb = model.run_with_cache(corp)
        dv = (ca[hk][:, -1, :] - cb[hk][:, -1, :]).reshape(-1).float().cpu()
        out.append((dv / torch.clamp(torch.norm(dv), min=1e-9)).detach())
    return out


def pairwise_unit_cosine_summary(
    vecs: Sequence[torch.Tensor],
    *,
    max_pairs_print: int = 36,
) -> None:
    """Print upper-triangle |cos| for unit-normalized CPU vectors."""
    n = len(vecs)
    printed = 0
    for i in range(n):
        for j in range(i + 1, n):
            if printed >= max_pairs_print:
                print(f"  … ({max_pairs_print} pair lines shown; tighten PROMPT_VARIANTS slice to see fewer)")
                return
            c = torch.dot(vecs[i], vecs[j]).abs().item()
            print(f"  [{i}] vs [{j}]  |cos Δ|={c:.3f}")
            printed += 1


def _pairwise_cos_enabled() -> bool:
    return (
        os.environ.get("PG_FRONTIER_PAIRWISE_COS", "1").strip().lower()
        not in ("0", "false", "no", "off")
    )


def print_mediation_scaffold(
    *,
    path_joint_nt: float,
    path_joint_nm: float,
    head_layer: int,
    head_idx: int,
    mlp_layer: int,
    resid_denoise_layer: int,
    resid_noise_layer: int,
    ko_layer: int,
    wash_triple_nt: float,
    solo_resid_nt: float,
) -> None:
    print("\n=== MEDIATION SCAFFOLD (Pearl-style narrative from same-run point estimates; not full ID) ===")
    print(
        "Forward flow sketch: attn.z → attn_out → residual_add → MLP_out → residual_add → ⋯ → logits.\n"
        "Hooks used here: attn.hook_z (single head), hook_mlp_out, hook_resid_post.\n"
        "This is **not** a fully identified Pearl DAG from observational data — it summarizes intervention contrasts."
    )
    print(
        f"  • best head_z        L{head_layer}H{head_idx}\n"
        f"  • best MLP donor     L{mlp_layer}\n"
        f"  • competitive resid  denoise≈L{resid_denoise_layer}; noise≈L{resid_noise_layer}\n"
        f"  • joint z+mlp patch  nt={path_joint_nt:.3f} nm={path_joint_nm:.3f}\n"
        f"  • knockout site      L{ko_layer} (then re-sweep under ablation)\n"
        f"  • washout check      stacked z+mlp+resid nt≈{wash_triple_nt:.3f} vs solo resid nt≈{solo_resid_nt:.3f}"
    )
    print(
        "**Qualitative blocked-path next step**: TL-native do-like ablations (freeze attention while routing MLP, "
        "or vice versa) to separate direct vs indirect paths from z → resid → logits."
    )


def run_frontier_suite(
    model: HookedTransformer,
    *,
    clean_prompt: str,
    corrupt_prompt: str,
    target_token: str,
    distractor_token: str,
    clean_act,
    corr_act,
    clean_tgt: float,
    corrupt_tgt: float,
    clean_margin: float,
    corrupt_margin: float,
    exclude_readout_layer: int,
    focal_resid_layer: int,
    prompt_pairs_for_contrast: Sequence[Tuple[str, str]],
    path_joint_nt: float,
    path_joint_nm: float,
    wash_triple_nt: float,
    solo_resid_nt: float,
    reroute_exclude_layer: Optional[int] = None,
) -> None:
    _ = (clean_prompt, path_joint_nt, path_joint_nm, wash_triple_nt, solo_resid_nt)

    hk = f"blocks.{focal_resid_layer}.hook_resid_post"
    exclude_fp: Set[int] = {exclude_readout_layer}

    holdout = max(0, int(os.environ.get("PG_SUBSPACE_HOLDOUT_PAIRS", "0")))
    pairs_all = list(prompt_pairs_for_contrast)
    pairs_fit = pairs_all
    if holdout > 0:
        if len(pairs_all) <= holdout:
            print(
                "\n[PG_SUBSPACE_HOLDOUT_PAIRS] ignored — not enough contrast pairs "
                f"({len(pairs_all)} pairs, holdout {holdout})."
            )
        else:
            pairs_fit = pairs_all[: len(pairs_all) - holdout]

    print("\n=== SUBSPACE RESIDUAL PATCH (Δ-SVD vs full donor vs random rank-matched) ===")
    if len(pairs_fit) < len(pairs_all):
        print(
            f"Δ-basis fit uses {len(pairs_fit)} / {len(pairs_all)} contrast pairs "
            f"(PG_SUBSPACE_HOLDOUT_PAIRS={holdout})."
        )

    delta_vecs = collect_residual_deltas(model, [focal_resid_layer], pairs_fit)
    rank_basis = max(2, min(8, len(delta_vecs)))
    q_rows = contrast_basis_rows(delta_vecs, rank=rank_basis)
    sub_hook = _make_subspace_residual_hook(clean_act, q_rows)

    logits_sub = model.run_with_hooks(corrupt_prompt, fwd_hooks=[(hk, sub_hook)])
    nt_s, nm_s = _nt_nm_denoise(
        model,
        logits_sub,
        target_token=target_token,
        distractor_token=distractor_token,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
    )

    logits_full = model.run_with_hooks(corrupt_prompt, fwd_hooks=[(hk, _fwd_full_residual_patch(clean_act))])
    nt_f, nm_f = _nt_nm_denoise(
        model,
        logits_full,
        target_token=target_token,
        distractor_token=distractor_token,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
    )

    nt_rand_summ = ""
    n_rand_trials = max(0, int(os.environ.get("PG_RANDOM_SUBSPACE_TRIALS", "3")))
    if n_rand_trials > 0:
        r_nt: List[float] = []
        r_nm: List[float] = []
        rseed = int(os.environ.get("PG_RANDOM_SUBSPACE_SEED", "0"))
        for tr in range(n_rand_trials):
            gen = torch.Generator()
            gen.manual_seed(rseed + tr)
            rq = random_row_orthonormal_basis(
                dim=model.cfg.d_model,
                rank=rank_basis,
                generator=gen,
            )
            hook_r = _make_subspace_residual_hook(clean_act, rq)
            logits_r = model.run_with_hooks(corrupt_prompt, fwd_hooks=[(hk, hook_r)])
            nt_ra, nm_ra = _nt_nm_denoise(
                model,
                logits_r,
                target_token=target_token,
                distractor_token=distractor_token,
                corrupt_tgt=corrupt_tgt,
                clean_tgt=clean_tgt,
                corrupt_margin=corrupt_margin,
                clean_margin=clean_margin,
            )
            if not math.isnan(nt_ra):
                r_nt.append(nt_ra)
            if not math.isnan(nm_ra):
                r_nm.append(nm_ra)
        if r_nt:
            mnt = statistics.mean(r_nt)
            s_nt = statistics.stdev(r_nt) if len(r_nt) > 1 else 0.0
            mnm = statistics.mean(r_nm) if r_nm else float("nan")
            snm = statistics.stdev(r_nm) if len(r_nm) > 1 else 0.0
            nt_rand_summ = (
                f"\n       random orthonormal rank≤{rank_basis} "
                f"(PG_RANDOM_SUBSPACE_TRIALS={n_rand_trials}) "
                f"nt≈{mnt:.3f}±{s_nt:.3f} nm≈{mnm:.3f}±{snm:.3f}"
            )
        else:
            nt_rand_summ = "\n       random baseline: nt all NaN (degenerate denominators)"

    print(
        f"L{focal_resid_layer} full donor nt={nt_f:.3f} nm={nm_f:.3f}\n"
        f"       contrast subspace rank≤{rank_basis} Δ-basis graft nt={nt_s:.3f} nm={nm_s:.3f}"
        f"{nt_rand_summ}"
    )

    if _pairwise_cos_enabled() and len(pairs_all) >= 2:
        print(
            "\nPairwise |cos Δ| across contrast surfaces (unit last-pos residual Δ; "
            f"L{focal_resid_layer})."
        )
        unit_vecs = per_pair_unit_deltas_lastpos(model, focal_resid_layer, pairs_all)
        pairwise_unit_cosine_summary(unit_vecs)

    print("\n=== TINY L1 AUTOENCODER + LATENT COORD TRANSFER ===")
    sae_steps = int(os.environ.get("PG_SAE_STEPS", "280"))
    latent_dim = min(512, max(96, model.cfg.d_model))
    samp: List[torch.Tensor] = []
    for clp, corp in prompt_pairs_for_contrast[: min(8, len(prompt_pairs_for_contrast))]:
        _, ca = model.run_with_cache(clp)
        _, cb = model.run_with_cache(corp)
        samp.extend(
            [
                ca[hk][:, -1, :].reshape(-1).detach().cpu().float(),
                cb[hk][:, -1, :].reshape(-1).detach().cpu().float(),
            ]
        )
    train_device = torch.device("cpu")

    dev_model = torch.device(next(model.parameters()).device)
    X = torch.stack(samp, dim=0).float().to(train_device)
    ae = TinySparseAE(model.cfg.d_model, latent_dim).to(train_device)
    train_tiny_sae(ae, X, steps=sae_steps, lr=2e-3, l1=float(os.environ.get("PG_SAE_L1", "5e-3")))
    ae = ae.to(dev_model)

    k_lat = int(os.environ.get("PG_SAE_TOPK", "32"))
    lat_hook = _make_latent_transfer_hook(clean_act, corr_act, ae, k_latents=k_lat)
    logits_lat = model.run_with_hooks(corrupt_prompt, fwd_hooks=[(hk, lat_hook)])
    nt_l, nm_l = _nt_nm_denoise(
        model,
        logits_lat,
        target_token=target_token,
        distractor_token=distractor_token,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
    )
    print(
        f"Toy SAE latent_dim={latent_dim} steps={sae_steps} topkXfer={k_lat} → nt={nt_l:.3f} nm={nm_l:.3f} "
        f"(vs full nt={nt_f:.3f})"
    )

    print("\n=== RESIDUAL-PATCH RANKING + COUNTERFACTUAL ROUTE (default corrupt) ===")
    ranked0 = fingerprint_ranked_competitive_residual(
        model,
        corrupt_prompt,
        donor_cache_clean=clean_act,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
        target_token=target_token,
        distractor_token=distractor_token,
        exclude_layers=exclude_fp,
    )
    top3 = ranked0[: min(3, len(ranked0))]
    t3s = ", ".join(f"L{L} nt={nt:.3f}" for L, nt, _nm in top3)
    print(f"Top competitive resid layers (readout L{exclude_readout_layer} excluded): {t3s}")
    if (
        reroute_exclude_layer is not None
        and reroute_exclude_layer != exclude_readout_layer
        and reroute_exclude_layer >= 0
    ):
        ex2: Set[int] = set(exclude_fp)
        ex2.add(reroute_exclude_layer)
        ranked1 = fingerprint_ranked_competitive_residual(
            model,
            corrupt_prompt,
            donor_cache_clean=clean_act,
            corrupt_tgt=corrupt_tgt,
            clean_tgt=clean_tgt,
            corrupt_margin=corrupt_margin,
            clean_margin=clean_margin,
            target_token=target_token,
            distractor_token=distractor_token,
            exclude_layers=ex2,
        )
        if ranked1 and top3:
            L1, nt1, nm1 = ranked1[0]
            print(
                f"Also excluding mediating resid site L{reroute_exclude_layer}: "
                f"best L{L1} nt={nt1:.3f} nm={nm1:.3f} "
                f"(was L{top3[0][0]} nt={top3[0][1]:.3f})"
            )

    print("\n=== HEURISTIC EXPLANATION ATTACK (layer fingerprint shift) ===")
    trials = int(os.environ.get("PG_ATTACK_TRIALS", "26"))
    seed = int(os.environ.get("PG_ATTACK_SEED", "0"))
    rng = random.Random(seed)
    base_fp = fingerprint_argmax_competitive_residual(
        model,
        corrupt_prompt,
        donor_cache_clean=clean_act,
        corrupt_tgt=corrupt_tgt,
        clean_tgt=clean_tgt,
        corrupt_margin=corrupt_margin,
        clean_margin=clean_margin,
        target_token=target_token,
        distractor_token=distractor_token,
        exclude_layers=exclude_fp,
    )

    mutations = (
        lambda t: t.replace("The capital", "The principal city acting as capital"),
        lambda t: t.replace("France", "the nation France"),
        lambda t: t.replace(". The correct answer is", "! The correct answer is"),
        lambda t: re.sub(r"\s+", "  ", t, count=1),
        lambda t: t.lower(),
        lambda t: t.replace("capital of France", "French capital seat"),
        lambda t: t.replace(
            "is Rome.",
            "is Rome geographically listed incorrectly among wrong choices as.",
        ),
    )

    best_drift = -1
    best_candidate = corrupt_prompt
    best_fplocal = base_fp
    best_mut_rt = corrupt_tgt
    best_mut_mr = corrupt_margin

    for _ in range(trials):
        mut_fn = rng.choice(mutations)
        cand = mut_fn(corrupt_prompt)
        if cand == corrupt_prompt:
            continue
        log_c, _ = model.run_with_cache(cand)
        rt = _token_logit(model, log_c, target_token)
        rd = _token_logit(model, log_c, distractor_token)
        mr = rt - rd
        if rt >= clean_tgt - 1e-9:
            continue

        tgt_gap_abs = abs(clean_tgt - rt)
        if tgt_gap_abs < 0.85:
            continue

        fp_here = fingerprint_argmax_competitive_residual(
            model,
            cand,
            donor_cache_clean=clean_act,
            corrupt_tgt=rt,
            clean_tgt=clean_tgt,
            corrupt_margin=mr,
            clean_margin=clean_margin,
            target_token=target_token,
            distractor_token=distractor_token,
            exclude_layers=exclude_fp,
        )

        drift = abs(fp_here - base_fp)

        if drift > best_drift:
            best_drift = drift
            best_candidate = cand
            best_fplocal = fp_here
            best_mut_rt = rt
            best_mut_mr = mr
            if drift >= max(6, trials // 3):
                break

    drift_show = (
        str(best_drift)
        if best_drift >= 0
        else "0 (no corruption-only fingerprint diverged)"
    )

    nt_shift = float("nan")
    nm_shift = float("nan")
    if best_candidate != corrupt_prompt and best_fplocal >= 0:
        hk_s = f"blocks.{best_fplocal}.hook_resid_post"
        logits_s = model.run_with_hooks(
            best_candidate,
            fwd_hooks=[(hk_s, _fwd_full_residual_patch(clean_act))],
        )
        nt_shift, nm_shift = _nt_nm_denoise(
            model,
            logits_s,
            target_token=target_token,
            distractor_token=distractor_token,
            corrupt_tgt=best_mut_rt,
            clean_tgt=clean_tgt,
            corrupt_margin=best_mut_mr,
            clean_margin=clean_margin,
        )

    clip = best_candidate[:160]
    mut_metrics = ""
    if best_candidate != corrupt_prompt and not math.isnan(nt_shift):
        mut_metrics = (
            f"\nPatch at shifted argmax L{best_fplocal} on best mutant "
            f"(donor=default clean run): nt={nt_shift:.3f} nm={nm_shift:.3f} "
            f"(mutant tgt={best_mut_rt:.3f}, margin={best_mut_mr:.3f})."
        )
    print(
        f"Baseline argmax L{base_fp}; best-shift argmax L{best_fplocal} "
        f"(Δ|layer−baseline|={drift_show}).\n"
        f"Mutated corrupt ({len(best_candidate)} chars): {clip}"
        + ("…" if len(best_candidate) > len(clip) else "")
        + mut_metrics
    )


def frontier_env_enabled() -> bool:
    v = os.environ.get("PG_FRONTIER", "").strip().lower()
    return v in ("1", "true", "yes", "on", "all")