"""第 13 章：加法、乘法和 ReLU 组成的标量反向自动微分。"""


class Value:
    def __init__(self, data, parents=()):
        self.data = float(data)
        self.grad = 0.0
        self.parents = parents
        self._backward = lambda: None

    @staticmethod
    def as_value(value):
        return value if isinstance(value, Value) else Value(value)

    def __add__(self, other):
        other = self.as_value(other)
        output = Value(self.data+other.data, (self, other))

        def backward():
            self.grad += output.grad
            other.grad += output.grad

        output._backward = backward
        return output

    __radd__ = __add__

    def __mul__(self, other):
        other = self.as_value(other)
        output = Value(self.data*other.data, (self, other))

        def backward():
            self.grad += other.data*output.grad
            other.grad += self.data*output.grad

        output._backward = backward
        return output

    __rmul__ = __mul__

    def __neg__(self):
        return self*-1

    def __sub__(self, other):
        return self + (-self.as_value(other))

    def __rsub__(self, other):
        return self.as_value(other) + (-self)

    def relu(self):
        output = Value(max(0, self.data), (self,))

        def backward():
            self.grad += (self.data > 0)*output.grad

        output._backward = backward
        return output

    def backward(self):
        order, visited = [], set()

        def visit(node):
            if node not in visited:
                visited.add(node)
                for parent in node.parents:
                    visit(parent)
                order.append(node)

        visit(self)
        # 为了让本例重复调用可预测，每次重置这张图上全部梯度。
        for node in order:
            node.grad = 0.0
        self.grad = 1.0
        for node in reversed(order):
            node._backward()


def network_loss(w, b, v):
    hidden = (w*2+b).relu()
    prediction = v*hidden
    error = prediction-4
    return error*error, prediction


def main():
    w, b, v = Value(0.5), Value(0.5), Value(2)
    loss, prediction = network_loss(w, b, v)
    loss.backward()
    print(f"预测={prediction.data:.3f}，损失={loss.data:.3f}")
    print(f"梯度：dw={w.grad:.3f}，db={b.grad:.3f}，dv={v.grad:.3f}")
    for parameter in [w, b, v]:
        parameter.data -= 0.01*parameter.grad
    updated, prediction = network_loss(w, b, v)
    print(f"更新后预测={prediction.data:.3f}，损失={updated.data:.6f}")
    x = Value(3)
    shared = x*x+x
    shared.backward()
    print(f"x=3 时，x*x+x 对 x 的梯度={x.grad:.1f}")


if __name__ == "__main__":
    main()