"""
=============================================================================
第 9 讲 · 工具模块：一个字符级 n-gram 语言模型 + canary / exposure 工具
=============================================================================

对应讲义：lectures/lecture09_privacy_attacks.md 第 9 节（05 号脚本导入本模块）

-----------------------------------------------------------------------------
【为什么用 n-gram，而不是 Transformer？ Why an n-gram model?】
-----------------------------------------------------------------------------
我们要演示的现象 —— 语言模型会逐字记住训练数据里的秘密、重复越多记得越牢、模型容量越大记得越多 ——
不依赖 Transformer 的任何细节。字符级 n-gram 模型是"记忆"最透明的语言模型：
    P(下一个字符 | 前面 n − 1 个字符) 直接由训练语料里的**计数**决定。
n（order，阶数）就是它的"容量"：n 越大，模型能区分的上下文越长、越具体，越接近"背下整句话"；
n 很小时，它只能学到"数字后面常跟连字符"这类统计规律，背不下任何具体的号码。
好处：训练 = 数数，几秒钟；同一份语料重复运行结果逐字相同；可以**精确地**枚举一个秘密的全部候选。

-----------------------------------------------------------------------------
【Witten–Bell 插值平滑 Witten–Bell interpolation】
-----------------------------------------------------------------------------
没见过的组合不能给概率 0（否则 log 概率是 −∞）。Witten–Bell 把长上下文和短上下文的估计混合起来：

    P(c | h) = [ C(h c) + T(h) · P(c | h') ] / [ C(h) + T(h) ]

    C(h c) = 上下文 h 后面跟着字符 c 的次数；C(h) = h 出现的总次数；
    T(h)   = h 后面出现过的**不同**字符的个数；h' = h 去掉最前面一个字符（更短的上下文）。
直觉：如果 h 后面出现过很多种不同的字符（T(h) 大），说明"还会出现新字符"的可能性大，就多借用短上下文的估计；
如果 h 后面总是同一个字符（比如 canary 的上下文），T(h) = 1，模型就几乎完全相信这个字符 —— 这正是"记住了"。
递归到空上下文时，用加一平滑（add-one）的单字符频率，保证每个字符的概率都 > 0。

【注释标签】同第 1 讲：【思路】【参数】【返回】【边界】【可信联系】【Python 小知识】
"""

import math

import numpy as np

NAMES = ["alice", "bob", "carol", "dave", "erin", "frank", "grace", "heidi", "ivan", "judy",
         "mallory", "nick", "olivia", "peggy", "rupert", "sybil", "trent", "victor", "walter", "yvonne"]
VERBS = ["paid", "sent", "lost", "found", "signed", "checked", "opened", "closed", "moved", "printed"]
OBJECTS = ["the bill", "the letter", "the form", "the card", "the report", "the key", "the box", "the file",
           "the ticket", "the note"]
DAYS = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]
THINGS = ["room", "code", "order", "ticket", "seat", "locker", "account", "pin"]
PIN_DIGITS = 4                                   # PIN 固定 4 位：候选空间 10⁴ = 10000 个
BOS, EOS = "^", "$"                              # 句首填充、句末标记
VOCAB = sorted(set("abcdefghijklmnopqrstuvwxyz0123456789 .-'") | {EOS})


def digit_string(digits):
    """[4, 8, 2, 9] → "4-8-2-9"（数字之间用连字符，这样每个数字都是一个独立的"字符"）。"""
    return "-".join(str(int(d)) for d in digits)


def make_corpus(n_sentences, rng):
    """
    生成一份合成语料：60% 是"某人在某天做了某事"，40% 是"我的某某号码是 x-x-x"。

    【思路】号码句子让模型学会"is 后面跟数字和连字符"的一般规律；其中 "my pin is" 的句子是**别人的 PIN**，
        每个只出现一次 —— 它们是语料里本来就有的"秘密"，04 演示看模型是否也记住了它们。
    【参数】n_sentences：句子数；rng：np.random.Generator
    【返回】(sentences, pins)：句子列表，以及语料中出现过的所有 PIN（每个是长度 4 的数字元组）
    【边界】随机生成的 PIN 可能偶然重复（10000 个候选里抽几百个），这里不去重 —— 真实语料也是这样。
    【Python 小知识】f"..." 里的 {表达式} 会被替换成表达式的值；int(rng.integers(3, 6)) 把 numpy 整数转成普通 int。
    """
    sentences, pins = [], []
    for _ in range(n_sentences):
        u = rng.random()
        if u < 0.6:
            sentences.append(f"{NAMES[rng.integers(len(NAMES))]} {VERBS[rng.integers(len(VERBS))]} "
                             f"{OBJECTS[rng.integers(len(OBJECTS))]} on {DAYS[rng.integers(len(DAYS))]}.")
            continue
        thing = THINGS[rng.integers(len(THINGS))]
        k = PIN_DIGITS if thing == "pin" else int(rng.integers(3, 6))
        digits = tuple(int(d) for d in rng.integers(0, 10, k))
        if thing == "pin":
            pins.append(digits)
        lead = "my" if u < 0.9 else f"{NAMES[rng.integers(len(NAMES))]} said my"
        sentences.append(f"{lead} {thing} is {digit_string(digits)}.")
    return sentences, pins


class CharNGram:
    """
    字符级 n-gram 语言模型（Witten–Bell 插值，见模块开头）。

    【用法】
        lm = CharNGram(order=7)
        for s in sentences: lm.add(s)
        lm.prob("my pin is ", "4")          # P("4" | 前文)
        lm.add(canary, weight=-8)            # 把之前加进去的 8 份 canary 再减掉（计数可以加减）
    【思路】对每个句子，在句首补 order − 1 个 "^"，句末加 "$"；对每个位置 i 和每个长度 k = 0..order−1，
        把"上下文 s[i−k:i] 后面跟 s[i]"计数一次。
    【边界】计数可以减（weight < 0）：用来在同一个基础语料上反复插入、删除 canary，不必每次重新数整份语料。
        减到 0 的条目会被删掉，保证 T(h)（不同后继字符数）正确。
    【Python 小知识】self.counts 是"字典套字典"：counts[上下文][字符] = 次数。d.get(c, 0) 在 c 不存在时返回 0。
    """

    def __init__(self, order):
        if order < 1:
            raise ValueError("order 至少是 1")
        self.order = order
        self.counts = {}                         # 上下文 → {下一个字符: 次数}
        self.totals = {}                         # 上下文 → 总次数 C(h)

    def add(self, sentence, weight=1):
        """把一句话的所有 n-gram 计数加 weight（可以是负数）。"""
        s = BOS * (self.order - 1) + sentence + EOS
        for i in range(self.order - 1, len(s)):
            ch = s[i]
            for k in range(self.order):
                h = s[i - k:i]
                d = self.counts.setdefault(h, {})
                new = d.get(ch, 0) + weight
                if new == 0:
                    d.pop(ch, None)
                else:
                    d[ch] = new
                tot = self.totals.get(h, 0) + weight
                if tot == 0:
                    self.totals.pop(h, None)
                    self.counts.pop(h, None)
                else:
                    self.totals[h] = tot

    def dist(self, context):
        """
        下一个字符的完整分布：返回长度 len(VOCAB) 的 numpy 数组（按 VOCAB 的顺序，和为 1）。

        【思路】从空上下文（add-one 的单字符频率）开始，逐步加长上下文，每一层按 Witten–Bell 公式混合。
            只用 context 的最后 order − 1 个字符；更长的部分 n-gram 模型看不见。
        【边界】context 比 order − 1 短时（比如提示词只有 3 个字符），就只用这么长的上下文 —— 这正是"已知前缀越短，
            模型越难回忆"的来源（05 演示 3）。
        """
        h = context[len(context) - (self.order - 1):] if self.order > 1 else ""
        base = self.counts.get("", {})
        p = np.array([base.get(c, 0) + 1.0 for c in VOCAB])
        p /= p.sum()
        for k in range(1, len(h) + 1):
            hk = h[-k:]
            d = self.counts.get(hk)
            if not d:
                continue                          # 没见过这个上下文：保持更短上下文的估计
            C, T = self.totals[hk], len(d)
            cnt = np.array([d.get(c, 0) for c in VOCAB], dtype=float)
            p = (cnt + T * p) / (C + T)
        return p

    def prob(self, context, ch):
        """
        P(ch | context)：只算一个字符的概率，公式和 dist 完全相同，但用标量运算 —— 比先算整个分布快很多。

        【思路】bits_per_char 和 score_all_pins 要调用它几万到几十万次，速度很重要；
            dist 适合"要看所有字符"的场合（解码、展开数字树的一层）。自检里会核对两者一致。
        """
        h = context[len(context) - (self.order - 1):] if self.order > 1 else ""
        base = self.counts.get("", {})
        p = (base.get(ch, 0) + 1.0) / (self.totals.get("", 0) + len(VOCAB))
        for k in range(1, len(h) + 1):
            hk = h[-k:]
            d = self.counts.get(hk)
            if not d:
                continue
            T = len(d)
            p = (d.get(ch, 0) + T * p) / (self.totals[hk] + T)
        return p

    def bits_per_char(self, sentences):
        """
        平均每个字符的负 log₂ 概率（bits per character）。越低 = 模型对这些句子越"不意外"。

        【可信联系】训练语料上的 bpc 远低于新句子上的 bpc = 过拟合 = 记忆（和第 1 讲的训练 / 测试 loss 差距是一回事）。
        """
        total, n = 0.0, 0
        for s in sentences:
            ctx = BOS * (self.order - 1)
            for ch in s + EOS:
                total -= math.log2(self.prob(ctx, ch))
                ctx += ch
                n += 1
        return total / n


def score_all_pins(lm, prefix):
    """
    对全部 10⁴ 个 4 位 PIN，计算 log P("d-d-d-d." | prefix)。

    【思路】按数字树（trie）逐层展开：第一层 10 个分支，第二层 100 个……每个节点只算一次下一字符的分布，
        共 1 + 10 + 100 + 1000 = 1111 次，而不是 10⁴ × 8 次。每一步要乘上"数字"的概率和后面"-"（或最后的"."）的概率。
    【参数】lm：CharNGram；prefix：提示词，例如 "my pin is "
    【返回】长度 10⁴ 的数组，下标 = PIN 的十进制值（例如 4-8-2-9 → 4829）
    【numpy 小知识】np.log 作用在整个数组上；np.repeat / np.tile 用来把"父节点的分数"展开到 10 个子节点。
    """
    digit_idx = [VOCAB.index(str(d)) for d in range(10)]
    scores = np.zeros(1)
    texts = [prefix]
    for pos in range(PIN_DIGITS):
        nxt_char = "." if pos == PIN_DIGITS - 1 else "-"
        new_scores, new_texts = [], []
        for sc, t in zip(scores, texts):
            p = lm.dist(t)                        # 这个节点的 10 个数字分支：一次算整个分布
            for d in range(10):
                t2 = t + str(d)
                q = lm.prob(t2, nxt_char)         # 数字后面那个固定的 "-" 或 "."：只算一个字符
                new_scores.append(sc + math.log(p[digit_idx[d]]) + math.log(q))
                new_texts.append(t2 + nxt_char)
        scores, texts = np.array(new_scores), new_texts
    return scores


def pin_index(digits):
    """(4, 8, 2, 9) → 4829。"""
    v = 0
    for d in digits:
        v = 10 * v + int(d)
    return v


def exposure_from_scores(scores, idx):
    """
    Exposure（Carlini et al. 2019，"The Secret Sharer"）：exposure = log₂|R| − log₂ rank。

    【思路】把真正的秘密放到全部候选 R 里，按模型给的概率排名（rank = 1 最可能）。
        · 模型完全没记住：秘密的排名是随机的，期望 exposure ≈ log₂ e ≈ 1.44 bit（噪声地板）。
        · 模型把秘密排第 1：exposure = log₂|R|（这里 log₂ 10⁴ ≈ 13.3 bit）—— 按概率从高到低猜，第一次就猜中。
        exposure 每多 1 bit，攻击者需要尝试的次数减半。
    【参数】scores：全部候选的 log 概率；idx：真实秘密的下标
    【返回】(exposure, rank)
    【边界】概率并列时，rank = "概率严格更高的候选个数 + 1"，即并列的候选都算排在秘密后面
        （对攻击者最有利的计法；n-gram 模型里完全没见过的候选经常并列）。
    """
    rank = 1 + int(np.sum(scores > scores[idx]))
    return math.log2(len(scores)) - math.log2(rank), rank


def greedy_decode(lm, prompt, n_chars):
    """
    贪心解码：每一步取概率最大的字符。【边界】概率并列时取 VOCAB 里靠前的字符（np.argmax 的规则），结果确定。
    """
    out = prompt
    for _ in range(n_chars):
        out += VOCAB[int(np.argmax(lm.dist(out)))]
    return out[len(prompt):]


def beam_decode(lm, prompt, n_chars, width):
    """
    束搜索（beam search）：每一步保留 log 概率之和最高的 width 个候选，最后返回最好的那个。

    【思路】贪心只看一步，可能第一步就走错（比如 canary 的第一个数字在"pin is "后面不是最常见的数字），
        而整条序列的概率其实是最高的。束搜索看得更远，更接近"找到整体概率最高的续写"。
    【Python 小知识】sorted(..., key=lambda b: -b[0])[:width] 按分数从高到低排序后取前 width 个。
    """
    beams = [(0.0, "")]
    for _ in range(n_chars):
        cand = []
        for sc, s in beams:
            p = lm.dist(prompt + s)
            for j, c in enumerate(VOCAB):
                cand.append((sc + math.log(p[j]), s + c))
        beams = sorted(cand, key=lambda b: (-b[0], b[1]))[:width]
    return beams[0][1]


if __name__ == "__main__":
    # 自检：分布和为 1；计数加减可以完全复原；exposure 的两个极端。
    rng = np.random.default_rng(0)
    sents, pins = make_corpus(500, rng)
    lm = CharNGram(5)
    for s in sents:
        lm.add(s)
    before = dict(lm.totals)
    lm.add("my pin is 1-2-3-4.", 3)
    lm.add("my pin is 1-2-3-4.", -3)
    print("自检 self-check")
    print(f"  分布之和 = {lm.dist('my pin is ').sum():.6f}（应为 1）")
    gap = max(abs(lm.prob("my pin is 4-", c) - lm.dist("my pin is 4-")[j]) for j, c in enumerate(VOCAB))
    print(f"  prob（标量）与 dist（向量）的最大差 = {gap:.1e}（应约为 0）")
    print(f"  插入再删除 canary 后计数完全复原：{before == lm.totals}")
    sc = score_all_pins(lm, "my pin is ")
    print(f"  全部 PIN 的概率之和 = {np.exp(sc).sum():.4f}（< 1：剩下的概率给了其他格式的续写）")
    print(f"  exposure(最可能的那个) = {exposure_from_scores(sc, int(np.argmax(sc)))[0]:.2f} bit"
          f"（应为 log₂ 10⁴ = {math.log2(1e4):.2f}）")