from node import *


def neural_network(n, classes=2, hiddens=(12,), activation="ReLU"):
    """
    构造一个多层全连接神经网络的计算图。
    """

    # x 是一个 n 维向量变量，不初始化，不参与训练
    x = Variable((n, 1), init=False, trainable=False)

    # 构造全连接层
    input_size = n
    input_vector = x
    for l, h_size in zip(np.arange(len(hiddens)), hiddens):

        output = Add(
            MatMul(
                Variable((h_size, input_size), True),
                input_vector
            ),
            Variable((h_size, 1), True)
        )

        # 隐藏层的输出
        if activation == "ReLU":
            output = ReLU(output)
        elif activation == "Logistic":
            output = Logistic(output)
        else:
            output = output

        input_size = h_size
        input_vector = output

    # 输出层的神经元
    logits = Add(
        MatMul(
            Variable((classes, input_size), True),
            input_vector
        ),
        Variable((classes, 1), True)
    )

    # 返回输入和 logits
    return x, logits