"""
icl_lab.py —— 一个可以算出"最优答案"的上下文学习小世界（第 11 讲）
=====================================================================

【这个文件要回答的问题】
    第 9 讲我们发现：农场语言的小 GPT 怎么提示都看不到上下文学习（in-context learning, ICL）。为什么？
    一个流行的解释（Xie et al., 2022；Chan et al., 2022）：只有当训练数据里**同一段上下文共享某个隐藏的"任务"**，
    而且这个任务只能从上下文里推断出来时，"从上下文里学"才是降低损失的好办法。农场语言的句子彼此独立，没有这种结构。

    所以我们造一个有这种结构的小世界：
        16 个"物品"= 4 种颜色 × 4 种形状，每个物品是一个 token（红圆、绿方、蓝角、黄十……）。
        一段序列 = 一串"物品 → 标签"的例子，**整段共用同一个任务**：
            - 任务先选一个**特征**：颜色或形状；
            - 再选**标签的写法**：
                自然标签（natural）：标签就是那个特征的名字（红 / 绿 / 蓝 / 黄，或 圆 / 方 / 角 / 十）——"模型早就认识的任务"；
                符号标签（symbolic）：标签是 A / B / C / D，特征值到符号的对应关系每段**随机重排**——"全新的任务"，
                只能从这一段的例子里推断。
            - 自然标签的序列有一半在开头带一个"任务说明"（颜色 / 形状，有几个出现频率不同的同义说法），
              用来做零样本和提示工程的实验。
        例：<s> 颜色 红圆 红 蓝方 蓝 黄十 黄 ……          （自然标签、带说明）
            <s> 红圆 C 蓝方 A 红角 C 绿圆 B ……             （符号标签：颜色 → 红=C、蓝=A、绿=B、……）

    这个世界的好处：任务只有 2 × 1 + 2 × 24 = 50 种，可以精确算出**贝叶斯最优**的预测（bayes_predict）——
    "一个完美的推理者看到这些例子后会怎么答"，拿来和 Transformer 比。

【本文件提供】
    ITEMS / COLORS / SHAPES / SYMBOLS / DESCRIPTORS / VOCAB / ID     词表
    Task / sample_task / label_of                                     任务
    make_sequence / sample_batch                                      训练数据（只在标签位置算损失）
    bayes_predict                                                     精确的贝叶斯预测（50 个假设）
    build_model / train_icl / load_or_train_icl                       模型（第 2 讲的 GPT）与训练、缓存
    query_probs / episode_ids                                         评估：给定例子和查询，模型对下一个词的分布
"""

import itertools
import math
import os
import random
import sys
import time

import torch
import torch.nn.functional as F

_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture02",):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from lm_models import GPT  # noqa: E402

CHECKPOINT_DIR = os.path.join(_HERE, "checkpoints")
IGNORE = -100

# =====================================================================
# 1. 词表
# =====================================================================

COLORS = ["红", "绿", "蓝", "黄"]
SHAPES = ["圆", "方", "角", "十"]
ITEMS = [c + s for c in COLORS for s in SHAPES]            # 16 个物品；下标 i 的颜色是 i // 4、形状是 i % 4
SYMBOLS = ["A", "B", "C", "D"]
# 任务说明（自然标签的序列有一半带一个）：每个特征三种说法，出现频率差别很大
DESCRIPTORS = {
    "color": [("颜色", 0.800), ("什么色", 0.199), ("色彩", 0.001)],
    "shape": [("形状", 0.800), ("什么形", 0.199), ("外形", 0.001)],
}
UNSEEN = "看这里"                                            # 词表里有、但训练时从来没出现过的"说明"
SPECIALS = ["<pad>", "<s>"]
VOCAB = SPECIALS + ITEMS + COLORS + SHAPES + SYMBOLS + \
    [w for f in ("color", "shape") for w, _ in DESCRIPTORS[f]] + [UNSEEN]
ID = {w: i for i, w in enumerate(VOCAB)}
PAD, BOS = ID["<pad>"], ID["<s>"]
LABEL_IDS = [ID[w] for w in COLORS + SHAPES + SYMBOLS]      # 所有可能的标签（12 个）
P_NATURAL = 0.2                                              # 训练数据里自然标签的比例（其余 80% 是符号标签）
P_DESCRIPTOR = 0.5                                           # 自然标签的序列带说明的比例
K_TRAIN = 12                                                 # 每段序列的例子数


def feature_value(item, feature):
    """物品（下标 0..15）在某个特征上的取值（0..3）。"""
    return item // 4 if feature == "color" else item % 4


# =====================================================================
# 2. 任务
# =====================================================================

class Task:
    """
    一个任务 = (标签写法 kind, 特征 feature, 符号的排列 perm)。
    【属性】kind："natural" 或 "symbolic"；feature："color" 或 "shape"；
        perm：符号任务里"特征值 v 对应符号 SYMBOLS[perm[v]]"，自然任务为 None。
    """
    __slots__ = ("kind", "feature", "perm")

    def __init__(self, kind, feature, perm=None):
        self.kind, self.feature, self.perm = kind, feature, perm

    def __repr__(self):
        if self.kind == "natural":
            return f"自然标签·{'颜色' if self.feature == 'color' else '形状'}"
        names = COLORS if self.feature == "color" else SHAPES
        return f"符号标签·{'颜色' if self.feature == 'color' else '形状'}（" + \
            "、".join(f"{names[v]}={SYMBOLS[self.perm[v]]}" for v in range(4)) + "）"


def sample_task(rng, p_natural=P_NATURAL):
    """按训练分布抽一个任务。【参数】rng：random.Random；p_natural：自然标签的概率。"""
    kind = "natural" if rng.random() < p_natural else "symbolic"
    feature = rng.choice(["color", "shape"])
    perm = None
    if kind == "symbolic":
        perm = list(range(4))
        rng.shuffle(perm)
    return Task(kind, feature, perm)


def label_of(task, item):
    """物品在这个任务下的标签（token 字符串）。"""
    v = feature_value(item, task.feature)
    if task.kind == "natural":
        return (COLORS if task.feature == "color" else SHAPES)[v]
    return SYMBOLS[task.perm[v]]


# =====================================================================
# 3. 训练数据
# =====================================================================

def make_sequence(task, items, descriptor=None, labels=None):
    """
    拼一段序列：[<s>] + [说明]? + [x1, y1, x2, y2, ...]。
    【参数】items：物品下标列表；descriptor：说明 token（None 表示不带）；
        labels：手动指定每个例子的标签（做"随机标签"实验用），None 表示用任务的正确标签
    【返回】(ids 列表, 要算损失的位置列表)——所有标签的位置，外加 <s> 之后的第一个词
    【为什么也算第一个词】这样模型会学到"一段序列以哪个说明开头"的概率 P(说明 | <s>)，03 号脚本拿它算提示的困惑度。
        后面的物品是均匀随机的，预测它们学不到任何东西，所以不算。
    """
    ids = [BOS] + ([ID[descriptor]] if descriptor else [])
    label_pos = [1]
    for j, it in enumerate(items):
        ids.append(ID[ITEMS[it]])
        label_pos.append(len(ids))
        ids.append(ID[labels[j] if labels is not None else label_of(task, it)])
    return ids, label_pos


def sample_descriptor(rng, task):
    """自然标签的任务按 P_DESCRIPTOR 的概率带一个说明（按同义词的频率抽）；符号任务从不带。"""
    if task.kind != "natural" or rng.random() >= P_DESCRIPTOR:
        return None
    words, weights = zip(*DESCRIPTORS[task.feature])
    return rng.choices(words, weights=weights)[0]


def sample_batch(rng, batch_size, p_natural=P_NATURAL, k=K_TRAIN, shared_task=True):
    """
    一批训练序列（输入、目标都是 (B, T) 的 LongTensor；目标只在标签位置有值，其余是 IGNORE）。
    【参数】shared_task=False 时，每个例子都换一个新任务（"例子之间没有共同的任务"——对照实验用）。
    【边界】带说明的序列比不带的长 1，短的在末尾补 <pad>。
    """
    seqs = []
    for _ in range(batch_size):
        task = sample_task(rng, p_natural)
        items = [rng.randrange(16) for _ in range(k)]
        if shared_task:
            ids, pos = make_sequence(task, items, sample_descriptor(rng, task))
        else:
            labels = [label_of(sample_task(rng, p_natural), it) for it in items]
            ids, pos = make_sequence(task, items, None, labels)
        seqs.append((ids, pos))
    T = max(len(s) for s, _ in seqs)
    inp = torch.full((batch_size, T - 1), PAD, dtype=torch.long)
    tgt = torch.full((batch_size, T - 1), IGNORE, dtype=torch.long)
    for b, (ids, pos) in enumerate(seqs):
        inp[b, :len(ids) - 1] = torch.tensor(ids[:-1])
        for p in pos:
            tgt[b, p - 1] = ids[p]                     # 位置 p 的 token 由位置 p−1 的输出预测
    return inp, tgt


# =====================================================================
# 4. 贝叶斯最优的预测
# =====================================================================

def all_tasks():
    """全部 50 个任务和它们在训练分布下的先验概率。"""
    out = []
    for f in ("color", "shape"):
        out.append((Task("natural", f), P_NATURAL / 2))
        for perm in itertools.permutations(range(4)):
            out.append((Task("symbolic", f, list(perm)), (1 - P_NATURAL) / 2 / 24))
    return out


_ALL_TASKS = all_tasks()


def descriptor_likelihood(task, descriptor):
    """P(说明 | 任务)：自然任务带说明的概率 × 这个说法的频率；符号任务从不带说明。"""
    if descriptor is None:
        return 1 - P_DESCRIPTOR if task.kind == "natural" else 1.0
    if task.kind != "natural":
        return 0.0
    for w, p in DESCRIPTORS[task.feature]:
        if w == descriptor:
            return P_DESCRIPTOR * p
    return 0.0                                          # 另一个特征的说法，或者没见过的说法


def consistent_tasks(demos, descriptor=None):
    """和所有例子（以及说明）都对得上的任务列表——贝叶斯推理里"还活着"的假设。"""
    return [t for t, prior in _ALL_TASKS
            if prior * descriptor_likelihood(t, descriptor) > 0 and all(label_of(t, it) == lab for it, lab in demos)]


def bayes_predict(demos, query, descriptor=None):
    """
    贝叶斯最优的预测：P(查询的标签 | 例子, 说明)。

    【思路】50 个任务假设，每个的后验 ∝ 先验 × P(说明 | 任务) × Π P(例子的标签 | 任务)。
        标签是任务的确定函数，所以 P(例子的标签 | 任务) 不是 1 就是 0："和所有例子都对得上"的任务才留下。
        预测 = 留下的任务按后验加权，各自给查询物品贴的标签。
    【参数】demos：[(物品下标, 标签字符串)]；query：物品下标；descriptor：说明（或 None）
    【返回】字典 {标签字符串: 概率}
    【边界】如果没有任何任务和例子对得上（比如例子的标签是随机打乱的），返回 None——训练分布下这种情况不会发生，
        贝叶斯推理"无话可说"；调用者自己决定怎么处理。
    """
    post = {}
    for task, prior in _ALL_TASKS:
        w = prior * descriptor_likelihood(task, descriptor)
        if w == 0:
            continue
        if all(label_of(task, it) == lab for it, lab in demos):
            y = label_of(task, query)
            post[y] = post.get(y, 0.0) + w
    z = sum(post.values())
    if z == 0:
        return None
    return {y: p / z for y, p in post.items()}


# =====================================================================
# 5. 模型、训练、缓存
# =====================================================================

MODEL_CONFIG = {"d_model": 128, "n_layers": 4, "n_heads": 4, "max_len": 32}
MAIN_STEPS = 10000          # 01 号脚本训练的主模型的步数；02、03 号脚本按同样的配置读取缓存


def build_model(seed=0, config=MODEL_CONFIG):
    torch.manual_seed(seed)
    return GPT(len(VOCAB), config["d_model"], config["n_layers"], config["n_heads"], max_len=config["max_len"])


def train_icl(model, steps, p_natural=P_NATURAL, batch_size=64, lr=3e-4, warmup=200, seed=0, shared_task=True,
              callback=None, every=0, log=print):
    """
    在 ICL 世界的数据上训练（每一步现造一批新序列，所以不存在"过拟合训练集"）。
    【参数】callback(step, model)：每 every 步调用一次（记录学习曲线用）
        lr：默认 3e-4。pilot 里学习率 1e-3 训练 6000 步，符号任务的上下文学习始终没有出现（一直是瞎猜）；
        3e-4 时在 3000–3500 步之间突然出现。
    【返回】最后 100 步的平均损失（只算标签位置）
    """
    rng = random.Random(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    losses = []
    model.train()
    for step in range(1, steps + 1):
        for gr in opt.param_groups:
            gr["lr"] = lr * min(1.0, step / warmup)
        inp, tgt = sample_batch(rng, batch_size, p_natural, shared_task=shared_task)
        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()
        losses.append(loss.item())
        if callback is not None and every and step % every == 0:
            model.eval()
            callback(step, model)
            model.train()
        if log and step % max(steps // 5, 1) == 0:
            log(f"    step {step:5d}  最近 100 步的平均损失 {sum(losses[-100:]) / len(losses[-100:]):.4f}")
    model.eval()
    return sum(losses[-100:]) / len(losses[-100:])


def load_or_train_icl(name, steps, p_natural=P_NATURAL, shared_task=True, log=print, callback=None, every=0):
    """
    有缓存就读取 checkpoints/<name>.pt，否则训练并保存。
    【返回】(model, info)：info 里有训练时记录的学习曲线（callback 的返回值列表）和用时
    【边界】缓存里存了配置；配置不一致时重新训练。steps 很小（冒烟测试）时既不读也不写缓存。
    """
    config = dict(MODEL_CONFIG, steps=steps, p_natural=p_natural, shared_task=shared_task)
    path = os.path.join(CHECKPOINT_DIR, f"{name}.pt")
    model = build_model()
    use_cache = steps >= 1000
    if use_cache and os.path.exists(path):
        blob = torch.load(path, weights_only=True)
        if blob["config"] == config:
            model.load_state_dict(blob["state_dict"])
            model.eval()
            if log:
                log(f"  （读取缓存：checkpoints/{name}.pt）")
            return model, blob["info"]
    curve = []

    def cb(step, m):
        curve.append(callback(step, m))

    t0 = time.perf_counter()
    loss = train_icl(model, steps, p_natural, shared_task=shared_task, callback=cb if callback else None,
                     every=every, log=log)
    info = {"curve": curve, "final_loss": loss, "seconds": time.perf_counter() - t0}
    if use_cache:
        os.makedirs(CHECKPOINT_DIR, exist_ok=True)
        torch.save({"config": config, "state_dict": model.state_dict(), "info": info}, path)
        if log:
            log(f"  （已保存到 checkpoints/{name}.pt）")
    return model, info


# =====================================================================
# 6. 评估
# =====================================================================

def episode_ids(demos, query, descriptor=None):
    """[<s>] + [说明]? + 例子 + [查询物品]：模型在最后一个位置的输出就是对查询标签的预测。"""
    ids = [BOS] + ([ID[descriptor]] if descriptor else [])
    for it, lab in demos:
        ids += [ID[ITEMS[it]], ID[lab]]
    return ids + [ID[ITEMS[query]]]


@torch.no_grad()
def query_probs(model, episodes):
    """
    一批"例子 + 查询"，返回模型在查询位置对整个词表的概率分布 (N, |V|)。
    【参数】episodes：[(demos, query, descriptor)]
    【思路】不同长度的提示不能简单地补齐：GPT 用的是可学习的绝对位置，左边补 <pad> 会改变每个词的位置。
        所以按长度分组，同样长度的一起算。
    """
    out = [None] * len(episodes)
    groups = {}
    for i, ep in enumerate(episodes):
        ids = episode_ids(*ep)
        groups.setdefault(len(ids), []).append((i, ids))
    for _, items in groups.items():
        x = torch.tensor([ids for _, ids in items])
        logits, _ = model(x)
        p = torch.softmax(logits[:, -1], dim=-1)
        for (i, _), row in zip(items, p):
            out[i] = row
    return torch.stack(out)


def entropy_bits(dist):
    """一个 {标签: 概率} 分布的熵（比特）。"""
    return -sum(p * math.log2(p) for p in dist.values() if p > 0)


def make_episodes(kind, k, n, seed, descriptor=False):
    """
    造 n 个评估用的"例子 + 查询"。
    【参数】kind："natural" / "symbolic"；k：例子个数；descriptor=True 时（只对自然任务有意义）带上最常见的说明
    【返回】字典列表，每个有 demos、query、descriptor、label（查询的正确标签）、task
    """
    rng = random.Random(seed)
    eps = []
    for _ in range(n):
        task = sample_task(rng, 1.0 if kind == "natural" else 0.0)
        items = [rng.randrange(16) for _ in range(k + 1)]
        desc = DESCRIPTORS[task.feature][0][0] if (descriptor and kind == "natural") else None
        eps.append({"demos": [(it, label_of(task, it)) for it in items[:k]], "query": items[k],
                    "descriptor": desc, "label": label_of(task, items[k]), "task": task})
    return eps


def evaluate(model, eps, bayes=True):
    """
    模型在一批评估题上的表现，和贝叶斯最优比较。
    【返回】字典：
        acc         模型的准确率（在整个词表上取概率最大的词）
        p_true      模型给正确标签的平均概率
        bayes_acc   贝叶斯最优的**期望**准确率 = 平均的 max_y P(y | 例子)（比"在这几百题上数对了几个"更稳定）
        bayes_p     贝叶斯后验给正确标签的平均概率
        kl          平均的 KL(贝叶斯后验 ‖ 模型)（nats）：模型的预测离"完美推理"有多远
    【边界】贝叶斯后验为 None（例子和任何任务都对不上）的题不计入 bayes_* 和 kl。
    """
    probs = query_probs(model, [(e["demos"], e["query"], e["descriptor"]) for e in eps])
    pred = probs.argmax(1)
    out = {"acc": sum(int(pred[i]) == ID[e["label"]] for i, e in enumerate(eps)) / len(eps),
           "p_true": sum(probs[i, ID[e["label"]]].item() for i, e in enumerate(eps)) / len(eps)}
    if bayes:
        ba, bp, kl, n = 0.0, 0.0, 0.0, 0
        for i, e in enumerate(eps):
            b = bayes_predict(e["demos"], e["query"], e["descriptor"])
            if b is None:
                continue
            n += 1
            ba += max(b.values())
            bp += b.get(e["label"], 0.0)
            kl += sum(p * math.log(p / max(probs[i, ID[y]].item(), 1e-12)) for y, p in b.items() if p > 0)
        out.update(bayes_acc=ba / max(n, 1), bayes_p=bp / max(n, 1), kl=kl / max(n, 1))
    return out