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

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()