"""
infer_world.py —— 第 19 讲：高效推理的公用工具——KV 缓存、在线 softmax、分页缓存、量化、推测解码
=====================================================================

【这个文件做什么】
    这一讲关心"生成一个 token 为什么慢、怎样在不改变输出的前提下变快"。所有实验都用花园世界上训练的小模型：
        target_lm()   "大"模型：d = 128、4 层（约 80 万参数），推测解码里的目标模型；
        draft_lm()    "小"模型：d = 32、1 层（约 1.6 万参数），推测解码里的草稿模型；
    以及几样从零实现的东西：
        CachedDecoder        带 KV 缓存的逐 token 解码（和第 2 讲 TransformerLM 共用同一套参数）；
        online_softmax / tiled_attention   FlashAttention 的核心：分块 + 在线 softmax；
        quantize             权重的对称量化（逐张量 / 逐通道 / 分组）；
        speculative_step     推测解码的一步（Leviathan et al. 2023 / Chen et al. 2023）。

【LLM 联系】
    本仓库的 CMU11-868（LLM 系统）有这些主题的系统实现细节（GPU kernel、vLLM、GPTQ）；这一讲讲"为什么"和
    "保证了什么"，并在花园世界里把"输出分布严格不变"之类的结论精确地验证一遍。
"""

import math
import os
import sys

import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import clm_batches, clm_eval, clm_loss, ext_vocab, load_or_train, train  # noqa: E402
from toy_world import standard_splits  # noqa: E402
from transformer import TransformerLM, sinusoidal_positions  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")
VOCAB = ext_vocab()                      # 第 2 讲的扩展词表（花园世界的 90 个 token + <pad> 等）
SPLITS = standard_splits()
BOS, EOS = VOCAB.stoi["<bos>"], VOCAB.stoi["<eos>"]

# =====================================================================
# 1. 两个模型
# =====================================================================

TARGET_ARCH = dict(d_model=128, n_heads=4, n_layers=4, d_ff=512)
DRAFT_ARCH = dict(d_model=32, n_heads=2, n_layers=1, d_ff=128)


def _lm(name, arch, epochs, seed=0):
    batches = clm_batches(SPLITS["train"], VOCAB, 32)

    def build():
        torch.manual_seed(seed)
        return TransformerLM(len(VOCAB), attn="causal", **arch)

    m, _ = load_or_train(name, {"arch": arch, "epochs": epochs, "seed": seed, "v": 1}, build,
                         lambda m: train(m, batches, clm_loss, epochs=epochs, lr=3e-3, warmup=200, seed=seed,
                                         verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


def target_lm():
    """目标模型：d = 128、4 层，在第 1 讲的 8000 篇花园文档上训练 6 遍。"""
    return _lm("target", TARGET_ARCH, 6)


def draft_lm(epochs=6):
    """草稿模型：d = 32、1 层。epochs 越少越"差"（推测解码里用来比较不同质量的草稿）。"""
    return _lm(f"draft_e{epochs}", DRAFT_ARCH, epochs)


def test_nll(model, n=500):
    """测试集上每个 token 的平均 NLL（第 1 讲的口径）。"""
    return clm_eval(model, SPLITS["test"][:n], VOCAB)[0]


# =====================================================================
# 2. KV 缓存（01 号脚本）
# =====================================================================


class CachedDecoder:
    """
    带 KV 缓存的逐 token 解码，参数直接取自一个（pre-norm、正弦位置编码、权重共享的）TransformerLM。

    【思路】因果注意力里，位置 t 的 key / value 只依赖 t 及之前的 token，生成后面的 token 时它们不会变——
            所以每层把算过的 k、v 存下来，新 token 只需要算自己的 q、k、v，再和缓存里的全部 k、v 做注意力。
            每步的计算量从 O(t·d²)（重算整个前缀）降到 O(d² + t·d)。
    【用法】dec = CachedDecoder(model); logits = dec.prefill(ids); logits = dec.step(next_ids) ...
    【边界】只支持 batch 内所有序列等长（不处理补齐）；缓存存在 Python 列表里，每步 torch.cat（简单但不高效：
            真实系统预先分配或分页，见 03 号脚本）。
    """

    def __init__(self, model):
        self.m = model
        self.d = model.d
        self.cache = None
        self.t = 0

    def _block_step(self, blk, x, layer):
        att = blk.self_attn
        h = blk.n1(x)
        B, T, _ = h.shape
        q, k, v = (att._split(f(h)) for f in (att.q_proj, att.k_proj, att.v_proj))   # (B, h, T, hd)
        if self.cache[layer] is not None:
            k = torch.cat([self.cache[layer][0], k], 2)
            v = torch.cat([self.cache[layer][1], v], 2)
        self.cache[layer] = (k, v)
        S = q @ k.transpose(-1, -2) / math.sqrt(att.hd)
        Tk = k.shape[2]
        # 新的 T 个 query 是位置 Tk−T .. Tk−1：每个只能看自己及之前的 key
        allowed = torch.arange(Tk)[None, :] <= (torch.arange(T)[:, None] + Tk - T)
        S = S.masked_fill(~allowed, float("-inf"))
        out = (torch.softmax(S, -1) @ v).transpose(1, 2).reshape(B, T, -1)
        x = x + att.o_proj(out)
        return x + blk.ffn(blk.n3(x))

    @torch.no_grad()
    def _forward(self, ids):
        T = ids.shape[1]
        pe = sinusoidal_positions(self.t + T, self.d)[self.t:]
        x = self.m.emb(ids) * math.sqrt(self.d) + pe
        for li, blk in enumerate(self.m.blocks):
            x = self._block_step(blk, x, li)
        self.t += T
        return self.m.norm_f(x) @ self.m.emb.weight.T

    def prefill(self, ids):
        """处理整段提示（一次前向），建立缓存。【返回】每个位置的 logits (B, T, V)。"""
        self.cache = [None] * len(self.m.blocks)
        self.t = 0
        return self._forward(ids)

    def step(self, ids):
        """处理新的 token（(B, 1) 或 (B, k)），追加缓存。【返回】这些位置的 logits。"""
        return self._forward(ids)

    def cache_bytes(self):
        return sum(k.numel() * k.element_size() + v.numel() * v.element_size() for k, v in self.cache)


@torch.no_grad()
def generate_no_cache(model, ids, n_new, seed=0):
    """不用缓存：每生成一个 token，都把整个前缀重新跑一遍。【返回】(B, T + n_new) 的 token。"""
    g = torch.Generator().manual_seed(seed)
    for _ in range(n_new):
        p = torch.softmax(model(ids)[:, -1], -1)
        ids = torch.cat([ids, torch.multinomial(p, 1, generator=g)], 1)
    return ids


@torch.no_grad()
def generate_cached(model, ids, n_new, seed=0):
    """用 KV 缓存：先 prefill，之后每步只算新 token。采样的随机数和 generate_no_cache 一样，结果应当相同。"""
    g = torch.Generator().manual_seed(seed)
    dec = CachedDecoder(model)
    logits = dec.prefill(ids)[:, -1]
    for _ in range(n_new):
        nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=g)
        ids = torch.cat([ids, nxt], 1)
        logits = dec.step(nxt)[:, -1]
    return ids


# =====================================================================
# 3. 在线 softmax 与分块注意力（02 号脚本）
# =====================================================================


def online_softmax_stats(scores_blocks):
    """
    一遍扫描地计算 softmax 的归一化常数（Milakov & Gimelshein 2018）：逐块维护"目前的最大值 m"和"目前的和 l"。
        新块到来：m' = max(m, max(块))；l' = l · exp(m − m') + Σ exp(块 − m')
    【返回】(m, l)：softmax(x)_i = exp(x_i − m) / l。不需要先扫一遍求最大值、再扫一遍求和。
    """
    m, l = -math.inf, 0.0
    for blk in scores_blocks:
        m_new = max(m, float(np.max(blk)))
        l = l * math.exp(m - m_new) + float(np.exp(blk - m_new).sum())
        m = m_new
    return m, l


def tiled_attention(q, k, v, block=16, causal=True):
    """
    FlashAttention 的前向（Dao et al. 2022）的 numpy 版本：按 key 分块，对每个 query 维护 (m, l, acc)。
        每来一块 key / value：s = q kᵀ/√d；m' = max(m, max s)；
        acc ← acc · exp(m − m') + exp(s − m') v；l ← l · exp(m − m') + Σ exp(s − m')
        最后输出 acc / l。
    【参数】q、k、v：(T, d) numpy 数组；block：每块的 key 数。
    【返回】(T, d) 输出——和"先算完整的 T×T 注意力矩阵"在浮点误差内相同，但任何时候只需要 T×block 的分数。
    【边界】真实的 kernel 还对 query 分块，并把块放在 GPU 的片上 SRAM 里；这里只演示数学上的等价性。
    """
    T, d = q.shape
    m = np.full(T, -np.inf)
    l = np.zeros(T)
    acc = np.zeros((T, v.shape[1]))
    for s0 in range(0, T, block):
        s = q @ k[s0:s0 + block].T / math.sqrt(d)                      # (T, block)
        if causal:
            j = np.arange(s0, min(s0 + block, T))
            s = np.where(j[None, :] <= np.arange(T)[:, None], s, -np.inf)
        m_new = np.maximum(m, s.max(1))
        scale = np.exp(np.where(np.isfinite(m), m - m_new, -np.inf))
        p = np.exp(s - m_new[:, None])
        acc = acc * scale[:, None] + p @ v[s0:s0 + block]
        l = l * scale + p.sum(1)
        m = m_new
    return acc / l[:, None]


def reference_attention(q, k, v, causal=True):
    T, d = q.shape
    s = q @ k.T / math.sqrt(d)
    if causal:
        s = np.where(np.tril(np.ones((T, T), bool)), s, -np.inf)
    s = s - s.max(1, keepdims=True)
    p = np.exp(s)
    return (p / p.sum(1, keepdims=True)) @ v


# =====================================================================
# 4. 量化（04 号脚本）
# =====================================================================


def quantize(w, bits, scheme="channel", group=32):
    """
    对称（absmax）量化后再反量化：q = round(w / s)，s = max|w| / (2^(bits−1) − 1)，返回 q · s。

    【参数】scheme："tensor"（整个矩阵一个 s）、"channel"（每个输出通道一行一个 s）、"group"（每行每 group 个一个 s）。
    【返回】和 w 同形状的张量（浮点，已经带上了量化误差）。
    【边界】只量化 2 维的权重矩阵；一个极大的值会让整组的 s 变大、其余的值都被"挤"到少数几个格子里——
            这就是离群值（outlier）的问题，也是分组越细越好的原因。
    """
    qmax = 2 ** (bits - 1) - 1
    if scheme == "tensor":
        s = w.abs().max().clamp(min=1e-8) / qmax
        return (w / s).round().clamp(-qmax, qmax) * s
    if scheme == "channel":
        s = w.abs().amax(1, keepdim=True).clamp(min=1e-8) / qmax
        return (w / s).round().clamp(-qmax, qmax) * s
    out, cols = w.clone(), w.shape[1]
    for c0 in range(0, cols, group):
        blk = w[:, c0:c0 + group]
        s = blk.abs().amax(1, keepdim=True).clamp(min=1e-8) / qmax
        out[:, c0:c0 + group] = (blk / s).round().clamp(-qmax, qmax) * s
    return out


def quantized_copy(model, bits, scheme="channel", group=32):
    """把模型里所有 nn.Linear 的权重量化（嵌入 / 输出层、LayerNorm、偏置保持浮点——常见的做法）。"""
    import copy
    m = copy.deepcopy(model)
    with torch.no_grad():
        for mod in m.modules():
            if isinstance(mod, torch.nn.Linear):
                mod.weight.copy_(quantize(mod.weight, bits, scheme, group))
    return m


def _ln_linears(model):
    """每个块里"LayerNorm → 若干 Linear"的对应关系：(n1 → q/k/v)，(n3 → ffn.up)。"""
    out = []
    for blk in model.blocks:
        out.append((blk.n1, [blk.self_attn.q_proj, blk.self_attn.k_proj, blk.self_attn.v_proj]))
        out.append((blk.n3, [blk.ffn.up]))
    return out


def inject_outlier(model, channel=5, factor=30.0):
    """
    制造一个"离群通道"，而**不改变模型的函数**：把每个 LayerNorm 第 channel 维的缩放和偏置乘以 factor，
    同时把后面 Linear 权重的第 channel 列除以 factor。Linear 的输出一模一样，但它的输入里这一维大了 factor 倍。
    【LLM 联系】Dettmers et al. (2022, LLM.int8()) 发现 6.7B 以上的模型里会"涌现"出少数几个特别大的激活维度
                （比其他维度大几十倍），正是它们让 8 比特的激活量化失效。我们的小模型没有，所以人为造一个。
    """
    import copy
    m = copy.deepcopy(model)
    with torch.no_grad():
        for ln, lins in _ln_linears(m):
            ln.weight[channel] *= factor
            ln.bias[channel] *= factor
            for lin in lins:
                lin.weight[:, channel] /= factor
    return m


class ActQuant:
    """
    给 LayerNorm 后面的 Linear 的**输入**（激活）做 bits 比特的对称量化（forward pre-hook）。

    【参数】scheme："tensor"（整批一个缩放因子）或 "token"（每个 token 一个，动态量化）；
            keep：不量化的通道集合（LLM.int8() 的混合精度分解：离群通道留在高精度）。
    【用法】with ActQuant(model, 8, "token"): nll = test_nll(model)
    """

    def __init__(self, model, bits=8, scheme="token", keep=()):
        self.model, self.bits, self.scheme, self.keep = model, bits, scheme, list(keep)
        self.handles = []

    def _hook(self, mod, inp):
        x = inp[0]
        qmax = 2 ** (self.bits - 1) - 1
        xq = x.clone()
        if self.keep:
            xq[..., self.keep] = 0
        if self.scheme == "tensor":
            s = xq.abs().max().clamp(min=1e-8) / qmax
        else:
            s = xq.abs().amax(-1, keepdim=True).clamp(min=1e-8) / qmax
        out = (xq / s).round().clamp(-qmax, qmax) * s
        if self.keep:
            out[..., self.keep] = x[..., self.keep]
        return (out,) + tuple(inp[1:])

    def __enter__(self):
        for _, lins in _ln_linears(self.model):
            for lin in lins:
                self.handles.append(lin.register_forward_pre_hook(self._hook))
        return self

    def __exit__(self, *a):
        for h in self.handles:
            h.remove()


@torch.no_grad()
def activation_absmax(model, docs):
    """每个"LayerNorm → Linear"位置、每个通道的激活绝对值最大值（校准数据上）。【返回】list[(LN, Linear 们, 最大值向量)]。"""
    stats = {}
    hs = []
    for i, (ln, lins) in enumerate(_ln_linears(model)):
        def hook(mod, inp, out, i=i):
            m = out.detach().abs().reshape(-1, out.shape[-1]).amax(0)
            stats[i] = torch.maximum(stats[i], m) if i in stats else m
        hs.append(ln.register_forward_hook(hook))
    for b in clm_batches(docs, VOCAB, 64):
        model(b[0][:, :-1], pad=b[1][:, :-1])
    for h in hs:
        h.remove()
    return [(ln, lins, stats[i]) for i, (ln, lins) in enumerate(_ln_linears(model))]


def smooth(model, calib_docs, alpha=0.5):
    """
    SmoothQuant（Xiao et al. 2023）：把激活的"难量化"迁移一部分给权重。每个通道 j 选 s_j = max|X_j|^α / max|W_j|^(1−α)，
    激活除以 s_j（改 LayerNorm 的缩放）、权重的第 j 列乘以 s_j——函数不变，激活的离群通道被压平。
    """
    import copy
    m = copy.deepcopy(model)
    with torch.no_grad():
        for ln, lins, amax in activation_absmax(m, calib_docs):
            wmax = torch.stack([lin.weight.abs().amax(0) for lin in lins]).amax(0)
            sj = (amax.clamp(min=1e-5) ** alpha / wmax.clamp(min=1e-5) ** (1 - alpha))
            ln.weight /= sj
            ln.bias /= sj
            for lin in lins:
                lin.weight *= sj
    return m


# =====================================================================
# 5. 推测解码（05 号脚本）
# =====================================================================


@torch.no_grad()
def next_dist(model, ids):
    """(B, T) → 每个位置的下一个 token 分布 (B, T, V)。"""
    return torch.softmax(model(ids), -1)


@torch.no_grad()
def speculative_generate(target, draft, n_tokens, gamma=4, seed=0, stop_at_eos=True):
    """
    推测解码（单个序列）：草稿模型先连续猜 gamma 个 token，目标模型**一次前向**同时检查它们。
        对第 i 个猜测 x：以概率 min(1, p(x)/q(x)) 接受（p 是目标模型、q 是草稿模型的分布）；
        第一次拒绝时，从修正分布 norm(max(0, p − q)) 里重新采一个 token，丢掉后面的猜测；
        全部接受时，再从目标模型的分布里多采一个 token（白送的）。
    【保证】每个 token 的分布严格等于目标模型的分布：P(输出 x) = min(p, q)(x) + max(0, p − q)(x) = p(x)。
    【返回】(token 列表, 目标模型的前向次数, 每轮接受的猜测数列表)。
    【边界】为了清楚，这里不用 KV 缓存（每次都把整段前缀交给模型），所以"前向次数"而不是"时间"是要看的成本。
    """
    g = torch.Generator().manual_seed(seed)
    ids = [BOS]
    calls, accepted = 0, []
    while len(ids) - 1 < n_tokens:
        # 1. 草稿模型连续猜 gamma 个
        guess, qs = [], []
        ctx = list(ids)
        for _ in range(gamma):
            q = next_dist(draft, torch.tensor([ctx]))[0, -1]
            x = int(torch.multinomial(q, 1, generator=g))
            guess.append(x)
            qs.append(q)
            ctx.append(x)
        # 2. 目标模型一次前向：得到每个猜测位置（以及最后多一个位置）的分布
        P = next_dist(target, torch.tensor([ids + guess]))[0, len(ids) - 1:]
        calls += 1
        n_acc = 0
        for i, x in enumerate(guess):
            p, q = P[i], qs[i]
            if float(torch.rand(1, generator=g)) < min(1.0, float(p[x] / q[x])):
                ids.append(x)
                n_acc += 1
                if stop_at_eos and x == EOS:
                    break
            else:
                resid = (p - q).clamp(min=0)
                ids.append(int(torch.multinomial(resid / resid.sum(), 1, generator=g)))
                break
        else:
            ids.append(int(torch.multinomial(P[gamma], 1, generator=g)))
        accepted.append(n_acc)
        if stop_at_eos and ids[-1] == EOS:
            break
    return ids[1:], calls, accepted


@torch.no_grad()
def sample_docs(model, n, seed=0, max_len=60):
    """
    直接从模型成批采样 n 篇文档（到 <eos> 为止）。已经结束的文档移出批次，不再参与前向。
    【返回】token id 列表的列表（含结尾的 <eos>）。
    """
    g = torch.Generator().manual_seed(seed)
    ids = torch.full((n, 1), BOS)
    out = [[] for _ in range(n)]
    active = list(range(n))
    for _ in range(max_len):
        x = torch.multinomial(next_dist(model, ids)[:, -1], 1, generator=g)
        keep = []
        for row, (j, t) in enumerate(zip(active, x[:, 0].tolist())):
            out[j].append(t)
            if t != EOS:
                keep.append(row)
        ids = torch.cat([ids, x], 1)[keep]
        active = [active[r] for r in keep]
        if not active:
            break
    return out


__all__ = ["VOCAB", "SPLITS", "BOS", "EOS", "TARGET_ARCH", "DRAFT_ARCH", "target_lm", "draft_lm", "test_nll",
           "CachedDecoder", "generate_no_cache", "generate_cached", "online_softmax_stats", "tiled_attention",
           "reference_attention", "quantize", "quantized_copy", "inject_outlier", "ActQuant", "activation_absmax", "smooth",
           "next_dist", "speculative_generate", "sample_docs"]