"""
=============================================================================
第 5 讲 · 代码 3：用 numpy 从零实现循环神经网络 —— 前向、时间反向传播（BPTT）
=============================================================================

对应讲义：lectures/lecture05_language_models_rnn.md 第 5、6 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
循环神经网络（RNN）的核心想法：**用同一组参数，一个词一个词地读下去**，
用一个"隐藏状态" h_t 总结到目前为止读过的全部内容：

    h_t = tanh(W_h h_{t-1} + W_e e_t + b)        e_t 是第 t 个词的词向量
    ŷ_t = softmax(U h_t + c)                     预测第 t+1 个词

无论句子多长，参数都只有 W_h、W_e、b、U、c、E 这几个，所有时间步**共享**。

训练用的是**时间反向传播**（Backpropagation Through Time，BPTT）：
把 RNN 沿时间"展开"成一个很深的前馈网络（每个时间步一层，层与层之间共享参数），
然后做普通的反向传播。关键一点：**一个被重复使用的参数，梯度是它在每个时间步的梯度之和**。

    演示 1：前向传播，逐步打印隐藏状态
    演示 2：BPTT 的梯度检验
    演示 3：在几个玩具句子上训练，看损失下降、生成句子
    演示 4：梯度怎样随时间传播？—— ‖∂L_T/∂h_t‖ 随 t 变化（第 6 讲"梯度消失"的预告）

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 03_rnn_from_scratch.py

需要 numpy。运行时间约 10 秒。
"""

import numpy as np

from toy_language import Vocab, generate_corpus


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


# =============================================================================
# 第 1 部分：RNN 语言模型
# =============================================================================
class NumpyRNNLM:
    """
    最简单的 RNN 语言模型（tanh 激活），所有前向、反向计算都用 numpy 手写。

    【参数】（构造函数）
        V      : 词表大小
        d      : 词向量维度
        H      : 隐藏状态维度
        seed   : 随机种子
        w_scale: W_h 的初始化尺度（演示 4 用它观察梯度消失 / 爆炸）

    【参数矩阵】
        E   (V, d)  词向量
        W_e (H, d)  输入 → 隐藏
        W_h (H, H)  隐藏 → 隐藏（"循环"的那条边）
        b   (H,)
        U   (V, H)  隐藏 → 输出
        c   (V,)
    """

    def __init__(self, V, d=8, H=16, seed=0, w_scale=None):
        rng = np.random.default_rng(seed)
        self.p = {
            "E": rng.normal(scale=0.1, size=(V, d)),
            "W_e": rng.normal(scale=1 / np.sqrt(d), size=(H, d)),
            "W_h": rng.normal(scale=(w_scale if w_scale is not None else 1 / np.sqrt(H)), size=(H, H)),
            "b": np.zeros(H),
            "U": rng.normal(scale=1 / np.sqrt(H), size=(V, H)),
            "c": np.zeros(V),
        }
        self.H = H

    def forward(self, ids):
        """
        前向传播：输入一个编号序列 [w_0, w_1, ..., w_T]，用 w_t 预测 w_{t+1}。

        【参数】ids : list[int]，长度 T+1（通常以 <s> 开头、</s> 结尾）

        【返回】
            (loss, cache)
                loss  : float，平均交叉熵（对 T 个预测取平均）
                cache : 反向传播需要的中间结果（每个时间步的 h、概率等）

        【Python 小知识】
            hs 用字典保存，hs[-1] 是初始隐藏状态 h_{-1} = 0，这样 hs[t-1] 在 t = 0 时也能用。
        """
        p = self.p
        hs = {-1: np.zeros(self.H)}
        probs, loss = {}, 0.0
        T = len(ids) - 1
        for t in range(T):
            e = p["E"][ids[t]]
            hs[t] = np.tanh(p["W_h"] @ hs[t - 1] + p["W_e"] @ e + p["b"])
            logits = p["U"] @ hs[t] + p["c"]
            z = np.exp(logits - logits.max())
            probs[t] = z / z.sum()
            loss -= np.log(probs[t][ids[t + 1]])
        return loss / T, (ids, hs, probs)

    def backward(self, cache):
        """
        时间反向传播（BPTT），返回所有参数的梯度。

        【思路】（讲义第 6.2 节）
            从最后一个时间步往前走。在时间步 t：
              1. 输出层：dlogits = (ŷ_t - y_t) / T  →  dU += dlogits h_tᵀ，dc += dlogits
              2. 传到 h_t 的梯度有**两个来源**：
                   来自这一步的输出：Uᵀ dlogits
                   来自下一个时间步：dh_next（h_t 被用来计算 h_{t+1}）
                 两者相加 = dh。
              3. 穿过 tanh：dz = dh ∘ (1 - h_t²)
              4. 这一步对共享参数的梯度**累加**：dW_h += dz h_{t-1}ᵀ，dW_e += dz e_tᵀ，db += dz，
                 dE[w_t] += W_eᵀ dz
              5. 传给上一个时间步：dh_next = W_hᵀ dz
            "W_h 的梯度 = 每个时间步上 W_h 的梯度之和"——这就是多元链式法则（第 3 讲第 5.4 节）。

        【返回】dict，和 self.p 同结构的梯度
        """
        ids, hs, probs = cache
        p = self.p
        g = {k: np.zeros_like(v) for k, v in p.items()}
        T = len(ids) - 1
        dh_next = np.zeros(self.H)
        for t in reversed(range(T)):
            dlogits = probs[t].copy()
            dlogits[ids[t + 1]] -= 1.0
            dlogits /= T
            g["U"] += np.outer(dlogits, hs[t])
            g["c"] += dlogits
            dh = p["U"].T @ dlogits + dh_next
            dz = dh * (1 - hs[t] ** 2)
            g["W_h"] += np.outer(dz, hs[t - 1])
            g["W_e"] += np.outer(dz, p["E"][ids[t]])
            g["b"] += dz
            g["E"][ids[t]] += p["W_e"].T @ dz
            dh_next = p["W_h"].T @ dz
        return g

    def hidden_grad_norms(self, ids):
        """
        只把最后一个时间步的损失 L_T 反向传播，返回 ‖∂L_T / ∂h_t‖ 对每个 t 的值。

        【思路】
            ∂L_T/∂h_t = ∂L_T/∂h_T · ∂h_T/∂h_{T-1} · ... · ∂h_{t+1}/∂h_t，
            每个因子 ∂h_{k+1}/∂h_k = diag(1 - h_{k+1}²) W_h。很多个这样的矩阵连乘，
            如果它们"平均来说"把向量缩小，梯度就会指数衰减（梯度消失）；放大则指数增长（梯度爆炸）。
        """
        _, (ids, hs, probs) = self.forward(ids)
        T = len(ids) - 1
        dlogits = probs[T - 1].copy()
        dlogits[ids[T]] -= 1.0
        dh = self.p["U"].T @ dlogits
        norms = {T - 1: np.linalg.norm(dh)}
        for t in reversed(range(T - 1)):
            dz = dh * (1 - hs[t + 1] ** 2)
            dh = self.p["W_h"].T @ dz
            norms[t] = np.linalg.norm(dh)
        return [norms[t] for t in range(T)]

    def generate(self, vocab, rng, max_len=20):
        """从 <s> 开始逐词采样，直到 </s>。返回词列表。"""
        p = self.p
        h = np.zeros(self.H)
        w = vocab.stoi[Vocab.BOS]
        out = []
        for _ in range(max_len):
            h = np.tanh(p["W_h"] @ h + p["W_e"] @ p["E"][w] + p["b"])
            logits = p["U"] @ h + p["c"]
            z = np.exp(logits - logits.max())
            w = rng.choice(len(z), p=z / z.sum())
            if vocab.itos[w] == Vocab.EOS:
                break
            out.append(vocab.itos[w])
        return out


# =============================================================================
# 演示
# =============================================================================
def demo_forward(vocab):
    """演示 1：前向传播。"""
    print_header("演示 1：前向传播 —— 同一组参数，一个词一个词地读")
    model = NumpyRNNLM(len(vocab), d=4, H=5, seed=0)
    sent = ["the", "dog", "runs", "."]
    ids = vocab.encode(sent)
    loss, (ids, hs, probs) = model.forward(ids)
    print(f"\n  句子：{sent}（前面加 <s>，后面加 </s>）；词向量 4 维，隐藏状态 5 维")
    for t in range(len(ids) - 1):
        print(f"    t = {t}：输入 {vocab.itos[ids[t]]:<5s} → h_{t} = {np.round(hs[t], 2)}，"
              f"P(下一个词 = {vocab.itos[ids[t + 1]]}) = {probs[t][ids[t + 1]]:.4f}")
    print(f"  平均交叉熵 = {loss:.4f}（随机初始化，所以接近 log V = {np.log(len(vocab)):.4f}）")
    print("  参数：" + "，".join(f"{k} {v.shape}" for k, v in model.p.items()) +
          f"，共 {sum(v.size for v in model.p.values())} 个 —— 和句子长度无关。")


def demo_gradient_check(vocab):
    """演示 2：BPTT 梯度检验。"""
    print_header("演示 2：BPTT 的梯度检验")
    model = NumpyRNNLM(len(vocab), d=3, H=4, seed=1)
    ids = vocab.encode(["the", "dogs", "that", "the", "cat", "likes", "run", "."])
    _, cache = model.forward(ids)
    grads = model.backward(cache)
    worst = 0.0
    for k, P in model.p.items():
        num = np.zeros_like(P)
        it = np.nditer(P, flags=["multi_index"])
        while not it.finished:
            i = it.multi_index
            old = P[i]
            P[i] = old + 1e-5
            lp, _ = model.forward(ids)
            P[i] = old - 1e-5
            lm, _ = model.forward(ids)
            P[i] = old
            num[i] = (lp - lm) / 2e-5
            it.iternext()
        rel = np.max(np.abs(num - grads[k])) / max(np.max(np.abs(num)), np.max(np.abs(grads[k])), 1e-12)
        worst = max(worst, rel)
        print(f"    {k:<4s}形状 {str(P.shape):<9s}解析 vs 数值的相对误差 {rel:.1e}")
    assert worst < 1e-6, "BPTT 梯度检验失败"
    print("  → 全部通过。注意 W_h 的梯度是 8 个时间步的梯度之和，E 中只有出现过的词那几行非零。")


def demo_train(vocab):
    """演示 3：在少量句子上训练。"""
    print_header("演示 3：用 SGD 在 200 个玩具句子上训练（numpy 手写的 RNN）")
    sents = generate_corpus(200, seed=5, modifier_probs=(0.6, 0.4))
    model = NumpyRNNLM(len(vocab), d=16, H=32, seed=0)
    rng = np.random.default_rng(0)
    lr = 0.5
    for epoch in range(15):
        total = 0.0
        for i in rng.permutation(len(sents)):
            loss, cache = model.forward(vocab.encode(sents[i]))
            g = model.backward(cache)
            for k in model.p:
                np.clip(g[k], -5, 5, out=g[k])              # 梯度裁剪（第 6 讲详细讲）
                model.p[k] -= lr * g[k]
            total += loss
        if epoch % 3 == 0 or epoch == 14:
            print(f"    第 {epoch + 1:>2d} 轮：平均交叉熵 {total / len(sents):.4f}")
    print("\n  训练后从模型中采样：")
    for _ in range(4):
        print("    " + " ".join(model.generate(vocab, rng)))
    print("""
  观察：损失稳定下降，生成的句子越来越像玩具语言。这里一次只处理一个句子、用纯 Python 循环，
    非常慢；代码 04 用 PyTorch 批量训练，快几百倍，原理完全相同。""")


def demo_gradient_flow(vocab):
    """演示 4：梯度随时间的传播。"""
    print_header("演示 4：‖∂L_T/∂h_t‖ 随 t 的变化（未训练的 RNN，隐藏状态 32 维）")
    rng = np.random.default_rng(3)
    ids = [vocab.stoi[Vocab.BOS]] + list(rng.integers(4, len(vocab), size=40)) + [vocab.stoi[Vocab.EOS]]
    print(f"\n  一个 41 个词的随机序列，只看最后一步的损失 L_T 对每个时间步隐藏状态的梯度大小：")
    print(f"    {'W_h 的初始化尺度':<18s}" + "".join(f"{'t=T-' + str(k):>10s}" for k in [0, 5, 10, 20, 30, 40]))
    for scale in [0.5 / np.sqrt(32), 1 / np.sqrt(32), 3 / np.sqrt(32)]:
        m = NumpyRNNLM(len(vocab), d=16, H=32, seed=0, w_scale=scale)
        norms = m.hidden_grad_norms(ids)
        T = len(norms)
        print(f"    {'σ = ' + f'{scale * np.sqrt(32):.1f}' + '/√H':<18s}" + "".join(f"{norms[T - 1 - k]:>10.2e}" for k in [0, 5, 10, 20, 30, 40]))
    print("""
  观察：梯度每往前传一步，都要乘一次 diag(1 - h²) W_h。
    - W_h 较小时（σ = 0.5/√H），梯度指数衰减：40 步之前的隐藏状态收到的梯度只有最后一步的 10⁻¹³ 量级，
      模型几乎不可能学会"40 个词之前的主语决定了现在的动词形式"。这就是**梯度消失**。
    - W_h 较大时（σ = 3/√H），梯度指数增长，40 步之前放大了上万倍 —— **梯度爆炸**：
      一次参数更新就可能把模型"炸飞"（演示 3 里的梯度裁剪就是为此准备的）。
    - 即使是"刚刚好"的尺度，梯度也在逐渐变小。（另外，W_h 很大时 tanh 会饱和、导数接近 0，
      所以训练过程中爆炸和消失可能交替出现。）
    第 6 讲会详细分析这个问题，并介绍 LSTM 这样的解决办法。""")


if __name__ == "__main__":
    vocab = Vocab(generate_corpus(20000, seed=0))
    demo_forward(vocab)
    demo_gradient_check(vocab)
    demo_train(vocab)
    demo_gradient_flow(vocab)