"""
=============================================================================
第 6 讲 · 代码 1：梯度消失与梯度爆炸 —— 原因、后果、梯度裁剪
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 5 讲看到：RNN 反向传播时，梯度每往前传一步都要乘一次雅可比矩阵
    ∂h_t / ∂h_{t-1} = diag(σ'(z_t)) · W_h
传 k 步就是 k 个矩阵连乘。这个文件把"为什么会消失 / 爆炸"拆开来看：

    演示 1：矩阵连乘 —— 最大特征值（谱半径）小于 1 就消失，大于 1 就爆炸
    演示 2：后果 —— W_h 的梯度几乎完全来自最近的几个时间步，远处的依赖"学不到"
    演示 3：梯度爆炸的"悬崖"—— 一个只有一个参数的 RNN，损失曲面上有一道陡壁
    演示 4：梯度裁剪 —— 在悬崖边上，裁剪与不裁剪的梯度下降有什么不同

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 01_vanishing_gradients.py

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

import numpy as np


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


# =============================================================================
# 演示 1：矩阵连乘
# =============================================================================
def matrix_with_spectral_radius(n, rho, seed=0):
    """
    构造一个 n × n 的随机矩阵，并把它缩放到谱半径（特征值绝对值的最大值）恰好为 rho。

    【思路】
        先生成随机矩阵 A，求出它的特征值，按 rho / max|λ| 缩放。
        讲义第 2.2 节的论证：如果 W = Q Λ Q⁻¹，那么 W^k = Q Λ^k Q⁻¹，
        当 k 很大时 W^k 的大小由 |λ_max|^k 决定 —— 小于 1 指数衰减，大于 1 指数增长。

    【参数】n : 维度；rho : 目标谱半径；seed : 随机种子
    【返回】np.ndarray (n, n)

    【Python 小知识】np.linalg.eigvals(A) 返回特征值（可能是复数），np.abs 取模。
    """
    A = np.random.default_rng(seed).normal(size=(n, n)) / np.sqrt(n)
    return A * rho / np.max(np.abs(np.linalg.eigvals(A)))


def demo_matrix_powers():
    """演示 1：‖W^k v‖ 随 k 的变化。"""
    print_header("演示 1：矩阵连乘 —— 谱半径决定梯度是消失还是爆炸")
    v = np.random.default_rng(1).normal(size=64)
    v /= np.linalg.norm(v)
    ks = [1, 10, 20, 50, 100]
    print(f"\n  64 维随机矩阵 W，单位向量 v，计算 ‖Wᵏ v‖：")
    print(f"    {'谱半径 ρ':<10s}" + "".join(f"{'k=' + str(k):>12s}" for k in ks))
    for rho in [0.9, 1.0, 1.1]:
        W = matrix_with_spectral_radius(64, rho)
        row, x = [], v.copy()
        for k in range(1, 101):
            x = W @ x
            if k in ks:
                row.append(np.linalg.norm(x))
        print(f"    {rho:<10.1f}" + "".join(f"{r:>12.2e}" for r in row))
    print("""
  观察：ρ = 0.9 时，100 次连乘后只剩 10⁻⁵ 量级（0.9¹⁰⁰ ≈ 2.7 × 10⁻⁵）；ρ = 1.1 时放大了上万倍（1.1¹⁰⁰ ≈ 1.4 × 10⁴）。
    RNN 的反向传播正是这样的连乘（再乘上 tanh 的导数，它不超过 1，会进一步加速衰减）。
    Pascanu 等人（2013）证明：如果 W_h 的最大特征值小于 1/γ（γ 是激活函数导数的上界，tanh 为 1），
    梯度一定会指数消失。""")


# =============================================================================
# 演示 2：后果 —— W_h 的梯度来自哪些时间步
# =============================================================================
def demo_gradient_contributions():
    """演示 2：最后一步的损失对 W_h 的梯度中，各个时间步各贡献了多少。"""
    print_header("演示 2：W_h 的梯度来自哪里？")
    rng = np.random.default_rng(0)
    H, T = 32, 40
    W_h = matrix_with_spectral_radius(H, 0.9, seed=2)
    W_x = rng.normal(size=(H, 16)) / 4
    xs = rng.normal(size=(T, 16))
    hs = [np.zeros(H)]
    for t in range(T):
        hs.append(np.tanh(W_h @ hs[-1] + W_x @ xs[t]))
    # 假设损失 L = uᵀ h_T，反向传播并记录每个时间步对 ∂L/∂W_h 的贡献 δ_t h_{t-1}ᵀ 的大小
    dh = rng.normal(size=H)
    contrib = {}
    for t in range(T, 0, -1):
        delta = dh * (1 - hs[t] ** 2)
        contrib[t] = np.linalg.norm(np.outer(delta, hs[t - 1]))
        dh = W_h.T @ delta
    total = sum(contrib.values())
    print(f"\n  RNN（隐藏 {H} 维，W_h 的谱半径 0.9），序列长度 {T}，只看最后一步的损失 L_T：")
    print("  ∂L_T/∂W_h = Σ_t δ_t h_{t-1}ᵀ，各时间步贡献的大小（占总和的比例）：")
    for t in [T, T - 1, T - 2, T - 5, T - 10, T - 20, 1]:
        print(f"    t = {t:>2d}（距离末尾 {T - t:>2d} 步）：{contrib[t]:.2e}（{contrib[t] / total:.2%}）")
    near = sum(contrib[t] for t in range(T - 4, T + 1)) / total
    print(f"\n  最近 5 步的贡献合计 {near:.1%}")
    print("""
  观察：W_h 的梯度几乎完全由最近几个时间步决定。即使"40 步之前的某个词决定了现在该输出什么"，
    这条信息带来的梯度也被淹没了 —— 模型只会根据**近处**的依赖调整参数。
    更糟的是，梯度很小时我们无法区分两种情况：
      (1) 第 t 步和第 T 步之间**真的没有依赖关系**；
      (2) 有依赖关系，但**参数还没学好**，梯度传不过来。
    这就是梯度消失在语言模型上的后果：Manning 课件中 "...she found that the printer was out of toner...
    she bought more ___" 这样的例子，模型很难学会利用很久以前出现的 toner。""")


# =============================================================================
# 演示 3、4：梯度爆炸的悬崖与梯度裁剪
# =============================================================================
def scalar_rnn_loss(w, b=-0.2, T=40, h0=0.5, target=-0.8):
    """
    一个只有一个可训练参数 w 的 RNN：h_t = tanh(w · h_{t-1} + b)，跑 T 步，损失 L = (h_T - target)²。

    【思路】
        w 较小时，h 收敛到一个负的固定点，离 target = -0.8 很近，损失很小；
        w 越过某个临界值后，系统变成"双稳态"，从正的 h0 出发会收敛到**正**的固定点，损失突然变大。
        在临界值附近，h_T 对 w 极其敏感 —— 损失曲面上出现一道"悬崖"，梯度在那里可以大到上千。

    【参数】w : 参数；b、T、h0、target : 固定的设置

    【返回】(L, dL/dw)：损失和它对 w 的导数（用前向模式求导：同时递推 dh_t/dw）

    【Python 小知识】
        前向模式求导：dh_t/dw = (1 - h_t²) · (h_{t-1} + w · dh_{t-1}/dw)，和 h_t 一起逐步计算。
        只有一个参数时，前向模式和反向模式一样高效（第 3 讲第 5.6 节）。
    """
    h, dh = h0, 0.0
    for _ in range(T):
        h_new = np.tanh(w * h + b)
        dh = (1 - h_new ** 2) * (h + w * dh)
        h = h_new
    return (h - target) ** 2, 2 * (h - target) * dh


def demo_cliff():
    """演示 3：损失曲面上的悬崖。"""
    print_header("演示 3：梯度爆炸的'悬崖'（只有一个参数 w 的 RNN）")
    ws = np.linspace(0.5, 3.0, 26)
    print(f"\n    {'w':>5s}{'损失 L(w)':>12s}{'|dL/dw|':>12s}   损失的示意图")
    for w in ws:
        L, g = scalar_rnn_loss(w)
        print(f"    {w:>5.2f}{L:>12.4f}{abs(g):>12.2e}   " + "#" * int(L * 12))
    fine = np.linspace(1.4, 1.6, 2001)
    gs = [abs(scalar_rnn_loss(w)[1]) for w in fine]
    i = int(np.argmax(gs))
    print(f"\n  在 w = {fine[i]:.4f} 附近，|dL/dw| 最大达到 {gs[i]:.2e}。")
    print("""
  观察：大部分地方损失曲面很平缓，梯度很小；但在 w ≈ 1.5 处有一道陡壁，梯度在那里大了几千倍。
    Pascanu 等人（2013）指出，RNN 的损失曲面上经常有这样的悬崖。""")


def gradient_descent(w0, lr, steps, clip=None):
    """
    在 scalar_rnn_loss 上做梯度下降，返回每一步的 (w, L)。

    【参数】
        w0   : 初始值
        lr   : 学习率
        steps: 步数
        clip : None 或阈值。梯度的绝对值超过阈值时，把它缩放到阈值（一维情形下就是截断）。

    【思路】（讲义第 2.4 节）
        梯度裁剪：if ‖g‖ > 阈值: g ← (阈值 / ‖g‖) · g。
        只改变更新的**步长**，不改变**方向**。
    """
    w, path = w0, []
    for _ in range(steps):
        L, g = scalar_rnn_loss(w)
        path.append((w, L))
        if clip is not None and abs(g) > clip:
            g = clip * np.sign(g)
        w = w - lr * g
    path.append((w, scalar_rnn_loss(w)[0]))
    return path


def demo_clipping():
    """演示 4：梯度裁剪。"""
    print_header("演示 4：落在悬崖壁上之后 —— 裁剪 vs 不裁剪（学习率 0.1）")
    print("""
  悬崖壁非常窄（宽度只有千分之几），梯度下降很少恰好落在上面；可一旦落上去，后果很严重。
  所以这里直接从悬崖壁上的 w = 1.498 出发（这里损失还很小，但梯度大约是 1700）：""")
    for name, clip in [("不裁剪", None), ("裁剪阈值 1.0", 1.0)]:
        path = gradient_descent(1.498, lr=0.1, steps=50, clip=clip)
        print(f"\n  {name}：")
        for step in [0, 1, 2, 5, 10, 20, 50]:
            w, L = path[step]
            print(f"    第 {step:>2d} 步：w = {w:>10.4f}，损失 {L:.4f}")
    print("""
  观察：
    - 不裁剪：一步就被巨大的梯度"弹飞"到 w ≈ -171，损失变成 3.24；那里的损失曲面是平的（梯度几乎为 0），
      再也回不来了 —— 之前的进展全部作废。在高维的真实网络里，这表现为损失突然飙升，甚至变成 nan。
    - 裁剪：梯度的方向不变，但步长被限制在 学习率 × 阈值 = 0.1 以内，w 平稳地滑下悬崖，损失继续下降。
    在有上百万个参数的网络里，"某个方向恰好落在悬崖上"的机会要多得多。
    梯度裁剪简单有效，是训练 RNN（以及 Transformer）的标准做法。注意它只能对付梯度**爆炸**，
    对付不了梯度**消失** —— 那需要改变网络结构（LSTM、残差连接，见代码 02 和讲义第 3、4 节）。""")


if __name__ == "__main__":
    demo_matrix_powers()
    demo_gradient_contributions()
    demo_cliff()
    demo_clipping()