"""
mini_tf.py —— TensorFlow 1 风格的静态计算图：占位符、变量、常量、运算节点、会话，以及"在图上加节点"的符号式自动微分（第 5 讲 §4）
=====================================================================
（对应原课第 5 讲 slides 第 39–53 页，以及课上的代码练习 llmsys_code_examples/mini_tensorflow）

【这个文件要回答的问题】
    slides 讲的框架设计原则是"两阶段、延迟执行"：
        1. 先定义程序：搭一张符号式的计算图（带占位符），此时什么都不算；
        2. 再在可用的设备上执行（优化过的）图：会话 Session 按拓扑序计算需要的节点。
    这和 PyTorch 的"边写边算"（动态图、即时执行）正好相反。这个文件实现前一种，tensor.py 是后一种，
    05 号脚本把两者放在一起比较。

【本文件提供】
    图与节点  : Graph、default_graph()、reset_default_graph()；Placeholder、Variable、Constant、Operation
    运算      : add、multiply、matmul、negative、sigmoid、log、reciprocal、relu、step、reduce_sum、transpose、
                sum_to_shape_of、broadcast_like（以及节点上的 + − × @ 运算符）
    自动微分  : gradients(loss, 变量列表) —— 不是"算出梯度的数值"，而是往图里添加计算梯度的新节点（slides 第 29–30 页）
    执行      : Session.run(fetches, feed_dict)，只计算 fetches 依赖的节点；evaluated 记录每次计算了几个节点
    优化      : GradientDescentOptimizer(lr).minimize(loss) → 一个"训练步"节点；constant_fold(fetch) 预先算掉只依赖常量的子图

【约定】所有值都是 NumPy float64 数组。节点在创建时自动加入默认图。
"""

import numpy as np


# =====================================================================
# 1. 图与节点
# =====================================================================

class Graph:
    """一张计算图：nodes 按创建顺序记录所有节点。"""

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


_DEFAULT = [Graph()]


def default_graph():
    return _DEFAULT[0]


def reset_default_graph():
    _DEFAULT[0] = Graph()
    return _DEFAULT[0]


class Node:
    """
    计算图的节点（slides 第 13 页：每个节点表示一个变量或一个运算，有向边表示运算的输入）。

    【属性】inputs（输入节点）、consumers（使用它的节点）、name、kind（类型名，打印用）
    【运算符】a + b、a − b、a * b、a @ b、−a 都会创建新的运算节点（不会立即计算）。
    """

    kind = "node"

    def __init__(self, inputs=(), name=None):
        self.inputs = list(inputs)
        self.consumers = []
        for i in self.inputs:
            i.consumers.append(self)
        g = default_graph()
        g.nodes.append(self)
        self.name = name or f"{self.kind}_{len(g.nodes)}"

    def __add__(self, b):
        return add(self, _as_node(b))

    __radd__ = __add__

    def __sub__(self, b):
        return add(self, negative(_as_node(b)))

    def __rsub__(self, b):
        return add(_as_node(b), negative(self))

    def __mul__(self, b):
        return multiply(self, _as_node(b))

    __rmul__ = __mul__

    def __neg__(self):
        return negative(self)

    def __matmul__(self, b):
        return matmul(self, b)

    def __repr__(self):
        return f"<{self.kind} {self.name}>"


def _as_node(x):
    return x if isinstance(x, Node) else Constant(np.asarray(x, dtype=np.float64))


class Placeholder(Node):
    """占位符（slides 第 43 页）：输入、标签等，值在执行时由 feed_dict 提供。"""
    kind = "placeholder"

    def __init__(self, shape=None, name=None):
        super().__init__((), name)
        self.shape = shape


class Variable(Node):
    """变量（slides 第 44 页）：有状态的节点，值在多次执行之间保留——通常是参数。"""
    kind = "variable"

    def __init__(self, initial_value, name=None):
        super().__init__((), name)
        self.value = np.array(initial_value, dtype=np.float64)


class Constant(Node):
    """常量：图里的固定数据。"""
    kind = "constant"

    def __init__(self, value, name=None):
        super().__init__((), name)
        self.value = np.array(value, dtype=np.float64)


class Operation(Node):
    """
    运算节点（slides 第 45–46 页）：每个运算要定义 forward（由输入的值算输出的值）
    和 gradient（给定"损失对输出的梯度"这个节点，返回"损失对每个输入的梯度"的节点——注意返回的是节点，不是数值）。
    """
    kind = "op"

    def forward(self, *values):
        raise NotImplementedError

    def gradient(self, g):
        raise NotImplementedError


# =====================================================================
# 2. 运算
# =====================================================================

class Add(Operation):
    kind = "add"

    def forward(self, a, b):
        return a + b

    def gradient(self, g):
        a, b = self.inputs
        return [sum_to_shape_of(g, a), sum_to_shape_of(g, b)]        # 广播过的输入要把梯度求和缩回去


class Multiply(Operation):
    kind = "multiply"

    def forward(self, a, b):
        return a * b

    def gradient(self, g):
        a, b = self.inputs
        return [sum_to_shape_of(multiply(g, b), a), sum_to_shape_of(multiply(g, a), b)]


class MatMul(Operation):
    kind = "matmul"

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

    def gradient(self, g):
        a, b = self.inputs
        return [matmul(g, transpose(b)), matmul(transpose(a), g)]


class Negative(Operation):
    kind = "negative"

    def forward(self, a):
        return -a

    def gradient(self, g):
        return [negative(g)]


class Sigmoid(Operation):
    kind = "sigmoid"

    def forward(self, a):
        return 1.0 / (1.0 + np.exp(-a))

    def gradient(self, g):
        return [multiply(g, multiply(self, add(Constant(1.0), negative(self))))]   # g · σ · (1 − σ)，σ 就是本节点


class Log(Operation):
    kind = "log"

    def forward(self, a):
        return np.log(a)

    def gradient(self, g):
        return [multiply(g, reciprocal(self.inputs[0]))]


class Reciprocal(Operation):
    kind = "reciprocal"

    def forward(self, a):
        return 1.0 / a

    def gradient(self, g):
        return [negative(multiply(g, multiply(self, self)))]                       # d(1/a) = −1/a²


class Relu(Operation):
    kind = "relu"

    def forward(self, a):
        return np.maximum(a, 0.0)

    def gradient(self, g):
        return [multiply(g, step(self.inputs[0]))]


class Step(Operation):
    """x > 0 时为 1，否则为 0（ReLU 的导数；它自己的导数处处为 0）。"""
    kind = "step"

    def forward(self, a):
        return (a > 0).astype(np.float64)

    def gradient(self, g):
        return [multiply(g, Constant(0.0))]


class ReduceSum(Operation):
    kind = "reduce_sum"

    def __init__(self, a, axis=None, name=None):
        super().__init__([a], name)
        self.axis = axis

    def forward(self, a):
        return np.sum(a, axis=self.axis, keepdims=self.axis is not None)

    def gradient(self, g):
        return [broadcast_like(g, self.inputs[0])]


class Transpose(Operation):
    kind = "transpose"

    def forward(self, a):
        return a.T

    def gradient(self, g):
        return [transpose(g)]


class SumToShapeOf(Operation):
    """运行时把 g 求和缩成 like 的形状（处理广播的反向）。第二个输入只用来取形状，不传梯度。"""
    kind = "sum_to_shape"

    def forward(self, g, like):
        like = np.asarray(like)
        while g.ndim > like.ndim:
            g = g.sum(axis=0)
        for k, s in enumerate(like.shape):
            if s == 1 and g.shape[k] != 1:
                g = g.sum(axis=k, keepdims=True)
        return g

    def gradient(self, g):
        return [broadcast_like(g, self.inputs[0]), multiply(g, Constant(0.0))]


class BroadcastLike(Operation):
    """运行时把 g 广播成 like 的形状（reduce_sum 的反向）。"""
    kind = "broadcast_like"

    def forward(self, g, like):
        return np.broadcast_to(g, np.shape(like)).copy()

    def gradient(self, g):
        return [sum_to_shape_of(g, self.inputs[0]), multiply(g, Constant(0.0))]


class AssignSub(Operation):
    """执行时把变量 var 的值减去 delta 的值（优化器的"更新"节点）。这是图里唯一有副作用的节点。"""
    kind = "assign_sub"

    def forward(self, var_value, delta):
        self.inputs[0].value = var_value - delta
        return self.inputs[0].value


class Group(Operation):
    """把若干节点打包成一个：执行它就会执行所有输入（训练步 = 所有参数的更新）。"""
    kind = "group"

    def forward(self, *values):
        return None


def add(a, b):
    return Add([a, b])


def multiply(a, b):
    return Multiply([a, b])


def matmul(a, b):
    return MatMul([a, b])


def negative(a):
    return Negative([a])


def sigmoid(a):
    return Sigmoid([a])


def log(a):
    return Log([a])


def reciprocal(a):
    return Reciprocal([a])


def relu(a):
    return Relu([a])


def step(a):
    return Step([a])


def reduce_sum(a, axis=None):
    return ReduceSum(a, axis)


def transpose(a):
    return Transpose([a])


def sum_to_shape_of(g, like):
    return SumToShapeOf([g, like])


def broadcast_like(g, like):
    return BroadcastLike([g, like])


# =====================================================================
# 3. 拓扑排序、会话、符号式求导
# =====================================================================

def topo_order(fetches):
    """
    只包含 fetches 所依赖的节点、并且每个节点排在它的输入之后的列表（slides 第 51 页：
    对每个运算节点，先递归访问它的输入，再把它加入已访问列表）。这里用显式栈避免递归太深。
    """
    order, seen = [], set()
    stack = [(f, False) for f in reversed(fetches)]
    while stack:
        n, expanded = stack.pop()
        if expanded:
            order.append(n)
            continue
        if id(n) in seen:
            continue
        seen.add(id(n))
        stack.append((n, True))
        for i in reversed(n.inputs):
            if id(i) not in seen:
                stack.append((i, False))
    return order


class Session:
    """
    会话（slides 第 49、52 页）：在一个执行环境里运行计算图。
        run(fetches, feed_dict)：对 fetches 做拓扑排序，按顺序计算每个节点的值——
            占位符从 feed_dict 取；变量、常量用自己的值；运算节点取输入的值再调用 forward。
    【属性】evaluated：上一次 run 计算了几个节点（只算需要的那些）。
    【边界】需要的占位符没有在 feed_dict 里给出 → KeyError。
    """

    def __init__(self):
        self.evaluated = 0

    def run(self, fetches, feed_dict=None):
        single = not isinstance(fetches, (list, tuple))
        fetches = [fetches] if single else list(fetches)
        feed = feed_dict or {}
        values = {}
        order = topo_order(fetches)
        for n in order:
            if isinstance(n, Placeholder):
                if n not in feed:
                    raise KeyError(f"占位符 {n.name} 没有在 feed_dict 里提供值")
                values[id(n)] = np.asarray(feed[n], dtype=np.float64)
            elif isinstance(n, (Variable, Constant)):
                values[id(n)] = n.value
            else:
                values[id(n)] = n.forward(*[values[id(i)] for i in n.inputs])
        self.evaluated = len(order)
        out = [values[id(f)] for f in fetches]
        return out[0] if single else out


def gradients(loss, wrt):
    """
    符号式反向模式自动微分（slides 第 28–30 页"为每个节点构建计算梯度的图"）：
    不计算任何数值，而是往图里添加新的运算节点，这些节点算出来的值就是 d loss / d 变量。

    【思路】和 autodiff.backpropagate 一模一样，只是"梯度"从数值换成了节点：
        1. 对 loss 的祖先做拓扑排序，倒过来处理；
        2. grad_table[节点] 收集它从各个使用者那里得到的梯度节点，多于一个时用 add 节点加起来（扇出求和）；
        3. 对每个运算节点调用 gradient(g)，得到每个输入的梯度节点。
    【返回】和 wrt 一一对应的梯度节点列表；与 loss 无关的变量返回 None。
    """
    order = topo_order([loss])
    grad_table = {id(loss): [Constant(1.0, name="d_loss")]}
    final = {}
    for n in reversed(order):
        contribs = grad_table.get(id(n))
        if not contribs:
            continue
        g = contribs[0]
        for c in contribs[1:]:
            g = add(g, c)
        final[id(n)] = g
        if isinstance(n, Operation):
            for inp, gi in zip(n.inputs, n.gradient(g)):
                if isinstance(inp, (Constant, Placeholder)):
                    continue
                grad_table.setdefault(id(inp), []).append(gi)
    return [final.get(id(w)) for w in wrt]


class GradientDescentOptimizer:
    """
    slides 第 48 页：GradientDescentOptimizer(lr).minimize(loss) 往图里加入"求梯度"和"更新参数"的节点，
    返回一个训练步节点；每次 session.run(train_step, feed) 就做一步梯度下降。
    """

    def __init__(self, lr):
        self.lr = lr

    def minimize(self, loss):
        variables = [n for n in topo_order([loss]) if isinstance(n, Variable)]
        grads = gradients(loss, variables)
        updates = [AssignSub([v, multiply(Constant(self.lr), g)]) for v, g in zip(variables, grads) if g is not None]
        return Group(updates, name="train_step")


def constant_fold(fetch):
    """
    一个最简单的图优化：只依赖常量的运算节点，结果在执行前就能算好，之后每次运行都不用再算。
    【返回】(可以预先算掉的运算节点个数, fetch 依赖的运算节点总数)
    """
    order = topo_order([fetch])
    const = set()
    folded = 0
    for n in order:
        if isinstance(n, Constant):
            const.add(id(n))
        elif isinstance(n, Operation) and n.inputs and all(id(i) in const for i in n.inputs):
            const.add(id(n))
            folded += 1
    return folded, sum(isinstance(n, Operation) for n in order)