Module 15. Reasoning Models and Test-Time Compute

# Install only missing dependencies; ordinary Colab already has them.
import importlib.util
import subprocess
import sys

required = {"numpy": "numpy", "matplotlib": "matplotlib"}
missing = [package for module, package in required.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import hashlib
import json
import math
import re
from dataclasses import dataclass
from itertools import product
from typing import Iterator, Sequence

import matplotlib.pyplot as plt
import numpy as np

np.set_printoptions(precision=4, suppress=True)
plt.rcParams["font.family"] = "DejaVu Sans"
plt.rcParams["mathtext.fontset"] = "dejavusans"
MSG = {'setup_ok': 'Environment ready.', 'reference_ok': 'Module-15 reference payload unpacked and verified', 'checker_ok': 'The strict checker accepts the contract and rejects ambiguous strings', 'vote_ok': 'The exact curve and Monte Carlo agree; the N=3 dip depends on tie policy', 'verifier_ok': 'Perfect coverage, a strong verifier, and a weak verifier are kept distinct', 'batch_ok': 'Shapes, selection indices, oracle gap, and full generation cost verified', 'adaptive_ok': 'On this mixture, adaptive selection inspected fewer candidates; quality and coverage were measured separately', 'cost_ok': 'The latency/FLOPs/KV ladder and two concurrency models reproduced', 'mix_ok': 'The adaptive mix and non-monotone toy curve verified', 'failure_ok': 'Expected errors caught; hidden NaNs and negative memory do not pass'}
print(MSG["setup_ok"])

INK="#2E3440"; MUTED="#4C566A"; BORDER="#D8DEE9"
def plot_style():
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "font.family": "DejaVu Sans", "font.size": 12,
        "text.color": INK, "axes.labelcolor": INK, "axes.titlecolor": INK,
        "axes.titlesize": 17, "axes.titleweight": "bold", "axes.titlelocation": "left",
        "axes.labelsize": 13, "xtick.labelsize": 11, "ytick.labelsize": 11,
        "xtick.color": MUTED, "ytick.color": MUTED,
        "axes.edgecolor": BORDER, "axes.linewidth": .7,
        "axes.facecolor": "white", "figure.facecolor": "white", "savefig.facecolor": "white",
        "grid.color": BORDER, "grid.linewidth": .7, "axes.axisbelow": True,
        "legend.fontsize": 11, "legend.frameon": False, "savefig.dpi": 150,
        "xtick.major.size": 0, "ytick.major.size": 0,
    })

plot_style()
plt.rcParams["axes.prop_cycle"] = plt.cycler(color=["#5E81AC","#D08770","#B48EAD","#A3BE8C"])

REFERENCE_JSON = '{\r\n  "adaptive_mix": {\r\n    "average_budget": 4350.0,\r\n    "budgets": [\r\n      500,\r\n      4000,\r\n      32000\r\n    ],\r\n    "classes": [\r\n      "simple",\r\n      "medium",\r\n      "hard"\r\n    ],\r\n    "saving_factor": 7.35632183908046,\r\n    "uniform_budget": 32000,\r\n    "weights": [\r\n      0.7,\r\n      0.2,\r\n      0.1\r\n    ]\r\n  },\r\n  "adaptive_search": {\r\n    "metrics": {\r\n      "adaptive_accuracy": 0.88556,\r\n      "adaptive_mean_candidates": 4.02068,\r\n      "fixed_best_of_n_accuracy": 0.9403,\r\n      "one_shot_accuracy": 0.63238,\r\n      "oracle_coverage": 0.99738,\r\n      "threshold_crossing_rate": 0.98076\r\n    },\r\n    "parameters": {\r\n      "difficulty_weights": [\r\n        0.7,\r\n        0.2,\r\n        0.1\r\n      ],\r\n      "max_candidates": 16,\r\n      "n_prompts": 50000,\r\n      "p_correct_by_difficulty": [\r\n        0.75,\r\n        0.45,\r\n        0.2\r\n      ],\r\n      "seed": 13,\r\n      "threshold": 0.82\r\n    }\r\n  },\r\n  "candidate_batch": {\r\n    "answer_ids": [\r\n      [\r\n        0,\r\n        1,\r\n        0,\r\n        2,\r\n        0\r\n      ],\r\n      [\r\n        3,\r\n        4,\r\n        5,\r\n        3,\r\n        6\r\n      ],\r\n      [\r\n        7,\r\n        8,\r\n        7,\r\n        8,\r\n        9\r\n      ],\r\n      [\r\n        10,\r\n        11,\r\n        12,\r\n        13,\r\n        14\r\n      ],\r\n      [\r\n        20,\r\n        21,\r\n        22,\r\n        23,\r\n        20\r\n      ]\r\n    ],\r\n    "correct_answer_ids": [\r\n      0,\r\n      3,\r\n      7,\r\n      14,\r\n      24\r\n    ],\r\n    "token_counts": [\r\n      [\r\n        120,\r\n        96,\r\n        144,\r\n        110,\r\n        132\r\n      ],\r\n      [\r\n        210,\r\n        170,\r\n        190,\r\n        225,\r\n        160\r\n      ],\r\n      [\r\n        88,\r\n        105,\r\n        92,\r\n        99,\r\n        80\r\n      ],\r\n      [\r\n        310,\r\n        280,\r\n        295,\r\n        260,\r\n        340\r\n      ],\r\n      [\r\n        130,\r\n        118,\r\n        140,\r\n        125,\r\n        150\r\n      ]\r\n    ],\r\n    "verifier_scores": [\r\n      [\r\n        0.84,\r\n        0.22,\r\n        0.92,\r\n        0.33,\r\n        0.76\r\n      ],\r\n      [\r\n        0.55,\r\n        0.9,\r\n        0.18,\r\n        0.62,\r\n        0.31\r\n      ],\r\n      [\r\n        0.7,\r\n        0.8,\r\n        0.65,\r\n        0.58,\r\n        0.25\r\n      ],\r\n      [\r\n        0.12,\r\n        0.2,\r\n        0.18,\r\n        0.24,\r\n        0.88\r\n      ],\r\n      [\r\n        0.48,\r\n        0.32,\r\n        0.29,\r\n        0.17,\r\n        0.51\r\n      ]\r\n    ]\r\n  },\r\n  "cost": {\r\n    "base_tokens": 512,\r\n    "costs": {\r\n      "high": {\r\n        "concurrency_upper_bound": 33,\r\n        "dense_decode_flops_floor": 256960000000000.0,\r\n        "kv_bytes": 2097152000,\r\n        "kv_gib": 1.953125,\r\n        "latency_seconds": 76.4087870105062,\r\n        "tokens": 16000\r\n      },\r\n      "low": {\r\n        "concurrency_upper_bound": 532,\r\n        "dense_decode_flops_floor": 16060000000000.0,\r\n        "kv_bytes": 131072000,\r\n        "kv_gib": 0.1220703125,\r\n        "latency_seconds": 4.775549188156638,\r\n        "tokens": 1000\r\n      },\r\n      "medium": {\r\n        "concurrency_upper_bound": 133,\r\n        "dense_decode_flops_floor": 64240000000000.0,\r\n        "kv_bytes": 524288000,\r\n        "kv_gib": 0.48828125,\r\n        "latency_seconds": 19.10219675262655,\r\n        "tokens": 4000\r\n      },\r\n      "xhigh": {\r\n        "concurrency_upper_bound": 8,\r\n        "dense_decode_flops_floor": 1027840000000000.0,\r\n        "kv_bytes": 8388608000,\r\n        "kv_gib": 7.8125,\r\n        "latency_seconds": 305.6351480420248,\r\n        "tokens": 64000\r\n      }\r\n    },\r\n    "forcing_extra_seconds": 28.653295128939828,\r\n    "ladder": {\r\n      "high": 16000,\r\n      "low": 1000,\r\n      "medium": 4000,\r\n      "xhigh": 64000\r\n    },\r\n    "reserved_concurrency_upper_bound": {\r\n      "high": 29,\r\n      "low": 467,\r\n      "medium": 116,\r\n      "xhigh": 7\r\n    },\r\n    "reserved_scenario": {\r\n      "bytes_per_element": 2,\r\n      "decode_tokens_per_second": 209.4,\r\n      "device_budget_gib": 80.0,\r\n      "kv_width": 1024,\r\n      "n_layers": 32,\r\n      "n_parameters": 8030000000.0,\r\n      "runtime_reserve_gib": 8.0,\r\n      "weight_bytes_per_parameter": 2\r\n    },\r\n    "scenario": {\r\n      "bytes_per_element": 2,\r\n      "decode_tokens_per_second": 209.4,\r\n      "device_budget_gib": 80.0,\r\n      "kv_width": 1024,\r\n      "n_layers": 32,\r\n      "n_parameters": 8030000000.0,\r\n      "runtime_reserve_gib": 0.0,\r\n      "weight_bytes_per_parameter": 2\r\n    },\r\n    "short_512": {\r\n      "concurrency_upper_bound": 1040,\r\n      "dense_decode_flops_floor": 8222720000000.0,\r\n      "kv_bytes": 67108864,\r\n      "kv_gib": 0.0625,\r\n      "latency_seconds": 2.445081184336199,\r\n      "tokens": 512\r\n    }\r\n  },\r\n  "overthinking_toy": {\r\n    "budgets": [\r\n      256,\r\n      1024,\r\n      4096,\r\n      16384\r\n    ],\r\n    "hard_accuracy": [\r\n      0.18,\r\n      0.31,\r\n      0.54,\r\n      0.68\r\n    ],\r\n    "note": "synthetic mechanism illustration, not a benchmark",\r\n    "simple_accuracy": [\r\n      0.93,\r\n      0.95,\r\n      0.94,\r\n      0.9\r\n    ]\r\n  },\r\n  "schema_version": 2,\r\n  "verifier": {\r\n    "noisy": {\r\n      "fallback": "random_candidate",\r\n      "false_positive_rate": 0.05,\r\n      "selection_accuracy": {\r\n        "1": 0.4,\r\n        "17": 0.9228846834475187,\r\n        "3": 0.7284400000000001,\r\n        "5": 0.8506525240000002,\r\n        "9": 0.9130491559355268\r\n      },\r\n      "true_positive_rate": 0.9\r\n    },\r\n    "perfect_coverage": {\r\n      "1": 0.4,\r\n      "17": 0.9998307334055526,\r\n      "3": 0.784,\r\n      "5": 0.92224,\r\n      "9": 0.989922304\r\n    }\r\n  },\r\n  "vote": {\r\n    "n_distractors": 6,\r\n    "p_correct": 0.4,\r\n    "sample_counts": [\r\n      1,\r\n      3,\r\n      5,\r\n      9,\r\n      17\r\n    ],\r\n    "statistics": {\r\n      "1": {\r\n        "any_mode_tie": 0.0,\r\n        "correct_in_top_tie": 0.0,\r\n        "random_tie_success": 0.4,\r\n        "unique_mode_success": 0.4\r\n      },\r\n      "17": {\r\n        "any_mode_tie": 0.07955325424125981,\r\n        "correct_in_top_tie": 0.07023020038594624,\r\n        "random_tie_success": 0.8893158767850492,\r\n        "unique_mode_success": 0.8566882832072792\r\n      },\r\n      "3": {\r\n        "any_mode_tie": 0.4800000000000002,\r\n        "correct_in_top_tie": 0.36000000000000004,\r\n        "random_tie_success": 0.4720000000000002,\r\n        "unique_mode_success": 0.35200000000000026\r\n      },\r\n      "5": {\r\n        "any_mode_tie": 0.25920000000000004,\r\n        "correct_in_top_tie": 0.21599999999999994,\r\n        "random_tie_success": 0.5958400000000001,\r\n        "unique_mode_success": 0.50944\r\n      },\r\n      "9": {\r\n        "any_mode_tie": 0.18846911999999855,\r\n        "correct_in_top_tie": 0.15789311999999908,\r\n        "random_tie_success": 0.737034879999996,\r\n        "unique_mode_success": 0.6710579199999984\r\n      }\r\n    }\r\n  }\r\n}\r\n'
REFERENCE_SHA256 = 'fa7119bbc6405d404b3e11270115f48b18ad9be058a9bf1a91862b439cb58a09'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == REFERENCE_SHA256
REF = json.loads(REFERENCE_JSON)
assert REF["schema_version"] == 2
print(MSG["reference_ok"], REFERENCE_SHA256[:12])

FINAL_INTEGER_RE = re.compile(r"^\s*FINAL:\s*([+-]?\d+)\s*$")

def verify_final_integer(candidate: str, expected: int) -> bool:
    """Verify an exact integer under the strict ``FINAL: <integer>`` protocol.

    The function never evaluates code and never searches arbitrary rationale text.
    A format mismatch is a normal failed verification rather than a parser crash.
    """

    if not isinstance(candidate, str):
        raise TypeError("candidate must be a string")
    if not isinstance(expected, (int, np.integer)) or isinstance(expected, bool):
        raise TypeError("expected must be an integer")
    match = FINAL_INTEGER_RE.fullmatch(candidate)
    if match is None:
        return False
    try:
        return int(match.group(1)) == int(expected)
    except ValueError:
        return False  # also reject integers beyond the interpreter parsing limit

cases = {
    "FINAL: 42": True,
    "  FINAL: -7  ": False,
    "42": False,
    "reasoning... FINAL: 42": False,
    "FINAL: 42 because it is obvious": False,
    "FINAL: 0042": True,
}
observed = {text: verify_final_integer(text, 42) for text in cases}
assert observed == cases
assert verify_final_integer("FINAL: -7", -7)
print(MSG["checker_ok"])
observed

def compositions(total: int, parts: int) -> Iterator[tuple[int, ...]]:
    """Yield weak compositions of ``total`` into ``parts`` counts."""

    if isinstance(total, bool) or not isinstance(total, int) or total < 0:
        raise ValueError("total must be a non-negative integer")
    if isinstance(parts, bool) or not isinstance(parts, int) or parts <= 0:
        raise ValueError("parts must be a positive integer")
    if parts == 1:
        yield (total,)
        return
    for first in range(total + 1):
        for tail in compositions(total - first, parts - 1):
            yield (first, *tail)

def multinomial_mass(counts: Sequence[int], probabilities: Sequence[float]) -> float:
    """Return one multinomial probability using log-gamma arithmetic."""

    if len(counts) != len(probabilities) or len(counts) == 0:
        raise ValueError("counts and probabilities must have matching lengths")
    probs = np.asarray(probabilities, dtype=np.float64)
    if probs.ndim != 1 or np.any(~np.isfinite(probs)) or np.any(probs < 0.0):
        raise ValueError("probabilities must be finite and non-negative")
    if not math.isclose(float(probs.sum()), 1.0, rel_tol=0.0, abs_tol=1e-12):
        raise ValueError("probabilities must sum to one")
    if any(isinstance(c, bool) or (not isinstance(c, int)) or c < 0 for c in counts):
        raise ValueError("counts must be non-negative integers")
    log_p = math.lgamma(sum(counts) + 1)
    for count, probability in zip(counts, probs):
        log_p -= math.lgamma(count + 1)
        if count:
            if probability == 0.0:
                return 0.0
            log_p += count * math.log(float(probability))
    return math.exp(log_p)

def exact_vote_statistics(n_samples: int, p_correct: float, n_distractors: int) -> dict[str, float]:
    """Return exact unique-mode, random-tie, and tie probabilities."""

    if any(isinstance(v, bool) or not isinstance(v, int) or v <= 0 for v in (n_samples, n_distractors)):
        raise ValueError("sample and distractor counts must be positive")
    if not 0.0 <= p_correct <= 1.0 or not math.isfinite(p_correct):
        raise ValueError("p_correct must lie in [0, 1]")
    probs = [p_correct] + [(1.0 - p_correct) / n_distractors] * n_distractors
    result = dict(unique_mode_success=0.0, random_tie_success=0.0,
                  any_mode_tie=0.0, correct_in_top_tie=0.0)
    total_mass = 0.0
    for counts in compositions(n_samples, n_distractors + 1):
        mass = multinomial_mass(counts, probs)
        total_mass += mass
        top = max(counts)
        modes = [i for i, count in enumerate(counts) if count == top]
        if len(modes) > 1:
            result["any_mode_tie"] += mass
        if counts[0] > max(counts[1:]):
            result["unique_mode_success"] += mass
        if 0 in modes:
            result["random_tie_success"] += mass / len(modes)
            if len(modes) > 1:
                result["correct_in_top_tie"] += mass
    assert math.isclose(total_mass, 1.0, abs_tol=1e-10)
    return result

def simulate_vote(n_samples: int, p_correct: float, n_distractors: int,
                  trials: int = 100_000, seed: int = 7) -> dict[str, float]:
    """Monte-Carlo check of unique-mode success and non-unique-mode frequency."""

    if any(isinstance(v, bool) or not isinstance(v, int) or v <= 0 for v in (trials, n_samples, n_distractors)):
        raise ValueError("trials must be positive")
    probs = np.array([p_correct] + [(1.0 - p_correct) / n_distractors] * n_distractors)
    rng = np.random.default_rng(seed)
    counts = rng.multinomial(n_samples, probs, size=trials)
    top = counts.max(axis=1)
    unique_correct = counts[:, 0] > counts[:, 1:].max(axis=1)
    any_tie = (counts == top[:, None]).sum(axis=1) > 1
    return {"unique_mode_success": float(unique_correct.mean()),
            "any_mode_tie": float(any_tie.mean())}

vote_cfg = REF["vote"]
Ns = vote_cfg["sample_counts"]
exact = {n: exact_vote_statistics(n, vote_cfg["p_correct"], vote_cfg["n_distractors"])
         for n in Ns}
mc = {n: simulate_vote(n, vote_cfg["p_correct"], vote_cfg["n_distractors"], seed=7+n)
      for n in Ns}
for n in Ns:
    frozen = vote_cfg["statistics"][str(n)]
    for key, value in exact[n].items():
        assert math.isclose(value, frozen[key], rel_tol=0.0, abs_tol=1e-12)
    assert abs(mc[n]["unique_mode_success"] - exact[n]["unique_mode_success"]) < 0.006
    assert abs(mc[n]["any_mode_tie"] - exact[n]["any_mode_tie"]) < 0.006
    print(f"N={n:>2}: unique={exact[n]['unique_mode_success']:.3f} · "
          f"random-tie={exact[n]['random_tie_success']:.3f} · "
          f"any tie={exact[n]['any_mode_tie']:.3f}")

assert exact[3]["unique_mode_success"] < exact[1]["unique_mode_success"]
assert exact[3]["random_tie_success"] > exact[1]["random_tie_success"]

fig, ax = plt.subplots(figsize=(11, 5.5), layout="constrained")
ax.plot(Ns, [exact[n]["unique_mode_success"] for n in Ns], "-o", label="tie → fail")
ax.plot(Ns, [exact[n]["random_tie_success"] for n in Ns], "-s", label="random among modes")
ax.axhline(vote_cfg["p_correct"], linestyle=":", linewidth=1, label="one-shot p")
ax.set(xlabel="N", ylabel="P(correct selection)", ylim=(0.3, 0.95))
ax.grid(alpha=0.3); ax.legend(); plt.show()
print(MSG["vote_ok"])

def perfect_verifier_coverage(p_correct: float, n_samples: int) -> float:
    """Return oracle coverage: the probability that the pool contains a correct candidate."""

    if not 0.0 <= p_correct <= 1.0 or not math.isfinite(p_correct):
        raise ValueError("p_correct must lie in [0, 1]")
    if isinstance(n_samples, bool) or not isinstance(n_samples, int) or n_samples <= 0:
        raise ValueError("n_samples must be positive")
    return 1.0 if p_correct == 1 else -math.expm1(n_samples * math.log1p(-p_correct))

def noisy_verifier_success(p_correct: float, n_samples: int, tpr: float, fpr: float,
                           fallback: str = "random_candidate") -> float:
    """Compute exact accuracy for uniform selection among accepted candidates.

    The verifier accepts each correct candidate with ``tpr`` and each incorrect
    candidate with ``fpr``. If none are accepted, the default fallback uniformly
    selects one of the generated candidates.
    """

    for name, value in {"p_correct": p_correct, "tpr": tpr, "fpr": fpr}.items():
        if not 0.0 <= value <= 1.0 or not math.isfinite(value):
            raise ValueError(f"{name} must lie in [0, 1]")
    if isinstance(n_samples, bool) or not isinstance(n_samples, int) or n_samples <= 0:
        raise ValueError("n_samples must be positive")
    if fallback not in {"random_candidate", "fail"}:
        raise ValueError("unsupported fallback")

    success = 0.0
    for n_correct in range(n_samples + 1):
        p_k = math.comb(n_samples, n_correct) * p_correct**n_correct * (1-p_correct)**(n_samples-n_correct)
        for accepted_correct in range(n_correct + 1):
            p_ac = math.comb(n_correct, accepted_correct) * tpr**accepted_correct * (1-tpr)**(n_correct-accepted_correct)
            n_wrong = n_samples - n_correct
            for accepted_wrong in range(n_wrong + 1):
                p_aw = math.comb(n_wrong, accepted_wrong) * fpr**accepted_wrong * (1-fpr)**(n_wrong-accepted_wrong)
                total = accepted_correct + accepted_wrong
                if total:
                    conditional = accepted_correct / total
                elif fallback == "random_candidate":
                    conditional = n_correct / n_samples
                else:
                    conditional = 0.0
                success += p_k * p_ac * p_aw * conditional
    return success

p = REF["vote"]["p_correct"]
perfect = {n: perfect_verifier_coverage(p, n) for n in Ns}
strong = {n: noisy_verifier_success(p, n, 0.90, 0.05) for n in Ns}
weak = {n: noisy_verifier_success(p, n, 0.75, 0.25) for n in Ns}
for n in Ns:
    assert math.isclose(perfect[n], REF["verifier"]["perfect_coverage"][str(n)], abs_tol=1e-12)
    assert math.isclose(strong[n], REF["verifier"]["noisy"]["selection_accuracy"][str(n)], abs_tol=1e-12)
    assert strong[n] <= perfect[n] + 1e-12
    print(f"N={n:>2}: oracle={perfect[n]:.3f} · strong={strong[n]:.3f} · weak={weak[n]:.3f} · mode={exact[n]['unique_mode_success']:.3f}")

assert weak[17] < exact[17]["unique_mode_success"]  # verifier quality can reverse the ranking
assert perfect[3] - exact[3]["unique_mode_success"] > 0.43

fig, ax = plt.subplots(figsize=(11, 5.5), layout="constrained")
ax.plot(Ns, [perfect[n] for n in Ns], "-o", label="oracle coverage")
ax.plot(Ns, [strong[n] for n in Ns], "-s", label="noisy verifier: TPR=.90 FPR=.05")
ax.plot(Ns, [weak[n] for n in Ns], "-^", label="weak verifier: TPR=.75 FPR=.25")
ax.plot(Ns, [exact[n]["unique_mode_success"] for n in Ns], "--", label="unique-mode vote")
ax.set(xlabel="N", ylabel="P(correct selection)", ylim=(0.3, 1.03))
ax.grid(alpha=0.3); ax.legend(fontsize=11); plt.show()
print(MSG["verifier_ok"])

def integer_array(value):
    raw = np.asarray(value)
    if raw.dtype.kind not in "iu":
        raise ValueError("IDs and token counts must be integer arrays")
    if raw.size and (np.any(raw > np.iinfo(np.int64).max)):
        raise ValueError("integer array exceeds int64")
    return raw.astype(np.int64)

@dataclass(frozen=True)
class CandidateBatch:
    """A validated Best-of-N batch with aligned ``[batch, candidates]`` tensors."""

    answer_ids: np.ndarray
    correct_answer_ids: np.ndarray
    verifier_scores: np.ndarray
    token_counts: np.ndarray

    def __post_init__(self) -> None:
        """Normalize arrays and enforce the shared ``[batch, candidates]`` contract."""

        answer_ids = integer_array(self.answer_ids)
        correct = integer_array(self.correct_answer_ids)
        scores = np.asarray(self.verifier_scores, dtype=np.float64)
        token_counts = integer_array(self.token_counts)
        if answer_ids.ndim != 2 or 0 in answer_ids.shape:
            raise ValueError("answer_ids must have non-empty [batch, candidates] shape")
        if scores.shape != answer_ids.shape or token_counts.shape != answer_ids.shape:
            raise ValueError("scores and token_counts must match answer_ids")
        if correct.shape != (answer_ids.shape[0],):
            raise ValueError("correct_answer_ids must have [batch] shape")
        if np.any(answer_ids < 0) or np.any(correct < 0):
            raise ValueError("answer IDs must be non-negative")
        if np.any(~np.isfinite(scores)):
            raise ValueError("verifier scores must be finite")
        if np.any(token_counts <= 0):
            raise ValueError("token counts must be positive")
        object.__setattr__(self, "answer_ids", answer_ids)
        object.__setattr__(self, "correct_answer_ids", correct)
        object.__setattr__(self, "verifier_scores", scores)
        object.__setattr__(self, "token_counts", token_counts)

    @property
    def is_correct(self) -> np.ndarray:
        """Return the boolean correctness matrix with shape ``[batch, candidates]``."""

        return self.answer_ids == self.correct_answer_ids[:, None]

def select_unique_mode(answer_ids: np.ndarray, tie_policy: str = "fail") -> np.ndarray:
    """Select the per-row unique mode; return ``-1`` on ties under ``fail``."""

    answers = integer_array(answer_ids)
    if answers.ndim != 2 or 0 in answers.shape:
        raise ValueError("answer_ids must be a non-empty matrix")
    if np.any(answers < 0):
        raise ValueError("answer IDs must be non-negative; -1 is reserved for abstention")
    if tie_policy not in {"fail", "lowest_id"}:
        raise ValueError("tie_policy must be 'fail' or 'lowest_id'")
    selected = np.full(answers.shape[0], -1, dtype=np.int64)
    for row, values in enumerate(answers):
        unique, counts = np.unique(values, return_counts=True)
        modes = unique[counts == counts.max()]
        if len(modes) == 1:
            selected[row] = modes[0]
        elif tie_policy == "lowest_id":
            selected[row] = modes.min()
    return selected

def select_by_score(scores: np.ndarray) -> np.ndarray:
    """Return one candidate index per prompt using deterministic first-argmax."""

    scores = np.asarray(scores, dtype=np.float64)
    if scores.ndim != 2 or 0 in scores.shape or np.any(~np.isfinite(scores)):
        raise ValueError("scores must be a finite non-empty matrix")
    return scores.argmax(axis=1)

def selection_accuracy(selected_answer_ids: np.ndarray, correct_answer_ids: np.ndarray) -> float:
    """Return answer-level accuracy; the sentinel ``-1`` counts as failure."""

    selected = integer_array(selected_answer_ids)
    correct = integer_array(correct_answer_ids)
    if selected.shape != correct.shape or selected.ndim != 1 or selected.size == 0 or np.any(correct < 0) or np.any(selected < -1):
        raise ValueError("selected and correct IDs must be matching vectors")
    return float((selected == correct).mean())

raw = REF["candidate_batch"]
batch = CandidateBatch(raw["answer_ids"], raw["correct_answer_ids"],
                       raw["verifier_scores"], raw["token_counts"])
mode_answers = select_unique_mode(batch.answer_ids, tie_policy="fail")
top_indices = select_by_score(batch.verifier_scores)
rows = np.arange(batch.answer_ids.shape[0])
top_answers = batch.answer_ids[rows, top_indices]

mode_accuracy = selection_accuracy(mode_answers, batch.correct_answer_ids)
verifier_accuracy = selection_accuracy(top_answers, batch.correct_answer_ids)
oracle_coverage = float(batch.is_correct.any(axis=1).mean())
generated_per_prompt = batch.token_counts.sum(axis=1)
selected_lengths = batch.token_counts[rows, top_indices]

assert batch.answer_ids.shape == (5, 5)
assert oracle_coverage == 0.8
assert mode_accuracy == 0.4 and verifier_accuracy == 0.4
assert np.all(generated_per_prompt >= selected_lengths)
assert np.any(batch.verifier_scores[~batch.is_correct] > 0.85)  # deliberate false positive
print("mode answers:", mode_answers)
print("top-score indices:", top_indices, "top answers:", top_answers)
print(f"mode acc={mode_accuracy:.2f} · verifier acc={verifier_accuracy:.2f} · oracle coverage={oracle_coverage:.2f}")
print(f"mean generated tokens/prompt={generated_per_prompt.mean():.1f}; selected-only mean={selected_lengths.mean():.1f}")
length_penalty = 0.001  # score units per generated token; a declared policy
penalized = batch.verifier_scores - length_penalty * batch.token_counts
penalized_indices = select_by_score(penalized)
penalized_accuracy = selection_accuracy(batch.answer_ids[rows, penalized_indices],
                                         batch.correct_answer_ids)
print("length-penalized indices:", penalized_indices, "accuracy:", penalized_accuracy)
assert batch.token_counts.sum() == generated_per_prompt.sum()
# Ranking already-generated answers by length does not refund their generation cost.
print(MSG["batch_ok"])

def make_synthetic_workload(n_prompts: int, max_candidates: int, seed: int,
                            difficulty_weights: Sequence[float],
                            p_correct_by_difficulty: Sequence[float]) -> tuple[np.ndarray, np.ndarray]:
    """Create labelled correctness and overlapping verifier-score matrices."""

    weights = np.asarray(difficulty_weights, dtype=np.float64)
    probabilities = np.asarray(p_correct_by_difficulty, dtype=np.float64)
    if any(isinstance(v, bool) or not isinstance(v, int) or v <= 0 for v in (n_prompts, max_candidates)):
        raise ValueError("workload dimensions must be positive")
    if weights.ndim != 1 or probabilities.shape != weights.shape or weights.size == 0:
        raise ValueError("difficulty vectors must have matching non-zero shapes")
    if np.any(weights < 0) or np.any(~np.isfinite(weights)) or not math.isclose(float(weights.sum()), 1.0, rel_tol=0.0, abs_tol=1e-12):
        raise ValueError("difficulty weights must be finite, non-negative, and sum to one")
    if np.any((probabilities < 0) | (probabilities > 1)) or np.any(~np.isfinite(probabilities)):
        raise ValueError("correctness probabilities must lie in [0, 1]")
    rng = np.random.default_rng(seed)
    difficulty = rng.choice(weights.size, size=n_prompts, p=weights)
    is_correct = rng.random((n_prompts, max_candidates)) < probabilities[difficulty, None]
    u = rng.random((n_prompts, max_candidates))
    scores = np.where(is_correct, 0.52 + 0.48*u, 0.03 + 0.90*u**2)
    return is_correct, scores

def evaluate_adaptive_threshold(is_correct: np.ndarray, scores: np.ndarray, threshold: float) -> dict[str, float]:
    """Stop at the first threshold crossing, else use the best score after max-N."""

    labels = np.asarray(is_correct)
    scores = np.asarray(scores, dtype=np.float64)
    if labels.shape != scores.shape or labels.ndim != 2 or 0 in labels.shape or labels.dtype != np.bool_:
        raise ValueError("labels and scores must be matching non-empty matrices")
    if np.any(~np.isfinite(scores)) or not 0.0 <= threshold <= 1.0:
        raise ValueError("scores must be finite and threshold must lie in [0, 1]")
    crosses = scores >= threshold
    has_crossing = crosses.any(axis=1)
    first_crossing = crosses.argmax(axis=1)
    best_index = scores.argmax(axis=1)
    selected_index = np.where(has_crossing, first_crossing, best_index)
    attempts = np.where(has_crossing, first_crossing + 1, labels.shape[1])
    rows = np.arange(labels.shape[0])
    return {
        "accuracy": float(labels[rows, selected_index].mean()),
        "mean_candidates": float(attempts.mean()),
        "crossing_rate": float(has_crossing.mean()),
    }

params = REF["adaptive_search"]["parameters"]
labels, scores = make_synthetic_workload(
    params["n_prompts"], params["max_candidates"], params["seed"],
    params["difficulty_weights"], params["p_correct_by_difficulty"])
rows = np.arange(labels.shape[0])
one_shot = float(labels[:, 0].mean())
fixed_index = scores.argmax(axis=1)
fixed_accuracy = float(labels[rows, fixed_index].mean())
oracle = float(labels.any(axis=1).mean())

thresholds = np.array([0.60, 0.68, 0.74, 0.78, 0.82, 0.86, 0.90, 0.94])
curve = [evaluate_adaptive_threshold(labels, scores, float(threshold)) for threshold in thresholds]
chosen = evaluate_adaptive_threshold(labels, scores, params["threshold"])
frozen = REF["adaptive_search"]["metrics"]
assert abs(one_shot - frozen["one_shot_accuracy"]) < 1e-12
assert abs(chosen["accuracy"] - frozen["adaptive_accuracy"]) < 1e-12
assert abs(chosen["mean_candidates"] - frozen["adaptive_mean_candidates"]) < 1e-12
assert one_shot < chosen["accuracy"] <= fixed_accuracy <= oracle  # observed here, not a general ordering
assert chosen["mean_candidates"] < params["max_candidates"]

for threshold, metrics in zip(thresholds, curve):
    print(f"threshold={threshold:.2f}: acc={metrics['accuracy']:.3f} · mean N={metrics['mean_candidates']:.2f}")
print(f"one-shot={one_shot:.3f} · fixed Best-of-16={fixed_accuracy:.3f} · oracle={oracle:.3f}")

fig, ax = plt.subplots(figsize=(11, 5.5), layout="constrained")
xs = [item["mean_candidates"] for item in curve]
ys = [item["accuracy"] for item in curve]
ax.plot(xs, ys, "-o")
for threshold, x, y in zip(thresholds, xs, ys):
    ax.annotate(f"{threshold:.2f}", (x, y), xytext=(4, 4), textcoords="offset points", fontsize=11)
ax.axhline(fixed_accuracy, linestyle="--", linewidth=1, label="fixed Best-of-16")
ax.axhline(one_shot, linestyle=":", linewidth=1, label="one-shot")
ax.set(xlabel="mean inspected candidates", ylabel="selection accuracy")
ax.grid(alpha=0.3); ax.legend(); plt.show()
print(MSG["adaptive_ok"])

@dataclass(frozen=True)
class InferenceCostConfig:
    """Transparent teaching assumptions for dense decode and KV accounting."""

    n_parameters: float = 8.03e9
    decode_tokens_per_second: float = 209.4
    n_layers: int = 32
    kv_width: int = 1024
    bytes_per_element: int = 2
    device_budget_gib: float = 80.0
    weight_bytes_per_parameter: int = 2
    runtime_reserve_gib: float = 0.0

    def validate(self) -> None:
        """Reject invalid assumptions before arithmetic hides them."""

        for name, value in {"n_parameters": self.n_parameters,
                            "decode_tokens_per_second": self.decode_tokens_per_second,
                            "device_budget_gib": self.device_budget_gib}.items():
            if not math.isfinite(value) or value <= 0:
                raise ValueError(f"{name} must be finite and positive")
        for name, value in {"n_layers": self.n_layers, "kv_width": self.kv_width,
                            "bytes_per_element": self.bytes_per_element,
                            "weight_bytes_per_parameter": self.weight_bytes_per_parameter}.items():
            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if not math.isfinite(self.runtime_reserve_gib) or self.runtime_reserve_gib < 0:
            raise ValueError("runtime reserve must be finite and non-negative")
        if self.available_kv_gib <= 0:
            raise ValueError("weights and reserve leave no KV budget")

    @property
    def weight_gib(self) -> float:
        """Return model-weight memory in GiB under the stated byte assumption."""

        return self.n_parameters * self.weight_bytes_per_parameter / 2**30

    @property
    def available_kv_gib(self) -> float:
        """Return the teaching budget left for KV after weights and reserve."""

        return self.device_budget_gib - self.weight_gib - self.runtime_reserve_gib

    @property
    def kv_bytes_per_token(self) -> int:
        """Return K+V bytes per generated token across all layers."""

        return 2 * self.bytes_per_element * self.kv_width * self.n_layers

def estimate_reasoning_cost(tokens: int, cfg: InferenceCostConfig) -> dict[str, float | int]:
    """Estimate fixed-rate latency, weight-matmul work, generated KV and concurrency."""

    if isinstance(tokens, bool) or not isinstance(tokens, int) or tokens <= 0:
        raise ValueError("tokens must be a positive integer")
    cfg.validate()
    kv_bytes = cfg.kv_bytes_per_token * tokens
    kv_gib = kv_bytes / 2**30
    return {
        "latency_s": tokens / cfg.decode_tokens_per_second,
        "dense_decode_flops_floor": 2.0 * cfg.n_parameters * tokens,
        "kv_bytes": kv_bytes,
        "kv_gib": kv_gib,
        "concurrency_upper_bound": int(cfg.available_kv_gib // kv_gib),
    }

ideal = InferenceCostConfig(runtime_reserve_gib=0.0)
reserved = InferenceCostConfig(runtime_reserve_gib=8.0)
assert ideal.kv_bytes_per_token / 1024 == 128  # KiB/token
print(f"weights={ideal.weight_gib:.2f} GiB · KV={ideal.kv_bytes_per_token/1024:.0f} KiB/token")
ladder = REF["cost"]["ladder"]
ladder_order = ("low", "medium", "high", "xhigh")
for name in ladder_order:
    tokens = ladder[name]
    a = estimate_reasoning_cost(tokens, ideal)
    b = estimate_reasoning_cost(tokens, reserved)
    frozen = REF["cost"]["costs"][name]
    assert math.isclose(a["latency_s"], frozen["latency_seconds"], abs_tol=1e-12)
    assert math.isclose(a["kv_gib"], frozen["kv_gib"], abs_tol=1e-12)
    assert a["concurrency_upper_bound"] == frozen["concurrency_upper_bound"]
    print(f"{name:>6}: {tokens:>6} tok · {a['latency_s']:6.1f} s · "
          f"{a['kv_gib']:5.2f} GiB KV · conc≤{a['concurrency_upper_bound']:>3} "
          f"(with 8-GiB reserve: ≤{b['concurrency_upper_bound']:>3})")

short = estimate_reasoning_cost(512, ideal)
assert short["concurrency_upper_bound"] == REF["cost"]["short_512"]["concurrency_upper_bound"] == 1040
forcing = (8_000 - 2_000) / ideal.decode_tokens_per_second
assert math.isclose(forcing, REF["cost"]["forcing_extra_seconds"], abs_tol=1e-12)
print(f"2K → 8K budget forcing adds {forcing:.1f} s under this scenario")
print(MSG["cost_ok"])

mix = REF["adaptive_mix"]
weights = np.asarray(mix["weights"], dtype=np.float64)
budgets = np.asarray(mix["budgets"], dtype=np.float64)
average_budget = float(weights @ budgets)
saving_factor = mix["uniform_budget"] / average_budget
assert math.isclose(weights.sum(), 1.0, rel_tol=0.0, abs_tol=1e-12)
assert average_budget == mix["average_budget"] == 4350.0
assert math.isclose(saving_factor, mix["saving_factor"], abs_tol=1e-12)
print(f"adaptive average={average_budget:,.0f} tokens · uniform={mix['uniform_budget']:,} · saving ×{saving_factor:.2f}")

print("hard-share sensitivity; simple/medium retain a 7:2 ratio:")
for hard_share in (0.05, 0.10, 0.25, 0.50, 0.75, 1.00):
    simple_share = (1-hard_share) * 7/9
    medium_share = (1-hard_share) * 2/9
    average = simple_share*500 + medium_share*4000 + hard_share*32000
    print(f"  hard={hard_share:>5.0%}: mean={average:>8,.0f} · saving ×{32000/average:.2f}")

toy = REF["overthinking_toy"]
toy_budgets = np.asarray(toy["budgets"])
simple = np.asarray(toy["simple_accuracy"])
hard = np.asarray(toy["hard_accuracy"])
assert simple.argmax() == 1 and simple[-1] < simple.max()
assert np.all(np.diff(hard) > 0)
fig, ax = plt.subplots(figsize=(11, 5.5), layout="constrained")
ax.plot(toy_budgets, simple, "-o", label="simple — synthetic")
ax.plot(toy_budgets, hard, "-s", label="hard — synthetic")
ax.set_xscale("log", base=2)
ax.set(xlabel="thinking budget, tokens (log2)", ylabel="toy accuracy", ylim=(0.1, 1.0))
ax.grid(alpha=0.3); ax.legend();
ax.text(0.01, 0.02, toy["note"], transform=ax.transAxes, fontsize=11)
plt.show()
print(MSG["mix_ok"])

def expect_error(error_type: type[BaseException], fn, *args, **kwargs) -> str:
    """Run a negative test and return the expected exception message."""

    try:
        fn(*args, **kwargs)
    except error_type as exc:
        return str(exc)
    raise AssertionError(f"expected {error_type.__name__}")

errors = [
    expect_error(ValueError, exact_vote_statistics, 3, 1.2, 6),
    expect_error(ValueError, noisy_verifier_success, 0.4, 0, 0.9, 0.1),
    expect_error(ValueError, select_by_score, np.empty((2, 0))),
    expect_error(ValueError, CandidateBatch,
                 np.array([[1, 2]]), np.array([1]), np.array([[0.1, np.nan]]), np.array([[4, 5]])),
    expect_error(ValueError, estimate_reasoning_cost, 1000,
                 InferenceCostConfig(device_budget_gib=8.0)),
]
assert len(errors) == 5 and all(errors)
for message in errors:
    print("expected:", message)
print(MSG["failure_ok"])

import types, sys
core_module=types.ModuleType('m15_core')
sys.modules['m15_core']=core_module
exec('"""Coverage, noisy selection and scoped inference accounting."""\nimport math\nfrom dataclasses import dataclass\n\n\ndef positive_count(value, name):\n    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:\n        raise ValueError(name + " must be a positive integer")\n\n\ndef probability(value, name):\n    if not math.isfinite(value) or not 0 <= value <= 1:\n        raise ValueError(name + " must be finite and lie in [0, 1]")\n\n\ndef oracle_coverage(p_correct, n):\n    """At least one success in n independent, identically distributed attempts."""\n    positive_count(n, "n")\n    probability(p_correct, "p_correct")\n    if p_correct == 1:\n        return 1.0\n    return -math.expm1(n * math.log1p(-p_correct))\n\n\ndef noisy_selection(p, n, tpr, fpr):\n    """Uniform accepted candidate; uniform pool fallback if none are accepted."""\n    positive_count(n, "n")\n    for name, value in (("p", p), ("tpr", tpr), ("fpr", fpr)):\n        probability(value, name)\n    acceptance = p*tpr + (1-p)*fpr\n    if acceptance == 0 or n == 1:\n        return p\n    precision = p*tpr/acceptance\n    if acceptance == 1:\n        return precision\n    log_survival = (n-1)*math.log1p(-acceptance)\n    # A non-negative mixture avoids cancellation for tiny probabilities.\n    return p*math.exp(log_survival) + precision*(-math.expm1(log_survival))\n\n\n@dataclass(frozen=True)\nclass ReasoningCost:\n    parameters: float = 8.03e9\n    tokens_per_second: float = 209.4\n    layers: int = 32\n    kv_width: int = 1024\n    bytes_per_element: int = 2\n\n    def __post_init__(self):\n        for name in ("parameters", "tokens_per_second"):\n            value = getattr(self, name)\n            if not math.isfinite(value) or value <= 0:\n                raise ValueError(name + " must be finite and positive")\n        for name in ("layers", "kv_width", "bytes_per_element"):\n            positive_count(getattr(self, name), name)\n\n    @property\n    def kv_bytes_per_token(self):\n        return 2*self.layers*self.kv_width*self.bytes_per_element\n\n    def estimate(self, tokens, prompt_tokens=0):\n        positive_count(tokens, "tokens")\n        if isinstance(prompt_tokens, bool) or not isinstance(prompt_tokens, int) or prompt_tokens < 0:\n            raise ValueError("prompt_tokens must be a non-negative integer")\n        return dict(decode_seconds=tokens/self.tokens_per_second,\n                    parameter_matmul_flops=2*self.parameters*tokens,\n                    generated_kv_gib=self.kv_bytes_per_token*tokens/2**30,\n                    total_kv_gib=self.kv_bytes_per_token*(tokens+prompt_tokens)/2**30)\n\n\nscenario = ReasoningCost()\nassert abs(noisy_selection(.4, 17, .9, .05)-.923)<.001\nassert oracle_coverage(1e-18, 1) > 0  # direct 1-(1-p)**n loses this value\nassert noisy_selection(1e-18, 1, 1, 0) == 1e-18\nassert math.isclose(noisy_selection(1e-18, 2, 1, 0), 2e-18, rel_tol=1e-15)\n', core_module.__dict__)

for p_check in (0., .4, 1.):
    for tpr_check, fpr_check in ((0,0), (1,1), (.9,.05), (0,1)):
        for n_check in (1,3,9):
            assert math.isclose(core_module.noisy_selection(p_check,n_check,tpr_check,fpr_check),
                                noisy_verifier_success(p_check,n_check,tpr_check,fpr_check),
                                rel_tol=0,abs_tol=1e-12)
assert perfect_verifier_coverage(1e-18,1)>0
assert core_module.ReasoningCost().estimate(4000,4096)['total_kv_gib']==.98828125
for bad_count in (True, 1.5, 0):
    expect_error(ValueError, exact_vote_statistics, bad_count, .4, 6)
    expect_error(ValueError, noisy_verifier_success, .4, bad_count, .9, .05)
expect_error(ValueError, CandidateBatch, [[1.5]], [1], [[.9]], [[10]])
expect_error(ValueError, CandidateBatch, [[1]], [1], [[.9]], [[1.5]])
expect_error(ValueError, select_unique_mode, np.empty((0,2),dtype=int))
expect_error(ValueError, selection_accuracy, [], [])
expect_error(ValueError, select_unique_mode, [[-1, -1]])
assert select_unique_mode([[0, 1], [0, 0]]).tolist() == [-1, 0]
expect_error(ValueError, multinomial_mass, (1,0), (.5,.5000000001))
assert not verify_final_integer("FINAL: "+"9"*10000,42)
assert math.isclose(.5*perfect_verifier_coverage(.8,3),.496,abs_tol=1e-14)
print("Closed form, boundary cases, input validation and prompt KV: PASS")

"""Generate, check, select and stop on a finite polynomial-program task.

No LLM or API is used. The generator samples two unknown coefficients; a visible
example fixes the constant term. Tests execute the candidates rather than drawing
precomputed correctness flags. This isolates selection and accounting.
"""
from dataclasses import dataclass, asdict
import itertools
import json
import random
import re

@dataclass(frozen=True)
class ProgramTask:
    task_id: int
    a: int
    b: int
    c: int
    def target(self, x):
        return self.a*x*x+self.b*x+self.c

TASKS = [ProgramTask(i, *abc) for i, abc in enumerate(itertools.product(range(-2, 3), repeat=3))]
CALIBRATION = [t for t in TASKS if t.c == 0]
HELD_OUT = [t for t in TASKS if t.c != 0]
assert len(CALIBRATION) == 25 and len(HELD_OUT) == 100
assert not {t.task_id for t in CALIBRATION} & {t.task_id for t in HELD_OUT}

def execute_program(source, x):
    # A deliberately restricted grammar; no eval or arbitrary code execution.
    match = re.fullmatch(r"(-?\d+)\*x\*x\+(-?\d+)\*x\+(-?\d+)", source)
    if match is None:
        raise ValueError("not a polynomial in the declared grammar")
    a, b, c = map(int, match.groups())
    return a*x*x+b*x+c

class ProposalGenerator:
    def __init__(self):
        self.calls = 0
    def propose(self, task_id, visible_f0, candidate_index):
        # Same deterministic random stream for each task/policy comparison.
        # Hidden coefficients, test outputs and future candidates are unavailable.
        rng = random.Random(150_000 + task_id*10_000 + candidate_index)
        a, b = rng.randrange(-2, 3), rng.randrange(-2, 3)
        self.calls += 1
        return f"{a}*x*x+{b}*x+{visible_f0}"

class ProgramVerifier:
    def __init__(self):
        self.calls = 0
        self.test_cases = 0
    def score(self, source, examples):
        self.calls += 1
        results = []
        for x, target in examples:
            results.append(execute_program(source, x) == target)
            self.test_cases += 1
        return sum(results)/len(results), results

def search_program(task, *, budget, threshold=None):
    if isinstance(budget, bool) or not isinstance(budget, int) or budget <= 0:
        raise ValueError("budget must be a positive integer")
    if threshold is not None and threshold not in (.5, 1.):
        raise ValueError("threshold must be .5, 1, or None")
    generator, verifier = ProposalGenerator(), ProgramVerifier()
    # The generator sees only f(0). The verifier receives two more examples.
    examples = [(x, task.target(x)) for x in (1, -1)]
    records = []
    for index in range(budget):
        source = generator.propose(task.task_id, task.target(0), index)
        score, checks = verifier.score(source, examples)
        records.append(dict(index=index, source=source, score=score, checks=checks))
        if threshold is not None and score >= threshold:
            break
    # Deterministic earliest candidate on a score tie. No extra generation or scoring.
    best = max(range(len(records)), key=lambda i: records[i]['score'])
    assert generator.calls == verifier.calls == len(records)
    assert verifier.test_cases == 2*len(records)
    if threshold is not None and len(records) < budget:
        assert records[-1]['score'] >= threshold
        assert all(r['score'] < threshold for r in records[:-1])
    # Hidden tests are evaluation only; they cannot influence selection or stopping.
    before = (generator.calls, verifier.calls, verifier.test_cases)
    hidden = [all(execute_program(r['source'], x) == task.target(x) for x in (2, 3))
              for r in records]
    assert before == (generator.calls, verifier.calls, verifier.test_cases)
    result = dict(task_id=task.task_id, records=records, selected_index=best,
                  correct=hidden[best], observed_coverage=any(hidden),
                  generated=generator.calls, verifier_calls=verifier.calls,
                  visible_test_cases=verifier.test_cases, hidden_tests_max=2*len(records))
    assert json.loads(json.dumps(result)) == result
    return result

def evaluate_program_policy(tasks, *, budget, threshold=None):
    runs = [search_program(t, budget=budget, threshold=threshold) for t in tasks]
    n = len(runs)
    assert n > 0
    return dict(accuracy=sum(r['correct'] for r in runs)/n,
                mean_candidates=sum(r['generated'] for r in runs)/n,
                observed_coverage=sum(r['observed_coverage'] for r in runs)/n,
                generator_calls=sum(r['generated'] for r in runs),
                verifier_calls=sum(r['verifier_calls'] for r in runs),
                visible_test_cases=sum(r['visible_test_cases'] for r in runs)), runs

def run_program_search():
    results = {}
    traces = {}
    for name, budget, threshold in [('one', 1, None), ('fixed16', 16, None),
                                    ('stop_partial', 16, .5), ('stop_complete', 16, 1.)]:
        metrics, records = evaluate_program_policy(HELD_OUT, budget=budget, threshold=threshold)
        results[name], traces[name] = metrics, records
    assert results['fixed16']['generator_calls'] == 1600
    assert results['one']['generator_calls'] == 100
    assert results['stop_complete']['accuracy'] == results['fixed16']['accuracy']
    assert results['stop_complete']['mean_candidates'] < 16
    assert results['stop_partial']['accuracy'] < results['stop_complete']['accuracy']
    # Passing all three x values (one visible to the generator, two to the verifier)
    # identifies a degree-at-most-two polynomial exactly. This guarantee is grammar-specific.
    for name in ('fixed16', 'stop_complete'):
        for record in traces[name]:
            assert (record['records'][record['selected_index']]['score'] == 1) == record['correct']
    for name, result in results.items(): print(name, result)
    return results, traces


program_results, program_traces = run_program_search()

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

"""Generate, check, select and stop on a finite polynomial-program task.

No LLM or API is used. The generator samples two unknown coefficients; a visible
example fixes the constant term. Tests execute the candidates rather than drawing
precomputed correctness flags. This isolates selection and accounting.
"""
from dataclasses import dataclass, asdict
import itertools
import json
import random
import re

@dataclass(frozen=True)
class ProgramTask:
    task_id: int
    a: int
    b: int
    c: int
    def target(self, x):
        return self.a*x*x+self.b*x+self.c

TASKS = [ProgramTask(i, *abc) for i, abc in enumerate(itertools.product(range(-2, 3), repeat=3))]
CALIBRATION = [t for t in TASKS if t.c == 0]
HELD_OUT = [t for t in TASKS if t.c != 0]
assert len(CALIBRATION) == 25 and len(HELD_OUT) == 100
assert not {t.task_id for t in CALIBRATION} & {t.task_id for t in HELD_OUT}

def execute_program(source, x):
    # A deliberately restricted grammar; no eval or arbitrary code execution.
    match = re.fullmatch(r"(-?\d+)\*x\*x\+(-?\d+)\*x\+(-?\d+)", source)
    if match is None:
        raise ValueError("not a polynomial in the declared grammar")
    a, b, c = map(int, match.groups())
    return a*x*x+b*x+c

class ProposalGenerator:
    def __init__(self):
        self.calls = 0
    def propose(self, task_id, visible_f0, candidate_index):
        # Same deterministic random stream for each task/policy comparison.
        # Hidden coefficients, test outputs and future candidates are unavailable.
        rng = random.Random(150_000 + task_id*10_000 + candidate_index)
        a, b = rng.randrange(-2, 3), rng.randrange(-2, 3)
        self.calls += 1
        return f"{a}*x*x+{b}*x+{visible_f0}"

class ProgramVerifier:
    def __init__(self):
        self.calls = 0
        self.test_cases = 0
    def score(self, source, examples):
        self.calls += 1
        results = []
        for x, target in examples:
            results.append(execute_program(source, x) == target)
            self.test_cases += 1
        return sum(results)/len(results), results

def search_program(task, *, budget, threshold=None):
    if isinstance(budget, bool) or not isinstance(budget, int) or budget <= 0:
        raise ValueError("budget must be a positive integer")
    if threshold is not None and threshold not in (.5, 1.):
        raise ValueError("threshold must be .5, 1, or None")
    generator, verifier = ProposalGenerator(), ProgramVerifier()
    # The generator sees only f(0). The verifier receives two more examples.
    examples = [(x, task.target(x)) for x in (1, -1)]
    records = []
    for index in range(budget):
        source = generator.propose(task.task_id, task.target(0), index)
        score, checks = verifier.score(source, examples)
        records.append(dict(index=index, source=source, score=score, checks=checks))
        if threshold is not None and score >= threshold:
            break
    # Deterministic earliest candidate on a score tie. No extra generation or scoring.
    best = max(range(len(records)), key=lambda i: records[i]['score'])
    assert generator.calls == verifier.calls == len(records)
    assert verifier.test_cases == 2*len(records)
    if threshold is not None and len(records) < budget:
        assert records[-1]['score'] >= threshold
        assert all(r['score'] < threshold for r in records[:-1])
    # Hidden tests are evaluation only; they cannot influence selection or stopping.
    before = (generator.calls, verifier.calls, verifier.test_cases)
    hidden = [all(execute_program(r['source'], x) == task.target(x) for x in (2, 3))
              for r in records]
    assert before == (generator.calls, verifier.calls, verifier.test_cases)
    result = dict(task_id=task.task_id, records=records, selected_index=best,
                  correct=hidden[best], observed_coverage=any(hidden),
                  generated=generator.calls, verifier_calls=verifier.calls,
                  visible_test_cases=verifier.test_cases, hidden_tests_max=2*len(records))
    assert json.loads(json.dumps(result)) == result
    return result

def evaluate_program_policy(tasks, *, budget, threshold=None):
    runs = [search_program(t, budget=budget, threshold=threshold) for t in tasks]
    n = len(runs)
    assert n > 0
    return dict(accuracy=sum(r['correct'] for r in runs)/n,
                mean_candidates=sum(r['generated'] for r in runs)/n,
                observed_coverage=sum(r['observed_coverage'] for r in runs)/n,
                generator_calls=sum(r['generated'] for r in runs),
                verifier_calls=sum(r['verifier_calls'] for r in runs),
                visible_test_cases=sum(r['visible_test_cases'] for r in runs)), runs

def run_program_search():
    results = {}
    traces = {}
    for name, budget, threshold in [('one', 1, None), ('fixed16', 16, None),
                                    ('stop_partial', 16, .5), ('stop_complete', 16, 1.)]:
        metrics, records = evaluate_program_policy(HELD_OUT, budget=budget, threshold=threshold)
        results[name], traces[name] = metrics, records
    assert results['fixed16']['generator_calls'] == 1600
    assert results['one']['generator_calls'] == 100
    assert results['stop_complete']['accuracy'] == results['fixed16']['accuracy']
    assert results['stop_complete']['mean_candidates'] < 16
    assert results['stop_partial']['accuracy'] < results['stop_complete']['accuracy']
    # Passing all three x values (one visible to the generator, two to the verifier)
    # identifies a degree-at-most-two polynomial exactly. This guarantee is grammar-specific.
    for name in ('fixed16', 'stop_complete'):
        for record in traces[name]:
            assert (record['records'][record['selected_index']]['score'] == 1) == record['correct']
    for name, result in results.items(): print(name, result)
    return results, traces


program_results, program_traces = run_program_search()

# Choose the cheapest candidate budget meeting a declared calibration target.
# The target is empirical accuracy, not a statistical lower confidence bound.
target_accuracy = .60
calibration_rows = []
for budget in (1, 4, 8, 16, 32):
    metrics, _ = evaluate_program_policy(CALIBRATION, budget=budget, threshold=1.)
    calibration_rows.append(dict(budget=budget, **metrics))
eligible = [r for r in calibration_rows if r['accuracy'] >= target_accuracy]
if not eligible:
    raise RuntimeError("No calibration budget meets the target")
chosen = min(eligible, key=lambda r: (r['mean_candidates'], r['budget']))
frozen_policy = dict(budget=chosen['budget'], threshold=1.)
held_out, held_out_runs = evaluate_program_policy(HELD_OUT, **frozen_policy)
assert chosen['budget'] == 32
assert held_out['generator_calls'] == held_out['verifier_calls']
assert held_out['visible_test_cases'] == 2*held_out['verifier_calls']
artifact = dict(policy=frozen_policy, calibration=calibration_rows, held_out=held_out)
assert json.loads(json.dumps(artifact)) == artifact
print("calibration:", calibration_rows)
print("frozen policy:", frozen_policy, "held-out:", held_out)
# Inspect an actual wrong selected program; never use this inspection to retune test results.
failure = next((r for r in held_out_runs if not r['correct']), None)
print("held-out failure:", failure)