"""
autodiff.py —— 从零实现"基于模块的自动微分"（NumPy 版，讲义 §5）
=====================================================================

【这个文件要回答的问题】
    神经网络的梯度是怎样"自动"算出来的？PyTorch 的 loss.backward() 背后在做什么？

【按 slides 的顺序，这里有三种写法】
    A. 过程式（procedural）：nn_forward / nn_backward
       把一个"一层隐藏层 + softmax"的网络从头到尾写成一长串公式（slides 里的 17 行代码）。
       能用，但难以复用、难以调试：梯度检查报错时，你只知道"这 17 行里某处错了"。
    B. 基于模块（module-based）：linear_forward / linear_backward / sigmoid_forward / ...
       每一层是一个"模块"，只负责两件事：
         forward ：给定输入 a，算输出 b = f(a)
         backward：给定"损失对输出的梯度" g_b，算"损失对输入的梯度" g_a（链式法则）
       整个网络的前向 = 依次调用各模块的 forward；反向 = 倒序调用各模块的 backward。
    C. 面向对象 + tape（磁带）：Tape、Leaf、Op 及其子类
       每次调用一个运算，就把它"录"到磁带上；反向时从磁带尾部一个个弹出、调用 backward。
       于是你不必再手写"反向的顺序"——**程序的控制流（for 循环、if）自动决定计算图**。
       这正是 PyTorch 的 define-by-run（边运行边建图）。

【记号（和 slides 一致）】
    g_x 表示"最终的标量损失 J 对 x 的梯度" ∂J/∂x，形状和 x 相同。
    向量默认是一维 NumPy 数组；矩阵 W 的形状是 (输出维度, 输入维度)，所以线性层是 b = W @ a。

【GenAI 联系】
    PyTorch、JAX、TensorFlow 的核心都是一个自动微分引擎。Bottou & Gallinari (1991)
    就提出了"把学习系统拆成能前向、能反向的模块"的框架；Torch（Lua）、PyTorch 的
    nn.Module 一脉相承。理解了这 300 行，你就理解了训练任何大模型时梯度的来源。
"""

import math

import numpy as np


# =====================================================================
# 0. 小工具：sigmoid、稳定的 softmax、梯度检查
# =====================================================================

def sigmoid(x):
    """
    σ(x) = 1 / (1 + e^{-x})，把任意实数压到 (0, 1)。

    【边界】x 很负时 e^{-x} 会溢出成 inf；这里分正负两种情况写，避免溢出警告。
    【Python 小知识】np.where(条件, A, B) 逐元素"条件为真取 A，否则取 B"。
        np.exp(-np.abs(x)) 永远不会溢出（指数部分 ≤ 0）。
    """
    x = np.asarray(x, dtype=np.float64)
    e = np.exp(-np.abs(x))
    return np.where(x >= 0, 1.0 / (1.0 + e), e / (1.0 + e))


def softmax(b):
    """
    softmax(b)_i = e^{b_i} / Σ_j e^{b_j}，把一组实数（logits）变成概率分布。

    【为什么减去最大值】e^{1000} 会溢出；但 softmax(b) = softmax(b - c) 对任何常数 c 成立，
        所以先减去 max(b)，最大的指数变成 e^0 = 1，就不会溢出了。
    【边界】b 必须是一维数组；全部相等时返回均匀分布。
    """
    b = np.asarray(b, dtype=np.float64)
    e = np.exp(b - np.max(b))
    return e / e.sum()


def numerical_gradient(f, x, eps=1e-6):
    """
    用中心差分近似梯度：∂f/∂x_i ≈ (f(x + ε e_i) - f(x - ε e_i)) / (2ε)。

    【思路】一次只动 x 的一个元素，看 f 变了多少。慢（要调用 2·x.size 次 f），但几乎不可能写错，
        所以拿它当"标准答案"来检查手写的 backward。
    【参数】
        f   : 无参数函数，返回标量（它内部会读取 x 的当前值）
        x   : 要求梯度的 NumPy 数组（会被临时修改，函数结束前恢复原值）
        eps : 差分步长。太大 → 截断误差大；太小 → 浮点舍入误差大。float64 下 1e-6 左右最好
    【返回】和 x 同形状的数组
    【边界】必须用 float64；float32 下 eps=1e-6 的差分几乎全是舍入噪声。
    【Python 小知识】np.nditer(..., flags=["multi_index"]) 可以遍历任意形状数组的每个下标。
    """
    grad = np.zeros_like(x, dtype=np.float64)
    it = np.nditer(x, flags=["multi_index"])
    while not it.finished:
        idx = it.multi_index
        old = x[idx]
        x[idx] = old + eps
        f_plus = f()
        x[idx] = old - eps
        f_minus = f()
        x[idx] = old                      # 一定要恢复！否则后面的元素都在错的点上求导
        grad[idx] = (f_plus - f_minus) / (2 * eps)
        it.iternext()
    return grad


def rel_error(a, b):
    """
    相对误差 max|a-b| / max(|a|+|b|)，用来比较两个梯度是否"相同"。

    【怎么读】float64 下：< 1e-7 基本可以认为写对了；> 1e-3 几乎一定写错了。
    【边界】两个都是全 0 时返回 0。
    """
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    denom = np.max(np.abs(a) + np.abs(b))
    return 0.0 if denom == 0 else float(np.max(np.abs(a - b)) / denom)


# =====================================================================
# A. 过程式实现：slides 里的 NNForward / NNBackward
# =====================================================================
#   网络（没有偏置项；需要偏置时可以在 x 末尾补一个常数 1）：
#     a = α x          线性层 1，α 形状 (D, M)
#     z = σ(a)         sigmoid 隐藏层
#     b = β z          线性层 2，β 形状 (K, D)
#     ŷ = softmax(b)   K 个类别的概率
#     J = -yᵀ log ŷ    交叉熵损失，y 是 one-hot 真实标签

def nn_forward(x, y, alpha, beta):
    """
    过程式前向计算。
    【返回】字典 o，存下所有中间量（反向时要用），o["J"] 是损失。
    【为什么要存中间量】反向公式里要用到 z、ŷ 等；不存就得重算一遍前向。
        这就是"训练比推理更费显存"的根本原因：激活值要一直留到反向结束（见前置讲 P1）。
    """
    a = alpha @ x
    z = sigmoid(a)
    b = beta @ z
    y_hat = softmax(b)
    J = -float(y @ np.log(y_hat))
    return {"x": x, "a": a, "z": z, "b": b, "y_hat": y_hat, "J": J}


def nn_backward(x, y, alpha, beta, o):
    """
    过程式反向计算（逐行对应 slides 的 Algorithm 2）。

    每一行都是链式法则：g_上游 = g_下游 × ∂下游/∂上游。
      g_ŷ = -y ÷ ŷ                          （交叉熵对 ŷ 的导数，÷ 是逐元素除）
      g_b = (diag(ŷ) - ŷŷᵀ) g_ŷ              （softmax 的雅可比矩阵是 diag(ŷ) - ŷŷᵀ，对称）
      g_β = g_b zᵀ                           （外积：形状 (K,) × (D,) → (K, D)）
      g_z = βᵀ g_b
      g_a = g_z ⊙ z ⊙ (1 - z)                （σ'(a) = σ(a)(1-σ(a)) = z(1-z)，⊙ 是逐元素乘）
      g_α = g_a xᵀ
    【返回】(g_alpha, g_beta)
    【边界】ŷ 中某个元素恰好为 0 时 -y ÷ ŷ 会除以 0。稳定 softmax 下只有 logits 差距 > 700 才会发生；
        真实框架会把 softmax 和交叉熵合并成一个模块（见下面的 SoftmaxCrossEntropy）来彻底避免。
    """
    z, y_hat = o["z"], o["y_hat"]
    g_yhat = -y / y_hat
    g_b = (np.diag(y_hat) - np.outer(y_hat, y_hat)) @ g_yhat
    g_beta = np.outer(g_b, z)
    g_z = beta.T @ g_b
    g_a = g_z * z * (1 - z)
    g_alpha = np.outer(g_a, x)
    return g_alpha, g_beta


# =====================================================================
# B. 基于模块的实现：每层一对 forward / backward 函数
# =====================================================================
#   约定：xxx_forward(输入...) -> 输出；
#         xxx_backward(输入..., 输出, g_输出) -> 各输入的梯度。
#   "输出"这个参数有的模块用不上，但为了接口统一照样传进去（slides 也是这么写的）。

def linear_forward(a, w):
    """b = w a。a 形状 (A,)，w 形状 (B, A)，b 形状 (B,)。"""
    return w @ a


def linear_backward(a, w, b, g_b):
    """
    【返回】(g_w, g_a)：g_w = g_b aᵀ（外积），g_a = wᵀ g_b
    【为什么 g_w 是外积】b_i = Σ_j w_ij a_j，所以 ∂J/∂w_ij = (∂J/∂b_i)(∂b_i/∂w_ij) = g_b[i] · a[j]。
    """
    return np.outer(g_b, a), w.T @ g_b


def sigmoid_forward(a):
    return sigmoid(a)


def sigmoid_backward(a, b, g_b):
    """g_a = g_b ⊙ b ⊙ (1 - b)。注意这里用的是输出 b = σ(a)，不必重算 σ。"""
    return g_b * b * (1 - b)


def tanh_forward(a):
    return np.tanh(a)


def tanh_backward(a, b, g_b):
    """tanh'(a) = 1 - tanh²(a) = 1 - b²。"""
    return g_b * (1 - b * b)


def softmax_forward(a):
    return softmax(a)


def softmax_backward(a, b, g_b):
    """g_a = (diag(b) - b bᵀ) g_b。"""
    return (np.diag(b) - np.outer(b, b)) @ g_b


def cross_entropy_forward(a, a_hat):
    """b = -aᵀ log â，a 是 one-hot 真实标签，â 是预测分布。输出是标量。"""
    return -float(a @ np.log(a_hat))


def cross_entropy_backward(a, a_hat, b, g_b):
    """g_â = -g_b (a ÷ â)。真实标签 a 是常数，不需要它的梯度。"""
    return -g_b * (a / a_hat)


def nn_forward_modular(x, y, alpha, beta):
    """和 nn_forward 完全相同的网络，用模块拼起来。"""
    a = linear_forward(x, alpha)
    z = sigmoid_forward(a)
    b = linear_forward(z, beta)
    y_hat = softmax_forward(b)
    J = cross_entropy_forward(y, y_hat)
    return {"x": x, "a": a, "z": z, "b": b, "y_hat": y_hat, "J": J}


def nn_backward_modular(x, y, alpha, beta, o, sigmoid_bwd=sigmoid_backward):
    """
    倒序调用各模块的 backward。
    【参数】sigmoid_bwd：可以换成一个"有 bug 的版本"，用来演示逐模块梯度检查怎样定位错误
    【返回】(g_alpha, g_beta)
    """
    g_J = 1.0                                          # 基本情况：dJ/dJ = 1
    g_yhat = cross_entropy_backward(y, o["y_hat"], o["J"], g_J)
    g_b = softmax_backward(o["b"], o["y_hat"], g_yhat)
    g_beta, g_z = linear_backward(o["z"], beta, o["b"], g_b)
    g_a = sigmoid_bwd(o["a"], o["z"], g_z)
    g_alpha, _g_x = linear_backward(x, alpha, o["a"], g_a)   # 输入 x 的梯度用不上，丢掉
    return g_alpha, g_beta


def check_module(forward, backward, inputs, wrt, seed=0, eps=1e-6):
    """
    对单个模块做有限差分检查。

    【思路】随便取一个"上游梯度" g_out（随机向量），定义标量 L = Σ g_out ⊙ forward(inputs)。
        那么 ∂L/∂输入 恰好就是 backward(inputs, out, g_out) 应该返回的东西。
        用 numerical_gradient 算出 ∂L/∂输入，和 backward 的结果比较。
    【参数】
        forward  : 模块的 forward 函数
        backward : 模块的 backward 函数，签名 backward(*inputs, out, g_out)
        inputs   : 输入数组的列表（float64）
        wrt      : 对哪些输入求梯度（下标列表）；backward 返回多个梯度时，顺序要和 wrt 对应
    【返回】每个 wrt 输入的相对误差列表
    【边界】标量输出（如交叉熵）时 g_out 是一个随机标量。
    """
    rng = np.random.default_rng(seed)
    out = forward(*inputs)
    g_out = rng.normal(size=np.shape(out)) if np.ndim(out) > 0 else float(rng.normal())
    analytic = backward(*inputs, out, g_out)
    if not isinstance(analytic, tuple):
        analytic = (analytic,)
    errors = []
    for k, idx in enumerate(wrt):
        x = inputs[idx]

        def L():
            return float(np.sum(g_out * forward(*inputs)))
        num = numerical_gradient(L, x, eps=eps)
        errors.append(rel_error(analytic[k], num))
    return errors


# =====================================================================
# C. 面向对象 + tape：边运行边建图（define-by-run）
# =====================================================================

class Tape:
    """
    磁带：按执行顺序记录每一个运算节点。

    【思路】前向时每执行一个 Op，就 push 到栈顶；于是栈里的顺序就是一个"拓扑序"
        （任何节点都在它的输入之后被压栈）。反向时从栈顶往下弹出，
        就是"逆拓扑序"——处理一个节点时，所有用到它输出的节点都已经处理完了，
        它的 out_gradient 已经累加完整。
    【边界】backward 会把栈清空（和 slides 的 tape_bwd 一样），所以同一条磁带只能反向一次。
        PyTorch 默认也这样：第二次 backward 会报错，除非 retain_graph=True。
    """

    def __init__(self):
        self.stack = []

    def push(self, node):
        self.stack.append(node)

    def backward(self, output):
        """从标量输出 output 开始反向传播：先令 g_output = 1，再逆序调用每个节点的 apply_bwd。"""
        output.out_gradient = np.ones_like(output.out_tensor)
        while self.stack:
            node = self.stack.pop()
            node.apply_bwd()


class Node:
    """计算图里的一个节点：保存前向的输出 out_tensor 和反向累加的梯度 out_gradient。"""

    def __init__(self):
        self.in_nodes = ()
        self.out_tensor = None
        self.out_gradient = None


class Leaf(Node):
    """
    叶子节点：输入数据或参数。它不是任何运算的输出，所以不进磁带。

    【参数】value：NumPy 数组（会转成 float64 并复制一份）；name：名字，方便打印
    【边界】它的 out_gradient 是"累加器"：同一个参数被用了几次，梯度就加几次。
        RNN 在每个时间步都用同一个 W_hh，所以 W_hh 的梯度 = 各时间步贡献之和（见 03 号脚本）。
        这也是为什么每一步训练前要清零梯度（zero_grad），PyTorch 同理。
    """

    def __init__(self, value, name=""):
        super().__init__()
        self.out_tensor = np.array(value, dtype=np.float64)
        self.out_gradient = np.zeros_like(self.out_tensor)
        self.name = name

    def zero_grad(self):
        self.out_gradient = np.zeros_like(self.out_tensor)


class Op(Node):
    """
    运算节点的基类。子类只需实现 forward(*输入数组) 和 backward(输入数组列表, 输出, g_输出)。

    【用法】y = Sigmoid().apply_fwd(tape, x)        # x 是 Node，返回的 y 也是 Node
    【边界】一个 Op 对象只能 apply_fwd 一次（它要记住自己的输入和输出）。
        要把同一种运算用很多次，就每次新建一个对象——下面的 linear()、tanh() 等小函数就是这么做的。
    """

    def apply_fwd(self, tape, *in_nodes):
        self.in_nodes = in_nodes
        in_tensors = [n.out_tensor for n in in_nodes]
        self.out_tensor = self.forward(*in_tensors)
        self.out_gradient = np.zeros_like(self.out_tensor)
        tape.push(self)
        return self

    def apply_bwd(self):
        in_tensors = [n.out_tensor for n in self.in_nodes]
        in_grads = self.backward(in_tensors, self.out_tensor, self.out_gradient)
        for node, g in zip(self.in_nodes, in_grads):
            if g is not None:
                # 累加而不是覆盖：一个节点可能被多个下游节点使用（slides 的 Version B）
                node.out_gradient = node.out_gradient + g

    def forward(self, *in_tensors):
        raise NotImplementedError

    def backward(self, in_tensors, out_tensor, out_gradient):
        raise NotImplementedError


class Linear(Op):
    """b = W a。输入顺序 (a, W)。"""

    def forward(self, a, w):
        return w @ a

    def backward(self, in_tensors, b, g_b):
        a, w = in_tensors
        g_w, g_a = linear_backward(a, w, b, g_b)
        return g_a, g_w


class Add(Op):
    """b = x + y（形状相同）。加法把上游梯度原样分给两个输入。"""

    def forward(self, x, y):
        return x + y

    def backward(self, in_tensors, b, g_b):
        return g_b, g_b


class Sigmoid(Op):
    def forward(self, a):
        return sigmoid(a)

    def backward(self, in_tensors, b, g_b):
        return (sigmoid_backward(in_tensors[0], b, g_b),)


class Tanh(Op):
    def forward(self, a):
        return np.tanh(a)

    def backward(self, in_tensors, b, g_b):
        return (tanh_backward(in_tensors[0], b, g_b),)


class Softmax(Op):
    def forward(self, a):
        return softmax(a)

    def backward(self, in_tensors, b, g_b):
        return (softmax_backward(in_tensors[0], b, g_b),)


class CrossEntropy(Op):
    """b = -yᵀ log ŷ。输入顺序 (y, ŷ)；y 是常数标签，不传梯度（返回 None）。"""

    def forward(self, y, y_hat):
        return np.array(cross_entropy_forward(y, y_hat))

    def backward(self, in_tensors, b, g_b):
        y, y_hat = in_tensors
        return None, cross_entropy_backward(y, y_hat, b, g_b)


class SoftmaxCrossEntropy(Op):
    """
    融合的 softmax + 交叉熵：J = -log softmax(b)[target]。

    【为什么要融合】分开算时梯度要经过 -y ÷ ŷ 再乘 softmax 的雅可比，ŷ 很小时数值很差；
        合起来求导，结果是一个极其简洁的式子：g_b = softmax(b) - onehot(target)。
        既快又稳定。PyTorch 的 F.cross_entropy 也是直接吃 logits 的。
    【参数】target：正确类别的下标（整数，构造时传入，不是图里的节点）
    """

    def __init__(self, target):
        super().__init__()
        self.target = int(target)

    def forward(self, b):
        m = np.max(b)
        log_z = m + np.log(np.sum(np.exp(b - m)))      # log Σ e^{b_j}，同样先减最大值
        return np.array(log_z - b[self.target])

    def backward(self, in_tensors, J, g_J):
        g = softmax(in_tensors[0])
        g[self.target] -= 1.0
        return (g_J * g,)


class Lookup(Op):
    """
    嵌入查表：out = E[index]，E 形状 (V, d)。

    【直觉】它等价于"one-hot 向量 × 矩阵"：onehot(index)ᵀ E 正好取出第 index 行。
        直接取行比真的做乘法快得多，但梯度一样：只有被取出的那一行有梯度。
    """

    def __init__(self, index):
        super().__init__()
        self.index = int(index)

    def forward(self, E):
        return E[self.index].copy()

    def backward(self, in_tensors, out, g_out):
        g_E = np.zeros_like(in_tensors[0])
        g_E[self.index] += g_out
        return (g_E,)


class Sum(Op):
    """把若干个标量节点加起来（用于把每个时间步的损失加成总损失）。"""

    def forward(self, *xs):
        # math.fsum 是精确求和，结果不依赖加法顺序和 Python 版本
        return np.array(math.fsum(float(x) for x in xs))

    def backward(self, in_tensors, b, g_b):
        return tuple(g_b for _ in in_tensors)


# ---- 函数式的小包装：每次调用新建一个 Op 对象（于是同一种运算可以用无数次） ----

def linear(tape, a, w):
    return Linear().apply_fwd(tape, a, w)


def add(tape, x, y):
    return Add().apply_fwd(tape, x, y)


def sigmoid_op(tape, a):
    return Sigmoid().apply_fwd(tape, a)


def tanh_op(tape, a):
    return Tanh().apply_fwd(tape, a)


def softmax_op(tape, a):
    return Softmax().apply_fwd(tape, a)


def cross_entropy(tape, y, y_hat):
    return CrossEntropy().apply_fwd(tape, y, y_hat)


def softmax_cross_entropy(tape, logits, target):
    return SoftmaxCrossEntropy(target).apply_fwd(tape, logits)


def lookup(tape, E, index):
    return Lookup(index).apply_fwd(tape, E)


def total(tape, *xs):
    return Sum().apply_fwd(tape, *xs)


def mlp_tape(tape, x, y, alpha, beta):
    """
    slides 里 NeuralNetwork.forward 的 tape 版：x、y、alpha、beta 都是 Leaf。
    【返回】损失节点 J（调用 tape.backward(J) 后，alpha.out_gradient 就是 g_α）
    """
    a = linear(tape, x, alpha)
    z = sigmoid_op(tape, a)
    b = linear(tape, z, beta)
    y_hat = softmax_op(tape, b)
    return cross_entropy(tape, y, y_hat)


def rnn_lm_tape(tape, ids, E, W_xh, W_hh, b_h, W_hy, b_y, h0):
    """
    把一个 Elman RNN 语言模型在一句话上"展开"，录到磁带上。

        h_t = tanh(W_xh x_t + W_hh h_{t-1} + b_h)
        logits_t = W_hy h_t + b_y
        J = Σ_t -log softmax(logits_t)[ids[t+1]]

    【参数】
        ids  : 整数列表，[<s>, w_1, ..., w_T, </s>]；输入是 ids[:-1]，目标是 ids[1:]
        E、W_xh、W_hh、b_h、W_hy、b_y、h0 : 都是 Leaf
    【返回】(总损失节点, 每个时间步的隐藏状态节点列表)
    【直觉】这里没有任何"反向传播穿越时间（BPTT）"的专门代码——for 循环把同一组参数
        用了 T 次，磁带就记下一张 T 层深的计算图；tape.backward 在这张图上做普通的反向模式，
        这就是 BPTT。
    """
    h = h0
    losses = []
    hs = []
    for t in range(len(ids) - 1):
        x = lookup(tape, E, ids[t])
        pre = add(tape, add(tape, linear(tape, x, W_xh), linear(tape, h, W_hh)), b_h)
        h = tanh_op(tape, pre)
        hs.append(h)
        logits = add(tape, linear(tape, h, W_hy), b_y)
        losses.append(softmax_cross_entropy(tape, logits, ids[t + 1]))
    return total(tape, *losses), hs


# =====================================================================
# D. 测试用的小数据集（03、04 号脚本共用）
# =====================================================================

def make_spiral(seed=0, n_per_class=100):
    """
    造一个二维三分类的"螺旋"数据集（CS231n 的经典例子）：三条旋臂互相缠绕，任何直线都切不开。
    【返回】(X 形状 (N, 3)，最后一列是常数 1 充当偏置；labels 形状 (N,))
    【Python 小知识】np.linspace(a, b, n) 生成从 a 到 b 均匀分布的 n 个数；
        np.stack([...], axis=1) 把两个长度为 n 的数组并排成 (n, 2) 的矩阵。
    """
    rng = np.random.default_rng(seed)
    xs, ys = [], []
    for c in range(3):
        r = np.linspace(0.05, 1.0, n_per_class)                                  # 半径从里到外
        t = np.linspace(c * 4, (c + 1) * 4, n_per_class) + rng.normal(0, 0.2, n_per_class)   # 角度 + 噪声
        xs.append(np.stack([r * np.sin(t), r * np.cos(t)], axis=1))
        ys.append(np.full(n_per_class, c))
    X = np.concatenate(xs)
    X = np.concatenate([X, np.ones((len(X), 1))], axis=1)      # 补常数 1 → 偏置
    return X, np.concatenate(ys)