"""后训练的奖励契约、组内信号与实测记录；默认只读记录，不训练。"""

from argparse import ArgumentParser
from hashlib import sha256
import json
import math
from pathlib import Path
import random

DATA = Path(__file__).with_name("data")
PLAN = DATA / "post-training-plan.json"
RECORD = DATA / "post-training-record.json"


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def make_plan():
    groups = [(a, b) for a in range(10) for b in range(a, 10)]
    random.Random(20261020).shuffle(groups)
    split = {}
    for name, selected in zip(("train", "dev", "test"),
                              (groups[:35], groups[35:45], groups[45:])):
        split[name] = sorted([[a, b] for a, b in selected]
                             + [[b, a] for a, b in selected if a != b])
    return {
        "version": 1, "date": "2026-10-20", "task": "single_digit_addition",
        "split": split, "seeds": [11, 22, 33],
        "vocab": "0123456789+=~_", "eos": "~", "pad": "_",
        "embedding": 24, "hidden": 48, "threads": 1,
        "sft_steps": 200, "sft_batch": 16, "sft_lr": .003,
        "rl_steps": 240, "prompts_per_step": 8, "group_size": 8,
        "updates_per_rollout": 2, "rl_lr": .0005, "clip": .2,
        "kl_beta": .03, "max_new_tokens": 6, "temperature": 1.,
        "eval_samples": 16, "gradient_clip": 1.,
        "arms": ["strict", "substring"],
        "loss_reduction": "mean over all valid response tokens, including EOS",
        "advantage": "group population std + 1e-8; no dynamic resampling",
        "kl": "exact categorical KL at old-policy sampled histories",
        "selection": "fixed last checkpoint; no tuning or seed selection",
        "rewards": {
            "strict": "terminated, nonempty ASCII digits, integer equals a+b; leading zeros allowed",
            "substring": "decimal str(a+b) occurs anywhere, even without EOS",
        },
    }


def reward(pair, text, terminated, arm):
    if arm == "substring":
        return int(str(sum(pair)) in text)
    if arm != "strict":
        raise ValueError(arm)
    return int(terminated and bool(text)
               and all(c in "0123456789" for c in text)
               and int(text) == sum(pair))


def contrast_probability(p, group=8):
    return 1 - p**group - (1-p)**group


def advantages(rewards):
    mean = sum(rewards) / len(rewards)
    std = math.sqrt(sum((r-mean)**2 for r in rewards) / len(rewards))
    return [(r-mean)/(std+1e-8) for r in rewards]


def metrics(rows):
    # Each row stores [a, b, visible text, emitted EOS].
    n = len(rows)
    strict = sum(reward(r[:2], r[2], r[3], "strict") for r in rows)
    weak = sum(reward(r[:2], r[2], r[3], "substring") for r in rows)
    return {"n": n, "strict": strict, "substring": weak,
            "strict_rate": strict/n, "substring_rate": weak/n,
            "mean_tokens": sum(len(r[2])+int(r[3]) for r in rows)/n,
            "truncated": sum(not r[3] for r in rows)}


def evaluate(rows, samples):
    assert len(rows) % samples == 0
    covered = 0
    for start in range(0, len(rows), samples):
        group = rows[start:start+samples]
        assert all(row[:2] == group[0][:2] for row in group)
        covered += any(reward(r[:2], r[2], r[3], "strict") for r in group)
    return dict(metrics(rows), covered=covered, questions=len(rows)//samples)


def summarize(record, plan):
    summaries = []
    for run in record["runs"]:
        for stage, ev in run["evaluations"].items():
            row = {"seed": run["seed"], "stage": stage}
            for split, samples in ev.items():
                row[split] = {"greedy": evaluate(samples["greedy"], 1),
                              "sampled": evaluate(samples["sampled"], plan["eval_samples"])}
            summaries.append(row)
    return summaries


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--freeze", type=Path)
    args = parser.parse_args()
    if args.freeze:
        # Refuse replacement: experiment rules are saved before training.
        with args.freeze.open("x") as stream:
            json.dump(make_plan(), stream, ensure_ascii=False, indent=2)
            stream.write("\n")
        print(f"已冻结计划：{args.freeze}")
        return
    plan = json.loads(PLAN.read_text())
    record = json.loads(RECORD.read_text())
    assert record["plan_sha256"] == digest(PLAN)
    print("组内出现不同奖励的概率：")
    for p in [.01, .5, .99]:
        print(f"p={p:.2f}: {contrast_probability(p):.6f}")
    print("留出测试：种子 / 阶段 / 贪心正确 / 采样正确 / 子串接受 / 16次覆盖")
    for row in summarize(record, plan):
        g, s = row["test"]["greedy"], row["test"]["sampled"]
        print(f"{row['seed']:2d} {row['stage']:9s} "
              f"{g['strict']}/{g['n']} "
              f"{s['strict']}/{s['n']} {s['substring']}/{s['n']} "
              f"{s['covered']}/{s['questions']}")


if __name__ == "__main__":
    main()