"""
=============================================================================
第 15 讲 · 工具模块：green-list 水印与 LLM 生成文本检测
=============================================================================

对应讲义：lectures/lecture15_llm_agent_security.md 第 10 节（代码 04 导入本模块）

-----------------------------------------------------------------------------
【这个模块做什么 What this module does】
-----------------------------------------------------------------------------
"这段文本是不是 LLM 生成的？"是 CS329T 2023 年第 6 讲列出的安全项目方向之一
（watermarking、DetectGPT、"Can AI-Generated Text be Reliably Detected?"）。
本模块实现 Kirchenbauer et al.（2023）的 **green-list 水印**，并用它说明：

    · 水印怎么嵌入（生成时偏向一个由上一个 token 决定的"绿名单"）、怎么检测（数绿 token，做 z 检验）；
    · 检测率随文本长度 T 怎么变（信号 ∝ √T）；
    · **释义攻击**（paraphrase：随机替换一部分 token，Sadasivan et al. 2023）怎么打断水印、检测率怎么掉；
    · 在**人类文本**（没有水印）上的误报率（false positive rate）—— z 检验的 I 类错误。

这是一个 toy LM（在 sec_world 的普通词上的 unigram 分布），不是真实模型；数字只说明机制。
"""

import hashlib
import math
import random

from sec_world import mean


_GREEN_CACHE = {}


def _green_set(prev, vocab, gamma, key):
    """
    由上一个 token `prev` 和密钥 `key` 决定的"绿名单"：vocab 的一个 γ 比例的随机子集。

    【思路】Kirchenbauer et al. 的做法：用前一个 token 做哈希种子，把词表伪随机地分成 green / red 两半
           （green 占 γ）。检测方只要知道 key 和规则，就能对任意文本重算同一个划分 —— 不需要模型、不需要概率。
    【参数】prev：上一个 token（整数）；vocab：候选 token 列表；gamma：绿名单比例；key：整数密钥
    【返回】一个 frozenset（不可修改的集合），装着绿名单里的 token
    【Python 小知识】hashlib + 固定字节串 → 跨平台、跨 Python 版本都一样的哈希（Python 内置 hash() 会随进程变化）。
           结果存进字典 _GREEN_CACHE（同一个 prev 的绿名单永远相同），避免重复计算 —— 这叫 memoization。
    """
    ck = (key, prev, gamma, tuple(vocab))
    if ck not in _GREEN_CACHE:
        h = hashlib.sha256(f"{key}:{prev}".encode()).digest()
        seed = int.from_bytes(h[:8], "big")
        rng = random.Random(seed)
        k = int(round(gamma * len(vocab)))
        _GREEN_CACHE[ck] = frozenset(rng.sample(vocab, k))
    return _GREEN_CACHE[ck]


class ToyLM:
    """
    一个 toy 语言模型：在一组 token 上的 unigram 分布（权重随机、固定）。可带 green-list 水印采样。

    【思路】真实水印作用在 next-token 分布上。这里的 unigram 分布就扮演"模型给出的 logits"：
           水印把绿名单 token 的 logit 加上 δ，再 softmax 采样。δ = 0 等于没有水印。
    【参数】vocab：token 列表；seed：决定 unigram 权重的种子
    """

    def __init__(self, vocab, seed):
        rng = random.Random(seed)
        self.vocab = vocab
        self.logit = {t: rng.gauss(0.0, 1.0) for t in vocab}

    def generate(self, rng, length, gamma=0.5, delta=0.0, key=7, prev=0):
        """
        采样 `length` 个 token。delta > 0 时启用 green-list 水印。

        【参数】rng：random.Random；length：长度 T；gamma：绿名单比例；delta：绿名单 logit 加成；
               key：水印密钥；prev：起始的"上一个 token"
        【返回】token 列表
        【边界】delta 很大时几乎只会从绿名单采样，文本质量会变差（讲义 §7：水印强度 vs 文本质量的权衡）。
        """
        out = []
        for _ in range(length):
            green = _green_set(prev, self.vocab, gamma, key)
            ws = []
            for t in self.vocab:
                z = self.logit[t] + (delta if t in green else 0.0)
                ws.append(math.exp(z))
            tot = math.fsum(ws)
            u = rng.random() * tot
            acc = 0.0
            chosen = self.vocab[-1]
            for t, w in zip(self.vocab, ws):
                acc += w
                if u < acc:
                    chosen = t
                    break
            out.append(chosen)
            prev = chosen
        return out


def green_fraction(tokens, vocab, gamma, key, prev0=0):
    """
    一段文本里"绿 token"的个数和总数（检测的原始统计量）。

    【思路】对每个位置，用它的**前一个 token** 重算绿名单，看当前 token 在不在里面。
           检测方不需要模型，只要会算 _green_set。
    【返回】(绿 token 数, 计分的 token 总数)
    """
    g, prev = 0, prev0
    for t in tokens:
        if t in _green_set(prev, vocab, gamma, key):
            g += 1
        prev = t
    return g, len(tokens)


def z_score(green, total, gamma):
    """
    绿 token 比例的 z 统计量：z = (green − γT) / √(T γ(1−γ))。

    【思路】零假设 = "文本和水印无关"，那么每个 token 落进绿名单的概率是 γ，绿 token 数 ~ Binomial(T, γ)。
           水印文本的绿比例会显著高于 γ，z 很大。阈值通常取 z > 4（对应极小的误报率）。
    【边界】T = 0 时返回 0。
    """
    if total == 0:
        return 0.0
    return (green - gamma * total) / math.sqrt(total * gamma * (1 - gamma))


def paraphrase(tokens, rng, vocab, rate):
    """
    释义攻击的玩具版：以概率 `rate` 把每个 token 替换成词表里随机的一个 token。

    【思路】真实的 paraphrase 用另一个模型改写（Sadasivan et al. 2023, "Can AI-Generated Text be
           Reliably Detected?"）。关键机制是：green-list 水印依赖"前一个 token → 绿名单"的链条，
           替换一个 token 会同时弄错**它自己**和**它后面那个**的绿名单判定，所以检测信号掉得比替换率还快。
    【参数】rate：每个 token 被替换的概率（substitution rate）
    【返回】新的 token 列表（长度不变）
    """
    return [rng.choice(vocab) if rng.random() < rate else t for t in tokens]


def detection_rate(lm, vocab, rng, n, length, gamma, delta, key, z_thr=4.0, para_rate=0.0):
    """
    检测率：生成 n 段水印文本（可选释义），z > z_thr 的比例。

    【参数】lm；n：样本数；length：每段长度；para_rate：释义替换率；z_thr：判为"有水印"的阈值
    【返回】(检测率, 平均 z)
    """
    hit, zs = 0, []
    for _ in range(n):
        toks = lm.generate(rng, length, gamma=gamma, delta=delta, key=key)
        if para_rate > 0:
            toks = paraphrase(toks, rng, vocab, para_rate)
        z = z_score(*green_fraction(toks, vocab, gamma, key), gamma)
        zs.append(z)
        hit += 1 if z > z_thr else 0
    return hit / n, mean(zs)


def false_positive_rate(lm, vocab, rng, n, length, gamma, key, z_thr=4.0):
    """
    误报率：对 **没有水印** 的文本（delta = 0），z > z_thr 的比例（应当接近 z 检验的名义水平）。

    【可信联系】第 1 讲 §8 的"噪声地板"：任何检测器都要报告它在"没有信号"的数据上的表现。
    """
    hit = 0
    for _ in range(n):
        toks = lm.generate(rng, length, gamma=gamma, delta=0.0, key=key)
        z = z_score(*green_fraction(toks, vocab, gamma, key), gamma)
        hit += 1 if z > z_thr else 0
    return hit / n


if __name__ == "__main__":
    vocab = list(range(60))
    lm = ToyLM(vocab, seed=1)
    rng = random.Random(0)
    print("detection @ T=200, delta=2:", detection_rate(lm, vocab, rng, 200, 200, 0.5, 2.0, 7))
    print("FPR @ T=200:", false_positive_rate(lm, vocab, rng, 1000, 200, 0.5, 7))
    print("after 30% paraphrase:", detection_rate(lm, vocab, rng, 200, 200, 0.5, 2.0, 7, para_rate=0.3))