"""
=============================================================================
第 3 讲 · 代码 4：PyTorch 入门 + 窗口分类器做命名实体识别
=============================================================================

对应讲义：lectures/lecture03_neural_networks_backprop.md 第 2、6、7 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
前三个文件把神经网络的数学讲清楚了：矩阵求导、反向传播、手写两层网络。
从这里开始，我们改用 **PyTorch**：张量运算 + 自动求导 + 现成的网络层和优化器，
让我们专心设计模型，而不用每次都手推梯度。

任务是**命名实体识别**（NER）：给每个词打上 PER / LOC / ORG / O 标签。
用的模型是 Collobert 等人（2011）的**窗口分类器**：

    把中心词和它左右各 m 个词的词向量**拼接**成一个长向量 x
    → 隐藏层 h = tanh(W x + b)
    → 输出层 logits = U h + c → softmax → 4 个类别的概率

    演示 1：PyTorch 入门 —— 张量、自动求导、nn.Embedding、nn.Linear、优化器
    演示 2：数据长什么样
    演示 3：训练窗口分类器，看每一类的精确率、召回率、F1
    演示 4：窗口大小 m 和"词丢弃"（word dropout）对**没见过的实体**的影响
    演示 5：歧义词（jordan、washington、apple、amazon）需要上下文
    演示 6：用无标注语料预训练词向量 —— 冻结还是微调？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 04_pytorch_window_ner.py

需要 numpy 和 PyTorch。演示 6 会导入第 2 讲的代码训练词向量（请保持 code/ 下的目录结构）。
运行时间约 1 分钟（CPU）。
"""

import importlib
import sys
import time
from collections import Counter
from pathlib import Path

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

from ner_data import AMBIGUOUS, LABELS, make_datasets, unseen_entity_words

PAD, UNK = "<pad>", "<unk>"


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


# =============================================================================
# 演示 1：PyTorch 入门
# =============================================================================
def demo_pytorch_basics():
    """演示 1：PyTorch 最基本的几样东西。"""
    print_header("演示 1：PyTorch 入门")

    # (a) 张量：和 numpy 数组几乎一样
    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
    b = torch.ones(2, 2)
    print(f"\n  (a) 张量 a 的形状 {tuple(a.shape)}，数据类型 {a.dtype}")
    print(f"      a @ b = {(a @ b).tolist()}，a * b（逐元素）= {(a * b).tolist()}")
    print(f"      和 numpy 互相转换：a.numpy() 是 {type(a.numpy()).__name__}，torch.from_numpy(...) 反过来")

    # (b) 自动求导
    x = torch.tensor([1.0, -2.0, 3.0], requires_grad=True)
    f = (x ** 2).sum()
    f.backward()
    print(f"\n  (b) 自动求导：f = Σ x²，x = {x.tolist()} → f.backward() 后 x.grad = {x.grad.tolist()}（= 2x）")

    # (c) softmax + 交叉熵的梯度就是 ŷ - y
    logits = torch.tensor([[2.0, 1.0, 0.1]], requires_grad=True)
    loss = F.cross_entropy(logits, torch.tensor([0]))
    loss.backward()
    expected = torch.softmax(logits.detach(), dim=1) - torch.tensor([[1.0, 0.0, 0.0]])
    fmt = lambda t: "[" + ", ".join(f"{v:+.3f}" for v in t.flatten().tolist()) + "]"
    print(f"  (c) F.cross_entropy 的梯度 {fmt(logits.grad)}，ŷ - y = {fmt(expected)}（第 1 讲的老朋友）")

    # (d) nn.Embedding = one-hot × 矩阵
    torch.manual_seed(0)
    emb = nn.Embedding(5, 3)                      # 词表 5 个词，每个 3 维
    ids = torch.tensor([4, 1])
    onehot = F.one_hot(ids, 5).float()
    print(f"\n  (d) nn.Embedding(5, 3)：权重形状 {tuple(emb.weight.shape)}；"
          f"emb(ids) 与 one_hot(ids) @ weight 相同？{torch.allclose(emb(ids), onehot @ emb.weight)}")

    # (e) nn.Linear：y = x Wᵀ + b（注意 PyTorch 的权重形状是 (输出维度, 输入维度)）
    lin = nn.Linear(3, 2)
    xx = torch.randn(4, 3)
    print(f"  (e) nn.Linear(3, 2)：weight 形状 {tuple(lin.weight.shape)}，"
          f"输入 (4, 3) → 输出 {tuple(lin(xx).shape)}；等于 x @ W.T + b？"
          f"{torch.allclose(lin(xx), xx @ lin.weight.T + lin.bias)}")

    # (f) 优化器：每一步 = 清零梯度 → 前向 → 反向 → 更新
    w = torch.tensor([5.0], requires_grad=True)
    opt = torch.optim.SGD([w], lr=0.1)
    for _ in range(30):
        opt.zero_grad()
        loss = (w - 2.0) ** 2                     # 最小值在 w = 2
        loss.backward()
        opt.step()
    print(f"\n  (f) 用 SGD 最小化 (w - 2)²，从 w = 5 出发走 30 步：w = {w.item():.4f}")
    print("""
  这就是 PyTorch 训练的全部套路：
      optimizer.zero_grad()   # 清零梯度（梯度是累加的，见代码 02 演示 3）
      loss = 模型前向(...)     # 一边算，一边自动记录计算图
      loss.backward()         # 反向传播，算出每个参数的 .grad
      optimizer.step()        # 用梯度更新参数""")


# =============================================================================
# 第 2 部分：把数据变成窗口
# =============================================================================
def build_vocab(sentences):
    """
    从句子中建立词表。编号 0 留给 <pad>（句子边界外的填充），1 留给 <unk>（没见过的词）。

    【参数】sentences : list[list[str]]
    【返回】dict[str, int]
    """
    words = sorted(set(w for s in sentences for w in s))
    word2id = {PAD: 0, UNK: 1}
    for w in words:
        word2id[w] = len(word2id)
    return word2id


def make_windows(data, word2id, m):
    """
    把每个词变成一个"窗口样本"：中心词和左右各 m 个词的编号，以及中心词的标签。

    【思路】
        句子 "the meeting was held in paris ."，m = 2，中心词 paris：
            窗口 = [held, in, paris, ., <pad>]（右边超出句子的位置用 <pad> 填充）
        不在词表里的词变成 <unk>。

    【参数】
        data    : list[(tokens, labels)]
        word2id : dict
        m       : int，窗口半径（m = 0 表示只看中心词本身）

    【返回】
        (X, y, words)
            X     : torch.LongTensor (N, 2m+1)
            y     : torch.LongTensor (N,)，标签编号（LABELS 中的位置）
            words : list[str]，每个样本的中心词（用于分析）

    【Python 小知识】
        torch.tensor(list_of_lists, dtype=torch.long)：整数张量，nn.Embedding 的输入必须是整数。
    """
    X, y, words = [], [], []
    for tokens, labels in data:
        ids = [word2id.get(t, word2id[UNK]) for t in tokens]
        padded = [word2id[PAD]] * m + ids + [word2id[PAD]] * m
        for t in range(len(tokens)):
            X.append(padded[t:t + 2 * m + 1])
            y.append(LABELS.index(labels[t]))
            words.append(tokens[t])
    return torch.tensor(X, dtype=torch.long), torch.tensor(y, dtype=torch.long), words


# =============================================================================
# 第 3 部分：模型、训练、评估
# =============================================================================
class WindowClassifier(nn.Module):
    """
    窗口分类器：嵌入 → 拼接 → 隐藏层（tanh）→ 输出层。

    【参数】（构造函数）
        vocab_size : 词表大小
        m          : 窗口半径
        embed_dim  : 词向量维度
        hidden     : 隐藏单元数
        n_classes  : 类别数
        pretrained : None 或 np.ndarray (vocab_size, embed_dim)，用来初始化词向量
        freeze     : bool，True 时词向量不参与训练（梯度不更新它）

    【Python 小知识】
        继承 nn.Module 的类：在 __init__ 里定义子模块（它们的参数会被自动登记），
        在 forward 里写前向计算。调用 model(x) 会自动调用 forward(x)。
        .view(B, -1)：改变形状，-1 表示"这一维由其他维自动推算"。
    """

    def __init__(self, vocab_size, m, embed_dim=30, hidden=100, n_classes=len(LABELS), pretrained=None, freeze=False):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        if pretrained is not None:
            self.emb.weight.data.copy_(torch.as_tensor(pretrained, dtype=torch.float32))
        self.emb.weight.requires_grad_(not freeze)
        self.hidden = nn.Linear((2 * m + 1) * embed_dim, hidden)
        self.out = nn.Linear(hidden, n_classes)

    def forward(self, X):
        """X: (B, 2m+1) 的词编号 → (B, n_classes) 的 logits。"""
        e = self.emb(X)                          # (B, 2m+1, embed_dim)
        x = e.view(X.shape[0], -1)               # 拼接：(B, (2m+1)·embed_dim)
        h = torch.tanh(self.hidden(x))
        return self.out(h)


def train_model(model, X, y, epochs=8, lr=2e-3, batch_size=64, word_dropout=0.0, seed=0, unk_id=1):
    """
    用 Adam 训练窗口分类器。

    【参数】
        model        : WindowClassifier
        X, y         : 训练数据
        epochs, lr, batch_size : 训练轮数、学习率、批大小
        word_dropout : float，训练时把窗口中每个词以这个概率替换成 <unk>
        seed         : 随机种子
        unk_id       : <unk> 的编号

    【返回】list[float]，每轮的平均损失

    【边界】
        - 词丢弃（word dropout）是一种数据增强 / 正则化：测试时没见过的词都会变成 <unk>，
          如果训练时 <unk> 从来没出现在实体的位置上，模型就不知道该怎么处理它。
          随机把一些词换成 <unk>，相当于让模型提前"练习"处理没见过的词，被迫更多地依靠上下文。
    """
    torch.manual_seed(seed)
    opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=lr)
    history = []
    for _ in range(epochs):
        perm = torch.randperm(len(X))
        total = 0.0
        for start in range(0, len(X), batch_size):
            idx = perm[start:start + batch_size]
            xb = X[idx].clone()
            if word_dropout > 0:
                drop = (torch.rand(xb.shape) < word_dropout) & (xb != 0)
                xb[drop] = unk_id
            opt.zero_grad()
            loss = F.cross_entropy(model(xb), y[idx])
            loss.backward()
            opt.step()
            total += loss.item() * len(idx)
        history.append(total / len(X))
    return history


@torch.no_grad()
def predict(model, X):
    """返回预测的标签编号（不需要梯度，用 @torch.no_grad() 省内存、提速）。"""
    model.eval()
    pred = model(X).argmax(dim=1)
    model.train()
    return pred


def prf(pred, y):
    """
    对 PER、LOC、ORG 三个实体类别分别计算词级别的精确率、召回率、F1，以及它们的宏平均 F1。

    【思路】（讲义第 2.4 节）
        对类别 c：TP = 预测为 c 且真的是 c；FP = 预测为 c 但不是 c；FN = 真的是 c 但没预测成 c。
        精确率 P = TP / (TP + FP)：说是 c 的里面，有多少真的是 c。
        召回率 R = TP / (TP + FN)：真的是 c 的里面，找出来了多少。
        F1 = 2PR / (P + R)：两者的调和平均，任何一个很低，F1 就很低。
        不统计 O 类：O 占了绝大多数，把它算进来会让分数虚高。

    【参数】pred, y : torch.LongTensor (N,)
    【返回】(dict[类别名, (P, R, F1)], macro_F1)

    【边界】
        - 真实的 NER 评估（CoNLL）用**实体级**F1：一个两词的人名只有两个词都对才算对。这里用更简单的词级别。
    """
    out = {}
    for c in range(1, len(LABELS)):
        tp = int(((pred == c) & (y == c)).sum())
        fp = int(((pred == c) & (y != c)).sum())
        fn = int(((pred != c) & (y == c)).sum())
        p = tp / (tp + fp) if tp + fp else 0.0
        r = tp / (tp + fn) if tp + fn else 0.0
        out[LABELS[c]] = (p, r, 2 * p * r / (p + r) if p + r else 0.0)
    return out, float(np.mean([v[2] for v in out.values()]))


def accuracy_on(pred, y, words, word_set):
    """只在中心词属于 word_set 的样本上计算准确率。返回 (准确率, 样本数)。"""
    mask = torch.tensor([w in word_set for w in words])
    n = int(mask.sum())
    return (float((pred[mask] == y[mask]).float().mean()) if n else float("nan")), n


# =============================================================================
# 演示 2 ~ 6
# =============================================================================
def main():
    demo_pytorch_basics()
    train, test_seen, test_unseen, unlabeled = make_datasets()
    unseen = unseen_entity_words()

    # ---------------- 演示 2：数据 ----------------
    print_header("演示 2：数据")
    label_counts = Counter(l for _, ls in train for l in ls)
    print(f"\n  训练集 {len(train)} 句，{sum(label_counts.values())} 个词；标签分布：" +
          "，".join(f"{k} {v}" for k, v in label_counts.items()))
    print(f"  只出现在'未见实体'测试集里的实体词（{len(unseen)} 个）：{sorted(unseen)[:12]} ……")
    for tokens, labels in test_unseen[:3]:
        print("    " + " ".join(f"{t}/{l}" for t, l in zip(tokens, labels)))
    word2id = build_vocab([t for t, _ in train])
    print(f"  词表（只用训练集建立）：{len(word2id)} 个词（含 <pad>、<unk>）")

    # ---------------- 演示 3：训练 ----------------
    print_header("演示 3：训练窗口分类器（m = 2，词向量 30 维，隐藏层 100，词丢弃 0.2）")
    m = 2
    Xtr, ytr, _ = make_windows(train, word2id, m)
    model = WindowClassifier(len(word2id), m)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"\n  训练样本（窗口）{len(Xtr)} 个；模型参数 {n_params} 个")
    t0 = time.perf_counter()
    hist = train_model(model, Xtr, ytr, word_dropout=0.2)
    print(f"  每轮平均损失：{[round(h, 4) for h in hist]}（用时 {time.perf_counter() - t0:.1f} 秒，取决于你的电脑）")
    for name, data in [("已见实体测试集", test_seen), ("未见实体测试集", test_unseen)]:
        X, y, words = make_windows(data, word2id, m)
        pred = predict(model, X)
        scores, macro = prf(pred, y)
        print(f"\n  {name}：宏平均 F1 = {macro:.3f}")
        for c, (p, r, f1) in scores.items():
            print(f"    {c}：精确率 {p:.3f}  召回率 {r:.3f}  F1 {f1:.3f}")
    X, y, words = make_windows(test_unseen, word2id, m)
    acc, n = accuracy_on(predict(model, X), y, words, unseen)
    print(f"\n  只看 {n} 个未见实体词：准确率 {acc:.1%}（这些词在模型眼里全是 <unk>，只能靠上下文判断）")

    # ---------------- 演示 4：窗口大小与词丢弃 ----------------
    print_header("演示 4：窗口大小 m 与词丢弃的影响（每个设置 3 个随机种子取平均）")
    print(f"\n    {'设置':<24s}{'已见实体 F1':>12s}{'未见实体 F1':>12s}{'未见实体词准确率':>16s}")
    for m in [0, 1, 2]:
        Xtr, ytr, _ = make_windows(train, word2id, m)
        Xs, ys, _ = make_windows(test_seen, word2id, m)
        Xu, yu, wu = make_windows(test_unseen, word2id, m)
        for wd in [0.0, 0.2]:
            rows = []
            for seed in [0, 1, 2]:
                torch.manual_seed(seed)
                mdl = WindowClassifier(len(word2id), m)
                train_model(mdl, Xtr, ytr, word_dropout=wd, seed=seed)
                pu = predict(mdl, Xu)
                rows.append((prf(predict(mdl, Xs), ys)[1], prf(pu, yu)[1], accuracy_on(pu, yu, wu, unseen)[0]))
            r = np.mean(rows, axis=0)
            print(f"    {'m = ' + str(m) + '，词丢弃 ' + str(wd):<24s}{r[0]:>12.3f}{r[1]:>12.3f}{r[2]:>16.1%}")
    print("""
  观察：
    - m = 0（只看这个词本身）：已见实体还行，未见实体完全没救 —— 所有没见过的词都是同一个 <unk>，
      模型没有任何信息可用。
    - m >= 1：上下文（"born in ___"、"shares of ___"、"mr. ___"）提供了判断依据。
    - 不加词丢弃时，训练集里实体位置上从来没出现过 <unk>，模型遇到 <unk> 时倾向于猜 O；
      加了词丢弃，模型学会了"中心词不认识就看上下文"，未见实体的表现大幅提高。
    - 但 m = 0 时词丢弃反而更糟：没有上下文可看，被换成 <unk> 的词大多是 O，
      模型只能学到"<unk> 就是 O"。词丢弃只有在模型有别的信息可依靠时才有用。""")

    # ---------------- 演示 5：歧义词 ----------------
    print_header("演示 5：歧义词需要上下文")
    ambiguous_data = [(t, l) for t, l in test_seen if any(w in AMBIGUOUS for w in t)]
    print(f"\n  已见实体测试集中含歧义词的句子 {len(ambiguous_data)} 句，例如：")
    shown = set()
    for t, l in ambiguous_data:
        for w, lab in zip(t, l):
            if w in AMBIGUOUS and (w, lab) not in shown and len(shown) < 6:
                print(f"    {' '.join(t):<50s} {w} → {lab}")
                shown.add((w, lab))
                break
    for m in [0, 2]:
        Xtr, ytr, _ = make_windows(train, word2id, m)
        mdl = WindowClassifier(len(word2id), m)
        train_model(mdl, Xtr, ytr, word_dropout=0.2)
        X, y, words = make_windows(test_seen, word2id, m)
        acc, n = accuracy_on(predict(mdl, X), y, words, AMBIGUOUS)
        print(f"  m = {m}：{n} 个歧义词的准确率 {acc:.1%}")
    print("  → 只看词本身（m = 0）时，同一个词只能给出同一个答案；有了上下文才能区分 jordan 是人还是国家。")

    # ---------------- 演示 6：预训练词向量 ----------------
    print_header("演示 6：用无标注语料预训练词向量 —— 冻结还是微调？")
    sgns = load_sgns()
    vocab_sents = unlabeled + [t for t, _ in train]
    word2id_big = build_vocab(vocab_sents)
    id_sents = [[word2id_big[w] for w in s] for s in unlabeled]
    counts = np.zeros(len(word2id_big))
    for s in id_sents:
        for w in s:
            counts[w] += 1
    counts[counts == 0] = 1.0                     # <pad>、<unk> 从不出现，给一个计数避免除以 0
    C, U, _ = sgns.train_sgns(id_sents, counts, dim=30, epochs=10, verbose=False)
    pretrained = C + U
    print(f"\n  在 {len(unlabeled)} 句无标注语料上训练 SGNS 词向量（30 维；无标注语料里包含所有实体，但没有标签）。")
    print("  （真实场景：无标注文本几乎无限，标注数据很少。预训练的词向量能把标注数据里学到的知识，")
    print("    迁移到只在无标注文本里出现过的词上。）\n")
    print(f"    {'设置':<34s}{'未见实体 F1':>12s}{'未见实体词准确率':>16s}")
    for m in [0, 2]:
        Xtr, ytr, _ = make_windows(train, word2id_big, m)
        Xu, yu, wu = make_windows(test_unseen, word2id_big, m)
        configs = [("随机初始化（未见实体 = <unk>）", None, False),
                   ("预训练词向量，冻结", pretrained, True),
                   ("预训练词向量，微调", pretrained, False)]
        for name, pre, freeze in configs:
            if pre is None:
                Xtr2, ytr2, _ = make_windows(train, word2id, m)
                Xu2, yu2, wu2 = make_windows(test_unseen, word2id, m)
                mdl = WindowClassifier(len(word2id), m)
                train_model(mdl, Xtr2, ytr2, word_dropout=0.2)
                pu, yy, ww = predict(mdl, Xu2), yu2, wu2
            else:
                mdl = WindowClassifier(len(word2id_big), m, pretrained=pre, freeze=freeze)
                train_model(mdl, Xtr, ytr, word_dropout=0.0)
                pu, yy, ww = predict(mdl, Xu), yu, wu
            print(f"    {'m = ' + str(m) + '，' + name:<34s}{prf(pu, yy)[1]:>12.3f}{accuracy_on(pu, yy, ww, unseen)[0]:>16.1%}")
    print("""
  观察（讲义第 4.6 节）：
    - 预训练词向量里，没在标注数据中出现过的人名（比如 emma）也和见过的人名（alice、bob）靠在一起，
      因为它们在无标注语料里出现在相似的上下文中。冻结词向量时，分类器学到的"人名区域"自然地覆盖了它们。
    - 微调时，只有标注数据里出现过的词向量会被更新、移动位置；没出现过的词留在原地，
      可能被分类器的新边界"抛下"。标注数据少时，冻结（或只用很小的学习率微调）往往更稳妥；
      标注数据多、覆盖了大部分词表时，微调通常更好。""")


def load_sgns():
    """导入第 2 讲的 02_sgns_training（它又会通过 shared.py 导入第 1 讲的代码）。"""
    lecture02 = Path(__file__).resolve().parent.parent / "lecture02"
    if str(lecture02) not in sys.path:
        sys.path.append(str(lecture02))
    return importlib.import_module("02_sgns_training")


if __name__ == "__main__":
    main()