"""
icl_world.py —— 第 9 讲：一个能精确算出"贝叶斯最优上下文学习"的小世界
=====================================================================

【这个文件做什么】
    第 6 讲看到：上下文学习（ICL）只有在预训练数据"奖励利用上下文"时才会出现。这里造一个专门的预训练语料：
    每篇文档由一个**隐藏的任务**生成，内容是一串"输入 → 标签"的示例，例如

        <bos> apple : ripe ; cat : hungry ; rose : pretty ; … <eos>          （任务 = 名词的类别）
        <bos> apple -> small , goat -> big , … <eos>                          （任务 = 大小，另一种格式）

    在这样的语料上训练的语言模型，读到几个示例后要预测下一个标签，就必须先"猜出这篇文档是哪个任务"——
    这正是 Xie et al. (2022) 把 ICL 解释为**隐式贝叶斯推断**的设定。好处是：任务的先验、每个任务的映射、
    标签噪声都是我们定的，所以**贝叶斯最优**的预测可以精确算出来，和训练出的 Transformer 逐项比较。

【世界设定】
    输入：花园世界的 43 个名词。属性（颜色 / 大小 / 价位）沿用第 8 讲花园小店的商品目录。
    任务（括号里是预训练时的先验概率）：
        cat     类别     → ripe / fresh / pretty / hungry（0.15）
        color   颜色     → red / green / yellow / white  （0.15）
        size    大小     → small / big                   （0.12）
        price   价位     → cheap / dear                  （0.12）
        animal  是动物吗 → yes / no                      （0.12）
        flower  是花吗   → yes / no                      （0.12）   ← 和 animal 共用标签空间
        size_flip / animal_flip  上面两个任务的"反着说"版本（各 0.02）：预训练里很少见的"违反常识"的映射
        symbol  抽象标签：每篇文档随机把 4 个类别对应到 8 个无意义符号 <s0>…<s7> 中的 4 个（0.18）
    标签噪声：每个示例有 10% 的概率换成同一标签空间里的另一个标签。
    格式（先验）：  x : y ;（0.70）   x -> y ,（0.25）   x = y |（0.05）；另有一种预训练从没出现过的 x ~ y &。
    每篇文档 2～24 个示例。

【LLM 联系】
    真实预训练数据里，"同一主题的一串例子"（词汇表、问答列表、对照翻译、表格）就是这样的"隐式任务文档"。
    Xie et al. (2022) 的 GINC 数据集用隐马尔可夫模型的混合做了同样的事；本讲用更简单的"名词 → 标签"映射，
    换来可以穷举的贝叶斯后验（8 个命名任务 + 1680 种符号映射）。
"""

import itertools
import os
import random
import sys

import numpy as np
import torch

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

from garden_lm import Vocab, clm_batches, clm_loss, load_or_train, train  # noqa: E402
from hai_world import CATALOG  # noqa: E402
from toy_world import NOUN_CAT, NOUNS  # noqa: E402

NOUN_LIST = [n for c in NOUNS for n in NOUNS[c]]                       # 43 个名词，固定顺序
CATS = list(NOUNS)                                                     # fruit veg flower animal
SYMBOLS = [f"<s{i}>" for i in range(8)]

# 任务：名字 -> (标签列表, 名词 -> 标签 的函数)
TASKS = {
    "cat": (["ripe", "fresh", "pretty", "hungry"], lambda n: ["ripe", "fresh", "pretty", "hungry"][CATS.index(NOUN_CAT[n])]),
    "color": (["red", "green", "yellow", "white"], lambda n: CATALOG[n]["color"]),
    "size": (["small", "big"], lambda n: CATALOG[n]["size"]),
    "price": (["cheap", "dear"], lambda n: CATALOG[n]["price"]),
    "animal": (["yes", "no"], lambda n: "yes" if NOUN_CAT[n] == "animal" else "no"),
    "flower": (["yes", "no"], lambda n: "yes" if NOUN_CAT[n] == "flower" else "no"),
    "size_flip": (["small", "big"], lambda n: "big" if CATALOG[n]["size"] == "small" else "small"),
    "animal_flip": (["yes", "no"], lambda n: "no" if NOUN_CAT[n] == "animal" else "yes"),
}
TASK_PRIOR = {"cat": 0.15, "color": 0.15, "size": 0.12, "price": 0.12, "animal": 0.12, "flower": 0.12,
              "size_flip": 0.02, "animal_flip": 0.02, "symbol": 0.18}
NOISE = 0.1
FORMATS = {"A": (":", ";"), "B": ("->", ","), "C": ("=", "|"), "D": ("~", "&")}   # D 从不出现在预训练里
FORMAT_PRIOR = {"A": 0.70, "B": 0.25, "C": 0.05}
NULL = "<null>"                                                       # 第 03 号脚本"内容无关的输入"
COT_SEP = "=>"                                                        # 第 05 号脚本的思维链格式

LABELS = sorted({l for labs, _ in TASKS.values() for l in labs}) + SYMBOLS
VERBS_2HOP = ["eat", "grow", "water", "feed"]                         # 第 05 号脚本：类别 -> 动词


def build_vocab():
    """<bos> <eos> <pad> + 43 个名词 + 全部标签 + 格式符号 + <null> + 思维链用的词。"""
    toks = ["<bos>", "<eos>", "<pad>"] + NOUN_LIST + LABELS
    for sep, end in FORMATS.values():
        toks += [sep, end]
    toks += [NULL, COT_SEP] + CATS + VERBS_2HOP
    seen, out = set(), []
    for t in toks:
        if t not in seen:
            seen.add(t)
            out.append(t)
    return Vocab(out)


VOCAB = build_vocab()
BOS, EOS, PAD = 0, 1, 2
LABEL_IDS = torch.tensor([VOCAB.stoi[l] for l in LABELS])


# =====================================================================
# 1. 预训练语料
# =====================================================================

def sample_concept(rng):
    """
    抽一个任务（"概念" θ）。【返回】(任务名, 名词 -> 标签 的字典, 标签列表)。
    symbol 任务会随机抽 4 个符号、随机对应到 4 个类别。
    """
    names = list(TASK_PRIOR)
    t = rng.choices(names, weights=[TASK_PRIOR[x] for x in names])[0]
    if t == "symbol":
        syms = rng.sample(SYMBOLS, 4)
        return t, {n: syms[CATS.index(NOUN_CAT[n])] for n in NOUN_LIST}, syms
    labs, f = TASKS[t]
    return t, {n: f(n) for n in NOUN_LIST}, labs


def noisy(label, labels, rng, noise=NOISE):
    """以 noise 的概率把标签换成同一标签空间里的另一个。"""
    if rng.random() < noise:
        return rng.choice([l for l in labels if l != label])
    return label


def format_pairs(pairs, fmt):
    """[(x, y), ...] → token 列表，例如 ['apple', ':', 'ripe', ';', ...]。"""
    sep, end = FORMATS[fmt]
    out = []
    for x, y in pairs:
        out += [x, sep, y, end]
    return out


def sample_doc(rng):
    """一篇预训练文档：一个隐藏任务、一种格式、2～24 个带噪声的示例。"""
    _, mapping, labels = sample_concept(rng)
    fmt = rng.choices(list(FORMAT_PRIOR), weights=list(FORMAT_PRIOR.values()))[0]
    k = rng.randint(2, 24)
    pairs = [(x, noisy(mapping[x], labels, rng)) for x in (rng.choice(NOUN_LIST) for _ in range(k))]
    return format_pairs(pairs, fmt) + ["<eos>"]


def corpus(n, seed):
    rng = random.Random(seed)
    return [sample_doc(rng) for _ in range(n)]


# =====================================================================
# 2. 训练（带缓存）
# =====================================================================
SIZES = {"small": dict(d_model=64, n_heads=4, n_layers=2, d_ff=256),
         "large": dict(d_model=128, n_heads=4, n_layers=4, d_ff=512)}
CKPT_DIR = os.path.join(HERE, "checkpoints")


def icl_model(size="small", n_docs=20000, epochs=4, extra_docs=None, name=None, verbose=False):
    """
    在 ICL 语料上从零预训练一个因果 Transformer（第 2 讲的 TransformerLM），读缓存或训练。

    【参数】size："small"（d=64、2 层，约 10 万参数）或 "large"（d=128、4 层，约 80 万参数）；
            n_docs / epochs：语料规模与遍数；extra_docs：额外混入的文档（第 05 号脚本的思维链语料），
            此时必须给 name 作为缓存名。
    【返回】模型（eval 模式）。第一次训练：small 约 30 秒，large 约 90 秒。
    """
    from transformer import TransformerLM
    docs = corpus(n_docs, seed=0) + (extra_docs or [])
    valid = corpus(500, seed=1)
    arch = SIZES[size]
    cfg = {"arch": arch, "n_docs": n_docs, "epochs": epochs, "data": name or "icl_v1", "extra": len(extra_docs or [])}

    def to_ids(ds):
        # clm_batches 会在前面加 <bos>，所以文档里不放 <bos>
        return ds

    batches = clm_batches(to_ids(docs), VOCAB, 32)

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

    def fit(m):
        from garden_lm import clm_eval
        return train(m, batches, clm_loss, epochs=epochs, lr=2e-3, warmup=200, seed=0, verbose=verbose,
                     tag=f"icl_{size}", valid_fn=lambda mm: clm_eval(mm, valid, VOCAB)[1])

    model, _ = load_or_train(name or f"icl_{size}", cfg, build, fit, verbose=verbose, ckpt_dir=CKPT_DIR)
    return model


# =====================================================================
# 3. 贝叶斯最优预测
# =====================================================================

def _concepts():
    """
    全部概念：8 个命名任务 + symbol 任务的 1680 种映射（8 个符号里选 4 个、排列到 4 个类别）。
    【返回】(table, prior, label_sets, names)：
        table[c, i]：概念 c 下第 i 个名词的标签（在 LABELS 里的下标）；
        prior[c]：先验概率；label_sets[c]：概念 c 的标签空间（下标集合）；names[c]：概念名。
    """
    rows, prior, lsets, names = [], [], [], []
    for t, (labs, f) in TASKS.items():
        rows.append([LABELS.index(f(n)) for n in NOUN_LIST])
        prior.append(TASK_PRIOR[t])
        lsets.append({LABELS.index(l) for l in labs})
        names.append(t)
    perms = list(itertools.permutations(range(8), 4))
    for p in perms:
        syms = [SYMBOLS[i] for i in p]
        rows.append([LABELS.index(syms[CATS.index(NOUN_CAT[n])]) for n in NOUN_LIST])
        prior.append(TASK_PRIOR["symbol"] / len(perms))
        lsets.append({LABELS.index(s) for s in syms})
        names.append("symbol")
    return np.array(rows), np.array(prior), lsets, names


TABLE, PRIOR, LSETS, CNAMES = _concepts()
_LMASK = np.zeros((len(TABLE), len(LABELS)), dtype=bool)
for _c, _s in enumerate(LSETS):
    _LMASK[_c, list(_s)] = True
_LSIZE = _LMASK.sum(1)


def bayes_posterior(pairs, noise=NOISE):
    """
    P(概念 | 示例)。每个示例 (x, y) 的似然：y 等于概念给出的标签 → 1 − noise；
    y 在概念的标签空间里但不相等 → noise / (|标签空间| − 1)；y 不在标签空间里 → 0。
    【返回】长度为"概念数"的后验数组。所有示例的似然都是 0 时（没有任何概念能解释），退回先验。
    """
    logp = np.log(PRIOR)
    for x, y in pairs:
        xi, yi = NOUN_LIST.index(x), LABELS.index(y)
        match = TABLE[:, xi] == yi
        lik = np.where(match, 1 - noise, np.where(_LMASK[:, yi], noise / (_LSIZE - 1), 0.0))
        with np.errstate(divide="ignore"):
            logp = logp + np.log(lik)
    if not np.isfinite(logp).any():
        logp = np.log(PRIOR)
    logp -= logp.max()
    p = np.exp(logp)
    return p / p.sum()


def bayes_predict(pairs, xq, noise=NOISE):
    """贝叶斯最优：P(y | 示例, x_q) = Σ_θ P(θ | 示例) P(y | x_q, θ)。【返回】长度为 len(LABELS) 的分布。"""
    post = bayes_posterior(pairs, noise)
    xi = NOUN_LIST.index(xq)
    dist = np.zeros(len(LABELS))
    for c in np.nonzero(post > 1e-12)[0]:
        labs = list(LSETS[c])
        right = TABLE[c, xi]
        dist[right] += post[c] * (1 - noise)
        for l in labs:
            if l != right:
                dist[l] += post[c] * noise / (len(labs) - 1)
    return dist


# =====================================================================
# 4. 模型的 ICL 预测
# =====================================================================

def prompt_ids(pairs, xq, fmt="A"):
    """<bos> + 示例 + 查询输入 + 分隔符（模型接下来该说标签）。"""
    sep, _ = FORMATS[fmt]
    return [BOS] + VOCAB.encode(format_pairs(pairs, fmt) + [xq, sep])


@torch.no_grad()
def model_label_probs(model, prompts):
    """
    一批提示（长度可以不同）之后，模型对"下一个 token 是哪个标签"的分布（只在全部标签 token 上归一化）。

    【返回】(B, len(LABELS)) 的 numpy 数组。
    【思路】按长度分组，同一组不需要补齐（和第 6 讲 greedy_generate 的理由一样）。
    """
    out = [None] * len(prompts)
    groups = {}
    for i, p in enumerate(prompts):
        groups.setdefault(len(p), []).append(i)
    for _, idx in groups.items():
        for s in range(0, len(idx), 512):
            chunk = idx[s:s + 512]
            ids = torch.tensor([prompts[i] for i in chunk])
            logits = model(ids)[:, -1, LABEL_IDS]
            probs = torch.softmax(logits, -1).numpy()
            for j, i in enumerate(chunk):
                out[i] = probs[j]
    return np.stack(out)


def make_episodes(task, k, n, seed, mapping=None, labels=None):
    """
    评估用的 ICL 题目：n 道，每道 k 个示例 + 1 个查询，示例和查询的输入都是随机名词，示例标签按任务给出（不加噪声）。

    【参数】task：TASKS 里的任务名，或 "symbol"（此时每道题随机抽一种符号映射）；mapping / labels：可直接指定映射。
    【返回】[(pairs, xq, 正确标签)]。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        if mapping is not None:
            mp = mapping
        elif task == "symbol":
            syms = rng.sample(SYMBOLS, 4)
            mp = {x: syms[CATS.index(NOUN_CAT[x])] for x in NOUN_LIST}
        else:
            labs, f = TASKS[task]
            mp = {x: f(x) for x in NOUN_LIST}
        xs = [rng.choice(NOUN_LIST) for _ in range(k + 1)]
        out.append(([(x, mp[x]) for x in xs[:k]], xs[k], mp[xs[k]]))
    return out


def accuracy_model(model, episodes, fmt="A"):
    probs = model_label_probs(model, [prompt_ids(p, xq, fmt) for p, xq, _ in episodes])
    return float(np.mean([LABELS[int(pr.argmax())] == y for pr, (_, _, y) in zip(probs, episodes)]))


def accuracy_bayes(episodes):
    return float(np.mean([LABELS[int(bayes_predict(p, xq).argmax())] == y for p, xq, y in episodes]))


# =====================================================================
# 5. 监督微调（第 04 号脚本的"对齐"实验）
# =====================================================================

def sft(model, examples, steps=300, lr=5e-4, batch_size=32, seed=0):
    """
    只对"最后一个 token"计算损失的监督微调：examples = [(prompt_ids, target_id)]，教模型在提示之后说出 target。

    【思路】右侧补齐、用补齐掩码屏蔽补齐位置；每条样本只在它自己的最后一个提示位置上算交叉熵。
    【返回】None（就地修改模型，结束后切回 eval 模式）。
    """
    import torch.nn.functional as F
    from garden_lm import lr_lambda
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(20, steps))
    rng = random.Random(seed)
    model.train()
    for _ in range(steps):
        batch = [examples[rng.randrange(len(examples))] for _ in range(batch_size)]
        L = max(len(p) for p, _ in batch)
        ids = torch.full((len(batch), L), PAD, dtype=torch.long)
        pad = torch.ones(len(batch), L, dtype=torch.bool)
        for i, (p, _) in enumerate(batch):
            ids[i, :len(p)] = torch.tensor(p)
            pad[i, :len(p)] = False
        logits = model(ids, pad=pad)
        last = torch.tensor([len(p) - 1 for p, _ in batch])
        tgt = torch.tensor([t for _, t in batch])
        loss = F.cross_entropy(logits[torch.arange(len(batch)), last], tgt)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        sched.step()
    model.eval()


@torch.no_grad()
def next_token_probs(model, prompts):
    """一批提示之后，下一个 token 在整个词表上的分布。【返回】(B, |V|) 的 numpy 数组。"""
    out = [None] * len(prompts)
    groups = {}
    for i, p in enumerate(prompts):
        groups.setdefault(len(p), []).append(i)
    for _, idx in groups.items():
        ids = torch.tensor([prompts[i] for i in idx])
        probs = torch.softmax(model(ids)[:, -1], -1).numpy()
        for j, i in enumerate(idx):
            out[i] = probs[j]
    return np.stack(out)


__all__ = ["VOCAB", "NOUN_LIST", "CATS", "SYMBOLS", "TASKS", "TASK_PRIOR", "NOISE", "FORMATS", "FORMAT_PRIOR",
           "LABELS", "NULL", "COT_SEP", "VERBS_2HOP", "sample_concept", "format_pairs", "sample_doc", "corpus",
           "icl_model", "bayes_posterior", "bayes_predict", "prompt_ids", "model_label_probs", "make_episodes",
           "accuracy_model", "accuracy_bayes", "CNAMES", "PRIOR", "LABEL_IDS", "noisy", "NOUN_CAT", "CATALOG", "sft",
           "next_token_probs"]