"""
neural_lm.py —— 固定窗口的神经语言模型（Bengio et al., 2003 的思路，带权重共享）
=====================================================================

【这个文件要回答的问题】
    第 1 讲的 slides 把语言模型画成一条流水线：
        分词 → 嵌入（embedding） → 神经网络 → logits → softmax → 解码
    这里实现其中最简单的一种"神经网络"：只看最近 C 个 token 的多层感知机（MLP）。

        输入：最近 C 个 token 的 id                      形状 (B, C)
          │  嵌入矩阵 E 查表（每个 id 取 E 的一行）
          ▼
        C 个嵌入向量拼成一个长向量                       形状 (B, C·d)
          │  线性层 + tanh（"隐藏状态" hidden state）
          ▼
        隐藏向量                                         形状 (B, h)
          │  线性层，投影回嵌入空间
          ▼
        "预测出来的输出嵌入" o                           形状 (B, d)
          │  和词表里每个词的嵌入做点积：logits = o · Eᵀ + b
          ▼
        logits（每个词一个分数）                          形状 (B, V)

    最后一步正是 slides 上说的："把预测出来的嵌入乘以词表的嵌入矩阵，得到每个词的分数"。
    输入和输出共用同一个矩阵 E，叫做**权重共享（weight tying）**。

【为什么神经网络比 n-gram 强】
    n-gram 把每个上下文当成一个孤立的"格子"：没见过 "my friend grows the white" 这 5 个词，
    它就对下一个词一无所知。神经网络把每个词变成一个向量，意思相近的词（she / he / my friend）
    向量也相近，所以没见过的上下文也能借用"相似上下文"的经验。这就是 Bengio et al. (2003)
    论文的核心论点：用**分布式表示**对抗**维度灾难**。

【怎样使用】
    from neural_lm import load_or_train, NextTokenModel
    model, info = load_or_train()                 # 第一次训练约 3 秒，之后读缓存
    ntm = NextTokenModel(model, garden_vocab())
    ntm.logprobs(["I", "eat", "the"])             # 下一个词的 log 概率（长度 V 的数组）

【LLM 联系】
    今天的 LLM 把中间那个 MLP 换成了几十层 Transformer（第 2 讲），窗口从 5 个 token
    变成几十万个 token，但"嵌入 → 网络 → 与嵌入矩阵做点积得到 logits → softmax"这条
    流水线一模一样。GPT-2、T5、Gemma 等模型都使用了权重共享。
"""

import hashlib
import json
import math
import os
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from toy_world import BOS, garden_vocab, standard_splits

HERE = os.path.dirname(os.path.abspath(__file__))
CKPT_DIR = os.path.join(HERE, "checkpoints")

# 默认配置。改任何一项都会得到一个新的缓存文件（文件名里带配置的哈希）。
DEFAULT_CFG = {
    "context": 5,        # 窗口大小 C：只看最近 5 个 token
    "d_emb": 32,         # 嵌入维度 d
    "d_hidden": 128,     # 隐藏层宽度 h
    "tie": True,         # 输入嵌入和输出嵌入是否共用一个矩阵
    "steps": 6000,       # 训练步数
    "batch": 256,        # 每步的样本数
    "lr": 3e-3,          # Adam 的初始学习率（之后按余弦曲线降到 0）
    "seed": 0,
}


class WindowLM(nn.Module):
    """
    固定窗口神经语言模型。

    【参数】
        vocab_size：词表大小 V。
        context：窗口大小 C（看最近几个 token）。
        d_emb：嵌入维度 d；d_hidden：隐藏层宽度 h。
        tie：True 时 logits = o · Eᵀ + b（与输入共用嵌入矩阵 E）；
             False 时用一个独立的输出矩阵 W_out（形状 V × d）。
    【边界】输入的 id 必须在 [0, V) 之内；文档开头不够 C 个 token 时，调用者要用 <bos> 补齐。
    【Python 小知识】继承 nn.Module 的类：在 __init__ 里把子层赋值给 self.xxx，
        PyTorch 就会自动登记它们的参数（model.parameters() 能找到）；
        forward 定义前向计算，调用 model(x) 时会自动执行 forward(x)。
    """

    def __init__(self, vocab_size, context=5, d_emb=32, d_hidden=128, tie=True):
        super().__init__()                         # 必须先调用父类的初始化
        self.context = context
        self.tie = tie
        self.emb = nn.Embedding(vocab_size, d_emb)          # 嵌入矩阵 E：V × d
        self.hidden = nn.Linear(context * d_emb, d_hidden)  # 第一层：C·d -> h
        self.to_emb = nn.Linear(d_hidden, d_emb)            # 第二层：h -> d（输出嵌入 o）
        if tie:
            self.out_bias = nn.Parameter(torch.zeros(vocab_size))   # 每个词一个偏置 b
        else:
            self.out = nn.Linear(d_emb, vocab_size)                 # 独立的输出矩阵（含偏置）

    def output_embedding(self, ctx_ids):
        """
        前半段：上下文 id -> 预测出来的输出嵌入 o。

        【参数】ctx_ids：整数张量，形状 (B, C)。
        【返回】形状 (B, d) 的张量。
        """
        e = self.emb(ctx_ids)                      # (B, C, d)：查表
        e = e.flatten(start_dim=1)                 # (B, C·d)：把 C 个向量首尾相接
        h = torch.tanh(self.hidden(e))             # (B, h)：隐藏状态
        return self.to_emb(h)                      # (B, d)

    def forward(self, ctx_ids):
        """
        【返回】logits，形状 (B, V)。注意返回的是 logits 而不是概率：
                损失函数 F.cross_entropy 内部会做（数值稳定的）log-softmax。
        """
        o = self.output_embedding(ctx_ids)
        if self.tie:
            # o: (B, d)，E: (V, d)，o @ Eᵀ: (B, V)——每个词的分数 = o 和该词嵌入的点积
            return o @ self.emb.weight.T + self.out_bias
        return self.out(o)


def make_windows(docs, vocab, context):
    """
    把文档切成 (上下文, 下一个 token) 训练样本。

    【参数】docs：文档列表（每篇是词列表，以 <eos> 结尾）；vocab：Vocab；context：窗口 C。
    【返回】(X, Y)：X 是 (N, C) 的整数张量，Y 是 (N,) 的整数张量。N = 所有文档的 token 总数。
    【例子】C = 3，文档 [I, eat, ., <eos>]：
            ([<bos>,<bos>,<bos>] -> I), ([<bos>,<bos>,I] -> eat),
            ([<bos>,I,eat] -> .), ([I,eat,.] -> <eos>)
    【边界】每篇文档独立补 <bos>，窗口不会跨越文档边界（文档之间互相看不见）。
    """
    bos = vocab.stoi[BOS]
    xs, ys = [], []
    for doc in docs:
        ids = [bos] * context + vocab.encode(doc)
        for i in range(context, len(ids)):
            xs.append(ids[i - context:i])
            ys.append(ids[i])
    return torch.tensor(xs, dtype=torch.long), torch.tensor(ys, dtype=torch.long)


@torch.no_grad()
def evaluate(model, X, Y, batch=8192):
    """
    计算平均 NLL 和困惑度。

    【返回】(平均 NLL, 困惑度)。
    【Python 小知识】@torch.no_grad() 是"装饰器"：这个函数里的计算不记录梯度，省内存也更快。
        评估时不需要反向传播，所以总是这样写。
    """
    model.eval()
    total = 0.0
    for i in range(0, len(X), batch):
        logits = model(X[i:i + batch])
        total += F.cross_entropy(logits, Y[i:i + batch], reduction="sum").item()
    avg = total / len(X)
    return avg, math.exp(avg)


def train_window_lm(train_docs, valid_docs, vocab, cfg, log_every=1000, verbose=True):
    """
    训练一个 WindowLM。

    【思路】标准的训练循环：随机抽一批样本 → 前向得到 logits → 交叉熵损失 →
            反向传播 → Adam 更新。学习率按余弦曲线从 lr 降到 0（训练后期步子变小，收敛得更好）。
    【参数】train_docs / valid_docs：文档列表；vocab：Vocab；cfg：配置字典（见 DEFAULT_CFG）；
            log_every：每隔多少步打印一次验证集困惑度。
    【返回】(model, history)。history 是 [(步数, 训练损失的滑动平均, 验证集困惑度), ...]。
    【边界】为了可复现：固定随机种子、单线程，批次的抽取用独立的 torch.Generator。
    """
    torch.manual_seed(cfg["seed"])
    torch.set_num_threads(1)
    C = cfg["context"]
    Xtr, Ytr = make_windows(train_docs, vocab, C)
    Xva, Yva = make_windows(valid_docs, vocab, C)
    model = WindowLM(len(vocab), C, cfg["d_emb"], cfg["d_hidden"], cfg["tie"])
    opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
    steps = cfg["steps"]
    # 【Python 小知识】LambdaLR 每步把学习率设成 lr × f(step)；这里 f 从 1 余弦下降到 0。
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 0.5 * (1 + math.cos(math.pi * s / steps)))
    gen = torch.Generator().manual_seed(cfg["seed"])
    history, running = [], None
    t0 = time.time()
    for step in range(1, steps + 1):
        model.train()
        idx = torch.randint(0, len(Xtr), (cfg["batch"],), generator=gen)
        loss = F.cross_entropy(model(Xtr[idx]), Ytr[idx])   # 平均每个 token 的 NLL
        opt.zero_grad()            # 清空上一步的梯度（PyTorch 默认会累加梯度）
        loss.backward()            # 反向传播：算出每个参数的梯度
        opt.step()                 # 按梯度更新参数
        sched.step()               # 更新学习率
        running = loss.item() if running is None else 0.98 * running + 0.02 * loss.item()
        if step % log_every == 0 or step == steps:
            _, vppl = evaluate(model, Xva, Yva)
            history.append((step, running, vppl))
            if verbose:
                print(f"    步数 {step:5d}  训练损失(滑动平均) {running:.3f}  "
                      f"验证集困惑度 {vppl:.3f}  用时 {time.time() - t0:.1f}s")
    model.eval()
    return model, history


def _cache_path(cfg, train_docs):
    """缓存文件名 = 配置 + 训练语料内容的哈希；任何一个变了都会重新训练。"""
    h = hashlib.sha1()
    h.update(json.dumps(cfg, sort_keys=True).encode())
    for d in train_docs:
        h.update(" ".join(d).encode())
        h.update(b"\n")
    return os.path.join(CKPT_DIR, f"window_lm_{h.hexdigest()[:10]}.pt")


def load_or_train(cfg=None, splits=None, verbose=True):
    """
    读取缓存的模型；没有缓存就训练一个并保存。

    【参数】cfg：配置（默认 DEFAULT_CFG，可以只传想改的几项）；splits：standard_splits() 的结果（可省略）。
    【返回】(model, info)。info 是字典：{"history": ..., "path": 缓存路径, "cached": 是否读的缓存}。
    【边界】缓存放在本目录的 checkpoints/ 下；删除这个目录即可强制重新训练。
    """
    full = dict(DEFAULT_CFG)
    full.update(cfg or {})
    splits = splits or standard_splits()
    vocab = garden_vocab()
    path = _cache_path(full, splits["train"])
    if os.path.exists(path):
        blob = torch.load(path, weights_only=False)
        model = WindowLM(len(vocab), full["context"], full["d_emb"], full["d_hidden"], full["tie"])
        model.load_state_dict(blob["state_dict"])
        model.eval()
        return model, {"history": blob["history"], "path": path, "cached": True}
    if verbose:
        print(f"  （没有找到缓存，开始训练：{full}）")
    model, history = train_window_lm(splits["train"], splits["valid"], vocab, full, verbose=verbose)
    os.makedirs(CKPT_DIR, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "history": history, "cfg": full}, path)
    return model, {"history": history, "path": path, "cached": False}


class NextTokenModel:
    """
    把 WindowLM 包装成"给一段历史，返回下一个 token 的 log 概率"的函数，供解码使用。

    【思路】窗口模型只看最近 C 个 token，所以只要最近 C 个 token 一样，答案就一样；
            用一个字典把算过的结果存起来（memoization，记忆化），解码时能快很多。
    【参数】model：WindowLM；vocab：Vocab。
    """

    def __init__(self, model, vocab):
        self.model = model
        self.vocab = vocab
        self.C = model.context
        self.bos = vocab.stoi[BOS]
        self._cache = {}

    def logprobs_ids(self, history_ids):
        """
        【参数】history_ids：已经生成的 token id 列表（可以为空，或比 C 短）。
        【返回】长度 V 的 NumPy 数组（float64），是下一个 token 的 log 概率。
        【边界】历史不够 C 个时左边补 <bos>，和训练时完全一致。
        """
        ctx = ([self.bos] * self.C + list(history_ids))[-self.C:]
        key = tuple(ctx)
        if key not in self._cache:
            with torch.no_grad():
                logits = self.model(torch.tensor([ctx]))[0]
                self._cache[key] = F.log_softmax(logits.double(), dim=-1).numpy()
        return self._cache[key]

    def logprobs(self, history_words):
        """同上，但输入是词列表。"""
        return self.logprobs_ids(self.vocab.encode(history_words))

    def probs(self, history_words):
        return np.exp(self.logprobs(history_words))