"""
=============================================================================
CS349E · 第 6 讲 · 03_speculative_sampling.py —— 投机采样为什么不改变输出分布
=============================================================================

【这个脚本要回答的问题】
    投机解码最让人惊讶的一点：草稿模型再差，输出的分布也和目标模型**完全一样**（Leviathan et al. 2023 §2.3、附录 A.1）。
    靠的是一个拒绝采样规则：
        草稿抽到 x（概率 q(x)）；以概率 min(1, p(x)/q(x)) 接受；
        被拒绝时，从 p'(x) = norm(max(0, p(x) − q(x))) 重新抽一个。
    这个脚本从三个角度检验它：

      演示 1：精确计算——不抽样，直接按概率算出一步的输出分布，和 p 比较；β = Σ min(p, q) = 1 − TV(p, q)。
              用第 3 讲的模型（QCFG）和几种草稿：INT4 量化的自己、只用前几层的自己、二元语法表、一个无关的小模型。
      演示 2：蒙特卡洛——在一个极小的模型（词表 64）上跑 4 万个独立的投机解码，统计前两个输出 token 的联合分布，
              和目标模型的精确分布比较；再看两个**错误**的变体偏了多少。
      演示 3：接受率 α 和草稿、温度的关系（以及这个玩具模型和真实模型不一样的地方）。
      演示 4：贪心解码的特例：投机解码的输出和目标模型自己贪心生成的逐 token 相同。

【怎么读输出】TV（总变差距离）= ½ Σ |P(x) − Q(x)|，0 表示两个分布完全相同、1 表示完全不重叠。
    蒙特卡洛的 TV 不会是 0：即使直接从正确的分布里抽 4 万个样本，经验分布和真分布也有抽样误差——
    表里给出"直接抽样的 TV"作为噪声水平，投机解码的 TV 应该落在这个水平上。

【运行】python3 03_speculative_sampling.py      （约 15 秒）
【依赖】numpy；本目录的 decode_lab；../lecture02 的 transformer_lab；../lecture03 的 quant_lab（它需要 ../prereq3）。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
for sub in ("lecture03", "prereq3"):
    sys.path.insert(0, os.path.join(HERE, "..", sub))

from decode_lab import (NGramLM, TransformerLM, adjust, autoregressive_generate, spec_output_dist,  # noqa: E402
                        spec_sample_batch, speculative_generate)
from console import print_header, print_observations, print_table  # noqa: E402
from quant_lab import QCFG, quantize_params, sample_tokens  # noqa: E402
from transformer_lab import Config, config_like, init_params  # noqa: E402


def build_models():
    """
    第 3 讲的模型 QCFG（4 层、d=256、词表 1024，随机初始化）作为目标模型，以及 5 个草稿模型：
        INT8 / INT4（每 32 个一组）量化的自己（含 lm_head）——"量化自草稿"；
        只用前 1 层 / 前 3 层（加上最后的 norm 和 lm_head）——"跳层自草稿"（self-speculative / LayerSkip 的思路）；
        二元语法表：从目标模型自己生成的 8,192 个 token 里数出来（几乎零成本的草稿）；
        一个无关的小模型（随机初始化、结构不同）——最差的草稿。
    【返回】(目标模型, {名字: 草稿模型})
    """
    P = init_params(QCFG, seed=0)
    target = TransformerLM(P, QCFG, "目标")
    drafts = {}
    for scheme, label in (("int8-channel", "INT8 量化的自己"), ("int4-g32", "INT4 量化的自己")):
        drafts[label] = TransformerLM(quantize_params(P, scheme, include_lm_head=True)[0], QCFG, label)
    for k in (1, 3):
        pk = dict(P)
        pk["layers"] = P["layers"][:k]
        drafts[f"只用前 {k} 层"] = TransformerLM(pk, config_like(QCFG, n_layers=k), f"前{k}层")
    corpus = sample_tokens(P, QCFG, n_seq=64, length=128, seed=1)
    drafts["二元语法表"] = NGramLM.fit(corpus, QCFG.vocab)
    small = config_like(QCFG, n_layers=1, d_model=64, n_heads=2, n_kv_heads=1, d_head=32, d_ff=192)
    drafts["无关的小模型"] = TransformerLM(init_params(small, seed=5), small, "无关")
    return target, drafts


def draft_logits(draft, toks):
    """草稿模型在一批序列 toks [N, L] 每个位置上的 logits（二元语法表直接查表）。"""
    if isinstance(draft, NGramLM):
        return draft.table[toks]
    return draft.full_logits(toks)


# -----------------------------------------------------------------------------
def demo1_exact(target, drafts, toks):
    print_header("演示 1：一步投机采样的输出分布——精确计算（不抽样）")
    lp = target.full_logits(toks)
    rows = []
    for name, dr in drafts.items():
        lq = draft_logits(dr, toks)
        worst = 0.0
        betas, tvs = [], []
        for (T, k) in ((1.0, None), (0.5, None), (1.0, 50)):
            p = adjust(lp, T, k).reshape(-1, QCFG.vocab)
            q = adjust(lq, T, k).reshape(-1, QCFG.vocab)
            out, beta = spec_output_dist(p, q)
            worst = max(worst, float(np.max(np.abs(out - p))))
            if T == 1.0 and k is None:
                betas = beta
                tvs = 0.5 * np.abs(p - q).sum(axis=1)
        rows.append([name, f"{worst:.1e}", f"{np.mean(betas):.3f}", f"{np.mean(1 - tvs):.3f}",
                     f"{np.max(np.abs(betas - (1 - tvs))):.1e}"])
    print(f"在 {toks.shape[0]} 条序列 × {toks.shape[1]} 个位置上逐位置计算（目标模型自己采样的文本）；"
          "三种采样设置：温度 1、温度 0.5、温度 1 + top-50（p 和 q 用同样的设置调整）")
    print_table(["草稿", "max|P_out − p|", "β 的平均（温度 1）", "1 − TV 的平均", "max|β − (1−TV)|"], rows)
    print_observations([
        "不管草稿多差（无关的小模型和目标几乎不重叠），按规则算出来的输出分布都和 p 一致到 1e-16：拒绝采样把 q 多给的部分"
        "退掉（min(1, p/q)），再用 p − q 的正部补上 q 少给的部分。",
        "每一步的接受概率 β = Σ min(p, q) 恰好等于 1 − TV(p, q)（论文定理 3.5）：草稿和目标的分布越接近，接受得越多。",
        "温度、top-k 这些采样设置只要同时作用在 p 和 q 上，结论不变（论文 §2.2 的\"标准化采样\"）。",
    ])


# -----------------------------------------------------------------------------
def demo2_monte_carlo():
    print_header("演示 2：蒙特卡洛检验——4 万个独立的投机解码（γ = 3，温度 1），前两个输出 token 的联合分布")
    cfg = Config(vocab=64, d_model=32, n_layers=2, n_heads=4, n_kv_heads=2, d_head=8, d_ff=96, max_seq=32,
                 name="micro")
    P = init_params(cfg, seed=3)
    tgt = TransformerLM(P, cfg, "目标")
    dp = dict(P)
    dp["layers"] = P["layers"][:1]
    drf = TransformerLM(dp, config_like(cfg, n_layers=1), "草稿")
    prompt = np.array([5, 17, 42, 8, 30, 1])
    p1 = adjust(tgt.full_logits(prompt[None])[0, -1], 1.0)
    q1 = adjust(drf.full_logits(prompt[None])[0, -1], 1.0)
    ext = np.concatenate([np.repeat(prompt[None], cfg.vocab, 0), np.arange(cfg.vocab)[:, None]], axis=1)
    p2 = adjust(tgt.full_logits(ext)[:, -1], 1.0)                       # p(x2 | prompt, x1)，对每个 x1
    joint = p1[:, None] * p2                                              # 精确的联合分布 [64, 64]
    N = 40000
    rng = np.random.default_rng(9)
    floor = [0.5 * np.abs(rng.multinomial(N, joint.ravel()) / N - joint.ravel()).sum() for _ in range(20)]
    print(f"极小模型：词表 {cfg.vocab}、2 层；草稿 = 只用第 1 层。第一个位置的接受概率 β = {np.minimum(p1, q1).sum():.3f}")
    rows = [["直接从精确分布抽 4 万个（噪声水平）", f"{np.mean(floor):.4f} ± {np.std(floor):.4f}", "-", "-"]]
    labels = {"correct": "投机解码（正确的规则）",
              "resample_p": "错误 1：被拒绝后直接从 p 重抽",
              "accept_if_p_ge_q": "错误 2：p(x) ≥ q(x) 才接受（不掷骰子）"}
    for var in ("correct", "resample_p", "accept_if_p_ge_q"):
        out = spec_sample_batch(tgt, drf, prompt, 2, 3, 1.0, N, np.random.default_rng(1), variant=var)
        emp = np.zeros((cfg.vocab, cfg.vocab))
        np.add.at(emp, (out[:, 0], out[:, 1]), 1.0)
        emp /= N
        tv = 0.5 * np.abs(emp - joint).sum()
        m1 = emp.sum(axis=1)
        worst = int(np.argmax(np.abs(m1 - p1)))
        rows.append([labels[var], f"{tv:.4f}", f"{0.5 * np.abs(m1 - p1).sum():.4f}",
                     f"token {worst}：{m1[worst]:.3f} vs {p1[worst]:.3f}"])
    print_table(["做法", "TV（联合分布，4096 格）", "TV（第一个 token）", "第一个 token 偏得最多的"], rows)
    print_observations([
        "正确的规则：联合分布的 TV 和\"直接从正确分布抽样\"的噪声水平相当（只高约 1 个标准差）——4 万个样本分辨不出它和目标模型的区别。",
        "错误 1（被拒绝后从 p 重抽）：q 偏爱的 token 先被接受了一部分，重抽时又按 p 给了一次，被\"多算\"了。",
        "错误 2（去掉随机性）：只要 q(x) > p(x) 就一律拒绝，q 偏爱的 token 几乎出不来，偏差最大。",
        "所以那个看起来很怪的 p − q 残差不是装饰：它正好补上\"q 给少了\"的那部分概率。",
    ])


# -----------------------------------------------------------------------------
def demo3_alpha(target, drafts, toks):
    print_header("演示 3：接受率 α = E[Σ min(p, q)] 和草稿、温度的关系")
    lp = target.full_logits(toks)
    rows = []
    for name, dr in drafts.items():
        lq = draft_logits(dr, toks)
        cells = [name]
        for T in (1.0, 0.7, 0.5, 0.0):
            cells.append(f"{np.minimum(adjust(lp, T), adjust(lq, T)).sum(-1).mean():.3f}")
        rows.append(cells)
    print("（温度 0 = 贪心：α 就是\"草稿和目标的第一名相同\"的比例。）")
    print_table(["草稿", "T = 1", "T = 0.7", "T = 0.5", "T = 0（贪心）"], rows)
    pr = adjust(lp, 1.0).reshape(-1, QCFG.vocab)
    print(f"\n目标模型在这些位置上的平均熵 {-(pr * np.log(pr)).sum(-1).mean():.2f} nat（均匀分布是 ln 1024 = 6.93），"
          f"第一名的平均概率只有 {pr.max(-1).mean():.3f}。")
    print_observations([
        "草稿越像目标，α 越高：INT8 量化的自己几乎全接受，INT4 约 0.9，跳层的 0.7–0.83，二元语法表约 0.34，无关的小模型 0.16。",
        "在这个随机初始化的模型里，**温度越低 α 越低**——和 Leviathan 表 3、EAGLE 表 2 在真实模型上看到的相反"
        "（它们是温度 0 时 α 更高）。原因：α(T=0) 是\"第一名相同\"的比例，α(T=1) 是 1 − TV。随机模型的分布很平"
        "（第一名只有约 0.1 的概率），量化 / 跳层的扰动很容易换掉第一名，但整体形状变化不大。",
        "真实语言模型相反：大部分位置有一个明显的\"容易的\"答案（标点、固定搭配、代码模板），大小模型都会选它；"
        "差别主要在尾部那些低概率的词。这时第一名的一致率高于 1 − TV。所以这里只看机制，不要把数字搬到真实模型上。",
    ])


# -----------------------------------------------------------------------------
def demo4_greedy(target, drafts):
    print_header("演示 4：贪心解码（温度 0）的特例——输出和目标模型自己贪心生成的完全相同")
    prompts = sample_tokens(target.params, QCFG, n_seq=4, length=16, seed=7)
    n_new = 64
    rows = []
    for name in ("INT4 量化的自己", "只用前 3 层", "二元语法表"):
        dr = drafts[name]
        same = 0
        tcalls = []
        for pr in prompts:
            ref, _ = autoregressive_generate(target, pr, n_new, 0, np.random.default_rng(0))
            out, st = speculative_generate(target, dr, pr, n_new, 4, 0, np.random.default_rng(0))
            same += int(ref == out)
            tcalls.append(st["target_calls"])
        rows.append([name, f"{same} / {len(prompts)}", f"{np.mean(tcalls):.1f}", f"{n_new / np.mean(tcalls):.2f}"])
    print(f"4 个提示，各生成 {n_new} 个 token，γ = 4；普通贪心解码要目标模型前向 {n_new} 次")
    print_table(["草稿", "输出完全相同的提示数", "目标模型前向次数（平均）", "每次前向产出的 token"], rows)
    print_observations([
        "贪心时 p、q 都是 one-hot：接受 ⇔ 草稿的 token 正好是目标的 argmax；被拒绝时 p − q 的正部就是目标的 argmax。"
        "所以输出逐 token 相同，这是最容易写单元测试的一种情况（真实引擎的测试也这么做）。",
        "INT4 和跳层草稿把目标模型的前向次数从 64 次降到 19–25 次（每次验证 5 个位置，平均留下 2.6–3.3 个）；"
        "二元语法表的第一名一致率只有 0.2，几乎没有帮助（53 次）。这些数和 04 脚本里的公式对得上。",
    ])


if __name__ == "__main__":
    TARGET, DRAFTS = build_models()
    TOKS = sample_tokens(TARGET.params, QCFG, n_seq=8, length=64, seed=0)
    demo1_exact(TARGET, DRAFTS, TOKS)
    demo2_monte_carlo()
    demo3_alpha(TARGET, DRAFTS, TOKS)
    demo4_greedy(TARGET, DRAFTS)