"""
=============================================================================
第 6 讲 · 代码 3：回到第 5 讲的主谓一致实验 —— RNN vs LSTM vs GRU
=============================================================================

对应讲义：lectures/lecture06_lstm_seq2seq_attention.md 第 3.5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 5 讲代码 04 的演示 3(b) 发现：如果训练语料里**只有短距离**的主谓一致（0~1 个修饰语），
普通 RNN 推广到更长的距离时准确率会逐渐下降。
这里用完全相同的设置训练 LSTM 和 GRU，看门控结构能不能让模型学到更稳健的"记忆"。

    演示 1：只见过短依赖时，三种模型在不同距离上的主谓一致准确率（每种 2 个随机种子）
    演示 2：三种模型的困惑度与参数量

本文件导入第 5 讲的 toy_language.py 和 04_rnn_lm.py（请保持 code/ 下的目录结构）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 03_agreement_lstm.py           # 完整运行
    python 03_agreement_lstm.py --quick   # 冒烟测试

需要 PyTorch。完整运行约 4~8 分钟（CPU，训练 6 个语言模型）。
"""

import os

# 必须在 import torch 之前设置：限制底层数学库（OpenMP）的线程数，原因见下面 set_num_threads 处的说明
os.environ.setdefault("OMP_NUM_THREADS", "1")
import importlib
import sys
from pathlib import Path

import numpy as np

# 【Python 小知识】把第 5 讲的文件夹加进搜索路径，才能 import 那里的模块
LECTURE05 = Path(__file__).resolve().parent.parent / "lecture05"
if str(LECTURE05) not in sys.path:
    sys.path.append(str(LECTURE05))

from toy_language import Vocab, agreement_test_set, generate_corpus  # noqa: E402

rnnlm = importlib.import_module("04_rnn_lm")
import torch  # noqa: E402

# 【Python 小知识】这些模型很小、而且是一步一步算的，多线程反而更慢（线程之间的协调开销大于收益，
#   实测单线程比 4 线程快约 2 倍），还会和同时运行的其他程序争抢 CPU。所以只用 1 个线程。
torch.set_num_threads(1)


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


def main(quick=False):
    n_train, epochs, seeds = (2000, 1, [0]) if quick else (20000, 6, [0, 1])
    vocab = Vocab(generate_corpus(20000, seed=0))
    short_train = generate_corpus(n_train, seed=0, modifier_probs=(0.6, 0.4))
    test = generate_corpus(2000, seed=2)
    ks = [1, 2, 3, 4, 6, 8]
    tests = {k: agreement_test_set(k, n=400) for k in ks}

    print_header(f"演示 1：只见过 0~1 个修饰语的训练语料，推广到更长的距离（每种模型 {len(seeds)} 个随机种子）")
    print(f"\n    {'模型':<6s}" + "".join(f"{'k=' + str(k):>9s}" for k in ks) + "   （k = 修饰语个数）")
    print(f"    {'距离':<6s}" + "".join(f"{np.mean([t['distance'] for t in tests[k]]):>9.1f}" for k in ks))
    models = {}
    for cell in ["rnn", "gru", "lstm"]:
        accs = []
        for seed in seeds:
            m = rnnlm.train_rnnlm(short_train, vocab, cell=cell, epochs=epochs, seed=seed, verbose=False)
            models[(cell, seed)] = m
            accs.append([rnnlm.agreement_accuracy(m, vocab, tests[k]) for k in ks])
        mean = np.mean(accs, axis=0)
        print(f"    {cell.upper():<6s}" + "".join(f"{a:>9.1%}" for a in mean))
        if len(seeds) > 1:
            print(f"    {'':<6s}" + "".join(f"{f'{min(x):.0%}~{max(x):.0%}':>9s}" for x in zip(*accs)) + "   （各种子的最低~最高）")
    print("""
  观察：三种模型在训练时见过的短距离上都几乎全对；距离一长，普通 RNN 的准确率逐渐下降。
    LSTM 几乎不受影响 —— 细胞状态让"主语是单数还是复数"这条信息能稳定地保存下去，
    不会每一步都被重新混合、逐渐损耗。
    GRU 平均来说比 RNN 好，但两个随机种子的结果差别很大（看"最低~最高"那一行）：
    有的种子学到了稳健的机制，有的没有。只跑一次实验很容易得出片面的结论。
    （具体数字以你运行的结果为准。）""")

    print_header("演示 2：困惑度与参数量")
    print()
    for cell in ["rnn", "gru", "lstm"]:
        m = models[(cell, seeds[0])]
        print(f"    {cell.upper():<6s}测试集困惑度 {rnnlm.perplexity(m, test, vocab):.2f}，"
              f"参数 {sum(p.numel() for p in m.parameters()):,} 个")
    print("""
  观察：在这门简单的玩具语言上，三种模型的困惑度很接近（大部分词只依赖近处的上下文），
    差别集中在少数需要长距离信息的位置上 —— 这也说明了为什么要专门设计主谓一致这样的测试。""")


if __name__ == "__main__":
    main(quick="--quick" in sys.argv)