mech_workbench/cli.py

"""Unified command-line interface for the Mechanistic Interpretability Workbench."""

from __future__ import annotations

import argparse
import csv
import dataclasses
import hashlib
import json
import math
import random
import sys
import zipfile
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from mech_workbench.model_utils import is_finite as _finite, safe_mean as _mean

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

VALID_PROBES = ("residual", "logit_lens", "gradient", "attention")
DEFAULT_ROBUSTNESS_PROBES = ("logit_lens", "gradient", "attention")
def _parse_probe_names(raw: str | None) -> list[str] | None:
    if not raw:
        return None
    names = [part.strip() for part in raw.split(",") if part.strip()]
    unknown = [name for name in names if name not in VALID_PROBES]
    if unknown:
        raise ValueError(f"Unknown probe(s): {unknown}. Valid: {list(VALID_PROBES)}")
    return names

def _cmd_run(args: argparse.Namespace) -> int:
    """Load the model once, run selected probes, and write numeric artifacts."""
    from mech_workbench.config import config_to_dict, load_config

    try:
        probe_names = _parse_probe_names(args.probe)
    except ValueError as exc:
        print(f"[error] {exc}")
        return 1

    config = load_config(args.config)
    if getattr(args, "limit", None):
        config = dataclasses.replace(config, prompt_pairs=config.prompt_pairs[: args.limit])

    if args.dry_run:
        _print_dry_run(config, probe_names)
        return 0

    from mech_workbench.logger import ResultLogger
    from mech_workbench.model_loader import load_model_and_tokenizer
    from mech_workbench.probes import run_all_probes

    selected = probe_names or list(DEFAULT_ROBUSTNESS_PROBES)
    logger = ResultLogger(args.results_root)
    logger.write_config(config_to_dict(config))

    print("")
    print("=" * 60)
    print(f"  mech run - {config.model_name}")
    print(f"  Probes  : {selected}")
    print(f"  Pairs   : {len(config.prompt_pairs)}")
    print("=" * 60)
    print("")

    print("[1/4] Loading model...")
    model, tokenizer = load_model_and_tokenizer(
        config.model_name, device=config.device, dtype=config.dtype
    )

    print("[2/4] Running probes...")
    try:
        results = run_all_probes(model, tokenizer, config, logger, selected)
    except Exception:
        logger.finalize()
        raise

    print("[3/4] Writing probe artifacts...")
    raw_dir = logger.run_dir / "raw"
    for probe_name, result in results.items():
        path = raw_dir / f"{probe_name}_results.json"
        path.write_text(
            json.dumps(result, indent=2, default=str) + "\n", encoding="utf-8"
        )
        print(f"  wrote raw/{path.name}")

    _write_probe_csvs(logger.run_dir, results)

    print("[4/4] Writing summary, manifest, integrity...")
    summary = _write_probe_run_artifacts(logger.run_dir, results, config)

    _write_run_card(logger.run_dir, config, ["discovery"])

    from mech_workbench.integrity import write_integrity, write_package_manifest

    write_integrity(logger.run_dir)
    write_package_manifest(logger.run_dir)

    archive_path = None
    if args.archive_root:
        archive_path = _create_run_zip(logger.run_dir, Path(args.archive_root).resolve())

    print("")
    print(f"Results directory: {logger.run_dir}")
    print(f"Run status: {summary['verdict']['status']} - {summary['verdict']['claim']}")
    if archive_path:
        print(f"Artifact zip: {archive_path}")
    print("")
    return 0

def _cmd_ablate(args: argparse.Namespace) -> int:
    """Run targeted attention-head ablations."""
    from mech_workbench.config import config_to_dict, load_config
    from mech_workbench.probes.ablation import (
        HeadAblationProbe,
        parse_ablation_modes,
        parse_head_targets,
    )

    try:
        targets = parse_head_targets(args.target)
        if args.also:
            seen = set(targets)
            for target in parse_head_targets(args.also):
                if target not in seen:
                    targets.append(target)
                    seen.add(target)
        modes = parse_ablation_modes(args.mode)
    except ValueError as exc:
        print(f"[error] {exc}")
        return 1

    config = load_config(args.config)
    if getattr(args, "limit", None):
        config = dataclasses.replace(config, prompt_pairs=config.prompt_pairs[: args.limit])
    combined = getattr(args, "combined", False)

    if args.dry_run:
        _print_ablation_dry_run(config, targets, modes)
        return 0

    from mech_workbench.integrity import write_integrity, write_package_manifest
    from mech_workbench.logger import ResultLogger
    from mech_workbench.model_loader import load_model_and_tokenizer

    run_name = f"mech_ablation_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    logger = ResultLogger(args.results_root, run_name=run_name)
    logger.write_config(config_to_dict(config))

    mode_label = "combined" if combined else "individual"
    print("")
    print("=" * 60)
    print(f"  mech ablate - {config.model_name}")
    print(f"  Targets : {[target.label for target in targets]}")
    print(f"  Modes   : {modes}  [{mode_label}]")
    print(f"  Pairs   : {len(config.prompt_pairs)}")
    print("=" * 60)
    print("")

    print("[1/4] Loading model...")
    model, tokenizer = load_model_and_tokenizer(
        config.model_name, device=config.device, dtype=config.dtype
    )

    print("[2/4] Running head ablations...")
    probe = HeadAblationProbe(targets=targets, modes=modes, combined=combined)
    try:
        result = probe.run(model, tokenizer, config, logger)
    except Exception:
        logger.finalize()
        raise

    print("[3/4] Writing ablation artifacts...")
    raw_dir = logger.run_dir / "raw"
    result_path = raw_dir / "ablation_results.json"
    result_path.write_text(
        json.dumps(result, indent=2, default=str) + "\n", encoding="utf-8"
    )
    print("  wrote raw/ablation_results.json")

    rows = _ablation_rows(result)
    if rows:
        _write_csv(logger.run_dir / "data" / "ablation_results.csv", rows)
        print("  wrote data/ablation_results.csv")

    print("[4/4] Writing summary, manifest, integrity...")
    summary = _write_ablation_run_artifacts(logger.run_dir, result, config)
    _write_run_card(logger.run_dir, config, ["ablation"])
    write_integrity(logger.run_dir)
    write_package_manifest(logger.run_dir)

    archive_path = None
    if args.archive_root:
        archive_path = _create_run_zip(logger.run_dir, Path(args.archive_root).resolve())

    print("")
    print(f"Results directory: {logger.run_dir}")
    print(f"Run status: {summary['verdict']['status']} - {summary['verdict']['claim']}")
    if archive_path:
        print(f"Artifact zip: {archive_path}")
    print("")
    return 0

def _cmd_source(args: argparse.Namespace) -> int:
    """Run source-layer probing for a target attention head."""
    from mech_workbench.config import config_to_dict, load_config
    from mech_workbench.probes.ablation import parse_head_target
    from mech_workbench.probes.source import (
        SourceProbe,
        parse_source_layers,
        parse_source_position,
    )

    try:
        target = parse_head_target(args.target)
        source_layers = parse_source_layers(args.source_layers)
        source_position = parse_source_position(args.source_position)
    except ValueError as exc:
        print(f"[error] {exc}")
        return 1

    config = load_config(args.config)
    if getattr(args, "limit", None):
        config = dataclasses.replace(config, prompt_pairs=config.prompt_pairs[: args.limit])

    if args.dry_run:
        _print_source_dry_run(config, target, source_layers, source_position)
        return 0

    from mech_workbench.integrity import write_integrity, write_package_manifest
    from mech_workbench.logger import ResultLogger
    from mech_workbench.model_loader import load_model_and_tokenizer

    run_name = f"mech_source_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    logger = ResultLogger(args.results_root, run_name=run_name)
    logger.write_config(config_to_dict(config))

    print("")
    print("=" * 60)
    print(f"  mech source - {config.model_name}")
    print(f"  Target  : {target.label}")
    print(f"  Layers  : {source_layers}")
    print(f"  Position: {source_position}")
    print(f"  Pairs   : {len(config.prompt_pairs)}")
    print("=" * 60)
    print("")

    print("[1/4] Loading model...")
    model, tokenizer = load_model_and_tokenizer(
        config.model_name, device=config.device, dtype=config.dtype
    )

    print("[2/4] Running source probing...")
    probe = SourceProbe(
        target=target,
        source_layers=source_layers,
        source_position=source_position,
    )
    try:
        result = probe.run(model, tokenizer, config, logger)
    except Exception:
        logger.finalize()
        raise

    print("[3/4] Writing source artifacts...")
    raw_dir = logger.run_dir / "raw"
    result_path = raw_dir / "source_results.json"
    result_path.write_text(
        json.dumps(result, indent=2, default=str) + "\n", encoding="utf-8"
    )
    print("  wrote raw/source_results.json")

    rows = _source_rows(result)
    if rows:
        _write_csv(logger.run_dir / "data" / "source_results.csv", rows)
        print("  wrote data/source_results.csv")

    print("[4/4] Writing summary, manifest, integrity...")
    summary = _write_source_run_artifacts(logger.run_dir, result, config)
    _write_run_card(logger.run_dir, config, ["source"])
    write_integrity(logger.run_dir)
    write_package_manifest(logger.run_dir)

    archive_path = None
    if args.archive_root:
        archive_path = _create_run_zip(logger.run_dir, Path(args.archive_root).resolve())

    print("")
    print(f"Results directory: {logger.run_dir}")
    print(f"Run status: {summary['verdict']['status']} - {summary['verdict']['claim']}")
    if archive_path:
        print(f"Artifact zip: {archive_path}")
    print("")
    return 0

def _print_ablation_dry_run(config: Any, targets: list[Any], modes: list[str]) -> None:
    layer_end = config.layer_end if config.layer_end is not None else "(all)"
    print("[dry-run] Ablation config validated:")
    print(f"  Model        : {config.model_name}")
    print(f"  Prompt pairs : {len(config.prompt_pairs)}")
    print(f"  Layer range  : {config.layer_start} - {layer_end}")
    print(f"  Patch pos    : {config.patch_position}")
    print(f"  Targets      : {[target.label for target in targets]}")
    print(f"  Modes        : {modes}")
    for pair in config.prompt_pairs:
        headline = "headline" if getattr(pair, "headline", True) else "specificity"
        print(
            f"  Pair '{pair.id}' [{pair.family}/{pair.variant}/{headline}]: "
            f"'{pair.correct_token}' vs '{pair.incorrect_token}'"
        )

def _print_source_dry_run(
    config: Any,
    target: Any,
    source_layers: list[int],
    source_position: str,
) -> None:
    print("[dry-run] Source config validated:")
    print(f"  Model        : {config.model_name}")
    print(f"  Prompt pairs : {len(config.prompt_pairs)}")
    print(f"  Patch pos    : {config.patch_position}")
    print(f"  Target       : {target.label}")
    print(f"  Source layers: {source_layers}")
    print(f"  Source pos   : {source_position}")
    for pair in config.prompt_pairs:
        headline = "headline" if getattr(pair, "headline", True) else "diagnostic"
        print(
            f"  Pair '{pair.id}' [{pair.family}/{pair.variant}/{headline}]: "
            f"'{pair.correct_token}' vs '{pair.incorrect_token}'"
        )

def _print_dry_run(config: Any, probe_names: list[str] | None) -> None:
    layer_end = config.layer_end if config.layer_end is not None else "(all)"
    probes = probe_names or list(DEFAULT_ROBUSTNESS_PROBES)
    print("[dry-run] Config validated:")
    print(f"  Model        : {config.model_name}")
    print(f"  Prompt pairs : {len(config.prompt_pairs)}")
    print(f"  Layer range  : {config.layer_start} - {layer_end}")
    print(f"  Patch pos    : {config.patch_position}")
    print(f"  Probes       : {probes}")
    print(f"  Random heads : {getattr(config, 'random_baseline_samples', 32)} samples/pair")
    for pair in config.prompt_pairs:
        headline = "headline" if getattr(pair, "headline", True) else "diagnostic"
        print(
            f"  Pair '{pair.id}' [{pair.family}/{pair.variant}/{headline}]: "
            f"'{pair.correct_token}' vs '{pair.incorrect_token}'"
        )

def _write_probe_csvs(run_dir: Path, results: dict[str, Any]) -> None:
    """Write per-probe CSV files to the data/ subdirectory."""
    data_dir = run_dir / "data"

    attn = results.get("attention", {})
    head_rows: list[dict[str, Any]] = []
    for pr in attn.get("pair_results", []):
        head_rows.extend(pr.get("head_rows", []))
    if head_rows:
        _write_csv(data_dir / "head_recovery.csv", head_rows)
        print("  wrote data/head_recovery.csv")

    lens = results.get("logit_lens", {})
    lens_rows: list[dict[str, Any]] = []
    for pr in lens.get("pair_results", []):
        for row in pr.get("rows", []):
            lens_rows.append({"pair_id": pr["pair_id"], **row})
    if lens_rows:
        _write_csv(data_dir / "logit_lens.csv", lens_rows)
        print("  wrote data/logit_lens.csv")

    grad = results.get("gradient", {})
    grad_rows: list[dict[str, Any]] = []
    for pr in grad.get("pair_results", []):
        for row in pr.get("layer_scores", []):
            grad_rows.append({"pair_id": pr["pair_id"], **row})
    if grad_rows:
        _write_csv(data_dir / "gradient_attribution.csv", grad_rows)
        print("  wrote data/gradient_attribution.csv")

def _ablation_rows(result: dict[str, Any]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for pair in result.get("pair_results", []):
        rows.extend(pair.get("head_results", []))
    return rows

def _source_rows(result: dict[str, Any]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for pair in result.get("pair_results", []):
        rows.extend(pair.get("source_rows", []))
    return rows

def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    fieldnames = list(rows[0].keys())
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

def _write_probe_run_artifacts(run_dir: Path, results: dict[str, Any], config: Any) -> dict[str, Any]:
    """Write discovery summary and log artifacts."""
    summary = _summarize_probe_results(run_dir, results, config)
    summaries_dir = run_dir / "summaries"
    (summaries_dir / "run_summary.json").write_text(
        json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8"
    )
    (run_dir / "logs" / "run_log.txt").write_text(_render_run_log(summary), encoding="utf-8")
    _write_candidate_heads_csv(summaries_dir, summary)
    return summary

def _write_candidate_heads_csv(summaries_dir: Path, summary: dict[str, Any]) -> None:
    attn = summary.get("attention", {})
    top = (
        attn.get("top_headline_control_adjusted_heads")
        or attn.get("top_mean_margin_heads")
        or attn.get("top_heads_by_mean_margin")
        or []
    )
    if top:
        rows = []
        for rank, row in enumerate(top, start=1):
            out = dict(row)
            out.setdefault("rank", rank)
            out.setdefault("head_id", f"L{out.get('layer')}H{out.get('head')}")
            rows.append(out)
        _write_csv(summaries_dir / "candidate_heads.csv", rows)
        print("  wrote summaries/candidate_heads.csv")

def _write_ablation_run_artifacts(
    run_dir: Path,
    result: dict[str, Any],
    config: Any,
) -> dict[str, Any]:
    summary = _summarize_ablation_result(run_dir, result, config)
    summaries_dir = run_dir / "summaries"
    (summaries_dir / "run_summary.json").write_text(
        json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8"
    )
    (run_dir / "logs" / "run_log.txt").write_text(_render_ablation_log(summary), encoding="utf-8")
    _write_ablation_summary_csv(summaries_dir, summary)
    return summary

def _write_ablation_summary_csv(summaries_dir: Path, summary: dict[str, Any]) -> None:
    rows = summary.get("aggregate", {}).get("by_target_mode_scope", [])
    if rows:
        _write_csv(summaries_dir / "ablation_summary.csv", rows)
        print("  wrote summaries/ablation_summary.csv")

def _write_source_run_artifacts(
    run_dir: Path,
    result: dict[str, Any],
    config: Any,
) -> dict[str, Any]:
    summary = _summarize_source_result(run_dir, result, config)
    summaries_dir = run_dir / "summaries"
    (summaries_dir / "run_summary.json").write_text(
        json.dumps(summary, indent=2, default=str) + "\n", encoding="utf-8"
    )
    (run_dir / "logs" / "run_log.txt").write_text(_render_source_log(summary), encoding="utf-8")
    _write_source_summary_csv(summaries_dir, summary)
    return summary

def _write_source_summary_csv(summaries_dir: Path, summary: dict[str, Any]) -> None:
    rows = summary.get("pair_summaries", [])
    if rows:
        _write_csv(summaries_dir / "source_summary.csv", rows)
        print("  wrote summaries/source_summary.csv")

def _summarize_probe_results(run_dir: Path, results: dict[str, Any], config: Any) -> dict[str, Any]:
    pair_meta = _pair_metadata(config)
    random_samples = int(getattr(config, "random_baseline_samples", 32) or 32)
    seed = int(getattr(config, "seed", 1234) or 1234)

    logit_lens = _summarize_logit_lens(results.get("logit_lens", {}), pair_meta)
    attention = _summarize_attention(
        results.get("attention", {}),
        pair_meta,
        seed=seed,
        random_baseline_samples=random_samples,
    )
    gradient = _summarize_gradient(results.get("gradient", {}), pair_meta)

    warnings: list[str] = []
    if logit_lens["control_bias_warnings"]:
        warnings.append(
            f"{len(logit_lens['control_bias_warnings'])} control prompt(s) have positive final logit bias."
        )
    weak_margin = [
        row["pair_id"]
        for row in attention["pair_summaries"]
        if row.get("best_margin_patch_recovery") is not None
        and row["best_margin_patch_recovery"] < 0
    ]
    if weak_margin:
        warnings.append(
            "Best-margin head has negative patch recovery for: " + ", ".join(weak_margin)
        )
    if gradient["nan_scores"] > 0:
        warnings.append(f"{gradient['nan_scores']} gradient score(s) are NaN.")
    unstable = [
        row["canonical_id"]
        for row in attention.get("paraphrase_stability", [])
        if row.get("num_variants", 0) > 1 and not row.get("stable_best_margin_heads")
    ]
    if unstable:
        warnings.append(
            "No repeated best-margin head across paraphrases for: " + ", ".join(unstable)
        )

    status = "pass_with_warnings" if warnings else "pass"
    if not results:
        status = "no_result"
        claim = "No probe result files were produced."
    elif not any(results.get(name) for name in DEFAULT_ROBUSTNESS_PROBES):
        claim = "Residual probe run completed."
    else:
        claim = (
            "Robustness run completed; interpret heads using control-adjusted "
            "recovery, paraphrase stability, family grouping, and random-head baselines."
        )

    return {
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "run_dir": str(run_dir),
        "model_name": getattr(config, "model_name", None),
        "prompt_pairs": len(getattr(config, "prompt_pairs", [])),
        "probes": sorted(results.keys()),
        "pair_metadata": list(pair_meta.values()),
        "random_baseline_samples": random_samples,
        "verdict": {
            "status": status,
            "claim": claim,
            "warnings": warnings,
        },
        "logit_lens": logit_lens,
        "attention": attention,
        "gradient": gradient,
    }

def _summarize_ablation_result(
    run_dir: Path,
    result: dict[str, Any],
    config: Any,
) -> dict[str, Any]:
    rows = _ablation_rows(result)
    aggregate = result.get("aggregate", {})
    targets = [row.get("label") for row in result.get("targets", [])]
    modes = result.get("modes", [])
    primary_target = _primary_ablation_target(aggregate, targets)

    warnings: list[str] = []
    primary_mean = _find_ablation_aggregate(
        aggregate,
        target_head=primary_target,
        mode="mean",
        headline=True,
    )
    primary_control = _find_specificity(
        aggregate,
        target_head=primary_target,
        mode="mean",
    )

    if primary_mean:
        mean_drop = primary_mean.get("mean_logit_diff_drop")
        rel_drop = primary_mean.get("mean_relative_drop_fraction")
        if not _finite(mean_drop) or float(mean_drop) <= 1.5:
            warnings.append(
                f"{primary_target} mean ablation headline drop did not clear 1.5 logits."
            )
        if not _finite(rel_drop) or float(rel_drop) <= 0.25:
            warnings.append(
                f"{primary_target} mean ablation headline relative drop did not clear 25%."
            )
    else:
        warnings.append(f"No headline mean-ablation aggregate found for {primary_target}.")

    if primary_control:
        ratio = primary_control.get("factual_to_control_drop_ratio")
        if _is_nan(ratio) or (_finite(ratio) and float(ratio) <= 2.0):
            warnings.append(
                f"{primary_target} factual/control specificity ratio did not clear 2x."
            )

    status = "pass_with_warnings" if warnings else "pass"
    claim = (
        "Targeted head ablation completed; interpret necessity using "
        "absolute drop, gap-normalized drop, and factual-vs-control specificity."
    )

    return {
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "run_dir": str(run_dir),
        "model_name": getattr(config, "model_name", None),
        "prompt_pairs": len(getattr(config, "prompt_pairs", [])),
        "probe": "ablation",
        "targets": targets,
        "modes": modes,
        "rows": len(rows),
        "verdict": {
            "status": status,
            "claim": claim,
            "warnings": warnings,
        },
        "aggregate": aggregate,
        "pair_summaries": _summarize_ablation_pairs(result),
    }

def _primary_ablation_target(
    aggregate: dict[str, Any],
    targets: list[Any],
) -> str | None:
    mean_headline_labels = {
        str(row.get("target_head"))
        for row in aggregate.get("by_target_mode_scope", [])
        if row.get("ablation_mode") == "mean" and bool(row.get("headline")) is True
    }
    target_labels = [str(target) for target in targets if target is not None]
    combined_label = "+".join(target_labels)
    if combined_label and combined_label in mean_headline_labels:
        return combined_label
    if target_labels and target_labels[0] in mean_headline_labels:
        return target_labels[0]
    if mean_headline_labels:
        return sorted(mean_headline_labels)[0]
    return target_labels[0] if target_labels else None

def _find_ablation_aggregate(
    aggregate: dict[str, Any],
    *,
    target_head: str | None,
    mode: str,
    headline: bool,
) -> dict[str, Any] | None:
    for row in aggregate.get("by_target_mode_scope", []):
        if (
            row.get("target_head") == target_head
            and row.get("ablation_mode") == mode
            and bool(row.get("headline")) == headline
        ):
            return row
    return None

def _find_specificity(
    aggregate: dict[str, Any],
    *,
    target_head: str | None,
    mode: str,
) -> dict[str, Any] | None:
    for row in aggregate.get("specificity", []):
        if row.get("target_head") == target_head and row.get("ablation_mode") == mode:
            return row
    return None

def _summarize_ablation_pairs(result: dict[str, Any]) -> list[dict[str, Any]]:
    summaries: list[dict[str, Any]] = []
    for pair in result.get("pair_results", []):
        head_rows = pair.get("head_results", [])
        best_drop = _best_ablation_row(head_rows, "logit_diff_drop")
        summaries.append(
            {
                "pair_id": pair.get("pair_id"),
                "family": pair.get("family"),
                "canonical_id": pair.get("canonical_id"),
                "variant": pair.get("variant"),
                "headline": pair.get("headline"),
                "clean_logit_diff": pair.get("clean_logit_diff"),
                "corrupted_logit_diff": pair.get("corrupted_logit_diff"),
                "control_logit_diff": pair.get("control_logit_diff"),
                "largest_drop": {
                    "target_head": best_drop.get("target_head"),
                    "ablation_mode": best_drop.get("ablation_mode"),
                    "logit_diff_drop": _float_or_none(best_drop.get("logit_diff_drop")),
                    "relative_drop_fraction": _float_or_none(
                        best_drop.get("relative_drop_fraction")
                    ),
                    "gap_normalized_drop": _float_or_none(
                        best_drop.get("gap_normalized_drop")
                    ),
                } if best_drop else {},
            }
        )
    return summaries

def _summarize_source_result(
    run_dir: Path,
    result: dict[str, Any],
    config: Any,
) -> dict[str, Any]:
    rows = _source_rows(result)
    aggregate = result.get("aggregate", {})
    source_position = str(result.get("source_position", "unknown"))
    best_layer_input = (
        aggregate.get("best_token_layer_summaries", [])
        if source_position == "sweep"
        else aggregate.get("source_layer_summaries", [])
    )
    best_layer = _best_source_layer(best_layer_input)

    warnings: list[str] = []
    if not best_layer:
        warnings.append("No finite source-layer recovery values were produced.")
    else:
        best_layer_score = _source_layer_target_score(best_layer)
        if not _finite(best_layer_score) or float(best_layer_score) <= 0.2:
            if source_position == "sweep":
                warnings.append("Best-token source-layer mean target recovery did not clear 0.2.")
            else:
                warnings.append("Best source-layer mean target recovery did not clear 0.2.")
    if source_position == "all":
        warnings.append(
            "All-position source patch is a ceiling sanity check, not source localization."
        )
        saturated = [
            row for row in aggregate.get("source_layer_summaries", [])
            if _finite(row.get("mean_target_recovery_fraction"))
            and float(row["mean_target_recovery_fraction"]) >= 0.99
            and _finite(row.get("mean_logit_recovery_fraction"))
            and float(row["mean_logit_recovery_fraction"]) >= 0.99
        ]
        if len(saturated) == len(aggregate.get("source_layer_summaries", [])) and saturated:
            warnings.append(
                "All candidate source layers saturated near 1.0; rerun with source_position=sweep."
            )

    status = "pass_with_warnings" if warnings else "pass"
    claim = (
        "Source probing completed; interpret source layers by target-head recovery "
        "and logit recovery after residual-stream patching."
    )
    target = result.get("target", {}).get("label")

    return {
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "run_dir": str(run_dir),
        "model_name": getattr(config, "model_name", None),
        "prompt_pairs": len(getattr(config, "prompt_pairs", [])),
        "probe": "source",
        "target": target,
        "source_layers": result.get("source_layers", []),
        "source_position": source_position,
        "rows": len(rows),
        "verdict": {
            "status": status,
            "claim": claim,
            "warnings": warnings,
        },
        "aggregate": aggregate,
        "best_source_layer": best_layer,
        "pair_summaries": _summarize_source_pairs(result),
    }

def _best_source_layer(source_summaries: list[dict[str, Any]]) -> dict[str, Any]:
    finite = [row for row in source_summaries if _finite(_source_layer_target_score(row))]
    if not finite:
        return {}
    return max(finite, key=lambda row: float(_source_layer_target_score(row)))

def _source_layer_target_score(row: dict[str, Any]) -> Any:
    if "mean_best_target_recovery_fraction" in row:
        return row.get("mean_best_target_recovery_fraction")
    return row.get("mean_target_recovery_fraction")

def _source_layer_logit_score(row: dict[str, Any]) -> Any:
    if "mean_best_logit_recovery_fraction" in row:
        return row.get("mean_best_logit_recovery_fraction")
    return row.get("mean_logit_recovery_fraction")

def _source_layer_shift_score(row: dict[str, Any]) -> Any:
    if "mean_best_target_output_shift" in row:
        return row.get("mean_best_target_output_shift")
    return row.get("mean_target_output_shift")

def _summarize_source_pairs(result: dict[str, Any]) -> list[dict[str, Any]]:
    summaries: list[dict[str, Any]] = []
    for pair in result.get("pair_results", []):
        best = pair.get("best_source_layer") or {}
        summaries.append(
            {
                "pair_id": pair.get("pair_id"),
                "family": pair.get("family"),
                "canonical_id": pair.get("canonical_id"),
                "variant": pair.get("variant"),
                "headline": pair.get("headline"),
                "clean_logit_diff": pair.get("clean_logit_diff"),
                "corrupted_logit_diff": pair.get("corrupted_logit_diff"),
                "best_source_layer": best.get("source_layer"),
                "source_position": best.get("source_position") or pair.get("source_position"),
                "source_position_index": best.get("source_position_index"),
                "source_positions_patched": best.get("source_positions_patched"),
                "clean_source_token": best.get("clean_source_token"),
                "corrupted_source_token": best.get("corrupted_source_token"),
                "target_recovery_fraction": _float_or_none(
                    best.get("target_recovery_fraction")
                ),
                "target_output_shift": _float_or_none(best.get("target_output_shift")),
                "logit_recovery_fraction": _float_or_none(
                    best.get("logit_recovery_fraction")
                ),
            }
        )
    return summaries

def _best_ablation_row(rows: list[dict[str, Any]], key: str) -> dict[str, Any]:
    finite = [row for row in rows if _finite(row.get(key))]
    if not finite:
        return {}
    return max(finite, key=lambda row: float(row[key]))

def _pair_metadata(config: Any) -> dict[str, dict[str, Any]]:
    metadata: dict[str, dict[str, Any]] = {}
    for pair in getattr(config, "prompt_pairs", []) or []:
        pair_id = getattr(pair, "id", None)
        if not pair_id:
            continue
        metadata[pair_id] = {
            "pair_id": pair_id,
            "family": getattr(pair, "family", "general"),
            "canonical_id": getattr(pair, "canonical_id", None) or pair_id,
            "variant": getattr(pair, "variant", "base"),
            "headline": bool(getattr(pair, "headline", True)),
        }
    return metadata

def _meta_for_pair(pair_id: Any, pair_meta: dict[str, dict[str, Any]]) -> dict[str, Any]:
    pid = str(pair_id)
    return pair_meta.get(
        pid,
        {
            "pair_id": pid,
            "family": "general",
            "canonical_id": pid,
            "variant": "base",
            "headline": True,
        },
    )

def _summarize_logit_lens(
    data: dict[str, Any],
    pair_meta: dict[str, dict[str, Any]],
    control_threshold: float = 0.5,
) -> dict[str, Any]:
    pair_summaries = []
    warnings = []
    token_bias_diagnostics = []
    for pr in data.get("pair_results", []):
        meta = _meta_for_pair(pr.get("pair_id"), pair_meta)
        rows_by_label: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for row in pr.get("rows", []):
            rows_by_label[str(row.get("prompt_label"))].append(row)

        clean = _last_logit_diff(rows_by_label.get("clean", []))
        corrupted = _last_logit_diff(rows_by_label.get("corrupted", []))
        control = pr.get("final_control_logit_diff")
        if control is None:
            control = _last_logit_diff(rows_by_label.get("control", []))

        summary = {
            "pair_id": pr.get("pair_id"),
            "family": meta["family"],
            "canonical_id": meta["canonical_id"],
            "variant": meta["variant"],
            "headline": meta["headline"],
            "first_positive_layer": pr.get("first_positive_layer"),
            "final_clean_logit_diff": clean,
            "final_corrupted_logit_diff": corrupted,
            "final_control_logit_diff": control,
        }
        pair_summaries.append(summary)

        severity = _control_bias_severity(control, control_threshold)
        token_bias_diagnostics.append(
            {
                "pair_id": pr.get("pair_id"),
                "family": meta["family"],
                "canonical_id": meta["canonical_id"],
                "variant": meta["variant"],
                "headline": meta["headline"],
                "final_control_logit_diff": _float_or_none(control),
                "severity": severity,
            }
        )

        if _finite(control) and float(control) > control_threshold:
            warnings.append(
                {
                    "pair_id": pr.get("pair_id"),
                    "final_control_logit_diff": float(control),
                    "threshold": control_threshold,
                }
            )

    return {
        "num_pairs": len(pair_summaries),
        "control_bias_threshold": control_threshold,
        "control_bias_warnings": warnings,
        "token_bias_diagnostics": token_bias_diagnostics,
        "pair_summaries": pair_summaries,
    }

def _control_bias_severity(value: Any, threshold: float) -> str:
    if not _finite(value):
        return "unknown"
    numeric = float(value)
    if numeric > threshold:
        return "biased_toward_correct"
    if numeric > 0:
        return "mild_correct_prior"
    return "neutral_or_opposite"

def _last_logit_diff(rows: list[dict[str, Any]]) -> float | None:
    if not rows:
        return None
    rows = sorted(rows, key=lambda row: int(row.get("layer", 0)))
    value = rows[-1].get("logit_diff")
    return float(value) if _finite(value) else None

def _summarize_attention(
    data: dict[str, Any],
    pair_meta: dict[str, dict[str, Any]],
    *,
    seed: int,
    random_baseline_samples: int,
) -> dict[str, Any]:
    pair_summaries = []
    all_rows: list[dict[str, Any]] = []
    margin_counter: Counter[str] = Counter()
    patch_counter: Counter[str] = Counter()
    headline_margin_counter: Counter[str] = Counter()
    headline_patch_counter: Counter[str] = Counter()

    for pr in data.get("pair_results", []):
        meta = _meta_for_pair(pr.get("pair_id"), pair_meta)
        rows = pr.get("head_rows", [])
        rows = [_row_with_metadata(row, pair_meta) for row in rows]
        all_rows.extend(rows)
        best_margin = pr.get("best_margin_head") or pr.get("best_head") or {}
        best_patch = pr.get("best_patch_head") or _best_head(rows, "patch_recovery")
        best_margin = _row_with_metadata(best_margin, pair_meta) if best_margin else {}
        best_patch = _row_with_metadata(best_patch, pair_meta) if best_patch else {}

        if best_margin:
            margin_counter[_head_id(best_margin)] += 1
        if best_patch:
            patch_counter[_head_id(best_patch)] += 1
        if meta["headline"] and best_margin:
            headline_margin_counter[_head_id(best_margin)] += 1
        if meta["headline"] and best_patch:
            headline_patch_counter[_head_id(best_patch)] += 1

        pair_summaries.append(
            {
                "pair_id": pr.get("pair_id"),
                "family": meta["family"],
                "canonical_id": meta["canonical_id"],
                "variant": meta["variant"],
                "headline": meta["headline"],
                "clean_logit_diff": pr.get("clean_logit_diff"),
                "corrupted_logit_diff": pr.get("corrupted_logit_diff"),
                "best_margin_head": _compact_head(best_margin),
                "best_patch_head": _compact_head(best_patch),
                "best_margin_patch_recovery": _float_or_none(best_margin.get("patch_recovery")),
                "best_margin_recovery_margin": _float_or_none(best_margin.get("recovery_margin")),
                "best_margin_percentile": _percentile_rank(
                    [row.get("recovery_margin") for row in rows],
                    best_margin.get("recovery_margin"),
                ),
                "random_head_baseline": _random_head_baseline(
                    rows,
                    pair_id=str(pr.get("pair_id")),
                    seed=seed,
                    samples=random_baseline_samples,
                ),
                "best_patch_recovery": _float_or_none(best_patch.get("patch_recovery")),
                "best_patch_recovery_margin": _float_or_none(best_patch.get("recovery_margin")),
            }
        )

    finite_rows = [
        row for row in all_rows if _finite(row.get("recovery_margin"))
    ]
    mean_heads = _top_mean_margin_heads(finite_rows)
    headline_rows = [row for row in finite_rows if bool(row.get("headline", True))]

    return {
        "num_pairs": len(pair_summaries),
        "total_head_rows": len(all_rows),
        "finite_margin_rows": len(finite_rows),
        "repeated_best_margin_heads": {
            head: count for head, count in margin_counter.items() if count > 1
        },
        "repeated_best_patch_heads": {
            head: count for head, count in patch_counter.items() if count > 1
        },
        "headline_repeated_best_margin_heads": {
            head: count for head, count in headline_margin_counter.items() if count > 1
        },
        "headline_repeated_best_patch_heads": {
            head: count for head, count in headline_patch_counter.items() if count > 1
        },
        "top_mean_margin_heads": mean_heads,
        "top_headline_control_adjusted_heads": _top_mean_margin_heads(headline_rows),
        "family_summaries": _summarize_attention_groups(pair_summaries, "family"),
        "paraphrase_stability": _summarize_attention_groups(pair_summaries, "canonical_id"),
        "pair_summaries": pair_summaries,
    }

def _row_with_metadata(
    row: dict[str, Any],
    pair_meta: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    if not row:
        return {}
    meta = _meta_for_pair(row.get("pair_id"), pair_meta)
    merged = dict(row)
    merged.setdefault("family", meta["family"])
    merged.setdefault("canonical_id", meta["canonical_id"])
    merged.setdefault("variant", meta["variant"])
    merged.setdefault("headline", meta["headline"])
    if "control_adjusted_recovery" not in merged:
        merged["control_adjusted_recovery"] = merged.get("recovery_margin")
    return merged

def _percentile_rank(values: list[Any], value: Any) -> float | None:
    finite = sorted(float(v) for v in values if _finite(v))
    if not finite or not _finite(value):
        return None
    numeric = float(value)
    below_or_equal = sum(1 for v in finite if v <= numeric)
    return 100.0 * below_or_equal / len(finite)

def _random_head_baseline(
    rows: list[dict[str, Any]],
    *,
    pair_id: str,
    seed: int,
    samples: int,
) -> dict[str, Any]:
    finite = [row for row in rows if _finite(row.get("recovery_margin"))]
    if not finite:
        return {
            "samples": 0,
            "mean_recovery_margin": None,
            "median_recovery_margin": None,
            "p95_recovery_margin": None,
            "max_recovery_margin": None,
        }

    digest = hashlib.sha256(f"{seed}:{pair_id}".encode("utf-8")).hexdigest()
    rng = random.Random(int(digest[:16], 16))
    sample_count = min(max(int(samples), 1), len(finite))
    sampled = rng.sample(finite, sample_count)
    margins = sorted(float(row["recovery_margin"]) for row in sampled)
    return {
        "samples": sample_count,
        "mean_recovery_margin": sum(margins) / len(margins),
        "median_recovery_margin": _percentile_value(margins, 0.50),
        "p95_recovery_margin": _percentile_value(margins, 0.95),
        "max_recovery_margin": max(margins),
    }

def _percentile_value(sorted_values: list[float], q: float) -> float | None:
    if not sorted_values:
        return None
    if len(sorted_values) == 1:
        return sorted_values[0]
    pos = _clamp_float(q, 0.0, 1.0) * (len(sorted_values) - 1)
    lower = int(math.floor(pos))
    upper = int(math.ceil(pos))
    if lower == upper:
        return sorted_values[lower]
    frac = pos - lower
    return sorted_values[lower] * (1 - frac) + sorted_values[upper] * frac

def _clamp_float(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))

def _summarize_attention_groups(
    pair_summaries: list[dict[str, Any]],
    key: str,
) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in pair_summaries:
        grouped[str(row.get(key) or "unknown")].append(row)

    summaries: list[dict[str, Any]] = []
    for group_id, rows in sorted(grouped.items()):
        margins = [
            float(row["best_margin_recovery_margin"])
            for row in rows
            if _finite(row.get("best_margin_recovery_margin"))
        ]
        patches = [
            float(row["best_patch_recovery"])
            for row in rows
            if _finite(row.get("best_patch_recovery"))
        ]
        margin_heads = Counter(
            _head_id(row["best_margin_head"])
            for row in rows
            if row.get("best_margin_head")
        )
        patch_heads = Counter(
            _head_id(row["best_patch_head"])
            for row in rows
            if row.get("best_patch_head")
        )
        summaries.append(
            {
                key: group_id,
                "family": rows[0].get("family", "general"),
                "num_pairs": len(rows),
                "num_variants": len({row.get("variant") for row in rows}),
                "headline_pairs": sum(1 for row in rows if row.get("headline")),
                "variants": sorted(str(row.get("variant", "base")) for row in rows),
                "mean_best_margin": _mean(margins),
                "mean_best_patch_recovery": _mean(patches),
                "best_margin_heads": dict(margin_heads),
                "best_patch_heads": dict(patch_heads),
                "stable_best_margin_heads": {
                    head: count for head, count in margin_heads.items() if count > 1
                },
                "stable_best_patch_heads": {
                    head: count for head, count in patch_heads.items() if count > 1
                },
            }
        )
    return summaries

def _best_head(rows: list[dict[str, Any]], key: str) -> dict[str, Any]:
    finite = [row for row in rows if _finite(row.get(key))]
    if not finite:
        return {}
    return max(finite, key=lambda row: float(row[key]))

def _compact_head(head: dict[str, Any]) -> dict[str, Any]:
    if not head:
        return {}
    return {
        "layer": head.get("layer"),
        "head": head.get("head"),
        "patch_recovery": _float_or_none(head.get("patch_recovery")),
        "control_recovery": _float_or_none(head.get("control_recovery")),
        "recovery_margin": _float_or_none(head.get("recovery_margin")),
    }

def _head_id(head: dict[str, Any]) -> str:
    return f"L{head.get('layer')}H{head.get('head')}"

def _top_mean_margin_heads(rows: list[dict[str, Any]], limit: int = 10) -> list[dict[str, Any]]:
    grouped: dict[tuple[int, int], list[float]] = defaultdict(list)
    for row in rows:
        grouped[(int(row["layer"]), int(row["head"]))].append(float(row["recovery_margin"]))

    ranked = []
    for (layer, head), values in grouped.items():
        ranked.append(
            {
                "layer": layer,
                "head": head,
                "mean_recovery_margin": sum(values) / len(values),
                "max_recovery_margin": max(values),
                "min_recovery_margin": min(values),
                "num_pairs": len(values),
            }
        )
    return sorted(ranked, key=lambda row: row["mean_recovery_margin"], reverse=True)[:limit]

def _summarize_gradient(
    data: dict[str, Any],
    pair_meta: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    pair_summaries = []
    total = 0
    finite = 0
    top1: Counter[int] = Counter()
    top3: Counter[int] = Counter()

    for pr in data.get("pair_results", []):
        meta = _meta_for_pair(pr.get("pair_id"), pair_meta)
        rows = pr.get("layer_scores", [])
        total += len(rows)
        finite_rows = [row for row in rows if _finite(row.get("attribution"))]
        finite += len(finite_rows)
        tops = pr.get("top_5_layers", [])
        if tops:
            top1[int(tops[0]["layer"])] += 1
        for row in tops[:3]:
            top3[int(row["layer"])] += 1
        pair_summaries.append(
            {
                "pair_id": pr.get("pair_id"),
                "family": meta["family"],
                "canonical_id": meta["canonical_id"],
                "variant": meta["variant"],
                "headline": meta["headline"],
                "finite_scores": len(finite_rows),
                "total_scores": len(rows),
                "top_5_layers": [int(row["layer"]) for row in tops[:5]],
            }
        )

    return {
        "num_pairs": len(pair_summaries),
        "total_scores": total,
        "finite_scores": finite,
        "nan_scores": total - finite,
        "top1_layer_counts": dict(sorted(top1.items())),
        "top3_layer_counts": dict(sorted(top3.items())),
        "pair_summaries": pair_summaries }

def _is_nan(value: Any) -> bool:
    try:
        return math.isnan(float(value))
    except (TypeError, ValueError):
        return False

def _float_or_none(value: Any) -> float | None:
    return float(value) if _finite(value) else None

def _render_run_log(summary: dict[str, Any]) -> str:
    lines = [
        "Mechanistic Interpretability Probe Run",
        "=====================================",
        "",
        f"Run directory: {summary['run_dir']}",
        f"Model: {summary.get('model_name')}",
        f"Probes: {', '.join(summary.get('probes', []))}",
        f"Verdict: {summary['verdict']['status']} - {summary['verdict']['claim']}",
        "",
    ]

    warnings = summary["verdict"].get("warnings", [])
    if warnings:
        lines.append("Warnings:")
        for warning in warnings:
            lines.append(f"  - {warning}")
        lines.append("")

    lens = summary["logit_lens"]
    lines.append("Control Quality:")
    if lens["control_bias_warnings"]:
        for row in lens["control_bias_warnings"]:
            lines.append(
                "  "
                f"{row['pair_id']}: final_control_logit_diff="
                f"{row['final_control_logit_diff']:.4f}"
            )
    else:
        lines.append("  No control bias warnings.")
    mild = [
        row for row in lens.get("token_bias_diagnostics", [])
        if row.get("severity") == "mild_correct_prior"
    ]
    if mild:
        lines.append(f"  Mild positive control priors: {len(mild)}")
    lines.append("")

    attention = summary["attention"]
    if attention.get("family_summaries"):
        lines.append("Pair Families:")
        for row in attention["family_summaries"]:
            repeated = _head_counts(row.get("stable_best_margin_heads", {}))
            lines.append(
                "  "
                f"{row['family']}: pairs={row['num_pairs']} "
                f"headline={row['headline_pairs']} "
                f"mean_margin={_fmt(row.get('mean_best_margin'))} "
                f"stable_margin={repeated or 'none'}"
            )
        lines.append("")

    paraphrases = [
        row for row in attention.get("paraphrase_stability", [])
        if row.get("num_variants", 0) > 1
    ]
    if paraphrases:
        lines.append("Paraphrase Stability:")
        for row in paraphrases:
            stable = _head_counts(row.get("stable_best_margin_heads", {}))
            margin_heads = _head_counts(row.get("best_margin_heads", {}))
            lines.append(
                "  "
                f"{row['canonical_id']}: variants={row['num_variants']} "
                f"stable_margin={stable or 'none'} "
                f"all_margin_heads={margin_heads or 'none'}"
            )
        lines.append("")

    lines.append("Attention Highlights:")
    lines.append(
        f"  finite margins: {attention['finite_margin_rows']}/"
        f"{attention['total_head_rows']}"
    )
    for row in attention["pair_summaries"]:
        bm = row.get("best_margin_head", {})
        bp = row.get("best_patch_head", {})
        lines.append(
            "  "
            f"{row['pair_id']}: "
            f"[{row.get('family', 'general')}/{row.get('variant', 'base')}] "
            f"best-margin L{bm.get('layer')}H{bm.get('head')} "
            f"margin={_fmt(bm.get('recovery_margin'))} "
            f"patch={_fmt(bm.get('patch_recovery'))} "
            f"pct={_fmt(row.get('best_margin_percentile'))}; "
            f"best-patch L{bp.get('layer')}H{bp.get('head')} "
            f"patch={_fmt(bp.get('patch_recovery'))} "
            f"margin={_fmt(bp.get('recovery_margin'))}"
        )
    lines.append("")

    gradient = summary["gradient"]
    lines.append("Gradient Attribution:")
    lines.append(
        f"  finite scores: {gradient['finite_scores']}/{gradient['total_scores']}"
    )
    lines.append(f"  top-1 layer counts: {gradient['top1_layer_counts']}")
    lines.append(f"  top-3 layer counts: {gradient['top3_layer_counts']}")
    lines.append("")

    lines.extend(
        [
            "Interpretation Notes:",
            "  - Best-margin means control-adjusted recovery: patch recovery minus control recovery.",
            "  - Best-patch heads show the largest raw causal recovery, even if controls also move.",
            "  - Headline claims should focus on repeated heads across factual paraphrases.",
        ]
    )
    return "\n".join(lines).rstrip() + "\n"

def _render_ablation_log(summary: dict[str, Any]) -> str:
    lines = [
        "Mechanistic Interpretability Ablation Run",
        "========================================",
        "",
        f"Run directory: {summary['run_dir']}",
        f"Model: {summary.get('model_name')}",
        f"Targets: {', '.join(str(target) for target in summary.get('targets', []))}",
        f"Modes: {', '.join(str(mode) for mode in summary.get('modes', []))}",
        f"Verdict: {summary['verdict']['status']} - {summary['verdict']['claim']}",
        "",
    ]

    warnings = summary["verdict"].get("warnings", [])
    if warnings:
        lines.append("Warnings:")
        for warning in warnings:
            lines.append(f"  - {warning}")
        lines.append("")

    aggregate = summary.get("aggregate", {})
    lines.append("Aggregate Drops:")
    for row in aggregate.get("by_target_mode_scope", []):
        scope = "headline" if row.get("headline") else "specificity"
        lines.append(
            "  "
            f"{row.get('target_head')} {row.get('ablation_mode')} {scope}: "
            f"pairs={row.get('num_pairs')} "
            f"mean_drop={_fmt(row.get('mean_logit_diff_drop'))} "
            f"mean_rel={_fmt(row.get('mean_relative_drop_fraction'))} "
            f"mean_gap={_fmt(row.get('mean_gap_normalized_drop'))}"
        )
    lines.append("")

    if aggregate.get("specificity"):
        lines.append("Specificity:")
        for row in aggregate["specificity"]:
            lines.append(
                "  "
                f"{row.get('target_head')} {row.get('ablation_mode')}: "
                f"headline_drop={_fmt(row.get('headline_mean_drop'))} "
                f"control_drop={_fmt(row.get('control_mean_drop'))} "
                f"ratio={_fmt(row.get('factual_to_control_drop_ratio'))}"
            )
        lines.append("")

    lines.append("Pair Highlights:")
    for row in summary.get("pair_summaries", []):
        best = row.get("largest_drop", {})
        scope = "headline" if row.get("headline") else "specificity"
        lines.append(
            "  "
            f"{row.get('pair_id')} [{row.get('family')}/{row.get('variant')}/{scope}]: "
            f"clean={_fmt(row.get('clean_logit_diff'))} "
            f"corrupted={_fmt(row.get('corrupted_logit_diff'))}; "
            f"largest_drop={best.get('target_head')} {best.get('ablation_mode')} "
            f"drop={_fmt(best.get('logit_diff_drop'))} "
            f"gap={_fmt(best.get('gap_normalized_drop'))}"
        )
    lines.append("")

    lines.extend(
        [
            "Interpretation Notes:",
            "  - Positive drop means ablation reduced the clean correct-minus-incorrect logit diff.",
            "  - Mean ablation is the primary intervention; zero ablation is a sanity check.",
            "  - Specificity increases when factual headline drops are much larger than non-factual control drops.",
        ]
    )
    return "\n".join(lines).rstrip() + "\n"

def _render_source_log(summary: dict[str, Any]) -> str:
    lines = [
        "Mechanistic Interpretability Source Run",
        "======================================",
        "",
        f"Run directory: {summary['run_dir']}",
        f"Model: {summary.get('model_name')}",
        f"Target: {summary.get('target')}",
        f"Source layers: {summary.get('source_layers')}",
        f"Source position: {summary.get('source_position')}",
        f"Verdict: {summary['verdict']['status']} - {summary['verdict']['claim']}",
        "",
    ]

    warnings = summary["verdict"].get("warnings", [])
    if warnings:
        lines.append("Warnings:")
        for warning in warnings:
            lines.append(f"  - {warning}")
        lines.append("")

    aggregate = summary.get("aggregate", {})
    if summary.get("source_position") == "sweep":
        lines.append("Best-Token Layer Summary:")
        for row in aggregate.get("best_token_layer_summaries", []):
            top = row.get("top_tokens", [])[:3]
            token_bits = []
            for token in top:
                token_bits.append(
                    f"{token.get('pair_id')} pos={token.get('source_position')} "
                    f"{token.get('clean_source_token')}->{token.get('corrupted_source_token')}"
                )
            lines.append(
                "  "
                f"L{row.get('source_layer')}: "
                f"pairs={row.get('num_pairs')} "
                f"mean_best_target_rec={_fmt(row.get('mean_best_target_recovery_fraction'))} "
                f"mean_best_shift={_fmt(row.get('mean_best_target_output_shift'))} "
                f"mean_best_logit_rec={_fmt(row.get('mean_best_logit_recovery_fraction'))}"
            )
            if token_bits:
                lines.append("    top_tokens: " + " | ".join(token_bits))
        lines.append("")
        lines.append("All-Token Mean Diagnostic:")
    else:
        lines.append("Source Layer Summary:")

    for row in aggregate.get("source_layer_summaries", []):
        lines.append(
            "  "
            f"L{row.get('source_layer')}: "
            f"mean_target_rec={_fmt(row.get('mean_target_recovery_fraction'))} "
            f"mean_shift={_fmt(row.get('mean_target_output_shift'))} "
            f"mean_logit_rec={_fmt(row.get('mean_logit_recovery_fraction'))}"
        )
    lines.append("")

    best = summary.get("best_source_layer", {})
    if best:
        lines.append(
            "Best Source Layer: "
            f"L{best.get('source_layer')} "
            f"mean_target_rec={_fmt(_source_layer_target_score(best))} "
            f"mean_logit_rec={_fmt(_source_layer_logit_score(best))}"
        )
        lines.append("")

    lines.append("Pair Highlights:")
    for row in summary.get("pair_summaries", []):
        lines.append(
            "  "
            f"{row.get('pair_id')} [{row.get('family')}/{row.get('variant')}]: "
            f"best_source=L{row.get('best_source_layer')} "
            f"source_pos={row.get('source_position')} "
            f"token={row.get('clean_source_token')}->{row.get('corrupted_source_token')} "
            f"patched_tokens={row.get('source_positions_patched')} "
            f"target_rec={_fmt(row.get('target_recovery_fraction'))} "
            f"shift={_fmt(row.get('target_output_shift'))} "
            f"logit_rec={_fmt(row.get('logit_recovery_fraction'))}"
        )
    lines.append("")

    lines.extend(
        [
            "Interpretation Notes:",
            "  - target_rec measures how much the target head output moves along its clean-minus-corrupted direction.",
            "  - logit_rec measures downstream logit recovery from the same source-layer patch.",
            "  - Source candidates should show positive target recovery across multiple factual variants.",
        ]
    )
    return "\n".join(lines).rstrip() + "\n"

def _head_counts(counts: dict[str, Any]) -> str:
    if not counts:
        return ""
    return ", ".join(
        f"{head}x{count}" for head, count in sorted(counts.items(), key=lambda item: (-int(item[1]), item[0]))
    )

def _fmt(value: Any) -> str:
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return "nan"
    if math.isinf(numeric):
        return "inf" if numeric > 0 else "-inf"
    return f"{numeric:.4f}" if math.isfinite(numeric) else "nan"

def _create_run_zip(run_dir: Path, archive_root: Path) -> Path:
    archive_root.mkdir(parents=True, exist_ok=True)
    archive_path = archive_root / f"{run_dir.name}.zip"
    with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
        for path in sorted(run_dir.rglob("*")):
            if path.is_file():
                archive.write(path, path.relative_to(run_dir).as_posix())
    return archive_path

def inspect_results_root(results_root: Path) -> int:
    packages = sorted(
        list(results_root.glob("mech_results_*"))
        + list(results_root.glob("mech_run_*"))
        + list(results_root.glob("mech_ablation_*"))
        + list(results_root.glob("mech_source_*")),
        key=lambda path: path.name,
    )
    if not packages:
        print(f"No result packages found in {results_root}")
        return 1

    latest = packages[-1]
    print(f"Latest package: {latest}")

    for log_name in ("logs/run_log.txt", "run_log.txt", "validation_log.txt"):
        log_path = latest / log_name
        if log_path.exists():
            print("")
            print(log_path.read_text(encoding="utf-8").rstrip())
            break

    _print_probe_summary(latest)

    report = latest / "report.html"
    if report.exists():
        print(f"\nHTML report: {report}")

    package_manifest = latest / "package_manifest.json"
    integrity = latest / "integrity.json"
    print("")
    print(f"Package manifest: {'present' if package_manifest.exists() else 'missing'}")
    print(f"Integrity: {'present' if integrity.exists() else 'missing'}")
    return 0

def _print_probe_summary(run_dir: Path) -> None:
    for probe in ("residual", "attention", "logit_lens", "gradient", "ablation", "source"):
        path = run_dir / "raw" / f"{probe}_results.json"
        if not path.exists():
            path = run_dir / f"{probe}_results.json"
        if not path.exists():
            continue
        data = json.loads(path.read_text(encoding="utf-8"))
        pair_results = data.get("pair_results", [])
        print(f"\n[{probe}] {len(pair_results)} pair(s) recorded")
        for pr in pair_results:
            pid = pr.get("pair_id", "?")
            if probe == "attention":
                _print_attention_pair(pid, pr)
            elif probe == "logit_lens":
                ctrl = pr.get("final_control_logit_diff")
                flag = " [WARN] control biased" if _finite(ctrl) and float(ctrl) > 0.5 else ""
                print(
                    f"  {pid}: first positive layer = {pr.get('first_positive_layer', '?')}, "
                    f"control L23 = {_fmt(ctrl)}{flag}"
                )
            elif probe == "gradient":
                top = pr.get("top_5_layers", [])
                top_str = ", ".join(f"L{row['layer']}" for row in top[:3])
                print(f"  {pid}: top layers = {top_str}")
            elif probe == "ablation":
                _print_ablation_pair(pid, pr)
            elif probe == "source":
                _print_source_pair(pid, pr)

def _print_attention_pair(pid: str, pr: dict[str, Any]) -> None:
    bm = pr.get("best_margin_head") or pr.get("best_head", {})
    bp = pr.get("best_patch_head", {})
    if bm:
        print(
            f"  {pid}: best-margin L{bm.get('layer', '?')}H{bm.get('head', '?')} "
            f"margin={_fmt(bm.get('recovery_margin'))} "
            f"patch={_fmt(bm.get('patch_recovery'))}"
        )
    if bp and bp != bm:
        print(
            f"  {pid}: best-patch  L{bp.get('layer', '?')}H{bp.get('head', '?')} "
            f"patch={_fmt(bp.get('patch_recovery'))} "
            f"margin={_fmt(bp.get('recovery_margin'))}"
        )

def _print_ablation_pair(pid: str, pr: dict[str, Any]) -> None:
    rows = pr.get("head_results", [])
    best = _best_ablation_row(rows, "logit_diff_drop")
    if not best:
        print(f"  {pid}: no finite ablation drops")
        return
    print(
        f"  {pid}: largest drop {best.get('target_head')} "
        f"{best.get('ablation_mode')} "
        f"drop={_fmt(best.get('logit_diff_drop'))} "
        f"gap={_fmt(best.get('gap_normalized_drop'))}"
    )

def _print_source_pair(pid: str, pr: dict[str, Any]) -> None:
    best = pr.get("best_source_layer") or {}
    if not best:
        print(f"  {pid}: no finite source recovery")
        return
    print(
        f"  {pid}: best source L{best.get('source_layer')} "
        f"target_rec={_fmt(best.get('target_recovery_fraction'))} "
        f"logit_rec={_fmt(best.get('logit_recovery_fraction'))}"
    )

def _cmd_inspect(args: argparse.Namespace) -> int:
    return inspect_results_root(Path(args.results_root))

def _cmd_package(args: argparse.Namespace) -> int:
    import zipfile as _zf

    out = Path(args.output or "mech_workbench_core.zip").resolve()
    out.parent.mkdir(parents=True, exist_ok=True)
    include_dirs = ["mech_workbench", "configs", "prompt curation", "scripts"]
    include_files = [
        "pyproject.toml",
        "README.md",
        "requirements.txt",
    ]

    with _zf.ZipFile(out, "w", compression=_zf.ZIP_DEFLATED) as zf:
        for d in include_dirs:
            p = ROOT / d
            if p.exists():
                for f in p.rglob("*"):
                    if f.is_file() and "__pycache__" not in str(f):
                        zf.write(f, f.relative_to(ROOT))
        for name in include_files:
            p = ROOT / name
            if p.exists():
                zf.write(p, name)
    print(f"Package written: {out}")
    return 0

def _cmd_prompt_validate(args: argparse.Namespace) -> int:
    """Validate prompt pairs against the loaded model."""
    from mech_workbench.config import load_config
    from mech_workbench.validation import validate_prompt_pairs

    config = load_config(args.config)
    if getattr(args, "limit", None):
        config = dataclasses.replace(config, prompt_pairs=config.prompt_pairs[: args.limit])

    from mech_workbench.model_loader import load_model_and_tokenizer
    model, tokenizer = load_model_and_tokenizer(
        config.model_name, device=config.device, dtype=config.dtype
    )
    results = validate_prompt_pairs(model, tokenizer, config.prompt_pairs)
    total = len(results)
    valid = sum(1 for r in results if r["status"] == "valid")
    flagged = sum(1 for r in results if r["status"] == "flagged")
    excluded = sum(1 for r in results if r["status"] == "excluded")
    print(f"\nValidation: {valid}/{total} valid, {flagged} flagged, {excluded} excluded\n")
    for r in results:
        sym = {"valid": "[OK]", "flagged": "[!!]", "excluded": "[XX]"}.get(r["status"], "[??]")
        print(f"  {sym} {r['pair_id']:30s} {r['status']:10s}  {r.get('reason', '')}")
    if excluded > 0 or flagged > 0:
        print(f"\n[FAIL] {flagged + excluded} prompt(s) did not pass validation.")
        print("Fix or remove flagged pairs before running discovery.")
        return 1
    print("\n[OK] All prompts passed validation.")
    return 0

def _cmd_verify(args: argparse.Namespace) -> int:
    """Verify a completed run bundle."""
    from mech_workbench.bundle import verify_bundle
    import glob

    patterns = args.run_path
    paths: list[Path] = []
    for pattern in patterns:
        matched = glob.glob(pattern)
        if matched:
            paths.extend(Path(m) for m in matched)
        else:
            paths.append(Path(pattern))

    if not paths:
        print("[error] No run paths specified.")
        return 1

    all_ok = True
    for path in paths:
        ok = verify_bundle(path)
        if not ok:
            all_ok = False
    return 0 if all_ok else 1

def _write_run_card(run_dir: Path, config: Any, stages: list[str]) -> None:
    card_path = run_dir / "RUN_CARD.json"
    existing: dict[str, Any] = {}
    if card_path.exists():
        try:
            existing = json.loads(card_path.read_text(encoding="utf-8"))
        except Exception:
            pass
    completed = list(existing.get("stages_completed", []))
    for s in stages:
        if s not in completed:
            completed.append(s)
    card = {
        "created_at_utc": existing.get("created_at_utc", datetime.now(timezone.utc).isoformat()),
        "updated_at_utc": datetime.now(timezone.utc).isoformat(),
        "model": getattr(config, "model_name", None),
        "prompt_pairs": len(getattr(config, "prompt_pairs", [])),
        "stages_completed": completed,
        "claim_status": "pending_review",
    }
    card_path.write_text(json.dumps(card, indent=2) + "\n", encoding="utf-8")

def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="mech",
        description="Mechanistic Interpretability Workbench CLI",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_run = sub.add_parser("run", help="Load model and run interpretability probes.")
    p_run.add_argument(
        "--config",
        default=str(ROOT / "configs" / "discovery_4b.yaml"),
        help="Path to experiment YAML config.",
    )
    p_run.add_argument(
        "--probe",
        default=None,
        metavar="PROBES",
        help=(
            "Comma-separated probe names. Choices: residual, attention, "
            "logit_lens, gradient. Default suite: logit_lens,gradient,attention."
        ),
    )
    p_run.add_argument(
        "--limit",
        type=int,
        default=None,
        metavar="N",
        help="Only process the first N prompt pairs.",
    )
    p_run.add_argument(
        "--results-root",
        default=str(ROOT / "results"),
        help="Root directory for timestamped run folders.",
    )
    p_run.add_argument(
        "--archive-root",
        default=None,
        help="Optional directory where a downloadable <run_name>.zip is written.",
    )
    p_run.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate config and print summary without loading the model.",
    )

    p_abl = sub.add_parser("ablate", help="Run targeted attention-head ablations.")
    p_abl.add_argument(
        "--config",
        default=str(ROOT / "configs" / "ablation_4b.yaml"),
        help="Path to ablation YAML config.",
    )
    p_abl.add_argument(
        "--target",
        default="L31H1",
        help="Primary head target, e.g. L31H1. Comma-separated targets are allowed.",
    )
    p_abl.add_argument(
        "--also",
        default="",
        help="Optional comma-separated comparison heads. Defaults to none.",
    )
    p_abl.add_argument(
        "--mode",
        default="both",
        help="Ablation mode: mean, zero, mean,zero, or both.",
    )
    p_abl.add_argument(
        "--combined",
        action="store_true",
        help="Ablate all targets simultaneously in one forward pass.",
    )
    p_abl.add_argument(
        "--limit",
        type=int,
        default=None,
        metavar="N",
        help="Only process the first N prompt pairs.",
    )
    p_abl.add_argument(
        "--results-root",
        default=str(ROOT / "results"),
        help="Root directory for timestamped ablation folders.",
    )
    p_abl.add_argument(
        "--archive-root",
        default=None,
        help="Optional directory where a downloadable <run_name>.zip is written.",
    )
    p_abl.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate config and ablation settings without loading the model.",
    )

    p_src = sub.add_parser("source", help="Probe source layers for a target attention head.")
    p_src.add_argument(
        "--config",
        default=str(ROOT / "configs" / "discovery_4b.yaml"),
        help="Path to source-probing YAML config.",
    )
    p_src.add_argument(
        "--target",
        default="L31H1",
        help="Target head, e.g. L31H1.",
    )
    p_src.add_argument(
        "--source-layers",
        default="15-30",
        help="Candidate source layers, e.g. 15-30 or 12,15,18-28.",
    )
    p_src.add_argument(
        "--source-position",
        default="sweep",
        help="Residual token positions to restore: sweep, all, last, first, or an integer.",
    )
    p_src.add_argument(
        "--limit",
        type=int,
        default=None,
        metavar="N",
        help="Only process the first N prompt pairs.",
    )
    p_src.add_argument(
        "--results-root",
        default=str(ROOT / "results"),
        help="Root directory for timestamped source-probing folders.",
    )
    p_src.add_argument(
        "--archive-root",
        default=None,
        help="Optional directory where a downloadable <run_name>.zip is written.",
    )
    p_src.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate config and source settings without loading the model.",
    )

    p_pval = sub.add_parser("validate", help="Validate prompt pairs against the model.")
    p_pval.add_argument(
        "--config",
        default=str(ROOT / "configs" / "discovery_4b.yaml"),
        help="Path to experiment YAML config.",
    )
    p_pval.add_argument(
        "--limit",
        type=int,
        default=None,
        metavar="N",
        help="Only validate the first N prompt pairs.",
    )

    p_vfy = sub.add_parser("verify", help="Verify a completed run bundle.")
    p_vfy.add_argument(
        "run_path",
        nargs="+",
        metavar="PATH",
        help="Path(s) to run directory/directories to verify.",
    )

    p_ins = sub.add_parser("inspect", help="Print the latest result summary.")
    p_ins.add_argument("--results-root", default=str(ROOT / "results"))

    p_pkg = sub.add_parser("package", help="Build a distributable zip of the codebase.")
    p_pkg.add_argument("--output", default=None)

    return parser

def main() -> int:
    parser = _build_parser()
    args = parser.parse_args()
    dispatch = {
        "run": _cmd_run,
        "ablate": _cmd_ablate,
        "source": _cmd_source,
        "validate": _cmd_prompt_validate,
        "verify": _cmd_verify,
        "inspect": _cmd_inspect,
        "package": _cmd_package,
    }
    return dispatch[args.command](args)

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