"""
=============================================================================
CS525 · 第 8 讲模块：一个"知道真相"的指令世界（instruct_world.py）
=============================================================================

对应讲义：lectures/lecture08_posttraining_data_i.md

-----------------------------------------------------------------------------
【为什么要造这个世界】
-----------------------------------------------------------------------------
后训练（post-training）要教会模型两件事：一是"听指令"——看到一个指令，去完成它，而不是接着往下续写；
二是"做得好"——完成得正确、有帮助、无害。真实的研究里，这两件事纠缠在一起，很难分开看。
这个模块把它们拆开：

    任务         12 个"任务簇"（像 FLAN 的翻译、问答、摘要、情感分析……），每簇 4 种任务，共 48 种任务。
    指令         每条指令是一个 24 维向量 = 一个所有指令共有的"这是一条指令"方向 u
                                         + 它所在任务簇的方向 + 任务自己的特点 + 措辞（模板）的噪声。
                 普通文本（"请续写我"）是没有 u 方向的随机向量。
    基础能力     预训练模型对每种任务有一个"会不会做"的能力 s（0–1）：如果它去做，做对的概率是 s。
                 大模型的能力高。指令微调主要改变的是"去不去做"，对能力本身改变不大（"表面对齐假说"）。
    学生模型     "去不去做" = sigmoid(w · 指令向量 + b)：一个逻辑回归，从指令微调的数据里学；
                 能力 s：指令微调的示范如果是对的，s 会往示范的正确率挪一点（挪多少由示范数量决定）。

    这样就能研究：指令数据覆盖了多少种任务（FLAN）？示范的质量和数量哪个重要（LIMA）？
    用模型自己生成的指令行不行（Self-Instruct、Alpaca）？

【运行自检】
    python3 instruct_world.py        # 打印基础模型和一个指令微调后的模型在各任务簇上的表现
"""

import numpy as np

DIM = 24


class InstructWorld:
    """
    指令世界：任务簇、任务、指令向量、普通文本、两个尺寸的预训练模型的基础能力。

    【参数】seed；n_clusters：任务簇数；per_cluster：每簇的任务数；
            common：所有指令共有的"指令方向"有多强（越强越容易举一反三）；
            cluster_spread：不同任务簇之间差别多大；task_spread：同一簇里任务之间差别多大；template_noise：措辞的噪声
    【属性】task_vec（48 × 24）、cluster_of（每个任务属于哪个簇）、skill（字典：模型尺寸 → 48 个基础能力）
    """

    def __init__(self, seed=0, n_clusters=12, per_cluster=4, common=3.0, cluster_spread=1.5, task_spread=0.7,
                 template_noise=0.6):
        rng = np.random.default_rng(seed)
        u = rng.normal(size=DIM)
        self.u = u / np.linalg.norm(u)
        centers = rng.normal(size=(n_clusters, DIM)) * cluster_spread
        self.cluster_of = np.repeat(np.arange(n_clusters), per_cluster)
        self.task_vec = common * self.u + centers[self.cluster_of] + rng.normal(size=(len(self.cluster_of), DIM)) * task_spread
        self.template_noise = template_noise
        self.n_tasks = len(self.cluster_of)
        self.n_clusters = n_clusters
        # 基础能力：同一个任务，大模型比小模型更会做。用 logit 表示：大模型 = 小模型 + 1.6
        base_logit = rng.normal(-0.2, 0.8, size=self.n_tasks)
        self.skill = {"small": 1 / (1 + np.exp(-base_logit)), "large": 1 / (1 + np.exp(-(base_logit + 1.6)))}

    def instructions(self, tasks, rng):
        """
        为给定的任务列表各生成一条指令（措辞随机）。

        【参数】tasks：任务编号的数组；rng
        【返回】(len(tasks), 24) 的指令向量
        """
        tasks = np.asarray(tasks)
        return self.task_vec[tasks] + rng.normal(size=(len(tasks), DIM)) * self.template_noise

    def plain_text(self, n, rng):
        """n 段普通文本（不是指令，应该续写而不是"回答"）。"""
        return rng.normal(size=(n, DIM)) * 1.5

    def tasks_in(self, clusters):
        """这些任务簇里的全部任务编号。"""
        return np.flatnonzero(np.isin(self.cluster_of, clusters))


class Student:
    """
    一个可以被"后训练"的模型：会不会去做（follow）× 做得对不对（skill）。

    【参数】world；size："small" 或 "large"（决定基础能力）；base_follow_logit：预训练模型"碰巧去做"的倾向
            （预训练数据里也有问答，所以基础模型不是完全不听指令——但大多数时候在续写）
    【属性】w, b：去不去做的逻辑回归；skill：48 个任务上的能力（会被示范改变）
    """

    def __init__(self, world, size="small", base_follow_logit=-1.5):
        self.world = world
        self.w = np.zeros(DIM)
        self.b = base_follow_logit
        self.skill = world.skill[size].copy()

    def follow_prob(self, X):
        """看到这些输入时去"回答"的概率。"""
        return 1 / (1 + np.exp(-(X @ self.w + self.b)))

    def sft(self, X, tasks, correct, rng, n_plain=None, l2=1e-2, steps=500, lr=0.5, skill_rate=0.3, kappa=200.0):
        """
        监督微调（SFT）：在 (指令, 示范回答) 上训练。

        【参数】X：指令向量 (n, 24)；tasks：每条指令的任务编号；correct：每条示范回答对不对（布尔数组）；rng；
                n_plain：同时混入多少段普通文本当"不要回答，要续写"的例子（默认和指令一样多）；
                l2 / steps / lr：逻辑回归的设置；skill_rate、kappa：示范对能力的影响有多大、要多少条示范才明显
        【返回】self
        【思路】(1) "去不去做"：指令是正例、普通文本是负例，训练逻辑回归——模型学会认出"这是一条指令"。
                (2) "做得好不好"：每种任务，能力往示范的正确率挪一步：
                        s ← s + skill_rate · (1 − e^{−n/kappa}) · (示范正确率 − s)
                    示范越多挪得越多，但 skill_rate 限制了上限——SFT 主要教格式和意愿，教不了多少新能力。
                    注意：如果示范的正确率低于模型本来的能力，能力会**下降**（模仿一个更弱的老师）。
        【边界】没有任何示范的任务，能力不变。
        """
        n_plain = len(X) if n_plain is None else n_plain
        Xn = self.world.plain_text(n_plain, rng)
        A = np.vstack([X, Xn])
        y = np.r_[np.ones(len(X)), np.zeros(n_plain)]
        w, b = self.w.copy(), self.b
        for _ in range(steps):
            p = 1 / (1 + np.exp(-(A @ w + b)))
            w -= lr * (A.T @ (p - y) / len(y) + l2 * w)
            b -= lr * np.mean(p - y)
        self.w, self.b = w, b
        tasks = np.asarray(tasks)
        correct = np.asarray(correct, dtype=float)
        for t in np.unique(tasks):
            m = tasks == t
            q = correct[m].mean()
            self.skill[t] += skill_rate * (1 - np.exp(-m.sum() / kappa)) * (q - self.skill[t])
        return self

    def evaluate(self, tasks, rng, n_per=40):
        """
        在给定任务上评测：每种任务 n_per 条新措辞的指令。

        【返回】字典：follow（去回答的比例）、success（回答而且答对的比例）、plain_follow（把普通文本当指令的比例，越低越好）
        """
        tasks = np.repeat(np.asarray(tasks), n_per)
        X = self.world.instructions(tasks, rng)
        pf = self.follow_prob(X)
        succ = pf * self.skill[tasks]
        pn = self.follow_prob(self.world.plain_text(500, rng))
        return {"follow": float(pf.mean()), "success": float(succ.mean()), "plain_follow": float(pn.mean())}


if __name__ == "__main__":
    W = InstructWorld()
    rng = np.random.default_rng(1)
    held = W.tasks_in([9, 10, 11])
    for size in ["small", "large"]:
        s = Student(W, size)
        print(size, "基础模型（留出的任务簇）：", {k: round(v, 3) for k, v in s.evaluate(held, rng).items()})
    s = Student(W, "small")
    tr = W.tasks_in(range(9))
    t = rng.choice(tr, size=2000)
    s.sft(W.instructions(t, rng), t, rng.random(len(t)) < 0.9, rng)
    print("small 指令微调后（留出的任务簇）：", {k: round(v, 3) for k, v in s.evaluate(held, rng).items()})