from sklearn.metrics import accuracy_score

from node import *
from util import mnist
from optimizer import *
from layer import *

print(__file__)
train_x, train_y, test_x, test_y = mnist('./data/MNIST')
test_x = test_x[:1000]
test_y = test_y[:1000]

# train_x = train_x[:10000]
# train_y = train_y[:10000]

img_shape = (28, 28)

img = Variable(img_shape, init=False, trainable=False)  # 占位符，28x28 的图像
conv1 = conv([img], img_shape, 3, (5, 5), "ReLU")  # 第一卷积层
pooling1 = pooling(conv1, (3, 3), (2, 2))  # 第一池化层

conv2 = conv(pooling1, (14, 14), 3, (3, 3), "ReLU")  # 第二卷积层
pooling2 = pooling(conv2, (3, 3), (2, 2))  # 第二池化层

fc1 = fc(Flatten(*pooling2), 147, 120, "ReLU")  # 第一全连接层
fc2 = fc(fc1, 120, 10, "None")  # 第二全连接层

# 分类概率
prob = SoftMax(fc2)

# 训练标签
label = Variable((10, 1), trainable=False)

# 交叉熵损失
loss = CrossEntropyWithSoftMax(fc2, label)

# Adam 优化器
optimizer = Adam(default_graph, loss, 0.005, batch_size=64)


# 训练
print("start training", flush=True)
for e in range(10):

    for i in range(len(train_x)):
        img.set_value(np.mat(train_x[i, :]).reshape(28, 28))
        label.set_value(np.mat(train_y[i, :]).T)

        # 执行一步优化
        optimizer.one_step()

        # 显示进度
        if i % 100 == 0:
            loss.forward()
            percent = int((i + 1) / len(train_x) * 100)
            print("".join(["="] * percent) + "> loss:{:.6f} {:d}({:.0f}%)".format(loss.value[0, 0], i + 1, percent), flush=True)

        if i > 1 and (i + 1) % 5000 == 0:

            # 在测试集上评估模型正确率
            probs = []
            losses = []
            for j in range(len(test_x)):
                img.set_value(np.mat(test_x[j, :]).reshape(28, 28))
                label.set_value(np.mat(test_y[j, :]).T)

                # 前向传播计算概率
                prob.forward()
                probs.append(prob.value.A1)

                # 计算损失值
                loss.forward()
                losses.append(loss.value[0, 0])

                # print("test instance: {:d}".format(j))

            # 取概率最大的类别为预测类别
            pred = np.argmax(np.array(probs), axis=1)
            truth = np.argmax(test_y, axis=1)
            accuracy = accuracy_score(truth, pred)

            default_graph.draw()
            print("epoch: {:d}, iter: {:d}, loss: {:.3f}, accuracy: {:.2f}%".format(e + 1, i + 1, np.mean(losses), accuracy * 100), flush=True)