"""
=============================================================================
CMU 17-630 · 第 2 讲共享模块：icl_lab.py —— 上下文学习（in-context learning）的实验台
=============================================================================

【这个模块是做什么的】
    第 2 讲研究同一个问题的很多个侧面：**参数不变，只往提示里放几个示例，模型为什么就"会"了一个任务？
    它到底从示例里学到了什么？为什么换几个示例、换个顺序，准确率就大起大落？**
    这个模块把所有实验共用的零件放在一起：

        数据      test_set, DemoPool, sample_demos       第 1 讲世界里的评论：情感二分类（positive / negative）
        模板      TEMPLATES, VERBALIZERS, render         提示怎么写：输入格式、标签字段、标签词（verbalizer）
        打分      label_probs, accuracy, by_kind          比较两个标签词作为"下一个词"的概率（批量计算，很快）
        校准      calibrate, CONTENT_FREE                 上下文校准（Zhao et al. 2021）与批量校准（Zhou et al. 2024）
        模型      get_base_model, get_bursty_model        两个真实的迷你 GPT：第 1 讲的 base 模型，
                                                         和一个只改了"评论页"写法的对照模型（见下）
        模拟器    SimICL                                  一个"大模型"的玩具模型：每种偏差都是一个旋钮

【两个真实的迷你 GPT】
    base   ：第 1 讲训练的那个。它的语料里，同一个评论页上的几条评论**互相独立**：上一条是好评，
             下一条是好评的概率仍然是一半。
    bursty ：除了评论页之外，语料和 base 一模一样（同一个世界、同样的文档配方、同样的训练步数），
             只是每个评论页属于一家"好店"或"差店"：好店的评论 85% 是好评，差店的 85% 是差评。
             于是在这份语料里，**前面几条评论的评分，是猜下一条评分的有用线索**。
             （"bursty"一词借自 Chan et al. 2022：训练数据里同类的东西成簇出现，是上下文学习能力出现的条件之一。）
    两个模型的对比回答一个问题：模型会不会"读"示例里的标签，是不是由预训练数据决定的？

【SimICL：一个"大模型"的玩具模型】
    迷你 GPT 只有两层、26 万个参数，有些只在大模型上出现的现象（例如被翻转的标签带着走，Wei et al. 2023）
    它表现不出来。SimICL 用一个公式代替神经网络：模型给"第一个标签词"的打分
        Δ = (语义先验 + 从示例里学到的映射) × 输入本身的证据 + 各种偏差（模板、常见词、多数标签、近因、顺序噪声）
    每一项都有一个旋钮。它能演示"如果真实模型是这样的，会看到什么"，但**不能证明**真实模型就是这样；
    讲义里每次用它都会给出论文里真实模型的数字作对照。

【运行】python3 icl_lab.py     # 自检：打印几个提示和它们的打分；第一次运行会训练 bursty 模型（约 45 秒）
【依赖】PyTorch；import 第 1 讲的 pretrain_world.py、tiny_gpt.py、console.py，前置讲 P2 的 llm_api.py（SimICL 的接口）。
"""

import hashlib
import json
import math
import os
import random
import re
import sys
import time
from dataclasses import dataclass

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "prereq2", "prereq3"):
    _p = os.path.join(HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)

import pretrain_world as pw  # noqa: E402
from llm_api import LLM, LLMResponse, Usage, count_request_tokens, validate_messages  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")
CLASSES = ("positive", "negative")      # 两个**类别**的名字（和"标签词"不是一回事，见 VERBALIZERS）


# ---------------------------------------------------------------------------
# 1. 数据：评论情感二分类
# ---------------------------------------------------------------------------

def test_set():
    """
    测试集：300 条新生成的评论 [(正文, 类别, 难度类型), ...]，和第 1 讲 §8 用的是同一批（seed=7）。

    【难度类型】simple（同极性的一两个分句）、but（以 but 后面为准）、not（带否定）——见 pretrain_world.make_review。
    【边界】类别不完全平衡：138 条 positive、162 条 negative（"总是猜 negative"能得 54%，"总是猜 positive"得 46%）。
    """
    return pw.review_probes(seed=7, n=300)


@dataclass
class DemoPool:
    """
    示例池：示例从这里抽，和测试集不重叠（不同的 seed）。

    【字段】all：全部 400 条；pos / neg：只含 1 个分句（≤ 4 个词）的短评论，按类别分开。
    【为什么示例只用短评论】迷你 GPT 的上下文窗口只有 64 个 token。一条短示例
        "review : the soup was cold . sentiment : negative ." 是 11 个 token，4 条示例 + 页眉 + 测试评论
        正好塞得下；用长评论的话，4 条示例就会把最前面的内容挤出窗口（01 号脚本演示 3）。
        真实模型的窗口有几十万个 token，但"示例占地方"这件事一样存在——只是换成了费用和延迟。
    """
    all: list
    pos: list
    neg: list


def demo_pool():
    """生成示例池（seed=11，400 条；其中短评论 126 条）。"""
    pool = pw.review_probes(seed=11, n=400)
    short = [r for r in pool if len(r[0].split()) <= 4]
    return DemoPool(all=pool, pos=[r for r in short if r[1] == "positive"],
                    neg=[r for r in short if r[1] == "negative"])


def sample_demos(pool, rng, k=None, pattern=None):
    """
    抽一组示例。

    【参数】pool：DemoPool；rng：random.Random；
        pattern：标签顺序，例如 "PPNN" 表示"好、好、差、差"（P = positive，N = negative）；
        k：不给 pattern 时，抽 k 个、两类各半（k 为奇数时多一个 negative），然后随机打乱顺序。
    【返回】[(正文, 类别, 难度类型), ...]，按出现在提示里的顺序
    【边界】同一组里不会重复抽到同一条；pattern 里 P 的个数不能超过 pool.pos 的大小。
    """
    if pattern is None:
        demos = rng.sample(pool.pos, k // 2) + rng.sample(pool.neg, k - k // 2)
        rng.shuffle(demos)
        return demos
    pos = iter(rng.sample(pool.pos, pattern.count("P")))
    neg = iter(rng.sample(pool.neg, pattern.count("N")))
    return [next(pos) if c == "P" else next(neg) for c in pattern]


def flip(c):
    """把类别翻过来：positive ↔ negative。"""
    return "negative" if c == "positive" else "positive"


# ---------------------------------------------------------------------------
# 2. 模板与标签词：迷你 GPT 的"提示词"
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class Template:
    """
    一个提示模板：每条示例（以及最后的测试输入）按 fmt 写出来，{x} 是评论正文，{y} 是标签词。

    【字段】name：名字；fmt：格式；note：一句话说明它和训练语料的关系
    【例子】Template("sentiment", "review : {x} . sentiment : {y} .")
        一条示例："review : the soup was cold . sentiment : negative ."
        测试输入："review : the food was great . sentiment :"   ← 模型接着写的那个词就是它的回答
    """
    name: str
    fmt: str
    note: str = ""

    def demo(self, x, y):
        """一条完整的示例。"""
        return self.fmt.format(x=x, y=y)

    def query(self, x):
        """测试输入：模板里 {y} 之前的部分（模型从这里开始续写）。"""
        return self.fmt.split("{y}")[0].format(x=x).strip()


# 迷你 GPT 的词表只有 290 个词（全部来自第 1 讲的语料），模板里的词表外词会变成 <unk>。
TEMPLATES = {
    "rating": Template("rating", "review : {x} . rating : {y} .", "和训练语料的评论页一模一样"),
    "sentiment": Template("sentiment", "review : {x} . sentiment : {y} .", "标签字段换成词表外的 sentiment（变成 <unk>）"),
    "qa": Template("qa", "q : {x} ? a : {y} .", "借用问答页的格式"),
    "arrow": Template("arrow", "{x} => {y} .", "箭头（=> 在词表外）"),
    "bare": Template("bare", "{x} . {y} .", "没有任何标记，评论后面直接跟标签"),
}

# 标签词（verbalizer）：每个类别用哪个词来"说出来"。模型的回答 = 两个标签词里概率更大的那个。
VERBALIZERS = {
    "positive/negative": {"positive": "positive", "negative": "negative"},
    "great/awful": {"positive": "great", "negative": "awful"},        # 语义一致，但训练时从没出现在标签位置
    "awful/great": {"positive": "awful", "negative": "great"},        # 语义反过来：好评要说 awful
    "copper/wheat": {"positive": "copper", "negative": "wheat"},      # 和情感无关的词（"语义无关标签"）
}

HEADER = "<doc> reviews :"     # 文档页眉：告诉模型"这是一个评论页"


def render(template, verbalizer, demos, x, header=HEADER, style="full"):
    """
    拼出一个完整的提示。

    【参数】template：Template；verbalizer：{类别: 标签词}；demos：[(正文, 类别, 类型), ...]（类别可以是被故意改过的）；
        x：测试评论；header：页眉（"" 表示不要）；
        style："full" 每条示例都有输入和标签；"input_only" 示例只有输入（没有标签）；
              "label_only" 示例只有标签（没有输入）——04 号脚本拆解"示例到底教了什么"时用
    【返回】字符串，例如 "<doc> reviews : review : the soup was cold . sentiment : negative . review : ... sentiment :"
    【Python 小知识】" ".join(列表) 用空格把字符串连起来；re.sub(r"\\s+", " ", s) 把连续的空白压成一个空格。
    """
    parts = [header] if header else []
    for dx, dy, _ in demos:
        if style == "full":
            parts.append(template.demo(dx, verbalizer[dy]))
        elif style == "input_only":
            parts.append(template.query(dx) + " .")
        elif style == "label_only":
            parts.append(template.demo("", verbalizer[dy]))
        else:
            raise ValueError(f"未知的 style：{style}")
    parts.append(template.query(x))
    return re.sub(r"\s+", " ", " ".join(parts)).strip()


# ---------------------------------------------------------------------------
# 3. 打分：两个标签词谁的概率大
# ---------------------------------------------------------------------------

def label_probs(model, tok, prompts, words, return_mass=False):
    """
    对很多个提示，计算"下一个词是 words[0] / words[1]"的概率，并在两者之间归一化。

    【参数】model, tok：迷你 GPT 和它的分词器；prompts：提示列表；words：两个标签词（必须在词表里）；
        return_mass：True 时另外返回 p(words[0]) + p(words[1])（模型有多大的把握"在这里写一个标签词"）
    【返回】[[p0, p1], ...]（每行和为 1），或者 (这个列表, mass 列表)
    【思路】逐个提示调用模型太慢（几千次前向）。这里把**长度相同**的提示叠成一个批次，一次前向算完：
        因果注意力只看左边，所以只要长度一样，叠在一起和分开算的结果完全相同；不需要填充（padding）。
    【边界】提示超过 64 个 token 时只保留最后 64 个（和第 1 讲 tiny_gpt 的做法一致，最前面的内容被截掉）。
    【提示工程联系】这就是 Zhao et al.（2021）、Min et al.（2022）等论文用的"比较标签词概率"的分类法。
        真实的 API 大多不返回概率（P2 §3.6），只能看模型生成了哪个词——07 号脚本就是这样做的。
    """
    import torch
    import torch.nn.functional as F
    ids = [tok.encode(p)[-model.cfg.block_size:] for p in prompts]
    wi = [tok.stoi[w] for w in words]
    norm, mass = [None] * len(prompts), [None] * len(prompts)
    groups = {}
    for i, x in enumerate(ids):
        groups.setdefault(len(x), []).append(i)          # 【Python 小知识】setdefault：键不存在时先放一个空列表
    with torch.no_grad():
        for _, idx in sorted(groups.items()):
            logits = model(torch.tensor([ids[i] for i in idx]))[0][:, -1]
            p = F.softmax(logits, dim=-1)[:, wi]
            for j, i in enumerate(idx):
                a, b = p[j].tolist()
                s = a + b
                norm[i] = [a / s, b / s] if s > 0 else [0.5, 0.5]
                mass[i] = s
    return (norm, mass) if return_mass else norm


def predictions(probs):
    """概率 → 预测的类别：probs[i][0] 是"第一个类别"（positive）的归一化概率。"""
    return [CLASSES[0] if p[0] > p[1] else CLASSES[1] for p in probs]


def accuracy(probs, items):
    """准确率：items 是 [(正文, 类别, 类型), ...]，和 probs 一一对应。"""
    return sum(pred == y for pred, (_, y, _) in zip(predictions(probs), items)) / len(items)


def correct_vector(probs, items):
    """每道题对不对（0/1 列表）——做配对检验（McNemar）时要用。"""
    return [1 if pred == y else 0 for pred, (_, y, _) in zip(predictions(probs), items)]


def by_kind(probs, items):
    """按难度类型分别算准确率：{"simple": ..., "but": ..., "not": ...}。"""
    out = {}
    for kind in ("simple", "but", "not"):
        sel = [(p, it) for p, it in zip(probs, items) if it[2] == kind]
        out[kind] = accuracy([p for p, _ in sel], [it for _, it in sel]) if sel else float("nan")
    return out


def mean_positive(probs):
    """平均的 P(positive)：一个提示整体上"偏向"哪个标签。测试集里 positive 占 46%。"""
    return sum(p[0] for p in probs) / len(probs)


# ---------------------------------------------------------------------------
# 4. 校准
# ---------------------------------------------------------------------------

# "无内容"的输入：放在测试评论的位置上，用来测量"不看输入时，模型偏向哪个标签"。
# Zhao et al.（2021）用的是 "N/A"、空字符串、"[MASK]"；迷你 GPT 的词表里没有这些词，
# "N/A" 进来就是 <unk>，所以这里用 <unk>。03 号脚本演示 3 会看到：换一个"无内容"输入，结果可能差很多。
CONTENT_FREE = ["<unk>"]


def calibrate(probs, ref):
    """
    校准：每个类别的概率除以一个"参考概率"，再重新归一化。

    【参数】probs：[[p0, p1], ...]；ref：[r0, r1]
    【公式】p̂_c ∝ p_c / r_c。等价于在对数空间里减掉 log r_c：如果模型的偏差是"给某个标签词整体加了一个常数分"，
        这一步正好把它减掉（Zhao et al. 2021 写成 W = diag(r)^{-1} 的仿射变换）。
    【ref 从哪里来】
        上下文校准（contextual calibration）：ref = 把"无内容"输入放进同一个提示时的概率；
        批量校准（batch calibration，Zhou et al. 2024）：ref = 整批测试输入的平均概率
        （假设"平均来看，模型不该偏向任何一类"——测试集严重不平衡时这个假设不成立）。
    【边界】ref 某一项接近 0 时，除法会把它放大很多倍——无内容输入的概率很极端时，校准会"矫枉过正"。
    """
    out = []
    for a, b in probs:
        a2, b2 = a / max(ref[0], 1e-12), b / max(ref[1], 1e-12)
        s = a2 + b2
        out.append([a2 / s, b2 / s])
    return out


def mean_probs(probs):
    """一组概率的平均：[[p0, p1], ...] → [mean p0, mean p1]。批量校准的参考值。"""
    n = len(probs)
    return [sum(p[0] for p in probs) / n, sum(p[1] for p in probs) / n]


# ---------------------------------------------------------------------------
# 5. 两个真实的迷你 GPT
# ---------------------------------------------------------------------------

def get_base_model():
    """第 1 讲的 base 模型（读缓存；没有缓存时训练约 50 秒）。返回 (model, tok)。"""
    import torch
    from tiny_gpt import get_model
    torch.set_num_threads(1)
    model, tok, _ = get_model("base", verbose=False)
    return model, tok


def review_with_polarity(rng, polarity):
    """
    生成一条**指定极性**的评论（pretrain_world.make_review 的极性是随机的，这里要能指定）。

    【参数】polarity：+1 好评 / −1 差评
    【返回】正文（三种难度类型的比例和 make_review 一样：60% simple、20% but、20% not）
    """
    r = rng.random()
    if r < 0.6:
        return " and ".join(pw._clause(rng, polarity, False)[0] for _ in range(rng.randint(1, 2)))
    if r < 0.8:
        return pw._clause(rng, -polarity, False)[0] + " but " + pw._clause(rng, polarity, False)[0]
    parts = [pw._clause(rng, polarity, True)[0]]
    if rng.random() < 0.5:
        parts.append(pw._clause(rng, polarity, False)[0])
        rng.shuffle(parts)
    return " and ".join(parts)


def bursty_reviews_page(rng, q=0.85):
    """
    一个"成簇"的评论页：先决定这是一家好店还是差店（各一半），好店的每条评论以概率 q 是好评，差店以概率 q 是差评。

    【和 pretrain_world.make_reviews_page 的唯一区别】那边每条评论的极性独立地随机；这里同一页的评论共享一个"店的好坏"。
        页面长度（2–4 条）、评论的写法、格式都一样。
    【提示工程联系】真实网页也是这样：同一家店的评论、同一个话题的帖子、同一份试卷的答案，彼此相关。
        Xie et al.（2022）把上下文学习解释成"推断这篇文档的隐含概念"：这里的隐含概念就是"店的好坏"。
    """
    good = rng.random() < 0.5
    items = []
    for _ in range(rng.randint(2, 4)):
        pol = 1 if rng.random() < (q if good else 1 - q) else -1
        items.append(f"review : {review_with_polarity(rng, pol)} . rating : {pw.LABELS[pol]} .")
    return "reviews : " + " ".join(items)


def build_bursty_corpus(world, seed=0, scale=1.0, q=0.85):
    """
    和 pretrain_world.build_corpus 相同的配方，只把评论页换成 bursty_reviews_page。

    【返回】Document 列表（顺序打乱）
    【边界】因为随机数的消耗顺序变了，百科、问答、故事的**具体句子**和 base 的语料不同，但分布相同。
    """
    rng = random.Random(1000 + seed)
    docs = []
    n_art = max(1, round(pw.CORPUS_MIX["article"] * scale))
    for c in world.countries:
        docs += [pw.Document("article", pw.make_article(rng, c)) for _ in range(n_art)]
    makers = [("faq", lambda: pw.make_faq(rng, world.seen)),
              ("quiz", lambda: pw.make_quiz(rng, world.countries)),
              ("reviews", lambda: bursty_reviews_page(rng, q)),
              ("story", lambda: pw.make_story(rng))]
    for kind, maker in makers:
        n = max(1, round(pw.CORPUS_MIX[kind] * scale))
        docs += [pw.Document(kind, maker()) for _ in range(n)]
    rng.shuffle(docs)
    return docs


BURSTY_RECIPE = dict(world_seed=0, corpus_seed=0, q=0.85, block_size=64, n_layer=2, n_head=4, d_model=96,
                     steps=3000, batch_size=32, lr=3e-3, seed=0)


def get_bursty_model(verbose=True):
    """
    取 bursty 模型：有缓存就读缓存（code/lecture02/checkpoints/），没有就训练（约 45 秒）并存下来。

    【返回】(model, tok)
    【边界】要求它的词表和 base 模型完全一样（两份语料用的词相同），这样两个模型可以用同一个分词器、
        同一批提示比较。不一样时会报错（说明 pretrain_world.py 被改过）。
    【和 tiny_gpt.get_model 的关系】结构、超参数、训练循环全部复用第 1 讲的代码，只换了语料。
    """
    import torch
    from tiny_gpt import GPTConfig, TinyGPT, _recipe_hash, train_model
    torch.set_num_threads(1)
    r = BURSTY_RECIPE
    world = pw.build_world(seed=r["world_seed"])
    train_text = pw.corpus_text(build_bursty_corpus(world, seed=r["corpus_seed"], q=r["q"]))
    tok = pw.WordTokenizer(train_text)
    _, base_tok = get_base_model()
    if tok.itos != base_tok.itos:
        raise RuntimeError("bursty 语料的词表和 base 模型不同：pretrain_world.py 是不是被改过？")
    cfg = GPTConfig(vocab_size=len(tok), block_size=r["block_size"], n_layer=r["n_layer"],
                    n_head=r["n_head"], d_model=r["d_model"])
    path = os.path.join(CKPT_DIR, f"bursty-{_recipe_hash(r, train_text)}.pt")
    torch.manual_seed(r["seed"])
    model = TinyGPT(cfg)
    if os.path.exists(path):
        ckpt = torch.load(path, weights_only=True)
        model.load_state_dict(ckpt["state_dict"])
        model.eval()
        return model, tok
    eval_text = pw.corpus_text(build_bursty_corpus(world, seed=r["corpus_seed"] + 1, scale=0.2, q=r["q"]))
    train_ids = torch.tensor(tok.encode(train_text), dtype=torch.long)
    eval_ids = torch.tensor(tok.encode(eval_text), dtype=torch.long)
    if verbose:
        print(f"  训练 bursty 模型（第一次运行，约 45 秒；之后读缓存）：{len(train_ids):,} 个训练 token")
    t0 = time.time()
    hist = train_model(model, train_ids, steps=r["steps"], batch_size=r["batch_size"], lr=r["lr"],
                       seed=r["seed"], eval_ids=eval_ids, log_every=1000, verbose=verbose)
    os.makedirs(CKPT_DIR, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "vocab": tok.itos, "history": hist,
                "train_seconds": time.time() - t0, "recipe": r}, path)
    if verbose:
        print(f"  已保存到 {os.path.relpath(path, HERE)}")
    model.eval()
    return model, tok


# ---------------------------------------------------------------------------
# 6. SimICL：一个"大模型"做上下文学习的玩具模型
# ---------------------------------------------------------------------------

def _hash_unit(*parts):
    """把若干个值确定性地映射到 [0, 1)（同样的输入永远得到同样的数；不用内置 hash()，原因见 P2 llm_api）。"""
    h = hashlib.sha256("|".join(str(p) for p in parts).encode("utf-8")).hexdigest()
    return int(h[:12], 16) / 16 ** 12


def _hash_gauss(*parts):
    """确定性的标准正态随机数（Box–Muller 变换，用两个哈希出来的均匀数）。"""
    u1 = max(_hash_unit("g1", *parts), 1e-12)
    u2 = _hash_unit("g2", *parts)
    return math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)


def review_polarity(text):
    """
    按评论的"语法"规则读出它的极性和难度类型（SimICL 用它来"理解"输入）。

    【返回】(极性 +1 / −1 / 0, 类型 "simple" / "but" / "not" / "none")
    【规则】有 " but " 时以后半句为准；每个分句 "the X was [not] ADJ" 看形容词在褒义表还是贬义表，有 not 就反过来；
        几个分句的极性相加。读不出任何形容词（例如 "N/A"）时返回 (0, "none")。
    【边界】只认第 1 讲世界里的形容词；真实评论当然复杂得多——这正是"模拟器只能演示机制"的原因之一。
    """
    t = text.lower().strip().rstrip(".")
    kind = "but" if " but " in t else ("not" if " not " in f" {t} " else "simple")
    if " but " in t:
        t = t.split(" but ")[-1]
    score = 0
    for clause in t.split(" and "):
        words = clause.split()
        for i, w in enumerate(words):
            if w in pw.POSITIVE or w in pw.NEGATIVE:
                s = 1 if w in pw.POSITIVE else -1
                if i > 0 and words[i - 1] == "not":
                    s = -s
                score += s
    if score == 0:
        return 0, "none"
    return (1 if score > 0 else -1), kind


# 标签词的语义方向：+1 表示"好"的意思，−1 表示"坏"的意思，不在表里的词是 0（和情感无关）
WORD_SENTIMENT = {"positive": 1, "great": 1, "good": 1, "negative": -1, "awful": -1, "bad": -1, "terrible": -1}
# 标签词的"常见程度"（对数频率，虚构的数字）：越常见的词，模型越容易把它说出来（常见词偏差）
WORD_LOGFREQ = {"positive": 0.0, "negative": -0.2, "great": 0.5, "awful": -0.6, "good": 0.8, "bad": 0.3,
                "terrible": -0.7, "foo": -1.0, "bar": -1.0}


class SimICL(LLM):
    """
    一个用公式代替神经网络的"语言模型分类器"。给定一个英文的少样本提示，它给两个标签词打分：

        Δ = S(k) · e(x)  +  u(x) · B(示例)  +  b_template + b_token
        S(k) = a_prior · sem · r(k)  +  a_icl · m̂ · k/(k+2)          ← "看输入"的部分
        B(示例) = β_maj·(多数标签) + β_rec·(近因) + σ_order·ε(这组示例的这个顺序)   ← 示例带来的偏差
        P(第一个标签词) = sigmoid(Δ)

    【每一项的意思（也就是旋钮）】
        e(x)       输入本身的证据：好评为正、差评为负；simple 句最清楚（1.5），not 句次之（0.9），but 句最难（0.7），
                   再加一点每条评论自己的噪声（有些评论本来就容易看错）
        sem        标签词的语义方向：positive/negative 是 +1，awful/great（反着用）是 −1，foo/bar 是 0
        a_prior    "语义先验"有多强：模型多相信标签词本来的意思
        r(k)       任务识别：r(k) = r0 + (1 − r0)·k/(k+1)。零样本时模型只有 r0 的把握"知道要做的是情感分类"，
                   示例越多越确定（Min et al. 2022 认为示例的主要作用之一就是让模型认出任务和格式）
        m̂          从示例里估计出的"输入 → 标签"映射：示例的输入是好评、标签是第一个标签词，就 +1；反之 −1；取平均。
                   标签全对时 m̂ ≈ +1，全部翻转时 ≈ −1，随机标签时 ≈ 0
        a_icl      模型有多会"从示例里学映射"（大模型强、小模型弱——Wei et al. 2023 的发现，这里变成一个旋钮）
        多数标签   β_maj × (示例里第一个标签词的个数 − 第二个的个数) / k
        近因       β_rec × Σ γ^(距离末尾的位置) × (这条示例是不是第一个标签词，是 +1 否 −1)，γ = 0.5
        顺序噪声   σ_order × 一个由"模板 + 这串示例"决定的随机数：同一组示例换个顺序，偏向就变了
                   （Lu et al. 2022 观察到的现象；我们不知道它的机制，所以用噪声代表）
        u(x)       每条输入对示例偏差的"易感度"：1 + u_sd × 随机数（有的输入更容易被前面的标签带着走）。
                   正因为 u(x) 因输入而异，偏差不是一个纯粹的常数，校准只能减掉它的大部分
        b_template 每个模板自带的偏向；b_token 常见词偏差 c_tok × (两个标签词的对数频率之差)
        无内容输入 "N/A" 不是完全中性的：它有一点模板相关的证据 cf_sd × 随机数，易感度取 1
    【规模预设】PRESETS 里 small / medium / large 只改 a_prior 和 a_icl（Wei et al. 2023 的两条主线）。
    【边界】只能读本模块 render_english 写出来的那种提示；读不懂时返回两个标签词各 0.5。
        它不生成文字，complete() 只回答概率更大的那个标签词。
    """

    PRESETS = {
        "small": dict(a_prior=1.6, a_icl=0.3),
        "medium": dict(a_prior=1.3, a_icl=1.2),
        "large": dict(a_prior=1.0, a_icl=3.0),
    }
    DEFAULTS = dict(item_sd=0.6, template_sd=0.8, c_tok=0.8, r0=0.4, maj=0.8, rec=1.2, gamma=0.5,
                    order_sd=0.8, u_sd=0.5, cf_sd=0.4)
    CLARITY = {"simple": 1.5, "not": 0.9, "but": 0.7, "none": 0.0}

    def __init__(self, size="large", name=None, **knobs):
        self.size = size
        self.name = name or f"sim-icl-{size}"
        self.k = dict(self.DEFAULTS, **self.PRESETS[size])
        self.k.update(knobs)

    # ------------------------------------------------------------ 读提示
    @staticmethod
    def parse(prompt):
        """
        把 render_english 写出的提示拆开：(两个标签词, [(示例输入, 示例标签词), ...], 测试输入, 模板签名)。

        【格式】第一行是指令 "... as W1 or W2."；然后是若干个用空行分开的块，每块最后一行是 "字段: 标签词"，
            最后一块的标签为空（等模型来写）。块的第一行是输入（可能带 "Review: " 这样的前缀）。
        """
        blocks = [b for b in prompt.strip().split("\n\n") if b.strip()]
        m = re.search(r"as (\w+) or (\w+)", blocks[0]) if blocks else None
        if not m:
            return None
        words = (m.group(1).lower(), m.group(2).lower())
        demos, query, sig = [], None, ""
        for b in blocks[1:]:
            lines = b.strip().split("\n")
            first = re.sub(r"^\w+:\s*", "", lines[0]).strip()
            last = lines[-1]
            sig = re.sub(r"[^\n:]*$", "", b.split("\n", 1)[-1]) if "\n" in b else ""
            value = last.split(":", 1)[1].strip().lower() if ":" in last else ""
            if value:
                demos.append((first, value))
            else:
                query = first
        return words, demos, query, sig

    # ------------------------------------------------------------ 打分
    def evidence(self, text, template_sig=""):
        """输入本身的证据 e(x)。"N/A" 之类读不出极性的输入：只有模板相关的一点偏向（cf_sd）。"""
        pol, kind = review_polarity(text)
        if pol == 0:
            return self.k["cf_sd"] * _hash_gauss("cf", self.name, template_sig)
        return pol * self.CLARITY[kind] + self.k["item_sd"] * _hash_gauss("item", text)

    def delta(self, prompt):
        """
        第一个标签词相对第二个的分数 Δ（对数几率）。读不懂提示时返回 0。

        【返回】(Δ, 一个分解字典 {"signal", "context_bias", "template", "token"})
        """
        parsed = self.parse(prompt)
        if parsed is None or parsed[2] is None:
            return 0.0, {}
        (w1, w2), demos, query, sig = parsed
        k = len(demos)
        kn = self.k
        sem = (WORD_SENTIMENT.get(w1, 0) - WORD_SENTIMENT.get(w2, 0)) / 2
        # 从示例里估计映射 m̂：示例输入的极性 × (标签是不是 w1)
        m_hat = 0.0
        for dx, dy in demos:
            pol, _ = review_polarity(dx)
            m_hat += pol * (1 if dy == w1 else -1)
        m_hat = m_hat / k if k else 0.0
        recog = kn["r0"] + (1 - kn["r0"]) * k / (k + 1)
        e = self.evidence(query, sig)
        signal = (kn["a_prior"] * sem * recog + kn["a_icl"] * m_hat * k / (k + 2)) * e
        bias = 0.0
        if k:
            bias = (kn["maj"] * (sum(1 for _, y in demos if y == w1) - sum(1 for _, y in demos if y == w2)) / k
                    + kn["rec"] * sum(kn["gamma"] ** (k - 1 - i) * (1 if y == w1 else -1) for i, (_, y) in enumerate(demos))
                    + kn["order_sd"] * _hash_gauss("order", self.name, sig, tuple(demos)))
        pol, _ = review_polarity(query)
        u = 1.0 if pol == 0 else max(0.0, 1 + kn["u_sd"] * _hash_gauss("u", query))
        parts = {
            "signal": signal,
            "context_bias": u * bias,
            "template": kn["template_sd"] * _hash_gauss("tpl", self.name, sig),
            "token": kn["c_tok"] * (WORD_LOGFREQ.get(w1, -1.0) - WORD_LOGFREQ.get(w2, -1.0)),
        }
        return sum(parts.values()), parts

    def label_probs(self, prompt):
        """[P(第一个标签词), P(第二个标签词)]。"""
        d, _ = self.delta(prompt)
        p = 1 / (1 + math.exp(-d))
        return [p, 1 - p]

    def complete(self, req):
        """llm_api 的接口：回答概率更大的那个标签词（只看最后一条用户消息；system 被忽略）。"""
        validate_messages(req.messages)
        prompt = req.messages[-1]["content"]
        parsed = self.parse(prompt)
        if parsed is None:
            text = "I am not sure what you want me to label."
        else:
            p = self.label_probs(prompt)
            text = parsed[0][0] if p[0] >= p[1] else parsed[0][1]
        return LLMResponse(text=text, stop_reason="end_turn", usage=Usage(count_request_tokens(req), 1),
                           model=self.name)


# 给 SimICL（以及 07 号脚本里的真实模型）用的英文模板：{x} 是评论，{y} 是标签词
ENGLISH_TEMPLATES = {
    "review": "Review: {x}\nSentiment: {y}",
    "input": "Input: {x}\nLabel: {y}",
    "question": "{x}\nQuestion: How was the experience?\nAnswer: {y}",
}


def render_english(template_name, words, demos, x):
    """
    英文的少样本提示：一行指令 + 若干个示例块 + 测试块（标签留空）。

    【参数】template_name：ENGLISH_TEMPLATES 的键；words：(第一个类别的标签词, 第二个的)，例如 ("positive", "negative")；
        demos：[(正文, 类别, 类型), ...]；x：测试评论
    【返回】字符串。评论首字母大写、结尾加句号，让它像真实的英文。
    【例子】
        Classify the sentiment of each review as positive or negative.

        Review: The soup was cold.
        Sentiment: negative

        Review: The food was great.
        Sentiment:
    """
    tpl = ENGLISH_TEMPLATES[template_name]
    verb = {"positive": words[0], "negative": words[1]}

    def sent(t):
        return t[0].upper() + t[1:] + "."

    blocks = [f"Classify the sentiment of each review as {words[0]} or {words[1]}."]
    blocks += [tpl.format(x=sent(dx), y=verb[dy]) for dx, dy, _ in demos]
    blocks.append(tpl.format(x=sent(x), y="").rstrip())
    return "\n\n".join(blocks)


if __name__ == "__main__":
    from console import print_header
    print_header("icl_lab 自检")
    test, pool = test_set(), demo_pool()
    print(f"  测试集 {len(test)} 条（positive {sum(1 for t in test if t[1] == 'positive')} 条）；"
          f"示例池 {len(pool.all)} 条，其中短评论 positive {len(pool.pos)} 条、negative {len(pool.neg)} 条")
    rng = random.Random(0)
    demos = sample_demos(pool, rng, pattern="PN")
    p = render(TEMPLATES["sentiment"], VERBALIZERS["positive/negative"], demos, test[0][0])
    print(f"  迷你 GPT 的提示：{p}")
    base, tok = get_base_model()
    bursty, _ = get_bursty_model()
    for name, m in [("base", base), ("bursty", bursty)]:
        print(f"  {name:6s} P(positive, negative) = {[round(v, 3) for v in label_probs(m, tok, [p], ['positive', 'negative'])[0]]}"
              f"（正确答案：{test[0][1]}）")
    e = render_english("review", ("positive", "negative"), demos, test[0][0])
    print("  英文提示：\n    " + e.replace("\n", "\n    "))
    for size in SimICL.PRESETS:
        print(f"  SimICL-{size:6s} P = {[round(v, 3) for v in SimICL(size).label_probs(e)]}")
    print(f"  review_polarity('the soup was not cold and the staff was friendly') = "
          f"{review_polarity('the soup was not cold and the staff was friendly')}")