"""
=============================================================================
CS324 · 第 5 讲模块：icl_world.py —— 一个可以精确算出"贝叶斯最优答案"的上下文学习世界
=============================================================================

【这个模块是干什么的】
    Xie et al. 2021 把 in-context learning（ICL）解释成**隐式的贝叶斯推断**：
        预训练文档是这样生成的——先抽一个隐藏的"概念"（concept，可以理解为"这篇文档在做什么任务"），
        再按这个概念生成整篇文档。要把下一个 token 预测好，模型就得从前文**推断**当前文档的概念。
        提示里的示例（demonstrations）就是推断概念的证据。
    真实的互联网里我们写不出"概念"的精确分布，所以这里造一个小世界，在里面**所有概率都能精确算出来**：

        一篇文档（16 对"输入 标签"）：   b1 a0 b1 d0 a0 h1 b1 ...      （每对 = 一个输入字母 + 一个标签数字）

        - 输入：8 个字母 a–h（相当于"句子"）。
        - 主题（topic）：4 个。主题 t 有两个"常用字母"（a/b、c/d、e/f、g/h），各占 0.3 的概率，
          其余 6 个字母共占 0.4——**输入的分布**带着主题的信息（就像影评和新闻用词不同）。
        - 标签空间：主题 t 只用两个标签 {t, t+1 mod 4}，例如主题 0 用 {0, 1}，主题 3 用 {3, 0}。
          相邻主题共享一个标签，所以只看到一个标签还不能确定主题。
        - 映射（mapping）：每个输入字母对应标签空间里的一个标签（"输入 → 标签"的对应关系）。
        - 噪声：每个标签有 ε = 0.1 的概率被换成 0–9 里随机的一个数字（数字 4–9 只会作为噪声出现）。

    两个"预训练世界"：
        世界 A（概念少）：每个主题只有 2 种固定的映射，一共 8 个概念。整个预训练语料就是这 8 个任务反复出现。
                          ——像"互联网上反复出现的少数几类任务"；模型可以把映射**记在权重里**，
                          示例只需要帮它**认出**是哪个任务。
        世界 B（概念多）：每篇文档的映射都是**新随机抽的**（每个字母独立地随机选标签空间里的一个标签）。
                          ——映射没法记在权重里，想预测对，只能**从同一篇文档前面的示例里照抄**。

    对任意一个提示（k 个示例 + 一个查询字母），我们都能精确算出"一个完美地学会了预训练分布的模型"
    会给出的答案——**贝叶斯后验预测（Bayes posterior predictive）**：

        p(y | 示例, x_查询) = Σ_概念 p(概念 | 示例, x_查询) · p(y | x_查询, 概念)

    这就是 Xie et al. 理论里的"理想的 in-context learner"。代码 02 训练一个小 Transformer，
    看它的 ICL 和这个理想答案差多少。

【主要接口】
    ALPHABET / make_tokenizer()          字母表与分词器
    CONCEPTS                             世界 A 的 8 个概念：(主题, 映射)
    make_doc(rng, world)                 生成一篇预训练文档
    doc_batch(tok, world, n, generator)  一批文档 → (x, y, mask)，交给 seq_tasks.train_masked_lm
    make_items(condition, k, n, seed)    生成 n 个评测提示（8 种条件，见 CONDITIONS）
    bayes_predict(world, demos, xq, evidence)  贝叶斯后验预测（evidence="full" 或 "bag"）
    bayes_accuracy(world, items, evidence)     贝叶斯预测在一组提示上的准确率
    model_accuracy(model, tok, items)          模型的准确率（取下一个字符概率最大的那个）

【运行方式】
    python3 icl_world.py      # 自检：打印几篇文档、8 个概念，以及一个手算可以核对的后验（< 1 秒）

只用 Python 标准库 + PyTorch（模型相关的两个函数）。

【Python 小知识】
    - math.log / math.exp：把很多概率相乘时，直接相乘会下溢成 0；改成"对数相加"，最后再减去最大值、取 exp，
      数值上就稳定了（"log-sum-exp 技巧"）。bayes_predict 里就是这么做的。
    - collections.namedtuple：一种"带字段名的元组"，item.prompt 比 item[0] 好读。
"""

import math
import os
import random
import sys
from collections import namedtuple

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from tiny_gpt import CharTokenizer  # noqa: E402

# =============================================================================
# 1. 世界的设定
# =============================================================================
ALPHABET = "\n abcdefgh0123456789"
INPUTS = "abcdefgh"
DIGITS = "0123456789"
N_TOPICS = 4
HOME_P = 0.3            # 主题的两个常用字母，各自的概率
EPS = 0.1               # 标签噪声：以 ε 的概率换成 0–9 里随机一个数字
N_PAIRS = 16            # 每篇文档 16 对；文档长度 = 16 × 3 − 1 个字符（对与对之间一个空格）
BLOCK_SIZE = 3 * N_PAIRS  # "\n" + 文档 + "\n" 共 3·16 + 1 个字符，模型一次看 3·16 个
LABEL_SETS = [(str(t), str((t + 1) % N_TOPICS)) for t in range(N_TOPICS)]


def make_tokenizer():
    """本世界的字符级分词器（20 个字符：换行、空格、8 个字母、10 个数字）。"""
    return CharTokenizer(ALPHABET)


def input_dist(t):
    """
    主题 t 的输入分布：dict 字母 → 概率。

    【思路】常用字母是 INPUTS[2t] 和 INPUTS[2t+1]，各 0.3；其余 6 个字母平分剩下的 0.4（各约 0.067）。
    【FM 联系】Min et al. 2022 发现：示例里的**输入**来自哪个分布很重要（换成别的语料里的句子会掉 3–16 个点）。
            在这个世界里，输入分布本身就是"这是哪个主题"的证据之一。
    """
    home = {INPUTS[2 * t], INPUTS[2 * t + 1]}
    return {x: (HOME_P if x in home else (1 - 2 * HOME_P) / (len(INPUTS) - 2)) for x in INPUTS}


def random_mapping(rng, t):
    """为主题 t 随机抽一个映射：每个字母独立地选标签空间里的一个标签。返回 dict 字母 → 标签。"""
    return {x: LABEL_SETS[t][rng.randrange(2)] for x in INPUTS}


def make_concepts(seed=0):
    """
    世界 A 的 8 个概念：每个主题 2 个固定映射。

    【思路】两个映射如果几乎一样，示例就分不出它们；如果完全相反，又太容易分。所以要求两者
            至少在 3 个字母上相同、至少在 3 个字母上不同（抽到不满足的就重抽）。
    【返回】list[(主题, 映射)]，长度 8，顺序固定（主题 0 的两个、主题 1 的两个……）
    """
    rng = random.Random(seed)
    concepts = []
    for t in range(N_TOPICS):
        m1 = random_mapping(rng, t)
        while True:
            m2 = random_mapping(rng, t)
            same = sum(m1[x] == m2[x] for x in INPUTS)
            if 3 <= same <= len(INPUTS) - 3:
                break
        concepts += [(t, m1), (t, m2)]
    return concepts


CONCEPTS = make_concepts()


def label_prob(y, x, mapping):
    """p(标签 y | 输入 x, 映射)：(1 − ε)·[y 是映射给出的标签] + ε / 10。"""
    return (1 - EPS) * (y == mapping[x]) + EPS / len(DIGITS)


def sample_input(rng, t):
    """按主题 t 的输入分布抽一个字母。"""
    r, acc = rng.random(), 0.0
    for x, p in input_dist(t).items():
        acc += p
        if r < acc:
            return x
    return INPUTS[-1]   # 浮点误差兜底


def sample_label(rng, x, mapping):
    """按 label_prob 抽一个标签（以 ε 的概率是 0–9 里的随机数字）。"""
    return DIGITS[rng.randrange(len(DIGITS))] if rng.random() < EPS else mapping[x]


def make_doc(rng, world):
    """
    生成一篇预训练文档（不含首尾的换行）。

    【参数】rng：random.Random；world："A"（8 个固定概念）或 "B"（每篇文档一个新的随机映射）
    【返回】字符串，例如 "b1 a0 b1 d0 a0 h1 ..."（16 对，对与对之间一个空格）
    """
    if world == "A":
        t, mapping = CONCEPTS[rng.randrange(len(CONCEPTS))]
    elif world == "B":
        t = rng.randrange(N_TOPICS)
        mapping = random_mapping(rng, t)
    else:
        raise ValueError(f"unknown world: {world}")
    pairs = []
    for _ in range(N_PAIRS):
        x = sample_input(rng, t)
        pairs.append(x + sample_label(rng, x, mapping))
    return " ".join(pairs)


def doc_batch(tok, world, n, generator):
    """
    一批预训练文档 → (x, y, mask)，可以直接交给 seq_tasks.train_masked_lm。

    【思路】每条序列 = "\\n" + 文档 + "\\n"，从位置 0 开始（和评测时的提示对齐）。
            mask 全是 1：**每个位置都预测**（输入字母、标签、空格、结尾换行），这就是普通的语言模型预训练，
            没有任何"任务"或"标签"的概念——ICL 只能从"预测下一个字符"里自己长出来。
    【参数】generator：torch.Generator，用它抽一个种子再造 random.Random（这样整批文档由训练循环的种子决定）
    """
    rng = random.Random(torch.randint(0, 2 ** 31 - 1, (1,), generator=generator).item())
    ids = torch.tensor([tok.encode("\n" + make_doc(rng, world) + "\n") for _ in range(n)], dtype=torch.long)
    return ids[:, :-1], ids[:, 1:], torch.ones(n, ids.shape[1] - 1)


# =============================================================================
# 2. 评测提示：8 种条件
# =============================================================================
Item = namedtuple("Item", "prompt demos xq target")

CONDITIONS = {
    "gold": "正确标签（标准的 few-shot 提示）",
    "random": "标签随机（但在正确的标签空间里）",
    "random_oodin": "标签随机 + 输入不按主题分布（输入分布错）",
    "random_space": "标签换成标签空间之外的数字 4–9（标签空间错）",
    "inputs_only": "只有输入、没有标签（格式错）",
    "flip": "标签翻转（标签空间里的两个标签对调）",
    "new": "没见过的映射（正确标签，但映射是新随机抽的）",
}


def make_items(condition, k, n, seed):
    """
    生成 n 个评测提示：k 个示例 + 一个查询字母。

    【思路】先抽一个"真实概念"：除了 "new" 以外都从世界 A 的 8 个概念里抽（世界 B 的模型同样用这些提示评测——
            对它来说，它们只是又一个随机映射）。示例按条件改造（见 CONDITIONS）；查询字母总是按主题分布抽。
            目标答案：
              - "flip"：翻转之后的标签（看模型是否**跟着示例走**）；
              - 其余：真实映射给出的标签。
    【参数】condition：CONDITIONS 的键；k：示例个数（0–15）；n：提示个数；seed：随机种子
    【返回】list[Item]。Item.prompt 形如 "\\nb1 c0 a1 b"（最后一个字母是查询，模型接着写的那个字符就是答案）
    【边界】k 不能超过 15，否则提示比预训练文档还长。
    """
    if condition not in CONDITIONS:
        raise ValueError(f"unknown condition: {condition}")
    if not 0 <= k < N_PAIRS:
        raise ValueError("k 必须在 0 到 15 之间")
    rng = random.Random(seed)
    items = []
    for _ in range(n):
        if condition == "new":
            t = rng.randrange(N_TOPICS)
            mapping = random_mapping(rng, t)
        else:
            t, mapping = CONCEPTS[rng.randrange(len(CONCEPTS))]
        demos = []
        for _ in range(k):
            x = INPUTS[rng.randrange(len(INPUTS))] if condition == "random_oodin" else sample_input(rng, t)
            y = mapping[x]
            if condition in ("random", "random_oodin"):
                y = LABEL_SETS[t][rng.randrange(2)]
            elif condition == "random_space":
                y = DIGITS[4 + rng.randrange(6)]
            elif condition == "flip":
                y = LABEL_SETS[t][1 - LABEL_SETS[t].index(mapping[x])]
            demos.append((x, y))
        xq = sample_input(rng, t)
        target = mapping[xq]
        if condition == "flip":
            target = LABEL_SETS[t][1 - LABEL_SETS[t].index(mapping[xq])]
        if condition == "inputs_only":
            prompt = "\n" + "".join(x + " " for x, _ in demos) + xq
        else:
            prompt = "\n" + "".join(x + y + " " for x, y in demos) + xq
        items.append(Item(prompt, demos, xq, target))
    return items


# =============================================================================
# 3. 贝叶斯后验预测
# =============================================================================
def _normalize_logs(logw):
    """一组对数权重 → 归一化的概率（先减最大值再 exp，避免下溢）。"""
    m = max(logw)
    w = [math.exp(v - m) for v in logw]
    z = sum(w)
    return [v / z for v in w]


def bayes_predict(world, demos, xq, evidence="full"):
    """
    贝叶斯后验预测：一个完美学会了"世界 world 的预训练分布"的模型，在这个提示后面给每个标签的概率。

    【思路】对每个可能的概念 c：
                log 权重(c) = log 先验(c) + log p(x_查询 | c) + Σ_示例 [ log p(x_i | c) + log p(y_i | x_i, c) ]
            归一化得到后验 p(c | 提示)，再对 p(y | x_查询, c) 加权平均。
            evidence = "full"：上面这个完整的式子——既看输入分布、标签空间，也看"哪个输入配哪个标签"。
            evidence = "bag" ：把 log p(y_i | x_i, c) 换成 log p(y_i | c)（对 x 求平均之后的标签分布）——
                               也就是**不看配对**，只看"出现了哪些输入、哪些标签"（像一个只数词频的读者）。
                               用来对照：模型到底有没有用上"输入 → 标签"的对应关系。
            世界 B 的"概念" = (主题, 每个字母各自的标签)，有 4 × 2^8 个；因为各字母的标签相互独立，
            后验可以按字母拆开算，不用真的枚举 1024 个概念。
    【参数】world："A" 或 "B"；demos：list[(x, y)]；xq：查询字母；evidence："full" 或 "bag"
    【返回】dict 标签 → 概率（10 个数字，和为 1）
    【边界】示例里的每个字母、每个数字在预训练分布下都有非零概率（噪声保证了这一点），所以后验总有定义；
            "只有输入、没有标签"的提示在预训练分布下概率为 0（预训练文档里字母后面一定跟着数字），
            这个函数不处理它——代码 02 在表里写 "—"。
    """
    if evidence not in ("full", "bag"):
        raise ValueError("evidence 必须是 'full' 或 'bag'")
    pred = {y: 0.0 for y in DIGITS}
    if world == "A":
        logw = []
        for t, mapping in CONCEPTS:
            q = input_dist(t)
            lw = math.log(1 / len(CONCEPTS)) + math.log(q[xq])
            for x, y in demos:
                lw += math.log(q[x])
                if evidence == "full":
                    lw += math.log(label_prob(y, x, mapping))
                else:
                    lw += math.log(sum(q[x2] * label_prob(y, x2, mapping) for x2 in INPUTS))
            logw.append(lw)
        for (t, mapping), w in zip(CONCEPTS, _normalize_logs(logw)):
            for y in DIGITS:
                pred[y] += w * label_prob(y, xq, mapping)
        return pred
    if world != "B":
        raise ValueError(f"unknown world: {world}")
    logw, per_topic = [], []
    for t in range(N_TOPICS):
        q = input_dist(t)
        lw = math.log(1 / N_TOPICS) + math.log(q[xq])
        post_q = [0.5, 0.5]                       # 查询字母的标签是标签空间里第 0 个 / 第 1 个的后验
        for x in INPUTS:
            obs = [y for x2, y in demos if x2 == x]
            lw += len(obs) * math.log(q[x])
            # 这个字母的标签 = LABEL_SETS[t][j]（j = 0 或 1，先验各一半）时，观察到 obs 的概率
            lik = [0.5, 0.5]
            for j in range(2):
                for y in obs:
                    lik[j] *= label_prob(y, x, {x: LABEL_SETS[t][j]})
            if evidence == "full":
                lw += math.log(lik[0] + lik[1])
                if x == xq:
                    post_q = [lik[0] / (lik[0] + lik[1]), lik[1] / (lik[0] + lik[1])]
            else:   # 不看配对：每个标签只按"这个主题的标签分布"计分，查询字母的标签保持先验
                for y in obs:
                    lw += math.log(0.5 * label_prob(y, x, {x: LABEL_SETS[t][0]})
                                   + 0.5 * label_prob(y, x, {x: LABEL_SETS[t][1]}))
        logw.append(lw)
        per_topic.append({y: sum(post_q[j] * label_prob(y, xq, {xq: LABEL_SETS[t][j]}) for j in range(2))
                          for y in DIGITS})
    for p, w in zip(per_topic, _normalize_logs(logw)):
        for y in DIGITS:
            pred[y] += w * p[y]
    return pred


def argmax_accuracy(pred, target):
    """
    预测分布的 argmax 是否等于目标。并列第一时按"随机挑一个"算期望（例如两个并列、其中一个对 → 0.5）。
    """
    best = max(pred.values())
    tops = [y for y, v in pred.items() if v > best - 1e-12]
    return (1.0 / len(tops)) if target in tops else 0.0


def bayes_accuracy(world, items, evidence="full"):
    """贝叶斯预测在一组提示上的平均准确率（math.fsum 求和，保证不同 Python 版本结果一致）。"""
    return math.fsum(argmax_accuracy(bayes_predict(world, it.demos, it.xq, evidence), it.target)
                     for it in items) / len(items)


# =============================================================================
# 4. 模型的预测
# =============================================================================
@torch.no_grad()
def model_accuracy(model, tok, items):
    """
    模型的 ICL 准确率：在提示末尾，模型给出的下一个字符里概率最大的那个是否等于目标。

    【思路】同一组提示的 k 相同、长度相同，可以放进一个 batch 一次算完。
    【边界】提示长度不同时（不会发生在 make_items 生成的同一组里）会被 torch.tensor 拒绝。
    """
    x = torch.tensor([tok.encode(it.prompt) for it in items], dtype=torch.long)
    pred = model(x)[:, -1].argmax(dim=-1).tolist()
    return sum(tok.alphabet[i] == it.target for i, it in zip(pred, items)) / len(items)


# =============================================================================
# 5. 自检
# =============================================================================
if __name__ == "__main__":
    rng = random.Random(0)
    for world in ("A", "B"):
        print(f"世界 {world} 的两篇文档：")
        for _ in range(2):
            print("   ", make_doc(rng, world))
    print("\n世界 A 的 8 个概念（主题：标签空间；映射 a–h）：")
    for t, m in CONCEPTS:
        print(f"    主题 {t} {LABEL_SETS[t]}：" + " ".join(f"{x}→{m[x]}" for x in INPUTS))
    # 手算核对：没有示例时，世界 A 的贝叶斯预测只依赖查询字母属于哪个主题的常用字母
    p0 = bayes_predict("A", [], "a")
    print("\n世界 A，0 个示例，查询 'a' 的后验预测（前 4 个标签）：", {y: round(p0[y], 3) for y in "0123"})
    assert abs(sum(p0.values()) - 1) < 1e-9
    items = make_items("gold", 4, 3, seed=1)
    for it in items:
        print("    提示", repr(it.prompt), "目标", it.target,
              "贝叶斯 argmax", max(bayes_predict("A", it.demos, it.xq).items(), key=lambda kv: kv[1])[0])
    print("自检通过。")