"""第 37 章：显式 --run 测量 CPU KV 缓存；默认读取随书保存的原始测量。"""

import argparse
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import platform
import statistics
import torch
from cached_decoder import TinyTransformer, generate

DATA = Path(__file__).with_name("data")/"inference-benchmark.json"
CONFIG = dict(vocab_size=128, max_context=512, width=64, heads=4, layers=2)
LENGTHS = [16, 128]
NEW_TOKENS, WARMUPS, REPEATS = 32, 3, 7


def model_and_inputs():
    torch.set_num_threads(1)
    torch.manual_seed(2036)
    model = TinyTransformer(**CONFIG).eval()
    inputs = [torch.randint(0, CONFIG["vocab_size"], (1, length)) for length in LENGTHS]
    return model, inputs


def model_hash(model):
    digest = hashlib.sha256()
    for key, value in model.state_dict().items():
        digest.update(key.encode())
        digest.update(value.detach().numpy().tobytes())
    return digest.hexdigest()


def summarize(rows):
    result = []
    for length in LENGTHS:
        for cached in [False, True]:
            selected = [r for r in rows if r["length"] == length and r["cached"] == cached]
            totals = [1000*r["total_seconds"] for r in selected]
            first = [1000*r["first_seconds"] for r in selected]
            speed = [(NEW_TOKENS-1)/(r["total_seconds"]-r["first_seconds"]) for r in selected]
            result.append(dict(length=length, cached=cached,
                               median_total_ms=statistics.median(totals),
                               median_first_ms=statistics.median(first),
                               median_decode_tps=statistics.median(speed),
                               min_total_ms=min(totals), max_total_ms=max(totals),
                               cache_bytes=selected[0]["cache_bytes"]))
    return result


def benchmark(path):
    model, inputs = model_and_inputs()
    rows = []
    for prompt in inputs:
        for _ in range(WARMUPS):
            for cached in [False, True]:
                generate(model, prompt, NEW_TOKENS, cached)
        expected = None
        for trial in range(REPEATS):
            # 交替顺序，减少总让一个方案先运行的偏差；仍不能消除系统噪声。
            for cached in ([False, True] if trial % 2 == 0 else [True, False]):
                row = generate(model, prompt, NEW_TOKENS, cached)
                ids = row.pop("ids").tolist()
                if expected is None:
                    expected = ids
                assert ids == expected, "缓存前后生成不同，不能直接比较速度"
                rows.append(dict(length=prompt.shape[1], cached=cached, trial=trial, ids=ids, **row))
    report = dict(complete=True, measured_at=datetime.now(timezone.utc).isoformat(),
                  environment=dict(python=platform.python_version(), torch=torch.__version__,
                                   platform=platform.platform(), processor=platform.processor(),
                                   device="cpu", dtype="float32", threads=1),
                  config=CONFIG, seed=2036, new_tokens=NEW_TOKENS, warmups=WARMUPS,
                  repeats=REPEATS, prompts=[p.tolist() for p in inputs],
                  model_sha256=model_hash(model), rows=rows)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
    return report


def checked_report(path=DATA):
    report = json.loads(path.read_text())
    assert report["complete"] and report["config"] == CONFIG
    assert report["new_tokens"] == NEW_TOKENS and report["repeats"] == REPEATS
    assert report["warmups"] == WARMUPS and report["seed"] == 2036
    model, prompts = model_and_inputs()
    assert report["model_sha256"] == model_hash(model)
    assert report["prompts"] == [p.tolist() for p in prompts]
    assert len(report["rows"]) == len(LENGTHS)*2*REPEATS
    keys = {(r["length"], r["cached"], r["trial"]) for r in report["rows"]}
    assert keys == {(length, cached, i) for length in LENGTHS
                    for cached in [False, True] for i in range(REPEATS)}
    for row in report["rows"]:
        stamps = row["token_seconds"]
        assert len(stamps) == NEW_TOKENS and 0 < stamps[0] < stamps[-1]
        assert stamps == sorted(stamps)
        assert row["first_seconds"] == stamps[0] and row["total_seconds"] == stamps[-1]
    # 用当前代码复核每种输入/路径的实际输出与计数，不把它当作重新测量原时间。
    for prompt in prompts:
        for cached in [False, True]:
            current = generate(model, prompt, NEW_TOKENS, cached)
            for row in report["rows"]:
                if row["length"] == prompt.shape[1] and row["cached"] == cached:
                    assert row["ids"] == current["ids"].tolist()
                    for key in ["projected_positions", "score_elements", "cache_bytes"]:
                        assert row[key] == current[key]
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    if args.run and args.out is None:
        parser.error("请用 --out 保存自己的测量，不覆盖随书记录")
    report = benchmark(args.out) if args.run else checked_report()
    print("以下为本次测量：" if args.run else "以下为随书保存的 CPU 单线程测量：")
    for row in summarize(report["rows"]):
        print(f"提示 {row['length']:3d}；缓存={row['cached']}；"
              f"首词元中位数={row['median_first_ms']:.2f} ms；"
              f"总时长中位数={row['median_total_ms']:.2f} ms；"
              f"后续速度中位数={row['median_decode_tps']:.1f} token/s；"
              f"KV={row['cache_bytes']} 字节")
    print("每种条件预热 3 次、测量 7 次；同条件的生成 ID 完全一致。")
    print("随机权重只用于性能路径比较，不代表语言能力；本地结果可能不同。")


if __name__ == "__main__":
    main()