"""
autodiff.py —— 自动微分的核心：计算图、拓扑排序、反向传播，以及一个标量版的自动微分引擎（第 5 讲）
=====================================================================
（对应原课 HW2 Problem 1：topological_sort 与 backpropagate；tensor.py 的张量版复用这里的两个函数）

【这个文件要回答的问题】
    loss.backward() 怎样对"任意网络"算出每个参数的梯度？答案分三步：
        1. 前向计算时，每做一次运算就记下"谁由谁算出来、用的是哪个函数"——这就是计算图（有向无环图 DAG）；
        2. 反向时，从输出开始，按"反向拓扑序"处理节点：一个节点的梯度必须等它所有的使用者都处理完才完整；
        3. 每个节点用链式法则把自己的梯度分给它的输入：d输入 = d输出 × ∂输出/∂输入（每个运算只需要知道自己的局部导数）。
    这个文件把这三步写成两个通用函数（只依赖一个很小的"变量协议"），再实现一个标量版的引擎 Scalar 来使用它们。

【本文件提供】
    变量协议（任何类只要有这些成员，就能被 topological_sort / backpropagate 处理）：
        unique_id               : 唯一编号
        is_leaf()               : 是否是叶子（用户创建、需要梯度的变量，例如参数）
        is_constant()           : 是否是常量（不需要梯度，反向时直接跳过）
        parents                 : 算出它的那些输入变量
        chain_rule(d_output)    : 返回 [(输入变量, 该输入的梯度贡献), ...]
        accumulate_derivative(d): 叶子节点把梯度累加到自己身上
    topological_sort(variable)  : 从输出出发的反向拓扑序（输出在最前，每个节点排在它的所有输入之前）
    backpropagate(variable, d)  : 反向传播，把梯度累加到所有叶子上
    backprop_naive(variable, d) : 反面教材——不排序、沿每条路径递归往回传（结果对，但代价可能是指数级的）
    Scalar、ScalarFunction 及其子类 Add、Mul、Neg、Inv、Exp、Log、ReLU、Sigmoid：标量版引擎
    COUNTERS                    : 计数器（前向运算次数、chain_rule 调用次数），用来比较不同算法的代价
"""

import math

COUNTERS = {"forward_ops": 0, "chain_rule_calls": 0}

_next_id = [0]


def _new_id():
    _next_id[0] += 1
    return _next_id[0]


# =====================================================================
# 1. 两个通用函数（HW2 Problem 1）
# =====================================================================

def topological_sort(variable):
    """
    从 variable（通常是损失）出发，返回计算图里所有非常量节点的"反向拓扑序"列表：
    variable 排在最前；任何一个节点都排在它的所有输入（parents）之前。

    【思路】原课 HW2 的提示：后序深度优先搜索（post-order DFS）——先访问完一个节点的所有输入，
        再把这个节点"放到结果的最前面"。这里用显式的栈代替递归：
        Python 默认最多递归约 1000 层，而一个深网络的计算图可以有成千上万层（03 号脚本会建一条 3000 层的链）。
    【边界】常量节点（不需要梯度）不进入结果；同一个节点被多个使用者引用时只出现一次（visited 集合）。
    【复杂度】每个节点、每条边各处理一次：O(节点数 + 边数)。
    """
    order = []
    visited = set()
    stack = [(variable, False)]
    while stack:
        v, expanded = stack.pop()
        if expanded:                          # 它的所有输入都已经处理完：现在把它记下（后序）
            order.append(v)
            continue
        if v.unique_id in visited or v.is_constant():
            continue
        visited.add(v.unique_id)
        stack.append((v, True))               # 先压回自己（标记为"输入处理完之后再记"）
        for p in v.parents:                   # 再压入输入（它们会先被处理）
            if p.unique_id not in visited and not p.is_constant():
                stack.append((p, False))
    order.reverse()                           # 后序是"输入在前"，反过来就是"输出在前"
    return order


def backpropagate(variable, deriv):
    """
    反向传播：从 variable 出发，把 d variable / d 叶子 累加到每个叶子节点上。

    【参数】variable：计算图的输出（例如损失）；deriv：输出本身的梯度（标量损失时是 1）
    【思路】
        1. 用 topological_sort 得到反向拓扑序；
        2. 用一个字典 grads 记录"每个节点目前收到的梯度之和"；
        3. 按顺序处理节点：轮到某个节点时，它的所有使用者都已经处理过了，所以 grads 里它的梯度已经完整；
           叶子：把梯度累加到它身上；中间节点：调用 chain_rule 把梯度分给它的输入，加到它们的 grads 上。
        "一个节点被多个地方使用（扇出），梯度要相加"——slides 第 23 页的公式——就是第 3 步里的加法。
    【边界】常量输入直接跳过；从 variable 出发到不了的节点不会被访问。
    """
    order = topological_sort(variable)
    grads = {variable.unique_id: deriv}
    for v in order:
        d = grads.pop(v.unique_id, None)
        if d is None:
            continue
        if v.is_leaf():
            v.accumulate_derivative(d)
            continue
        for p, dp in v.chain_rule(d):
            if p.is_constant():
                continue
            grads[p.unique_id] = grads[p.unique_id] + dp if p.unique_id in grads else dp


def backprop_naive(variable, deriv):
    """
    反面教材：不做拓扑排序，直接沿每条路径递归往回传，到叶子就累加。

    【结果】对的——链式法则对"所有路径求和"本来就成立。
    【问题】一个节点有 k 个使用者时，它（以及它上游的整个子图）会被处理 k 次；
        层层扇出时，路径数是指数级的（03 号脚本第 3 部分）。拓扑排序保证每个节点只处理一次。
    """
    if variable.is_constant():
        return
    if variable.is_leaf():
        variable.accumulate_derivative(deriv)
        return
    for p, dp in variable.chain_rule(deriv):
        backprop_naive(p, dp)


# =====================================================================
# 2. 标量引擎
# =====================================================================

class Context:
    """前向时保存反向需要的值（例如 Mul 要记住两个输入）。"""

    def __init__(self):
        self.saved_values = ()

    def save_for_backward(self, *values):
        self.saved_values = values


class ScalarHistory:
    """
    一个 Scalar 的"出身"：由哪个函数（last_fn）、用哪些输入（inputs）算出来的，以及前向保存的上下文（ctx）。
    叶子变量的 history 存在但 last_fn 为 None；常量的 history 是 None。
    """

    def __init__(self, last_fn=None, ctx=None, inputs=()):
        self.last_fn = last_fn
        self.ctx = ctx
        self.inputs = inputs


class Scalar:
    """
    一个参与自动微分的实数。

    【创建】Scalar(3.0) 是常量（不求梯度）；Scalar(3.0, requires_grad=True) 是叶子（会累加梯度到 .derivative）。
    【运算】+ - * / 、负号、.exp() .log() .relu() .sigmoid()：每次运算返回一个新的 Scalar，并在它的 history 里记下出身。
    【反向】loss.backward() 之后，每个叶子的 .derivative 就是 d loss / d 叶子。
    """

    def __init__(self, value, history=None, name=None, requires_grad=False):
        self.data = float(value)
        self.history = history if history is not None else (ScalarHistory() if requires_grad else None)
        self.derivative = None
        self.unique_id = _new_id()
        self.name = name or f"s{self.unique_id}"

    # ---- 变量协议 ----
    def is_leaf(self):
        return self.history is not None and self.history.last_fn is None

    def is_constant(self):
        return self.history is None

    @property
    def parents(self):
        return self.history.inputs if self.history is not None else ()

    def chain_rule(self, d_output):
        """调用产生自己的那个函数的 backward，得到对每个输入的梯度。"""
        COUNTERS["chain_rule_calls"] += 1
        h = self.history
        grads = h.last_fn.backward(h.ctx, d_output)
        return list(zip(h.inputs, grads))

    def accumulate_derivative(self, d):
        """叶子的梯度是"累加"的：同一个参数被多次 backward（或在图里出现多次）时，贡献相加。"""
        self.derivative = d if self.derivative is None else self.derivative + d

    def backward(self, d_output=1.0):
        backpropagate(self, d_output)

    def zero_grad(self):
        self.derivative = None

    # ---- 运算符 ----
    def __add__(self, b):
        return Add.apply(self, b)

    __radd__ = __add__

    def __mul__(self, b):
        return Mul.apply(self, b)

    __rmul__ = __mul__

    def __neg__(self):
        return Neg.apply(self)

    def __sub__(self, b):
        return Add.apply(self, Neg.apply(b))

    def __rsub__(self, b):
        return Add.apply(b, Neg.apply(self))

    def __truediv__(self, b):
        return Mul.apply(self, Inv.apply(b))

    def __rtruediv__(self, b):
        return Mul.apply(b, Inv.apply(self))

    def exp(self):
        return Exp.apply(self)

    def log(self):
        return Log.apply(self)

    def relu(self):
        return ReLU.apply(self)

    def sigmoid(self):
        return Sigmoid.apply(self)

    def __repr__(self):
        return f"Scalar({self.data:.6g})"


class ScalarFunction:
    """
    所有标量运算的基类。子类只需要写两个静态方法：
        forward(ctx, *输入的 float)          → 输出的 float（需要时用 ctx.save_for_backward 保存值）
        backward(ctx, d_output)              → 每个输入的梯度组成的元组（局部导数 × d_output）
    apply 负责把输入包装成 Scalar、调用 forward、并在结果上记下 history（构建计算图）。
    """

    @classmethod
    def apply(cls, *vals):
        vals = [v if isinstance(v, Scalar) else Scalar(v) for v in vals]
        ctx = Context()
        out = cls.forward(ctx, *[v.data for v in vals])
        COUNTERS["forward_ops"] += 1
        need_grad = any(not v.is_constant() for v in vals)
        hist = ScalarHistory(cls, ctx, tuple(vals)) if need_grad else None
        return Scalar(out, hist)


class Add(ScalarFunction):
    @staticmethod
    def forward(ctx, a, b):
        return a + b

    @staticmethod
    def backward(ctx, d):
        return d, d                                   # ∂(a+b)/∂a = ∂(a+b)/∂b = 1


class Mul(ScalarFunction):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a * b

    @staticmethod
    def backward(ctx, d):
        a, b = ctx.saved_values
        return d * b, d * a                           # ∂(ab)/∂a = b，∂(ab)/∂b = a


class Neg(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        return -a

    @staticmethod
    def backward(ctx, d):
        return (-d,)


class Inv(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return 1.0 / a

    @staticmethod
    def backward(ctx, d):
        (a,) = ctx.saved_values
        return (-d / (a * a),)                        # d(1/a)/da = −1/a²


class Exp(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        out = math.exp(a)
        ctx.save_for_backward(out)                    # 保存输出：exp 的导数就是它自己
        return out

    @staticmethod
    def backward(ctx, d):
        (out,) = ctx.saved_values
        return (d * out,)


class Log(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return math.log(a)

    @staticmethod
    def backward(ctx, d):
        (a,) = ctx.saved_values
        return (d / a,)


class ReLU(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return a if a > 0 else 0.0

    @staticmethod
    def backward(ctx, d):
        (a,) = ctx.saved_values
        return (d if a > 0 else 0.0,)                 # 在 0 处取 0（次梯度的一个选择，和 PyTorch 一致）


class Sigmoid(ScalarFunction):
    @staticmethod
    def forward(ctx, a):
        out = 1.0 / (1.0 + math.exp(-a)) if a >= 0 else math.exp(a) / (1.0 + math.exp(a))   # 两种写法避免 exp 溢出
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, d):
        (s,) = ctx.saved_values
        return (d * s * (1.0 - s),)                   # σ' = σ(1 − σ)


def central_difference(f, xs, i, h=1e-6):
    """
    数值梯度（slides 第 34 页）：∂f/∂x_i ≈ [f(…, x_i + h, …) − f(…, x_i − h, …)] / (2h)。
    【参数】f：接受一串 float、返回 float 的普通函数；xs：输入列表；i：对第几个输入求导
    【边界】h 太大有截断误差（O(h²)），太小有舍入误差（约 机器精度 / h）——slides 建议用 float64、h = 1e-6。
    """
    up = list(xs)
    dn = list(xs)
    up[i] += h
    dn[i] -= h
    return (f(*up) - f(*dn)) / (2 * h)