class Node(object):
    def __init__(self):
        """
            base for all types of node in the static calculation graph
        """
        self.next_nodes = []    # self.next_nodes中存储了该节点的所有后继节点
        self.data = None        # 代表该节点存储的数据，在第一次前向运算之前，所有的节点的data都是None。
        _default_graph.append(self) # 保存计算图上的所有node
    
    def __neg__(self):          # python 魔术方法，自定义运算符的行为
        return negative(self) 
    def __add__(self, node : Node):
        return add(self, node)
    def __mul__(self, node : Node): 
        return multiply(self, node)  # 乘法运算改为 multiply.__init__ + multiply.compute
    ...

class multiply(Operation):
    def __init__(self, x : Node, y : Node, node_name: str=""):
        super().__init__(input_nodes=[x, y], node_name=node_name)
    
    def compute(self, x_v : np.ndarray, y_v : np.ndarray):
        return x_v * y_v

class Operation(Node):
    def __init__(self, input_nodes : List[Node] = []):
        super().__init__()
        self.input_nodes = input_nodes
        for node in input_nodes:
            node.next_nodes.append(self)
    
    def compute(self, *args): # compute方法是所有具体的算子类必须实现的方法
        pass

class Placeholder(Node):
    def __init__(self):
        super().__init__()

class Variable(Node):
    def __init__(self, init_value : Union[np.ndarray, list] = None):
        super().__init__()
        self.data = init_value

if __name__ == "__main__":
    X = Placeholder()
    W = Variable(np.random.randn(13, 1))
    b = Variable(np.random.randn(1, 1))

    out = X @ W + b
    from pprint import pprint
    pprint(_default_graph)

def Linear(input_dim : int, output_dim : int, bias : bool = True):
    W = Variable(np.random.randn(input_dim, output_dim))
    if bias:
        b = Variable(np.random.randn(1, output_dim))
        return lambda X : X @ W + b
    else:
        return lambda X : X @ W
if __name__ == "__main__":
    X = Placeholder()
    out = Linear(13, 1, bias=True)(X)

    from pprint import pprint
    pprint(_default_graph)

_register_act_functions = Register()
# 在已经实现的激活函数的上面使用装饰器装饰，装饰函数参数填对应的激活函数的名字
@_register_act_functions("sigmoid")
class sigmoid(Operation):
    def __init__(self, x : Node):
        super().__init__(input_nodes=[x])
    
    def compute(self, x_v : np.ndarray):
        return 1 / (1 + np.exp(-1. * x_v))
# 我们的Linear算子修改为一个class，并且在__call__方法中添加一个act参数，用来为线性层增加激活函数
class Linear(object):
    def __init__(self, input_dim : int, output_dim : int, bias : bool = True, act : str = None):
        self.input_dim = input_dim
        self.output_dim = output_dim
        if act and act not in _register_act_functions:
            raise ValueError(f"input activate function '{act}' is not in registered activate function list:{list(_register_act_functions.keys())}")
        self.act = act
        self.W = Variable(np.random.randn(input_dim, output_dim))
        if bias:
            self.b = Variable(np.random.randn(1, output_dim))
    
    def __call__(self, X : Node):
        if not isinstance(X, Node):
            raise ValueError("Linear's parameter X must be a Node!")
        out = X @ self.W + self.b
        if self.act:
            act_func = _register_act_functions[self.act]
            return act_func(out)
        else:
            return out
if __name__ == "__main__":
    X = Placeholder()
    fc = Linear(13, 1, bias=True, act="relu")
    out = fc(X)

    from pprint import pprint
    pprint(_default_graph)

# create session to update nodes' data in the graph
class Session(object):
    def run(self, root_op : Operation, feed_dict : dict = {}):
        for node in _default_graph:
            if isinstance(node, Variable):
                node.data = np.array(node.data)     # 本身的data就是data
            elif isinstance(node, Placeholder):
                node.data = np.array(feed_dict[node])   # 从输入的字典中获取对应的值
            else:
                input_datas = [n.data for n in node.input_nodes]
                node.data = node.compute(*input_datas) # 获取其input_nodes的data，并作为参数送入Operation的compute方法中，compute返回的值就是Operation的data
        return root_op

class negative(Operation):
    def __init__(self, x : Node):
        super().__init__(input_nodes=[x])
    
    def compute(self, x_v : np.ndarray):
        return -1. * x_v
# negative对应的梯度函数，op_node为当前需要更新的节点的后继节点（所有节点的后继节点一定是Operation），grad为目标节点关于op_node节点值的梯度
@_register_grad_functions("negative")
def __negative_gradient(op_node : Operation, grad : np.ndarray):
    return np.array([-1. * grad])

# optimizer
class Optimizer(object):
    def __init__(self, learning_rate : float = 1e-3):
        """
            base for all the optimizer
        """
        self.learning_rate = learning_rate
    # BP算法的实现，大致过程是从目标节点（一般是损失函数节点，后面统称为loss节点）往前的一个BFS。
    def __backwards(self, op_node : Operation):
        """
            do the BP from the op_node, 
            return a gradient dict including op_node's gradients with respect to all the nodes before op_node
        """
        grad_table = {} # 创建一个梯度字典，用来存放各个节点的梯度。
        grad_table[op_node] = 1.    # 首先将loss节点加入梯度字典中，值为1，毕竟自己对自己的梯度肯定是1嘛
        visit_nodes = set()
        queue = Queue()
        visit_nodes.add(op_node)
        queue.put(op_node)

        while not queue.empty():    # 然后开始从loss往前BFS，在BFS的过程中，用当前节点的后继节点来更新该点的梯度值
            cur_node = queue.get()

            if cur_node != op_node:
                grad_table[cur_node] = 0.
                for next_node in cur_node.next_nodes:
                    grad_loss_wrt_next_node = grad_table[next_node]    # loss gradient of next_node
                    next_node_op_name = next_node.__class__.__name__   # next_node must be an Operation, we get its name
                    gradient_func = _register_grad_functions[next_node_op_name]  # get next_node's corresponding gradient function
                    grad_loss_wrt_cur_node = gradient_func(next_node, grad_loss_wrt_next_node) # call the gradient function to get the sub-gradient
                    
                    if len(next_node.input_nodes) == 1: # if next_node represents monocular operators, then add to total gradient directly
                        grad_table[cur_node] += grad_loss_wrt_cur_node
                    else:                               # else get the portion size of gradient
                        cur_node_in_next_node_index = next_node.input_nodes.index(cur_node)
                        grad_table[cur_node] += grad_loss_wrt_cur_node[cur_node_in_next_node_index]

            if isinstance(cur_node, Operation):         # put next op node into queue to do the BFS
                for input_node in cur_node.input_nodes:
                    if input_node not in visit_nodes:   # only add nodes which haven't been updated/visited yet
                        visit_nodes.add(input_node)
                        queue.put(input_node)

        return grad_table
    # 优化算法实现的地方，会涉及到一些最优化的理论
    def minimize(self, loss_node : Operation):
        """
            concrete optimizer method, 
            this method will update parameters before "loss" node(include loss)
        """
        pass
class SGD(Optimizer):   # Stochastic gradient descent 
    def __init__(self, learning_rate : float = 1e-3):
        super().__init__(learning_rate=learning_rate)
    
    def minimize(self, loss_node : Operation):
        lr = self.learning_rate
        grad_table = self._Optimizer__backwards(op_node=loss_node)
        for node in grad_table:
            if isinstance(node, Variable):
                grad = grad_table[node]
                node.data -= lr * grad # 根据当前node 梯度更新当前 Variable node.data
        return grad_table

# define model
net = Net([layer1, layer2, ...])
model = Model(net, loss_fn, optimizer)
# training，将 net、loss、optimizer 一起传给 model，model 实现了 forward、backward 和 apply_grad 三个接口分别对应前向传播、反向传播和参数更新三个功能
pred = model.forward(train_X)
loss, grads = model.backward(pred, train_Y)
model.apply_grad(grads)
# inference
test_pred = model.forward(test_X)

class Layer(object):
  def __init__(self, name):
      self.name = name
      self.params, self.grads = None, None
  def forward(self, inputs):
      raise NotImplementedError
  def backward(self, grad):
      raise NotImplementedError

class Dense(Layer):
  def __init__(self, num_in, num_out,w_init=XavierUniformInit(),b_init=ZerosInit()):
      super().__init__("Linear")
      self.params = {
          "w": w_init([num_in, num_out]),
          "b": b_init([1, num_out])}
      self.inputs = None
  # forward 方法接收上层的输入 inputs，实现  的运算
  def forward(self, inputs):
      self.inputs = inputs
      return inputs @ self.params["w"] + self.params["b"]
  # backward 的方法接收来自上层的梯度，计算关于参数  和输入的梯度，然后返回关于输入的梯度
  def backward(self, grad):
      self.grads["w"] = self.inputs.T @ grad
      self.grads["b"] = np.sum(grad, axis=0)
      return grad @ self.params["w"].T

class Activation(Layer):
  """Base activation layer"""
  def __init__(self, name):
      super().__init__(name)
      self.inputs = None
  def forward(self, inputs):
      self.inputs = inputs
      return self.func(inputs)
  def backward(self, grad):
      return self.derivative_func(self.inputs) * grad
  def func(self, x):
      raise NotImplementedError
  def derivative_func(self, x):
      raise NotImplementedError

class Net(object):
  def __init__(self, layers):
      self.layers = layers
  # 按顺序遍历所有层，每层计算的输出作为下一层的输入
  def forward(self, inputs):
      for layer in self.layers:
          inputs = layer.forward(inputs)
      return inputs
  # 逆序遍历所有层，将每层的梯度作为下一层的输入
  def backward(self, grad):
      all_grads = [] # 将每个网络层参数的梯度保存下来返回，后面参数更新需要用到
      for layer in reversed(self.layers):
          grad = layer.backward(grad)
          all_grads.append(layer.grads)
      return all_grads[::-1]

  def get_params_and_grads(self):
      for layer in self.layers:
          yield layer.params, layer.grads
  def get_parameters(self):
      return [layer.params for layer in self.layers]
  def set_parameters(self, params):
      for i, layer in enumerate(self.layers):
          for key in layer.params.keys():
              layer.params[key] = params[i][key]

class BaseLoss(object):
    # 计算损失值
    def loss(self, predicted, actual):
        raise NotImplementedError
    # 计算损失值和关于预测值的梯度
    def grad(self, predicted, actual):
        raise NotImplementedError

class BaseOptimizer(object):
    def __init__(self, lr, weight_decay):
        self.lr = lr
        self.weight_decay = weight_decay
    def compute_step(self, grads, params):
        step = list()
        # flatten all gradients
        flatten_grads = np.concatenate([np.ravel(v) for grad in grads for v in grad.values()])
        # compute step
        flatten_step = self._compute_step(flatten_grads)
        # reshape gradients
        p = 0
        for param in params:
            layer = dict()
            for k, v in param.items():
                block = np.prod(v.shape)
                _step = flatten_step[p:p+block].reshape(v.shape)
                _step -= self.weight_decay * v
                layer[k] = _step
                p += block
            step.append(layer)
        return step

    def _compute_step(self, grad):
        raise NotImplementedError

class Model(object):
  def __init__(self, net, loss, optimizer):
      self.net = net
      self.loss = loss
      self.optimizer = optimizer

  def forward(self, inputs):
      return self.net.forward(inputs)

  def backward(self, preds, targets):
      loss = self.loss.loss(preds, targets)
      grad = self.loss.grad(preds, targets)
      grads = self.net.backward(grad)
      params = self.net.get_parameters()
      step = self.optimizer.compute_step(grads, params)
      return loss, step

  def apply_grad(self, grads):
    for grad, (param, _) in zip(grads, self.net.get_params_and_grads()):
      for k, v in param.items():
          param[k] += grad[k]