evidence/source/e3_rollout_v3/scripts/eval_pairs.py

"""
Construct provenance-aware paired examples for the acceptance experiment.

This module intentionally has no MLX import. Pair construction, null matching,
and holdout validation can therefore be tested without a Metal device.
"""

import hashlib
import json
from collections import Counter

import numpy as np


def bootstrap_mean_ci(values, n_boot=2000, seed=0):
    """Percentile interval for a mean, resampling the supplied cluster units."""
    values = np.asarray(values, dtype=np.float64)
    if values.ndim != 1:
        raise ValueError("bootstrap values must be one-dimensional")
    if values.size == 0:
        return (float("nan"), float("nan"))
    if not np.all(np.isfinite(values)):
        raise ValueError("bootstrap values must be finite")
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, values.size, size=(n_boot, values.size))
    means = values[idx].mean(axis=1)
    return float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))


def paired_ratio_ci(num_by_doc, den_by_doc, n_boot=4000, seed=0):
    """
    Bootstrap a ratio of arm means while preserving document-level pairing.

    Each input element is one document's mean acceptance for an arm. Token
    positions inside a document are correlated and are not resampled as if they
    were independent.
    """
    num = np.asarray(num_by_doc, dtype=np.float64)
    den = np.asarray(den_by_doc, dtype=np.float64)
    if num.ndim != 1 or den.ndim != 1:
        raise ValueError("paired bootstrap inputs must be one-dimensional")
    if num.size != den.size:
        raise ValueError(
            f"paired bootstrap length mismatch: {num.size} versus {den.size}"
        )
    if num.size == 0:
        return (float("nan"), float("nan"), float("nan"))
    if not np.all(np.isfinite(num)) or not np.all(np.isfinite(den)):
        raise ValueError("paired bootstrap inputs must be finite")
    den_mean = den.mean()
    point = float(num.mean() / den_mean) if den_mean else float("nan")

    rng = np.random.default_rng(seed)
    idx = rng.integers(0, num.size, size=(n_boot, num.size))
    num_means = num[idx].mean(axis=1)
    den_means = den[idx].mean(axis=1)
    with np.errstate(divide="ignore", invalid="ignore"):
        ratios = num_means / den_means
    ratios = ratios[np.isfinite(ratios)]
    if ratios.size == 0:
        return (point, float("nan"), float("nan"))
    return (
        point,
        float(np.percentile(ratios, 2.5)),
        float(np.percentile(ratios, 97.5)),
    )


def paired_ratio_difference_ci(
    trained_num_by_doc,
    trained_den_by_doc,
    control_num_by_doc,
    control_den_by_doc,
    n_boot=4000,
    seed=0,
):
    """Bootstrap trained minus control ratios over the same target documents."""
    arrays = [
        np.asarray(values, dtype=np.float64)
        for values in (
            trained_num_by_doc,
            trained_den_by_doc,
            control_num_by_doc,
            control_den_by_doc,
        )
    ]
    if any(values.ndim != 1 for values in arrays):
        raise ValueError("control comparison inputs must be one-dimensional")
    sizes = {values.size for values in arrays}
    if len(sizes) != 1:
        raise ValueError("control comparison inputs must have equal length")
    size = arrays[0].size
    if size < 2:
        raise ValueError("control comparison requires at least two documents")
    if any(not np.all(np.isfinite(values)) for values in arrays):
        raise ValueError("control comparison inputs must be finite")
    trained_num, trained_den, control_num, control_den = arrays
    if trained_den.mean() <= 0 or control_den.mean() <= 0:
        raise ValueError("control comparison denominators must have positive means")

    trained_ratio = float(trained_num.mean() / trained_den.mean())
    control_ratio = float(control_num.mean() / control_den.mean())
    difference = trained_ratio - control_ratio
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, size, size=(n_boot, size))
    trained_boot = trained_num[idx].mean(axis=1) / trained_den[idx].mean(axis=1)
    control_boot = control_num[idx].mean(axis=1) / control_den[idx].mean(axis=1)
    differences = trained_boot - control_boot
    differences = differences[np.isfinite(differences)]
    if differences.size != n_boot:
        raise ValueError("control comparison bootstrap produced non-finite values")
    return (
        trained_ratio,
        control_ratio,
        difference,
        float(np.percentile(differences, 2.5)),
        float(np.percentile(differences, 97.5)),
    )


def iter_holdout(path):
    """Read provenance-bearing holdout records from JSONL."""
    document_ids = set()
    content_hashes = set()
    with open(path, encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_no}: invalid JSON: {exc}") from exc
            missing = [
                key for key in ("text", "language", "repo", "document_id")
                if not row.get(key)
            ]
            if missing:
                raise ValueError(
                    f"{path}:{line_no}: missing required fields {missing}"
                )
            digest = hashlib.sha256(row["text"].encode("utf-8")).hexdigest()
            declared = row.get("content_sha256")
            if declared and declared != digest:
                raise ValueError(
                    f"{path}:{line_no}: content_sha256 does not match text"
                )
            if row["document_id"] in document_ids:
                raise ValueError(
                    f"{path}:{line_no}: duplicate document_id "
                    f"{row['document_id']}"
                )
            if digest in content_hashes:
                raise ValueError(
                    f"{path}:{line_no}: duplicate document content"
                )
            document_ids.add(row["document_id"])
            content_hashes.add(digest)
            row["content_sha256"] = digest
            yield row


def iter_local_records(texts):
    """Label legacy local text input so reports cannot mistake it for a holdout."""
    for i, text in enumerate(texts):
        yield {
            "text": text,
            "language": "unknown",
            "repo": "local-untracked",
            "document_id": f"local-{i}",
        }


def file_sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def validate_publication_inputs(
    receipt_path,
    holdout_path,
    tokenizer_path,
    settings,
    checkpoint_meta,
    role,
):
    """Fail closed unless files, settings, and checkpoint match the contract."""
    with open(receipt_path, encoding="utf-8") as f:
        receipt = json.load(f)
    if receipt.get("schema_version") != 2:
        raise ValueError("publication receipt must use schema_version 2")
    contract = receipt.get("analysis_contract")
    if not isinstance(contract, dict):
        raise ValueError("receipt has no analysis_contract")
    if role not in ("trained", "untrained-control"):
        raise ValueError(f"invalid publication role {role!r}")

    problems = []
    holdout_hash = file_sha256(holdout_path)
    tokenizer_hash = file_sha256(tokenizer_path)
    expected_holdout = receipt.get("clean_holdout", {}).get("sha256")
    expected_tokenizer = receipt.get("pair_readiness", {}).get(
        "tokenizer_sha256"
    )
    if holdout_hash != expected_holdout:
        problems.append(
            f"holdout sha256 {holdout_hash} != registered {expected_holdout}"
        )
    if tokenizer_hash != expected_tokenizer:
        problems.append(
            f"tokenizer sha256 {tokenizer_hash} != registered {expected_tokenizer}"
        )

    expected_settings = contract.get("settings", {})
    for key, expected in expected_settings.items():
        actual = settings.get(key)
        if actual != expected:
            problems.append(f"{key} {actual!r} != registered {expected!r}")
    for key in (
        "instrument_version",
        "primary_comparison",
        "primary_depth",
        "bootstrap_unit",
        "sensitivity_subset",
    ):
        actual = settings.get(key)
        expected = contract.get(key)
        if actual != expected:
            problems.append(f"{key} {actual!r} != registered {expected!r}")

    expected_args = contract.get("model_args", {})
    actual_args = checkpoint_meta.get("model_args", {})
    for key, expected in expected_args.items():
        actual = actual_args.get(key)
        if actual != expected:
            problems.append(
                f"model_args.{key} {actual!r} != registered {expected!r}"
            )

    checkpoint_cfg = checkpoint_meta.get("config", {})
    if checkpoint_cfg.get("seed") != contract.get("model_seed"):
        problems.append(
            f"model seed {checkpoint_cfg.get('seed')!r} != registered "
            f"{contract.get('model_seed')!r}"
        )
    if role == "trained":
        expected_step = contract.get("trained_checkpoint_step")
        if checkpoint_meta.get("step") != expected_step:
            problems.append(
                f"trained checkpoint step {checkpoint_meta.get('step')!r} "
                f"!= registered {expected_step!r}"
            )
        expected_run = contract.get("trained_run_name")
        if checkpoint_cfg.get("run_name") != expected_run:
            problems.append(
                f"run name {checkpoint_cfg.get('run_name')!r} "
                f"!= registered {expected_run!r}"
            )
    else:
        control = contract.get("untrained_control", {})
        if checkpoint_meta.get("step") != control.get("step"):
            problems.append(
                f"untrained control step {checkpoint_meta.get('step')!r} "
                f"!= registered {control.get('step')!r}"
            )
        if checkpoint_cfg.get("lr") != control.get("learning_rate"):
            problems.append(
                f"untrained control lr {checkpoint_cfg.get('lr')!r} "
                f"!= registered {control.get('learning_rate')!r}"
            )
        if checkpoint_cfg.get("run_name") != control.get("run_name"):
            problems.append(
                f"untrained control run name "
                f"{checkpoint_cfg.get('run_name')!r} "
                f"!= registered {control.get('run_name')!r}"
            )
        if checkpoint_cfg.get("initialization_only") is not control.get(
            "initialization_only"
        ):
            problems.append(
                f"initialization_only "
                f"{checkpoint_cfg.get('initialization_only')!r} "
                f"!= registered {control.get('initialization_only')!r}"
            )
        if checkpoint_meta.get("optimizer_state_included") is not control.get(
            "optimizer_state_included"
        ):
            problems.append(
                f"optimizer_state_included "
                f"{checkpoint_meta.get('optimizer_state_included')!r} "
                f"!= registered {control.get('optimizer_state_included')!r}"
            )

    if problems:
        raise ValueError(
            "publication contract mismatch:\n- " + "\n- ".join(problems)
        )
    return {
        "receipt_path": receipt_path,
        "receipt_sha256": file_sha256(receipt_path),
        "role": role,
        "holdout_sha256": holdout_hash,
        "tokenizer_sha256": tokenizer_hash,
        "registered_at": receipt.get("analysis_contract_registered_at"),
    }, receipt


def guard_frozen_holdout_exploration(
    receipt_path, holdout_path, post_training_acknowledged
):
    """Prevent an accidental blind-holdout peek during model development."""
    with open(receipt_path, encoding="utf-8") as f:
        receipt = json.load(f)
    frozen_hash = receipt.get("clean_holdout", {}).get("sha256")
    supplied_hash = file_sha256(holdout_path)
    if supplied_hash == frozen_hash and not post_training_acknowledged:
        raise ValueError(
            "refusing to expose the frozen publication holdout to an "
            "exploratory checkpoint before training is frozen; after all "
            "training decisions are final, add "
            "--post-training-frozen-holdout"
        )


def validate_pair_summary(receipt, summary):
    """Check the constructed target and decoy set against its frozen receipt."""
    expected = receipt.get("pair_readiness", {})
    fields = (
        "paired_examples",
        "repositories",
        "decoy_match",
        "language_counts",
        "unique_decoy_documents",
        "max_decoy_reuse",
        "decoy_reuse_histogram",
        "disjoint_sensitivity_pairs",
    )
    problems = []
    for key in fields:
        if summary.get(key) != expected.get(key):
            problems.append(
                f"{key} {summary.get(key)!r} != registered {expected.get(key)!r}"
            )
    if problems:
        raise ValueError(
            "constructed pair set does not match receipt:\n- "
            + "\n- ".join(problems)
        )


def disjoint_pair_indices(metadata):
    """
    Greedily retain pairs whose target and decoy documents have not appeared.

    Input order is frozen by the registered holdout seed. The resulting
    sensitivity subset contains no document in more than one target-decoy pair.
    """
    used = set()
    selected = []
    for index, row in enumerate(metadata):
        target = row["document_id"]
        decoy = row["decoy_document_id"]
        if target in used or decoy in used:
            continue
        selected.append(index)
        used.update((target, decoy))
    return selected


def summarize_pair_dependencies(metadata):
    """Report shuffled-decoy reuse and the document-disjoint subset size."""
    reuse = Counter(row["decoy_document_id"] for row in metadata)
    histogram = Counter(reuse.values())
    return {
        "unique_decoy_documents": len(reuse),
        "max_decoy_reuse": max(reuse.values(), default=0),
        "decoy_reuse_histogram": {
            str(count): documents
            for count, documents in sorted(histogram.items())
        },
        "disjoint_sensitivity_pairs": len(disjoint_pair_indices(metadata)),
    }


def lexical_profile(tok, ids):
    """
    Small format profile used only to choose a null suffix.

    Suffix token length is already exact. These character ratios keep the decoy
    close in layout and lexical texture so the primary comparison is less able
    to win merely because the null suffix looks like a different kind of file.
    """
    text = tok.decode(ids)
    n = max(len(text), 1)
    return np.asarray([
        text.count("\n") / n,
        sum(c.isspace() for c in text) / n,
        sum(c.isalnum() or c == "_" for c in text) / n,
        sum(c.isdigit() for c in text) / n,
        sum(c in "{}[]();,:.=+-*/" for c in text) / n,
        len(set(ids)) / max(len(ids), 1),
    ], dtype=np.float64)


def choose_decoy(staged, index):
    """
    Match a decoy by language, repository independence, and lexical profile.

    The filters are relaxed only when the available pool cannot satisfy them.
    Every relaxation is recorded in pair metadata and summarized in the report.
    """
    target = staged[index]
    others = [(j, item) for j, item in enumerate(staged) if j != index]
    same_language = [
        (j, item) for j, item in others
        if item["language"] == target["language"]
    ]
    language_pool = same_language or others
    different_repo = [
        (j, item) for j, item in language_pool
        if item["repo"] != target["repo"]
    ]
    pool = different_repo or language_pool
    if not pool:
        raise ValueError("a shuffled-suffix control requires two documents")

    _, decoy = min(
        pool,
        key=lambda candidate: (
            float(np.linalg.norm(
                target["suffix_profile"] - candidate[1]["suffix_profile"]
            )),
            candidate[0],
        ),
    )
    if decoy["language"] != target["language"]:
        quality = "different_language"
    elif decoy["repo"] == target["repo"]:
        quality = "same_repo"
    else:
        quality = "matched"
    distance = float(np.linalg.norm(
        target["suffix_profile"] - decoy["suffix_profile"]
    ))
    return decoy, quality, distance


def build_pairs(records, tok, sentinels, n_examples, prefix_len, span_len,
                suffix_len, rng):
    """
    One document yields one L2R, FIM, and shuffled-suffix triple.

    The decoy has the exact same token length as the true suffix. It is selected
    from the same language, from a different repository when possible, and by
    nearest lexical profile. A positive claim still requires both FIM over the
    shuffled null and FIM over L2R.
    """
    pairs = []
    need = prefix_len + span_len + suffix_len + 8
    staged = []
    pool_limit = max(n_examples + 1, n_examples * 4)
    for record in records:
        if len(staged) >= pool_limit:
            break
        ids = tok.encode(record["text"]).ids
        if len(ids) < need:
            continue
        p0 = int(rng.integers(0, len(ids) - need + 1))
        a = p0 + prefix_len
        b = a + span_len
        c = b + suffix_len
        suffix = ids[b:c]
        staged.append({
            "prefix": ids[p0:a],
            "middle": ids[a:b],
            "suffix": suffix,
            "suffix_profile": lexical_profile(tok, suffix),
            "language": str(record.get("language") or "unknown"),
            "repo": str(record.get("repo") or "unknown"),
            "document_id": str(record.get("document_id") or len(staged)),
            "content_sha256": record.get("content_sha256"),
            "token_offset": p0,
        })

    if len(staged) < 2:
        return []

    for i, target in enumerate(staged[:n_examples]):
        decoy_item, match_quality, match_distance = choose_decoy(staged, i)
        prefix = target["prefix"]
        middle = target["middle"]
        suffix = target["suffix"]
        decoy = decoy_item["suffix"]

        l2r = prefix + middle
        fim = ([sentinels["prefix"]] + prefix
               + [sentinels["suffix"]] + suffix
               + [sentinels["middle"]] + middle)
        shuf = ([sentinels["prefix"]] + prefix
                + [sentinels["suffix"]] + decoy
                + [sentinels["middle"]] + middle)

        pairs.append({
            "l2r": (l2r, (len(prefix), len(prefix) + len(middle))),
            "fim": (fim, (len(fim) - len(middle), len(fim))),
            "fim_shuf": (shuf, (len(shuf) - len(middle), len(shuf))),
            "_meta": {
                "language": target["language"],
                "repo": target["repo"],
                "document_id": target["document_id"],
                "content_sha256": target["content_sha256"],
                "token_offset": target["token_offset"],
                "decoy_language": decoy_item["language"],
                "decoy_repo": decoy_item["repo"],
                "decoy_document_id": decoy_item["document_id"],
                "decoy_match": match_quality,
                "decoy_lexical_distance": match_distance,
            },
        })
    return pairs


def bin_acceptance_by_baseline_nll(acceptance_by_depth, baseline_nlls, edges):
    """
    Bin every treatment arm by the same corresponding L2R token difficulty.

    At draft depth j+1, acceptance position t targets the token whose baseline
    NLL is at t+j+1. Using an arm's own NLL would condition on a post-treatment
    variable because suffix visibility changes that distribution.
    """
    rows = []
    for j, per_document in enumerate(acceptance_by_depth):
        bins = []
        for k in range(len(edges) - 1):
            values = []
            for acceptance, baseline_nll in zip(
                    per_document, baseline_nlls):
                aligned = np.asarray(baseline_nll)[j + 1:]
                acceptance = np.asarray(acceptance)
                n = min(len(acceptance), len(aligned))
                if n <= 0:
                    continue
                keep = (
                    (aligned[:n] >= edges[k])
                    & (aligned[:n] < edges[k + 1])
                )
                values.append(acceptance[:n][keep])
            selected = np.concatenate(values) if values else np.zeros(0)
            bins.append({
                "bin": k,
                "nll_range": [
                    float(edges[k]) if np.isfinite(edges[k]) else None,
                    (
                        float(edges[k + 1])
                        if np.isfinite(edges[k + 1])
                        else None
                    ),
                ],
                "acceptance": (
                    float(selected.mean()) if selected.size else None
                ),
                "n": int(selected.size),
            })
        rows.append({"depth": j + 1, "bins": bins})
    return rows