"""推理与测试时计算：固定采样方案、投票计分及精确有理数搜索。"""

import argparse
import ast
from collections import Counter
from fractions import Fraction
import hashlib
from itertools import combinations
import json
from math import comb
from pathlib import Path


DATA = Path(__file__).resolve().parent / "data"
MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
REVISION = "7ae557604adf67be50417f59c2c2f167def9a775"
SYSTEM = ("Solve the arithmetic question. Give a brief calculation, then end "
          "with one line exactly FINAL: <integer>. Do not include units or "
          "any text after that final line.")


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


def fingerprint(value):
    raw = json.dumps(value, sort_keys=True, ensure_ascii=False,
                     separators=(",", ":")).encode()
    return hashlib.sha256(raw).hexdigest()


def save_new(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("x") as stream:
        json.dump(value, stream, ensure_ascii=False, indent=2)
        stream.write("\n")


def make_plan():
    problems = [
        ("pages", "Printing takes 3 seconds per page plus 20 seconds of setup. "
         "How many seconds does one 17-page job take?", 71, "3*17+20"),
        ("color", "A document has 10 black-and-white pages and 2 color pages. "
         "Each black-and-white page takes 3 seconds and each color page 5 seconds. "
         "Add 20 seconds of setup once. How many seconds in total?", 60, "10*3+2*5+20"),
        ("packs", "There are 3 trays with 12 parts each. Discard 5 defective parts. "
         "Put the remaining parts into packs of 4. How many full packs can be made?",
         7, "(3*12-5)//4"),
        ("discount", "Buy 4 spools at 18 coins each. Apply a 25 percent discount "
         "to the spool subtotal, then add a shipping fee of 6 coins. "
         "How many coins are paid?", 60, "4*18*3/4+6"),
        ("queue", "Three jobs run one after another, taking 30, 45 and 25 seconds. "
         "There is a 5-second pause between adjacent jobs, but none before the "
         "first or after the last. How many seconds from start to finish?", 110, "30+45+25+2*5"),
        ("mixture", "A 45 mL mixture has red and blue liquid in the ratio 3:2. "
         "Add 15 mL of blue liquid. How many mL of blue liquid are now present?",
         33, "45*2/5+15"),
        ("rate", "Printer A makes 12 parts per hour and printer B makes 18 parts "
         "per hour. Both run for 2 hours at constant rates. Exactly one fifth "
         "of all produced parts are defective. How many nondefective parts remain?",
         48, "(12+18)*2*4/5"),
        ("cycles", "Two signals flash together now. One repeats every 18 minutes, "
         "the other every 24 minutes. After how many minutes will they first "
         "flash together again?", 72, "lcm(18,24)"),
    ]
    return dict(
        schema=1, model=MODEL, revision=REVISION, system=SYSTEM,
        problems=[dict(id=i, question=q, answer=a, reference=r) for i, q, a, r in problems],
        generation=dict(samples=4, max_new_tokens=192, do_sample=True,
                        temperature=0.8, top_p=0.95, top_k=0,
                        attention="eager", device="cpu", dtype="float32", threads=4,
                        input_plus_output_cap=2048, first_seed=20261020),
        evaluation=dict(prefix_budgets=[1, 2, 4], tie_break="earliest valid candidate",
                        invalid="abstain from voting; the problem remains in the denominator"),
        puzzles=[[1, 2, 3, 4], [2, 3, 4, 5], [3, 3, 8, 8], [1, 5, 5, 5],
                 [2, 2, 6, 6], [1, 3, 4, 6], [4, 7, 8, 8], [1, 2, 5, 8]],
    )


def messages(problem):
    return [{"role": "system", "content": SYSTEM},
            {"role": "user", "content": problem["question"]}]


def parse_answer(text):
    import re
    lines = text.strip().splitlines()
    if not lines:
        return None
    match = re.fullmatch(r"FINAL:\s*(-?[0-9]+)", lines[-1])
    return int(match.group(1)) if match else None


def majority(rows):
    answers = [parse_answer(row["text"]) for row in rows]
    answers = [a for a in answers if a is not None]
    if not answers:
        return None
    counts = Counter(answers)
    return max(answers, key=counts.get)  # 并列时保持最早出现者。


def verify_expression(expression, numbers):
    """限定二元四则运算和正整数叶子；不使用 eval，也不执行任何代码。"""
    leaves = []
    if len(expression) > 500:
        return False

    def visit(node):
        if isinstance(node, ast.Constant) and type(node.value) is int and node.value > 0:
            leaves.append(node.value)
            return Fraction(node.value)
        if not isinstance(node, ast.BinOp):
            raise ValueError("只接受二元四则运算")
        left, right = visit(node.left), visit(node.right)
        if isinstance(node.op, ast.Add):
            return left + right
        if isinstance(node.op, ast.Sub):
            return left - right
        if isinstance(node.op, ast.Mult):
            return left * right
        if isinstance(node.op, ast.Div):
            return left / right
        raise ValueError("运算符不在规则中")

    try:
        value = visit(ast.parse(expression, mode="eval").body)
    except (SyntaxError, ValueError, ZeroDivisionError, RecursionError):
        return False
    return value == 24 and Counter(leaves) == Counter(numbers)


def state_key(items):
    return tuple(sorted(value for value, _ in items))


def children(items):
    for i, j in combinations(range(len(items)), 2):
        a, ae = items[i]
        b, be = items[j]
        rest = [item for k, item in enumerate(items) if k not in (i, j)]
        variants = [(a+b, f"({ae}+{be})"), (a*b, f"({ae}*{be})"),
                    (a-b, f"({ae}-{be})"), (b-a, f"({be}-{ae})")]
        if b:
            variants.append((a/b, f"({ae}/{be})"))
        if a:
            variants.append((b/a, f"({be}/{ae})"))
        for value, expression in variants:
            yield sorted(rest + [(value, expression)], key=lambda x: (x[0], x[1]))


def search_24(numbers, width=None):
    # 状态相同的精确值有同样的后续可达性；保留其中一条构造表达式的路径。
    frontier = [[(Fraction(n), str(n)) for n in sorted(numbers)]]
    expanded = generated = 0
    trace = []
    for depth in range(1, len(numbers)):
        unique = {}
        for items in frontier:
            expanded += 1
            for child in children(items):
                generated += 1
                unique.setdefault(state_key(child), child)
        ordered = sorted(unique.values(),
                         key=lambda items: (min(abs(v-24) for v, _ in items),
                                            state_key(items)))
        kept = ordered if width is None else ordered[:width]
        trace.append(dict(depth=depth, unique=len(ordered), kept=len(kept),
                          values=[[str(v) for v, _ in items] for items in kept]))
        frontier = kept
    solutions = [items[0][1] for items in frontier
                 if len(items) == 1 and items[0][0] == 24]
    expression = solutions[0] if solutions else None
    if expression is not None and not verify_expression(expression, numbers):
        raise AssertionError("搜索输出没有通过独立表达式检查")
    return dict(numbers=numbers, width=width, success=expression is not None,
                expression=expression, expanded=expanded, generated=generated, trace=trace)


def evaluate_record(record):
    plan = make_plan()
    if not record["complete"] or record["plan_sha256"] != fingerprint(plan):
        raise ValueError("模型记录未完成或方案不同")
    keys = {(r["id"], r["sample"]) for r in record["rows"]}
    expected = {(p["id"], i) for p in plan["problems"] for i in range(4)}
    if keys != expected or len(record["rows"]) != len(expected):
        raise ValueError("缺失或重复的模型调用")
    by_key = {(r["id"], r["sample"]): r for r in record["rows"]}
    results = []
    for budget in plan["evaluation"]["prefix_budgets"]:
        details = []
        for problem in plan["problems"]:
            rows = [by_key[problem["id"], i] for i in range(budget)]
            if any(row["messages"] != messages(problem) for row in rows):
                raise ValueError("提示与固定方案不同")
            answers = [parse_answer(row["text"]) for row in rows]
            voted = majority(rows)
            details.append(dict(id=problem["id"], answers=answers, voted=voted,
                                correct=voted == problem["answer"],
                                oracle=problem["answer"] in answers,
                                input_tokens=sum(r["input_tokens"] for r in rows),
                                output_tokens=sum(r["output_tokens"] for r in rows),
                                invalid=sum(a is None for a in answers),
                                capped=sum(not r["stopped_on_eos"] for r in rows)))
        results.append(dict(budget=budget, rows=details))
    return results


def theory():
    return dict(
        heterogeneous_success=[0.5 for n in (1, 3, 5)],
        wrong_average_substitution=[1-0.5**n for n in (1, 3, 5)],
        majority_p04=[sum(comb(n, k)*0.4**k*0.6**(n-k)
                          for k in range(n//2+1, n+1)) for n in (1, 3, 5)],
        false_positive_16=1-0.95**16,
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--record", type=Path, default=DATA / "test-time-model.json")
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    record = json.loads(args.record.read_text())
    results = evaluate_record(record)
    searches = [search_24(puzzle, width) for width in (1, 4, None)
                for puzzle in make_plan()["puzzles"]]
    report = dict(model_record_sha256=digest(args.record),
                  source_sha256=digest(__file__), sampling=results,
                  searches=searches, theory=theory())
    if args.out:
        save_new(args.out, report)
    for result in results:
        rows = result["rows"]
        print(f"每题 {result['budget']} 次：投票正确 {sum(r['correct'] for r in rows)}/8；"
              f"候选中存在正确 {sum(r['oracle'] for r in rows)}/8；"
              f"输入/输出词元 {sum(r['input_tokens'] for r in rows)}/"
              f"{sum(r['output_tokens'] for r in rows)}；"
              f"无合规答案 {sum(r['invalid'] for r in rows)}，触顶 {sum(r['capped'] for r in rows)}。")
    for width in (1, 4, None):
        subset = [r for r in searches if r["width"] == width]
        print(f"24 点搜索，宽度 {width if width else '全部'}："
              f"成功 {sum(r['success'] for r in subset)}/8；"
              f"扩展状态 {sum(r['expanded'] for r in subset)}，"
              f"生成候选状态 {sum(r['generated'] for r in subset)}。")
    print("正确概率 0.4、错误都投同一答案时，1/3/5 票多数正确率："
          + " / ".join(f"{p:.5f}" for p in theory()["majority_p04"]))
    print(f"16 个错误候选各有独立 5% 误放率，至少误放一个："
          f"{theory()['false_positive_16']:.5f}。")
    print("默认重算保存的生成记录与精确搜索，不调用语言模型。")


if __name__ == "__main__":
    main()

"""显式 --run 才执行测试时采样实验；先固定方案，再生成全部 32 个回答。"""

import argparse
from datetime import datetime, timezone
from importlib.metadata import version
import json
from pathlib import Path
import platform
import time

import test_time_compute as lab


def experiment(directory):
    directory.mkdir(parents=True, exist_ok=False)
    plan = lab.make_plan()
    lab.save_new(directory / "plan.json", plan)
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer
    settings = plan["generation"]
    torch.set_num_threads(settings["threads"])
    snapshot = Path(snapshot_download(
        plan["model"], revision=plan["revision"],
        allow_patterns=["*.json", "*.safetensors", "merges.txt", "README.md", "LICENSE"]))
    tokenizer = AutoTokenizer.from_pretrained(snapshot, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        snapshot, dtype=torch.float32, attn_implementation=settings["attention"],
        trust_remote_code=False).eval()
    report = dict(
        complete=False, started_at=datetime.now(timezone.utc).isoformat(),
        model=plan["model"], revision=plan["revision"], generation=settings,
        plan_sha256=lab.fingerprint(plan),
        source_sha256={path.name: lab.digest(path) for path in (Path(lab.__file__), Path(__file__))},
        model_files_sha256={path.name: lab.digest(path) for path in snapshot.iterdir()
                            if path.is_file() and (path.suffix in {".json", ".safetensors"}
                                                   or path.name == "merges.txt")},
        environment=dict(python=platform.python_version(), platform=platform.platform(),
                         packages={name: version(name) for name in
                                   ("torch", "transformers", "huggingface-hub",
                                    "tokenizers", "safetensors")}),
        rows=[],
    )
    target = directory / "model.json"

    def checkpoint():
        temporary = directory / "model.pending.json"
        temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
        temporary.replace(target)

    checkpoint()
    eos = model.generation_config.eos_token_id
    eos = [eos] if isinstance(eos, int) else eos
    for problem_index, problem in enumerate(plan["problems"]):
        conversation = lab.messages(problem)
        rendered = tokenizer.apply_chat_template(
            conversation, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer.apply_chat_template(
            conversation, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt")
        length = inputs["input_ids"].shape[1]
        if length + settings["max_new_tokens"] > settings["input_plus_output_cap"]:
            raise ValueError("超过固定预算，不静默截断")
        for sample in range(settings["samples"]):
            seed = settings["first_seed"] + problem_index * settings["samples"] + sample
            torch.manual_seed(seed)
            started = time.perf_counter()
            with torch.inference_mode():
                output = model.generate(
                    **inputs, max_new_tokens=settings["max_new_tokens"],
                    do_sample=True, temperature=settings["temperature"],
                    top_p=settings["top_p"], top_k=settings["top_k"],
                    pad_token_id=tokenizer.eos_token_id)
            seconds = time.perf_counter() - started
            ids = output[0, length:].tolist()
            text = tokenizer.decode(ids, skip_special_tokens=True)
            report["rows"].append(dict(
                id=problem["id"], sample=sample, seed=seed,
                messages=conversation, rendered_prompt=rendered,
                input_ids=inputs["input_ids"][0].tolist(), input_tokens=length,
                output_ids=ids, output_tokens=len(ids), text=text,
                stopped_on_eos=bool(ids and ids[-1] in eos), seconds=seconds,
                parsed=lab.parse_answer(text)))
            checkpoint()
            print(f"{problem['id']}/{sample}: answer={lab.parse_answer(text)}, "
                  f"output={len(ids)}, eos={report['rows'][-1]['stopped_on_eos']}", flush=True)
    report["complete"] = True
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    checkpoint()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("显式传入 --run --out-dir 新目录，才执行固定的 32 次真实模型采样。")
        return
    if args.out_dir is None:
        parser.error("请指定尚不存在的输出目录")
    experiment(args.out_dir)


if __name__ == "__main__":
    main()

1. $1-0.8^4=0.5904$。它只要求四份候选中有一份正确；投票还取决于其余答案的分布和并列规则。
2. 存在正确候选，但投票会选 70。这就是候选覆盖与最终选择的差别。
3. 仍为 0.5。误代入平均值会得到 $1-0.5^3=0.875$，因为它把必对与必错的题误当成单题随机成功率为 0.5。
4. 只有 FINAL: 48。可以同时保留契约正确率和事后内容复核，明确两个规则及复核时机；不能把后者悄悄替换成原始实验得分。
5. 只有 24 没有使用四张牌。正确表达式使用两次 8、两次 3，且分母为 $1/3$，不为零。
6. 还要使用所有剩余牌。当前八组的扩展状态从 24 增加到 72，生成候选从 479 增加到 1055，成功仍为 1/8。这里没有测量语言模型成本。
7. 这八组已参与阅读和方法调整，属于开发材料。先固定新评分、宽度、停止规则和计数口径，再对未用于调整的新输入运行；仍需说明新输入怎样产生，不只保留成功牌组。
8. 中间步骤、单位或错误推导中的整数，可能被误当成最终答案。被截断的 rate 回答最后提到了 1，却没有完整结论；“最后一个整数”会让没有答案的轨迹看起来有答案。
9. 不能，分组中的题目难度不同。应把相同的留出题分别交给各模式，或做预先设计的随机分配，再比较正确率与成本；还要把路由成本计入自适应方案。
10. 保存每题每轮答案、正确到错误、错误到正确、始终正确、始终错误的数量，以及输入/输出词元、验证器调用、结束原因和实际耗时。只看最后准确率，会掩盖纠错与破坏同时发生的情况。