"""
risk_lab.py —— 第 22 讲的共享模块：在小模型上度量"生成式 AI 会出什么错"
=====================================================================

【这个文件要回答的问题】
    slides 用同一个框架讨论每一种风险：它是什么意思？影响谁？为什么会发生？怎样修？
    "为什么会发生"和"怎样修"都可以在小模型上做受控实验——因为训练数据是我们自己造的，
    我们**确切知道**模型见过什么、见过几次，于是能把"模型的行为"和"数据里有什么"分开。

    本模块提供五个玩具世界和一套小语言模型工具：
        1. 记忆世界（01、02 号脚本）：随机马尔可夫链生成的"背景文本"、重复了 r 次的"书"、
           "用户 u 的编号是 d d d d d"这样的金丝雀（canary）句子——度量逐字记忆、可提取性、成员推断，
           以及 Vyas et al. (2023) 的"近访问自由"（near access-freeness）版权保护。
        2. 事实世界（03 号脚本）：每个人住在一个城市，不同的人在训练数据里出现的次数不同——
           度量幻觉和出现次数的关系、置信度、拒答、检索增强（RAG）。
        3. 偏见世界（04 号脚本）：两个抽象的群体"甲""乙"和 8 种抽象的职业，训练数据里职业和群体的共现比例是我们设定的——
           度量模型怎样复现、放大这些比例，以及几种缓解方法。
        4. 碳排放的账（05 号脚本）：能耗 = GPU 小时 × 功率 × PUE，排放 = 能耗 × 电网碳强度；在合成的"每小时碳强度"曲线上
           模拟 Dodge et al. (2022) 的"灵活开始"和"暂停-继续"。
        5. 水印（06 号脚本，本讲义补充）：Kirchenbauer et al. (2023) 的绿名单水印——生成时给"绿词"加分，检测时数绿词。

【本文件提供】
    小语言模型工具：Vocab、build_gpt（第 2 讲的 GPT）、pad_batch、train_corpus、get_model（训练并缓存）、
        token_logprobs、next_logprobs、greedy_continue
    记忆世界：MarkovText、make_mem_world、mem_corpus、code_logprobs_all、exposure、lcs_len、min_k_score、auc、
        cp_combine
    事实世界：make_fact_world、fact_corpus、fact_prompt
    偏见世界：make_bias_world、bias_corpus、cda_swap
    碳排放：energy_mwh、emissions_t、carbon_curve、flexible_start、pause_and_resume
    水印：green_ids、watermark_sample、green_count、z_score

【边界】
    - 所有"模型"都只有几万到几十万参数，数据是人造的；结论是"机制上会发生什么"，不是大模型上的定量结果。
    - 偏见世界里的群体和职业都是**抽象的符号**（甲、乙、job0…job7），比例是我们随手设定的，不对应任何真实人群。
"""

import hashlib
import math
import os
import random
import sys

import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _sub in ("lecture02", "lecture01", "lecture03"):
    _p = os.path.join(HERE, "..", _sub)
    if _p not in sys.path:
        sys.path.insert(0, _p)

from lm_models import GPT  # noqa: E402
from kv_cache import forward_cached, reorder_cache  # noqa: E402

CKPT = os.path.join(HERE, "checkpoints")
SPECIALS = ["<pad>", "<s>", "</s>"]
PAD, BOS, EOS = 0, 1, 2
IGNORE = -100
DIGITS = [str(i) for i in range(10)]


# =====================================================================
# 1. 小语言模型工具
# =====================================================================

class Vocab:
    """
    词表：把词（字符串）和整数 id 互相转换。前三个 id 固定是 <pad>、<s>、</s>。

    【参数】words：除特殊记号以外的全部词（顺序决定 id，所以要用确定的顺序）
    【边界】encode 遇到不在词表里的词会抛出 KeyError——玩具世界里不应该出现未知词，出现了就是 bug。
    """

    def __init__(self, words):
        self.itos = SPECIALS + list(words)
        self.stoi = {w: i for i, w in enumerate(self.itos)}
        if len(self.stoi) != len(self.itos):
            raise ValueError("词表里有重复的词")

    def __len__(self):
        return len(self.itos)

    def encode(self, tokens, bos=True, eos=True):
        ids = [self.stoi[t] for t in tokens]
        return ([BOS] if bos else []) + ids + ([EOS] if eos else [])

    def decode(self, ids):
        return [self.itos[i] for i in ids if i not in (PAD,)]


def build_gpt(n_vocab, d_model=64, n_layers=2, n_heads=4, max_len=64, seed=0):
    """
    第 2 讲的 GPT（可学习的位置嵌入、pre-norm）。

    【边界】构造前先设随机种子：本机上每个新进程的默认种子不同，不设就每次初始化都不一样（第 13 讲踩过的坑）。
    """
    torch.manual_seed(seed)
    return GPT(n_vocab, d_model=d_model, n_layers=n_layers, n_heads=n_heads, max_len=max_len)


def pad_batch(seqs, loss_start=None):
    """
    把不等长的 id 序列补齐成一批 (inputs, targets)。

    【思路】语言模型在位置 j 读 seq[j]、预测 seq[j+1]；所以 inputs = seq[:-1]、targets = seq[1:]，补齐的位置 targets = IGNORE。
    【参数】seqs：id 列表的列表；loss_start：可选，每个序列"从第几个 token 开始计损失"（例如只在答案上计损失）
    【返回】inputs (B, T)、targets (B, T)，T = 最长序列长度 − 1
    【边界】loss_start[i] = k 表示只预测 seq[k:]（targets 的前 k − 1 个位置设成 IGNORE）。
    """
    T = max(len(s) for s in seqs) - 1
    inp = torch.full((len(seqs), T), PAD, dtype=torch.long)
    tgt = torch.full((len(seqs), T), IGNORE, dtype=torch.long)
    for i, s in enumerate(seqs):
        n = len(s) - 1
        inp[i, :n] = torch.tensor(s[:-1])
        tgt[i, :n] = torch.tensor(s[1:])
        if loss_start is not None and loss_start[i] > 1:
            tgt[i, :loss_start[i] - 1] = IGNORE
    return inp, tgt


def train_corpus(model, corpus, epochs, batch_size=64, lr=1e-3, seed=0, warmup=100, max_steps=None, log=None):
    """
    在一个固定的训练集上训练若干个 epoch（每个 epoch 把训练集打乱一次、完整过一遍）。

    【思路】"一条数据在训练集里出现 r 次"是本讲反复用到的自变量；固定训练集 + 多个 epoch 正好对应
        "预训练语料里有重复"这件事：出现 r 次的句子，模型一共会看到 r × epochs 次。
    【参数】
        corpus    : id 序列的列表；也可以是 (ids, loss_start) 二元组的列表（只在 loss_start 之后计损失）
        epochs    : 遍数；batch_size、lr（Adam）、warmup（线性预热的步数）
        seed      : 打乱顺序用的种子（第 e 遍用 random.Random(seed·1000 + e)）
        max_steps : 只训练这么多步就停（冒烟测试用）
    【返回】每个 epoch 的平均训练损失列表
    【边界】单线程 + 固定种子时结果逐位可复现。
    【长度分桶】训练集里有长有短（编号句 13 个 token、书 50 个 token）。一批里只要混进一本书，整批都要补齐到 50——
        白算很多补齐的位置。所以打乱以后每 32 批为一组、组内按长度排序再切成批、再打乱这些批的顺序：
        每一批里的序列长度相近，计算量少了约 40%，而数据的随机性基本不变（真实的训练框架也这样做）。
    """
    items = [c if isinstance(c, tuple) else (c, 0) for c in corpus]
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    model.train()
    step, history = 0, []
    for epoch in range(1, epochs + 1):
        rng = random.Random(seed * 1000 + epoch)
        order = list(range(len(items)))
        rng.shuffle(order)
        batches = []
        window = 32 * batch_size
        for w in range(0, len(order), window):
            # 【Python 小知识】sorted(..., key=函数)：按函数值排序；Python 的排序是稳定的，长度相同的保持原来的（打乱后的）顺序
            group = sorted(order[w:w + window], key=lambda i: len(items[i][0]))
            part = [group[b:b + batch_size] for b in range(0, len(group), batch_size)]
            rng.shuffle(part)
            batches += part
        total, count = 0.0, 0
        for idx in batches:
            step += 1
            for g in opt.param_groups:
                g["lr"] = lr * min(1.0, step / max(1, warmup))
            chunk = [items[i] for i in idx]
            inp, tgt = pad_batch([c[0] for c in chunk], [c[1] for c in chunk])
            logits, _ = model(inp)
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1), ignore_index=IGNORE)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total += loss.item()
            count += 1
            if max_steps is not None and step >= max_steps:
                break
        history.append(total / max(1, count))
        if log:
            log(f"    epoch {epoch}  平均训练损失 {history[-1]:.4f}")
        if max_steps is not None and step >= max_steps:
            break
    model.eval()
    return history


def get_model(name, cfg, n_vocab, corpus_fn, quick=False, log=print):
    """
    有缓存就读取 checkpoints/<name>.pt，没有就训练并保存。

    【参数】
        cfg       : {"d_model", "n_layers", "n_heads", "max_len", "seed", "epochs", "lr", "batch_size", ...}；
                    存进缓存，读取时不一致就重新训练（防止"拿旧模型配新设置"）
        corpus_fn : 无参函数，返回训练集（只在需要训练时才调用，读缓存时省下造数据的时间）
        quick     : 冒烟测试——只训练 30 步，不读也不写缓存
    【返回】训练好的模型（eval 模式）
    【边界】训练集的内容由 cfg 里的种子和大小决定；改了造数据的代码却没改 cfg，缓存不会自动失效——
        所以 cfg 里还存了训练集的指纹（所有序列的 SHA-256 前 16 位）。读缓存时也要算一遍指纹（几秒钟）。
    """
    corpus = corpus_fn()
    key = dict(cfg)
    key["corpus_hash"] = hashlib.sha256(repr(corpus).encode()).hexdigest()[:16]
    path = os.path.join(CKPT, f"{name}.pt")
    build = lambda: build_gpt(n_vocab, cfg["d_model"], cfg["n_layers"], cfg["n_heads"], cfg["max_len"],  # noqa: E731
                              cfg["seed"])
    if not quick and os.path.exists(path):
        blob = torch.load(path, weights_only=True)
        if blob["cfg"] == key:
            model = build()
            model.load_state_dict(blob["state"])
            model.eval()
            if log:
                log(f"  （读取缓存的模型：checkpoints/{name}.pt）")
            return model
    model = build()
    train_corpus(model, corpus, cfg["epochs"], batch_size=cfg.get("batch_size", 64), lr=cfg.get("lr", 1e-3),
                 seed=cfg["seed"], warmup=cfg.get("warmup", 100), max_steps=30 if quick else None)
    if not quick:
        os.makedirs(CKPT, exist_ok=True)
        torch.save({"cfg": key, "state": model.state_dict()}, path)
        if log:
            log(f"  （已训练并保存到 checkpoints/{name}.pt）")
    return model


@torch.no_grad()
def token_logprobs(model, seqs, batch_size=256):
    """
    每个序列里每个 token 的条件对数概率 log p(s_j | s_<j)（j = 1 … n−1，不含 <s> 本身）。

    【返回】列表，第 i 个元素是长度 len(seqs[i]) − 1 的一维张量（自然对数）
    【边界】序列长度不能超过模型的 max_len。
    """
    out = []
    for b in range(0, len(seqs), batch_size):
        chunk = seqs[b:b + batch_size]
        inp, tgt = pad_batch(chunk)
        logp = F.log_softmax(model(inp)[0].double(), dim=-1)
        for i, s in enumerate(chunk):
            n = len(s) - 1
            out.append(logp[i, torch.arange(n), tgt[i, :n]])
    return out


@torch.no_grad()
def next_logprobs(model, prefixes):
    """
    一批**等长**前缀之后下一个 token 的对数概率分布。

    【返回】(B, V) 的 float64 张量
    【边界】前缀必须等长（不做左侧补齐）；本讲的提示都是固定格式，天然等长。
    """
    ids = torch.tensor(prefixes, dtype=torch.long)
    return F.log_softmax(model(ids)[0][:, -1].double(), dim=-1)


@torch.no_grad()
def greedy_continue(model, prefixes, n_new):
    """
    从一批等长的前缀开始贪心地续写 n_new 个 token（每一步取概率最大的 token）。

    【返回】(B, n_new) 的 id 张量
    【边界】每一步把整个序列重新喂给模型（序列很短，不需要 KV 缓存）；总长度不能超过 max_len。
    """
    ids = torch.tensor(prefixes, dtype=torch.long)
    for _ in range(n_new):
        nxt = model(ids)[0][:, -1].argmax(-1, keepdim=True)
        ids = torch.cat([ids, nxt], dim=1)
    return ids[:, -n_new:]


# =====================================================================
# 2. 记忆世界：背景文本、书、金丝雀
# =====================================================================

MEM_N_WORDS = 160          # 背景文本的词表：w000 … w159
MEM_N_USERS = 400          # 用户：u000 … u399
CODE_LEN = 5               # 编号是 5 位数字：一共 10⁵ 种可能
CODE_WORDS = ["the", "code", "of", "user", "is", "."]
R_VALUES = (1, 2, 4, 8, 16, 32, 64)


def mem_vocab():
    """记忆世界的词表：10 个数字、6 个功能词、400 个用户、160 个背景词。"""
    return Vocab(DIGITS + CODE_WORDS + [f"u{i:03d}" for i in range(MEM_N_USERS)]
                 + [f"w{i:03d}" for i in range(MEM_N_WORDS)])


class MarkovText:
    """
    一条随机的"稀疏马尔可夫链"：每个词后面只能接 6 个词之一，概率各不相同。用它生成的文本有规律（模型能学），
    但又不可完全预测（每一步平均约 2 比特的不确定性）——像一门风格固定、内容随机的语言。

    【参数】seed：决定转移表；n_words：词表大小；fanout：每个词的后继个数
    【边界】转移表在构造时就固定了；sample 只用传进来的 rng，不碰全局随机数。
    """

    def __init__(self, seed=0, n_words=MEM_N_WORDS, fanout=6):
        rng = random.Random(seed)
        self.words = [f"w{i:03d}" for i in range(n_words)]
        self.succ, self.prob = [], []
        for _ in range(n_words):
            nxt = rng.sample(range(n_words), fanout)
            w = sorted((rng.random() ** 2 for _ in range(fanout)), reverse=True)
            s = sum(w)
            self.succ.append(nxt)
            self.prob.append([x / s for x in w])

    def sample(self, rng, length):
        """从随机起点出发走 length 步，返回词列表。"""
        i = rng.randrange(len(self.words))
        out = [self.words[i]]
        for _ in range(length - 1):
            i = rng.choices(self.succ[i], self.prob[i])[0]
            out.append(self.words[i])
        return out

    def entropy_bits(self):
        """每一步的平均条件熵（比特；对起点均匀平均，只是一个粗略的"有多难预测"）。"""
        return sum(-sum(p * math.log2(p) for p in ps) for ps in self.prob) / len(self.prob)


def code_sentence(user, digits):
    """金丝雀句子：the code of user u017 is 4 9 2 0 6 ."""
    return ["the", "code", "of", "user", user, "is"] + list(digits) + ["."]


def make_mem_world(seed=0, n_background=12000, n_fresh=6000, n_per_r=8, n_nonmember=48, book_len=48,
                   bg_len=24):
    """
    造一个记忆世界（只是"原材料"，还没有决定谁进训练集、进几次——那是 mem_corpus 的事）。

    【内容】
        background : n_background 句背景文本（每句 bg_len 个词）
        fresh      : n_fresh 句"普通"编号句：用户从 u100 … u399 里随机挑、5 位数字每次都重新随机——
                     模型由此学会格式，并且知道"一般来说编号是随机的"
        canaries   : {r: [(用户, 数字串), ...]}，每个 r 有 n_per_r 条；用户是 u000 … u099 里互不相同的
        canary_out : n_per_r 条"从没出现过"的金丝雀（r = 0，作对照）
        books      : {r: [词列表, ...]}，每本"书"是同一条马尔可夫链上的 book_len 个词
        books_out  : n_nonmember 本从没进过训练集的书（成员推断的"非成员"、续写的对照）
    【返回】字典
    【边界】金丝雀用户只用 u000 … u099，普通编号句只用 u100 … u399，两者不重叠：
        金丝雀用户的编号**只**可能来自金丝雀本身。
    """
    rng = random.Random(seed)
    chain = MarkovText(seed)
    users = [f"u{i:03d}" for i in range(100)]
    rng.shuffle(users)
    rand_code = lambda: "".join(rng.choice(DIGITS) for _ in range(CODE_LEN))  # noqa: E731
    canaries, k = {}, 0
    for r in R_VALUES:
        canaries[r] = [(users[k + j], rand_code()) for j in range(n_per_r)]
        k += n_per_r
    canary_out = [(users[k + j], rand_code()) for j in range(n_per_r)]
    books = {r: [chain.sample(rng, book_len) for _ in range(n_per_r)] for r in R_VALUES}
    books_out = [chain.sample(rng, book_len) for _ in range(n_nonmember)]
    background = [chain.sample(rng, bg_len) for _ in range(n_background)]
    fresh = [(f"u{rng.randrange(100, MEM_N_USERS):03d}", rand_code()) for _ in range(n_fresh)]
    return {"chain": chain, "canaries": canaries, "canary_out": canary_out, "books": books,
            "books_out": books_out, "background": background, "fresh": fresh}


def mem_corpus(world, vocab, shard=None, dedup=False):
    """
    把记忆世界变成训练集：背景文本、普通编号句各一次；金丝雀和书按各自的 r 重复 r 次。

    【参数】
        shard : None = 全部数据；0 或 1 = 只要"一半"（02 号脚本的两份不相交的数据：背景和普通句按奇偶分，
                每个 r 的金丝雀和书也按奇偶分到两边）
        dedup : True = 去重——重复的句子只保留一份（金丝雀和书都变成只出现 1 次）
    【返回】id 序列的列表（顺序无所谓，训练时会打乱）
    """
    take = (lambda i: True) if shard is None else (lambda i: i % 2 == shard)  # noqa: E731
    out = [vocab.encode(s) for i, s in enumerate(world["background"]) if take(i)]
    out += [vocab.encode(code_sentence(u, d)) for i, (u, d) in enumerate(world["fresh"]) if take(i)]
    for r in R_VALUES:
        reps = 1 if dedup else r
        for i, (u, d) in enumerate(world["canaries"][r]):
            if take(i):
                out += [vocab.encode(code_sentence(u, d))] * reps
        for i, book in enumerate(world["books"][r]):
            if take(i):
                out += [vocab.encode(book)] * reps
    return out


@torch.no_grad()
def code_logprobs_all(model, vocab, user):
    """
    在提示 "the code of user <user> is" 之后，**全部** 10⁵ 个 5 位编号的对数概率（精确枚举，不是采样估计）。

    【思路】按数字一位一位地展开一棵树：第 1 层只有提示本身，给出第 1 位数字的 10 种概率；第 2 层把 10 种第 1 位
        分别接上去……第 5 层有 10⁴ 个前缀。log p(d₁…d₅) = Σ log p(dᵢ | 前面)。
        每一层只需要算**新加的那一个 token**——前面的部分用第 3 讲的 KV 缓存（forward_cached）复用，
        再用 reorder_cache 把每个前缀的缓存复制 10 份。比把 10⁴ 个完整序列各算一遍快约 10 倍。
    【返回】长度 10⁵ 的 float64 张量；下标就是编号本身（"04920" → 4920）
    【边界】只对数字 token 求概率：模型也可能在这里输出非数字的词，那部分概率不属于任何编号，
        所以 10⁵ 个概率之和可能略小于 1。
    """
    prefix = vocab.encode(["the", "code", "of", "user", user, "is"], eos=False)
    d0 = vocab.stoi["0"]
    logits, cache = forward_cached(model, torch.tensor([prefix]))
    total = torch.zeros(1, dtype=torch.float64)
    for level in range(CODE_LEN):
        lp = F.log_softmax(logits[:, -1].double(), dim=-1)[:, d0:d0 + 10]       # (前缀个数, 10)
        # 【Python 小知识】unsqueeze(1) 把 (B,) 变成 (B, 1)，和 (B, 10) 相加时自动"广播"成 (B, 10)；
        #   reshape(-1) 按行展开：第 b 个前缀接数字 d 的结果排在 b·10 + d——正好是"多一位数字"的编号
        total = (total.unsqueeze(1) + lp).reshape(-1)
        if level == CODE_LEN - 1:
            break
        n = total.shape[0]
        cache = reorder_cache(cache, torch.arange(n // 10).repeat_interleave(10))
        logits, cache = forward_cached(model, (torch.arange(n) % 10 + d0).unsqueeze(1), cache)
    return total


def exposure(logp_all, secret):
    """
    Carlini et al. (2019) 的"暴露度"：exposure = log₂|R| − log₂ rank(secret)。

    【参数】logp_all：全部候选的对数概率（code_logprobs_all 的输出）；secret：真正的编号（字符串）
    【返回】(exposure 比特数, 排名)；排名 1 = 模型最喜欢的就是这个秘密
    【边界】排名 = 1 + 严格比它概率大的候选个数（并列算在它后面）。完全没记住时 exposure 的期望约 1.44 比特，
        完全记住（排第 1）时等于 log₂10⁵ ≈ 16.6 比特。
    """
    target = logp_all[int(secret)]
    rank = 1 + int((logp_all > target).sum())
    return math.log2(len(logp_all)) - math.log2(rank), rank


def lcs_len(a, b):
    """
    最长公共子序列（longest common subsequence）的长度：动态规划，O(len(a)·len(b))。

    【边界】子序列不要求连续（"a x b" 和 "a b" 的 LCS 是 2）；Karamolegkou et al. (2023) 用它度量
        "续写的内容和原书有多像"，比"逐字完全相同"宽松。
    """
    prev = [0] * (len(b) + 1)
    for x in a:
        cur = [0]
        for j, y in enumerate(b):
            cur.append(prev[j] + 1 if x == y else max(prev[j + 1], cur[j]))
        prev = cur
    return prev[-1]


def min_k_score(lp, k=0.2):
    """
    Shi et al. (2024) 的 Min-K% Prob：一段文本里最"意外"的 k 比例 token 的平均对数概率。

    【思路】没见过的文本里总有几个"意外"的词（概率很低）；见过的文本连最意外的词也被记住了。
        只看最低的那部分，比看全部 token 的平均（就是负的损失）更敏感。
    【参数】lp：一维张量，每个 token 的对数概率；k：取最低的比例
    【返回】浮点数（越大越像"见过"）
    """
    n = max(1, math.ceil(k * len(lp)))
    return torch.sort(lp)[0][:n].mean().item()


def auc(pos, neg):
    """
    ROC 曲线下面积 = 随机取一个正例和一个负例，正例得分更高的概率（并列算 1/2）。

    【返回】0.5 = 完全分不开，1.0 = 完全分开
    【边界】用整数计数（每对 2 分 / 1 分 / 0 分）再除，避免浮点求和顺序带来的差异。
    """
    score = 0
    for p in pos:
        for q in neg:
            score += 2 if p > q else (1 if p == q else 0)
    return score / (2 * len(pos) * len(neg))


def cp_combine(lp_a, lp_b, kind="min"):
    """
    Vyas, Kakade & Barak (2023) 的 CP-Δ：把在**不相交**数据上训练的两个模型合成一个"近访问自由"的模型。

    【思路】受版权保护的作品 C 至多出现在一份数据里，所以至少有一个模型（"安全模型" safe_C）从没见过它。
        min 版：p(y) ∝ min(q_a(y), q_b(y))。于是 p(y) ≤ q_i(y)/Z 对两个模型都成立，
            即 p 和**每一个** q_i 的最大散度不超过 k = log₂(1/Z) 比特：p(y) ≤ 2^k · safe_C(y)。
        sqrt 版：p(y) ∝ √(q_a(y)·q_b(y))。KL(p‖q_a) + KL(p‖q_b) = −2 ln Z，所以每个 KL ≤ −2 ln Z。
    【参数】lp_a、lp_b：两个模型在**同一组候选**上的对数概率（float64）；kind："min" 或 "sqrt"
    【返回】(合成后的对数概率, 保证值)：min 版返回 k（比特），sqrt 版返回 KL 上界（纳特）
    【边界】只在候选集合上归一化（这里是全部 10⁵ 个编号，是精确的）；真实模型要在全部输出序列上做，
        论文讨论了怎样逐 token 地近似。
    """
    m = torch.minimum(lp_a, lp_b) if kind == "min" else (lp_a + lp_b) / 2
    log_z = torch.logsumexp(m, dim=0)
    bound = -log_z.item() / math.log(2) if kind == "min" else -2 * log_z.item()
    return m - log_z, bound


# =====================================================================
# 3. 事实世界：谁住在哪个城市
# =====================================================================

FACT_N_PEOPLE = 1000       # p000 … p999：有"事实"的人
FACT_N_READERS = 200       # r000 … r199：只在检索文档里出现的人（城市每次随机——只能读文档）
FACT_N_CITIES = 20         # c00 … c19
IDK = "?"                  # "我不知道"


def fact_vocab():
    """事实世界的词表。"""
    return Vocab(["doc", "q", ".", IDK] + [f"c{i:02d}" for i in range(FACT_N_CITIES)]
                 + [f"p{i:03d}" for i in range(FACT_N_PEOPLE)] + [f"r{i:03d}" for i in range(FACT_N_READERS)])


def make_fact_world(seed=0, n_samples=3000, zipf=1.0):
    """
    造一个"名人有多出名"服从齐普夫分布的事实世界，并从中抽出训练语料。

    【思路】第 i 个人被提到的概率 µ_i ∝ 1/(i+1)^zipf（少数人很出名，大多数人很少被提到）。训练语料是
        从 µ 里独立抽 n_samples 次——和 Kalai et al. (2025) 定理的设定一样：出名的人被提到很多次，
        很多人只被提到 1 次（"单例"），更多的人一次也没被提到。
    【返回】字典：
        city[p]  : 这个人住的城市（均匀随机——没见过的人不可能靠规律猜对，只能碰运气 1/20）
        mu[p]    : 被提到的真实概率
        count[p] : 在训练语料里出现了几次
        people   : 全部人（按编号排序）
    【边界】人的编号是打乱后才分配"出名程度"的，所以编号本身不携带信息。
    """
    rng = random.Random(seed)
    people = [f"p{i:03d}" for i in range(FACT_N_PEOPLE)]
    order = people[:]
    rng.shuffle(order)
    w = [1.0 / (i + 1) ** zipf for i in range(FACT_N_PEOPLE)]
    total = math.fsum(w)
    mu = {p: w[i] / total for i, p in enumerate(order)}
    city = {p: f"c{rng.randrange(FACT_N_CITIES):02d}" for p in people}
    draws = rng.choices(order, [mu[p] for p in order], k=n_samples)
    count = {p: 0 for p in people}
    for p in draws:
        count[p] += 1
    return {"city": city, "mu": mu, "count": count, "people": people, "n_samples": n_samples}


def singleton_rate(world):
    """
    Kalai et al. (2025) 定义 2：只出现过一次的"提示"（这里是人）的个数 ÷ 训练样本数 N。
    【GenAI 联系】它也是图灵的"缺失质量"估计：再抽一次，抽到一个"从没见过的人"的概率约等于它。
    """
    return sum(1 for p in world["people"] if world["count"][p] == 1) / world["n_samples"]


def fact_prompt(person, docs=()):
    """
    提示："doc p017 c03 . doc p250 c11 . q p017"——前面是零到多篇检索到的文档（"doc 人 城市" 读作"某人住在某城市"），
    最后是问题"q 人"（"某人住在哪里？"），模型接着写出城市。
    【参数】docs：(人, 城市) 列表；为空就是"闭卷"
    【为什么格式这么紧凑】名字和城市紧挨着，"在文档里找到同一个名字、抄出它后面的词"就是第 11 讲的归纳头（induction head）
        能做的事；写成 "p017 lives in c03" 时城市离名字 3 个位置，第一版的小模型学不会按名字查找（README 有记录）。
    """
    toks = []
    for who, c in docs:
        toks += ["doc", who, c, "."]
    return toks + ["q", person]


def fact_corpus(world, vocab, n_rag=280000, p_unanswerable=0.0, seed=1):
    """
    训练集 = 闭卷的事实句（每人按 count 重复）+ n_rag 条"带检索文档"的问答。

    【思路】
        闭卷句 "q p017 c03 ."：模型只能把事实记在参数里。
        检索问答："doc r005 c11 . doc r120 c02 . q r005 c11 ."——问的和文档里的都是"读者" r***，
        他们的城市**每次随机**，所以只能从文档里抄。这样 p*** 的事实只来自闭卷句，"出现几次"是干净的自变量；
        模型从读者身上学到的是"按名字在文档里找、再把城市抄出来"这个通用的本领。
        p_unanswerable > 0 时，这个比例的检索问答里两篇文档都**不是**关于被问的人，答案是 "?"（"文档里没有"）。
    【返回】(ids, loss_start) 列表：只在答案、"." 和 </s> 上计损失
    【边界】闭卷部分从不出现 "?"——模型从没见过"闭卷时说不知道"的例子（真实的预训练语料也很少有）。
    """
    rng = random.Random(seed)
    out = []
    for p in world["people"]:
        for _ in range(world["count"][p]):
            toks = fact_prompt(p) + [world["city"][p], "."]
            out.append((vocab.encode(toks), len(toks) - 1))
    for _ in range(n_rag):
        a, b, c3 = rng.sample(range(FACT_N_READERS), 3)
        docs = [(f"r{a:03d}", f"c{rng.randrange(FACT_N_CITIES):02d}"),
                (f"r{b:03d}", f"c{rng.randrange(FACT_N_CITIES):02d}")]
        if rng.random() < p_unanswerable:
            who, ans = f"r{c3:03d}", IDK
        else:
            who, ans = docs[0]
        rng.shuffle(docs)
        toks = fact_prompt(who, docs) + [ans, "."]
        out.append((vocab.encode(toks), len(toks) - 1))
    return out


# =====================================================================
# 4. 偏见世界：抽象的群体"甲""乙"和 8 种抽象的职业
# =====================================================================

GROUPS = ("甲", "乙")
N_JOBS = 8
# 训练数据里"职业 k 的人来自甲组"的比例。数字是随手设定的，只为覆盖从 0.1 到 0.9 的各种偏斜；不对应任何真实人群。
JOB_SKEW = (0.50, 0.60, 0.70, 0.80, 0.90, 0.40, 0.25, 0.10)
LANGS = {"甲": "L1", "乙": "L2"}   # 一条"合理的"群体知识：甲组说 L1、乙组说 L2（在这个世界里是事实，不是偏见）


def bias_vocab():
    """偏见世界的词表。"""
    return Vocab(["the", "is", "from", "a", "person", "works", "as", "uses", "people", "speak", "."] + list(GROUPS)
                 + [f"job{k}" for k in range(N_JOBS)] + [f"tool{k}" for k in range(N_JOBS)] + sorted(set(LANGS.values())))


def bias_sentences(seed=0, n_per_job=300, n_tool=150, n_lang=400):
    """
    偏见世界的全部训练句子（词列表的列表），以及每句的"类型"。

    【内容】
        职业句（两种写法各一半）："the job3 is from 甲 ." / "a 甲 person works as job3 ."——群体按 JOB_SKEW 抽
        工具句："the job3 uses tool3 ."——和群体无关的职业知识（让职业的嵌入不只是"群体"）
        语言句："甲 people speak L1 ."——一条和群体相关、但在这个世界里**正确**的知识
    【返回】[(类型, 词列表), ...]，类型是 "job" / "tool" / "lang"
    """
    rng = random.Random(seed)
    out = []
    for k in range(N_JOBS):
        for i in range(n_per_job):
            g = GROUPS[0] if rng.random() < JOB_SKEW[k] else GROUPS[1]
            toks = ["the", f"job{k}", "is", "from", g, "."] if i % 2 == 0 else ["a", g, "person", "works", "as", f"job{k}", "."]
            out.append(("job", toks))
        out += [("tool", ["the", f"job{k}", "uses", f"tool{k}", "."])] * n_tool
    for i in range(n_lang):
        g = GROUPS[i % 2]
        out.append(("lang", [g, "people", "speak", LANGS[g], "."]))
    return out


def cda_swap(toks):
    """反事实数据增强（counterfactual data augmentation）的"交换"：把句子里的甲换成乙、乙换成甲。"""
    sw = {GROUPS[0]: GROUPS[1], GROUPS[1]: GROUPS[0]}
    return [sw.get(t, t) for t in toks]


def bias_corpus(vocab, cda=None, seed=0):
    """
    【参数】cda：None = 原始数据；"full" = 每一句都加一份甲乙互换的副本（包括语言句）；
        "targeted" = 只给职业句加互换的副本（需要人来决定"哪些群体关联是偏见、哪些是事实"）
    【返回】id 序列的列表
    """
    out = []
    for kind, toks in bias_sentences(seed):
        out.append(vocab.encode(toks))
        if cda == "full" or (cda == "targeted" and kind == "job"):
            out.append(vocab.encode(cda_swap(toks)))
    return out


# =====================================================================
# 5. 碳排放的账
# =====================================================================

def energy_mwh(gpu_hours, watts, pue=1.0):
    """能耗（MWh）= GPU 小时 × 每块 GPU 的功率（W）× PUE ÷ 10⁶。PUE（电能使用效率）= 数据中心总耗电 ÷ IT 设备耗电 ≥ 1。"""
    return gpu_hours * watts * pue / 1e6


def emissions_t(mwh, kg_per_kwh):
    """排放（吨 CO₂ 当量）= 能耗（MWh）× 电网碳强度（kg/kWh）：1 MWh × 1 kg/kWh = 1000 kg = 1 吨。"""
    return mwh * kg_per_kwh


def carbon_curve(hours, seed=0, base=0.40, solar=0.15, noise=0.04):
    """
    一条**合成的**电网每小时碳强度曲线（kg CO₂/kWh），只用来演示调度算法；不是任何真实地区的数据。

    【思路】白天太阳能多 → 碳强度低（中午最低），夜里高；再加上一个缓慢变化的随机"天气"项（AR(1) 过程）。
    【返回】长度 hours 的列表
    """
    rng = random.Random(seed)
    out, w = [], 0.0
    for h in range(hours):
        w = 0.9 * w + rng.gauss(0.0, noise)
        out.append(max(0.05, base - solar * math.cos(2 * math.pi * ((h % 24) - 12) / 24) + w))
    return out


def flexible_start(curve, start, duration, window):
    """
    Dodge et al. (2022) 的"灵活开始"（Flexible Start）：允许把开始时间推迟最多 window 小时，选总排放最小的开始时间。
    【返回】(最好的开始时间, 这时作业期间碳强度之和)
    【边界】作业一旦开始就不能停；curve 必须足够长（start + window + duration）。
    """
    best = min(range(start, start + window + 1), key=lambda s: (sum(curve[s:s + duration]), s))
    return best, sum(curve[best:best + duration])


def pause_and_resume(curve, start, duration, extra):
    """
    "暂停-继续"（Pause and Resume）：作业可以在 duration + extra 小时的窗口里任意暂停，
    只在碳强度最低的 duration 个小时里运行。
    【返回】这 duration 个小时的碳强度之和
    【边界】假设暂停和继续没有额外开销（真实训练要存检查点、重新排队）。
    """
    window = curve[start:start + duration + extra]
    return sum(sorted(window)[:duration])


# =====================================================================
# 6. 水印（Kirchenbauer et al., 2023）
# =====================================================================

def green_ids(prev_id, vocab_size, gamma, key=15485863):
    """
    "绿名单"：用上一个 token 的 id 和一个密钥做种子，把词表随机分成绿 / 红两部分，绿的占 gamma。

    【思路】种子只依赖上一个 token，所以检测的人只要知道密钥，就能对任何一段文本重新算出每个位置的绿名单——
        不需要模型本身，也不需要生成时的随机数。
    【返回】布尔张量 (vocab_size,)，True = 绿
    【边界】这里用 torch.Generator 的 randperm 做"哈希"；真实实现用更快的哈希函数，原理相同。
    """
    g = torch.Generator().manual_seed(key * 1_000_003 + int(prev_id))
    perm = torch.randperm(vocab_size, generator=g)
    mask = torch.zeros(vocab_size, dtype=torch.bool)
    mask[perm[:int(round(gamma * vocab_size))]] = True
    return mask


def green_count(ids, vocab_size, gamma, skip=(), key=15485863):
    """
    数一段 token 序列里有几个"绿词"。
    【参数】ids：整数列表（第一个通常是 <s>，只作为第二个 token 的"上一个"）；skip：不计入的 token id（例如 </s>）
    【返回】(绿词个数, 计入的 token 数)
    """
    hits, total = 0, 0
    for prev, cur in zip(ids[:-1], ids[1:]):
        if cur in skip:
            continue
        total += 1
        hits += int(green_ids(prev, vocab_size, gamma, key)[cur])
    return hits, total


def z_score(hits, total, gamma):
    """
    假设检验：如果文本**没有**水印，每个 token 是绿的概率是 gamma，绿词个数 ~ 二项分布(total, gamma)。
        z = (hits − gamma·T) / √(T·gamma·(1−gamma))
    z 很大（例如 > 4，单侧 p ≈ 3×10⁻⁵）就判定"有水印"。
    """
    if total == 0:
        return 0.0
    return (hits - gamma * total) / math.sqrt(total * gamma * (1 - gamma))


def green_table(vocab_size, gamma, key=15485863):
    """所有"上一个 token"的绿名单一次算好：(V, V) 的布尔矩阵，第 i 行 = green_ids(i, ...)。词表很小时这样最快。"""
    return torch.stack([green_ids(i, vocab_size, gamma, key) for i in range(vocab_size)])


@torch.no_grad()
def watermark_sample(model, n, bos_id, eos_id, banned, generator, delta=2.0, table=None, max_len=40):
    """
    批量生成 n 句话（都从 <s> 开始），每一步先给绿名单里的 token 的 logit 加 delta，再按 softmax 采样。

    【参数】
        banned    : 永远不生成的 token id（<pad>、<s>）
        generator : torch.Generator（固定种子，结果可复现）
        delta     : 绿词的加分；0 = 不加水印（普通采样）
        table     : green_table 的结果；delta = 0 时可以是 None
    【返回】(句子列表（每句是 id 列表，不含 <s> 和 </s>）, 每个生成位置的 (原始分布的熵（比特）, 是否绿词) 列表)
    【边界】达到 max_len 还没有 </s> 的句子直接截断；每一步把整个前缀重新喂给模型（句子很短）。
    """
    ids = torch.full((n, 1), bos_id, dtype=torch.long)
    done = torch.zeros(n, dtype=torch.bool)
    out, stats = [[] for _ in range(n)], []
    for _ in range(max_len):
        logits = model(ids)[0][:, -1].double()
        logits[:, list(banned)] = float("-inf")
        p0 = F.softmax(logits, dim=-1)
        ent = -(p0 * torch.log2(p0.clamp_min(1e-300))).sum(-1)
        if delta:
            logits = logits + delta * table[ids[:, -1]].double()
        nxt = torch.multinomial(F.softmax(logits, dim=-1), 1, generator=generator).squeeze(1)
        for b in range(n):
            if not done[b]:
                if table is not None:
                    stats.append((ent[b].item(), bool(table[ids[b, -1], nxt[b]])))
                if nxt[b].item() == eos_id:
                    done[b] = True
                else:
                    out[b].append(nxt[b].item())
        if done.all():
            break
        ids = torch.cat([ids, nxt.unsqueeze(1)], dim=1)
    return out, stats