"""
=============================================================================
第 14 讲 · 工具模块：一个用纯 Python 写的迷你语言模型 —— 预测下一个词、梯度下降、生成
=============================================================================

对应讲义：lectures/lecture14_training_llms.md 第 1、2 节

-----------------------------------------------------------------------------
【这个模型是什么？】
-----------------------------------------------------------------------------
真实的 LLM 是一个几十亿参数的 Transformer。这里用一个小得多、但**训练目标完全相同**的模型：

    给定前面的词，预测下一个词的概率分布；训练 = 让训练文本里"实际出现的下一个词"的概率变大。

模型是一个"对数线性"模型（也叫多类逻辑回归、softmax 回归）：

    特征（从上下文里提取）：
        p1=前一个词、p2=前两个词（局部的语法："was born in" 后面是城市）
        bag=这段文字里到目前为止出现过的"内容词"（名字、"born"、"plays"……）—— 让模型能把"mira"和
            "veloria"联系起来（真实的 Transformer 用注意力机制做到这一点）
    打分：每个候选词 v 的分数 = b[v] + Σ（每个特征 f 的权重 W[f][v]）
    概率：softmax(分数)
    损失：交叉熵 = −log P(实际的下一个词)
    梯度：对分数的梯度 = 概率 − 独热向量（正确的词减 1）—— 正确的词的分数被推高，其他词被压低

它没有 Transformer 的表达能力，但足以演示本讲的现象：预训练学到知识、指令微调学到格式、蒸馏。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from tiny_lm import TinyLM
    lm = TinyLM()
    lm.train(texts, epochs=10)                       # 在每个位置上预测下一个词
    lm.generate("question : where was mira born ? answer :")
"""

import math
import random

EOS = "<eos>"
# 不算"内容词"的常见词（不放进 bag 特征）
STOP = {"was", "in", "as", "a", "the", "is", ".", ",", "?", ":", "by", "of", "up", "can", "does", "what", "where",
        "which", "question", "answer", EOS, "on", "all", "this", "every", "last", "near", "early", "new", "old"}


def tokenize(text):
    """按空格切分成词（语料里标点前后都已经加了空格）。"""
    return text.split()


def features(tokens, i):
    """
    第 i 个位置（预测 tokens[i]）可以用的特征：只能看 tokens[0..i-1]。

    【返回】特征名的列表，例如 ["p1=in", "p2=born in", "bag=mira", "bag=born"]
    【边界】序列开头没有前一个词时用 "<s>" 代替
    """
    p1 = tokens[i - 1] if i >= 1 else "<s>"
    p2 = tokens[i - 2] if i >= 2 else "<s>"
    feats = ["bias", f"p1={p1}", f"p2={p2} {p1}"]
    seen = set()
    for t in tokens[:i]:
        if t not in STOP and t not in seen:
            seen.add(t)
            feats.append(f"bag={t}")
    return feats


class TinyLM:
    """
    迷你语言模型。

    【参数】lr：学习率；seed：随机种子（只影响训练时样本的顺序）
    【属性】vocab（词表）、W（特征 → 每个词的权重列表）
    """

    def __init__(self, lr=0.3, seed=0):
        self.lr, self.rng = lr, random.Random(f"lm|{seed}")
        self.vocab, self.index = [], {}
        self.W = {}
        self.updates = 0

    def _add_words(self, tokens):
        """把新词加进词表；已有的权重向量在末尾补 0。"""
        for t in tokens:
            if t not in self.index:
                self.index[t] = len(self.vocab)
                self.vocab.append(t)
                for row in self.W.values():
                    row.append(0.0)

    def probs(self, feats):
        """
        给定特征，下一个词的概率分布（softmax）。

        【Python 小知识】减去最大值再取 exp，避免数值溢出（数学上结果不变）。求和用 math.fsum：
            Python 3.12 起内置的 sum() 对小数做了补偿求和，和旧版本差最后一位；训练要做几万次更新，
            最后一位的差别可能累积起来改变某次生成的结果。math.fsum 在所有版本上结果相同。
        """
        V = len(self.vocab)
        scores = [0.0] * V
        for f in feats:
            row = self.W.get(f)
            if row is not None:
                scores = [s + w for s, w in zip(scores, row)]
        m = max(scores)
        exps = [math.exp(s - m) for s in scores]
        z = math.fsum(exps)
        return [e / z for e in exps]

    def step(self, feats, target):
        """
        一次梯度下降：在这个位置上，让 target 的概率变大。

        【返回】这个位置的损失 −log P(target)
        【边界】没见过的特征会新建一行权重（全 0）
        """
        p = self.probs(feats)
        y = self.index[target]
        loss = -math.log(max(p[y], 1e-12))
        p[y] -= 1.0                                    # 梯度 = 概率 − 独热向量
        lr = self.lr
        for f in feats:
            row = self.W.get(f)
            if row is None:
                row = self.W[f] = [0.0] * len(self.vocab)
            self.W[f] = [w - lr * g for w, g in zip(row, p)]
        self.updates += 1
        return loss

    def train(self, texts, epochs=5, mask_prefix=None):
        """
        在一组文本上训练若干轮。每段文本末尾加上 <eos>，在每个位置上预测下一个词。

        【参数】mask_prefix：一个函数 text → 前缀的词数。只在前缀之后的位置计算损失（指令微调时
            只学"答案"部分，不学"问题"部分；None 表示所有位置都学）
        【返回】每一轮的平均损失（列表）
        """
        seqs = []
        for t in texts:
            toks = tokenize(t) + [EOS]
            self._add_words(toks)
            start = mask_prefix(t) if mask_prefix else 1
            seqs.append((toks, max(1, start)))
        history = []
        for _ in range(epochs):
            self.rng.shuffle(seqs)
            total, n = 0.0, 0
            for toks, start in seqs:
                for i in range(start, len(toks)):
                    total += self.step(features(toks, i), toks[i])
                    n += 1
            history.append(total / max(1, n))
        return history

    def next_dist(self, tokens):
        """接在 tokens 后面的下一个词的概率分布：[(词, 概率)]，从大到小。"""
        p = self.probs(features(tokens + ["?"], len(tokens)))
        return sorted(zip(self.vocab, p), key=lambda vp: -vp[1])

    def generate(self, prompt, max_new=12):
        """
        贪心生成：每次取概率最大的词，直到 <eos> 或者达到 max_new 个词。

        【返回】生成的词（字符串，不含提示部分）
        """
        toks = tokenize(prompt)
        for t in toks:
            if t not in self.index:
                self._add_words([t])
        out = []
        for _ in range(max_new):
            word = self.next_dist(toks)[0][0]
            if word == EOS:
                break
            out.append(word)
            toks.append(word)
        return " ".join(out)

    def loss_on(self, texts):
        """在一组文本上的平均交叉熵（不更新参数）。"""
        total, n = 0.0, 0
        for t in texts:
            toks = tokenize(t) + [EOS]
            for i in range(1, len(toks)):
                if toks[i] in self.index:
                    p = self.probs(features(toks, i))[self.index[toks[i]]]
                    total += -math.log(max(p, 1e-12))
                    n += 1
        return total / max(1, n)

    def copy(self, seed=0):
        """
        复制一个模型（继续训练副本，不影响原来的模型）。

        【参数】seed：副本训练时打乱样本顺序用的随机种子。固定下来，同样的副本、同样的数据总是得到同样的结果
            （不管之前复制过几次）。
        """
        m = TinyLM(self.lr, seed=f"copy|{seed}")
        m.vocab, m.index = list(self.vocab), dict(self.index)
        m.W = {f: list(row) for f, row in self.W.items()}
        return m